[df] Handle containers in RHist* filling

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

@@ -325,6 +325,44 @@ std::size_t GetSize(const T &val)
    }
 }
 
+// trait class to implement looping over data containers
+template <typename Helper>
+class R__CLING_PTRCHECK(off) ExecLoopTrait {
+private:
+   template <typename... Iterators>
+   void ExecLoop(unsigned int slot, std::size_t elements, Iterators... its)
+   {
+      for (std::size_t i = 0; i < elements; i++) {
+         Exec(slot, *its...);
+         (std::advance(its, 1), ...);
+      }
+   }
+
+public:
+   template <typename... ColumnTypes>
+   void Exec(unsigned int slot, const ColumnTypes &...columnValues)
+   {
+      if constexpr (std::disjunction_v<IsDataContainer<ColumnTypes>...>) {
+         constexpr std::array<bool, sizeof...(ColumnTypes)> isContainer{IsDataContainer<ColumnTypes>::value...};
+         constexpr std::size_t firstContainerIdx = FindIdxTrue(isContainer);
+         std::array<std::size_t, sizeof...(columnValues)> sizes = {{GetSize(columnValues)...}};
+         std::size_t elements = 0;
+         for (std::size_t i = 0; i < isContainer.size(); i++) {
+            if (isContainer[i]) {
+               if (i == firstContainerIdx) {
+                  elements = sizes[i];
+               } else if (elements != sizes[i]) {
+                  throw std::runtime_error("Cannot fill values in containers of different sizes.");
+               }
+            }
+         }
+         ExecLoop(slot, elements, MakeBegin(columnValues)...);
+      } else {
+         static_cast<Helper *>(this)->ExecSingle(slot, columnValues...);
+      }
+   }
+};
+
 // Helpers for dealing with histograms and similar:
 template <typename H, typename = decltype(std::declval<H>().Reset())>
 void ResetIfPossible(H *h)
@@ -485,8 +523,8 @@ public:
 
 #ifdef R__HAS_ROOT7
 template <typename BinContentType, bool WithWeight = false>
-class R__CLING_PTRCHECK(off) RHistFillHelper
-   : public ROOT::Detail::RDF::RActionImpl<RHistFillHelper<BinContentType, WithWeight>> {
+class R__CLING_PTRCHECK(off) RHistFillHelper : public RActionImpl<RHistFillHelper<BinContentType, WithWeight>>,
+                                               public ExecLoopTrait<RHistFillHelper<BinContentType, WithWeight>> {
 public:
    using Result_t = ROOT::Experimental::RHist<BinContentType>;
 
@@ -524,7 +562,7 @@ public:
    }
 
    template <typename... ColumnTypes>
-   void Exec(unsigned int slot, const ColumnTypes &...columnValues)
+   void ExecSingle(unsigned int slot, const ColumnTypes &...columnValues)
    {
       if constexpr (WithWeight) {
          auto t = std::forward_as_tuple(columnValues...);
@@ -546,7 +584,8 @@ public:
 
 template <typename BinContentType, bool WithWeight = false>
 class R__CLING_PTRCHECK(off) RHistEngineFillHelper
-   : public ROOT::Detail::RDF::RActionImpl<RHistEngineFillHelper<BinContentType, WithWeight>> {
+   : public RActionImpl<RHistEngineFillHelper<BinContentType, WithWeight>>,
+     public ExecLoopTrait<RHistEngineFillHelper<BinContentType, WithWeight>> {
 public:
    using Result_t = ROOT::Experimental::RHistEngine<BinContentType>;
 
@@ -576,7 +615,7 @@ public:
    }
 
    template <typename... ColumnTypes>
-   void Exec(unsigned int, const ColumnTypes &...columnValues)
+   void ExecSingle(unsigned int, const ColumnTypes &...columnValues)
    {
       if constexpr (WithWeight) {
          auto t = std::forward_as_tuple(columnValues...);

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

@@ -2032,6 +2032,9 @@ public:
    /// \param[in] vName The name of the column that will fill the histogram.
    /// \return the histogram wrapped in a RResultPtr.
    ///
+   /// The column can be of a container type (e.g. `std::vector` or `RVec`), in which case the histogram is filled with
+   /// each one of the elements of the container, and the Container can also be nested.
+   ///
    /// This action is *lazy*: upon invocation of this method the calculation is
    /// booked but not executed. Also see RResultPtr.
    ///
@@ -2057,6 +2060,11 @@ public:
    /// \param[in] columnList A list containing the names of the columns that will be passed when calling `Fill`
    /// \return the histogram wrapped in a RResultPtr.
    ///
+   /// Columns can be of a container type (e.g. `std::vector` or `RVec`), in which case the histogram is filled with
+   /// each one of the elements of the container. In case multiple columns of container type are provided, they must
+   /// have the same length for each event (but possibly different lengths between events). Containers can be nested,
+   /// in which case their sizes must match recursively. Scalars are broadcasted to match container columns.
+   ///
    /// This action is *lazy*: upon invocation of this method the calculation is
    /// booked but not executed. Also see RResultPtr.
    ///
@@ -2086,6 +2094,11 @@ public:
    /// \param[in] columnList A list containing the names of the columns that will be passed when calling `Fill`
    /// \return the histogram wrapped in a RResultPtr.
    ///
+   /// Columns can be of a container type (e.g. `std::vector` or `RVec`), in which case the histogram is filled with
+   /// each one of the elements of the container. In case multiple columns of container type are provided, they must
+   /// have the same length for each event (but possibly different lengths between events). Containers can be nested,
+   /// in which case their sizes must match recursively. Scalars are broadcasted to match container columns.
+   ///
    /// This action is *lazy*: upon invocation of this method the calculation is
    /// booked but not executed. Also see RResultPtr.
    ///
@@ -2122,6 +2135,11 @@ public:
    /// \param[in] wName The name of the column that will provide the weights.
    /// \return the histogram wrapped in a RResultPtr.
    ///
+   /// Columns can be of a container type (e.g. `std::vector` or `RVec`), in which case the histogram is filled with
+   /// each one of the elements of the container. In case multiple columns of container type are provided, they must
+   /// have the same length for each event (but possibly different lengths between events). Containers can be nested,
+   /// in which case their sizes must match recursively. Scalars are broadcasted to match container columns.
+   ///
    /// This action is *lazy*: upon invocation of this method the calculation is
    /// booked but not executed. Also see RResultPtr.
    ///
@@ -2149,6 +2167,11 @@ public:
    /// \param[in] wName The name of the column that will provide the weights.
    /// \return the histogram wrapped in a RResultPtr.
    ///
+   /// Columns can be of a container type (e.g. `std::vector` or `RVec`), in which case the histogram is filled with
+   /// each one of the elements of the container. In case multiple columns of container type are provided, they must
+   /// have the same length for each event (but possibly different lengths between events). Containers can be nested,
+   /// in which case their sizes must match recursively. Scalars are broadcasted to match container columns.
+   ///
    /// This action is *lazy*: upon invocation of this method the calculation is
    /// booked but not executed. Also see RResultPtr.
    ///
@@ -2185,6 +2208,11 @@ public:
    /// \param[in] wName The name of the column that will provide the weights.
    /// \return the histogram wrapped in a RResultPtr.
    ///
+   /// Columns can be of a container type (e.g. `std::vector` or `RVec`), in which case the histogram is filled with
+   /// each one of the elements of the container. In case multiple columns of container type are provided, they must
+   /// have the same length for each event (but possibly different lengths between events). Containers can be nested,
+   /// in which case their sizes must match recursively. Scalars are broadcasted to match container columns.
+   ///
    /// This action is *lazy*: upon invocation of this method the calculation is
    /// booked but not executed. Also see RResultPtr.
    ///
@@ -2228,6 +2256,11 @@ public:
    /// \param[in] columnList A list containing the names of the columns that will be passed when calling `Fill`
    /// \return the histogram wrapped in a RResultPtr.
    ///
+   /// Columns can be of a container type (e.g. `std::vector` or `RVec`), in which case the histogram is filled with
+   /// each one of the elements of the container. In case multiple columns of container type are provided, they must
+   /// have the same length for each event (but possibly different lengths between events). Containers can be nested,
+   /// in which case their sizes must match recursively. Scalars are broadcasted to match container columns.
+   ///
    /// This action is *lazy*: upon invocation of this method the calculation is
    /// booked but not executed. Also see RResultPtr.
    ///
@@ -2262,6 +2295,11 @@ public:
    /// \param[in] wName The name of the column that will provide the weights.
    /// \return the histogram wrapped in a RResultPtr.
    ///
+   /// Columns can be of a container type (e.g. `std::vector` or `RVec`), in which case the histogram is filled with
+   /// each one of the elements of the container. In case multiple columns of container type are provided, they must
+   /// have the same length for each event (but possibly different lengths between events). Containers can be nested,
+   /// in which case their sizes must match recursively. Scalars are broadcasted to match container columns.
+   ///
    /// This action is *lazy*: upon invocation of this method the calculation is
    /// booked but not executed. Also see RResultPtr.
    ///

tree/dataframe/test/dataframe_helpers.cxx

@@ -10,8 +10,10 @@
 #include <algorithm>
 #include <deque>
 #include <iomanip>
-#include <vector>
+#include <stdexcept>
 #include <string>
+#include <utility>
+#include <vector>
 
 #include "gtest/gtest.h"
 
@@ -813,6 +815,147 @@ TEST(RDFHelpers, Cleanup_After_Exception)
       << "The Finalize method should have changed the value of testVal during the post-exception cleanup." << std::endl;
 }
 
+struct SimpleExecLoopAction : public ROOT::Internal::RDF::ExecLoopTrait<SimpleExecLoopAction> {
+   std::vector<double> fValues;
+
+   template <typename... ColumnTypes>
+   void ExecSingle(unsigned int, const ColumnTypes &...columnValues)
+   {
+      fValues.emplace_back(columnValues...);
+   }
+};
+
+TEST(RDFHelpers, ExecLoopTrait)
+{
+   SimpleExecLoopAction execLoop;
+
+   execLoop.Exec(/*slot=*/0, 42.0);
+   ASSERT_EQ(execLoop.fValues.size(), 1);
+   EXPECT_EQ(execLoop.fValues.front(), 42.0);
+   execLoop.fValues.clear();
+
+   // One-dimensional containers are processed in order.
+   std::vector<double> v = {43.0};
+   execLoop.Exec(/*slot=*/0, v);
+   EXPECT_EQ(execLoop.fValues.size(), 1);
+   EXPECT_EQ(execLoop.fValues, v);
+   execLoop.fValues.clear();
+
+   v = {44.0, 45.0};
+   execLoop.Exec(/*slot=*/0, v);
+   EXPECT_EQ(execLoop.fValues.size(), 2);
+   EXPECT_EQ(execLoop.fValues, v);
+   execLoop.fValues.clear();
+
+   v = {};
+   execLoop.Exec(/*slot=*/0, v);
+   EXPECT_TRUE(execLoop.fValues.empty());
+   execLoop.fValues.clear();
+
+   // Multi-dimensional containers are traversed recursively.
+   std::vector<std::vector<double>> v2 = {{1.0}, {}, {2.0, 3.0}, {4.0}};
+   execLoop.Exec(/*slot=*/0, v2);
+   EXPECT_EQ(execLoop.fValues.size(), 4);
+   v = {1.0, 2.0, 3.0, 4.0};
+   EXPECT_EQ(execLoop.fValues, v);
+   execLoop.fValues.clear();
+
+   std::vector<std::vector<std::vector<double>>> v3 = {{{1.0}, {2.0}}, {}, {{3.0, 4.0}}, {{5.0}}};
+   execLoop.Exec(/*slot=*/0, v3);
+   EXPECT_EQ(execLoop.fValues.size(), 5);
+   v = {1.0, 2.0, 3.0, 4.0, 5.0};
+   EXPECT_EQ(execLoop.fValues, v);
+   execLoop.fValues.clear();
+}
+
+struct ExecLoopActionTwo : public ROOT::Internal::RDF::ExecLoopTrait<ExecLoopActionTwo> {
+   std::vector<std::pair<double, double>> fValues;
+
+   template <typename... ColumnTypes>
+   void ExecSingle(unsigned int, const ColumnTypes &...columnValues)
+   {
+      fValues.emplace_back(columnValues...);
+   }
+};
+
+TEST(RDFHelpers, ExecLoopTraitTwo)
+{
+   ExecLoopActionTwo execLoop;
+
+   execLoop.Exec(/*slot=*/0, 42.0, 43.0);
+   ASSERT_EQ(execLoop.fValues.size(), 1);
+   EXPECT_EQ(execLoop.fValues.front(), std::make_pair(42.0, 43.0));
+   execLoop.fValues.clear();
+
+   // Two containers are traversed in sync.
+   std::vector<double> v1 = {1.0, 2.0};
+   std::vector<double> v2 = {3.0, 4.0};
+   execLoop.Exec(/*slot=*/0, v1, v2);
+   ASSERT_EQ(execLoop.fValues.size(), 2);
+   EXPECT_EQ(execLoop.fValues.at(0), std::make_pair(1.0, 3.0));
+   EXPECT_EQ(execLoop.fValues.at(1), std::make_pair(2.0, 4.0));
+   execLoop.fValues.clear();
+
+   // A scalar is broadcasted to match the containers.
+   execLoop.Exec(/*slot=*/0, v1, 42.0);
+   ASSERT_EQ(execLoop.fValues.size(), 2);
+   EXPECT_EQ(execLoop.fValues.at(0), std::make_pair(1.0, 42.0));
+   EXPECT_EQ(execLoop.fValues.at(1), std::make_pair(2.0, 42.0));
+   execLoop.fValues.clear();
+
+   execLoop.Exec(/*slot=*/0, 42.0, v2);
+   ASSERT_EQ(execLoop.fValues.size(), 2);
+   EXPECT_EQ(execLoop.fValues.at(0), std::make_pair(42.0, 3.0));
+   EXPECT_EQ(execLoop.fValues.at(1), std::make_pair(42.0, 4.0));
+   execLoop.fValues.clear();
+
+   // Two nested containers are traversed in sync.
+   std::vector<std::vector<double>> v3 = {{1.0, 2.0}, {3.0}};
+   std::vector<std::vector<double>> v4 = {{4.0, 5.0}, {6.0}};
+   execLoop.Exec(/*slot=*/0, v3, v4);
+   ASSERT_EQ(execLoop.fValues.size(), 3);
+   EXPECT_EQ(execLoop.fValues.at(0), std::make_pair(1.0, 4.0));
+   EXPECT_EQ(execLoop.fValues.at(1), std::make_pair(2.0, 5.0));
+   EXPECT_EQ(execLoop.fValues.at(2), std::make_pair(3.0, 6.0));
+   execLoop.fValues.clear();
+
+   // Broadcasting works recursively in multiple dimensions.
+   execLoop.Exec(/*slot=*/0, v3, 42.0);
+   ASSERT_EQ(execLoop.fValues.size(), 3);
+   EXPECT_EQ(execLoop.fValues.at(0), std::make_pair(1.0, 42.0));
+   EXPECT_EQ(execLoop.fValues.at(1), std::make_pair(2.0, 42.0));
+   EXPECT_EQ(execLoop.fValues.at(2), std::make_pair(3.0, 42.0));
+   execLoop.fValues.clear();
+
+   execLoop.Exec(/*slot=*/0, v1, v3);
+   ASSERT_EQ(execLoop.fValues.size(), 3);
+   EXPECT_EQ(execLoop.fValues.at(0), std::make_pair(1.0, 1.0));
+   EXPECT_EQ(execLoop.fValues.at(1), std::make_pair(1.0, 2.0));
+   EXPECT_EQ(execLoop.fValues.at(2), std::make_pair(2.0, 3.0));
+   execLoop.fValues.clear();
+}
+
+TEST(RDFHelpers, ExecLoopTraitInvalid)
+{
+   ExecLoopActionTwo execLoop;
+
+   // The number of elements has to match for all columns.
+   std::vector<double> v0;
+   std::vector<double> v1 = {1.0};
+   std::vector<double> v2 = {2.0, 3.0};
+   EXPECT_THROW(execLoop.Exec(/*slot=*/0, v0, v1), std::runtime_error);
+   EXPECT_THROW(execLoop.Exec(/*slot=*/0, v0, v2), std::runtime_error);
+   EXPECT_THROW(execLoop.Exec(/*slot=*/0, v2, v1), std::runtime_error);
+
+   // For multiple dimensions, the number of elements has to match recursively.
+   std::vector<std::vector<double>> v4 = {{1.0}, {2.0}};
+   std::vector<std::vector<double>> v5 = {{3.0, 4.0}};
+   std::vector<std::vector<double>> v6 = {{3.0, 4.0}, {}};
+   EXPECT_THROW(execLoop.Exec(/*slot=*/0, v4, v5), std::runtime_error);
+   EXPECT_THROW(execLoop.Exec(/*slot=*/0, v4, v6), std::runtime_error);
+   EXPECT_THROW(execLoop.Exec(/*slot=*/0, v6, v5), std::runtime_error);
+}
+
 TEST(RDFHelpers, ProgressBarRestorePrecision)
 {
    // Regression test for https://github.com/root-project/root/issues/21165