RCombinedDS.h

Go to the documentation of this file.

// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
// All rights not expressly granted are reserved.
//
// This software is distributed under the terms of the GNU General Public
// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
//
// In applying this license CERN does not waive the privileges and immunities
// granted to it by virtue of its status as an Intergovernmental Organization
// or submit itself to any jurisdiction.
#ifndef ROOT_RCOMBINEDDS
#define ROOT_RCOMBINEDDS
 
#include "ROOT/RDataFrame.hxx"
#include "ROOT/RDataSource.hxx"
 
#include <memory>
#include <string>
#include <vector>
 
namespace ROOT
{
 
namespace RDF
{
 
class RCombinedDSIndex
{
 public:
  virtual ~RCombinedDSIndex() = default;
  virtual std::vector<std::pair<ULong64_t, ULong64_t>> BuildIndex(std::unique_ptr<RDataFrame>& left,
                                                                  std::unique_ptr<RDataFrame>& right) = 0;
  virtual std::pair<ULong64_t, ULong64_t> GetAssociatedEntries(ULong64_t entry) = 0;
};
 
class RCombinedDSFriendIndex : public RCombinedDSIndex
{
 public:
  std::pair<ULong64_t, ULong64_t> GetAssociatedEntries(ULong64_t entry) final
  {
    return std::pair<ULong64_t, ULong64_t>(entry, entry);
  }
  std::vector<std::pair<ULong64_t, ULong64_t>> BuildIndex(std::unique_ptr<RDataFrame>& left,
                                                          std::unique_ptr<RDataFrame>& right) final;
};
 
class RCombinedDSCrossJoinIndex : public RCombinedDSIndex
{
 public:
  std::pair<ULong64_t, ULong64_t> GetAssociatedEntries(ULong64_t entry) final
  {
    return std::make_pair<ULong64_t, ULong64_t>(entry / fRightCount, entry % fRightCount);
  }
  std::vector<std::pair<ULong64_t, ULong64_t>> BuildIndex(std::unique_ptr<RDataFrame>& left,
                                                          std::unique_ptr<RDataFrame>& right) final;
 
 private:
  ULong64_t fLeftCount;
  ULong64_t fRightCount;
};
 
template <typename INDEX_TYPE = int>
class RCombinedDSColumnJoinIndex : public RCombinedDSIndex
{
 public:
  RCombinedDSColumnJoinIndex(std::string const& indexColumnName)
    : fIndexColumnName{indexColumnName}
  {
  }
 
  std::pair<ULong64_t, ULong64_t> GetAssociatedEntries(ULong64_t entry) final
  {
    auto left = fAssociations[entry];
    return std::pair<ULong64_t, ULong64_t>(left, entry);
  }
 
  std::vector<std::pair<ULong64_t, ULong64_t>>
    BuildIndex(std::unique_ptr<RDataFrame>& left,
               std::unique_ptr<RDataFrame>& right) final
  {
    std::vector<std::pair<ULong64_t, ULong64_t>> ranges;
    auto nEntries = *right->Count();
    fAssociations.reserve(nEntries);
    // Fill the index with the associations
    auto filler = [&assoc = fAssociations](INDEX_TYPE ri) { assoc.push_back(ri); };
    right->Foreach(filler, std::vector<std::string>{fIndexColumnName});
 
    // Create the ranges by processing 64 entries per range
    auto deltaRange = 64;
    ranges.reserve(nEntries / deltaRange + 1);
    ULong64_t i = 0;
    while (deltaRange * (i + 1) < nEntries) {
      ranges.emplace_back(std::pair<ULong64_t, ULong64_t>(deltaRange * i, deltaRange * (i + 1)));
    }
    ranges.emplace_back(std::pair<ULong64_t, ULong64_t>(deltaRange * i, nEntries)); // Last entry
    return ranges;
  }
 
 private:
  std::string fIndexColumnName;
  std::vector<INDEX_TYPE> fAssociations;
};
 
enum struct BlockCombinationRule {
  Full,
  Upper,
  StrictlyUpper,
  Diagonal,
  Anti
};
 
struct RCombinedDSIndexHelpers {
  static char const* combinationRuleAsString(BlockCombinationRule ruleType);
};
 
template <typename INDEX_TYPE = int>
class RCombinedDSBlockJoinIndex : public RCombinedDSIndex
{
 
  using Association = std::pair<ULong64_t, ULong64_t>;
 
 public:
  RCombinedDSBlockJoinIndex(std::string const& leftCategoryColumn,
                            bool self = true,
                            BlockCombinationRule combinationType = BlockCombinationRule::Anti,
                            std::string const& rightCategoryColumn = "")
    : fLeftCategoryColumn{leftCategoryColumn},
      fRightCategoryColumn{rightCategoryColumn.empty() ? leftCategoryColumn : rightCategoryColumn},
      fSelf{self},
      fCombinationType{combinationType}
  {
  }
 
  std::pair<ULong64_t, ULong64_t> GetAssociatedEntries(ULong64_t entry) final
  {
    return fAssociations[entry];
  }
 
  std::vector<std::pair<ULong64_t, ULong64_t>>
    BuildIndex(std::unique_ptr<RDataFrame>& left,
               std::unique_ptr<RDataFrame>& right) final
  {
    std::vector<std::pair<ULong64_t, ULong64_t>> ranges;
    std::vector<INDEX_TYPE> leftCategories;
    std::vector<INDEX_TYPE> rightCategories;
    std::vector<Association> leftPairs;
    std::vector<Association> rightPairs;
 
    computePairsAndCategories(left, leftCategories, leftPairs, fLeftCategoryColumn);
    if (fSelf) {
      rightCategories = leftCategories;
      rightPairs = leftPairs;
    } else {
      computePairsAndCategories(right, rightCategories, rightPairs, fRightCategoryColumn);
    }
 
    auto same = [](std::pair<ULong64_t, ULong64_t> const& a, std::pair<ULong64_t, ULong64_t> const& b) {
      return a.first < b.first;
    };
 
    int startSize = fAssociations.size();
    for (auto categoryValue : leftCategories) {
      std::pair<ULong64_t, ULong64_t> p{categoryValue, 0};
      auto outerRange = std::equal_range(leftPairs.begin(), leftPairs.end(), p, same);
      decltype(outerRange) innerRange;
      if (fSelf) {
        innerRange = outerRange;
      } else {
        innerRange = std::equal_range(rightPairs.begin(), rightPairs.end(), p, same);
      }
      int offset = 0;
      switch (fCombinationType) {
        case BlockCombinationRule::Full:
          for (auto out = outerRange.first; out != outerRange.second; ++out) {
            for (auto in = innerRange.first; in != innerRange.second; ++in) {
              fAssociations.emplace_back(Association{out->second, in->second});
            }
          }
          break;
        case BlockCombinationRule::Upper:
          offset = 0;
          for (auto out = outerRange.first; out != outerRange.second; ++out) {
            if (innerRange.first == innerRange.second) {
              break;
            }
            for (auto in = innerRange.first + offset; in != innerRange.second; ++in) {
              fAssociations.emplace_back(Association{out->second, in->second});
            }
            offset++;
          }
          break;
        case BlockCombinationRule::StrictlyUpper:
          offset = 1;
          for (auto out = outerRange.first; out != outerRange.second; ++out) {
            if (innerRange.first == innerRange.second || innerRange.first + 1 == innerRange.second) {
              break;
            }
            for (auto in = innerRange.first + offset; in != innerRange.second; ++in) {
              fAssociations.emplace_back(Association{out->second, in->second});
            }
            offset++;
          }
          break;
        case BlockCombinationRule::Anti:
          for (auto out = outerRange.first; out != outerRange.second; ++out) {
            for (auto in = innerRange.first; in != innerRange.second; ++in) {
              if (std::distance(innerRange.first, in) == std::distance(outerRange.first, out)) {
                continue;
              }
              fAssociations.emplace_back(Association{out->second, in->second});
            }
            offset++;
          }
          break;
        case BlockCombinationRule::Diagonal:
          auto sizeRow = std::distance(outerRange.first, outerRange.second);
          auto sizeCol = std::distance(innerRange.first, innerRange.second);
          for (size_t i = 0, e = std::min(sizeRow, sizeCol); i < e; ++i) {
            fAssociations.emplace_back(Association{(outerRange.first + i)->second, (innerRange.first + i)->second});
          }
          break;
      }
      auto rangeFirst = startSize;
      auto rangeSecond = fAssociations.size();
      startSize = fAssociations.size();
      ranges.emplace_back(std::make_pair<ULong64_t, ULong64_t>(rangeFirst, rangeSecond));
    }
    return ranges;
  }
 
 private:
  std::string fLeftCategoryColumn;
  std::string fRightCategoryColumn;
  bool fSelf;
  BlockCombinationRule fCombinationType;
  std::vector<Association> fAssociations;
  void computePairsAndCategories(std::unique_ptr<RDataFrame>& df,
                                 std::vector<INDEX_TYPE>& categories,
                                 std::vector<Association>& pairs,
                                 std::string const& column)
  {
    categories = *df->template Take<INDEX_TYPE>(column);
    pairs.reserve(categories.size());
    // Fill the pairs according tho the actual category
    for (size_t i = 0; i < categories.size(); ++i) {
      pairs.emplace_back(categories[i], i);
    }
    // Do a stable sort so that same categories entries are
    // grouped together.
    std::stable_sort(pairs.begin(), pairs.end());
    // Keep only the categories.
    std::stable_sort(categories.begin(), categories.end());
    auto last = std::unique(categories.begin(), categories.end());
    categories.erase(last, categories.end());
  }
};
 
class RCombinedDS final : public ROOT::RDF::RDataSource
{
 private:
  RDataSource* fLeft;
  RDataSource* fRight;
  std::string fLeftPrefix;
  std::string fRightPrefix;
  std::unique_ptr<RDataFrame> fLeftDF;
  std::unique_ptr<RDataFrame> fRightDF;
  ULong64_t fLeftCount;
  ULong64_t fRightCount;
  size_t fNSlots = 0U;
  std::vector<std::string> fColumnNames;
  std::vector<std::pair<ULong64_t, ULong64_t>> fEntryRanges;
  std::unique_ptr<RCombinedDSIndex> fIndex;
 
 protected:
  std::vector<void*> GetColumnReadersImpl(std::string_view colName, const std::type_info& info) override;
 
 public:
  RCombinedDS(std::unique_ptr<RDataSource> left,
              std::unique_ptr<RDataSource> right,
              std::unique_ptr<RCombinedDSIndex> index = std::make_unique<RCombinedDSFriendIndex>(),
              std::string leftPrefix = std::string{"left_"},
              std::string rightPrefix = std::string{"right_"});
  ~RCombinedDS() override;
 
  template <typename T>
  std::vector<T**> GetColumnReaders(std::string_view colName)
  {
    if (colName.compare(0, fLeftPrefix.size(), fLeftPrefix)) {
      colName.remove_prefix(fLeftPrefix.size());
      return fLeft->GetColumnReaders<T>(colName);
    }
    if (colName.compare(0, fRightPrefix.size(), fRightPrefix)) {
      colName.remove_prefix(fRightPrefix.size());
      return fRight->GetColumnReaders<T>(colName);
    }
    std::string dummy("Column not found: ");
    dummy += colName.data();
    throw std::runtime_error(dummy);
  }
  [[nodiscard]] const std::vector<std::string>& GetColumnNames() const override;
  std::vector<std::pair<ULong64_t, ULong64_t>> GetEntryRanges() override;
  [[nodiscard]] std::string GetTypeName(std::string_view colName) const override;
  [[nodiscard]] bool HasColumn(std::string_view colName) const override;
  bool SetEntry(unsigned int slot, ULong64_t entry) override;
  void InitSlot(unsigned int slot, ULong64_t firstEntry) override;
  void SetNSlots(unsigned int nSlots) override;
  void Initialize() override;
};
 
RDataFrame MakeCombinedDataFrame(std::unique_ptr<RDataSource> left, std::unique_ptr<RDataSource>, std::unique_ptr<RCombinedDSIndex> index, std::string leftPrefix = "left_", std::string rightPrefix = "right_");
RDataFrame MakeCrossProductDataFrame(std::unique_ptr<RDataSource> left, std::unique_ptr<RDataSource>, std::string leftPrefix = "left_", std::string rightPrefix = "right_");
RDataFrame MakeColumnIndexedDataFrame(std::unique_ptr<RDataSource> left, std::unique_ptr<RDataSource>, std::string indexColName, std::string leftPrefix = "left_", std::string rightPrefix = "right_");
RDataFrame MakeFriendDataFrame(std::unique_ptr<RDataSource> left, std::unique_ptr<RDataSource> right, std::string leftPrefix = "left_", std::string rightPrefix = "right_");
RDataFrame MakeBlockAntiDataFrame(std::unique_ptr<RDataSource> left, std::unique_ptr<RDataSource> right, std::string indexColumnName, std::string leftPrefix = "left_", std::string rightPrefix = "right_");
 
} // namespace RDF
 
} // namespace ROOT
 
#endif