diff --git a/tree/dataframe/inc/ROOT/RDF/RActionSnapshot.hxx b/tree/dataframe/inc/ROOT/RDF/RActionSnapshot.hxx index 19da04127dbeb..510b8b5ec273b 100644 --- a/tree/dataframe/inc/ROOT/RDF/RActionSnapshot.hxx +++ b/tree/dataframe/inc/ROOT/RDF/RActionSnapshot.hxx @@ -74,9 +74,6 @@ class R__CLING_PTRCHECK(off) RActionSnapshot final : public RActionBase { const auto &prevVariations = nominalPrevNode->GetVariations(); fPrevNodes.reserve(1 + currentVariations.size()); - // Need to populate parts of the computation graph for which we have empty shells, e.g. RJittedFilters - if (!currentVariations.empty()) - fLoopManager->Jit(); for (const auto &variation : currentVariations) { if (IsStrInVec(variation, prevVariations)) { fPrevNodes.emplace_back( @@ -105,7 +102,19 @@ public: fIsDefine.push_back(colRegister.IsDefineOrAlias(columns[i])); if constexpr (std::is_same_v) { + // Need to populate parts of the computation graph for which we have empty shells, e.g. RJittedFilters and + // varied Defines + if (!GetVariations().empty()) + fLoopManager->Jit(); + AppendVariedPrevNodes(); + + for (auto i = 0u; i < nColumns; ++i) { + if (fIsDefine[i]) { + auto define = colRegister.GetDefine(columns[i]); + define->MakeVariations(GetVariations()); + } + } } } diff --git a/tree/dataframe/test/dataframe_snapshotWithVariations.cxx b/tree/dataframe/test/dataframe_snapshotWithVariations.cxx index a4ef10fb390ca..ab69fc11fb45e 100644 --- a/tree/dataframe/test/dataframe_snapshotWithVariations.cxx +++ b/tree/dataframe/test/dataframe_snapshotWithVariations.cxx @@ -529,3 +529,51 @@ TEST(RDFVarySnapshot, GH20320) EXPECT_EQ(take_val->front(), 2); EXPECT_EQ(take_var_up->front(), 3); } + +// https://github.com/root-project/root/issues/20506 +// When snapshot is invoked on JITted Defines with variations, the Snapshot +// Action needs to JIT the computation graph before creating outputs. +// When this wasn't happening, the varied Defines be identical to nominal and +// not be snapshot. +TEST(RDFVarySnapshot, IncludeDependentColumns_JIT) +{ + const char *fileName{"dataframe_snapshot_include_dependent_columns.root"}; + RemoveFileRAII(fileName); + constexpr unsigned int N = 10; + + ROOT::RDataFrame rdf(N); + auto df = rdf.Define("Muon_pt", "return static_cast(rdfentry_)") + .Vary("Muon_pt", "ROOT::VecOps::RVec({Muon_pt*0.5, Muon_pt*2})", {"down", "up"}, "muon_unc") + .Define("Muon_2pt", "return 2. * Muon_pt;"); + + ROOT::RDF::RSnapshotOptions opts; + opts.fOverwriteIfExists = true; + opts.fIncludeVariations = true; + + auto snapshot = df.Snapshot("Events", fileName, {"Muon_pt", "Muon_2pt"}, opts); + + TFile file(fileName); + auto tree = file.Get("Events"); + ASSERT_NE(tree, nullptr); + tree->Scan(); + + double Muon_pt, Muon_pt_up, Muon_pt_down; + double Muon_2pt, Muon_2pt_up, Muon_2pt_down; + ASSERT_EQ(TTree::kMatch, tree->SetBranchAddress("Muon_pt", &Muon_pt)); + ASSERT_EQ(TTree::kMatch, tree->SetBranchAddress("Muon_pt__muon_unc_up", &Muon_pt_up)); + ASSERT_EQ(TTree::kMatch, tree->SetBranchAddress("Muon_pt__muon_unc_down", &Muon_pt_down)); + ASSERT_EQ(TTree::kMatch, tree->SetBranchAddress("Muon_2pt", &Muon_2pt)); + ASSERT_EQ(TTree::kMatch, tree->SetBranchAddress("Muon_2pt__muon_unc_up", &Muon_2pt_up)); + ASSERT_EQ(TTree::kMatch, tree->SetBranchAddress("Muon_2pt__muon_unc_down", &Muon_2pt_down)); + + EXPECT_EQ(tree->GetEntries(), N); + for (unsigned int i = 0; i < tree->GetEntries(); ++i) { + ASSERT_GT(tree->GetEntry(i), 0); + EXPECT_EQ(2. * Muon_pt, Muon_2pt); + + EXPECT_EQ(0.5 * Muon_pt, Muon_pt_down); + EXPECT_EQ(0.5 * Muon_2pt, Muon_2pt_down); + EXPECT_EQ(2. * Muon_pt, Muon_pt_up); + EXPECT_EQ(2. * Muon_2pt, Muon_2pt_up); + } +}