diff --git a/cpp/src/arrow/acero/hash_aggregate_test.cc b/cpp/src/arrow/acero/hash_aggregate_test.cc index dce0e44eb135..de414219eb42 100644 --- a/cpp/src/arrow/acero/hash_aggregate_test.cc +++ b/cpp/src/arrow/acero/hash_aggregate_test.cc @@ -937,8 +937,8 @@ TEST_P(GroupBy, SumMeanProductDecimal) { AssertDatumsEqual(ArrayFromJSON(struct_({ field("key_0", int64()), - field("hash_sum", decimal128(3, 2)), - field("hash_sum", decimal256(3, 2)), + field("hash_sum", decimal128(38, 2)), + field("hash_sum", decimal256(76, 2)), field("hash_mean", decimal128(3, 2)), field("hash_mean", decimal256(3, 2)), field("hash_product", decimal128(3, 2)), diff --git a/cpp/src/arrow/compute/kernels/aggregate_basic.cc b/cpp/src/arrow/compute/kernels/aggregate_basic.cc index ee2c615bbfb7..b19a9f58e5b8 100644 --- a/cpp/src/arrow/compute/kernels/aggregate_basic.cc +++ b/cpp/src/arrow/compute/kernels/aggregate_basic.cc @@ -368,9 +368,9 @@ struct ProductImpl : public ScalarAggregator { Status Finalize(KernelContext*, Datum* out) override { if ((!options.skip_nulls && this->nulls_observed) || (this->count < options.min_count)) { - out->value = std::make_shared(out_type); + out->value = std::make_shared(this->out_type); } else { - out->value = std::make_shared(this->product, out_type); + out->value = std::make_shared(this->product, this->out_type); } return Status::OK(); } @@ -1049,10 +1049,10 @@ void RegisterScalarAggregateBasic(FunctionRegistry* registry) { func = std::make_shared("sum", Arity::Unary(), sum_doc, &default_scalar_aggregate_options); AddArrayScalarAggKernels(SumInit, {boolean()}, uint64(), func.get()); - AddAggKernel(KernelSignature::Make({Type::DECIMAL128}, FirstType), SumInit, func.get(), - SimdLevel::NONE); - AddAggKernel(KernelSignature::Make({Type::DECIMAL256}, FirstType), SumInit, func.get(), - SimdLevel::NONE); + AddAggKernel(KernelSignature::Make({Type::DECIMAL128}, MaxPrecisionDecimalType), + SumInit, func.get(), SimdLevel::NONE); + AddAggKernel(KernelSignature::Make({Type::DECIMAL256}, MaxPrecisionDecimalType), + SumInit, func.get(), SimdLevel::NONE); AddArrayScalarAggKernels(SumInit, SignedIntTypes(), int64(), func.get()); AddArrayScalarAggKernels(SumInit, UnsignedIntTypes(), uint64(), func.get()); AddArrayScalarAggKernels(SumInit, FloatingPointTypes(), float64(), func.get()); diff --git a/cpp/src/arrow/compute/kernels/aggregate_basic.inc.cc b/cpp/src/arrow/compute/kernels/aggregate_basic.inc.cc index 49010d182cd6..7f2bce4063d7 100644 --- a/cpp/src/arrow/compute/kernels/aggregate_basic.inc.cc +++ b/cpp/src/arrow/compute/kernels/aggregate_basic.inc.cc @@ -95,9 +95,9 @@ struct SumImpl : public ScalarAggregator { Status Finalize(KernelContext*, Datum* out) override { if ((!options.skip_nulls && this->nulls_observed) || (this->count < options.min_count)) { - out->value = std::make_shared(out_type); + out->value = std::make_shared(this->out_type); } else { - out->value = std::make_shared(this->sum, out_type); + out->value = std::make_shared(this->sum, this->out_type); } return Status::OK(); } @@ -168,6 +168,14 @@ struct SumLikeInit { const ScalarAggregateOptions& options) : ctx(ctx), type(type), options(options) {} + // If this returns true, then the aggregator will promote a decimal to the maximum + // precision for that type. For instance, a decimal128(3, 2) will be promoted to a + // decimal128(38, 2) + // + // TODO: Ideally this should be configurable via the function options with an enum + // PrecisionPolicy { PROMOTE_TO_MAX, DEMOTE_TO_DOUBLE, NO_PROMOTION } + virtual bool PromoteDecimal() const { return true; } + Status Visit(const DataType&) { return Status::NotImplemented("No sum implemented"); } Status Visit(const HalfFloatType&) { @@ -187,10 +195,18 @@ struct SumLikeInit { return Status::OK(); } + // By default, we widen the decimal to max precision for SumLikes + // However, this may not be the desired behaviour (see, e.g., MeanKernelInit) template enable_if_decimal Visit(const Type&) { - state.reset(new KernelClass(type, options)); - return Status::OK(); + if (PromoteDecimal()) { + ARROW_ASSIGN_OR_RAISE(auto ty, WidenDecimalToMaxPrecision(type)); + state.reset(new KernelClass(ty, options)); + return Status::OK(); + } else { + state.reset(new KernelClass(type, options)); + return Status::OK(); + } } virtual Status Visit(const NullType&) { @@ -275,6 +291,8 @@ struct MeanKernelInit : public SumLikeInit { const ScalarAggregateOptions& options) : SumLikeInit(ctx, type, options) {} + bool PromoteDecimal() const override { return false; } + Status Visit(const NullType&) override { this->state.reset(new NullSumImpl(this->options)); return Status::OK(); diff --git a/cpp/src/arrow/compute/kernels/aggregate_test.cc b/cpp/src/arrow/compute/kernels/aggregate_test.cc index d821fc7e2c50..171aa17cc8b8 100644 --- a/cpp/src/arrow/compute/kernels/aggregate_test.cc +++ b/cpp/src/arrow/compute/kernels/aggregate_test.cc @@ -495,51 +495,67 @@ TEST_F(TestSumKernelRoundOff, Basics) { } TEST(TestDecimalSumKernel, SimpleSum) { - for (const auto& ty : {decimal128(3, 2), decimal256(3, 2)}) { + std::vector> init_types = {decimal128(3, 2), + decimal256(3, 2)}; + std::vector> out_types = {decimal128(38, 2), + decimal256(76, 2)}; + + for (size_t i = 0; i < init_types.size(); ++i) { + const auto& ty = init_types[i]; + const auto& out_ty = out_types[i]; + EXPECT_THAT(Sum(ArrayFromJSON(ty, R"([])")), - ResultWith(ScalarFromJSON(ty, R"(null)"))); + ResultWith(ScalarFromJSON(out_ty, R"(null)"))); EXPECT_THAT(Sum(ArrayFromJSON(ty, R"([null])")), - ResultWith(ScalarFromJSON(ty, R"(null)"))); + ResultWith(ScalarFromJSON(out_ty, R"(null)"))); EXPECT_THAT( Sum(ArrayFromJSON(ty, R"(["0.00", "1.01", "2.02", "3.03", "4.04", "5.05"])")), - ResultWith(ScalarFromJSON(ty, R"("15.15")"))); + ResultWith(ScalarFromJSON(out_ty, R"("15.15")"))); Datum chunks = ChunkedArrayFromJSON(ty, {R"(["0.00", "1.01", "2.02", "3.03", "4.04", "5.05"])"}); - EXPECT_THAT(Sum(chunks), ResultWith(ScalarFromJSON(ty, R"("15.15")"))); + EXPECT_THAT(Sum(chunks), ResultWith(ScalarFromJSON(out_ty, R"("15.15")"))); chunks = ChunkedArrayFromJSON( ty, {R"(["0.00", "1.01", "2.02"])", R"(["3.03", "4.04", "5.05"])"}); - EXPECT_THAT(Sum(chunks), ResultWith(ScalarFromJSON(ty, R"("15.15")"))); + EXPECT_THAT(Sum(chunks), ResultWith(ScalarFromJSON(out_ty, R"("15.15")"))); chunks = ChunkedArrayFromJSON( ty, {R"(["0.00", "1.01", "2.02"])", "[]", R"(["3.03", "4.04", "5.05"])"}); - EXPECT_THAT(Sum(chunks), ResultWith(ScalarFromJSON(ty, R"("15.15")"))); + EXPECT_THAT(Sum(chunks), ResultWith(ScalarFromJSON(out_ty, R"("15.15")"))); ScalarAggregateOptions options(/*skip_nulls=*/true, /*min_count=*/0); EXPECT_THAT(Sum(ArrayFromJSON(ty, R"([])"), options), - ResultWith(ScalarFromJSON(ty, R"("0.00")"))); + ResultWith(ScalarFromJSON(out_ty, R"("0.00")"))); EXPECT_THAT(Sum(ArrayFromJSON(ty, R"([null])"), options), - ResultWith(ScalarFromJSON(ty, R"("0.00")"))); + ResultWith(ScalarFromJSON(out_ty, R"("0.00")"))); chunks = ChunkedArrayFromJSON(ty, {}); - EXPECT_THAT(Sum(chunks, options), ResultWith(ScalarFromJSON(ty, R"("0.00")"))); + EXPECT_THAT(Sum(chunks, options), ResultWith(ScalarFromJSON(out_ty, R"("0.00")"))); EXPECT_THAT( Sum(ArrayFromJSON(ty, R"(["1.01", null, "3.03", null, "5.05", null, "7.07"])"), options), - ResultWith(ScalarFromJSON(ty, R"("16.16")"))); + ResultWith(ScalarFromJSON(out_ty, R"("16.16")"))); EXPECT_THAT(Sum(ScalarFromJSON(ty, R"("5.05")")), - ResultWith(ScalarFromJSON(ty, R"("5.05")"))); + ResultWith(ScalarFromJSON(out_ty, R"("5.05")"))); EXPECT_THAT(Sum(ScalarFromJSON(ty, R"(null)")), - ResultWith(ScalarFromJSON(ty, R"(null)"))); + ResultWith(ScalarFromJSON(out_ty, R"(null)"))); EXPECT_THAT(Sum(ScalarFromJSON(ty, R"(null)"), options), - ResultWith(ScalarFromJSON(ty, R"("0.00")"))); + ResultWith(ScalarFromJSON(out_ty, R"("0.00")"))); } } TEST(TestDecimalSumKernel, ScalarAggregateOptions) { - for (const auto& ty : {decimal128(3, 2), decimal256(3, 2)}) { - Datum null = ScalarFromJSON(ty, R"(null)"); - Datum zero = ScalarFromJSON(ty, R"("0.00")"); - Datum result = ScalarFromJSON(ty, R"("14.14")"); + std::vector> init_types = {decimal128(3, 2), + decimal256(3, 2)}; + std::vector> out_types = {decimal128(38, 2), + decimal256(76, 2)}; + + for (size_t i = 0; i < init_types.size(); ++i) { + auto& ty = init_types[i]; + auto& out_ty = out_types[i]; + + Datum null = ScalarFromJSON(out_ty, R"(null)"); + Datum zero = ScalarFromJSON(out_ty, R"("0.00")"); + Datum result = ScalarFromJSON(out_ty, R"("14.14")"); Datum arr = ArrayFromJSON(ty, R"(["1.01", null, "3.03", null, "3.03", null, "7.07"])"); @@ -579,7 +595,7 @@ TEST(TestDecimalSumKernel, ScalarAggregateOptions) { EXPECT_THAT(Sum(ScalarFromJSON(ty, R"("5.05")"), ScalarAggregateOptions(/*skip_nulls=*/false)), - ResultWith(ScalarFromJSON(ty, R"("5.05")"))); + ResultWith(ScalarFromJSON(out_ty, R"("5.05")"))); EXPECT_THAT(Sum(ScalarFromJSON(ty, R"("5.05")"), ScalarAggregateOptions(/*skip_nulls=*/true, /*min_count=*/2)), ResultWith(null)); @@ -712,49 +728,64 @@ TYPED_TEST(TestNumericProductKernel, ScalarAggregateOptions) { } TEST(TestDecimalProductKernel, SimpleProduct) { - for (const auto& ty : {decimal128(3, 2), decimal256(3, 2)}) { - Datum null = ScalarFromJSON(ty, R"(null)"); + std::vector> init_types = {decimal128(3, 2), + decimal256(3, 2)}; + std::vector> out_types = {decimal128(3, 2), decimal256(3, 2)}; + + for (size_t i = 0; i < init_types.size(); ++i) { + auto& ty = init_types[i]; + auto& out_ty = out_types[i]; + + Datum null = ScalarFromJSON(out_ty, R"(null)"); EXPECT_THAT(Product(ArrayFromJSON(ty, R"([])")), ResultWith(null)); EXPECT_THAT(Product(ArrayFromJSON(ty, R"([null])")), ResultWith(null)); EXPECT_THAT( Product(ArrayFromJSON(ty, R"(["0.00", "1.00", "2.00", "3.00", "4.00", "5.00"])")), - ResultWith(ScalarFromJSON(ty, R"("0.00")"))); + ResultWith(ScalarFromJSON(out_ty, R"("0.00")"))); Datum chunks = ChunkedArrayFromJSON(ty, {R"(["1.00", "2.00", "3.00", "4.00", "5.00"])"}); - EXPECT_THAT(Product(chunks), ResultWith(ScalarFromJSON(ty, R"("120.00")"))); + EXPECT_THAT(Product(chunks), ResultWith(ScalarFromJSON(out_ty, R"("120.00")"))); chunks = ChunkedArrayFromJSON(ty, {R"(["1.00", "2.00"])", R"(["-3.00", "4.00", "5.00"])"}); - EXPECT_THAT(Product(chunks), ResultWith(ScalarFromJSON(ty, R"("-120.00")"))); + EXPECT_THAT(Product(chunks), ResultWith(ScalarFromJSON(out_ty, R"("-120.00")"))); chunks = ChunkedArrayFromJSON( ty, {R"(["1.00", "2.00"])", R"([])", R"(["-3.00", "4.00", "-5.00"])"}); - EXPECT_THAT(Product(chunks), ResultWith(ScalarFromJSON(ty, R"("120.00")"))); + EXPECT_THAT(Product(chunks), ResultWith(ScalarFromJSON(out_ty, R"("120.00")"))); const ScalarAggregateOptions options(/*skip_nulls=*/true, /*min_count=*/0); EXPECT_THAT(Product(ArrayFromJSON(ty, R"([])"), options), - ResultWith(ScalarFromJSON(ty, R"("1.00")"))); + ResultWith(ScalarFromJSON(out_ty, R"("1.00")"))); EXPECT_THAT(Product(ArrayFromJSON(ty, R"([null])"), options), - ResultWith(ScalarFromJSON(ty, R"("1.00")"))); + ResultWith(ScalarFromJSON(out_ty, R"("1.00")"))); chunks = ChunkedArrayFromJSON(ty, {}); - EXPECT_THAT(Product(chunks, options), ResultWith(ScalarFromJSON(ty, R"("1.00")"))); + EXPECT_THAT(Product(chunks, options), + ResultWith(ScalarFromJSON(out_ty, R"("1.00")"))); EXPECT_THAT(Product(ArrayFromJSON( ty, R"(["1.00", null, "-3.00", null, "3.00", null, "7.00"])"), options), - ResultWith(ScalarFromJSON(ty, R"("-63.00")"))); + ResultWith(ScalarFromJSON(out_ty, R"("-63.00")"))); EXPECT_THAT(Product(ScalarFromJSON(ty, R"("5.00")")), - ResultWith(ScalarFromJSON(ty, R"("5.00")"))); + ResultWith(ScalarFromJSON(out_ty, R"("5.00")"))); EXPECT_THAT(Product(null), ResultWith(null)); } } TEST(TestDecimalProductKernel, ScalarAggregateOptions) { - for (const auto& ty : {decimal128(3, 2), decimal256(3, 2)}) { - Datum null = ScalarFromJSON(ty, R"(null)"); - Datum one = ScalarFromJSON(ty, R"("1.00")"); - Datum result = ScalarFromJSON(ty, R"("63.00")"); + std::vector> init_types = {decimal128(3, 2), + decimal256(3, 2)}; + std::vector> out_types = {decimal128(3, 2), decimal256(3, 2)}; + + for (size_t i = 0; i < init_types.size(); ++i) { + auto& ty = init_types[i]; + auto& out_ty = out_types[i]; + + Datum null = ScalarFromJSON(out_ty, R"(null)"); + Datum one = ScalarFromJSON(out_ty, R"("1.00")"); + Datum result = ScalarFromJSON(out_ty, R"("63.00")"); Datum empty = ArrayFromJSON(ty, R"([])"); Datum null_arr = ArrayFromJSON(ty, R"([null])"); @@ -806,7 +837,7 @@ TEST(TestDecimalProductKernel, ScalarAggregateOptions) { EXPECT_THAT(Product(ScalarFromJSON(ty, R"("5.00")"), ScalarAggregateOptions(/*skip_nulls=*/false)), - ResultWith(ScalarFromJSON(ty, R"("5.00")"))); + ResultWith(ScalarFromJSON(out_ty, R"("5.00")"))); EXPECT_THAT(Product(ScalarFromJSON(ty, R"("5.00")"), ScalarAggregateOptions(/*skip_nulls=*/true, /*min_count=*/2)), ResultWith(null)); @@ -1336,94 +1367,116 @@ TYPED_TEST(TestRandomNumericMeanKernel, RandomArrayMeanOverflow) { TEST(TestDecimalMeanKernel, SimpleMean) { ScalarAggregateOptions options(/*skip_nulls=*/true, /*min_count=*/0); + std::vector> init_types = {decimal128(3, 2), + decimal256(3, 2)}; + std::vector> out_types = {decimal128(3, 2), decimal256(3, 2)}; + + for (size_t i = 0; i < init_types.size(); ++i) { + auto& ty = init_types[i]; + auto& out_ty = out_types[i]; - for (const auto& ty : {decimal128(3, 2), decimal256(3, 2)}) { // Decimal doesn't have NaN EXPECT_THAT(Mean(ArrayFromJSON(ty, R"([])"), options), - ResultWith(ScalarFromJSON(ty, R"(null)"))); + ResultWith(ScalarFromJSON(out_ty, R"(null)"))); EXPECT_THAT(Mean(ArrayFromJSON(ty, R"([null])"), options), - ResultWith(ScalarFromJSON(ty, R"(null)"))); + ResultWith(ScalarFromJSON(out_ty, R"(null)"))); EXPECT_THAT(Mean(ArrayFromJSON(ty, R"([])")), - ResultWith(ScalarFromJSON(ty, R"(null)"))); + ResultWith(ScalarFromJSON(out_ty, R"(null)"))); EXPECT_THAT(Mean(ArrayFromJSON(ty, R"([null])")), - ResultWith(ScalarFromJSON(ty, R"(null)"))); + ResultWith(ScalarFromJSON(out_ty, R"(null)"))); EXPECT_THAT(Mean(ArrayFromJSON(ty, R"(["1.01", null, "1.01"])")), - ResultWith(ScalarFromJSON(ty, R"("1.01")"))); + ResultWith(ScalarFromJSON(out_ty, R"("1.01")"))); // Check rounding EXPECT_THAT( Mean(ArrayFromJSON( ty, R"(["1.01", "2.02", "3.03", "4.04", "5.05", "6.06", "7.07", "8.08"])")), // 4.545 unrounded - ResultWith(ScalarFromJSON(ty, R"("4.55")"))); + ResultWith(ScalarFromJSON(out_ty, R"("4.55")"))); EXPECT_THAT( Mean(ArrayFromJSON( ty, R"(["-1.01", "-2.02", "-3.03", "-4.04", "-5.05", "-6.06", "-7.07", "-8.08"])")), // -4.545 unrounded - ResultWith(ScalarFromJSON(ty, R"("-4.55")"))); + ResultWith(ScalarFromJSON(out_ty, R"("-4.55")"))); EXPECT_THAT( Mean(ArrayFromJSON( ty, R"(["1.01", "2.02", "3.00", "4.04", "5.05", "6.06", "7.07", "8.08"])")), // 4.54125 unrounded - ResultWith(ScalarFromJSON(ty, R"("4.54")"))); + ResultWith(ScalarFromJSON(out_ty, R"("4.54")"))); EXPECT_THAT( Mean(ArrayFromJSON( ty, R"(["-1.01", "-2.02", "-3.00", "-4.04", "-5.05", "-6.06", "-7.07", "-8.08"])")), // -4.54125 unrounded - ResultWith(ScalarFromJSON(ty, R"("-4.54")"))); + ResultWith(ScalarFromJSON(out_ty, R"("-4.54")"))); EXPECT_THAT( Mean(ArrayFromJSON( ty, R"(["0.00", "0.00", "0.00", "0.00", "0.00", "0.00", "0.00", "0.00"])")), - ResultWith(ScalarFromJSON(ty, R"("0.00")"))); + ResultWith(ScalarFromJSON(out_ty, R"("0.00")"))); EXPECT_THAT( Mean(ArrayFromJSON( ty, R"(["1.01", "1.01", "1.01", "1.01", "1.01", "1.01", "1.01", "1.01"])")), - ResultWith(ScalarFromJSON(ty, R"("1.01")"))); + ResultWith(ScalarFromJSON(out_ty, R"("1.01")"))); EXPECT_THAT(Mean(ScalarFromJSON(ty, R"("5.05")")), - ResultWith(ScalarFromJSON(ty, R"("5.05")"))); + ResultWith(ScalarFromJSON(out_ty, R"("5.05")"))); EXPECT_THAT(Mean(ScalarFromJSON(ty, R"(null)")), - ResultWith(ScalarFromJSON(ty, R"(null)"))); + ResultWith(ScalarFromJSON(out_ty, R"(null)"))); } - for (const auto& ty : {decimal128(3, -2), decimal256(3, -2)}) { + init_types = {decimal128(3, -2), decimal256(3, -2)}; + out_types = {decimal128(3, -2), decimal256(3, -2)}; + + for (size_t i = 0; i < init_types.size(); ++i) { + auto& ty = init_types[i]; + auto& out_ty = out_types[i]; + // Check rounding + // + // N.B. In what follows, the additional Cast is due to the implementation of + // DecimalScalarFromJSON, which will try to construct a decimal with too big precision EXPECT_THAT( Mean(DecimalArrayFromJSON( ty, R"(["101E2", "202E2", "303E2", "404E2", "505E2", "606E2", "707E2", "808E2"])")), // 45450 unrounded - ResultWith(DecimalScalarFromJSON(ty, R"("455E2")"))); + ResultWith(Cast(DecimalScalarFromJSON(ty, R"("455E2")"), out_ty))); EXPECT_THAT( Mean(DecimalArrayFromJSON( ty, R"(["-101E2", "-202E2", "-303E2", "-404E2", "-505E2", "-606E2", "-707E2", "-808E2"])")), // -45450 unrounded - ResultWith(DecimalScalarFromJSON(ty, R"("-455E2")"))); + ResultWith(Cast(DecimalScalarFromJSON(ty, R"("-455E2")"), out_ty))); EXPECT_THAT( Mean(DecimalArrayFromJSON( ty, R"(["101E2", "202E2", "300E2", "404E2", "505E2", "606E2", "707E2", "808E2"])")), // 45412.5 unrounded - ResultWith(DecimalScalarFromJSON(ty, R"("454E2")"))); + ResultWith(Cast(DecimalScalarFromJSON(ty, R"("454E2")"), out_ty))); EXPECT_THAT( Mean(DecimalArrayFromJSON( ty, R"(["-101E2", "-202E2", "-300E2", "-404E2", "-505E2", "-606E2", "-707E2", "-808E2"])")), // -45412.5 unrounded - ResultWith(DecimalScalarFromJSON(ty, R"("-454E2")"))); + ResultWith(Cast(DecimalScalarFromJSON(ty, R"("-454E2")"), out_ty))); } } TEST(TestDecimalMeanKernel, ScalarAggregateOptions) { - for (const auto& ty : {decimal128(3, 2), decimal256(3, 2)}) { - Datum result = ScalarFromJSON(ty, R"("3.03")"); - Datum null = ScalarFromJSON(ty, R"(null)"); + std::vector> init_types = {decimal128(3, 2), + decimal256(3, 2)}; + std::vector> out_types = {decimal128(3, 2), decimal256(3, 2)}; + + for (size_t i = 0; i < init_types.size(); ++i) { + auto& ty = init_types[i]; + auto& out_ty = out_types[i]; + + Datum result = ScalarFromJSON(out_ty, R"("3.03")"); + Datum null = ScalarFromJSON(out_ty, R"(null)"); Datum arr = ArrayFromJSON(ty, R"(["1.01", null, "2.02", "2.02", null, "7.07"])"); EXPECT_THAT(Mean(ArrayFromJSON(ty, "[]"), @@ -1481,8 +1534,8 @@ TEST(TestDecimalMeanKernel, ScalarAggregateOptions) { EXPECT_THAT(Mean(ScalarFromJSON(ty, R"("5.05")"), ScalarAggregateOptions(/*skip_nulls=*/false)), - ResultWith(ScalarFromJSON(ty, R"("5.05")"))); - EXPECT_THAT(Mean(ScalarFromJSON(ty, R"("5.05")"), + ResultWith(ScalarFromJSON(out_ty, R"("5.05")"))); + EXPECT_THAT(Mean(ScalarFromJSON(out_ty, R"("5.05")"), ScalarAggregateOptions(/*skip_nulls=*/true, /*min_count=*/2)), ResultWith(null)); EXPECT_THAT(Mean(null, ScalarAggregateOptions(/*skip_nulls=*/false)), diff --git a/cpp/src/arrow/compute/kernels/codegen_internal.cc b/cpp/src/arrow/compute/kernels/codegen_internal.cc index 5f3efedec352..10ed9344d976 100644 --- a/cpp/src/arrow/compute/kernels/codegen_internal.cc +++ b/cpp/src/arrow/compute/kernels/codegen_internal.cc @@ -77,6 +77,11 @@ Result ListValuesType(KernelContext* ctx, return value_type; } +Result MaxPrecisionDecimalType(KernelContext*, + const std::vector& args) { + return WidenDecimalToMaxPrecision(args[0].GetSharedPtr()); +} + void EnsureDictionaryDecoded(std::vector* types) { EnsureDictionaryDecoded(types->data(), types->size()); } @@ -528,6 +533,28 @@ Status CastDecimalArgs(TypeHolder* begin, size_t count) { return Status::OK(); } +Result> WidenDecimalToMaxPrecision( + std::shared_ptr type) { + DCHECK(is_decimal(type->id())); + auto cast_type = checked_pointer_cast(type); + switch (type->id()) { + case Type::DECIMAL32: + return Decimal32Type::Make(Decimal32Type::kMaxPrecision, cast_type->scale()); + case Type::DECIMAL64: + return Decimal64Type::Make(Decimal64Type::kMaxPrecision, cast_type->scale()); + case Type::DECIMAL128: + return Decimal128Type::Make(Decimal128Type::kMaxPrecision, cast_type->scale()); + case Type::DECIMAL256: + return Decimal256Type::Make(Decimal256Type::kMaxPrecision, cast_type->scale()); + default: + DCHECK(false) << "An unknown DecimalType was passed to WidenDecimalToMaxPrecision: " + << type->ToString(); + return Status::TypeError( + "An unknown DecimalType was passed to WidenDecimalToMaxPrecision: " + + type->ToString()); + } +} + bool HasDecimal(const std::vector& types) { for (const auto& th : types) { if (is_decimal(th.id())) { diff --git a/cpp/src/arrow/compute/kernels/codegen_internal.h b/cpp/src/arrow/compute/kernels/codegen_internal.h index 289ba25f0b78..0190629691fe 100644 --- a/cpp/src/arrow/compute/kernels/codegen_internal.h +++ b/cpp/src/arrow/compute/kernels/codegen_internal.h @@ -55,6 +55,7 @@ using internal::BinaryBitBlockCounter; using internal::BitBlockCount; using internal::BitmapReader; using internal::checked_cast; +using internal::checked_pointer_cast; using internal::FirstTimeBitmapWriter; using internal::GenerateBitsUnrolled; using internal::VisitBitBlocks; @@ -480,6 +481,8 @@ ARROW_EXPORT Result LastType(KernelContext*, const std::vector& types); ARROW_EXPORT Result ListValuesType(KernelContext* ctx, const std::vector& types); +ARROW_EXPORT Result MaxPrecisionDecimalType( + KernelContext*, const std::vector& types); // ---------------------------------------------------------------------- // Helpers for iterating over common DataType instances for adding kernels to @@ -1444,6 +1447,12 @@ Status CastBinaryDecimalArgs(DecimalPromotion promotion, std::vector ARROW_EXPORT Status CastDecimalArgs(TypeHolder* begin, size_t count); +/// Given a DataType, if it is a DecimalType, return a DecimalType with the same scale +/// and the maximum precision for that DecimalType. +ARROW_EXPORT +Result> WidenDecimalToMaxPrecision( + std::shared_ptr type); + ARROW_EXPORT bool HasDecimal(const std::vector& types); diff --git a/cpp/src/arrow/compute/kernels/codegen_internal_test.cc b/cpp/src/arrow/compute/kernels/codegen_internal_test.cc index 6bb5568d2ff3..8aa90823c1dd 100644 --- a/cpp/src/arrow/compute/kernels/codegen_internal_test.cc +++ b/cpp/src/arrow/compute/kernels/codegen_internal_test.cc @@ -129,6 +129,32 @@ TEST(TestDispatchBest, CastDecimalArgs) { AssertTypeEqual(*args[2], *utf8()); } +TEST(TestDecimalPromotion, WidenDecimalToMaxPrecision) { + std::shared_ptr arg; + std::shared_ptr expected; + std::shared_ptr unwrapped; + + arg = decimal32(3, 2); + expected = decimal32(9, 2); + ASSERT_OK_AND_ASSIGN(unwrapped, WidenDecimalToMaxPrecision(arg)); + AssertTypeEqual(*unwrapped, *expected); + + arg = decimal64(3, 2); + expected = decimal64(18, 2); + ASSERT_OK_AND_ASSIGN(unwrapped, WidenDecimalToMaxPrecision(arg)); + AssertTypeEqual(*unwrapped, *expected); + + arg = decimal128(3, 2); + expected = decimal128(38, 2); + ASSERT_OK_AND_ASSIGN(unwrapped, WidenDecimalToMaxPrecision(arg)); + AssertTypeEqual(*unwrapped, *expected); + + arg = decimal256(3, 2); + expected = decimal256(76, 2); + ASSERT_OK_AND_ASSIGN(unwrapped, WidenDecimalToMaxPrecision(arg)); + AssertTypeEqual(*unwrapped, *expected); +} + TEST(TestDispatchBest, CommonTemporal) { std::vector args; diff --git a/cpp/src/arrow/compute/kernels/hash_aggregate_numeric.cc b/cpp/src/arrow/compute/kernels/hash_aggregate_numeric.cc index 6e0652efe693..a46f790548c3 100644 --- a/cpp/src/arrow/compute/kernels/hash_aggregate_numeric.cc +++ b/cpp/src/arrow/compute/kernels/hash_aggregate_numeric.cc @@ -54,7 +54,7 @@ struct GroupedReducingAggregator : public GroupedAggregator { reduced_ = TypedBufferBuilder(pool_); counts_ = TypedBufferBuilder(pool_); no_nulls_ = TypedBufferBuilder(pool_); - out_type_ = GetOutType(args.inputs[0].GetSharedPtr()); + ARROW_ASSIGN_OR_RAISE(out_type_, GetOutType(args.inputs[0].GetSharedPtr())); return Status::OK(); } @@ -154,17 +154,29 @@ struct GroupedReducingAggregator : public GroupedAggregator { std::shared_ptr out_type() const override { return out_type_; } template - static enable_if_t::value, std::shared_ptr> GetOutType( + enable_if_t::value, Result>> GetOutType( const std::shared_ptr& in_type) { return TypeTraits::type_singleton(); } template - static enable_if_decimal> GetOutType( + enable_if_decimal>> GetOutType( const std::shared_ptr& in_type) { - return in_type; + if (PromoteDecimal()) { + return WidenDecimalToMaxPrecision(in_type); + } else { + return in_type; + } } + // If this returns true, then the aggregator will promote a decimal to the maximum + // precision for that type. For instance, a decimal128(3, 2) will be promoted to a + // decimal128(38, 2) + // + // TODO: Ideally this should be configurable via the function options with an enum + // PrecisionPolicy { PROMOTE_TO_MAX, DEMOTE_TO_DOUBLE, NO_PROMOTION } + virtual bool PromoteDecimal() const { return true; } + int64_t num_groups_ = 0; ScalarAggregateOptions options_; TypedBufferBuilder reduced_; @@ -317,6 +329,8 @@ struct GroupedProductImpl final return MultiplyTraits::Multiply(out_type, u, v); } + bool PromoteDecimal() const override { return false; } + using Base::Finish; }; @@ -415,6 +429,8 @@ struct GroupedMeanImpl return values; } + bool PromoteDecimal() const override { return false; } + std::shared_ptr out_type() const override { if (is_decimal_type::value) return this->out_type_; return float64(); diff --git a/docs/source/cpp/compute.rst b/docs/source/cpp/compute.rst index b25ece967c13..448dd972e6ef 100644 --- a/docs/source/cpp/compute.rst +++ b/docs/source/cpp/compute.rst @@ -183,8 +183,7 @@ To avoid exhaustively listing supported types, the tables below use a number of general type categories: * "Numeric": Integer types (Int8, etc.) and Floating-point types (Float32, - Float64, sometimes Float16). Some functions also accept Decimal128 and - Decimal256 input. + Float64, sometimes Float16). Some functions also accept Decimal input. * "Temporal": Date types (Date32, Date64), Time types (Time32, Time64), Timestamp, Duration, Interval. @@ -211,57 +210,57 @@ Aggregations Scalar aggregations operate on a (chunked) array or scalar value and reduce the input to a single output value. -+--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+-------+ -| Function name | Arity | Input types | Output type | Options class | Notes | -+====================+=========+===============================================+========================+==================================+=======+ -| all | Unary | Boolean | Scalar Boolean | :struct:`ScalarAggregateOptions` | \(1) | -+--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+-------+ -| any | Unary | Boolean | Scalar Boolean | :struct:`ScalarAggregateOptions` | \(1) | -+--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+-------+ -| approximate_median | Unary | Numeric | Scalar Float64 | :struct:`ScalarAggregateOptions` | | -+--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+-------+ -| count | Unary | Any | Scalar Int64 | :struct:`CountOptions` | \(2) | -+--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+-------+ -| count_all | Nullary | | Scalar Int64 | | | -+--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+-------+ -| count_distinct | Unary | Non-nested types | Scalar Int64 | :struct:`CountOptions` | \(2) | -+--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+-------+ -| first | Unary | Numeric, Binary | Scalar Input type | :struct:`ScalarAggregateOptions` | \(3) | -+--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+-------+ -| first_last | Unary | Numeric, Binary | Scalar Struct | :struct:`ScalarAggregateOptions` | \(3) | -+--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+-------+ -| index | Unary | Any | Scalar Int64 | :struct:`IndexOptions` | \(4) | -+--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+-------+ -| kurtosis | Unary | Numeric | Scalar Float64 | :struct:`SkewOptions` | \(11) | -+--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+-------+ -| last | Unary | Numeric, Binary | Scalar Input type | :struct:`ScalarAggregateOptions` | \(3) | -+--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+-------+ -| max | Unary | Non-nested types | Scalar Input type | :struct:`ScalarAggregateOptions` | | -+--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+-------+ -| mean | Unary | Numeric | Scalar Decimal/Float64 | :struct:`ScalarAggregateOptions` | \(5) | -+--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+-------+ -| min | Unary | Non-nested types | Scalar Input type | :struct:`ScalarAggregateOptions` | | -+--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+-------+ -| min_max | Unary | Non-nested types | Scalar Struct | :struct:`ScalarAggregateOptions` | \(6) | -+--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+-------+ -| mode | Unary | Numeric | Struct | :struct:`ModeOptions` | \(7) | -+--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+-------+ -| pivot_wider | Binary | Binary, String, Integer (Arg 0); Any (Arg 1) | Scalar Struct | :struct:`PivotWiderOptions` | \(8) | -+--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+-------+ -| product | Unary | Numeric | Scalar Numeric | :struct:`ScalarAggregateOptions` | \(9) | -+--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+-------+ -| quantile | Unary | Numeric | Scalar Numeric | :struct:`QuantileOptions` | \(10) | -+--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+-------+ -| skew | Unary | Numeric | Scalar Float64 | :struct:`SkewOptions` | \(11) | -+--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+-------+ -| stddev | Unary | Numeric | Scalar Float64 | :struct:`VarianceOptions` | \(11) | -+--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+-------+ -| sum | Unary | Numeric | Scalar Numeric | :struct:`ScalarAggregateOptions` | \(9) | -+--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+-------+ -| tdigest | Unary | Numeric | Float64 | :struct:`TDigestOptions` | \(12) | -+--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+-------+ -| variance | Unary | Numeric | Scalar Float64 | :struct:`VarianceOptions` | \(11) | -+--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+-------+ ++--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+------------+ +| Function name | Arity | Input types | Output type | Options class | Notes | ++====================+=========+===============================================+========================+==================================+============+ +| all | Unary | Boolean | Scalar Boolean | :struct:`ScalarAggregateOptions` | \(1) | ++--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+------------+ +| any | Unary | Boolean | Scalar Boolean | :struct:`ScalarAggregateOptions` | \(1) | ++--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+------------+ +| approximate_median | Unary | Numeric | Scalar Float64 | :struct:`ScalarAggregateOptions` | | ++--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+------------+ +| count | Unary | Any | Scalar Int64 | :struct:`CountOptions` | \(2) | ++--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+------------+ +| count_all | Nullary | | Scalar Int64 | | | ++--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+------------+ +| count_distinct | Unary | Non-nested types | Scalar Int64 | :struct:`CountOptions` | \(2) | ++--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+------------+ +| first | Unary | Numeric, Binary | Scalar Input type | :struct:`ScalarAggregateOptions` | \(3) | ++--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+------------+ +| first_last | Unary | Numeric, Binary | Scalar Struct | :struct:`ScalarAggregateOptions` | \(3) | ++--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+------------+ +| index | Unary | Any | Scalar Int64 | :struct:`IndexOptions` | \(4) | ++--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+------------+ +| kurtosis | Unary | Numeric | Scalar Float64 | :struct:`SkewOptions` | \(12) | ++--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+------------+ +| last | Unary | Numeric, Binary | Scalar Input type | :struct:`ScalarAggregateOptions` | \(3) | ++--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+------------+ +| max | Unary | Non-nested types | Scalar Input type | :struct:`ScalarAggregateOptions` | | ++--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+------------+ +| mean | Unary | Numeric | Scalar Decimal/Float64 | :struct:`ScalarAggregateOptions` | \(5) | ++--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+------------+ +| min | Unary | Non-nested types | Scalar Input type | :struct:`ScalarAggregateOptions` | | ++--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+------------+ +| min_max | Unary | Non-nested types | Scalar Struct | :struct:`ScalarAggregateOptions` | \(6) | ++--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+------------+ +| mode | Unary | Numeric | Struct | :struct:`ModeOptions` | \(7) | ++--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+------------+ +| pivot_wider | Binary | Binary, String, Integer (Arg 0); Any (Arg 1) | Scalar Struct | :struct:`PivotWiderOptions` | \(8) | ++--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+------------+ +| product | Unary | Numeric | Scalar Numeric | :struct:`ScalarAggregateOptions` | \(9) | ++--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+------------+ +| quantile | Unary | Numeric | Scalar Numeric | :struct:`QuantileOptions` | \(11) | ++--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+------------+ +| skew | Unary | Numeric | Scalar Float64 | :struct:`SkewOptions` | \(12) | ++--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+------------+ +| stddev | Unary | Numeric | Scalar Float64 | :struct:`VarianceOptions` | \(12) | ++--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+------------+ +| sum | Unary | Numeric | Scalar Numeric | :struct:`ScalarAggregateOptions` | \(9) \(10) | ++--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+------------+ +| tdigest | Unary | Numeric | Float64 | :struct:`TDigestOptions` | \(13) | ++--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+------------+ +| variance | Unary | Numeric | Scalar Float64 | :struct:`VarianceOptions` | \(12) | ++--------------------+---------+-----------------------------------------------+------------------------+----------------------------------+------------+ * \(1) If null values are taken into account, by setting the ScalarAggregateOptions parameter skip_nulls = false, then `Kleene logic`_ @@ -297,11 +296,15 @@ the input to a single output value. * \(9) Output is Int64, UInt64, Float64, or Decimal128/256, depending on the input type. -* \(10) Output is Float64 or input type, depending on QuantileOptions. +* \(10) For Decimal input, the output precision is increased to the maximum + precision for the input type's width. For instance, an array of + ``decimal128(3, 2)`` will return a ``decimal128(38, 2)`` scalar. -* \(11) Decimal arguments are cast to Float64 first. +* \(11) Output is Float64 or input type, depending on QuantileOptions. -* \(12) tdigest/t-digest computes approximate quantiles, and so only needs a +* \(12) Decimal arguments are cast to Float64 first. + +* \(13) tdigest/t-digest computes approximate quantiles, and so only needs a fixed amount of memory. See the `reference implementation `_ for details. diff --git a/python/pyarrow/tests/test_compute.py b/python/pyarrow/tests/test_compute.py index 643b5c27eb96..ac2b18669265 100644 --- a/python/pyarrow/tests/test_compute.py +++ b/python/pyarrow/tests/test_compute.py @@ -359,6 +359,49 @@ def test_sum_array(arrow_type): assert pc.sum(arr, min_count=0).as_py() == 0 +@pytest.mark.parametrize("arrow_type", [pa.decimal128(3, 2), pa.decimal256(3, 2)]) +def test_sum_decimal_array(arrow_type): + from decimal import Decimal + max_precision_type = ( + pa.decimal128(38, arrow_type.scale) + if pa.types.is_decimal128(arrow_type) + else pa.decimal256(76, arrow_type.scale) + ) + expected_sum = Decimal("5.79") + expected_sum_overflow = Decimal("10.00") + zero = Decimal("0.00") + + # No overflow + arr = pa.array([Decimal("1.23"), Decimal("4.56")], type=arrow_type) + assert arr.sum().as_py() == expected_sum + assert arr.sum().type == max_precision_type + + arr = pa.array([Decimal("1.23"), Decimal("4.56"), None], type=arrow_type) + assert arr.sum().as_py() == expected_sum + assert arr.sum().type == max_precision_type + + # With overflow + arr = pa.array([Decimal("1.23"), Decimal("8.77")], type=arrow_type) + assert arr.sum().as_py() == expected_sum_overflow + assert arr.sum().type == max_precision_type + + arr = pa.array([Decimal("1.23"), Decimal("8.77"), None], type=arrow_type) + assert arr.sum().as_py() == expected_sum_overflow + assert arr.sum().type == max_precision_type + + arr = pa.array([None], type=arrow_type) + assert arr.sum().as_py() is None # noqa: E711 + assert arr.sum().type == max_precision_type # noqa: E711 + assert arr.sum(min_count=0).as_py() == zero + assert arr.sum(min_count=0).type == max_precision_type + + arr = pa.array([], type=arrow_type) + assert arr.sum().as_py() is None # noqa: E711 + assert arr.sum().type == max_precision_type # noqa: E711 + assert arr.sum(min_count=0).as_py() == zero + assert arr.sum(min_count=0).type == max_precision_type + + @pytest.mark.parametrize('arrow_type', numerical_arrow_types) def test_sum_chunked_array(arrow_type): arr = pa.chunked_array([pa.array([1, 2, 3, 4], type=arrow_type)]) @@ -382,6 +425,48 @@ def test_sum_chunked_array(arrow_type): assert pc.sum(arr, min_count=0).as_py() == 0 +@pytest.mark.parametrize('arrow_type', [pa.decimal128(3, 2), pa.decimal256(3, 2)]) +def test_sum_chunked_array_decimal_type(arrow_type): + from decimal import Decimal + max_precision_type = ( + pa.decimal128(38, arrow_type.scale) + if pa.types.is_decimal128(arrow_type) + else pa.decimal256(76, arrow_type.scale) + ) + expected_sum = Decimal("5.79") + zero = Decimal("0.00") + + arr = pa.chunked_array( + [ + pa.array([Decimal("1.23"), Decimal("4.56")], type=arrow_type) + ] + ) + assert pc.sum(arr).as_py() == expected_sum + assert pc.sum(arr).type == max_precision_type + + arr = pa.chunked_array([ + pa.array([Decimal("1.23")], type=arrow_type), + pa.array([Decimal("4.56")], type=arrow_type) + ]) + assert pc.sum(arr).as_py() == expected_sum + assert pc.sum(arr).type == max_precision_type + + arr = pa.chunked_array([ + pa.array([Decimal("1.23")], type=arrow_type), + pa.array([], type=arrow_type), + pa.array([Decimal("4.56")], type=arrow_type) + ]) + assert pc.sum(arr).as_py() == expected_sum + assert pc.sum(arr).type == max_precision_type + + arr = pa.chunked_array((), type=arrow_type) + assert arr.num_chunks == 0 + assert pc.sum(arr).as_py() is None # noqa: E711 + assert pc.sum(arr).type == max_precision_type + assert pc.sum(arr, min_count=0).as_py() == zero + assert pc.sum(arr, min_count=0).type == max_precision_type + + def test_mode_array(): # ARROW-9917 arr = pa.array([1, 1, 3, 4, 3, 5], type='int64')