From aa03186df056879ba49e3255e20ab5bf2d5ddaf9 Mon Sep 17 00:00:00 2001
From: Vincenzo Eduardo Padulano <vincenzo.eduardo.padulano@cern.ch>
Date: Thu, 21 May 2026 10:18:36 +0200
Subject: [PATCH] [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 aecad0a9570f7113392f43542f4a51eb9158e0d8)
---
 .../inc/ROOT/RDF/RDefinePerSample.hxx         |  5 ++-
 tree/dataframe/test/dataframe_vary.cxx        | 34 +++++++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/tree/dataframe/inc/ROOT/RDF/RDefinePerSample.hxx b/tree/dataframe/inc/ROOT/RDF/RDefinePerSample.hxx
index 0302113345774..7f0fb27af918b 100644
--- a/tree/dataframe/inc/ROOT/RDF/RDefinePerSample.hxx
+++ b/tree/dataframe/inc/ROOT/RDF/RDefinePerSample.hxx
@@ -81,7 +81,10 @@ public:
 
    RDefineBase &GetVariedDefine(const std::string &) final
    {
-      R__ASSERT(false && "This should never be called");
+      // RDefinePerSample cannot depend on varied columns, so we return itself.
+      // This supports the use case of a downstream defined variable that depends on variations and also on a column
+      // created via DefinePerSample. The request for an action depending on that defined variable will end up here when
+      // looking for a variation of the dependant column.
       return *this;
    }
 };
diff --git a/tree/dataframe/test/dataframe_vary.cxx b/tree/dataframe/test/dataframe_vary.cxx
index eda767df9a3e7..d0fd00bf43469 100644
--- a/tree/dataframe/test/dataframe_vary.cxx
+++ b/tree/dataframe/test/dataframe_vary.cxx
@@ -355,6 +355,40 @@ TEST_P(RDFVary, VaryDefinePerSample)
    EXPECT_EQ(ss["x:1"], 2 * 10);
 }
 
+TEST_P(RDFVary, VaryDefinePerSampleDownstreamVariedAction)
+{
+   // Regression test for https://github.com/root-project/root/issues/22367
+   struct DataRAII {
+      const char *fFileName{"VaryDefinePerSampleDownstreamVariedAction.root"};
+      const char *fTreeName{"VaryDefinePerSampleDownstreamVariedAction"};
+      DataRAII()
+      {
+         auto f = std::make_unique<TFile>(fFileName, "recreate");
+         auto t = std::make_unique<TTree>(fTreeName, fTreeName);
+
+         float qcd_scale{};
+         t->Branch("qcd_scale", &qcd_scale);
+         std::vector<float> vals{0.5f, 1.f, 1.5f, 2.f, 2.5f};
+         for (auto val : vals) {
+            qcd_scale = val;
+            t->Fill();
+         }
+         f->Write();
+      }
+      ~DataRAII() { std::remove(fFileName); }
+   } dataset;
+
+   ROOT::RDF::RNode df = ROOT::RDataFrame(dataset.fTreeName, dataset.fFileName);
+   df = df.DefinePerSample("xs", [](unsigned, const ROOT::RDF::RSampleInfo &) { return 0.5f; });
+   df = df.Vary("qcd_scale", [](float s) { return ROOT::RVecF{s * 1.1f, s * 0.9f}; }, {"qcd_scale"}, {"up", "down"});
+   df = df.Define("weight", [](float scale, float xs) { return scale * xs; }, {"qcd_scale", "xs"});
+   auto nominal = df.Sum<float>("weight");
+   auto vars = ROOT::RDF::Experimental::VariationsFor(nominal);
+   EXPECT_FLOAT_EQ(vars["nominal"], 0.5f * 7.5f);
+   EXPECT_FLOAT_EQ(vars["qcd_scale:up"], 0.5f * 1.1f * 7.5f);
+   EXPECT_FLOAT_EQ(vars["qcd_scale:down"], 0.5f * 0.9f * 7.5f);
+}
+
 TEST(RDFVary, SaveGraph)
 {
    ROOT::RDataFrame df(1);