[df] RedefinePerSample

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

@@ -749,25 +749,16 @@ public:
    template <typename F, typename RetType_t = typename TTraits::CallableTraits<F>::ret_type>
    RInterface<Proxied> 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<RDFDetail::RDefinePerSample<F>>(name, retTypeName, std::move(expression), *fLoopManager);
+      return DefinePerSampleImpl<F, RetType_t>(name, std::move(expression), false);
+   }
 
-      RDFInternal::RColumnRegister newCols(fColRegister);
-      newCols.AddDefine(std::move(newColumn));
-      RInterface<Proxied> 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 <typename F, typename RetType_t = typename TTraits::CallableTraits<F>::ret_type>
+   RInterface<Proxied> RedefinePerSample(std::string_view name, F expression)
+   {
+      return DefinePerSampleImpl<F, RetType_t>(name, std::move(expression), true);
    }
 
    // clang-format off
@@ -810,20 +801,18 @@ public:
    // clang-format on
    RInterface<Proxied> 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<Proxied> RedefinePerSample(std::string_view name, std::string_view expression)
+   {
+      return DefinePerSampleJitImpl(name, expression, true);
+   }
 
-      RInterface<Proxied> 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 <typename F, typename RetType_t = typename TTraits::CallableTraits<F>::ret_type>
+   RInterface<Proxied> 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<RDFDetail::RDefinePerSample<F>>(name, retTypeName, std::move(expression), *fLoopManager);
+
+      RDFInternal::RColumnRegister newCols(fColRegister);
+      newCols.AddDefine(std::move(newColumn));
+      RInterface<Proxied> newInterface(fProxiedPtr, *fLoopManager, std::move(newCols));
+      return newInterface;
+   }
+
+   ////////////////////////////////////////////////////////////////////////////
+   /// \brief Implementation of DefinePerSample and RedefinePerSample (jitted).
+   RInterface<Proxied> 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<Proxied> newInterface(fProxiedPtr, *fLoopManager, std::move(newCols));
+
+      return newInterface;
+   }
+
    ////////////////////////////////////////////////////////////////////////////
    /// \brief Implementation of cache.
    template <typename... ColTypes, std::size_t... S>

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<int>("x");
+   auto xmax = df2.Max<int>("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; });