diff --git a/cpp/src/arrow/compute/expression_test.cc b/cpp/src/arrow/compute/expression_test.cc index 5e1f3c093ee2..6d64ccd388cc 100644 --- a/cpp/src/arrow/compute/expression_test.cc +++ b/cpp/src/arrow/compute/expression_test.cc @@ -938,6 +938,28 @@ TEST(Expression, BindWithImplicitCastsForCaseWhenOnDecimal) { /*bound_out=*/nullptr, *exciting_schema); } +TEST(Expression, BindWithImplicitCastsForIfElseOnDecimal) { + auto exciting_schema = schema({field("cond", boolean()), + field("dec128_3_2", decimal128(3, 2)), + field("dec128_4_3", decimal128(4, 3)), + field("dec256_3_2", decimal256(3, 2))}); + + ExpectBindsTo( + call("if_else", + {field_ref("cond"), field_ref("dec128_3_2"), field_ref("dec128_4_3")}), + call("if_else", + {field_ref("cond"), cast(field_ref("dec128_3_2"), decimal128(4, 3)), + field_ref("dec128_4_3")}), + /*bound_out=*/nullptr, *exciting_schema); + ExpectBindsTo( + call("if_else", + {field_ref("cond"), field_ref("dec128_3_2"), field_ref("dec256_3_2")}), + call("if_else", + {field_ref("cond"), cast(field_ref("dec128_3_2"), decimal256(3, 2)), + field_ref("dec256_3_2")}), + /*bound_out=*/nullptr, *exciting_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..6aa610bff8c1 100644 --- a/cpp/src/arrow/compute/kernels/scalar_if_else.cc +++ b/cpp/src/arrow/compute/kernels/scalar_if_else.cc @@ -1285,6 +1285,15 @@ struct IfElseFunction : ScalarFunction { return arrow::compute::detail::NoMatchingKernel(this, *types); } + + static std::shared_ptr ValueTypesMatchConstraint() { + static auto constraint = + MatchConstraint::Make([](const std::vector& types) -> bool { + DCHECK_EQ(types.size(), 3); + return types[1] == types[2]; + }); + return constraint; + } }; void AddNullIfElseKernel(const std::shared_ptr& scalar_function) { @@ -1314,6 +1323,8 @@ void AddPrimitiveIfElseKernels(const std::shared_ptr& scalar_fun } else { sig = KernelSignature::Make({boolean(), type, type}, type); } + sig = KernelSignature::Make(sig->in_types(), sig->out_type(), sig->is_varargs(), + IfElseFunction::ValueTypesMatchConstraint()); ScalarKernel kernel(std::move(sig), exec); kernel.null_handling = NullHandling::COMPUTED_PREALLOCATE; kernel.mem_allocation = MemAllocation::PREALLOCATE; @@ -1331,7 +1342,10 @@ void AddBinaryIfElseKernels(const std::shared_ptr& scalar_functi /*AllocateMem=*/std::true_type>( *type); // cond array needs to be boolean always - ScalarKernel kernel({boolean(), type, type}, type, exec); + ScalarKernel kernel( + KernelSignature::Make({boolean(), type, type}, type, /*is_varargs=*/false, + IfElseFunction::ValueTypesMatchConstraint()), + exec); kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE; kernel.mem_allocation = MemAllocation::NO_PREALLOCATE; kernel.can_write_into_slices = false; @@ -1343,8 +1357,11 @@ void AddBinaryIfElseKernels(const std::shared_ptr& scalar_functi template void AddFixedWidthIfElseKernel(const std::shared_ptr& scalar_function) { auto type_id = T::type_id; - ScalarKernel kernel({boolean(), InputType(type_id), InputType(type_id)}, LastType, - ResolveIfElseExec::Exec); + ScalarKernel kernel( + KernelSignature::Make({boolean(), InputType(type_id), InputType(type_id)}, LastType, + /*is_varargs=*/false, + IfElseFunction::ValueTypesMatchConstraint()), + ResolveIfElseExec::Exec); kernel.null_handling = NullHandling::COMPUTED_PREALLOCATE; kernel.mem_allocation = MemAllocation::PREALLOCATE; kernel.can_write_into_slices = true; @@ -1357,8 +1374,11 @@ void AddNestedIfElseKernels(const std::shared_ptr& scalar_functi {Type::LIST, Type::LARGE_LIST, Type::LIST_VIEW, Type::LARGE_LIST_VIEW, Type::FIXED_SIZE_LIST, Type::MAP, Type::STRUCT, Type::DENSE_UNION, Type::SPARSE_UNION, Type::DICTIONARY}) { - ScalarKernel kernel({boolean(), InputType(type_id), InputType(type_id)}, LastType, - NestedIfElseExec::Exec); + ScalarKernel kernel( + KernelSignature::Make({boolean(), InputType(type_id), InputType(type_id)}, + LastType, /*is_varargs=*/false, + IfElseFunction::ValueTypesMatchConstraint()), + NestedIfElseExec::Exec); kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE; kernel.mem_allocation = MemAllocation::NO_PREALLOCATE; kernel.can_write_into_slices = false; 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..ebadcc4e3184 100644 --- a/cpp/src/arrow/compute/kernels/scalar_if_else_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_if_else_test.cc @@ -500,6 +500,8 @@ TEST_F(TestIfElseKernel, IfElseDispatchBest) { {boolean(), float32(), float32()}); CheckDispatchBest(name, {boolean(), float64(), int32()}, {boolean(), float64(), float64()}); + CheckDispatchBest(name, {boolean(), decimal128(3, 2), decimal128(4, 3)}, + {boolean(), decimal128(4, 3), decimal128(4, 3)}); CheckDispatchBest(name, {null(), uint8(), int8()}, {boolean(), int16(), int16()}); @@ -517,6 +519,21 @@ TEST_F(TestIfElseKernel, IfElseDispatchBest) { {boolean(), date32(), date32()}); } +TEST_F(TestIfElseKernel, IfElseDispatchExact) { + CheckDispatchExact("if_else", {boolean(), decimal128(3, 2), decimal128(3, 2)}); + CheckDispatchExactFails("if_else", + {boolean(), decimal128(3, 2), decimal128(4, 3)}); + CheckDispatchExactFails( + "if_else", + {boolean(), timestamp(TimeUnit::SECOND), timestamp(TimeUnit::SECOND, "UTC")}); + CheckDispatchExactFails("if_else", + {boolean(), fixed_size_binary(3), fixed_size_binary(4)}); + CheckDispatchExactFails("if_else", {boolean(), list(int32()), list(int64())}); + CheckDispatchExactFails("if_else", + {boolean(), dictionary(int8(), int32()), + dictionary(int8(), int64())}); +} + template class TestIfElseBaseBinary : public ::testing::Test {};