ARROW-13295: [C++] add hash_mean, hash_variance, hash_stddev kernels

cpp/src/arrow/compute/kernels/aggregate_internal.h

@@ -55,6 +55,10 @@ struct ScalarAggregator : public KernelState {
   virtual Status Finalize(KernelContext* ctx, Datum* out) = 0;
 };
 
+// Helper to differentiate between var/std calculation so we can fold
+// kernel implementations together
+enum class VarOrStd : bool { Var, Std };
+
 void AddAggKernel(std::shared_ptr<KernelSignature> sig, KernelInit init,
                   ScalarAggregateFunction* func,
                   SimdLevel::type simd_level = SimdLevel::NONE);

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

@@ -19,6 +19,7 @@
 
 #include "arrow/compute/api_aggregate.h"
 #include "arrow/compute/kernels/aggregate_internal.h"
+#include "arrow/compute/kernels/aggregate_var_std_internal.h"
 #include "arrow/compute/kernels/common.h"
 #include "arrow/util/bit_run_reader.h"
 #include "arrow/util/int128_internal.h"
@@ -85,32 +86,22 @@ struct VarStdState {
       valid_count -= count;
 
       if (count > 0) {
-        int64_t sum = 0;
-        int128_t square_sum = 0;
+        IntegerVarStd<ArrowType> var_std;
         const ArrayData& data = *slice->data();
         const CType* values = data.GetValues<CType>(1);
         VisitSetBitRunsVoid(data.buffers[0], data.offset, data.length,
                             [&](int64_t pos, int64_t len) {
                               for (int64_t i = 0; i < len; ++i) {
                                 const auto value = values[pos + i];
-                                sum += value;
-                                square_sum += static_cast<uint64_t>(value) * value;
+                                var_std.ConsumeOne(value);
                               }
                             });
 
-        const double mean = static_cast<double>(sum) / count;
-        // calculate m2 = square_sum - sum * sum / count
-        // decompose `sum * sum / count` into integers and fractions
-        const int128_t sum_square = static_cast<int128_t>(sum) * sum;
-        const int128_t integers = sum_square / count;
-        const double fractions = static_cast<double>(sum_square % count) / count;
-        const double m2 = static_cast<double>(square_sum - integers) - fractions;
-
         // merge variance
         ThisType state;
-        state.count = count;
-        state.mean = mean;
-        state.m2 = m2;
+        state.count = var_std.count;
+        state.mean = var_std.mean();
+        state.m2 = var_std.m2();
         this->MergeFrom(state);
       }
     }
@@ -128,21 +119,15 @@ struct VarStdState {
       this->m2 = state.m2;
       return;
     }
-    double mean = (this->mean * this->count + state.mean * state.count) /
-                  (this->count + state.count);
-    this->m2 += state.m2 + this->count * (this->mean - mean) * (this->mean - mean) +
-                state.count * (state.mean - mean) * (state.mean - mean);
-    this->count += state.count;
-    this->mean = mean;
+    MergeVarStd(this->count, this->mean, state.count, state.mean, state.m2, &this->count,
+                &this->mean, &this->m2);
   }
 
   int64_t count = 0;
   double mean = 0;
   double m2 = 0;  // m2 = count*s2 = sum((X-mean)^2)
 };
 
-enum class VarOrStd : bool { Var, Std };
-
 template <typename ArrowType>
 struct VarStdImpl : public ScalarAggregator {
   using ThisType = VarStdImpl<ArrowType>;

cpp/src/arrow/compute/kernels/aggregate_var_std_internal.h

@@ -0,0 +1,68 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include "arrow/util/int128_internal.h"
+
+namespace arrow {
+namespace compute {
+namespace internal {
+
+using arrow::internal::int128_t;
+
+// Accumulate sum/squared sum (using naive summation)
+// Shared implementation between scalar/hash aggregate variance/stddev kernels
+template <typename ArrowType>
+struct IntegerVarStd {
+  using c_type = typename ArrowType::c_type;
+
+  int64_t count = 0;
+  int64_t sum = 0;
+  int128_t square_sum = 0;
+
+  void ConsumeOne(const c_type value) {
+    sum += value;
+    square_sum += static_cast<uint64_t>(value) * value;
+    count++;
+  }
+
+  double mean() const { return static_cast<double>(sum) / count; }
+
+  double m2() const {
+    // calculate m2 = square_sum - sum * sum / count
+    // decompose `sum * sum / count` into integers and fractions
+    const int128_t sum_square = static_cast<int128_t>(sum) * sum;
+    const int128_t integers = sum_square / count;
+    const double fractions = static_cast<double>(sum_square % count) / count;
+    return static_cast<double>(square_sum - integers) - fractions;
+  }
+};
+
+static inline void MergeVarStd(int64_t count1, double mean1, int64_t count2, double mean2,
+                               double m22, int64_t* out_count, double* out_mean,
+                               double* out_m2) {
+  double mean = (mean1 * count1 + mean2 * count2) / (count1 + count2);
+  *out_m2 += m22 + count1 * (mean1 - mean) * (mean1 - mean) +
+             count2 * (mean2 - mean) * (mean2 - mean);
+  *out_count += count2;
+  *out_mean = mean;
+}
+
+}  // namespace internal
+}  // namespace compute
+}  // namespace arrow