Skip to content

Commit 23e9ff6

Browse files
zanmato1984zanmato1984
committed
GH-50869: [C++][Compute] Tighten coalesce exact dispatch for decimal varargs
Signed-off-by: Rossi Sun <zanmato1984@gmail.com>
1 parent e611f48 commit 23e9ff6

3 files changed

Lines changed: 104 additions & 4 deletions

File tree

‎cpp/src/arrow/compute/expression_test.cc‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -938,6 +938,62 @@ TEST(Expression, BindWithImplicitCastsForCaseWhenOnDecimal) {
938938
/*bound_out=*/nullptr, *exciting_schema);
939939
}
940940

941+
TEST(Expression, BindWithImplicitCastsForCoalesceOnDecimal) {
942+
auto exciting_schema = schema(
943+
{field("dec128_3_2", decimal128(3, 2)), field("dec128_4_1", decimal128(4, 1)),
944+
field("dec128_4_2", decimal128(4, 2)), field("dec128_4_3", decimal128(4, 3)),
945+
field("dec256_3_2", decimal256(3, 2))});
946+
947+
ExpectBindsTo(call("coalesce", {field_ref("dec128_3_2"), field_ref("dec128_4_2")}),
948+
call("coalesce", {cast(field_ref("dec128_3_2"), decimal128(4, 2)),
949+
field_ref("dec128_4_2")}),
950+
/*bound_out=*/nullptr, *exciting_schema);
951+
ExpectBindsTo(call("coalesce", {field_ref("dec128_4_2"), field_ref("dec128_3_2")}),
952+
call("coalesce", {field_ref("dec128_4_2"),
953+
cast(field_ref("dec128_3_2"), decimal128(4, 2))}),
954+
/*bound_out=*/nullptr, *exciting_schema);
955+
ExpectBindsTo(call("coalesce", {field_ref("dec128_4_1"), field_ref("dec128_3_2")}),
956+
call("coalesce", {cast(field_ref("dec128_4_1"), decimal128(5, 2)),
957+
cast(field_ref("dec128_3_2"), decimal128(5, 2))}),
958+
/*bound_out=*/nullptr, *exciting_schema);
959+
ExpectBindsTo(call("coalesce", {field_ref("dec128_3_2"), field_ref("dec128_4_1")}),
960+
call("coalesce", {cast(field_ref("dec128_3_2"), decimal128(5, 2)),
961+
cast(field_ref("dec128_4_1"), decimal128(5, 2))}),
962+
/*bound_out=*/nullptr, *exciting_schema);
963+
ExpectBindsTo(call("coalesce", {field_ref("dec128_3_2"), field_ref("dec128_4_3")}),
964+
call("coalesce", {cast(field_ref("dec128_3_2"), decimal128(4, 3)),
965+
field_ref("dec128_4_3")}),
966+
/*bound_out=*/nullptr, *exciting_schema);
967+
ExpectBindsTo(call("coalesce", {field_ref("dec128_4_3"), field_ref("dec128_3_2")}),
968+
call("coalesce", {field_ref("dec128_4_3"),
969+
cast(field_ref("dec128_3_2"), decimal128(4, 3))}),
970+
/*bound_out=*/nullptr, *exciting_schema);
971+
ExpectBindsTo(call("coalesce", {field_ref("dec128_3_2"), field_ref("dec256_3_2")}),
972+
call("coalesce", {cast(field_ref("dec128_3_2"), decimal256(3, 2)),
973+
field_ref("dec256_3_2")}),
974+
/*bound_out=*/nullptr, *exciting_schema);
975+
ExpectBindsTo(call("coalesce", {field_ref("dec256_3_2"), field_ref("dec128_3_2")}),
976+
call("coalesce", {field_ref("dec256_3_2"),
977+
cast(field_ref("dec128_3_2"), decimal256(3, 2))}),
978+
/*bound_out=*/nullptr, *exciting_schema);
979+
}
980+
981+
TEST(Expression, ExecuteCoalesceOnMixedDecimalTypes) {
982+
ASSERT_OK_AND_ASSIGN(
983+
auto input,
984+
StructArray::Make({ArrayFromJSON(decimal128(3, 2), R"(["1.23", null])"),
985+
ArrayFromJSON(decimal128(4, 3), R"([null, "2.345"])")},
986+
{"left", "right"}));
987+
Schema input_schema(input->type()->fields());
988+
auto expr = call("coalesce", {field_ref("left"), field_ref("right")});
989+
990+
ASSERT_OK_AND_ASSIGN(expr, expr.Bind(input_schema));
991+
ASSERT_OK_AND_ASSIGN(auto actual,
992+
ExecuteScalarExpression(expr, input_schema, Datum(input)));
993+
994+
AssertDatumsEqual(actual, ArrayFromJSON(decimal128(4, 3), R"(["1.230", "2.345"])"));
995+
}
996+
941997
TEST(Expression, BindNestedCall) {
942998
auto expr = add(field_ref("a"),
943999
call("subtract", {call("multiply", {field_ref("b"), field_ref("c")}),

‎cpp/src/arrow/compute/kernels/scalar_if_else.cc‎

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2035,6 +2035,20 @@ struct CoalesceFunction : ScalarFunction {
20352035
if (auto kernel = DispatchExactImpl(this, *types)) return kernel;
20362036
return arrow::compute::detail::NoMatchingKernel(this, *types);
20372037
}
2038+
2039+
static std::shared_ptr<MatchConstraint> DecimalMatchConstraint() {
2040+
static auto constraint =
2041+
MatchConstraint::Make([](const std::vector<TypeHolder>& types) -> bool {
2042+
DCHECK_GE(types.size(), 1);
2043+
DCHECK(std::all_of(types.begin(), types.end(), [](const TypeHolder& type) {
2044+
return is_decimal(type.id());
2045+
}));
2046+
return std::all_of(
2047+
types.begin() + 1, types.end(),
2048+
[&types](const TypeHolder& type) { return type == types[0]; });
2049+
});
2050+
return constraint;
2051+
}
20382052
};
20392053

20402054
// Helper: copy from a source value into all null slots of the output
@@ -2793,9 +2807,10 @@ void AddNestedCaseWhenKernels(const std::shared_ptr<CaseWhenFunction>& scalar_fu
27932807
}
27942808

27952809
void AddCoalesceKernel(const std::shared_ptr<ScalarFunction>& scalar_function,
2796-
detail::GetTypeId get_id, ArrayKernelExec exec) {
2810+
detail::GetTypeId get_id, ArrayKernelExec exec,
2811+
std::shared_ptr<MatchConstraint> constraint = nullptr) {
27972812
ScalarKernel kernel(KernelSignature::Make({InputType(get_id.id)}, FirstType,
2798-
/*is_varargs=*/true),
2813+
/*is_varargs=*/true, std::move(constraint)),
27992814
exec);
28002815
kernel.null_handling = NullHandling::COMPUTED_PREALLOCATE;
28012816
kernel.mem_allocation = MemAllocation::PREALLOCATE;
@@ -2938,8 +2953,10 @@ void RegisterScalarIfElse(FunctionRegistry* registry) {
29382953
AddPrimitiveCoalesceKernels(func, {boolean(), null(), float16()});
29392954
AddCoalesceKernel(func, Type::FIXED_SIZE_BINARY,
29402955
CoalesceFunctor<FixedSizeBinaryType>::Exec);
2941-
AddCoalesceKernel(func, Type::DECIMAL128, CoalesceFunctor<FixedSizeBinaryType>::Exec);
2942-
AddCoalesceKernel(func, Type::DECIMAL256, CoalesceFunctor<FixedSizeBinaryType>::Exec);
2956+
AddCoalesceKernel(func, Type::DECIMAL128, CoalesceFunctor<FixedSizeBinaryType>::Exec,
2957+
CoalesceFunction::DecimalMatchConstraint());
2958+
AddCoalesceKernel(func, Type::DECIMAL256, CoalesceFunctor<FixedSizeBinaryType>::Exec,
2959+
CoalesceFunction::DecimalMatchConstraint());
29432960
for (const auto& ty : BaseBinaryTypes()) {
29442961
AddCoalesceKernel(func, ty, GenerateTypeAgnosticVarBinaryBase<CoalesceFunctor>(ty));
29452962
}

‎cpp/src/arrow/compute/kernels/scalar_if_else_test.cc‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3693,8 +3693,22 @@ TEST(TestCoalesce, DispatchBest) {
36933693
CheckDispatchBest("coalesce", {int32(), decimal128(3, 2)},
36943694
{decimal128(12, 2), decimal128(12, 2)});
36953695
CheckDispatchBest("coalesce", {float32(), decimal128(3, 2)}, {float64(), float64()});
3696+
CheckDispatchBest("coalesce", {decimal128(3, 2), decimal128(4, 2)},
3697+
{decimal128(4, 2), decimal128(4, 2)});
3698+
CheckDispatchBest("coalesce", {decimal128(4, 2), decimal128(3, 2)},
3699+
{decimal128(4, 2), decimal128(4, 2)});
3700+
CheckDispatchBest("coalesce", {decimal128(4, 1), decimal128(3, 2)},
3701+
{decimal128(5, 2), decimal128(5, 2)});
3702+
CheckDispatchBest("coalesce", {decimal128(3, 2), decimal128(4, 1)},
3703+
{decimal128(5, 2), decimal128(5, 2)});
3704+
CheckDispatchBest("coalesce", {decimal128(3, 2), decimal128(4, 3)},
3705+
{decimal128(4, 3), decimal128(4, 3)});
3706+
CheckDispatchBest("coalesce", {decimal128(4, 3), decimal128(3, 2)},
3707+
{decimal128(4, 3), decimal128(4, 3)});
36963708
CheckDispatchBest("coalesce", {decimal128(3, 2), decimal256(3, 2)},
36973709
{decimal256(3, 2), decimal256(3, 2)});
3710+
CheckDispatchBest("coalesce", {decimal256(3, 2), decimal128(3, 2)},
3711+
{decimal256(3, 2), decimal256(3, 2)});
36983712
CheckDispatchBest("coalesce", {timestamp(TimeUnit::SECOND), date32()},
36993713
{timestamp(TimeUnit::SECOND), timestamp(TimeUnit::SECOND)});
37003714
CheckDispatchBest("coalesce", {timestamp(TimeUnit::SECOND), timestamp(TimeUnit::MILLI)},
@@ -3710,6 +3724,19 @@ TEST(TestCoalesce, DispatchBest) {
37103724
{large_binary(), large_binary()});
37113725
}
37123726

3727+
TEST(TestCoalesce, DispatchExact) {
3728+
CheckDispatchExact("coalesce", {decimal128(3, 2), decimal128(3, 2)});
3729+
CheckDispatchExact("coalesce", {decimal256(3, 2), decimal256(3, 2)});
3730+
CheckDispatchExactFails("coalesce", {decimal128(3, 2), decimal128(4, 2)});
3731+
CheckDispatchExactFails("coalesce", {decimal128(4, 2), decimal128(3, 2)});
3732+
CheckDispatchExactFails("coalesce", {decimal128(4, 1), decimal128(3, 2)});
3733+
CheckDispatchExactFails("coalesce", {decimal128(3, 2), decimal128(4, 1)});
3734+
CheckDispatchExactFails("coalesce", {decimal128(3, 2), decimal128(4, 3)});
3735+
CheckDispatchExactFails("coalesce", {decimal128(4, 3), decimal128(3, 2)});
3736+
CheckDispatchExactFails("coalesce", {decimal128(3, 2), decimal256(3, 2)});
3737+
CheckDispatchExactFails("coalesce", {decimal256(3, 2), decimal128(3, 2)});
3738+
}
3739+
37133740
template <typename Type>
37143741
class TestChooseNumeric : public ::testing::Test {};
37153742
template <typename Type>

0 commit comments

Comments
 (0)