[TDF] Add the TTrivialDS and its unit tests

The TTrivialDS is a very simple TDataSource, de facto a generator.
It is useful as example and for testing the behaviour of the TDF in combination with
external data sources.

Author: dpiparo

tree/treeplayer/inc/ROOT/TTrivialDS.hxx

@@ -0,0 +1,76 @@
+#ifndef ROOT_TTRIVIALTDS
+#define ROOT_TTRIVIALTDS
+
+#include "ROOT/TDataFrame.hxx"
+#include "ROOT/TDataSource.hxx"
+#include "ROOT/TSeq.hxx"
+
+namespace ROOT {
+namespace Experimental {
+namespace TDF {
+
+class TTrivialDS final : public ROOT::Experimental::TDF::TDataSource {
+private:
+   unsigned int fNSlots = 0U;
+   ULong64_t fSize = 0ULL;
+   std::vector<std::pair<ULong64_t, ULong64_t>> fEntryRanges;
+   std::vector<std::string> fColNames{"col0"};
+   std::vector<ULong64_t> fCounter;
+   std::vector<ULong64_t *> fCounterAddr;
+   std::vector<void *> GetColumnReadersImpl(std::string_view name, const std::type_info &)
+   {
+      std::vector<void *> ret;
+      for (auto i : ROOT::TSeqU(fNSlots)) {
+         fCounterAddr[i] = &fCounter[i];
+         ret.emplace_back((void *)(&fCounterAddr[i]));
+      }
+      return ret;
+   }
+
+public:
+   TTrivialDS(ULong64_t size) : fSize(size) {}
+
+   ~TTrivialDS() {}
+
+   const std::vector<std::string> &GetColumnNames() const { return fColNames; }
+
+   bool HasColumn(std::string_view colName) const { return colName == fColNames[0]; }
+
+   std::string GetTypeName(std::string_view) const { return "ULong64_t"; }
+
+   const std::vector<std::pair<ULong64_t, ULong64_t>> &GetEntryRanges() const
+   {
+      if (fEntryRanges.empty()) {
+         throw std::runtime_error("No ranges are available. Did you set the number of slots?");
+      }
+      return fEntryRanges;
+   }
+   void SetEntry(ULong64_t entry, unsigned int slot) { fCounter[slot] = entry; }
+
+   void SetNSlots(unsigned int nSlots)
+   {
+      if (0U != fNSlots) {
+         return;
+      }
+
+      fNSlots = nSlots;
+      fCounter.resize(fNSlots);
+      fCounterAddr.resize(fNSlots);
+
+      auto chunckSize = fSize / fNSlots;
+      auto start = 0UL;
+      auto end = 0UL;
+      for (auto i : ROOT::TSeqUL(fNSlots)) {
+         start = end;
+         end += chunckSize;
+         fEntryRanges.emplace_back(start, end);
+      }
+      // TODO: redistribute reminder to all slots
+      fEntryRanges.back().second += fSize % fNSlots;
+   };
+};
+}
+}
+}
+
+#endif

tree/treeplayer/test/CMakeLists.txt

@@ -7,3 +7,4 @@ ROOT_ADD_GTEST(dataframe_regression dataframe/dataframe_regression.cxx LIBRARIES
 ROOT_ADD_GTEST(dataframe_interface dataframe/dataframe_interface.cxx LIBRARIES TreePlayer)
 ROOT_ADD_GTEST(dataframe_utils dataframe/dataframe_utils.cxx LIBRARIES TreePlayer)
 ROOT_ADD_GTEST(dataframe_nodes dataframe/dataframe_nodes.cxx LIBRARIES TreePlayer)
+ROOT_ADD_GTEST(datasource_trivial dataframe/datasource_trivial.cxx LIBRARIES TreePlayer)

tree/treeplayer/test/dataframe/datasource_trivial.cxx

@@ -0,0 +1,69 @@
+#include "ROOT/TTrivialDS.hxx"
+
+#include "gtest/gtest.h"
+
+using namespace ROOT::Experimental;
+using namespace ROOT::Experimental::TDF;
+
+TEST(TTrivialDS, ColTypeNames)
+{
+   TTrivialDS tds(32);
+   tds.SetNSlots(1);
+
+   auto colName = tds.GetColumnNames()[0]; // We know it's one.
+   EXPECT_STREQ("col0", colName.c_str());
+   EXPECT_STREQ("ULong64_t", tds.GetTypeName("col0").c_str());
+
+   EXPECT_TRUE(tds.HasColumn("col0"));
+   EXPECT_FALSE(tds.HasColumn("col1"));
+}
+
+TEST(TTrivialDS, EntryRanges)
+{
+   TTrivialDS tds(32);
+   const auto nSlots = 4U;
+   tds.SetNSlots(nSlots);
+
+   auto ranges = tds.GetEntryRanges();
+
+   EXPECT_EQ(4U, ranges.size());
+   EXPECT_EQ(0U, ranges[0].first);
+   EXPECT_EQ(8U, ranges[0].second);
+   EXPECT_EQ(8U, ranges[1].first);
+   EXPECT_EQ(16U, ranges[1].second);
+   EXPECT_EQ(16U, ranges[2].first);
+   EXPECT_EQ(24U, ranges[2].second);
+   EXPECT_EQ(24U, ranges[3].first);
+   EXPECT_EQ(32U, ranges[3].second);
+}
+
+TEST(TTrivialDS, ColumnReaders)
+{
+   TTrivialDS tds(32);
+   const auto nSlots = 4U;
+   tds.SetNSlots(nSlots);
+   auto vals = tds.GetColumnReaders<ULong64_t>("col0");
+   auto ranges = tds.GetEntryRanges();
+   auto slot = 0U;
+   for (auto &&range : ranges) {
+      for (auto i : ROOT::TSeq<ULong64_t>(range.first, range.second)) {
+      tds.SetEntry(i, slot);
+      auto val = **vals[slot];
+      EXPECT_EQ(i, val);
+      }
+      slot++;
+  }
+}
+
+TEST(TTrivialDS, FromATDF)
+{
+  std::unique_ptr<TDataSource> tds(new TTrivialDS(32));
+  TDataFrame tdf(std::move(tds));
+  auto max = tdf.Max<ULong64_t>("col0");
+  auto min = tdf.Min<ULong64_t>("col0");
+  auto c = tdf.Count();
+
+  EXPECT_EQ(32U, *c);
+  EXPECT_DOUBLE_EQ(31., *max);
+  EXPECT_DOUBLE_EQ(0., *min);
+}