Instantiate tdigest per exec

Author: Rafał Hibner

cpp/src/arrow/compute/kernels/aggregate_tdigest.cc

@@ -35,7 +35,7 @@ using arrow::internal::TDigestScalerK1;
 using arrow::internal::VisitSetBitRunsVoid;
 
 struct TDigestBaseImpl : public ScalarAggregator {
-  explicit TDigestBaseImpl(std::unique_ptr<TDigest::Scaler> scaler, uint32_t buffer_size)
+  explicit TDigestBaseImpl(std::shared_ptr<TDigest::Scaler> scaler, uint32_t buffer_size)
       : tdigest{std::move(scaler), buffer_size}, count{0}, all_valid{true} {
     out_type = struct_({
         field("mean", list(field("item", float64(), false)), false),
@@ -57,7 +57,7 @@ struct TDigestBaseImpl : public ScalarAggregator {
     return Status::OK();
   }
 
-  static Result<std::unique_ptr<TDigest::Scaler>> MakeScaler(
+  static Result<std::shared_ptr<TDigest::Scaler>> MakeScaler(
       TDigestOptions::Scaler scaler, uint32_t delta) {
     switch (scaler) {
       case TDigestOptions::K0:
@@ -288,7 +288,7 @@ struct TDigestImpl
   // using TDigestBaseImpl::tdigest;
 
   explicit TDigestImpl(const TDigestOptions& options, const DataType& in_type,
-                       std::unique_ptr<TDigest::Scaler> scaler)
+                       std::shared_ptr<TDigest::Scaler> scaler)
       : TDigestInputConsumerImpl<ArrowType, TDigestQuantileFinalizer>(
             // TDigestInputConsumerImpl
             options.skip_nulls, in_type,
@@ -302,7 +302,7 @@ template <typename ArrowType>
 struct TDigestMapImpl
     : public TDigestInputConsumerImpl<ArrowType, TDigestCentroidFinalizer> {
   explicit TDigestMapImpl(const TDigestMapOptions& options, const DataType& in_type,
-                          std::unique_ptr<TDigest::Scaler> scaler)
+                          std::shared_ptr<TDigest::Scaler> scaler)
       : TDigestInputConsumerImpl<ArrowType, TDigestCentroidFinalizer>(
 
             // TDigestInputConsumerImpl
@@ -314,7 +314,7 @@ struct TDigestMapImpl
 
 struct TDigestReduceImpl : public TDigestCentroidConsumerImpl<TDigestCentroidFinalizer> {
   explicit TDigestReduceImpl(const TDigestReduceOptions& options,
-                             std::unique_ptr<TDigest::Scaler> scaler)
+                             std::shared_ptr<TDigest::Scaler> scaler)
       : TDigestCentroidConsumerImpl<TDigestCentroidFinalizer>(
             // TDigestCentroidConsumerImpl
             // TDigestCentroidFinalizer
@@ -325,7 +325,7 @@ struct TDigestReduceImpl : public TDigestCentroidConsumerImpl<TDigestCentroidFin
 struct TDigestQuantileImpl
     : public TDigestCentroidConsumerImpl<TDigestQuantileFinalizer> {
   explicit TDigestQuantileImpl(const TDigestQuantileOptions& options,
-                               std::unique_ptr<TDigest::Scaler> scaler)
+                               std::shared_ptr<TDigest::Scaler> scaler)
       : TDigestCentroidConsumerImpl<TDigestQuantileFinalizer>(
 
             // TDigestCentroidConsumerImpl
@@ -335,10 +335,10 @@ struct TDigestQuantileImpl
             std::move(scaler), options.delta) {}
 };
 
-struct TDigestQuantileScalarImpl : public TDigestQuantileImpl {
+struct TDigestQuantileScalarImpl : public KernelState {
   explicit TDigestQuantileScalarImpl(const TDigestQuantileOptions& options,
-                                     std::unique_ptr<TDigest::Scaler> scaler)
-      : TDigestQuantileImpl(options, std::move(scaler)) {}
+                                     std::shared_ptr<TDigest::Scaler> scaler)
+      : options(options), scaler(std::move(scaler)) {}
 
   static Result<std::unique_ptr<KernelState>> Init(KernelContext* ctx,
                                                    const KernelInitArgs& args) {
@@ -354,33 +354,35 @@ struct TDigestQuantileScalarImpl : public TDigestQuantileImpl {
     return state->OutputType();
   }
 
-  size_t OutputSize() const { return this->q.size(); }
+  size_t OutputSize() const { return options.q.size(); }
 
   TypeHolder OutputType() const {
     return fixed_size_list(field("item", float64()), OutputSize());
   }
 
   static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
     auto state = checked_cast<TDigestQuantileScalarImpl*>(ctx->state());
+    TDigestQuantileImpl tdigest(state->options, state->scaler);
     auto value_builder = std::make_shared<DoubleBuilder>(ctx->memory_pool());
     const auto output_size = state->OutputSize();
     FixedSizeListBuilder fsl_builder(
         ctx->memory_pool(), checked_pointer_cast<arrow::ArrayBuilder>(value_builder),
         output_size);
 
     std::shared_ptr<Array> array = MakeArray(batch[0].array.ToArrayData());
+
     for (int i = 0; i < array->length(); ++i) {
       if (array->IsValid(i)) {
         ARROW_RETURN_NOT_OK(fsl_builder.Append());
         ARROW_ASSIGN_OR_RAISE(auto scalar, array->GetScalar(i));
-        state->Reset();
-        ARROW_RETURN_NOT_OK(state->Consume(scalar.get()));
+        tdigest.Reset();
+        ARROW_RETURN_NOT_OK(tdigest.Consume(scalar.get()));
 
-        if (state->isNull()) {
+        if (tdigest.isNull()) {
           ARROW_RETURN_NOT_OK(value_builder->AppendNulls(output_size));
         } else {
           for (size_t i = 0; i < output_size; ++i) {
-            ARROW_RETURN_NOT_OK(value_builder->Append(state->Quantile(i)));
+            ARROW_RETURN_NOT_OK(value_builder->Append(tdigest.Quantile(i)));
           }
         }
       } else {
@@ -392,6 +394,10 @@ struct TDigestQuantileScalarImpl : public TDigestQuantileImpl {
     out->value = std::move(out_array->data());
     return Status::OK();
   }
+
+ private:
+  TDigestQuantileOptions options;
+  std::shared_ptr<TDigest::Scaler> scaler;
 };
 
 template <template <typename> typename TDigestImpl_T, typename TDigestOptions_T>

cpp/src/arrow/util/tdigest.cc

@@ -58,7 +58,7 @@ struct Centroid {
 // implements t-digest merging algorithm
 class TDigestMerger {
  public:
-  explicit TDigestMerger(std::unique_ptr<TDigest::Scaler> scaler)
+  explicit TDigestMerger(std::shared_ptr<TDigest::Scaler> scaler)
       : scaler_{std::move(scaler)} {
     Reset(0, nullptr);
   }
@@ -111,7 +111,7 @@ class TDigestMerger {
   uint32_t delta() const { return scaler_->delta_; }
 
  private:
-  std::unique_ptr<TDigest::Scaler> scaler_;
+  std::shared_ptr<TDigest::Scaler> scaler_;
   double total_weight_;   // total weight of this tdigest
   double weight_so_far_;  // accumulated weight till current bin
   double weight_limit_;   // max accumulated weight to move to next bin
@@ -122,7 +122,7 @@ class TDigestMerger {
 
 class TDigest::TDigestImpl {
  public:
-  explicit TDigestImpl(std::unique_ptr<Scaler> scaler) : merger_(std::move(scaler)) {
+  explicit TDigestImpl(std::shared_ptr<Scaler> scaler) : merger_(std::move(scaler)) {
     tdigests_[0].reserve(merger_.delta());
     tdigests_[1].reserve(merger_.delta());
     Reset();
@@ -305,7 +305,7 @@ class TDigest::TDigestImpl {
     if (diff > 0) {
       if (ci_right == td.size() - 1) {
         // index larger than center of last bin
-        DCHECK_EQ(weight_sum, total_weight_);
+        DCHECK_LE(std::abs(weight_sum - total_weight_), (weight_sum * 1e-9));
         const Centroid* c = &td[ci_right];
         DCHECK_GE(c->weight, 2);
         return Lerp(c->mean, max_, diff / (c->weight / 2));
@@ -363,7 +363,7 @@ class TDigest::TDigestImpl {
 TDigest::TDigest(uint32_t delta, uint32_t buffer_size)
     : TDigest(std::make_unique<TDigestScalerK1>(delta), buffer_size) {}
 
-TDigest::TDigest(std::unique_ptr<Scaler> scaler, uint32_t buffer_size)
+TDigest::TDigest(std::shared_ptr<Scaler> scaler, uint32_t buffer_size)
     : impl_(new TDigestImpl(std::move(scaler))) {
   input_.reserve(buffer_size);
   Reset();

cpp/src/arrow/util/tdigest_internal.h

@@ -49,7 +49,7 @@ class ARROW_EXPORT TDigest {
   };
 
   explicit TDigest(uint32_t delta = 100, uint32_t buffer_size = 500);
-  explicit TDigest(std::unique_ptr<Scaler> scaler, uint32_t buffer_size = 500);
+  explicit TDigest(std::shared_ptr<Scaler> scaler, uint32_t buffer_size = 500);
   ~TDigest();
   TDigest(TDigest&&);
   TDigest& operator=(TDigest&&);