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);