[DF] Add median support for RDataFrame

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

@@ -30,6 +30,9 @@
 #include "TObject.h"
 #include "ROOT/RDF/RActionImpl.hxx"
 #include "ROOT/RDF/RMergeableValue.hxx"
+#include <cmath>
+#include <cstddef>
+#include <cstdio>
 
 #include "RConfigure.h" // for R__HAS_ROOT7
 #ifdef R__HAS_ROOT7
@@ -1493,6 +1496,43 @@ public:
    }
 };
 
+class R__CLING_PTRCHECK(off) MedianHelper : public RActionImpl<MedianHelper> {
+   std::shared_ptr<double> fResult;
+   std::vector<std::vector<double>> fBuffers;
+
+public:
+   MedianHelper(const std::shared_ptr<double> &meanVPtr, const unsigned int nSlots);
+   MedianHelper(MedianHelper &&) = default;
+   MedianHelper &operator=(MedianHelper &&) = default;
+   MedianHelper(const MedianHelper &) = delete;
+   MedianHelper &operator=(const MedianHelper &other) = delete;
+
+   void InitTask(TTreeReader *, unsigned int) {}
+   void Exec(unsigned int slot, double v);
+
+   template <typename T, std::enable_if_t<IsDataContainer<T>::value, int> = 0>
+   void Exec(unsigned int slot, const T &vs)
+   {
+      fBuffers[slot].insert(fBuffers[slot].end(), std::begin(vs), std::end(vs));
+   }
+
+   void Initialize() { /* noop */ }
+
+   void Finalize();
+
+   std::string GetActionName() { return "Median"; }
+
+   std::shared_ptr<double> GetResultPtr() const { return fResult; }
+
+   MedianHelper MakeNew(void *newResult, std::string_view /*variation*/ = "nominal")
+   {
+      auto &result = *static_cast<std::shared_ptr<double> *>(newResult);
+      return MedianHelper(result, fBuffers.size());
+   }
+
+   ~MedianHelper() = default;
+};
+
 template <typename PrevNodeType>
 class R__CLING_PTRCHECK(off) DisplayHelper : public RActionImpl<DisplayHelper<PrevNodeType>> {
 private:

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

@@ -103,6 +103,7 @@ struct Sum{};
 struct Mean{};
 struct Fill{};
 struct StdDev{};
+struct Median{};
 struct Display{};
 struct Snapshot{};
 struct Book{};
@@ -299,6 +300,17 @@ BuildAction(const ColumnNames_t &bl, const std::shared_ptr<double> &stdDeviation
    return std::make_unique<Action_t>(Helper_t(stdDeviationV, nSlots), bl, prevNode, colRegister);
 }
 
+// Median action
+template <typename ColType, typename PrevNodeType>
+std::unique_ptr<RActionBase>
+BuildAction(const ColumnNames_t &bl, const std::shared_ptr<double> &meanV, const unsigned int nSlots,
+            std::shared_ptr<PrevNodeType> prevNode, ActionTags::Median, const RColumnRegister &colRegister)
+{
+   using Helper_t = MedianHelper;
+   using Action_t = RAction<Helper_t, PrevNodeType, TTraits::TypeList<ColType>>;
+   return std::make_unique<Action_t>(Helper_t(meanV, nSlots), bl, std::move(prevNode), colRegister);
+}
+
 using displayHelperArgs_t = std::pair<size_t, std::shared_ptr<ROOT::RDF::RDisplay>>;
 
 // Display action

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

@@ -2877,6 +2877,42 @@ public:
       return CreateAction<RDFInternal::ActionTags::StdDev, T>(userColumns, stdDeviationV, stdDeviationV, fProxiedPtr);
    }
 
+   ////////////////////////////////////////////////////////////////////////////
+   /// \brief Return the median value of the input column (*lazy action*).
+   /// \tparam T The type of the column.
+   /// \param[in] columnName The name of the column to be treated.
+   /// \return The median value of the selected column wrapped in a RResultPtr.
+   ///
+   /// If T is not specified, RDataFrame will infer it from the data and just-in-time compile the correct
+   /// template specialization of this method.
+   /// The result is always a double, irrespective of the type of column that is read. For an even
+   /// number of entries, the median is the average of the two middle values.
+   ///
+   /// \note If the column is empty, a signaling NaN is returned.
+   ///
+   /// \note Computing the exact median requires all column values to be held in memory at once, so the
+   /// memory used by this action grows with the number of processed entries. Consider this when running
+   /// over large datasets.
+   ///
+   /// This action is *lazy*: upon invocation of this method the calculation is
+   /// booked but not executed. Also see RResultPtr.
+   ///
+   /// ### Example usage:
+   /// ~~~{.cpp}
+   /// // Deduce column type (this invocation needs jitting internally)
+   /// auto medianVal0 = myDf.Median("values");
+   /// // Explicit column type
+   /// auto medianVal1 = myDf.Median<double>("values");
+   /// ~~~
+   ///
+   template <typename T = RDFDetail::RInferredType>
+   RResultPtr<double> Median(std::string_view columnName = "")
+   {
+      const auto userColumns = columnName.empty() ? ColumnNames_t() : ColumnNames_t({std::string(columnName)});
+      auto medianV = std::make_shared<double>(0);
+      return CreateAction<RDFInternal::ActionTags::Median, T>(userColumns, medianV, medianV, fProxiedPtr);
+   }
+
    // clang-format off
    ////////////////////////////////////////////////////////////////////////////
    /// \brief Return the sum of processed column values (*lazy action*).

tree/dataframe/src/RDFActionHelpers.cxx

@@ -13,6 +13,8 @@
 #include "ROOT/RDF/Utils.hxx" // CacheLineStep
 
 #include <RtypesCore.h>
+#include <cmath>
+#include <cstddef>
 #include <TStatistic.h>
 
 namespace ROOT {
@@ -220,6 +222,46 @@ void StdDevHelper::Finalize()
    *fResultStdDev = std::sqrt(variance);
 }
 
+MedianHelper::MedianHelper(const std::shared_ptr<double> &meanVPtr, const unsigned int nSlots)
+   : fResult(meanVPtr), fBuffers(nSlots, std::vector<double>())
+{
+}
+
+void MedianHelper::Exec(unsigned int slot, double v)
+{
+   fBuffers[slot].push_back(v);
+}
+
+void MedianHelper::Finalize()
+{
+   std::vector<double> fullBuffer;
+   std::size_t fullBufferSize = 0;
+   for (const auto &buffer : fBuffers) {
+      fullBufferSize += buffer.size();
+   }
+
+   if (fullBufferSize == 0) {
+      *fResult = std::numeric_limits<double>::signaling_NaN();
+      return;
+   }
+
+   fullBuffer.reserve(fullBufferSize);
+
+   for (const auto &buffer : fBuffers)
+      fullBuffer.insert(fullBuffer.end(), buffer.begin(), buffer.end());
+
+   const std::size_t k = fullBufferSize / 2;
+   std::nth_element(fullBuffer.begin(), fullBuffer.begin() + k, fullBuffer.end());
+   const double upper = fullBuffer[k];
+
+   if (fullBuffer.size() % 2) {
+      *fResult = upper;
+   } else {
+      const double lower = *std::max_element(fullBuffer.begin(), fullBuffer.begin() + k);
+      *fResult = 0.5 * (lower + upper);
+   }
+}
+
 // External templates are disabled for gcc5 since this version wrongly omits the C++11 ABI attribute
 #if __GNUC__ > 5
 template class TakeHelper<bool, bool, std::vector<bool>>;

tree/dataframe/test/dataframe_cloning.cxx

@@ -105,6 +105,18 @@ TEST(RDataFrameCloning, Mean)
    EXPECT_EQ(df.GetNRuns(), 1);
 }
 
+TEST(RDataFrameCloning, Median)
+{
+   ROOT::RDataFrame df{5};
+   // 1, 1, 1, 100, 100 -> median 1, mean 40.6
+   auto col1 = df.Define("x", [](ULong64_t e) { return e < 3 ? 1. : 100.; }, {"rdfentry_"});
+   auto median = col1.Median<double>("x");
+   auto clone = CloneResultAndAction(median);
+   EXPECT_DOUBLE_EQ(*median, 1.);
+   EXPECT_DOUBLE_EQ(*clone, 1.);
+   EXPECT_EQ(df.GetNRuns(), 1);
+}
+
 TEST(RDataFrameCloning, Histo1DNoModel)
 {
    ROOT::RDataFrame df{100};

tree/dataframe/test/dataframe_interface.cxx

@@ -3,9 +3,11 @@
 
 #include "ROOT/RCsvDS.hxx"
 #include "ROOT/RDataFrame.hxx"
+#include <cmath>
 #include <string_view>
 #include "ROOT/RTrivialDS.hxx"
 #include "ROOT/TestSupport.hxx"
+#include "RtypesCore.h"
 #include "TMemFile.h"
 #include "TSystem.h"
 #include "TTree.h"
@@ -532,7 +534,7 @@ TEST(RDataFrameInterface, ColumnWithSimpleStruct)
 }
 
 // Issue #6435
-TEST(RDataFrameInterface, MinMaxSumMeanStdDevOfScalar)
+TEST(RDataFrameInterface, MinMaxSumMeanStdDevMedianOfScalar)
 {
    auto df = ROOT::RDataFrame(4).Range(1, 0).Define("x", [](ULong64_t e) { return int(e); }, {"rdfentry_"});
    auto max = df.Max<int>("x");
@@ -545,6 +547,8 @@ TEST(RDataFrameInterface, MinMaxSumMeanStdDevOfScalar)
    auto jit_mean = df.Mean("x");
    auto stddev = df.StdDev<int>("x");
    auto jit_stddev = df.StdDev("x");
+   auto median = df.Median<int>("x");
+   auto jit_median = df.Median("x");
 
    EXPECT_EQ(*max, 3);
    EXPECT_DOUBLE_EQ(*jit_max, 3.f);
@@ -556,9 +560,11 @@ TEST(RDataFrameInterface, MinMaxSumMeanStdDevOfScalar)
    EXPECT_DOUBLE_EQ(*jit_mean, 2);
    EXPECT_DOUBLE_EQ(*stddev, 1.f);
    EXPECT_DOUBLE_EQ(*jit_stddev, 1.f);
+   EXPECT_DOUBLE_EQ(*median, 2);
+   EXPECT_DOUBLE_EQ(*jit_median, 2.);
 }
 
-TEST(RDataFrameInterface, MinMaxSumMeanStdDevOfRVec)
+TEST(RDataFrameInterface, MinMaxSumMeanStdDevMedianOfRVec)
 {
    auto df = ROOT::RDataFrame(1).Define("x", [] { return ROOT::RVec<int>{1,2,3}; });
    auto max = df.Max<ROOT::RVec<int>>("x");
@@ -571,6 +577,8 @@ TEST(RDataFrameInterface, MinMaxSumMeanStdDevOfRVec)
    auto jit_mean = df.Mean("x");
    auto stddev = df.StdDev<ROOT::RVec<int>>("x");
    auto jit_stddev = df.StdDev("x");
+   auto median = df.Median<ROOT::RVec<int>>("x");
+   auto jit_median = df.Median("x");
 
    EXPECT_EQ(*max, 3);
    EXPECT_DOUBLE_EQ(*jit_max, 3.f);
@@ -582,6 +590,59 @@ TEST(RDataFrameInterface, MinMaxSumMeanStdDevOfRVec)
    EXPECT_DOUBLE_EQ(*jit_mean, 2);
    EXPECT_DOUBLE_EQ(*stddev, 1.f);
    EXPECT_DOUBLE_EQ(*jit_stddev, 1.f);
+   EXPECT_DOUBLE_EQ(*median, 2);
+   EXPECT_DOUBLE_EQ(*jit_median, 2);
+}
+
+TEST(RDataFrameInterface, MedianOfScalarAndEvenNumberOfEntries)
+{
+   // 1, 1, 2, 100 -> two middle values are 1 and 2, so median is 1.5
+   auto df = ROOT::RDataFrame(4).Define("x",
+                                        [](ULong64_t e) {
+                                           if (e == 2)
+                                              return 2;
+                                           if (e == 3)
+                                              return 100;
+                                           return 1;
+                                        },
+                                        {"rdfentry_"});
+   EXPECT_DOUBLE_EQ(*df.Median<int>("x"), 1.5);
+   EXPECT_DOUBLE_EQ(*df.Median("x"), 1.5);
+}
+
+TEST(RDataFrameInterface, MedianOfOddNumberOfEntries)
+{
+   // 1, 1, 1, 100, 100 -> median 1, mean 40.6
+   auto df = ROOT::RDataFrame(5).Define("x", [](ULong64_t e) { return e < 3 ? 1 : 100; }, {"rdfentry_"});
+   EXPECT_DOUBLE_EQ(*df.Median<int>("x"), 1.);
+   EXPECT_DOUBLE_EQ(*df.Median("x"), 1.);
+}
+
+TEST(RDataFrameInterface, MedianOfEmptyDataFrameScalar)
+{
+   auto df = ROOT::RDataFrame(0).Define("x", "1.0");
+   EXPECT_TRUE(std::isnan(*df.Median("x")));
+}
+
+TEST(RDataFrameInterface, MedianOfEmptyDataFrameRVec)
+{
+   auto df = ROOT::RDataFrame(0).Define("x", [] { return ROOT::RVec<int>{}; });
+   EXPECT_TRUE(std::isnan(*df.Median("x")));
+}
+
+TEST(RDataFrameInterface, MedianWithImplicitMT)
+{
+#ifdef R__USE_IMT
+   const unsigned int nslots = std::min(4U, std::thread::hardware_concurrency());
+   if (ROOT::IsImplicitMTEnabled())
+      ROOT::DisableImplicitMT();
+   ROOT::EnableImplicitMT(nslots);
+   // 80% ones, 20% hundreds -> median 1, mean 20.8
+   auto df = ROOT::RDataFrame(100000).Define("x", [](ULong64_t e) { return e < 80000 ? 1. : 100.; }, {"rdfentry_"});
+   EXPECT_DOUBLE_EQ(*df.Median<double>("x"), 1.);
+   EXPECT_DOUBLE_EQ(*df.Median("x"), 1.);
+   ROOT::DisableImplicitMT();
+#endif
 }
 
 class Product {

tree/dataframe/test/dataframe_simple.cxx

@@ -189,8 +189,22 @@ TEST_P(RDFSimpleTests, Define_jitted)
    EXPECT_DOUBLE_EQ(1., *m);
 }
 
+TEST_P(RDFSimpleTests, Median)
+{
+   // 80% ones, 20% hundreds -> median 1, mean 20.8
+   auto df = ROOT::RDataFrame(100000).Define("x", [](ULong64_t e) { return e < 80000 ? 1. : 100.; }, {"rdfentry_"});
+   EXPECT_DOUBLE_EQ(df.Median<double>("x").GetValue(), 1.);
+}
+
 struct RFoo {};
 
+TEST_P(RDFSimpleTests, MedianOfNonArithmeticColumn)
+{
+   RDataFrame tdf(1);
+   auto d = tdf.Define("foo", []() { return RFoo(); });
+   EXPECT_ANY_THROW(d.Median("foo").GetValue());
+}
+
 TEST_P(RDFSimpleTests, Define_jitted_type_unknown_to_interpreter)
 {
    RDataFrame tdf(10);