From dee38e4679d1a917bef39d9a28c8abe6c055bc13 Mon Sep 17 00:00:00 2001 From: Rossi Sun Date: Wed, 4 Nov 2026 10:30:12 -0800 Subject: [PATCH] GH-55404: [C++] Fix coalesce exact dispatch for mixed dictionary types Signed-off-by: Rossi Sun --- cpp/src/arrow/compute/expression_test.cc | 52 +++++++++++++++++++ .../arrow/compute/kernels/scalar_if_else.cc | 3 +- .../compute/kernels/scalar_if_else_test.cc | 28 ++++++++++ 3 files changed, 82 insertions(+), 1 deletion(-) diff --git a/cpp/src/arrow/compute/expression_test.cc b/cpp/src/arrow/compute/expression_test.cc index 10c1b348e62d..e123411a3c7a 100644 --- a/cpp/src/arrow/compute/expression_test.cc +++ b/cpp/src/arrow/compute/expression_test.cc @@ -1063,6 +1063,58 @@ TEST(Expression, ExecuteCoalesceOnMixedDecimalTypes) { AssertDatumsEqual(actual, ArrayFromJSON(decimal128(4, 3), R"(["1.230", "2.345"])")); } +TEST(Expression, CoalesceDictionaryTypes) { + auto check = [](const ArrayVector& arguments, + const std::shared_ptr& output_type, + const std::shared_ptr& expected) { + FieldVector fields; + std::vector args, expected_args; + for (size_t i = 0; i < arguments.size(); ++i) { + auto name = "arg" + std::to_string(i); + fields.push_back(field(name, arguments[i]->type())); + args.push_back(field_ref(name)); + expected_args.push_back(arguments[i]->type()->Equals(output_type) + ? args.back() + : cast(args.back(), output_type)); + } + Schema input_schema(fields); + Expression bound; + ExpectBindsTo(call("coalesce", args), call("coalesce", expected_args), &bound, + input_schema); + ASSERT_NE(bound.type(), nullptr); + EXPECT_TRUE(bound.type()->Equals(output_type)); + ASSERT_OK_AND_ASSIGN(auto input, StructArray::Make(arguments, fields)); + ASSERT_OK_AND_ASSIGN(auto actual, + ExecuteScalarExpression(bound, input_schema, Datum(input))); + ASSERT_OK(actual.make_array()->ValidateFull()); + // Also check logical values when the result stays dictionary-encoded. + ASSERT_OK_AND_ASSIGN(actual, Cast(actual, expected->type())); + AssertDatumsEqual(actual, expected); + }; + + auto type = dictionary(int8(), int32()); + auto left = DictArrayFromJSON(type, "[0, null, 1, null, null]", "[10, 30]"); + auto right = DictArrayFromJSON(type, "[null, 0, 1, null, null]", "[20, 40]"); + auto third = DictArrayFromJSON(type, "[null, null, null, 0, null]", "[50]"); + auto expected = ArrayFromJSON(int32(), "[10, 20, 30, null, null]"); + check({left}, type, ArrayFromJSON(int32(), "[10, null, 30, null, null]")); + check({left, right}, type, expected); + check({left, right, third}, type, ArrayFromJSON(int32(), "[10, 20, 30, 50, null]")); + + for (const auto& other : {dictionary(int16(), int32()), dictionary(int8(), int64()), + dictionary(int8(), int32(), /*ordered=*/true)}) { + auto output_type = checked_cast(*other).value_type(); + auto mixed_right = DictArrayFromJSON(other, "[null, 0, 1, null, null]", "[20, 40]"); + check({left, mixed_right}, output_type, + ArrayFromJSON(output_type, "[10, 20, 30, null, null]")); + check({mixed_right, left}, output_type, + ArrayFromJSON(output_type, "[10, 20, 40, null, null]")); + auto mixed_third = DictArrayFromJSON(other, "[null, null, null, 0, null]", "[50]"); + check({left, right, mixed_third}, output_type, + ArrayFromJSON(output_type, "[10, 20, 30, 50, null]")); + } +} + 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 1d8e6c1f62ef..f3d6fbca68cb 100644 --- a/cpp/src/arrow/compute/kernels/scalar_if_else.cc +++ b/cpp/src/arrow/compute/kernels/scalar_if_else.cc @@ -2816,7 +2816,8 @@ void AddNestedCoalesceKernels(const std::shared_ptr& scalar_func AddNestedCoalesceKernel(scalar_function); AddNestedCoalesceKernel(scalar_function); AddNestedCoalesceKernel(scalar_function); - AddNestedCoalesceKernel(scalar_function); + AddCoalesceKernel(scalar_function, Type::DICTIONARY, + CoalesceFunctor::Exec, AllTypesAreIdentical()); } void AddChooseKernel(const std::shared_ptr& scalar_function, 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 e05a1f081beb..ac3a8bd4706d 100644 --- a/cpp/src/arrow/compute/kernels/scalar_if_else_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_if_else_test.cc @@ -3743,6 +3743,34 @@ TEST(TestCoalesce, DispatchExact) { CheckDispatchExactFails("coalesce", {decimal128(3, 2), decimal256(4, 1)}); } +TEST(TestCoalesce, DictionaryDispatchExact) { + auto type = dictionary(int8(), int32()); + CheckDispatchExact("coalesce", {type}); + CheckDispatchExact("coalesce", {type, dictionary(int8(), int32())}); + CheckDispatchExact("coalesce", {type, type, type}); + for (const auto& other : {dictionary(int16(), int32()), dictionary(int8(), int64()), + dictionary(int8(), int32(), /*ordered=*/true)}) { + CheckDispatchExactFails("coalesce", {type, other}); + CheckDispatchExactFails("coalesce", {other, type}); + CheckDispatchExactFails("coalesce", {type, type, other}); + } +} + +TEST(TestCoalesce, DictionaryDispatchBest) { + auto type = dictionary(int8(), int32()); + CheckDispatchBest("coalesce", {type}, {type}); + CheckDispatchBest("coalesce", {type, dictionary(int8(), int32())}, {type, type}); + CheckDispatchBest("coalesce", {type, type, type}, {type, type, type}); + for (const auto& other : {dictionary(int16(), int32()), dictionary(int8(), int64()), + dictionary(int8(), int32(), /*ordered=*/true)}) { + auto value_type = checked_cast(*other).value_type(); + CheckDispatchBest("coalesce", {type, other}, {value_type, value_type}); + CheckDispatchBest("coalesce", {other, type}, {value_type, value_type}); + CheckDispatchBest("coalesce", {type, type, other}, + {value_type, value_type, value_type}); + } +} + template class TestChooseNumeric : public ::testing::Test {}; template