[RDF] Add tests for GetClusterRanges

Author: siliataider

tree/dataframe/test/dataframe_utils.cxx

@@ -1,5 +1,6 @@
 #include "ROOT/RDataFrame.hxx"
 #include "ROOT/RDF/Utils.hxx"
+#include "ROOT/RNTupleDS.hxx"
 #include "TTree.h"
 
 #include "gtest/gtest.h"
@@ -8,6 +9,7 @@
 #include <vector>
 
 #include <TBranch.h>
+#include <TSystem.h>
 
 namespace RDFInt = ROOT::Internal::RDF;
 
@@ -237,3 +239,85 @@ TEST(RDataFrameUtils, TypeName2TypeID)
    EXPECT_THROW(RDFInt::TypeName2TypeID("float *"), std::runtime_error);
    EXPECT_THROW(RDFInt::TypeName2TypeID("float &"), std::runtime_error);
 }
+
+TEST(RDataFrameUtils, GetClusterRanges)
+{
+   const std::string fnameTTree1 = "dataframe_interfaceAndUtils_2_tree_1.root";
+   const std::string fnameTTree2 = "dataframe_interfaceAndUtils_2_tree_2.root";
+   const std::string fnameRNTuple1 = "dataframe_interfaceAndUtils_2_rntuple_1.root";
+   const std::string fnameRNTuple2 = "dataframe_interfaceAndUtils_2_rntuple_2.root";
+   const int nEntries = 1000;
+   const int clusterSize = 250;
+
+   // Test TTree single file
+   {
+      ROOT::RDF::RSnapshotOptions opts;
+      opts.fAutoFlush = clusterSize;
+      auto df = ROOT::RDataFrame(nEntries).Define("i", "rdfentry_");
+      df.Snapshot("t", fnameTTree1, {"i"}, opts);
+      df.Snapshot("t", fnameTTree2, {"i"}, opts);
+   }
+   {
+      ROOT::RDataFrame dfTTree("t", fnameTTree1);
+      auto *lmTTree = dfTTree.GetLoopManager();
+      auto rangesTTree = RDFInt::GetClusterRanges(*lmTTree);
+
+      EXPECT_EQ(rangesTTree.size(), nEntries / clusterSize);
+      for (size_t i = 0; i < rangesTTree.size(); ++i) {
+         EXPECT_EQ(rangesTTree[i].first, i * clusterSize);
+         EXPECT_EQ(rangesTTree[i].second, (i + 1) * clusterSize);
+      }
+   }
+
+   // Test TTree multiple files (chain)
+   {
+      ROOT::RDataFrame dfTTreeChain("t", {fnameTTree1, fnameTTree2});
+      auto *lmTTreeChain = dfTTreeChain.GetLoopManager();
+      auto rangesTTreeChain = RDFInt::GetClusterRanges(*lmTTreeChain);
+
+      EXPECT_EQ(rangesTTreeChain.size(), 2 * nEntries / clusterSize);
+      for (size_t i = 0; i < rangesTTreeChain.size(); ++i) {
+         EXPECT_EQ(rangesTTreeChain[i].first, i * clusterSize);
+         EXPECT_EQ(rangesTTreeChain[i].second, (i + 1) * clusterSize);
+      }
+   }
+   gSystem->Unlink(fnameTTree1.c_str());
+   gSystem->Unlink(fnameTTree2.c_str());
+
+   // Test RNTuple single file
+   {
+      ROOT::RDF::RSnapshotOptions opts;
+      opts.fOutputFormat = ROOT::RDF::ESnapshotOutputFormat::kRNTuple;
+      auto df = ROOT::RDataFrame(nEntries).Define("i", "rdfentry_");
+      df.Snapshot("nt", fnameRNTuple1, {"i"}, opts);
+      df.Snapshot("nt", fnameRNTuple2, {"i"}, opts);
+   }
+   {
+      auto dfRNTuple = ROOT::RDF::FromRNTuple("nt", fnameRNTuple1);
+      auto *lmRNTuple = dfRNTuple.GetLoopManager();
+      auto rangesRNTuple = RDFInt::GetClusterRanges(*lmRNTuple);
+
+      EXPECT_FALSE(rangesRNTuple.empty());
+      EXPECT_EQ(rangesRNTuple.front().first, 0u);
+      EXPECT_EQ(rangesRNTuple.back().second, ULong64_t(nEntries));
+      for (size_t i = 1; i < rangesRNTuple.size(); ++i) {
+         EXPECT_EQ(rangesRNTuple[i].first, rangesRNTuple[i - 1].second);
+      }
+   }
+
+   // Test RNTuple multiple files
+   {
+      auto dfRNTupleChain = ROOT::RDF::FromRNTuple("nt", {fnameRNTuple1, fnameRNTuple2});
+      auto *lmRNTupleChain = dfRNTupleChain.GetLoopManager();
+      auto rangesRNTupleChain = RDFInt::GetClusterRanges(*lmRNTupleChain);
+
+      EXPECT_FALSE(rangesRNTupleChain.empty());
+      EXPECT_EQ(rangesRNTupleChain.front().first, 0u);
+      EXPECT_EQ(rangesRNTupleChain.back().second, ULong64_t(2 * nEntries));
+      for (size_t i = 1; i < rangesRNTupleChain.size(); ++i) {
+         EXPECT_EQ(rangesRNTupleChain[i].first, rangesRNTupleChain[i - 1].second);
+      }
+   }
+   gSystem->Unlink(fnameRNTuple1.c_str());
+   gSystem->Unlink(fnameRNTuple2.c_str());
+}