diff --git a/cpp/src/arrow/compute/expression_test.cc b/cpp/src/arrow/compute/expression_test.cc index 5e1f3c093ee2..2699bfe33516 100644 --- a/cpp/src/arrow/compute/expression_test.cc +++ b/cpp/src/arrow/compute/expression_test.cc @@ -938,6 +938,20 @@ TEST(Expression, BindWithImplicitCastsForCaseWhenOnDecimal) { /*bound_out=*/nullptr, *exciting_schema); } +TEST(Expression, BindWithImplicitCastsForChooseOnDecimal) { + auto schema = arrow::schema({field("indices", int64()), + field("dec128_3_2", decimal128(3, 2)), + field("dec128_4_3", decimal128(4, 3))}); + + ExpectBindsTo( + call("choose", + {field_ref("indices"), field_ref("dec128_3_2"), field_ref("dec128_4_3")}), + call("choose", + {field_ref("indices"), cast(field_ref("dec128_3_2"), decimal128(4, 3)), + field_ref("dec128_4_3")}), + /*bound_out=*/nullptr, *schema); +} + TEST(Expression, BindNestedCall) { auto expr = add(field_ref("a"), call("subtract", {call("multiply", {field_ref("b"), field_ref("c")}), diff --git a/cpp/src/arrow/compute/kernels/scalar_if_else.cc b/cpp/src/arrow/compute/kernels/scalar_if_else.cc index 0c5cfe1c9037..e2eeba536e1c 100644 --- a/cpp/src/arrow/compute/kernels/scalar_if_else.cc +++ b/cpp/src/arrow/compute/kernels/scalar_if_else.cc @@ -2752,11 +2752,28 @@ struct ChooseFunction : ScalarFunction { *it = type; } } + if (HasDecimal(*types)) { + RETURN_NOT_OK(CastDecimalArgs(types->data() + 1, types->size() - 1)); + } if (auto kernel = DispatchExactImpl(this, {types->front(), types->back()})) { return kernel; } return arrow::compute::detail::NoMatchingKernel(this, *types); } + + static std::shared_ptr DecimalMatchConstraint() { + static auto constraint = + MatchConstraint::Make([](const std::vector& types) -> bool { + DCHECK_GE(types.size(), 2); + DCHECK(std::all_of(types.begin() + 1, types.end(), [](const TypeHolder& type) { + return is_decimal(type.id()); + })); + return std::all_of( + types.begin() + 2, types.end(), + [&types](const TypeHolder& type) { return type == types[1]; }); + }); + return constraint; + } }; void AddCaseWhenKernel(const std::shared_ptr& scalar_function, @@ -2852,9 +2869,10 @@ void AddNestedCoalesceKernels(const std::shared_ptr& scalar_func } void AddChooseKernel(const std::shared_ptr& scalar_function, - detail::GetTypeId get_id, ArrayKernelExec exec) { + detail::GetTypeId get_id, ArrayKernelExec exec, + std::shared_ptr constraint = nullptr) { ScalarKernel kernel(KernelSignature::Make({Type::INT64, InputType(get_id.id)}, LastType, - /*is_varargs=*/true), + /*is_varargs=*/true, std::move(constraint)), exec); kernel.null_handling = NullHandling::COMPUTED_PREALLOCATE; kernel.mem_allocation = MemAllocation::PREALLOCATE; @@ -2975,8 +2993,10 @@ void RegisterScalarIfElse(FunctionRegistry* registry) { AddPrimitiveChooseKernels(func, {boolean(), null(), float16()}); AddChooseKernel(func, Type::FIXED_SIZE_BINARY, ChooseFunctor::Exec); - AddChooseKernel(func, Type::DECIMAL128, ChooseFunctor::Exec); - AddChooseKernel(func, Type::DECIMAL256, ChooseFunctor::Exec); + AddChooseKernel(func, Type::DECIMAL128, ChooseFunctor::Exec, + ChooseFunction::DecimalMatchConstraint()); + AddChooseKernel(func, Type::DECIMAL256, ChooseFunctor::Exec, + ChooseFunction::DecimalMatchConstraint()); for (const auto& ty : BaseBinaryTypes()) { AddChooseKernel(func, ty, GenerateTypeAgnosticVarBinaryBase(ty)); } diff --git a/cpp/src/arrow/compute/kernels/scalar_if_else_test.cc b/cpp/src/arrow/compute/kernels/scalar_if_else_test.cc index 6fdcff8d9703..249947210c5e 100644 --- a/cpp/src/arrow/compute/kernels/scalar_if_else_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_if_else_test.cc @@ -3899,6 +3899,8 @@ TEST(TestChooseKernel, DispatchBest) { // Other arguments promoted separately from index EXPECT_EQ((std::vector{int64(), int32(), int32()}), Check({int8(), int32(), uint8()})); + EXPECT_EQ((std::vector{int64(), decimal128(4, 3), decimal128(4, 3)}), + Check({int8(), decimal128(3, 2), decimal128(4, 3)})); } TEST(TestChooseKernel, Errors) {