[RDF] Add tests for GetClusterRanges

Author: siliataider

tree/dataframe/test/dataframe_utils.cxx

@@ -8,6 +8,7 @@
 #include <vector>
 
 #include <TBranch.h>
+#include <TSystem.h>
 
 namespace RDFInt = ROOT::Internal::RDF;
 
@@ -237,3 +238,30 @@ 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 fname = "dataframe_interfaceAndUtils_2.root";
+   const int nEntries = 1000;
+   const int clusterSize = 250;
+
+   {
+      ROOT::RDF::RSnapshotOptions opts;
+      opts.fAutoFlush = clusterSize;
+      ROOT::RDataFrame(nEntries)
+         .Define("i", "rdfentry_")
+         .Snapshot("t", fname, {"i"}, opts);
+   }
+
+   ROOT::RDataFrame df("t", fname);
+   auto *lm = df.GetLoopManager();
+   auto ranges = RDFInt::GetClusterRanges(*lm);
+
+   EXPECT_EQ(ranges.size(), nEntries / clusterSize);
+   for (size_t i = 0; i < ranges.size(); ++i) {
+      EXPECT_EQ(ranges[i].first, i * clusterSize);
+      EXPECT_EQ(ranges[i].second, (i + 1) * clusterSize);
+   }
+
+   gSystem->Unlink(fname.c_str());
+}