ARROW-3121: [C++] Mean aggregate kernel

cpp/src/arrow/CMakeLists.txt

@@ -147,6 +147,7 @@ if(ARROW_COMPUTE)
  compute/kernels/boolean.cc
  compute/kernels/cast.cc
  compute/kernels/hash.cc
+ compute/kernels/mean.cc
  compute/kernels/sum.cc
  compute/kernels/util-internal.cc)
 endif()

cpp/src/arrow/compute/kernels/aggregate-test.cc

@@ -15,13 +15,18 @@
 // specific language governing permissions and limitations
 // under the License.
 
+#include <algorithm>
+#include <memory>
 #include <string>
 #include <type_traits>
+#include <utility>
 
 #include <gtest/gtest.h>
 
 #include "arrow/array.h"
 #include "arrow/compute/kernel.h"
+#include "arrow/compute/kernels/mean.h"
+#include "arrow/compute/kernels/sum-internal.h"
 #include "arrow/compute/kernels/sum.h"
 #include "arrow/compute/test-util.h"
 #include "arrow/type.h"
@@ -38,94 +43,69 @@ using std::vector;
 namespace arrow {
 namespace compute {
 
-template <typename Type, typename Enable = void>
-struct DatumEqual {
- static void EnsureEqual(const Datum& lhs, const Datum& rhs) {}
-};
-
-template <typename Type>
-struct DatumEqual<Type, typename std::enable_if<IsFloatingPoint<Type>::Value>::type> {
- static constexpr double kArbitraryDoubleErrorBound = 1.0;
- using ScalarType = typename TypeTraits<Type>::ScalarType;
-
- static void EnsureEqual(const Datum& lhs, const Datum& rhs) {
- ASSERT_EQ(lhs.kind(), rhs.kind());
- if (lhs.kind() == Datum::SCALAR) {
- auto left = static_cast<const ScalarType*>(lhs.scalar().get());
- auto right = static_cast<const ScalarType*>(rhs.scalar().get());
- ASSERT_EQ(left->type->id(), right->type->id());
- ASSERT_NEAR(left->value, right->value, kArbitraryDoubleErrorBound);
- }
- }
-};
-
-template <typename Type>
-struct DatumEqual<Type, typename std::enable_if<!IsFloatingPoint<Type>::value>::type> {
- using ScalarType = typename TypeTraits<Type>::ScalarType;
- static void EnsureEqual(const Datum& lhs, const Datum& rhs) {
- ASSERT_EQ(lhs.kind(), rhs.kind());
- if (lhs.kind() == Datum::SCALAR) {
- auto left = static_cast<const ScalarType*>(lhs.scalar().get());
- auto right = static_cast<const ScalarType*>(rhs.scalar().get());
- ASSERT_EQ(left->type->id(), right->type->id());
- ASSERT_EQ(left->value, right->value);
- }
- }
-};
-
 template <typename ArrowType>
-void ValidateSum(FunctionContext* ctx, const Array& input, Datum expected) {
- using OutputType = typename FindAccumulatorType<ArrowType>::Type;
- Datum result;
- ASSERT_OK(Sum(ctx, input, &result));
- DatumEqual<OutputType>::EnsureEqual(result, expected);
-}
+using SumResult =
+ std::pair<typename FindAccumulatorType<ArrowType>::Type::c_type, size_t>;
 
 template <typename ArrowType>
-void ValidateSum(FunctionContext* ctx, const char* json, Datum expected) {
- auto array = ArrayFromJSON(TypeTraits<ArrowType>::type_singleton(), json);
- ValidateSum<ArrowType>(ctx, *array, expected);
-}
-
-template <typename ArrowType>
-static Datum DummySum(const Array& array) {
+static SumResult<ArrowType> NaiveSumPartial(const Array& array) {
  using ArrayType = typename TypeTraits<ArrowType>::ArrayType;
- using SumType = typename FindAccumulatorType<ArrowType>::Type;
- using SumScalarType = typename TypeTraits<SumType>::ScalarType;
+ using ResultType = SumResult<ArrowType>;
 
- typename SumType::c_type sum = 0;
- int64_t count = 0;
+ ResultType result;
 
  auto data = array.data();
  internal::BitmapReader reader(array.null_bitmap_data(), array.offset(), array.length());
  const auto& array_numeric = reinterpret_cast<const ArrayType&>(array);
  const auto values = array_numeric.raw_values();
  for (int64_t i = 0; i < array.length(); i++) {
  if (reader.IsSet()) {
-sum += values[i];
-count++;
+result.first += values[i];
+result.second++;
  }
 
  reader.Next();
  }
 
- if (count > 0) {
- return Datum(std::make_shared<SumScalarType>(sum));
- } else {
- return Datum(std::make_shared<SumScalarType>(0, false));
- }
+ return result;
+}
+
+template <typename ArrowType>
+static Datum NaiveSum(const Array& array) {
+ using SumType = typename FindAccumulatorType<ArrowType>::Type;
+ using SumScalarType = typename TypeTraits<SumType>::ScalarType;
+
+ auto result = NaiveSumPartial<ArrowType>(array);
+ bool is_valid = result.second > 0;
+
+ return Datum(std::make_shared<SumScalarType>(result.first, is_valid));
+}
+
+template <typename ArrowType>
+void ValidateSum(FunctionContext* ctx, const Array& input, Datum expected) {
+ using OutputType = typename FindAccumulatorType<ArrowType>::Type;
+
+ Datum result;
+ ASSERT_OK(Sum(ctx, input, &result));
+ DatumEqual<OutputType>::EnsureEqual(result, expected);
+}
+
+template <typename ArrowType>
+void ValidateSum(FunctionContext* ctx, const char* json, Datum expected) {
+ auto array = ArrayFromJSON(TypeTraits<ArrowType>::type_singleton(), json);
+ ValidateSum<ArrowType>(ctx, *array, expected);
 }
 
 template <typename ArrowType>
 void ValidateSum(FunctionContext* ctx, const Array& array) {
- ValidateSum<ArrowType>(ctx, array, DummySum<ArrowType>(array));
+ ValidateSum<ArrowType>(ctx, array, NaiveSum<ArrowType>(array));
 }
 
 template <typename ArrowType>
-class TestSumKernelNumeric : public ComputeFixture, public TestBase {};
+class TestNumericSumKernel : public ComputeFixture, public TestBase {};
 
-TYPED_TEST_CASE(TestSumKernelNumeric, NumericArrowTypes);
-TYPED_TEST(TestSumKernelNumeric, SimpleSum) {
+TYPED_TEST_CASE(TestNumericSumKernel, NumericArrowTypes);
+TYPED_TEST(TestNumericSumKernel, SimpleSum) {
  using SumType = typename FindAccumulatorType<TypeParam>::Type;
  using ScalarType = typename TypeTraits<SumType>::ScalarType;
  using T = typename TypeParam::c_type;
@@ -145,10 +125,10 @@ TYPED_TEST(TestSumKernelNumeric, SimpleSum) {
 }
 
 template <typename ArrowType>
-class TestRandomSumKernelNumeric : public ComputeFixture, public TestBase {};
+class TestRandomNumericSumKernel : public ComputeFixture, public TestBase {};
 
-TYPED_TEST_CASE(TestRandomSumKernelNumeric, NumericArrowTypes);
-TYPED_TEST(TestRandomSumKernelNumeric, RandomArraySum) {
+TYPED_TEST_CASE(TestRandomNumericSumKernel, NumericArrowTypes);
+TYPED_TEST(TestRandomNumericSumKernel, RandomArraySum) {
  auto rand = random::RandomArrayGenerator(0x5487655);
  for (size_t i = 3; i < 14; i++) {
  for (auto null_probability : {0.0, 0.01, 0.1, 0.25, 0.5, 1.0}) {
@@ -161,7 +141,7 @@ TYPED_TEST(TestRandomSumKernelNumeric, RandomArraySum) {
  }
 }
 
-TYPED_TEST(TestRandomSumKernelNumeric, RandomSliceArraySum) {
+TYPED_TEST(TestRandomNumericSumKernel, RandomSliceArraySum) {
  auto arithmetic = ArrayFromJSON(TypeTraits<TypeParam>::type_singleton(),
  "[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16]");
  ValidateSum<TypeParam>(&this->ctx_, *arithmetic);
@@ -175,12 +155,87 @@ TYPED_TEST(TestRandomSumKernelNumeric, RandomSliceArraySum) {
  const int64_t length = 1U << 6;
  auto array = rand.Numeric<TypeParam>(length, 0, 10, 0.5);
  for (size_t i = 1; i < 16; i++) {
- for (size_t j = 1; i < 16; i++) {
+ for (size_t j = 1; j < 16; j++) {
  auto slice = array->Slice(i, length - j);
  ValidateSum<TypeParam>(&this->ctx_, *slice);
  }
  }
 }
 
+template <typename ArrowType>
+static Datum NaiveMean(const Array& array) {
+ using MeanScalarType = typename TypeTraits<DoubleType>::ScalarType;
+
+ const auto result = NaiveSumPartial<ArrowType>(array);
+ const double mean = static_cast<double>(result.first) /
+ static_cast<double>(result.second ? result.second : 1UL);
+ const bool is_valid = result.second > 0;
+
+ return Datum(std::make_shared<MeanScalarType>(mean, is_valid));
+}
+
+template <typename ArrowType>
+void ValidateMean(FunctionContext* ctx, const Array& input, Datum expected) {
+ using OutputType = typename FindAccumulatorType<DoubleType>::Type;
+
+ Datum result;
+ ASSERT_OK(Mean(ctx, input, &result));
+ DatumEqual<OutputType>::EnsureEqual(result, expected);
+}
+
+template <typename ArrowType>
+void ValidateMean(FunctionContext* ctx, const char* json, Datum expected) {
+ auto array = ArrayFromJSON(TypeTraits<ArrowType>::type_singleton(), json);
+ ValidateMean<ArrowType>(ctx, *array, expected);
+}
+
+template <typename ArrowType>
+void ValidateMean(FunctionContext* ctx, const Array& array) {
+ ValidateMean<ArrowType>(ctx, array, NaiveMean<ArrowType>(array));
+}
+
+template <typename ArrowType>
+class TestMeanKernelNumeric : public ComputeFixture, public TestBase {};
+
+TYPED_TEST_CASE(TestMeanKernelNumeric, NumericArrowTypes);
+TYPED_TEST(TestMeanKernelNumeric, SimpleMean) {
+ using ScalarType = typename TypeTraits<DoubleType>::ScalarType;
+
+ ValidateMean<TypeParam>(&this->ctx_, "[]",
+ Datum(std::make_shared<ScalarType>(0.0, false)));
+
+ ValidateMean<TypeParam>(&this->ctx_, "[null]",
+ Datum(std::make_shared<ScalarType>(0.0, false)));
+
+ ValidateMean<TypeParam>(&this->ctx_, "[1, null, 1]",
+ Datum(std::make_shared<ScalarType>(1.0)));
+
+ ValidateMean<TypeParam>(&this->ctx_, "[1, 2, 3, 4, 5, 6, 7, 8]",
+ Datum(std::make_shared<ScalarType>(4.5)));
+
+ ValidateMean<TypeParam>(&this->ctx_, "[0, 0, 0, 0, 0, 0, 0, 0]",
+ Datum(std::make_shared<ScalarType>(0.0)));
+
+ ValidateMean<TypeParam>(&this->ctx_, "[1, 1, 1, 1, 1, 1, 1, 1]",
+ Datum(std::make_shared<ScalarType>(1.0)));
+}
+
+template <typename ArrowType>
+class TestRandomNumericMeanKernel : public ComputeFixture, public TestBase {};
+
+TYPED_TEST_CASE(TestRandomNumericMeanKernel, NumericArrowTypes);
+TYPED_TEST(TestRandomNumericMeanKernel, RandomArrayMean) {
+ auto rand = random::RandomArrayGenerator(0x8afc055);
+ for (size_t i = 3; i < 14; i++) {
+ for (auto null_probability : {0.0, 0.01, 0.1, 0.25, 0.5, 1.0}) {
+ for (auto length_adjust : {-2, -1, 0, 1, 2}) {
+ int64_t length = (1UL << i) + length_adjust;
+ auto array = rand.Numeric<TypeParam>(length, 0, 100, null_probability);
+ ValidateMean<TypeParam>(&this->ctx_, *array);
+ }
+ }
+ }
+}
+
 } // namespace compute
 } // namespace arrow

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

@@ -0,0 +1,115 @@
+// 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.
+
+#include "arrow/compute/kernels/mean.h"
+
+#include <algorithm>
+
+#include "arrow/compute/kernels/sum-internal.h"
+
+namespace arrow {
+namespace compute {
+
+template <typename ArrowType,
+ typename SumType = typename FindAccumulatorType<ArrowType>::Type>
+struct MeanState {
+ using ThisType = MeanState<ArrowType, SumType>;
+
+ ThisType operator+(const ThisType& rhs) const {
+ return ThisType(this->count + rhs.count, this->sum + rhs.sum);
+ }
+
+ ThisType& operator+=(const ThisType& rhs) {
+ this->count += rhs.count;
+ this->sum += rhs.sum;
+
+ return *this;
+ }
+
+ std::shared_ptr<Scalar> Finalize() const {
+ using ScalarType = typename TypeTraits<DoubleType>::ScalarType;
+
+ const bool is_valid = count > 0;
+ const double divisor = static_cast<double>(is_valid ? count : 1UL);
+ const double mean = static_cast<double>(sum) / divisor;
+
+ return std::make_shared<ScalarType>(mean, is_valid);
+ }
+
+ static std::shared_ptr<DataType> out_type() {
+ return TypeTraits<DoubleType>::type_singleton();
+ }
+
+ size_t count = 0;
+ typename SumType::c_type sum = 0;
+};
+
+#define MEAN_AGG_FN_CASE(T) \
+ case T::type_id: \
+ return std::static_pointer_cast<AggregateFunction>( \
+ std::make_shared<SumAggregateFunction<T, MeanState<T>>>());
+
+std::shared_ptr<AggregateFunction> MakeMeanAggregateFunction(const DataType& type,
+ FunctionContext* ctx) {
+ switch (type.id()) {
+ MEAN_AGG_FN_CASE(UInt8Type);
+ MEAN_AGG_FN_CASE(Int8Type);
+ MEAN_AGG_FN_CASE(UInt16Type);
+ MEAN_AGG_FN_CASE(Int16Type);
+ MEAN_AGG_FN_CASE(UInt32Type);
+ MEAN_AGG_FN_CASE(Int32Type);
+ MEAN_AGG_FN_CASE(UInt64Type);
+ MEAN_AGG_FN_CASE(Int64Type);
+ MEAN_AGG_FN_CASE(FloatType);
+ MEAN_AGG_FN_CASE(DoubleType);
+ default:
+ return nullptr;
+ }
+
+#undef MEAN_AGG_FN_CASE
+}
+
+static Status GetMeanKernel(FunctionContext* ctx, const DataType& type,
+ std::shared_ptr<AggregateUnaryKernel>& kernel) {
+ std::shared_ptr<AggregateFunction> aggregate = MakeMeanAggregateFunction(type, ctx);
+ if (!aggregate) return Status::Invalid("No mean for type ", type);
+
+ kernel = std::make_shared<AggregateUnaryKernel>(aggregate);
+
+ return Status::OK();
+}
+
+Status Mean(FunctionContext* ctx, const Datum& value, Datum* out) {
+ std::shared_ptr<AggregateUnaryKernel> kernel;
+
+ auto data_type = value.type();
+ if (data_type == nullptr)
+ return Status::Invalid("Datum must be array-like");
+ else if (!is_integer(data_type->id()) && !is_floating(data_type->id()))
+ return Status::Invalid("Datum must contain a NumericType");
+
+ RETURN_NOT_OK(GetMeanKernel(ctx, *data_type, kernel));
+
+ return kernel->Call(ctx, value, out);
+}
+
+Status Mean(FunctionContext* ctx, const Array& array, Datum* out) {
+ return Mean(ctx, array.data(), out);
+}
+
+} // namespace compute
+} // namespace arrow