From 0ad7d4b421f7af3c2a44beee7efb8dad81362030 Mon Sep 17 00:00:00 2001 From: Giovanni Date: Thu, 30 Jul 2026 11:47:11 +0200 Subject: [PATCH 1/2] [df] Add RedefineForSample --- tree/dataframe/inc/ROOT/RDF/RInterface.hxx | 104 +++++++++++++++------ 1 file changed, 75 insertions(+), 29 deletions(-) diff --git a/tree/dataframe/inc/ROOT/RDF/RInterface.hxx b/tree/dataframe/inc/ROOT/RDF/RInterface.hxx index aa3bb93285d00..4552f7b553811 100644 --- a/tree/dataframe/inc/ROOT/RDF/RInterface.hxx +++ b/tree/dataframe/inc/ROOT/RDF/RInterface.hxx @@ -749,25 +749,16 @@ public: template ::ret_type> RInterface DefinePerSample(std::string_view name, F expression) { - RDFInternal::CheckValidCppVarName(name, "DefinePerSample"); - RDFInternal::CheckForRedefinition("DefinePerSample", name, fColRegister, - GetDataSource() ? GetDataSource()->GetColumnNames() : ColumnNames_t{}); - - auto retTypeName = RDFInternal::TypeID2TypeName(typeid(RetType_t)); - if (retTypeName.empty()) { - // The type is not known to the interpreter. - // We must not error out here, but if/when this column is used in jitted code - const auto demangledType = RDFInternal::DemangleTypeIdName(typeid(RetType_t)); - retTypeName = "CLING_UNKNOWN_TYPE_" + demangledType; - } - - auto newColumn = - std::make_shared>(name, retTypeName, std::move(expression), *fLoopManager); + return DefinePerSampleImpl(name, std::move(expression), false); + } - RDFInternal::RColumnRegister newCols(fColRegister); - newCols.AddDefine(std::move(newColumn)); - RInterface newInterface(fProxiedPtr, *fLoopManager, std::move(newCols)); - return newInterface; + //////////////////////////////////////////////////////////////////////////// + /// \brief Redefine an existing column that is updated when the input sample changes. + /// As for DefinePerSample, but the column must already exist and will be overwritten. + template ::ret_type> + RInterface RedefinePerSample(std::string_view name, F expression) + { + return DefinePerSampleImpl(name, std::move(expression), true); } // clang-format off @@ -810,20 +801,18 @@ public: // clang-format on RInterface DefinePerSample(std::string_view name, std::string_view expression) { - RDFInternal::CheckValidCppVarName(name, "DefinePerSample"); - // these checks must be done before jitting lest we throw exceptions in jitted code - RDFInternal::CheckForRedefinition("DefinePerSample", name, fColRegister, - GetDataSource() ? GetDataSource()->GetColumnNames() : ColumnNames_t{}); - - auto jittedDefine = RDFInternal::BookDefinePerSampleJit(name, expression, *fLoopManager, fColRegister); + return DefinePerSampleJitImpl(name, expression, false); + } - RDFInternal::RColumnRegister newCols(fColRegister); - newCols.AddDefine(std::move(jittedDefine)); + //////////////////////////////////////////////////////////////////////////// + /// \brief Redefine an existing column that is updated when the input sample changes. + /// As for DefinePerSample, but the column must already exist and will be overwritten. + RInterface RedefinePerSample(std::string_view name, std::string_view expression) + { + return DefinePerSampleJitImpl(name, expression, true); + } - RInterface newInterface(fProxiedPtr, *fLoopManager, std::move(newCols)); - return newInterface; - } /// \brief Register systematic variations for a single existing column using custom variation tags. /// \param[in] colName name of the column for which varied values are provided. @@ -3832,6 +3821,63 @@ private: return *this; // never reached } + //////////////////////////////////////////////////////////////////////////// + /// \brief Implementation of DefinePerSample and RedefinePerSample (non-jitted). + template ::ret_type> + RInterface DefinePerSampleImpl(std::string_view name, F expression, bool redefine) + { + if (!redefine) { + RDFInternal::CheckValidCppVarName(name, "DefinePerSample"); + RDFInternal::CheckForRedefinition("DefinePerSample", name, fColRegister, + GetDataSource() ? GetDataSource()->GetColumnNames() : ColumnNames_t{}); + } else { + RDFInternal::CheckForDefinition("RedefinePerSample", name, fColRegister, + GetDataSource() ? GetDataSource()->GetColumnNames() : ColumnNames_t{}); + RDFInternal::CheckForNoVariations("RedefinePerSample", name, fColRegister); + } + + auto retTypeName = RDFInternal::TypeID2TypeName(typeid(RetType_t)); + if (retTypeName.empty()) { + // The type is not known to the interpreter. + // We must not error out here, but if/when this column is used in jitted code + const auto demangledType = RDFInternal::DemangleTypeIdName(typeid(RetType_t)); + retTypeName = "CLING_UNKNOWN_TYPE_" + demangledType; + } + + auto newColumn = + std::make_shared>(name, retTypeName, std::move(expression), *fLoopManager); + + RDFInternal::RColumnRegister newCols(fColRegister); + newCols.AddDefine(std::move(newColumn)); + RInterface newInterface(fProxiedPtr, *fLoopManager, std::move(newCols)); + return newInterface; + } + + //////////////////////////////////////////////////////////////////////////// + /// \brief Implementation of DefinePerSample and RedefinePerSample (jitted). + RInterface DefinePerSampleJitImpl(std::string_view name, std::string_view expression, bool redefine) + { + // these checks must be done before jitting lest we throw exceptions in jitted code + if (!redefine) { + RDFInternal::CheckValidCppVarName(name, redefine ? "RedefinePerSample" : "DefinePerSample"); + RDFInternal::CheckForRedefinition("DefinePerSample", name, fColRegister, + GetDataSource() ? GetDataSource()->GetColumnNames() : ColumnNames_t{}); + } else { + RDFInternal::CheckForDefinition("RedefinePerSample", name, fColRegister, + GetDataSource() ? GetDataSource()->GetColumnNames() : ColumnNames_t{}); + RDFInternal::CheckForNoVariations("RedefinePerSample", name, fColRegister); + } + + auto jittedDefine = RDFInternal::BookDefinePerSampleJit(name, expression, *fLoopManager, fColRegister); + + RDFInternal::RColumnRegister newCols(fColRegister); + newCols.AddDefine(std::move(jittedDefine)); + + RInterface newInterface(fProxiedPtr, *fLoopManager, std::move(newCols)); + + return newInterface; + } + //////////////////////////////////////////////////////////////////////////// /// \brief Implementation of cache. template From ab1ac9c6eac73bb83b509f4c15ad47e330577640 Mon Sep 17 00:00:00 2001 From: Giovanni Date: Thu, 30 Jul 2026 12:06:14 +0200 Subject: [PATCH 2/2] [df] unit tests for RedefinePerSample --- .../test/dataframe_definepersample.cxx | 37 +++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/tree/dataframe/test/dataframe_definepersample.cxx b/tree/dataframe/test/dataframe_definepersample.cxx index bc276766026af..a71d3b7ddf8ca 100644 --- a/tree/dataframe/test/dataframe_definepersample.cxx +++ b/tree/dataframe/test/dataframe_definepersample.cxx @@ -140,6 +140,43 @@ TEST(DefinePerSampleMore, ThrowOnRedefinition) std::runtime_error); } +TEST_P(DefinePerSample, ThrowOnRedefinitionExistingTree) +{ + const std::string prefix = "rdfdefinepersample_tree"; + InputFilesRAII file(1u, prefix); + ROOT::RDataFrame df("t", prefix + "*"); + EXPECT_THROW(df.DefinePerSample("x", [](unsigned, const ROOT::RDF::RSampleInfo &) { return 42; }), + std::runtime_error); +} + +TEST_P(DefinePerSample, CheckRedefinitionTree) +{ + const std::string prefix = "rdfdefinepersample_tree"; + InputFilesRAII file(1u, prefix); + ROOT::RDataFrame df("t", prefix + "*"); + + std::atomic_int counter{0}; + auto df2 = df.RedefinePerSample("x", [&counter](unsigned int, const ROOT::RDF::RSampleInfo &db) { + EXPECT_EQ(db.EntryRange(), std::make_pair(0ull, 1ull)); + ++counter; + return 42; + }); + auto xmin = df2.Min("x"); + auto xmax = df2.Max("x"); + EXPECT_EQ(*xmin, 42); + EXPECT_EQ(*xmax, 42); + const auto expected = 1u; // as the TTree only contains one cluster, we only have one "data-block" + EXPECT_EQ(counter, expected); +} + +TEST(DefinePerSampleMore, ThrowOnNonRedefinition) +{ + auto df = ROOT::RDataFrame(1) + .Define("x", [] { return 42; }); + EXPECT_THROW(df.RedefinePerSample("y", [](unsigned, const ROOT::RDF::RSampleInfo &) { return 42; }), + std::runtime_error); +} + TEST(DefinePerSampleMore, GetColumnType) { auto df = ROOT::RDataFrame(1).DefinePerSample("x", [](unsigned, const ROOT::RDF::RSampleInfo &) { return 42; });