diff --git a/README/ReleaseNotes/v642/index.md b/README/ReleaseNotes/v642/index.md index 99308b1bc9626..9f9bd6323b2f3 100644 --- a/README/ReleaseNotes/v642/index.md +++ b/README/ReleaseNotes/v642/index.md @@ -77,6 +77,11 @@ maps) will now obtain different, mathematically consistent values. ## Math +## RDataFrame + +* Added `RedefinePerSample` transformation. Works similarly to `DefinePerSample`, but allows to redefine existing values + of a column on a per-sample basis. This operation is supported in local and distributed mode. + ## RooFit ### Small changes diff --git a/bindings/distrdf/python/DistRDF/Operation.py b/bindings/distrdf/python/DistRDF/Operation.py index a699192ce5cae..577fa711cb99d 100644 --- a/bindings/distrdf/python/DistRDF/Operation.py +++ b/bindings/distrdf/python/DistRDF/Operation.py @@ -27,6 +27,7 @@ def __init__(self, name: str, *args, **kwargs): class Action(Operation): """An action attached to a distributed RDataFrame graph node.""" + pass @@ -65,7 +66,8 @@ def __init__(self, name: str, *args, **kwargs): "Creating a histogram without a model is not supported in distributed mode. Please make sure to " "specify the histogram model when rerunning the distributed RDataFrame application. For example:\n\n" "\tHisto1D('mycolumn') --> Histo1D(('myhist', 'myhist', 100, 0, 10), 'mycolumn')\n\n" - "See the RDataFrame documentation for more details.") + "See the RDataFrame documentation for more details." + ) raise ValueError(message) @@ -74,26 +76,31 @@ class VariationsFor(Action): DistRDF.VariationsFor creates a specific node in the distributed RDataFrame graph. This acts as an action node. """ + pass class InstantAction(Operation): """An instant action attached to a distributed RDataFrame graph node.""" + pass class AsNumpy(InstantAction): """An 'AsNumpy' instant action attached to a distributed RDataFrame graph node.""" + pass class Snapshot(InstantAction): """A 'Snapshot' instant action attached to a distributed RDataFrame graph node.""" + pass class Transformation(Operation): """A trasformation attached to a distributed RDataFrame graph node.""" + pass @@ -121,6 +128,7 @@ class Transformation(Operation): "Profile2D": Action, "Profile3D": Action, "Redefine": Transformation, + "RedefinePerSample": Transformation, "Snapshot": Snapshot, "Stats": Action, "StdDev": Action, @@ -135,5 +143,7 @@ def create_op(name: str, *args, **kwargs) -> Union[Action, InstantAction, Transf try: return SUPPORTED_OPERATIONS[name](name, *args, **kwargs) except KeyError as e: - raise ValueError(f"Operation '{name}' is either invalid or not supported in distributed mode. " - "See the documentation for a list of supported operations.") from e + raise ValueError( + f"Operation '{name}' is either invalid or not supported in distributed mode. " + "See the documentation for a list of supported operations." + ) from e diff --git a/roottest/python/distrdf/backends/check_definepersample.py b/roottest/python/distrdf/backends/check_definepersample.py index ebfe2cd425a44..0d05f107aa62e 100644 --- a/roottest/python/distrdf/backends/check_definepersample.py +++ b/roottest/python/distrdf/backends/check_definepersample.py @@ -7,8 +7,7 @@ class TestDefinePerSample: """Check the working of merge operations in the reducer function.""" samples = ["sample1", "sample2", "sample3"] - filenames = [ - f"../data/ttree/distrdf_roottest_definepersample_{sample}.root" for sample in samples] + filenames = [f"../data/ttree/distrdf_roottest_definepersample_{sample}.root" for sample in samples] maintreename = "Events" def test_definepersample_simple(self, payload): @@ -47,7 +46,7 @@ def test_definepersample_withinitialization(self, payload): # needed functions available def declare_definepersample_code(): ROOT.gInterpreter.Declare( - ''' + """ #ifndef distrdf_test_definepersample_withinitialization #define distrdf_test_definepersample_withinitialization float sample1_weight(){ @@ -77,24 +76,61 @@ def declare_definepersample_code(): return id.AsString(); } #endif // distrdf_test_definepersample_withinitialization - ''') + """ + ) ROOT._distrdf.initialize(declare_definepersample_code) connection, _ = payload df = ROOT.RDataFrame(self.maintreename, self.filenames, executor=connection) - df1 = df.DefinePerSample("sample_weight", "samples_weights(rdfslot_, rdfsampleinfo_)")\ - .DefinePerSample("sample_name", "samples_names(rdfslot_, rdfsampleinfo_)") + df1 = df.DefinePerSample("sample_weight", "samples_weights(rdfslot_, rdfsampleinfo_)").DefinePerSample( + "sample_name", "samples_names(rdfslot_, rdfsampleinfo_)" + ) # Filter by the two defined columns per sample: a weight and the sample string representation # Each filtered dataset should have 10 entries, equal to the number of entries per sample weightsandnames = [ ("1.0f", f"{self.filenames[0]}/{self.maintreename}"), ("2.0f", f"{self.filenames[1]}/{self.maintreename}"), - ("3.0f", f"{self.filenames[2]}/{self.maintreename}") + ("3.0f", f"{self.filenames[2]}/{self.maintreename}"), ] samplescounts = [ - df1.Filter("sample_weight == {} && sample_name == \"{}\"".format(weight, name)).Count() - for (weight, name) in weightsandnames] + df1.Filter('sample_weight == {} && sample_name == "{}"'.format(weight, name)).Count() + for (weight, name) in weightsandnames + ] + + for count in samplescounts: + assert count.GetValue() == 10, f"{count.GetValue()=}" + + def test_redefinepersample_simple(self, payload): + """ + Test RedefinePerSample operation on three samples using a predefined + string of operations. + """ + + connection, _ = payload + df = ROOT.RDataFrame(self.maintreename, self.filenames, executor=connection) + + # Associate a number to each sample + definepersample_code = """ + if(rdfsampleinfo_.Contains(\"{}\")) return 1; + else if (rdfsampleinfo_.Contains(\"{}\")) return 2; + else if (rdfsampleinfo_.Contains(\"{}\")) return 3; + else return 0; + """.format(*self.samples) + + # Redefine the values with the same column name + redefinepersample_code = """ + if(rdfsampleinfo_.Contains(\"{}\")) return 11; + else if (rdfsampleinfo_.Contains(\"{}\")) return 22; + else if (rdfsampleinfo_.Contains(\"{}\")) return 33; + else return 0; + """.format(*self.samples) + + df1 = df.DefinePerSample("sampleid", definepersample_code).RedefinePerSample("sampleid", redefinepersample_code) + + # Filter by the sample number. Each filtered dataframe should contain + # 10 entries, equal to the number of entries per sample + samplescounts = [df1.Filter("sampleid == {}".format(id)).Count() for id in [11, 22, 33]] for count in samplescounts: assert count.GetValue() == 10, f"{count.GetValue()=}" diff --git a/tree/dataframe/inc/ROOT/RDF/RInterface.hxx b/tree/dataframe/inc/ROOT/RDF/RInterface.hxx index aa3bb93285d00..44bae42348e63 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. + /// \sa DefinePerSample. Works similarly, 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,19 +801,15 @@ 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); - - RDFInternal::RColumnRegister newCols(fColRegister); - newCols.AddDefine(std::move(jittedDefine)); - - RInterface newInterface(fProxiedPtr, *fLoopManager, std::move(newCols)); + return DefinePerSampleJitImpl(name, expression, false); + } - return newInterface; + //////////////////////////////////////////////////////////////////////////// + /// \brief Redefine an existing column that is updated when the input sample changes. + /// \sa DefinePerSample. Works similarly, 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); } /// \brief Register systematic variations for a single existing column using custom variation tags. @@ -3832,6 +3819,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 diff --git a/tree/dataframe/src/RDataFrame.cxx b/tree/dataframe/src/RDataFrame.cxx index ee6758cdc4f7d..254ff4d7996da 100644 --- a/tree/dataframe/src/RDataFrame.cxx +++ b/tree/dataframe/src/RDataFrame.cxx @@ -748,6 +748,7 @@ parts of the RDataFrame API currently work with this package. The subset that is - Min - Profile[1,2,3]D - Redefine +- RedefinePerSample - Snapshot - Stats - StdDev diff --git a/tree/dataframe/test/dataframe_definepersample.cxx b/tree/dataframe/test/dataframe_definepersample.cxx index bc276766026af..6f1fdfafdf382 100644 --- a/tree/dataframe/test/dataframe_definepersample.cxx +++ b/tree/dataframe/test/dataframe_definepersample.cxx @@ -140,6 +140,42 @@ 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; });