From b8f1d4b707037125b81fd0a33db1895acf3dabc2 Mon Sep 17 00:00:00 2001 From: Kevin H Wilson Date: Tue, 17 Sep 2024 14:55:41 -0400 Subject: [PATCH 01/38] Increase the precision of decimals during compute This is an initial pass whereby a scalar aggregate of a Decimal type increases its precision to the maximum. That is, a sum of an array of decimal128(3, 2)'s becomes a decimal128(38, 2). Previously, the exact decimal type was preserved (e.g., a sum of decimal128(3, 2)'s was a decimal128(3, 2)) *regardless* of whether that was enough precision to capture the full decimal value. --- .../arrow/compute/kernels/aggregate_basic.cc | 38 ++++++++-- .../compute/kernels/aggregate_basic.inc.cc | 30 +++++++- python/pyarrow/tests/test_compute.py | 74 +++++++++++++++++++ 3 files changed, 130 insertions(+), 12 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/aggregate_basic.cc b/cpp/src/arrow/compute/kernels/aggregate_basic.cc index ee2c615bbfb7..b5f746a18a14 100644 --- a/cpp/src/arrow/compute/kernels/aggregate_basic.cc +++ b/cpp/src/arrow/compute/kernels/aggregate_basic.cc @@ -366,11 +366,20 @@ struct ProductImpl : public ScalarAggregator { } Status Finalize(KernelContext*, Datum* out) override { + std::shared_ptr out_type_; + if (auto decimal128_type = std::dynamic_pointer_cast(this->out_type)) { + ARROW_ASSIGN_OR_RAISE(out_type_, Decimal128Type::Make(Decimal128Type::kMaxPrecision, decimal128_type->scale())); + } else if (auto decimal256_type = std::dynamic_pointer_cast(this->out_type)) { + ARROW_ASSIGN_OR_RAISE(out_type_, Decimal256Type::Make(Decimal256Type::kMaxPrecision, decimal256_type->scale())); + } else { + out_type_ = out_type; + } + if ((!options.skip_nulls && this->nulls_observed) || (this->count < options.min_count)) { - out->value = std::make_shared(out_type); + out->value = std::make_shared(out_type_); } else { - out->value = std::make_shared(this->product, out_type); + out->value = std::make_shared(this->product, out_type_); } return Status::OK(); } @@ -1021,6 +1030,19 @@ const FunctionDoc index_doc{"Find the index of the first occurrence of a given v } // namespace + +Result MaxPrecisionDecimalType(KernelContext*, const std::vector& types) { + std::shared_ptr out_type_; + if (auto decimal128_type = std::dynamic_pointer_cast(types.front().GetSharedPtr())) { + ARROW_ASSIGN_OR_RAISE(out_type_, Decimal128Type::Make(Decimal128Type::kMaxPrecision, decimal128_type->scale())); + } else if (auto decimal256_type = std::dynamic_pointer_cast(types.front().GetSharedPtr())) { + ARROW_ASSIGN_OR_RAISE(out_type_, Decimal256Type::Make(Decimal256Type::kMaxPrecision, decimal256_type->scale())); + } else { + return Status::TypeError("Bad call"); + } + return TypeHolder(out_type_); +} + void RegisterScalarAggregateBasic(FunctionRegistry* registry) { static auto default_scalar_aggregate_options = ScalarAggregateOptions::Defaults(); static auto default_count_options = CountOptions::Defaults(); @@ -1049,9 +1071,9 @@ 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(), + AddAggKernel(KernelSignature::Make({Type::DECIMAL128}, MaxPrecisionDecimalType), SumInit, func.get(), SimdLevel::NONE); - AddAggKernel(KernelSignature::Make({Type::DECIMAL256}, FirstType), SumInit, func.get(), + AddAggKernel(KernelSignature::Make({Type::DECIMAL256}, MaxPrecisionDecimalType), SumInit, func.get(), SimdLevel::NONE); AddArrayScalarAggKernels(SumInit, SignedIntTypes(), int64(), func.get()); AddArrayScalarAggKernels(SumInit, UnsignedIntTypes(), uint64(), func.get()); @@ -1077,9 +1099,9 @@ void RegisterScalarAggregateBasic(FunctionRegistry* registry) { &default_scalar_aggregate_options); AddArrayScalarAggKernels(MeanInit, {boolean()}, float64(), func.get()); AddArrayScalarAggKernels(MeanInit, NumericTypes(), float64(), func.get()); - AddAggKernel(KernelSignature::Make({Type::DECIMAL128}, FirstType), MeanInit, func.get(), + AddAggKernel(KernelSignature::Make({Type::DECIMAL128}, MaxPrecisionDecimalType), MeanInit, func.get(), SimdLevel::NONE); - AddAggKernel(KernelSignature::Make({Type::DECIMAL256}, FirstType), MeanInit, func.get(), + AddAggKernel(KernelSignature::Make({Type::DECIMAL256}, MaxPrecisionDecimalType), MeanInit, func.get(), SimdLevel::NONE); AddArrayScalarAggKernels(MeanInit, {null()}, float64(), func.get()); // Add the SIMD variants for mean @@ -1161,9 +1183,9 @@ void RegisterScalarAggregateBasic(FunctionRegistry* registry) { AddArrayScalarAggKernels(ProductInit::Init, UnsignedIntTypes(), uint64(), func.get()); AddArrayScalarAggKernels(ProductInit::Init, FloatingPointTypes(), float64(), func.get()); - AddAggKernel(KernelSignature::Make({Type::DECIMAL128}, FirstType), ProductInit::Init, + AddAggKernel(KernelSignature::Make({Type::DECIMAL128}, MaxPrecisionDecimalType), ProductInit::Init, func.get(), SimdLevel::NONE); - AddAggKernel(KernelSignature::Make({Type::DECIMAL256}, FirstType), ProductInit::Init, + AddAggKernel(KernelSignature::Make({Type::DECIMAL256}, MaxPrecisionDecimalType), ProductInit::Init, func.get(), SimdLevel::NONE); AddArrayScalarAggKernels(ProductInit::Init, {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 index 49010d182cd6..d1b173f90398 100644 --- a/cpp/src/arrow/compute/kernels/aggregate_basic.inc.cc +++ b/cpp/src/arrow/compute/kernels/aggregate_basic.inc.cc @@ -93,11 +93,20 @@ struct SumImpl : public ScalarAggregator { } Status Finalize(KernelContext*, Datum* out) override { + std::shared_ptr out_type_; + if (auto decimal128_type = std::dynamic_pointer_cast(out_type)) { + ARROW_ASSIGN_OR_RAISE(out_type_, Decimal128Type::Make(Decimal128Type::kMaxPrecision, decimal128_type->scale())); + } else if (auto decimal256_type = std::dynamic_pointer_cast(out_type)) { + ARROW_ASSIGN_OR_RAISE(out_type_, Decimal256Type::Make(Decimal256Type::kMaxPrecision, decimal256_type->scale())); + } else { + out_type_ = out_type; + } + if ((!options.skip_nulls && this->nulls_observed) || (this->count < options.min_count)) { - out->value = std::make_shared(out_type); + out->value = std::make_shared(out_type_); } else { - out->value = std::make_shared(this->sum, out_type); + out->value = std::make_shared(this->sum, out_type_); } return Status::OK(); } @@ -220,9 +229,22 @@ struct MeanImpl> template Status FinalizeImpl(Datum* out) { + std::shared_ptr out_type_; + if (auto decimal128_type = std::dynamic_pointer_cast(this->out_type)) { + ARROW_ASSIGN_OR_RAISE(out_type_, Decimal128Type::Make(Decimal128Type::kMaxPrecision, decimal128_type->scale())); + } else if (auto decimal256_type = std::dynamic_pointer_cast(this->out_type)) { + ARROW_ASSIGN_OR_RAISE(out_type_, Decimal256Type::Make(Decimal256Type::kMaxPrecision, decimal256_type->scale())); + } else { + return Status::TypeError( + "The decimal specialization of MeanImpl was passed a type ", + this->out_type->ToString(), + " and not a decimal type" + ); + } + if ((!options.skip_nulls && this->nulls_observed) || (this->count < options.min_count) || (this->count == 0)) { - out->value = std::make_shared(this->out_type); + out->value = std::make_shared(out_type_); } else { SumCType quotient, remainder; ARROW_ASSIGN_OR_RAISE(std::tie(quotient, remainder), this->sum.Divide(this->count)); @@ -235,7 +257,7 @@ struct MeanImpl> quotient -= 1; } } - out->value = std::make_shared(quotient, this->out_type); + out->value = std::make_shared(quotient, out_type_); } return Status::OK(); } diff --git a/python/pyarrow/tests/test_compute.py b/python/pyarrow/tests/test_compute.py index 643b5c27eb96..4f142e40eef8 100644 --- a/python/pyarrow/tests/test_compute.py +++ b/python/pyarrow/tests/test_compute.py @@ -355,10 +355,51 @@ def test_sum_array(arrow_type): arr = pa.array([], type=arrow_type) assert arr.sum().as_py() is None # noqa: E711 + assert pc.sum(arr).as_py() is None # noqa: E711 assert arr.sum(min_count=0).as_py() == 0 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") + zero = Decimal("0.00") + + 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 + assert pc.sum(arr).as_py() == expected_sum + assert pc.sum(arr).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 + assert pc.sum(arr).as_py() == expected_sum + assert pc.sum(arr).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 pc.sum(arr).as_py() is None # noqa: E711 + assert pc.sum(arr).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 + assert pc.sum(arr, min_count=0).as_py() == zero + assert pc.sum(arr, 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 pc.sum(arr).as_py() is None # noqa: E711 + assert pc.sum(arr).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 + assert pc.sum(arr, min_count=0).as_py() == zero + assert pc.sum(arr, 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 +423,39 @@ 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') From 89f743de0e4d0748984807cb5c076a219ff0a065 Mon Sep 17 00:00:00 2001 From: Kevin H Wilson Date: Sun, 22 Sep 2024 11:30:36 -0400 Subject: [PATCH 02/38] migrate utility function to codegen --- cpp/src/arrow/compute/kernels/aggregate_basic.cc | 13 ------------- .../arrow/compute/kernels/codegen_internal.cc | 16 ++++++++++++++++ cpp/src/arrow/compute/kernels/codegen_internal.h | 3 ++- 3 files changed, 18 insertions(+), 14 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/aggregate_basic.cc b/cpp/src/arrow/compute/kernels/aggregate_basic.cc index b5f746a18a14..f598caca8611 100644 --- a/cpp/src/arrow/compute/kernels/aggregate_basic.cc +++ b/cpp/src/arrow/compute/kernels/aggregate_basic.cc @@ -1030,19 +1030,6 @@ const FunctionDoc index_doc{"Find the index of the first occurrence of a given v } // namespace - -Result MaxPrecisionDecimalType(KernelContext*, const std::vector& types) { - std::shared_ptr out_type_; - if (auto decimal128_type = std::dynamic_pointer_cast(types.front().GetSharedPtr())) { - ARROW_ASSIGN_OR_RAISE(out_type_, Decimal128Type::Make(Decimal128Type::kMaxPrecision, decimal128_type->scale())); - } else if (auto decimal256_type = std::dynamic_pointer_cast(types.front().GetSharedPtr())) { - ARROW_ASSIGN_OR_RAISE(out_type_, Decimal256Type::Make(Decimal256Type::kMaxPrecision, decimal256_type->scale())); - } else { - return Status::TypeError("Bad call"); - } - return TypeHolder(out_type_); -} - void RegisterScalarAggregateBasic(FunctionRegistry* registry) { static auto default_scalar_aggregate_options = ScalarAggregateOptions::Defaults(); static auto default_count_options = CountOptions::Defaults(); diff --git a/cpp/src/arrow/compute/kernels/codegen_internal.cc b/cpp/src/arrow/compute/kernels/codegen_internal.cc index 5f3efedec352..a0b243c66834 100644 --- a/cpp/src/arrow/compute/kernels/codegen_internal.cc +++ b/cpp/src/arrow/compute/kernels/codegen_internal.cc @@ -77,6 +77,22 @@ Result ListValuesType(KernelContext* ctx, return value_type; } +Result MaxPrecisionDecimalType(KernelContext*, const std::vector& args) { + std::shared_ptr out_type_; + auto type_id = args[0].type.id; + if (type_id == Type::DECIMAL128 || type_id == Type::DECIMAL256) { + auto base_type_ = checked_cast(args[0].type); + if (type_id == Type::DECIMAL128) { + ARROW_ASSIGN_OR_RAISE(out_type_, Decimal128Type::Make(Decimal128Type::kMaxPrecision, base_type_->scale())); + } else (base_type_->id() == Type::DECIMAL256) { + ARROW_ASSIGN_OR_RAISE(out_type_, Decimal256Type::Make(Decimal256Type::kMaxPrecision, base_type_->scale())); + } + } else { + return Status::TypeError("A call to MaxPrecisionDecimalType was made with a non-DecimalType"); + } + return out_type_; +} + void EnsureDictionaryDecoded(std::vector* types) { EnsureDictionaryDecoded(types->data(), types->size()); } diff --git a/cpp/src/arrow/compute/kernels/codegen_internal.h b/cpp/src/arrow/compute/kernels/codegen_internal.h index 289ba25f0b78..d776a163e330 100644 --- a/cpp/src/arrow/compute/kernels/codegen_internal.h +++ b/cpp/src/arrow/compute/kernels/codegen_internal.h @@ -480,7 +480,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 // functions From 2564dcc33621123407d9dc7c1828bfd7d4392363 Mon Sep 17 00:00:00 2001 From: Kevin H Wilson Date: Sun, 22 Sep 2024 11:59:23 -0400 Subject: [PATCH 03/38] fix else condition --- cpp/src/arrow/compute/kernels/codegen_internal.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/codegen_internal.cc b/cpp/src/arrow/compute/kernels/codegen_internal.cc index a0b243c66834..44be3fe7e558 100644 --- a/cpp/src/arrow/compute/kernels/codegen_internal.cc +++ b/cpp/src/arrow/compute/kernels/codegen_internal.cc @@ -79,12 +79,12 @@ Result ListValuesType(KernelContext* ctx, Result MaxPrecisionDecimalType(KernelContext*, const std::vector& args) { std::shared_ptr out_type_; - auto type_id = args[0].type.id; + auto type_id = args[0].type->id(); if (type_id == Type::DECIMAL128 || type_id == Type::DECIMAL256) { auto base_type_ = checked_cast(args[0].type); if (type_id == Type::DECIMAL128) { ARROW_ASSIGN_OR_RAISE(out_type_, Decimal128Type::Make(Decimal128Type::kMaxPrecision, base_type_->scale())); - } else (base_type_->id() == Type::DECIMAL256) { + } else { ARROW_ASSIGN_OR_RAISE(out_type_, Decimal256Type::Make(Decimal256Type::kMaxPrecision, base_type_->scale())); } } else { From 63340d787b394088d54194da2833722a7e333f9a Mon Sep 17 00:00:00 2001 From: Kevin H Wilson Date: Sun, 22 Sep 2024 12:44:21 -0400 Subject: [PATCH 04/38] improve casts --- .../compute/kernels/aggregate_basic.inc.cc | 23 ++++++++----------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/aggregate_basic.inc.cc b/cpp/src/arrow/compute/kernels/aggregate_basic.inc.cc index d1b173f90398..192d0d2a6d56 100644 --- a/cpp/src/arrow/compute/kernels/aggregate_basic.inc.cc +++ b/cpp/src/arrow/compute/kernels/aggregate_basic.inc.cc @@ -94,10 +94,12 @@ struct SumImpl : public ScalarAggregator { Status Finalize(KernelContext*, Datum* out) override { std::shared_ptr out_type_; - if (auto decimal128_type = std::dynamic_pointer_cast(out_type)) { - ARROW_ASSIGN_OR_RAISE(out_type_, Decimal128Type::Make(Decimal128Type::kMaxPrecision, decimal128_type->scale())); - } else if (auto decimal256_type = std::dynamic_pointer_cast(out_type)) { - ARROW_ASSIGN_OR_RAISE(out_type_, Decimal256Type::Make(Decimal256Type::kMaxPrecision, decimal256_type->scale())); + if (out_type->id() == Type::DECIMAL128) { + auto cast_type = checked_cast(out_type); + ARROW_ASSIGN_OR_RAISE(out_type_, Decimal128Type::Make(Decimal128Type::kMaxPrecision, cast_type->scale())); + } else if (out_type->id() == Type::DECIMAL256) { + auto cast_type = checked_cast(out_type); + ARROW_ASSIGN_OR_RAISE(out_type_, Decimal256Type::Make(Decimal256Type::kMaxPrecision, cast_type->scale())); } else { out_type_ = out_type; } @@ -230,16 +232,11 @@ struct MeanImpl> template Status FinalizeImpl(Datum* out) { std::shared_ptr out_type_; - if (auto decimal128_type = std::dynamic_pointer_cast(this->out_type)) { - ARROW_ASSIGN_OR_RAISE(out_type_, Decimal128Type::Make(Decimal128Type::kMaxPrecision, decimal128_type->scale())); - } else if (auto decimal256_type = std::dynamic_pointer_cast(this->out_type)) { - ARROW_ASSIGN_OR_RAISE(out_type_, Decimal256Type::Make(Decimal256Type::kMaxPrecision, decimal256_type->scale())); + auto decimal_type = checked_cast(this->out_type); + if (decimal_type->id() == Type::DECIMAL128) { + ARROW_ASSIGN_OR_RAISE(out_type_, Decimal128Type::Make(Decimal128Type::kMaxPrecision, decimal_type->scale())); } else { - return Status::TypeError( - "The decimal specialization of MeanImpl was passed a type ", - this->out_type->ToString(), - " and not a decimal type" - ); + ARROW_ASSIGN_OR_RAISE(out_type_, Decimal256Type::Make(Decimal256Type::kMaxPrecision, decimal_type->scale())); } if ((!options.skip_nulls && this->nulls_observed) || From 8a6959b6dd9c25f4a2679556b5bdf06e9ec8a1be Mon Sep 17 00:00:00 2001 From: Kevin H Wilson Date: Sun, 22 Sep 2024 21:50:02 -0400 Subject: [PATCH 05/38] linting --- .../arrow/compute/kernels/aggregate_basic.cc | 36 ++++++++++--------- .../compute/kernels/aggregate_basic.inc.cc | 12 ++++--- .../arrow/compute/kernels/codegen_internal.cc | 12 ++++--- 3 files changed, 36 insertions(+), 24 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/aggregate_basic.cc b/cpp/src/arrow/compute/kernels/aggregate_basic.cc index f598caca8611..f9db9bf107cc 100644 --- a/cpp/src/arrow/compute/kernels/aggregate_basic.cc +++ b/cpp/src/arrow/compute/kernels/aggregate_basic.cc @@ -367,10 +367,14 @@ struct ProductImpl : public ScalarAggregator { Status Finalize(KernelContext*, Datum* out) override { std::shared_ptr out_type_; - if (auto decimal128_type = std::dynamic_pointer_cast(this->out_type)) { - ARROW_ASSIGN_OR_RAISE(out_type_, Decimal128Type::Make(Decimal128Type::kMaxPrecision, decimal128_type->scale())); - } else if (auto decimal256_type = std::dynamic_pointer_cast(this->out_type)) { - ARROW_ASSIGN_OR_RAISE(out_type_, Decimal256Type::Make(Decimal256Type::kMaxPrecision, decimal256_type->scale())); + if (out_type->id() == Type::DECIMAL128) { + auto cast_type = checked_cast(out_type); + ARROW_ASSIGN_OR_RAISE(out_type_, Decimal128Type::Make(Decimal128Type::kMaxPrecision, + cast_type->scale())); + } else if (out_type->id() == Type::DECIMAL256) { + auto cast_type = checked_cast(out_type); + ARROW_ASSIGN_OR_RAISE(out_type_, Decimal256Type::Make(Decimal256Type::kMaxPrecision, + cast_type->scale())); } else { out_type_ = out_type; } @@ -1058,10 +1062,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}, MaxPrecisionDecimalType), SumInit, func.get(), - SimdLevel::NONE); - AddAggKernel(KernelSignature::Make({Type::DECIMAL256}, MaxPrecisionDecimalType), 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()); @@ -1086,10 +1090,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::DECIMAL128}, MaxPrecisionDecimalType), MeanInit, func.get(), - SimdLevel::NONE); - AddAggKernel(KernelSignature::Make({Type::DECIMAL256}, MaxPrecisionDecimalType), MeanInit, func.get(), - SimdLevel::NONE); + AddAggKernel(KernelSignature::Make({Type::DECIMAL128}, MaxPrecisionDecimalType), + MeanInit, func.get(), SimdLevel::NONE); + AddAggKernel(KernelSignature::Make({Type::DECIMAL256}, MaxPrecisionDecimalType), + MeanInit, func.get(), SimdLevel::NONE); AddArrayScalarAggKernels(MeanInit, {null()}, float64(), func.get()); // Add the SIMD variants for mean #if defined(ARROW_HAVE_RUNTIME_AVX2) @@ -1170,10 +1174,10 @@ void RegisterScalarAggregateBasic(FunctionRegistry* registry) { AddArrayScalarAggKernels(ProductInit::Init, UnsignedIntTypes(), uint64(), func.get()); AddArrayScalarAggKernels(ProductInit::Init, FloatingPointTypes(), float64(), func.get()); - AddAggKernel(KernelSignature::Make({Type::DECIMAL128}, MaxPrecisionDecimalType), ProductInit::Init, - func.get(), SimdLevel::NONE); - AddAggKernel(KernelSignature::Make({Type::DECIMAL256}, 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), + ProductInit::Init, func.get(), SimdLevel::NONE); AddArrayScalarAggKernels(ProductInit::Init, {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 index 192d0d2a6d56..23ad6aea8093 100644 --- a/cpp/src/arrow/compute/kernels/aggregate_basic.inc.cc +++ b/cpp/src/arrow/compute/kernels/aggregate_basic.inc.cc @@ -96,10 +96,12 @@ struct SumImpl : public ScalarAggregator { std::shared_ptr out_type_; if (out_type->id() == Type::DECIMAL128) { auto cast_type = checked_cast(out_type); - ARROW_ASSIGN_OR_RAISE(out_type_, Decimal128Type::Make(Decimal128Type::kMaxPrecision, cast_type->scale())); + ARROW_ASSIGN_OR_RAISE(out_type_, Decimal128Type::Make(Decimal128Type::kMaxPrecision, + cast_type->scale())); } else if (out_type->id() == Type::DECIMAL256) { auto cast_type = checked_cast(out_type); - ARROW_ASSIGN_OR_RAISE(out_type_, Decimal256Type::Make(Decimal256Type::kMaxPrecision, cast_type->scale())); + ARROW_ASSIGN_OR_RAISE(out_type_, Decimal256Type::Make(Decimal256Type::kMaxPrecision, + cast_type->scale())); } else { out_type_ = out_type; } @@ -234,9 +236,11 @@ struct MeanImpl> std::shared_ptr out_type_; auto decimal_type = checked_cast(this->out_type); if (decimal_type->id() == Type::DECIMAL128) { - ARROW_ASSIGN_OR_RAISE(out_type_, Decimal128Type::Make(Decimal128Type::kMaxPrecision, decimal_type->scale())); + ARROW_ASSIGN_OR_RAISE(out_type_, Decimal128Type::Make(Decimal128Type::kMaxPrecision, + decimal_type->scale())); } else { - ARROW_ASSIGN_OR_RAISE(out_type_, Decimal256Type::Make(Decimal256Type::kMaxPrecision, decimal_type->scale())); + ARROW_ASSIGN_OR_RAISE(out_type_, Decimal256Type::Make(Decimal256Type::kMaxPrecision, + decimal_type->scale())); } if ((!options.skip_nulls && this->nulls_observed) || diff --git a/cpp/src/arrow/compute/kernels/codegen_internal.cc b/cpp/src/arrow/compute/kernels/codegen_internal.cc index 44be3fe7e558..32c6ccc9c98e 100644 --- a/cpp/src/arrow/compute/kernels/codegen_internal.cc +++ b/cpp/src/arrow/compute/kernels/codegen_internal.cc @@ -77,18 +77,22 @@ Result ListValuesType(KernelContext* ctx, return value_type; } -Result MaxPrecisionDecimalType(KernelContext*, const std::vector& args) { +Result MaxPrecisionDecimalType(KernelContext*, + const std::vector& args) { std::shared_ptr out_type_; auto type_id = args[0].type->id(); if (type_id == Type::DECIMAL128 || type_id == Type::DECIMAL256) { auto base_type_ = checked_cast(args[0].type); if (type_id == Type::DECIMAL128) { - ARROW_ASSIGN_OR_RAISE(out_type_, Decimal128Type::Make(Decimal128Type::kMaxPrecision, base_type_->scale())); + ARROW_ASSIGN_OR_RAISE(out_type_, Decimal128Type::Make(Decimal128Type::kMaxPrecision, + base_type_->scale())); } else { - ARROW_ASSIGN_OR_RAISE(out_type_, Decimal256Type::Make(Decimal256Type::kMaxPrecision, base_type_->scale())); + ARROW_ASSIGN_OR_RAISE(out_type_, Decimal256Type::Make(Decimal256Type::kMaxPrecision, + base_type_->scale())); } } else { - return Status::TypeError("A call to MaxPrecisionDecimalType was made with a non-DecimalType"); + return Status::TypeError( + "A call to MaxPrecisionDecimalType was made with a non-DecimalType"); } return out_type_; } From d98ad4aef5774fb336dbfc58244e790988a27c89 Mon Sep 17 00:00:00 2001 From: Kevin H Wilson Date: Mon, 23 Sep 2024 22:54:25 -0400 Subject: [PATCH 06/38] fix _pointer_ in casts --- cpp/src/arrow/compute/kernels/aggregate_basic.cc | 4 ++-- cpp/src/arrow/compute/kernels/aggregate_basic.inc.cc | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/aggregate_basic.cc b/cpp/src/arrow/compute/kernels/aggregate_basic.cc index f9db9bf107cc..1eed9ff5f89e 100644 --- a/cpp/src/arrow/compute/kernels/aggregate_basic.cc +++ b/cpp/src/arrow/compute/kernels/aggregate_basic.cc @@ -368,11 +368,11 @@ struct ProductImpl : public ScalarAggregator { Status Finalize(KernelContext*, Datum* out) override { std::shared_ptr out_type_; if (out_type->id() == Type::DECIMAL128) { - auto cast_type = checked_cast(out_type); + auto cast_type = checked_pointer_cast(this->out_type); ARROW_ASSIGN_OR_RAISE(out_type_, Decimal128Type::Make(Decimal128Type::kMaxPrecision, cast_type->scale())); } else if (out_type->id() == Type::DECIMAL256) { - auto cast_type = checked_cast(out_type); + auto cast_type = checked_pointer_cast(this->out_type); ARROW_ASSIGN_OR_RAISE(out_type_, Decimal256Type::Make(Decimal256Type::kMaxPrecision, cast_type->scale())); } else { diff --git a/cpp/src/arrow/compute/kernels/aggregate_basic.inc.cc b/cpp/src/arrow/compute/kernels/aggregate_basic.inc.cc index 23ad6aea8093..6db2c9e63521 100644 --- a/cpp/src/arrow/compute/kernels/aggregate_basic.inc.cc +++ b/cpp/src/arrow/compute/kernels/aggregate_basic.inc.cc @@ -95,11 +95,11 @@ struct SumImpl : public ScalarAggregator { Status Finalize(KernelContext*, Datum* out) override { std::shared_ptr out_type_; if (out_type->id() == Type::DECIMAL128) { - auto cast_type = checked_cast(out_type); + auto cast_type = checked_pointer_cast(this->out_type); ARROW_ASSIGN_OR_RAISE(out_type_, Decimal128Type::Make(Decimal128Type::kMaxPrecision, cast_type->scale())); } else if (out_type->id() == Type::DECIMAL256) { - auto cast_type = checked_cast(out_type); + auto cast_type = checked_pointer_cast(this->out_type); ARROW_ASSIGN_OR_RAISE(out_type_, Decimal256Type::Make(Decimal256Type::kMaxPrecision, cast_type->scale())); } else { @@ -234,7 +234,7 @@ struct MeanImpl> template Status FinalizeImpl(Datum* out) { std::shared_ptr out_type_; - auto decimal_type = checked_cast(this->out_type); + auto decimal_type = checked_pointer_cast(this->out_type); if (decimal_type->id() == Type::DECIMAL128) { ARROW_ASSIGN_OR_RAISE(out_type_, Decimal128Type::Make(Decimal128Type::kMaxPrecision, decimal_type->scale())); From 8c92b746b02f44996eb73ce50f684edd537e3eea Mon Sep 17 00:00:00 2001 From: Kevin H Wilson Date: Mon, 30 Sep 2024 10:04:31 -0400 Subject: [PATCH 07/38] extract common code --- .../arrow/compute/kernels/aggregate_basic.cc | 12 +---------- .../compute/kernels/aggregate_basic.inc.cc | 21 ++----------------- .../arrow/compute/kernels/codegen_internal.cc | 11 ++++++++++ .../arrow/compute/kernels/codegen_internal.h | 3 +++ 4 files changed, 17 insertions(+), 30 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/aggregate_basic.cc b/cpp/src/arrow/compute/kernels/aggregate_basic.cc index 1eed9ff5f89e..c65c6f291a32 100644 --- a/cpp/src/arrow/compute/kernels/aggregate_basic.cc +++ b/cpp/src/arrow/compute/kernels/aggregate_basic.cc @@ -367,17 +367,7 @@ struct ProductImpl : public ScalarAggregator { Status Finalize(KernelContext*, Datum* out) override { std::shared_ptr out_type_; - if (out_type->id() == Type::DECIMAL128) { - auto cast_type = checked_pointer_cast(this->out_type); - ARROW_ASSIGN_OR_RAISE(out_type_, Decimal128Type::Make(Decimal128Type::kMaxPrecision, - cast_type->scale())); - } else if (out_type->id() == Type::DECIMAL256) { - auto cast_type = checked_pointer_cast(this->out_type); - ARROW_ASSIGN_OR_RAISE(out_type_, Decimal256Type::Make(Decimal256Type::kMaxPrecision, - cast_type->scale())); - } else { - out_type_ = out_type; - } + ARROW_ASSIGN_OR_RAISE(out_type_, WidenDecimalToMaxPrecision(this->out_type)); if ((!options.skip_nulls && this->nulls_observed) || (this->count < options.min_count)) { diff --git a/cpp/src/arrow/compute/kernels/aggregate_basic.inc.cc b/cpp/src/arrow/compute/kernels/aggregate_basic.inc.cc index 6db2c9e63521..66d797d4dc17 100644 --- a/cpp/src/arrow/compute/kernels/aggregate_basic.inc.cc +++ b/cpp/src/arrow/compute/kernels/aggregate_basic.inc.cc @@ -94,17 +94,7 @@ struct SumImpl : public ScalarAggregator { Status Finalize(KernelContext*, Datum* out) override { std::shared_ptr out_type_; - if (out_type->id() == Type::DECIMAL128) { - auto cast_type = checked_pointer_cast(this->out_type); - ARROW_ASSIGN_OR_RAISE(out_type_, Decimal128Type::Make(Decimal128Type::kMaxPrecision, - cast_type->scale())); - } else if (out_type->id() == Type::DECIMAL256) { - auto cast_type = checked_pointer_cast(this->out_type); - ARROW_ASSIGN_OR_RAISE(out_type_, Decimal256Type::Make(Decimal256Type::kMaxPrecision, - cast_type->scale())); - } else { - out_type_ = out_type; - } + ARROW_ASSIGN_OR_RAISE(out_type_, WidenDecimalToMaxPrecision(this->out_type)); if ((!options.skip_nulls && this->nulls_observed) || (this->count < options.min_count)) { @@ -234,14 +224,7 @@ struct MeanImpl> template Status FinalizeImpl(Datum* out) { std::shared_ptr out_type_; - auto decimal_type = checked_pointer_cast(this->out_type); - if (decimal_type->id() == Type::DECIMAL128) { - ARROW_ASSIGN_OR_RAISE(out_type_, Decimal128Type::Make(Decimal128Type::kMaxPrecision, - decimal_type->scale())); - } else { - ARROW_ASSIGN_OR_RAISE(out_type_, Decimal256Type::Make(Decimal256Type::kMaxPrecision, - decimal_type->scale())); - } + ARROW_ASSIGN_OR_RAISE(out_type_, WidenDecimalToMaxPrecision(this->out_type)); if ((!options.skip_nulls && this->nulls_observed) || (this->count < options.min_count) || (this->count == 0)) { diff --git a/cpp/src/arrow/compute/kernels/codegen_internal.cc b/cpp/src/arrow/compute/kernels/codegen_internal.cc index 32c6ccc9c98e..eac924bf7c83 100644 --- a/cpp/src/arrow/compute/kernels/codegen_internal.cc +++ b/cpp/src/arrow/compute/kernels/codegen_internal.cc @@ -548,6 +548,17 @@ Status CastDecimalArgs(TypeHolder* begin, size_t count) { return Status::OK(); } +Result> WidenDecimalToMaxPrecision(std::shared_ptr type) { + if (type->id() == Type::DECIMAL128) { + auto cast_type = checked_pointer_cast(type); + return Decimal128Type::Make(Decimal128Type::kMaxPrecision, cast_type->scale()); + } else if (type->id() == Type::DECIMAL256) { + auto cast_type = checked_pointer_cast(type); + return Decimal256Type::Make(Decimal256Type::kMaxPrecision, cast_type->scale()); + } + return type; +} + 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 d776a163e330..33e9f1f8c2b8 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; @@ -1445,6 +1446,8 @@ Status CastBinaryDecimalArgs(DecimalPromotion promotion, std::vector ARROW_EXPORT Status CastDecimalArgs(TypeHolder* begin, size_t count); +Result> WidenDecimalToMaxPrecision(std::shared_ptr type); + ARROW_EXPORT bool HasDecimal(const std::vector& types); From cb6e0ede402f80f49fb1357bb4b5f80c2eb7548e Mon Sep 17 00:00:00 2001 From: Kevin H Wilson Date: Tue, 1 Oct 2024 19:38:29 -0400 Subject: [PATCH 08/38] continue simplifying code --- .../arrow/compute/kernels/codegen_internal.cc | 16 +--------------- 1 file changed, 1 insertion(+), 15 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/codegen_internal.cc b/cpp/src/arrow/compute/kernels/codegen_internal.cc index eac924bf7c83..53837bba8fce 100644 --- a/cpp/src/arrow/compute/kernels/codegen_internal.cc +++ b/cpp/src/arrow/compute/kernels/codegen_internal.cc @@ -79,21 +79,7 @@ Result ListValuesType(KernelContext* ctx, Result MaxPrecisionDecimalType(KernelContext*, const std::vector& args) { - std::shared_ptr out_type_; - auto type_id = args[0].type->id(); - if (type_id == Type::DECIMAL128 || type_id == Type::DECIMAL256) { - auto base_type_ = checked_cast(args[0].type); - if (type_id == Type::DECIMAL128) { - ARROW_ASSIGN_OR_RAISE(out_type_, Decimal128Type::Make(Decimal128Type::kMaxPrecision, - base_type_->scale())); - } else { - ARROW_ASSIGN_OR_RAISE(out_type_, Decimal256Type::Make(Decimal256Type::kMaxPrecision, - base_type_->scale())); - } - } else { - return Status::TypeError( - "A call to MaxPrecisionDecimalType was made with a non-DecimalType"); - } + ARROW_ASSIGN_OR_RAISE(auto out_type_, WidenDecimalToMaxPrecision(args[0].GetSharedPtr())); return out_type_; } From f59a4211d6912b04fb5f2c53c36821d6022c7731 Mon Sep 17 00:00:00 2001 From: Kevin H Wilson Date: Wed, 9 Oct 2024 21:58:17 -0400 Subject: [PATCH 09/38] add doc and export --- cpp/src/arrow/compute/kernels/codegen_internal.h | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/cpp/src/arrow/compute/kernels/codegen_internal.h b/cpp/src/arrow/compute/kernels/codegen_internal.h index 33e9f1f8c2b8..1f74ec61111e 100644 --- a/cpp/src/arrow/compute/kernels/codegen_internal.h +++ b/cpp/src/arrow/compute/kernels/codegen_internal.h @@ -1446,6 +1446,11 @@ 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. +/// +/// If it is _not_ a DecimalType, return type +ARROW_EXPORT Result> WidenDecimalToMaxPrecision(std::shared_ptr type); ARROW_EXPORT From a426640636e2ecfdc35431ec0144853680c29352 Mon Sep 17 00:00:00 2001 From: Kevin H Wilson Date: Sun, 13 Oct 2024 13:02:04 -0400 Subject: [PATCH 10/38] fix c++ tests --- .../arrow/compute/kernels/aggregate_basic.cc | 22 ++- .../compute/kernels/aggregate_basic.inc.cc | 6 +- .../arrow/compute/kernels/aggregate_test.cc | 167 +++++++++++------- .../arrow/compute/kernels/codegen_internal.cc | 24 ++- .../compute/kernels/test_util_internal.h | 12 +- 5 files changed, 157 insertions(+), 74 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/aggregate_basic.cc b/cpp/src/arrow/compute/kernels/aggregate_basic.cc index c65c6f291a32..44711184f3db 100644 --- a/cpp/src/arrow/compute/kernels/aggregate_basic.cc +++ b/cpp/src/arrow/compute/kernels/aggregate_basic.cc @@ -366,8 +366,10 @@ struct ProductImpl : public ScalarAggregator { } Status Finalize(KernelContext*, Datum* out) override { - std::shared_ptr out_type_; - ARROW_ASSIGN_OR_RAISE(out_type_, WidenDecimalToMaxPrecision(this->out_type)); + std::shared_ptr 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)) { @@ -1052,6 +1054,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::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), @@ -1080,6 +1086,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), @@ -1129,6 +1139,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 @@ -1164,6 +1176,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), @@ -1189,7 +1205,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 index 66d797d4dc17..b6781214e2df 100644 --- 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 out_type_; - ARROW_ASSIGN_OR_RAISE(out_type_, WidenDecimalToMaxPrecision(this->out_type)); + std::shared_ptr 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)) { diff --git a/cpp/src/arrow/compute/kernels/aggregate_test.cc b/cpp/src/arrow/compute/kernels/aggregate_test.cc index d821fc7e2c50..9a7ad482b279 100644 --- a/cpp/src/arrow/compute/kernels/aggregate_test.cc +++ b/cpp/src/arrow/compute/kernels/aggregate_test.cc @@ -495,51 +495,63 @@ TEST_F(TestSumKernelRoundOff, Basics) { } TEST(TestDecimalSumKernel, SimpleSum) { - for (const auto& ty : {decimal128(3, 2), decimal256(3, 2)}) { + std::vector> init_types = {decimal32(3, 2), decimal64(3, 2), decimal128(3, 2), decimal256(3, 2)}; + std::vector> out_types = {decimal32(9, 2), decimal64(18, 2), 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]; + 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 = {decimal32(3, 2), decimal64(3, 2), decimal128(3, 2), decimal256(3, 2)}; + std::vector> out_types = {decimal32(9, 2), decimal64(18, 2), 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 +591,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 +724,61 @@ 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 = {decimal32(3, 2), decimal64(3, 2), decimal128(3, 2), decimal256(3, 2)}; + std::vector> out_types = {decimal32(9, 2), decimal64(18, 2), 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)"); 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 = {decimal32(3, 2), decimal64(3, 2), decimal128(3, 2), decimal256(3, 2)}; + std::vector> out_types = {decimal32(9, 2), decimal64(18, 2), 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 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 +830,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)); @@ -1337,93 +1361,118 @@ TYPED_TEST(TestRandomNumericMeanKernel, RandomArrayMeanOverflow) { TEST(TestDecimalMeanKernel, SimpleMean) { ScalarAggregateOptions options(/*skip_nulls=*/true, /*min_count=*/0); - for (const auto& ty : {decimal128(3, 2), decimal256(3, 2)}) { + std::vector> init_types = {decimal32(3, 2), decimal64(3, 2), decimal128(3, 2), decimal256(3, 2)}; + std::vector> out_types = {decimal32(9, 2), decimal64(18, 2), 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]; + // 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)}) { + // TODO: Currently, casts are not implemented for decimal32/64 so we ignore that for now + // init_types = {decimal32(3, -2), decimal64(3, -2), decimal128(3, -2), decimal256(3, -2)}; + // out_types = {decimal32(9, -2), decimal64(18, -2), decimal128(38, -2), decimal256(76, -2)}; + + init_types = {decimal128(3, -2), decimal256(3, -2)}; + 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]; + // 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 = {decimal32(3, 2), decimal64(3, 2), decimal128(3, 2), decimal256(3, 2)}; + std::vector> out_types = {decimal32(9, 2), decimal64(18, 2), 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 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 +1530,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 53837bba8fce..7d37e5419a00 100644 --- a/cpp/src/arrow/compute/kernels/codegen_internal.cc +++ b/cpp/src/arrow/compute/kernels/codegen_internal.cc @@ -535,14 +535,22 @@ Status CastDecimalArgs(TypeHolder* begin, size_t count) { } Result> WidenDecimalToMaxPrecision(std::shared_ptr type) { - if (type->id() == Type::DECIMAL128) { - auto cast_type = checked_pointer_cast(type); - return Decimal128Type::Make(Decimal128Type::kMaxPrecision, cast_type->scale()); - } else if (type->id() == Type::DECIMAL256) { - auto cast_type = checked_pointer_cast(type); - return Decimal256Type::Make(Decimal256Type::kMaxPrecision, cast_type->scale()); - } - return type; + if (!is_decimal(type->id())) { + return Status::TypeError("Non-DecimalType passed to WidenDecimalToMaxPrecision: " + type->ToString()); + } + 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: + return Status::TypeError("An unknown DecimalType was passed to WidenDecimalToMaxPrecision: " + type->ToString()); + }; } bool HasDecimal(const std::vector& types) { diff --git a/cpp/src/arrow/compute/kernels/test_util_internal.h b/cpp/src/arrow/compute/kernels/test_util_internal.h index e3a27ab9addb..e51ce479ca5b 100644 --- a/cpp/src/arrow/compute/kernels/test_util_internal.h +++ b/cpp/src/arrow/compute/kernels/test_util_internal.h @@ -79,7 +79,11 @@ inline std::shared_ptr DecimalArrayFromJSON(const std::shared_ptr(*type); if (ty.scale() >= 0) return ArrayFromJSON(type, json); auto p = ty.precision() - ty.scale(); - auto adjusted_ty = ty.id() == Type::DECIMAL128 ? decimal128(p, 0) : decimal256(p, 0); + std::shared_ptr adjusted_ty = + ty.id() == Type::DECIMAL32 ? decimal32(p, 0) : + ty.id() == Type::DECIMAL64 ? decimal64(p, 0) : + ty.id() == Type::DECIMAL128 ? decimal128(p, 0) : + decimal256(p, 0); return Cast(ArrayFromJSON(adjusted_ty, json), type).ValueOrDie().make_array(); } @@ -88,7 +92,11 @@ inline std::shared_ptr DecimalScalarFromJSON( const auto& ty = checked_cast(*type); if (ty.scale() >= 0) return ScalarFromJSON(type, json); auto p = ty.precision() - ty.scale(); - auto adjusted_ty = ty.id() == Type::DECIMAL128 ? decimal128(p, 0) : decimal256(p, 0); + std::shared_ptr adjusted_ty = + ty.id() == Type::DECIMAL32 ? decimal32(p, 0) : + ty.id() == Type::DECIMAL64 ? decimal64(p, 0) : + ty.id() == Type::DECIMAL128 ? decimal128(p, 0) : + decimal256(p, 0); return Cast(ScalarFromJSON(adjusted_ty, json), type).ValueOrDie().scalar(); } From 094e7e4f898cdd162cfda371397865bb62ddf1fb Mon Sep 17 00:00:00 2001 From: Kevin H Wilson Date: Sun, 13 Oct 2024 13:13:57 -0400 Subject: [PATCH 11/38] linting and doc fixes --- .../arrow/compute/kernels/aggregate_basic.cc | 11 ++-- .../arrow/compute/kernels/aggregate_test.cc | 50 ++++++++++++------- .../arrow/compute/kernels/codegen_internal.cc | 13 +++-- .../arrow/compute/kernels/codegen_internal.h | 5 +- .../compute/kernels/test_util_internal.h | 20 ++++---- 5 files changed, 60 insertions(+), 39 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/aggregate_basic.cc b/cpp/src/arrow/compute/kernels/aggregate_basic.cc index 44711184f3db..5336a9d8b3d5 100644 --- a/cpp/src/arrow/compute/kernels/aggregate_basic.cc +++ b/cpp/src/arrow/compute/kernels/aggregate_basic.cc @@ -1054,10 +1054,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::DECIMAL32}, MaxPrecisionDecimalType), - SumInit, func.get(), SimdLevel::NONE); - AddAggKernel(KernelSignature::Make({Type::DECIMAL64}, MaxPrecisionDecimalType), - SumInit, func.get(), SimdLevel::NONE); + 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), @@ -1205,7 +1205,8 @@ 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), decimal32(1, 0), decimal64(1, 0), 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_test.cc b/cpp/src/arrow/compute/kernels/aggregate_test.cc index 9a7ad482b279..bbbf1d8b87c5 100644 --- a/cpp/src/arrow/compute/kernels/aggregate_test.cc +++ b/cpp/src/arrow/compute/kernels/aggregate_test.cc @@ -495,8 +495,10 @@ TEST_F(TestSumKernelRoundOff, Basics) { } TEST(TestDecimalSumKernel, SimpleSum) { - std::vector> init_types = {decimal32(3, 2), decimal64(3, 2), decimal128(3, 2), decimal256(3, 2)}; - std::vector> out_types = {decimal32(9, 2), decimal64(18, 2), decimal128(38, 2), decimal256(76, 2)}; + std::vector> init_types = { + decimal32(3, 2), decimal64(3, 2), decimal128(3, 2), decimal256(3, 2)}; + std::vector> out_types = { + decimal32(9, 2), decimal64(18, 2), decimal128(38, 2), decimal256(76, 2)}; for (size_t i = 0; i < init_types.size(); ++i) { auto& ty = init_types[i]; @@ -542,12 +544,14 @@ TEST(TestDecimalSumKernel, SimpleSum) { } TEST(TestDecimalSumKernel, ScalarAggregateOptions) { - std::vector> init_types = {decimal32(3, 2), decimal64(3, 2), decimal128(3, 2), decimal256(3, 2)}; - std::vector> out_types = {decimal32(9, 2), decimal64(18, 2), decimal128(38, 2), decimal256(76, 2)}; + std::vector> init_types = { + decimal32(3, 2), decimal64(3, 2), decimal128(3, 2), decimal256(3, 2)}; + std::vector> out_types = { + decimal32(9, 2), decimal64(18, 2), 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 ]; + auto& out_ty = out_types[i]; Datum null = ScalarFromJSON(out_ty, R"(null)"); Datum zero = ScalarFromJSON(out_ty, R"("0.00")"); @@ -724,12 +728,14 @@ TYPED_TEST(TestNumericProductKernel, ScalarAggregateOptions) { } TEST(TestDecimalProductKernel, SimpleProduct) { - std::vector> init_types = {decimal32(3, 2), decimal64(3, 2), decimal128(3, 2), decimal256(3, 2)}; - std::vector> out_types = {decimal32(9, 2), decimal64(18, 2), decimal128(38, 2), decimal256(76, 2)}; + std::vector> init_types = { + decimal32(3, 2), decimal64(3, 2), decimal128(3, 2), decimal256(3, 2)}; + std::vector> out_types = { + decimal32(9, 2), decimal64(18, 2), 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 ]; + auto& out_ty = out_types[i]; Datum null = ScalarFromJSON(out_ty, R"(null)"); @@ -755,7 +761,8 @@ TEST(TestDecimalProductKernel, SimpleProduct) { EXPECT_THAT(Product(ArrayFromJSON(ty, R"([null])"), options), ResultWith(ScalarFromJSON(out_ty, R"("1.00")"))); chunks = ChunkedArrayFromJSON(ty, {}); - EXPECT_THAT(Product(chunks, options), ResultWith(ScalarFromJSON(out_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"])"), @@ -769,12 +776,14 @@ TEST(TestDecimalProductKernel, SimpleProduct) { } TEST(TestDecimalProductKernel, ScalarAggregateOptions) { - std::vector> init_types = {decimal32(3, 2), decimal64(3, 2), decimal128(3, 2), decimal256(3, 2)}; - std::vector> out_types = {decimal32(9, 2), decimal64(18, 2), decimal128(38, 2), decimal256(76, 2)}; + std::vector> init_types = { + decimal32(3, 2), decimal64(3, 2), decimal128(3, 2), decimal256(3, 2)}; + std::vector> out_types = { + decimal32(9, 2), decimal64(18, 2), 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 ]; + auto& out_ty = out_types[i]; Datum null = ScalarFromJSON(out_ty, R"(null)"); Datum one = ScalarFromJSON(out_ty, R"("1.00")"); @@ -1361,8 +1370,10 @@ TYPED_TEST(TestRandomNumericMeanKernel, RandomArrayMeanOverflow) { TEST(TestDecimalMeanKernel, SimpleMean) { ScalarAggregateOptions options(/*skip_nulls=*/true, /*min_count=*/0); - std::vector> init_types = {decimal32(3, 2), decimal64(3, 2), decimal128(3, 2), decimal256(3, 2)}; - std::vector> out_types = {decimal32(9, 2), decimal64(18, 2), decimal128(38, 2), decimal256(76, 2)}; + std::vector> init_types = { + decimal32(3, 2), decimal64(3, 2), decimal128(3, 2), decimal256(3, 2)}; + std::vector> out_types = { + decimal32(9, 2), decimal64(18, 2), decimal128(38, 2), decimal256(76, 2)}; for (size_t i = 0; i < init_types.size(); ++i) { auto& ty = init_types[i]; @@ -1422,8 +1433,9 @@ TEST(TestDecimalMeanKernel, SimpleMean) { } // TODO: Currently, casts are not implemented for decimal32/64 so we ignore that for now - // init_types = {decimal32(3, -2), decimal64(3, -2), decimal128(3, -2), decimal256(3, -2)}; - // out_types = {decimal32(9, -2), decimal64(18, -2), decimal128(38, -2), decimal256(76, -2)}; + // init_types = {decimal32(3, -2), decimal64(3, -2), decimal128(3, -2), decimal256(3, + // -2)}; out_types = {decimal32(9, -2), decimal64(18, -2), decimal128(38, -2), + // decimal256(76, -2)}; init_types = {decimal128(3, -2), decimal256(3, -2)}; out_types = {decimal128(38, -2), decimal256(76, -2)}; @@ -1464,8 +1476,10 @@ TEST(TestDecimalMeanKernel, SimpleMean) { } TEST(TestDecimalMeanKernel, ScalarAggregateOptions) { - std::vector> init_types = {decimal32(3, 2), decimal64(3, 2), decimal128(3, 2), decimal256(3, 2)}; - std::vector> out_types = {decimal32(9, 2), decimal64(18, 2), decimal128(38, 2), decimal256(76, 2)}; + std::vector> init_types = { + decimal32(3, 2), decimal64(3, 2), decimal128(3, 2), decimal256(3, 2)}; + std::vector> out_types = { + decimal32(9, 2), decimal64(18, 2), decimal128(38, 2), decimal256(76, 2)}; for (size_t i = 0; i < init_types.size(); ++i) { auto& ty = init_types[i]; diff --git a/cpp/src/arrow/compute/kernels/codegen_internal.cc b/cpp/src/arrow/compute/kernels/codegen_internal.cc index 7d37e5419a00..2e01c84bf0f4 100644 --- a/cpp/src/arrow/compute/kernels/codegen_internal.cc +++ b/cpp/src/arrow/compute/kernels/codegen_internal.cc @@ -79,7 +79,8 @@ Result ListValuesType(KernelContext* ctx, Result MaxPrecisionDecimalType(KernelContext*, const std::vector& args) { - ARROW_ASSIGN_OR_RAISE(auto out_type_, WidenDecimalToMaxPrecision(args[0].GetSharedPtr())); + ARROW_ASSIGN_OR_RAISE(auto out_type_, + WidenDecimalToMaxPrecision(args[0].GetSharedPtr())); return out_type_; } @@ -534,9 +535,11 @@ Status CastDecimalArgs(TypeHolder* begin, size_t count) { return Status::OK(); } -Result> WidenDecimalToMaxPrecision(std::shared_ptr type) { +Result> WidenDecimalToMaxPrecision( + std::shared_ptr type) { if (!is_decimal(type->id())) { - return Status::TypeError("Non-DecimalType passed to WidenDecimalToMaxPrecision: " + type->ToString()); + return Status::TypeError("Non-DecimalType passed to WidenDecimalToMaxPrecision: " + + type->ToString()); } auto cast_type = checked_pointer_cast(type); switch (type->id()) { @@ -549,7 +552,9 @@ Result> WidenDecimalToMaxPrecision(std::shared_ptrscale()); default: - return Status::TypeError("An unknown DecimalType was passed to WidenDecimalToMaxPrecision: " + type->ToString()); + return Status::TypeError( + "An unknown DecimalType was passed to WidenDecimalToMaxPrecision: " + + type->ToString()); }; } diff --git a/cpp/src/arrow/compute/kernels/codegen_internal.h b/cpp/src/arrow/compute/kernels/codegen_internal.h index 1f74ec61111e..5fcd88455dcc 100644 --- a/cpp/src/arrow/compute/kernels/codegen_internal.h +++ b/cpp/src/arrow/compute/kernels/codegen_internal.h @@ -1449,9 +1449,10 @@ 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. /// -/// If it is _not_ a DecimalType, return type +/// If it is _not_ a DecimalType, return a Status::TypeError ARROW_EXPORT -Result> WidenDecimalToMaxPrecision(std::shared_ptr type); +Result> WidenDecimalToMaxPrecision( + std::shared_ptr type); ARROW_EXPORT bool HasDecimal(const std::vector& types); diff --git a/cpp/src/arrow/compute/kernels/test_util_internal.h b/cpp/src/arrow/compute/kernels/test_util_internal.h index e51ce479ca5b..3be569fa773d 100644 --- a/cpp/src/arrow/compute/kernels/test_util_internal.h +++ b/cpp/src/arrow/compute/kernels/test_util_internal.h @@ -79,11 +79,11 @@ inline std::shared_ptr DecimalArrayFromJSON(const std::shared_ptr(*type); if (ty.scale() >= 0) return ArrayFromJSON(type, json); auto p = ty.precision() - ty.scale(); - std::shared_ptr adjusted_ty = - ty.id() == Type::DECIMAL32 ? decimal32(p, 0) : - ty.id() == Type::DECIMAL64 ? decimal64(p, 0) : - ty.id() == Type::DECIMAL128 ? decimal128(p, 0) : - decimal256(p, 0); + std::shared_ptr adjusted_ty = ty.id() == Type::DECIMAL32 ? decimal32(p, 0) + : ty.id() == Type::DECIMAL64 ? decimal64(p, 0) + : ty.id() == Type::DECIMAL128 + ? decimal128(p, 0) + : decimal256(p, 0); return Cast(ArrayFromJSON(adjusted_ty, json), type).ValueOrDie().make_array(); } @@ -92,11 +92,11 @@ inline std::shared_ptr DecimalScalarFromJSON( const auto& ty = checked_cast(*type); if (ty.scale() >= 0) return ScalarFromJSON(type, json); auto p = ty.precision() - ty.scale(); - std::shared_ptr adjusted_ty = - ty.id() == Type::DECIMAL32 ? decimal32(p, 0) : - ty.id() == Type::DECIMAL64 ? decimal64(p, 0) : - ty.id() == Type::DECIMAL128 ? decimal128(p, 0) : - decimal256(p, 0); + std::shared_ptr adjusted_ty = ty.id() == Type::DECIMAL32 ? decimal32(p, 0) + : ty.id() == Type::DECIMAL64 ? decimal64(p, 0) + : ty.id() == Type::DECIMAL128 + ? decimal128(p, 0) + : decimal256(p, 0); return Cast(ScalarFromJSON(adjusted_ty, json), type).ValueOrDie().scalar(); } From 1b455834831ad68d7ddb40f65d255254d3fe6080 Mon Sep 17 00:00:00 2001 From: Kevin H Wilson Date: Sun, 13 Oct 2024 14:59:26 -0400 Subject: [PATCH 12/38] include hash aggregate sums --- .../arrow/compute/kernels/hash_aggregate.cc | 904 ++++++++++++++++++ 1 file changed, 904 insertions(+) diff --git a/cpp/src/arrow/compute/kernels/hash_aggregate.cc b/cpp/src/arrow/compute/kernels/hash_aggregate.cc index 0e3e359bde11..71774ff2c885 100644 --- a/cpp/src/arrow/compute/kernels/hash_aggregate.cc +++ b/cpp/src/arrow/compute/kernels/hash_aggregate.cc @@ -266,6 +266,910 @@ struct GroupedCountImpl : public GroupedAggregator { BufferBuilder counts_; }; +// ---------------------------------------------------------------------- +// Sum/Mean/Product implementation + +template ::Type> +struct GroupedReducingAggregator : public GroupedAggregator { + using AccType = AccumulateType; + using CType = typename TypeTraits::CType; + using InputCType = typename TypeTraits::CType; + + Status Init(ExecContext* ctx, const KernelInitArgs& args) override { + pool_ = ctx->memory_pool(); + options_ = checked_cast(*args.options); + reduced_ = TypedBufferBuilder(pool_); + counts_ = TypedBufferBuilder(pool_); + no_nulls_ = TypedBufferBuilder(pool_); + out_type_ = GetOutType(args.inputs[0].GetSharedPtr()); + return Status::OK(); + } + + Status Resize(int64_t new_num_groups) override { + auto added_groups = new_num_groups - num_groups_; + num_groups_ = new_num_groups; + RETURN_NOT_OK(reduced_.Append(added_groups, Impl::NullValue(*out_type_))); + RETURN_NOT_OK(counts_.Append(added_groups, 0)); + RETURN_NOT_OK(no_nulls_.Append(added_groups, true)); + return Status::OK(); + } + + Status Consume(const ExecSpan& batch) override { + CType* reduced = reduced_.mutable_data(); + int64_t* counts = counts_.mutable_data(); + uint8_t* no_nulls = no_nulls_.mutable_data(); + + VisitGroupedValues( + batch, + [&](uint32_t g, InputCType value) { + reduced[g] = Impl::Reduce(*out_type_, reduced[g], static_cast(value)); + counts[g]++; + }, + [&](uint32_t g) { bit_util::SetBitTo(no_nulls, g, false); }); + return Status::OK(); + } + + Status Merge(GroupedAggregator&& raw_other, + const ArrayData& group_id_mapping) override { + auto other = + checked_cast*>(&raw_other); + + CType* reduced = reduced_.mutable_data(); + int64_t* counts = counts_.mutable_data(); + uint8_t* no_nulls = no_nulls_.mutable_data(); + + const CType* other_reduced = other->reduced_.data(); + const int64_t* other_counts = other->counts_.data(); + const uint8_t* other_no_nulls = other->no_nulls_.data(); + + auto g = group_id_mapping.GetValues(1); + for (int64_t other_g = 0; other_g < group_id_mapping.length; ++other_g, ++g) { + counts[*g] += other_counts[other_g]; + reduced[*g] = Impl::Reduce(*out_type_, reduced[*g], other_reduced[other_g]); + bit_util::SetBitTo( + no_nulls, *g, + bit_util::GetBit(no_nulls, *g) && bit_util::GetBit(other_no_nulls, other_g)); + } + return Status::OK(); + } + + // Generate the values/nulls buffers + static Result> Finish(MemoryPool* pool, + const ScalarAggregateOptions& options, + const int64_t* counts, + TypedBufferBuilder* reduced, + int64_t num_groups, int64_t* null_count, + std::shared_ptr* null_bitmap) { + for (int64_t i = 0; i < num_groups; ++i) { + if (counts[i] >= options.min_count) continue; + + if ((*null_bitmap) == nullptr) { + ARROW_ASSIGN_OR_RAISE(*null_bitmap, AllocateBitmap(num_groups, pool)); + bit_util::SetBitsTo((*null_bitmap)->mutable_data(), 0, num_groups, true); + } + + (*null_count)++; + bit_util::SetBitTo((*null_bitmap)->mutable_data(), i, false); + } + return reduced->Finish(); + } + + Result Finalize() override { + std::shared_ptr null_bitmap = nullptr; + const int64_t* counts = counts_.data(); + int64_t null_count = 0; + + ARROW_ASSIGN_OR_RAISE(auto values, + Impl::Finish(pool_, options_, counts, &reduced_, num_groups_, + &null_count, &null_bitmap)); + + if (!options_.skip_nulls) { + null_count = kUnknownNullCount; + if (null_bitmap) { + arrow::internal::BitmapAnd(null_bitmap->data(), /*left_offset=*/0, + no_nulls_.data(), /*right_offset=*/0, num_groups_, + /*out_offset=*/0, null_bitmap->mutable_data()); + } else { + ARROW_ASSIGN_OR_RAISE(null_bitmap, no_nulls_.Finish()); + } + } + + return ArrayData::Make(out_type(), num_groups_, + {std::move(null_bitmap), std::move(values)}, null_count); + } + + std::shared_ptr out_type() const override { return out_type_; } + + template + static enable_if_t::value, std::shared_ptr> GetOutType( + const std::shared_ptr& in_type) { + return TypeTraits::type_singleton(); + } + + template + static enable_if_decimal> GetOutType( + const std::shared_ptr& in_type) { + return WidenDecimalToMaxPrecision(in_type).ValueOrDie(); + } + + int64_t num_groups_ = 0; + ScalarAggregateOptions options_; + TypedBufferBuilder reduced_; + TypedBufferBuilder counts_; + TypedBufferBuilder no_nulls_; + std::shared_ptr out_type_; + MemoryPool* pool_; +}; + +struct GroupedNullImpl : public GroupedAggregator { + Status Init(ExecContext* ctx, const KernelInitArgs& args) override { + pool_ = ctx->memory_pool(); + options_ = checked_cast(*args.options); + return Status::OK(); + } + + Status Resize(int64_t new_num_groups) override { + num_groups_ = new_num_groups; + return Status::OK(); + } + + Status Consume(const ExecSpan& batch) override { return Status::OK(); } + + Status Merge(GroupedAggregator&& raw_other, + const ArrayData& group_id_mapping) override { + return Status::OK(); + } + + Result Finalize() override { + if (options_.skip_nulls && options_.min_count == 0) { + ARROW_ASSIGN_OR_RAISE(std::shared_ptr data, + AllocateBuffer(num_groups_ * sizeof(int64_t), pool_)); + output_empty(data); + return ArrayData::Make(out_type(), num_groups_, {nullptr, std::move(data)}); + } else { + return MakeArrayOfNull(out_type(), num_groups_, pool_); + } + } + + virtual void output_empty(const std::shared_ptr& data) = 0; + + int64_t num_groups_; + ScalarAggregateOptions options_; + MemoryPool* pool_; +}; + +template