From 97cca4cdcfe444179e1d53f6732f07337cdad39e Mon Sep 17 00:00:00 2001 From: Arjun Varma <‘arj25835@gmail.com’> Date: Sat, 15 Aug 2026 15:52:49 +0530 Subject: [PATCH] [df] Introduce Median operation This commit introduces the Median lazy action, which computes the median of the values in the input column. Median always returns an RResultPtr irrespective of the input column type. If the input column is of a type which does not make sense to compute the median on, the action will throw a std::runtime_error during execution when the input value is tried to be coerced into a double. The current implementation works for local RDataFrame in single- and multi-threaded mode. It stores all column values in a buffer and computes the exact median at the end of processing. --- tree/dataframe/inc/ROOT/RDF/ActionHelpers.hxx | 40 ++++++++++++ .../dataframe/inc/ROOT/RDF/InterfaceUtils.hxx | 12 ++++ tree/dataframe/inc/ROOT/RDF/RInterface.hxx | 36 ++++++++++ tree/dataframe/src/RDFActionHelpers.cxx | 42 ++++++++++++ tree/dataframe/test/dataframe_cloning.cxx | 12 ++++ tree/dataframe/test/dataframe_interface.cxx | 65 ++++++++++++++++++- tree/dataframe/test/dataframe_simple.cxx | 14 ++++ 7 files changed, 219 insertions(+), 2 deletions(-) diff --git a/tree/dataframe/inc/ROOT/RDF/ActionHelpers.hxx b/tree/dataframe/inc/ROOT/RDF/ActionHelpers.hxx index 941576427c63e..99ae2942e7e56 100644 --- a/tree/dataframe/inc/ROOT/RDF/ActionHelpers.hxx +++ b/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 +#include +#include #include "RConfigure.h" // for R__HAS_ROOT7 #ifdef R__HAS_ROOT7 @@ -1493,6 +1496,43 @@ public: } }; +class R__CLING_PTRCHECK(off) MedianHelper : public RActionImpl { + std::shared_ptr fResult; + std::vector> fBuffers; + +public: + MedianHelper(const std::shared_ptr &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 ::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 GetResultPtr() const { return fResult; } + + MedianHelper MakeNew(void *newResult, std::string_view /*variation*/ = "nominal") + { + auto &result = *static_cast *>(newResult); + return MedianHelper(result, fBuffers.size()); + } + + ~MedianHelper() = default; +}; + template class R__CLING_PTRCHECK(off) DisplayHelper : public RActionImpl> { private: diff --git a/tree/dataframe/inc/ROOT/RDF/InterfaceUtils.hxx b/tree/dataframe/inc/ROOT/RDF/InterfaceUtils.hxx index 7b241022750d7..68bc9d48f9492 100644 --- a/tree/dataframe/inc/ROOT/RDF/InterfaceUtils.hxx +++ b/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 &stdDeviation return std::make_unique(Helper_t(stdDeviationV, nSlots), bl, prevNode, colRegister); } +// Median action +template +std::unique_ptr +BuildAction(const ColumnNames_t &bl, const std::shared_ptr &meanV, const unsigned int nSlots, + std::shared_ptr prevNode, ActionTags::Median, const RColumnRegister &colRegister) +{ + using Helper_t = MedianHelper; + using Action_t = RAction>; + return std::make_unique(Helper_t(meanV, nSlots), bl, std::move(prevNode), colRegister); +} + using displayHelperArgs_t = std::pair>; // Display action diff --git a/tree/dataframe/inc/ROOT/RDF/RInterface.hxx b/tree/dataframe/inc/ROOT/RDF/RInterface.hxx index 92348b0cac694..3a7257b6df869 100644 --- a/tree/dataframe/inc/ROOT/RDF/RInterface.hxx +++ b/tree/dataframe/inc/ROOT/RDF/RInterface.hxx @@ -2877,6 +2877,42 @@ public: return CreateAction(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("values"); + /// ~~~ + /// + template + RResultPtr Median(std::string_view columnName = "") + { + const auto userColumns = columnName.empty() ? ColumnNames_t() : ColumnNames_t({std::string(columnName)}); + auto medianV = std::make_shared(0); + return CreateAction(userColumns, medianV, medianV, fProxiedPtr); + } + // clang-format off //////////////////////////////////////////////////////////////////////////// /// \brief Return the sum of processed column values (*lazy action*). diff --git a/tree/dataframe/src/RDFActionHelpers.cxx b/tree/dataframe/src/RDFActionHelpers.cxx index dc88ab2711896..ff93c65402077 100644 --- a/tree/dataframe/src/RDFActionHelpers.cxx +++ b/tree/dataframe/src/RDFActionHelpers.cxx @@ -13,6 +13,8 @@ #include "ROOT/RDF/Utils.hxx" // CacheLineStep #include +#include +#include #include namespace ROOT { @@ -220,6 +222,46 @@ void StdDevHelper::Finalize() *fResultStdDev = std::sqrt(variance); } +MedianHelper::MedianHelper(const std::shared_ptr &meanVPtr, const unsigned int nSlots) + : fResult(meanVPtr), fBuffers(nSlots, std::vector()) +{ +} + +void MedianHelper::Exec(unsigned int slot, double v) +{ + fBuffers[slot].push_back(v); +} + +void MedianHelper::Finalize() +{ + std::vector fullBuffer; + std::size_t fullBufferSize = 0; + for (const auto &buffer : fBuffers) { + fullBufferSize += buffer.size(); + } + + if (fullBufferSize == 0) { + *fResult = std::numeric_limits::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>; diff --git a/tree/dataframe/test/dataframe_cloning.cxx b/tree/dataframe/test/dataframe_cloning.cxx index 843bd39589d3c..e0da18c2af847 100644 --- a/tree/dataframe/test/dataframe_cloning.cxx +++ b/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("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}; diff --git a/tree/dataframe/test/dataframe_interface.cxx b/tree/dataframe/test/dataframe_interface.cxx index 3411403ffcba5..2a304fb7947b8 100644 --- a/tree/dataframe/test/dataframe_interface.cxx +++ b/tree/dataframe/test/dataframe_interface.cxx @@ -3,9 +3,11 @@ #include "ROOT/RCsvDS.hxx" #include "ROOT/RDataFrame.hxx" +#include #include #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("x"); @@ -545,6 +547,8 @@ TEST(RDataFrameInterface, MinMaxSumMeanStdDevOfScalar) auto jit_mean = df.Mean("x"); auto stddev = df.StdDev("x"); auto jit_stddev = df.StdDev("x"); + auto median = df.Median("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{1,2,3}; }); auto max = df.Max>("x"); @@ -571,6 +577,8 @@ TEST(RDataFrameInterface, MinMaxSumMeanStdDevOfRVec) auto jit_mean = df.Mean("x"); auto stddev = df.StdDev>("x"); auto jit_stddev = df.StdDev("x"); + auto median = df.Median>("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("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("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{}; }); + 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("x"), 1.); + EXPECT_DOUBLE_EQ(*df.Median("x"), 1.); + ROOT::DisableImplicitMT(); +#endif } class Product { diff --git a/tree/dataframe/test/dataframe_simple.cxx b/tree/dataframe/test/dataframe_simple.cxx index 6e1461e3f9a7b..e5e5f94b98cac 100644 --- a/tree/dataframe/test/dataframe_simple.cxx +++ b/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("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);