Merge branch 'tdigest_map' into combined3

Author: Rafał Hibner

cpp/src/arrow/compute/api_aggregate.cc

@@ -141,6 +141,17 @@ static auto kTDigestOptionsType = GetFunctionOptionsType<TDigestOptions>(
     DataMember("skip_nulls", &TDigestOptions::skip_nulls),
     DataMember("min_count", &TDigestOptions::min_count),
     DataMember("scaler", &TDigestOptions::scaler));
+static auto kTDigestMapOptionsType = GetFunctionOptionsType<TDigestMapOptions>(
+    DataMember("delta", &TDigestMapOptions::delta),
+    DataMember("buffer_size", &TDigestMapOptions::buffer_size),
+    DataMember("skip_nulls", &TDigestMapOptions::skip_nulls),
+    DataMember("scaler", &TDigestMapOptions::scaler));
+static auto kTDigestReduceOptionsType = GetFunctionOptionsType<TDigestReduceOptions>(
+    DataMember("scaler", &TDigestReduceOptions::scaler));
+static auto kTDigestQuantileOptionsType = GetFunctionOptionsType<TDigestQuantileOptions>(
+    DataMember("q", &TDigestQuantileOptions::q),
+    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));
@@ -216,6 +227,34 @@ TDigestOptions::TDigestOptions(std::vector<double> q, uint32_t delta,
       scaler{scaler} {}
 constexpr char TDigestOptions::kTypeName[];
 
+TDigestMapOptions::TDigestMapOptions(uint32_t delta, uint32_t buffer_size,
+                                     bool skip_nulls, Scaler scaler)
+    : FunctionOptions(internal::kTDigestMapOptionsType),
+      delta{delta},
+      buffer_size{buffer_size},
+      skip_nulls{skip_nulls},
+      scaler{scaler} {}
+constexpr char TDigestMapOptions::kTypeName[];
+
+TDigestReduceOptions::TDigestReduceOptions(Scaler scaler)
+    : FunctionOptions(internal::kTDigestReduceOptionsType), scaler{scaler} {}
+constexpr char TDigestReduceOptions::kTypeName[];
+
+TDigestQuantileOptions::TDigestQuantileOptions(double q, uint32_t min_count,
+                                               Scaler scaler)
+    : FunctionOptions(internal::kTDigestQuantileOptionsType),
+      q{q},
+      min_count{min_count},
+      scaler{scaler} {}
+
+TDigestQuantileOptions::TDigestQuantileOptions(std::vector<double> q, uint32_t min_count,
+                                               Scaler scaler)
+    : FunctionOptions(internal::kTDigestQuantileOptionsType),
+      q{std::move(q)},
+      min_count{min_count},
+      scaler{scaler} {}
+constexpr char TDigestReduceOptions::kTypeName[];
+
 PivotWiderOptions::PivotWiderOptions(std::vector<std::string> key_names,
                                      UnexpectedKeyBehavior unexpected_key_behavior)
     : FunctionOptions(internal::kPivotOptionsType),
@@ -237,6 +276,9 @@ void RegisterAggregateOptions(FunctionRegistry* registry) {
   DCHECK_OK(registry->AddFunctionOptionsType(kSkewOptionsType));
   DCHECK_OK(registry->AddFunctionOptionsType(kQuantileOptionsType));
   DCHECK_OK(registry->AddFunctionOptionsType(kTDigestOptionsType));
+  DCHECK_OK(registry->AddFunctionOptionsType(kTDigestMapOptionsType));
+  DCHECK_OK(registry->AddFunctionOptionsType(kTDigestReduceOptionsType));
+  DCHECK_OK(registry->AddFunctionOptionsType(kTDigestQuantileOptionsType));
   DCHECK_OK(registry->AddFunctionOptionsType(kPivotOptionsType));
   DCHECK_OK(registry->AddFunctionOptionsType(kIndexOptionsType));
 }
@@ -321,6 +363,21 @@ Result<Datum> TDigest(const Datum& value, const TDigestOptions& options,
   return CallFunction("tdigest", {value}, &options, ctx);
 }
 
+Result<Datum> TDigestMap(const Datum& value, const TDigestMapOptions& options,
+                         ExecContext* ctx) {
+  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

@@ -179,10 +179,10 @@ class ARROW_EXPORT TDigestOptions : public FunctionOptions {
 
   explicit TDigestOptions(double q = 0.5, uint32_t delta = 100,
                           uint32_t buffer_size = 500, bool skip_nulls = true,
-                          uint32_t min_count = 0, enum Scaler scaler = K0);
+                          uint32_t min_count = 0, enum Scaler scaler = K1);
   explicit TDigestOptions(std::vector<double> q, uint32_t delta = 100,
                           uint32_t buffer_size = 500, bool skip_nulls = true,
-                          uint32_t min_count = 0, enum Scaler scaler = K0);
+                          uint32_t min_count = 0, enum Scaler scaler = K1);
   static constexpr char const kTypeName[] = "TDigestOptions";
   static TDigestOptions Defaults() { return TDigestOptions{}; }
 
@@ -201,6 +201,66 @@ class ARROW_EXPORT TDigestOptions : public FunctionOptions {
   enum Scaler scaler;
 };
 
+/// \brief Control TDigest controid calculation
+///
+/// By default, returns the median value.
+class ARROW_EXPORT TDigestMapOptions : public FunctionOptions {
+ public:
+  using Scaler = TDigestOptions::Scaler;
+
+  explicit TDigestMapOptions(uint32_t delta = 100, uint32_t buffer_size = 500,
+                             bool skip_nulls = true, Scaler scaler = Scaler::K1);
+  static constexpr char const kTypeName[] = "TDigestMapOptions";
+  static TDigestMapOptions Defaults() { return TDigestMapOptions{}; }
+
+  /// compression parameter, default 100
+  uint32_t delta;
+  /// input buffer size, default 500
+  uint32_t buffer_size;
+  /// If true (the default), null values are ignored. Otherwise, if any value is null,
+  /// emit null.
+  bool skip_nulls;
+  /// select scaler implementation
+  Scaler scaler;
+};
+
+/// \brief Control TDigest reduce behavior
+///
+/// By default, returns the median value.
+class ARROW_EXPORT TDigestReduceOptions : public FunctionOptions {
+ public:
+  using Scaler = TDigestOptions::Scaler;
+
+  explicit TDigestReduceOptions(Scaler scaler = Scaler::K1);
+  static constexpr char const kTypeName[] = "TDigestReduceOptions";
+  static TDigestReduceOptions Defaults() { return TDigestReduceOptions{}; }
+
+  /// select scaler implementation
+  Scaler scaler;
+};
+
+/// \brief Control TDigest approximate quantile kernel behavior
+///
+/// By default, returns the median value.
+class ARROW_EXPORT TDigestQuantileOptions : public FunctionOptions {
+ public:
+  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
 ///
 /// These options apply to the "pivot_wider" and "hash_pivot_wider" functions.
@@ -586,6 +646,50 @@ Result<Datum> TDigest(const Datum& value,
                       const TDigestOptions& options = TDigestOptions::Defaults(),
                       ExecContext* ctx = NULLPTR);
 
+/// \brief Calculate centroids of a numeric array with T-Digest algorithm
+///
+/// \param[in] value input datum, expecting Array or ChunkedArray
+/// \param[in] options see TDigestMapOptions 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> 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.