From 23e9ff63d1e1dd20f3334dab7d302edaf0feebd7 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Fri, 5 Jun 2026 11:52:06 -0700 Subject: [PATCH] GH-50869: [C++][Compute] Tighten coalesce exact dispatch for decimal varargs Signed-off-by: Rossi Sun --- cpp/src/arrow/compute/expression_test.cc | 56 +++++++++++++++++++ .../arrow/compute/kernels/scalar_if_else.cc | 25 +++++++-- .../compute/kernels/scalar_if_else_test.cc | 27 +++++++++ 3 files changed, 104 insertions(+), 4 deletions(-) diff --git a/cpp/src/arrow/compute/expression_test.cc b/cpp/src/arrow/compute/expression_test.cc index 5e1f3c093ee2..a7e279c164b3 100644 --- a/cpp/src/arrow/compute/expression_test.cc +++ b/cpp/src/arrow/compute/expression_test.cc @@ -938,6 +938,62 @@ TEST(Expression, BindWithImplicitCastsForCaseWhenOnDecimal) { /*bound_out=*/nullptr, *exciting_schema); } +TEST(Expression, BindWithImplicitCastsForCoalesceOnDecimal) { + auto exciting_schema = schema( + {field("dec128_3_2", decimal128(3, 2)), field("dec128_4_1", decimal128(4, 1)), + field("dec128_4_2", decimal128(4, 2)), field("dec128_4_3", decimal128(4, 3)), + field("dec256_3_2", decimal256(3, 2))}); + + ExpectBindsTo(call("coalesce", {field_ref("dec128_3_2"), field_ref("dec128_4_2")}), + call("coalesce", {cast(field_ref("dec128_3_2"), decimal128(4, 2)), + field_ref("dec128_4_2")}), + /*bound_out=*/nullptr, *exciting_schema); + ExpectBindsTo(call("coalesce", {field_ref("dec128_4_2"), field_ref("dec128_3_2")}), + call("coalesce", {field_ref("dec128_4_2"), + cast(field_ref("dec128_3_2"), decimal128(4, 2))}), + /*bound_out=*/nullptr, *exciting_schema); + ExpectBindsTo(call("coalesce", {field_ref("dec128_4_1"), field_ref("dec128_3_2")}), + call("coalesce", {cast(field_ref("dec128_4_1"), decimal128(5, 2)), + cast(field_ref("dec128_3_2"), decimal128(5, 2))}), + /*bound_out=*/nullptr, *exciting_schema); + ExpectBindsTo(call("coalesce", {field_ref("dec128_3_2"), field_ref("dec128_4_1")}), + call("coalesce", {cast(field_ref("dec128_3_2"), decimal128(5, 2)), + cast(field_ref("dec128_4_1"), decimal128(5, 2))}), + /*bound_out=*/nullptr, *exciting_schema); + ExpectBindsTo(call("coalesce", {field_ref("dec128_3_2"), field_ref("dec128_4_3")}), + call("coalesce", {cast(field_ref("dec128_3_2"), decimal128(4, 3)), + field_ref("dec128_4_3")}), + /*bound_out=*/nullptr, *exciting_schema); + ExpectBindsTo(call("coalesce", {field_ref("dec128_4_3"), field_ref("dec128_3_2")}), + call("coalesce", {field_ref("dec128_4_3"), + cast(field_ref("dec128_3_2"), decimal128(4, 3))}), + /*bound_out=*/nullptr, *exciting_schema); + ExpectBindsTo(call("coalesce", {field_ref("dec128_3_2"), field_ref("dec256_3_2")}), + call("coalesce", {cast(field_ref("dec128_3_2"), decimal256(3, 2)), + field_ref("dec256_3_2")}), + /*bound_out=*/nullptr, *exciting_schema); + ExpectBindsTo(call("coalesce", {field_ref("dec256_3_2"), field_ref("dec128_3_2")}), + call("coalesce", {field_ref("dec256_3_2"), + cast(field_ref("dec128_3_2"), decimal256(3, 2))}), + /*bound_out=*/nullptr, *exciting_schema); +} + +TEST(Expression, ExecuteCoalesceOnMixedDecimalTypes) { + ASSERT_OK_AND_ASSIGN( + auto input, + StructArray::Make({ArrayFromJSON(decimal128(3, 2), R"(["1.23", null])"), + ArrayFromJSON(decimal128(4, 3), R"([null, "2.345"])")}, + {"left", "right"})); + Schema input_schema(input->type()->fields()); + auto expr = call("coalesce", {field_ref("left"), field_ref("right")}); + + ASSERT_OK_AND_ASSIGN(expr, expr.Bind(input_schema)); + ASSERT_OK_AND_ASSIGN(auto actual, + ExecuteScalarExpression(expr, input_schema, Datum(input))); + + AssertDatumsEqual(actual, ArrayFromJSON(decimal128(4, 3), R"(["1.230", "2.345"])")); +} + 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 1510dd9fc83a..0193fd4f5d80 100644 --- a/cpp/src/arrow/compute/kernels/scalar_if_else.cc +++ b/cpp/src/arrow/compute/kernels/scalar_if_else.cc @@ -2035,6 +2035,20 @@ struct CoalesceFunction : ScalarFunction { if (auto kernel = DispatchExactImpl(this, *types)) 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(), 1); + DCHECK(std::all_of(types.begin(), types.end(), [](const TypeHolder& type) { + return is_decimal(type.id()); + })); + return std::all_of( + types.begin() + 1, types.end(), + [&types](const TypeHolder& type) { return type == types[0]; }); + }); + return constraint; + } }; // Helper: copy from a source value into all null slots of the output @@ -2793,9 +2807,10 @@ void AddNestedCaseWhenKernels(const std::shared_ptr& scalar_fu } void AddCoalesceKernel(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({InputType(get_id.id)}, FirstType, - /*is_varargs=*/true), + /*is_varargs=*/true, std::move(constraint)), exec); kernel.null_handling = NullHandling::COMPUTED_PREALLOCATE; kernel.mem_allocation = MemAllocation::PREALLOCATE; @@ -2938,8 +2953,10 @@ void RegisterScalarIfElse(FunctionRegistry* registry) { AddPrimitiveCoalesceKernels(func, {boolean(), null(), float16()}); AddCoalesceKernel(func, Type::FIXED_SIZE_BINARY, CoalesceFunctor::Exec); - AddCoalesceKernel(func, Type::DECIMAL128, CoalesceFunctor::Exec); - AddCoalesceKernel(func, Type::DECIMAL256, CoalesceFunctor::Exec); + AddCoalesceKernel(func, Type::DECIMAL128, CoalesceFunctor::Exec, + CoalesceFunction::DecimalMatchConstraint()); + AddCoalesceKernel(func, Type::DECIMAL256, CoalesceFunctor::Exec, + CoalesceFunction::DecimalMatchConstraint()); for (const auto& ty : BaseBinaryTypes()) { AddCoalesceKernel(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 a1ef82383e29..c9f0e7d7dc63 100644 --- a/cpp/src/arrow/compute/kernels/scalar_if_else_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_if_else_test.cc @@ -3693,8 +3693,22 @@ TEST(TestCoalesce, DispatchBest) { CheckDispatchBest("coalesce", {int32(), decimal128(3, 2)}, {decimal128(12, 2), decimal128(12, 2)}); CheckDispatchBest("coalesce", {float32(), decimal128(3, 2)}, {float64(), float64()}); + CheckDispatchBest("coalesce", {decimal128(3, 2), decimal128(4, 2)}, + {decimal128(4, 2), decimal128(4, 2)}); + CheckDispatchBest("coalesce", {decimal128(4, 2), decimal128(3, 2)}, + {decimal128(4, 2), decimal128(4, 2)}); + CheckDispatchBest("coalesce", {decimal128(4, 1), decimal128(3, 2)}, + {decimal128(5, 2), decimal128(5, 2)}); + CheckDispatchBest("coalesce", {decimal128(3, 2), decimal128(4, 1)}, + {decimal128(5, 2), decimal128(5, 2)}); + CheckDispatchBest("coalesce", {decimal128(3, 2), decimal128(4, 3)}, + {decimal128(4, 3), decimal128(4, 3)}); + CheckDispatchBest("coalesce", {decimal128(4, 3), decimal128(3, 2)}, + {decimal128(4, 3), decimal128(4, 3)}); CheckDispatchBest("coalesce", {decimal128(3, 2), decimal256(3, 2)}, {decimal256(3, 2), decimal256(3, 2)}); + CheckDispatchBest("coalesce", {decimal256(3, 2), decimal128(3, 2)}, + {decimal256(3, 2), decimal256(3, 2)}); CheckDispatchBest("coalesce", {timestamp(TimeUnit::SECOND), date32()}, {timestamp(TimeUnit::SECOND), timestamp(TimeUnit::SECOND)}); CheckDispatchBest("coalesce", {timestamp(TimeUnit::SECOND), timestamp(TimeUnit::MILLI)}, @@ -3710,6 +3724,19 @@ TEST(TestCoalesce, DispatchBest) { {large_binary(), large_binary()}); } +TEST(TestCoalesce, DispatchExact) { + CheckDispatchExact("coalesce", {decimal128(3, 2), decimal128(3, 2)}); + CheckDispatchExact("coalesce", {decimal256(3, 2), decimal256(3, 2)}); + CheckDispatchExactFails("coalesce", {decimal128(3, 2), decimal128(4, 2)}); + CheckDispatchExactFails("coalesce", {decimal128(4, 2), decimal128(3, 2)}); + CheckDispatchExactFails("coalesce", {decimal128(4, 1), decimal128(3, 2)}); + CheckDispatchExactFails("coalesce", {decimal128(3, 2), decimal128(4, 1)}); + CheckDispatchExactFails("coalesce", {decimal128(3, 2), decimal128(4, 3)}); + CheckDispatchExactFails("coalesce", {decimal128(4, 3), decimal128(3, 2)}); + CheckDispatchExactFails("coalesce", {decimal128(3, 2), decimal256(3, 2)}); + CheckDispatchExactFails("coalesce", {decimal256(3, 2), decimal128(3, 2)}); +} + template class TestChooseNumeric : public ::testing::Test {}; template