From dabd77b669cf41aaf3262ea5b5bf356c29b4b661 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
---
 .../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 4d4a990eb2ed7..f666b1f37090e 100644
--- a/tree/dataframe/test/dataframe_vary.cxx
+++ b/tree/dataframe/test/dataframe_vary.cxx
@@ -385,6 +385,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);