From 6c48513c0ff0c28a3463739e497c843430d4af92 Mon Sep 17 00:00:00 2001 From: Vincenzo Eduardo Padulano Date: Wed, 7 May 2025 15:49:10 +0200 Subject: [PATCH] [df] Support DefinePerSample for RNTuple The information of each sample is populated at every new `GetEntryRanges` call. For the single thread case, every call corresponds to a new file. For the MT case, different ranges might be accessing the same file. Extra logic is needed to expose the storage location of the RNTuple of the current range, since this is not available directly from the RPageSource. --- tree/dataframe/inc/ROOT/RDataSource.hxx | 6 +- tree/dataframe/inc/ROOT/RNTupleDS.hxx | 7 + tree/dataframe/inc/ROOT/RTTreeDS.hxx | 3 +- tree/dataframe/src/RDFUtils.cxx | 4 +- tree/dataframe/src/RDataSource.cxx | 2 +- tree/dataframe/src/RLoopManager.cxx | 4 +- tree/dataframe/src/RNTupleDS.cxx | 50 ++++- tree/dataframe/src/RTTreeDS.cxx | 2 +- .../test/dataframe_definepersample.cxx | 178 +++++++++++++++++- 9 files changed, 239 insertions(+), 17 deletions(-) diff --git a/tree/dataframe/inc/ROOT/RDataSource.hxx b/tree/dataframe/inc/ROOT/RDataSource.hxx index c457f99e1b503..e94eded844d9e 100644 --- a/tree/dataframe/inc/ROOT/RDataSource.hxx +++ b/tree/dataframe/inc/ROOT/RDataSource.hxx @@ -94,7 +94,7 @@ const std::vector &GetColumnNamesNoDuplicates(const ROOT::RDF::RDat void CallInitializeWithOpts(ROOT::RDF::RDataSource &ds, const std::set &suppressErrorsForMissingColumns); std::string DescribeDataset(ROOT::RDF::RDataSource &ds); ROOT::RDF::RSampleInfo -CreateSampleInfo(const ROOT::RDF::RDataSource &ds, +CreateSampleInfo(const ROOT::RDF::RDataSource &ds, unsigned int slot, const std::unordered_map &sampleMap); void RunFinalChecks(const ROOT::RDF::RDataSource &ds, bool nodesLeftNotRun); void ProcessMT(ROOT::RDF::RDataSource &ds, ROOT::Detail::RDF::RLoopManager &lm); @@ -168,10 +168,10 @@ protected: virtual std::string DescribeDataset() { return "Dataframe from datasource " + GetLabel(); } friend ROOT::RDF::RSampleInfo - ROOT::Internal::RDF::CreateSampleInfo(const ROOT::RDF::RDataSource &, + ROOT::Internal::RDF::CreateSampleInfo(const ROOT::RDF::RDataSource &, unsigned int, const std::unordered_map &); virtual ROOT::RDF::RSampleInfo - CreateSampleInfo(const std::unordered_map &) const; + CreateSampleInfo(unsigned int, const std::unordered_map &) const; friend void ROOT::Internal::RDF::RunFinalChecks(const ROOT::RDF::RDataSource &, bool); virtual void RunFinalChecks(bool) const {} diff --git a/tree/dataframe/inc/ROOT/RNTupleDS.hxx b/tree/dataframe/inc/ROOT/RNTupleDS.hxx index 70fca11375deb..544209ccc13b2 100644 --- a/tree/dataframe/inc/ROOT/RNTupleDS.hxx +++ b/tree/dataframe/inc/ROOT/RNTupleDS.hxx @@ -55,6 +55,7 @@ class RNTupleDS final : public ROOT::RDF::RDataSource { ULong64_t fFirstEntry = 0; ///< First entry index in fSource /// End entry index in fSource, e.g. the number of entries in the range is fLastEntry - fFirstEntry ULong64_t fLastEntry = 0; + std::string_view fFileName; ///< Storage location of the current RNTuple }; /// A clone of the first pages source's descriptor. @@ -109,6 +110,8 @@ class RNTupleDS final : public ROOT::RDF::RDataSource { /// the fCurrentRanges vectors. This is necessary because the returned ranges get distributed arbitrarily /// onto slots. In the InitSlot method, the column readers use this map to find the correct range to connect to. std::unordered_map fFirstEntry2RangeIdx; + /// One element per slot, corresponding to the current range index for that slot, as filled by InitSlot + std::vector fSlotsToRangeIdxs; /// The background thread that runs StageNextSources() std::thread fThreadStaging; @@ -194,6 +197,10 @@ public: std::unique_ptr GetColumnReaders(unsigned int /*slot*/, std::string_view /*name*/, const std::type_info &) final; + ROOT::RDF::RSampleInfo + CreateSampleInfo(unsigned int, + const std::unordered_map &) const final; + // Old API, unused bool SetEntry(unsigned int, ULong64_t) final { return true; } diff --git a/tree/dataframe/inc/ROOT/RTTreeDS.hxx b/tree/dataframe/inc/ROOT/RTTreeDS.hxx index d4b9f87389b5e..5a75a76f78008 100644 --- a/tree/dataframe/inc/ROOT/RTTreeDS.hxx +++ b/tree/dataframe/inc/ROOT/RTTreeDS.hxx @@ -75,7 +75,8 @@ class RTTreeDS final : public ROOT::RDF::RDataSource { std::unique_ptr fNoCleanupNotifier; ROOT::RDF::RSampleInfo - CreateSampleInfo(const std::unordered_map &sampleMap) const final; + CreateSampleInfo(unsigned int, + const std::unordered_map &sampleMap) const final; void RunFinalChecks(bool nodesLeftNotRun) const final; diff --git a/tree/dataframe/src/RDFUtils.cxx b/tree/dataframe/src/RDFUtils.cxx index 0c78b0ecdcbd4..b818ed814e6e2 100644 --- a/tree/dataframe/src/RDFUtils.cxx +++ b/tree/dataframe/src/RDFUtils.cxx @@ -572,10 +572,10 @@ std::string ROOT::Internal::RDF::DescribeDataset(ROOT::RDF::RDataSource &ds) } ROOT::RDF::RSampleInfo ROOT::Internal::RDF::CreateSampleInfo( - const ROOT::RDF::RDataSource &ds, + const ROOT::RDF::RDataSource &ds, unsigned int slot, const std::unordered_map &sampleMap) { - return ds.CreateSampleInfo(sampleMap); + return ds.CreateSampleInfo(slot, sampleMap); } void ROOT::Internal::RDF::RunFinalChecks(const ROOT::RDF::RDataSource &ds, bool nodesLeftNotRun) diff --git a/tree/dataframe/src/RDataSource.cxx b/tree/dataframe/src/RDataSource.cxx index 702593c283ef8..009d8f2c574e2 100644 --- a/tree/dataframe/src/RDataSource.cxx +++ b/tree/dataframe/src/RDataSource.cxx @@ -9,7 +9,7 @@ #endif ROOT::RDF::RSampleInfo ROOT::RDF::RDataSource::CreateSampleInfo( - const std::unordered_map &) const + unsigned int, const std::unordered_map &) const { // Currently not implemented for the generic data source, only works correctly for TTree. // TODO: Implement the feature also for the generic data source. diff --git a/tree/dataframe/src/RLoopManager.cxx b/tree/dataframe/src/RLoopManager.cxx index 8d56081346485..24f4a025153af 100644 --- a/tree/dataframe/src/RLoopManager.cxx +++ b/tree/dataframe/src/RLoopManager.cxx @@ -614,7 +614,7 @@ void RLoopManager::RunDataSource() ranges = fDataSource->GetEntryRanges(); - fSampleInfos[0] = ROOT::Internal::RDF::CreateSampleInfo(*fDataSource, fSampleMap); + fSampleInfos[0] = ROOT::Internal::RDF::CreateSampleInfo(*fDataSource, /*slot*/ 0, fSampleMap); try { for (const auto &range : ranges) { @@ -1326,6 +1326,8 @@ void ROOT::Detail::RDF::RLoopManager::DataSourceThreadTask(const std::pairGetLabel(), start, end, slot}); try { diff --git a/tree/dataframe/src/RNTupleDS.cxx b/tree/dataframe/src/RNTupleDS.cxx index 80197eb7e7ef5..1937c234588ff 100644 --- a/tree/dataframe/src/RNTupleDS.cxx +++ b/tree/dataframe/src/RNTupleDS.cxx @@ -13,6 +13,7 @@ #include #include +#include #include #include #include @@ -411,6 +412,7 @@ std::unique_ptr CreatePageSource(std::string_view n ROOT::RDF::RNTupleDS::RNTupleDS(std::string_view ntupleName, std::string_view fileName) : RNTupleDS(CreatePageSource(ntupleName, fileName)) { + fFileNames = std::vector{std::string{fileName}}; } ROOT::RDF::RNTupleDS::RNTupleDS(RNTuple *ntuple) : RNTupleDS(ROOT::Internal::RPageSourceFile::CreateFromAnchor(*ntuple)) @@ -534,6 +536,7 @@ void ROOT::RDF::RNTupleDS::PrepareNextRanges() // to open and attach files here. range.fSource = CreatePageSource(fNTupleName, fFileNames[fNextFileIndex]); } + range.fFileName = fFileNames[fNextFileIndex]; range.fSource->Attach(); fNextFileIndex++; @@ -553,6 +556,8 @@ void ROOT::RDF::RNTupleDS::PrepareNextRanges() unsigned int nSlotsPerFile = fNSlots / nRemainingFiles; for (std::size_t i = 0; (fNextRanges.size() < fNSlots) && (fNextFileIndex < nFiles); ++i) { std::unique_ptr source; + // Need to look for the file name to populate the sample info later + const auto &sourceFileName = fFileNames[fNextFileIndex]; std::swap(fStagingArea[fNextFileIndex], source); if (!source) { // Empty files trigger this condition @@ -595,6 +600,7 @@ void ROOT::RDF::RNTupleDS::PrepareNextRanges() auto end = rangesByCluster[iRange - 1].second; REntryRangeDS range; + range.fFileName = sourceFileName; // The last range for this file just takes the already opened page source. All previous ranges clone. if (iSlot == N - 1) { range.fSource = std::move(source); @@ -689,10 +695,20 @@ std::vector> ROOT::RDF::RNTupleDS::GetEntryRange void ROOT::RDF::RNTupleDS::InitSlot(unsigned int slot, ULong64_t firstEntry) { - if (fNSlots == 1) + if (fNSlots == 1) { + // Ensure the connection between slot and range is valid also in single-thread mode + fSlotsToRangeIdxs[0] = 0; return; + } + // The same slot ID could be picked multiple times in the same execution, thus + // ending up processing different page sources. Here we re-establish the + // connection between the slot and the correct page source by finding which + // range index corresponds to the first entry passed. auto idxRange = fFirstEntry2RangeIdx.at(firstEntry); + // We also remember this connection so it can later be retrieved in CreateSampleInfo + fSlotsToRangeIdxs[slot * ROOT::Internal::RDF::CacheLineStep()] = idxRange; + for (auto r : fActiveColumnReaders[slot]) { r->Connect(*fCurrentRanges[idxRange].fSource, firstEntry - fCurrentRanges[idxRange].fFirstEntry); } @@ -777,6 +793,7 @@ void ROOT::RDF::RNTupleDS::SetNSlots(unsigned int nSlots) assert(nSlots > 0); fNSlots = nSlots; fActiveColumnReaders.resize(fNSlots); + fSlotsToRangeIdxs.resize(fNSlots * ROOT::Internal::RDF::CacheLineStep()); } ROOT::RDataFrame ROOT::RDF::FromRNTuple(std::string_view ntupleName, std::string_view fileName) @@ -788,3 +805,34 @@ ROOT::RDataFrame ROOT::RDF::FromRNTuple(std::string_view ntupleName, const std:: { return ROOT::RDataFrame(std::make_unique(ntupleName, fileNames)); } + +ROOT::RDF::RSampleInfo ROOT::Internal::RDF::RNTupleDS::CreateSampleInfo( + unsigned int slot, const std::unordered_map &sampleMap) const +{ + // The same slot ID could be picked multiple times in the same execution, thus + // ending up processing different page sources. Here we re-establish the + // connection between the slot and the correct page source by retrieving + // which range is connected currently to the slot + const auto &rangeIdx = fSlotsToRangeIdxs.at(slot * ROOT::Internal::RDF::CacheLineStep()); + + // Missing source if a file does not exist + if (!fCurrentRanges[rangeIdx].fSource) + return ROOT::RDF::RSampleInfo{}; + + const auto &ntupleName = fCurrentRanges[rangeIdx].fSource->GetNTupleName(); + const auto &ntuplePath = fCurrentRanges[rangeIdx].fFileName; + const auto ntupleID = std::string(ntuplePath) + '/' + ntupleName; + + // TODO: There is no support for RNTuple in RDatasetSpec, thus the sample map + // is always empty at the moment. + if (sampleMap.empty()) + return ROOT::RDF::RSampleInfo( + ntupleID, std::make_pair(fCurrentRanges[rangeIdx].fFirstEntry, fCurrentRanges[rangeIdx].fLastEntry)); + + if (sampleMap.find(ntupleID) == sampleMap.end()) + throw std::runtime_error("Full sample identifier '" + ntupleID + "' cannot be found in the available samples."); + + return ROOT::RDF::RSampleInfo( + ntupleID, std::make_pair(fCurrentRanges[rangeIdx].fFirstEntry, fCurrentRanges[rangeIdx].fLastEntry), + sampleMap.at(ntupleName)); +} diff --git a/tree/dataframe/src/RTTreeDS.cxx b/tree/dataframe/src/RTTreeDS.cxx index 35e230d27a444..38284196c42db 100644 --- a/tree/dataframe/src/RTTreeDS.cxx +++ b/tree/dataframe/src/RTTreeDS.cxx @@ -155,7 +155,7 @@ ROOT::Internal::RDF::FromTTree(std::string_view treeName, const std::vector &sampleMap) const + unsigned int, const std::unordered_map &sampleMap) const { // one GetTree to retrieve the TChain, another to retrieve the underlying TTree auto *tree = fTreeReader->GetTree()->GetTree(); diff --git a/tree/dataframe/test/dataframe_definepersample.cxx b/tree/dataframe/test/dataframe_definepersample.cxx index 61c9fca902c6f..61358e2ef5376 100644 --- a/tree/dataframe/test/dataframe_definepersample.cxx +++ b/tree/dataframe/test/dataframe_definepersample.cxx @@ -4,9 +4,13 @@ #include #include +#include +#include + #include #include +#include #include #include // std::thread::hardware_concurrency @@ -170,20 +174,180 @@ TEST(DefinePerSample, TwoExecutions) EXPECT_TRUE(flag); } -/* TODO -// Not supported yet -TEST(DefinePerSample, DataSource) +struct InputRNTuplesRAII { + unsigned int fNFiles = 0; + std::string fPrefix; + + InputRNTuplesRAII(unsigned int nFiles, std::string prefix) : fNFiles(nFiles), fPrefix(std::move(prefix)) + { + for (auto i = 0u; i < fNFiles; ++i) { + auto model = ROOT::RNTupleModel::Create(); + auto fldX = model->MakeField("x"); + auto fn = fPrefix + std::to_string(i) + ".root"; + auto ntpl = ROOT::RNTupleWriter::Recreate(std::move(model), "ntuple", fn); + *fldX = i; + ntpl->Fill(); + } + } + + ~InputRNTuplesRAII() + { + for (auto i = 0u; i < fNFiles; ++i) + std::remove((fPrefix + std::to_string(i) + ".root").c_str()); + } +}; + +TEST_P(DefinePerSample, RNTupleSingle) +{ + const std::string prefix = "rdfdefinepersample_rntuple"; + InputRNTuplesRAII file(1u, prefix); + ROOT::RDataFrame df("ntuple", prefix + "0.root"); + + std::atomic_int counter{0}; + auto df2 = df.DefinePerSample("y", [&counter](unsigned int, const ROOT::RDF::RSampleInfo &db) { + EXPECT_EQ(db.EntryRange(), std::make_pair(0ull, 1ull)); + ++counter; + return 42; + }); + auto xmin = df2.Min("y"); + auto xmax = df2.Max("y"); + EXPECT_EQ(*xmin, 42); + EXPECT_EQ(*xmax, 42); + const auto expected = 1u; // as the RNTuple only contains one cluster, we only have one "data-block" + EXPECT_EQ(counter, expected); +} + +TEST_P(DefinePerSample, RNTupleMany) +{ + const std::vector fileNames{ + "rdfdefinepersample_rntuple_many0.root", "rdfdefinepersample_rntuple_many1.root", + "rdfdefinepersample_rntuple_many2.root", "rdfdefinepersample_rntuple_many3.root", + "rdfdefinepersample_rntuple_many4.root"}; + const std::string prefix{"rdfdefinepersample_rntuple_many"}; + InputRNTuplesRAII file(5u, prefix); + ROOT::RDataFrame df("ntuple", fileNames); + + std::atomic_int counter{0}; + auto df2 = df.DefinePerSample("y", [&counter](unsigned int, const ROOT::RDF::RSampleInfo &db) { + EXPECT_EQ(db.EntryRange(), std::make_pair(0ull, 1ull)); + ++counter; + return 42; + }); + auto xmin = df2.Min("y"); + auto xmax = df2.Max("y"); + EXPECT_EQ(*xmin, 42); + EXPECT_EQ(*xmax, 42); + const auto expected = 5u; // one "data-block" per RNTuple (because each RNTuple only has one cluster) + EXPECT_EQ(counter, expected); +} + +struct DefinePerSampleWithDataset : ::testing::TestWithParam> { + unsigned int fNSlots{}; + unsigned int fNEntries{5}; + std::string fDatasetName{"dataset"}; + std::vector fFileNames{ + "rdf_definepersample_with_dataset_0.root", "rdf_definepersample_with_dataset_1.root", + "rdf_definepersample_with_dataset_2.root", "rdf_definepersample_with_dataset_3.root", + "rdf_definepersample_with_dataset_4.root"}; + + void CreateRNTupleDataset() + { + for (const auto &fn : fFileNames) { + auto model = ROOT::RNTupleModel::Create(); + auto fldX = model->MakeField("x"); + auto ntpl = ROOT::RNTupleWriter::Recreate(std::move(model), fDatasetName, fn); + for (ULong64_t entry = 0; entry < fNEntries; entry++) { + *fldX = entry; + ntpl->Fill(); + if (entry % 2 == 0) + ntpl->CommitCluster(); + } + } + } + + void CreateTTreeDataset() + { + for (const auto &fn : fFileNames) { + auto f = std::make_unique(fn.c_str(), "recreate"); + auto t = std::make_unique(fDatasetName.c_str(), fDatasetName.c_str()); + ULong64_t i{}; + t->Branch("x", &i); + for (ULong64_t entry = 0; entry < fNEntries; entry++) { + i = entry; + t->Fill(); + if (entry % 2 == 0) + t->FlushBaskets(); + } + f->Write(); + } + } + + DefinePerSampleWithDataset() : fNSlots(GetParam().first ? std::min(4u, std::thread::hardware_concurrency()) : 1u) + { + if (GetParam().first) + ROOT::EnableImplicitMT(); + + if (GetParam().second) + CreateRNTupleDataset(); + else + CreateTTreeDataset(); + } + + ~DefinePerSampleWithDataset() override + { + if (GetParam().first) + ROOT::DisableImplicitMT(); + + for (const auto &fn : fFileNames) + std::remove(fn.c_str()); + } +}; + +template +void expect_vec_eq(const T0 &v1, const T1 &v2) +{ + ASSERT_EQ(v1.size(), v2.size()) << "Vectors 'v1' and 'v2' are of unequal length"; + for (std::size_t i = 0ull; i < v1.size(); ++i) { + EXPECT_EQ(v1[i], v2[i]) << "Vectors 'v1' and 'v2' differ at index " << i; + } +} + +TEST_P(DefinePerSampleWithDataset, CorrectSampleID) +{ + // Check that the sample id returned by the RSampleInfo is as expected + ROOT::RDataFrame df{fDatasetName, fFileNames}; + auto df1 = + df.DefinePerSample("sampleID", [](unsigned int, const ROOT::RDF::RSampleInfo &si) { return si.AsString(); }); + + auto take = df1.Take("sampleID"); + auto &sampleIDs = *take; + // 5 files, 5 entries per file + EXPECT_EQ(sampleIDs.size(), 25); + // Sort to have consistent results also in MT case + std::sort(sampleIDs.begin(), sampleIDs.end()); + std::vector expectedIDs; + expectedIDs.reserve(25); + for (auto i = 0; i < 25; i++) { + expectedIDs.push_back(fFileNames[i / 5] + '/' + fDatasetName); + } + expect_vec_eq(sampleIDs, expectedIDs); +} + +std::string GetTestLabel(const testing::TestParamInfo &testInfo) { - ROOT::RDataFrame df(std::make_unique(1)); - auto r = df.DefinePerSample("col0", [] { return 42; }).Max("col0"); - EXPECT_EQ(*r, 42); + if (testInfo.param.second) + return "RNTuple"; + return "TTree"; } -*/ // instantiate single-thread tests INSTANTIATE_TEST_SUITE_P(Seq, DefinePerSample, ::testing::Values(false)); +INSTANTIATE_TEST_SUITE_P(Seq, DefinePerSampleWithDataset, + ::testing::Values(std::make_pair(false, false), std::make_pair(false, true)), GetTestLabel); #ifdef R__USE_IMT // instantiate multi-thread tests INSTANTIATE_TEST_SUITE_P(MT, DefinePerSample, ::testing::Values(true)); +INSTANTIATE_TEST_SUITE_P(MT, DefinePerSampleWithDataset, + ::testing::Values(std::make_pair(true, false), std::make_pair(true, true)), GetTestLabel); #endif