fix c++ tests

khwilson · Oct 13, 2024 · 01e53b3 · 01e53b3
1 parent 12a2749
commit 01e53b3
Show file tree

Hide file tree

Showing 5 changed files with 157 additions and 74 deletions.
diff --git a/cpp/src/arrow/compute/kernels/aggregate_basic.cc b/cpp/src/arrow/compute/kernels/aggregate_basic.cc
@@ -365,8 +365,10 @@ struct ProductImpl : public ScalarAggregator {
   }
 
   Status Finalize(KernelContext*, Datum* out) override {
-    std::shared_ptr<DataType> out_type_;
-    ARROW_ASSIGN_OR_RAISE(out_type_, WidenDecimalToMaxPrecision(this->out_type));
+    std::shared_ptr<DataType> out_type_ = this->out_type;
+    if (is_decimal(this->out_type->id())) {
+      ARROW_ASSIGN_OR_RAISE(out_type_, WidenDecimalToMaxPrecision(this->out_type));
+    }
 
     if ((!options.skip_nulls && this->nulls_observed) ||
         (this->count < options.min_count)) {
@@ -1051,6 +1053,10 @@ void RegisterScalarAggregateBasic(FunctionRegistry* registry) {
   func = std::make_shared<ScalarAggregateFunction>("sum", Arity::Unary(), sum_doc,
                                                    &default_scalar_aggregate_options);
   AddArrayScalarAggKernels(SumInit, {boolean()}, uint64(), func.get());
+  AddAggKernel(KernelSignature::Make({Type::DECIMAL32}, MaxPrecisionDecimalType),
+               SumInit, func.get(), SimdLevel::NONE);
+  AddAggKernel(KernelSignature::Make({Type::DECIMAL64}, MaxPrecisionDecimalType),
+               SumInit, func.get(), SimdLevel::NONE);
   AddAggKernel(KernelSignature::Make({Type::DECIMAL128}, MaxPrecisionDecimalType),
                SumInit, func.get(), SimdLevel::NONE);
   AddAggKernel(KernelSignature::Make({Type::DECIMAL256}, MaxPrecisionDecimalType),
@@ -1079,6 +1085,10 @@ void RegisterScalarAggregateBasic(FunctionRegistry* registry) {
                                                    &default_scalar_aggregate_options);
   AddArrayScalarAggKernels(MeanInit, {boolean()}, float64(), func.get());
   AddArrayScalarAggKernels(MeanInit, NumericTypes(), float64(), func.get());
+  AddAggKernel(KernelSignature::Make({Type::DECIMAL32}, MaxPrecisionDecimalType),
+               MeanInit, func.get(), SimdLevel::NONE);
+  AddAggKernel(KernelSignature::Make({Type::DECIMAL64}, MaxPrecisionDecimalType),
+               MeanInit, func.get(), SimdLevel::NONE);
   AddAggKernel(KernelSignature::Make({Type::DECIMAL128}, MaxPrecisionDecimalType),
                MeanInit, func.get(), SimdLevel::NONE);
   AddAggKernel(KernelSignature::Make({Type::DECIMAL256}, MaxPrecisionDecimalType),
@@ -1128,6 +1138,8 @@ void RegisterScalarAggregateBasic(FunctionRegistry* registry) {
   AddMinMaxKernels(MinMaxInitDefault, BaseBinaryTypes(), func.get());
   AddMinMaxKernel(MinMaxInitDefault, Type::FIXED_SIZE_BINARY, func.get());
   AddMinMaxKernel(MinMaxInitDefault, Type::INTERVAL_MONTHS, func.get());
+  AddMinMaxKernel(MinMaxInitDefault, Type::DECIMAL32, func.get());
+  AddMinMaxKernel(MinMaxInitDefault, Type::DECIMAL64, func.get());
   AddMinMaxKernel(MinMaxInitDefault, Type::DECIMAL128, func.get());
   AddMinMaxKernel(MinMaxInitDefault, Type::DECIMAL256, func.get());
   // Add the SIMD variants for min max
@@ -1163,6 +1175,10 @@ void RegisterScalarAggregateBasic(FunctionRegistry* registry) {
   AddArrayScalarAggKernels(ProductInit::Init, UnsignedIntTypes(), uint64(), func.get());
   AddArrayScalarAggKernels(ProductInit::Init, FloatingPointTypes(), float64(),
                            func.get());
+  AddAggKernel(KernelSignature::Make({Type::DECIMAL32}, MaxPrecisionDecimalType),
+               ProductInit::Init, func.get(), SimdLevel::NONE);
+  AddAggKernel(KernelSignature::Make({Type::DECIMAL64}, MaxPrecisionDecimalType),
+               ProductInit::Init, func.get(), SimdLevel::NONE);
   AddAggKernel(KernelSignature::Make({Type::DECIMAL128}, MaxPrecisionDecimalType),
                ProductInit::Init, func.get(), SimdLevel::NONE);
   AddAggKernel(KernelSignature::Make({Type::DECIMAL256}, MaxPrecisionDecimalType),
@@ -1188,7 +1204,7 @@ void RegisterScalarAggregateBasic(FunctionRegistry* registry) {
   AddBasicAggKernels(IndexInit::Init, PrimitiveTypes(), int64(), func.get());
   AddBasicAggKernels(IndexInit::Init, TemporalTypes(), int64(), func.get());
   AddBasicAggKernels(IndexInit::Init,
-                     {fixed_size_binary(1), decimal128(1, 0), decimal256(1, 0), null()},
+                     {fixed_size_binary(1), decimal32(1, 0), decimal64(1, 0), decimal128(1, 0), decimal256(1, 0), null()},
                      int64(), func.get());
   DCHECK_OK(registry->AddFunction(std::move(func)));
 }

diff --git a/cpp/src/arrow/compute/kernels/aggregate_basic.inc.cc b/cpp/src/arrow/compute/kernels/aggregate_basic.inc.cc
@@ -93,8 +93,10 @@ struct SumImpl : public ScalarAggregator {
   }
 
   Status Finalize(KernelContext*, Datum* out) override {
-    std::shared_ptr<DataType> out_type_;
-    ARROW_ASSIGN_OR_RAISE(out_type_, WidenDecimalToMaxPrecision(this->out_type));
+    std::shared_ptr<DataType> out_type_ = this->out_type;
+    if (is_decimal(this->out_type->id())) {
+      ARROW_ASSIGN_OR_RAISE(out_type_, WidenDecimalToMaxPrecision(this->out_type));
+    }
 
     if ((!options.skip_nulls && this->nulls_observed) ||
         (this->count < options.min_count)) {