Skip to content

Commit 8bd0913

Browse files
vepadulanodpiparo
authored andcommitted
[df] Pass DefinePerSample column to varied actions on request
A varied action will be looking for column readers through its current variation name. In the case of an upstream Define, the node will return the correct reader based on the variation. In the case of an upstream DefinePerSample, given it can't depend on variations and it can't vary itself, we now return just the node itself. Fixes #22367 (cherry picked from commit aecad0a)
1 parent 0f5ce86 commit 8bd0913

2 files changed

Lines changed: 38 additions & 1 deletion

File tree

tree/dataframe/inc/ROOT/RDF/RDefinePerSample.hxx

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,10 @@ public:
8181

8282
RDefineBase &GetVariedDefine(const std::string &) final
8383
{
84-
R__ASSERT(false && "This should never be called");
84+
// RDefinePerSample cannot depend on varied columns, so we return itself.
85+
// This supports the use case of a downstream defined variable that depends on variations and also on a column
86+
// created via DefinePerSample. The request for an action depending on that defined variable will end up here when
87+
// looking for a variation of the dependant column.
8588
return *this;
8689
}
8790
};

tree/dataframe/test/dataframe_vary.cxx

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -355,6 +355,40 @@ TEST_P(RDFVary, VaryDefinePerSample)
355355
EXPECT_EQ(ss["x:1"], 2 * 10);
356356
}
357357

358+
TEST_P(RDFVary, VaryDefinePerSampleDownstreamVariedAction)
359+
{
360+
// Regression test for https://github.com/root-project/root/issues/22367
361+
struct DataRAII {
362+
const char *fFileName{"VaryDefinePerSampleDownstreamVariedAction.root"};
363+
const char *fTreeName{"VaryDefinePerSampleDownstreamVariedAction"};
364+
DataRAII()
365+
{
366+
auto f = std::make_unique<TFile>(fFileName, "recreate");
367+
auto t = std::make_unique<TTree>(fTreeName, fTreeName);
368+
369+
float qcd_scale{};
370+
t->Branch("qcd_scale", &qcd_scale);
371+
std::vector<float> vals{0.5f, 1.f, 1.5f, 2.f, 2.5f};
372+
for (auto val : vals) {
373+
qcd_scale = val;
374+
t->Fill();
375+
}
376+
f->Write();
377+
}
378+
~DataRAII() { std::remove(fFileName); }
379+
} dataset;
380+
381+
ROOT::RDF::RNode df = ROOT::RDataFrame(dataset.fTreeName, dataset.fFileName);
382+
df = df.DefinePerSample("xs", [](unsigned, const ROOT::RDF::RSampleInfo &) { return 0.5f; });
383+
df = df.Vary("qcd_scale", [](float s) { return ROOT::RVecF{s * 1.1f, s * 0.9f}; }, {"qcd_scale"}, {"up", "down"});
384+
df = df.Define("weight", [](float scale, float xs) { return scale * xs; }, {"qcd_scale", "xs"});
385+
auto nominal = df.Sum<float>("weight");
386+
auto vars = ROOT::RDF::Experimental::VariationsFor(nominal);
387+
EXPECT_FLOAT_EQ(vars["nominal"], 0.5f * 7.5f);
388+
EXPECT_FLOAT_EQ(vars["qcd_scale:up"], 0.5f * 1.1f * 7.5f);
389+
EXPECT_FLOAT_EQ(vars["qcd_scale:down"], 0.5f * 0.9f * 7.5f);
390+
}
391+
358392
TEST(RDFVary, SaveGraph)
359393
{
360394
ROOT::RDataFrame df(1);

0 commit comments

Comments
 (0)