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