Add tdigest_reduce and tdigest_quantile

Author: Rafał Hibner

cpp/src/arrow/compute/api_aggregate.cc

@@ -150,7 +150,8 @@ static auto kTDigestReduceOptionsType = GetFunctionOptionsType<TDigestReduceOpti
     DataMember("scaler", &TDigestReduceOptions::scaler));
 static auto kTDigestQuantileOptionsType = GetFunctionOptionsType<TDigestQuantileOptions>(
     DataMember("q", &TDigestQuantileOptions::q),
-    DataMember("min_count", &TDigestQuantileOptions::min_count));
+    DataMember("min_count", &TDigestQuantileOptions::min_count),
+    DataMember("scaler", &TDigestQuantileOptions::scaler));
 static auto kPivotOptionsType = GetFunctionOptionsType<PivotWiderOptions>(
     DataMember("key_names", &PivotWiderOptions::key_names),
     DataMember("unexpected_key_behavior", &PivotWiderOptions::unexpected_key_behavior));
@@ -239,15 +240,19 @@ TDigestReduceOptions::TDigestReduceOptions(Scaler scaler)
     : FunctionOptions(internal::kTDigestReduceOptionsType), scaler{scaler} {}
 constexpr char TDigestReduceOptions::kTypeName[];
 
-TDigestQuantileOptions::TDigestQuantileOptions(double q, uint32_t min_count)
+TDigestQuantileOptions::TDigestQuantileOptions(double q, uint32_t min_count,
+                                               Scaler scaler)
     : FunctionOptions(internal::kTDigestQuantileOptionsType),
       q{q},
-      min_count{min_count} {}
+      min_count{min_count},
+      scaler{scaler} {}
 
-TDigestQuantileOptions::TDigestQuantileOptions(std::vector<double> q, uint32_t min_count)
+TDigestQuantileOptions::TDigestQuantileOptions(std::vector<double> q, uint32_t min_count,
+                                               Scaler scaler)
     : FunctionOptions(internal::kTDigestQuantileOptionsType),
       q{std::move(q)},
-      min_count{min_count} {}
+      min_count{min_count},
+      scaler{scaler} {}
 constexpr char TDigestReduceOptions::kTypeName[];
 
 PivotWiderOptions::PivotWiderOptions(std::vector<std::string> key_names,
@@ -363,6 +368,16 @@ Result<Datum> TDigestMap(const Datum& value, const TDigestMapOptions& options,
   return CallFunction("tdigest_map", {value}, &options, ctx);
 }
 
+Result<Datum> TDigestReduce(const Datum& value, const TDigestReduceOptions& options,
+                            ExecContext* ctx) {
+  return CallFunction("tdigest_reduce", {value}, &options, ctx);
+}
+
+Result<Datum> TDigestQuantile(const Datum& value, const TDigestQuantileOptions& options,
+                              ExecContext* ctx) {
+  return CallFunction("tdigest_quantile", {value}, &options, ctx);
+}
+
 Result<Datum> Index(const Datum& value, const IndexOptions& options, ExecContext* ctx) {
   return CallFunction("index", {value}, &options, ctx);
 }

cpp/src/arrow/compute/api_aggregate.h

@@ -244,15 +244,21 @@ class ARROW_EXPORT TDigestReduceOptions : public FunctionOptions {
 /// By default, returns the median value.
 class ARROW_EXPORT TDigestQuantileOptions : public FunctionOptions {
  public:
-  explicit TDigestQuantileOptions(double q = 0.5, uint32_t min_count = 0);
-  explicit TDigestQuantileOptions(std::vector<double> q, uint32_t min_count = 0);
+  using Scaler = TDigestOptions::Scaler;
+
+  explicit TDigestQuantileOptions(double q = 0.5, uint32_t min_count = 0,
+                                  Scaler scaler = Scaler::K1);
+  explicit TDigestQuantileOptions(std::vector<double> q, uint32_t min_count = 0,
+                                  Scaler scaler = Scaler::K1);
   static constexpr char const kTypeName[] = "TDigestQuantileOptions";
   static TDigestQuantileOptions Defaults() { return TDigestQuantileOptions{}; }
 
   /// probability level of quantile must be between 0 and 1 inclusive
   std::vector<double> q;
   /// If less than this many non-null values are observed, emit null.
   uint32_t min_count;
+  /// select scaler implementation
+  Scaler scaler;
 };
 
 /// \brief Control Pivot kernel behavior
@@ -643,7 +649,7 @@ Result<Datum> TDigest(const Datum& value,
 /// \brief Calculate centroids of a numeric array with T-Digest algorithm
 ///
 /// \param[in] value input datum, expecting Array or ChunkedArray
-/// \param[in] options see TDigestOptions for more information
+/// \param[in] options see TDigestMapOptions for more information
 /// \param[in] ctx the function execution context, optional
 /// \return resulting struct of mean and weight arrays
 ///
@@ -654,6 +660,36 @@ Result<Datum> TDigestMap(const Datum& value,
                          const TDigestMapOptions& options = TDigestMapOptions::Defaults(),
                          ExecContext* ctx = NULLPTR);
 
+/// \brief Merge multiple centroid sets into one
+///
+/// \param[in] value input centroid sets, expecting Scalar, Array or ChunkedArray of
+/// centroid structs \param[in] options see TDigestReduceOptions for more information
+/// \param[in] ctx the function execution context, optional
+/// \return resulting struct of mean and weight arrays
+///
+/// \since 22.0.0
+/// \note API not yet finalized
+ARROW_EXPORT
+Result<Datum> TDigestReduce(
+    const Datum& value,
+    const TDigestReduceOptions& options = TDigestReduceOptions::Defaults(),
+    ExecContext* ctx = NULLPTR);
+
+/// \brief Calculate the approximate quantiles using centroids with T-Digest algorithm
+///
+/// \param[in] value input centroid sets, expecting Scalar, Array or ChunkedArray of
+/// centroid structs \param[in] options see TDigestQuantileOptions for more information
+/// \param[in] ctx the function execution context, optional
+/// \return resulting struct of mean and weight arrays
+///
+/// \since 22.0.0
+/// \note API not yet finalized
+ARROW_EXPORT
+Result<Datum> TDigestQuantile(
+    const Datum& value,
+    const TDigestQuantileOptions& options = TDigestQuantileOptions::Defaults(),
+    ExecContext* ctx = NULLPTR);
+
 /// \brief Find the first index of a value in an array.
 ///
 /// \param[in] value The array to search.