diff --git a/cpp/src/arrow/compute/kernels/scalar_if_else.cc b/cpp/src/arrow/compute/kernels/scalar_if_else.cc index 5618f24e5ce..a4bccfef108 100644 --- a/cpp/src/arrow/compute/kernels/scalar_if_else.cc +++ b/cpp/src/arrow/compute/kernels/scalar_if_else.cc @@ -1495,13 +1495,11 @@ struct CaseWhenFunction : ScalarFunction { return arrow::compute::detail::NoMatchingKernel(this, *types); } - static std::shared_ptr DecimalMatchConstraint() { + // For case_when exact dispatch, all value arguments must have identical DataType. + static std::shared_ptr AllValueTypesMatchConstraint() { 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]; }); @@ -2738,7 +2736,7 @@ struct ChooseFunction : ScalarFunction { void AddCaseWhenKernel(const std::shared_ptr& scalar_function, detail::GetTypeId get_id, ArrayKernelExec exec, - std::shared_ptr constraint = nullptr) { + std::shared_ptr constraint) { ScalarKernel kernel( KernelSignature::Make({InputType(Type::STRUCT), InputType(get_id.id)}, LastType, /*is_varargs=*/true, std::move(constraint)), @@ -2756,38 +2754,42 @@ void AddCaseWhenKernel(const std::shared_ptr& scalar_function, } void AddPrimitiveCaseWhenKernels(const std::shared_ptr& scalar_function, - const std::vector>& types) { + const std::vector>& types, + std::shared_ptr constraint) { for (auto&& type : types) { auto exec = GenerateTypeAgnosticPrimitive(*type); - AddCaseWhenKernel(scalar_function, type, std::move(exec)); + AddCaseWhenKernel(scalar_function, type, std::move(exec), constraint); } } void AddBinaryCaseWhenKernels(const std::shared_ptr& scalar_function, - const std::vector>& types) { + const std::vector>& types, + std::shared_ptr constraint) { for (auto&& type : types) { auto exec = GenerateTypeAgnosticVarBinaryBase(*type); - AddCaseWhenKernel(scalar_function, type, std::move(exec)); + AddCaseWhenKernel(scalar_function, type, std::move(exec), constraint); } } template -void AddNestedCaseWhenKernel(const std::shared_ptr& scalar_function) { +void AddNestedCaseWhenKernel(const std::shared_ptr& scalar_function, + std::shared_ptr constraint) { AddCaseWhenKernel(scalar_function, ArrowNestedType::type_id, - CaseWhenFunctor::Exec); + CaseWhenFunctor::Exec, constraint); } -void AddNestedCaseWhenKernels(const std::shared_ptr& scalar_function) { - AddNestedCaseWhenKernel(scalar_function); - AddNestedCaseWhenKernel(scalar_function); - AddNestedCaseWhenKernel(scalar_function); - AddNestedCaseWhenKernel(scalar_function); - AddNestedCaseWhenKernel(scalar_function); - AddNestedCaseWhenKernel(scalar_function); - AddNestedCaseWhenKernel(scalar_function); - AddNestedCaseWhenKernel(scalar_function); - AddNestedCaseWhenKernel(scalar_function); - AddNestedCaseWhenKernel(scalar_function); +void AddNestedCaseWhenKernels(const std::shared_ptr& scalar_function, + std::shared_ptr constraint) { + AddNestedCaseWhenKernel(scalar_function, constraint); + AddNestedCaseWhenKernel(scalar_function, constraint); + AddNestedCaseWhenKernel(scalar_function, constraint); + AddNestedCaseWhenKernel(scalar_function, constraint); + AddNestedCaseWhenKernel(scalar_function, constraint); + AddNestedCaseWhenKernel(scalar_function, constraint); + AddNestedCaseWhenKernel(scalar_function, constraint); + AddNestedCaseWhenKernel(scalar_function, constraint); + AddNestedCaseWhenKernel(scalar_function, constraint); + AddNestedCaseWhenKernel(scalar_function, constraint); } void AddCoalesceKernel(const std::shared_ptr& scalar_function, @@ -2909,19 +2911,21 @@ void RegisterScalarIfElse(FunctionRegistry* registry) { { auto func = std::make_shared( "case_when", Arity::VarArgs(/*min_args=*/2), case_when_doc); - AddPrimitiveCaseWhenKernels(func, NumericTypes()); - AddPrimitiveCaseWhenKernels(func, TemporalTypes()); - AddPrimitiveCaseWhenKernels(func, IntervalTypes()); - AddPrimitiveCaseWhenKernels(func, DurationTypes()); - AddPrimitiveCaseWhenKernels(func, {boolean(), null(), float16()}); + auto all_value_types_match = CaseWhenFunction::AllValueTypesMatchConstraint(); + AddPrimitiveCaseWhenKernels(func, NumericTypes(), all_value_types_match); + AddPrimitiveCaseWhenKernels(func, TemporalTypes(), all_value_types_match); + AddPrimitiveCaseWhenKernels(func, IntervalTypes(), all_value_types_match); + AddPrimitiveCaseWhenKernels(func, DurationTypes(), all_value_types_match); + AddPrimitiveCaseWhenKernels(func, {boolean(), null(), float16()}, + all_value_types_match); AddCaseWhenKernel(func, Type::FIXED_SIZE_BINARY, - CaseWhenFunctor::Exec); + CaseWhenFunctor::Exec, all_value_types_match); AddCaseWhenKernel(func, Type::DECIMAL128, CaseWhenFunctor::Exec, - CaseWhenFunction::DecimalMatchConstraint()); + all_value_types_match); AddCaseWhenKernel(func, Type::DECIMAL256, CaseWhenFunctor::Exec, - CaseWhenFunction::DecimalMatchConstraint()); - AddBinaryCaseWhenKernels(func, BaseBinaryTypes()); - AddNestedCaseWhenKernels(func); + all_value_types_match); + AddBinaryCaseWhenKernels(func, BaseBinaryTypes(), all_value_types_match); + AddNestedCaseWhenKernels(func, all_value_types_match); DCHECK_OK(registry->AddFunction(std::move(func))); } { 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 6fdcff8d970..a1ef82383e2 100644 --- a/cpp/src/arrow/compute/kernels/scalar_if_else_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_if_else_test.cc @@ -2726,6 +2726,36 @@ TEST(TestCaseWhen, UnionBoolStringRandom) { } TEST(TestCaseWhen, DispatchExact) { + // Matching parameterized types should exact-match. + CheckDispatchExact("case_when", {struct_({field("", boolean())}), fixed_size_binary(4), + fixed_size_binary(4)}); + CheckDispatchExact("case_when", + {struct_({field("", boolean())}), list(int32()), list(int32())}); + CheckDispatchExact("case_when", + {struct_({field("", boolean())}), fixed_size_list(int32(), 2), + fixed_size_list(int32(), 2)}); + CheckDispatchExact("case_when", + {struct_({field("", boolean())}), dictionary(int8(), utf8()), + dictionary(int8(), utf8())}); + + // Mismatched parameterized types should not exact-match. + CheckDispatchExactFails("case_when", {struct_({field("", boolean())}), + fixed_size_binary(4), fixed_size_binary(5)}); + CheckDispatchExactFails( + "case_when", {struct_({field("", boolean())}), list(int16()), list(int32())}); + CheckDispatchExactFails("case_when", + {struct_({field("", boolean())}), fixed_size_list(int32(), 2), + fixed_size_list(int32(), 3)}); + CheckDispatchExactFails( + "case_when", {struct_({field("", boolean())}), struct_({field("a", int32())}), + struct_({field("a", int64())})}); + CheckDispatchExactFails("case_when", + {struct_({field("", boolean())}), dictionary(int8(), utf8()), + dictionary(int8(), large_utf8())}); + CheckDispatchExactFails("case_when", + {struct_({field("", boolean())}), dictionary(int8(), utf8()), + dictionary(int16(), utf8())}); + // Decimal types with same (p, s) CheckDispatchExact("case_when", {struct_({field("", boolean())}), decimal128(20, 3), decimal128(20, 3)}); @@ -2825,6 +2855,28 @@ TEST(TestCaseWhen, DispatchBest) { {struct_({field("", boolean())}), decimal256(23, 3), decimal256(23, 3)}); } +TEST(TestCaseWhen, ParameterizedValueTypeMismatch) { + auto cond = MakeStruct({ArrayFromJSON(boolean(), "[true]")}); + + ASSERT_RAISES( + NotImplemented, + CallFunction("case_when", {cond, ArrayFromJSON(fixed_size_binary(4), R"(["abcd"])"), + ArrayFromJSON(fixed_size_binary(5), R"(["efghi"])")})); + ASSERT_RAISES(NotImplemented, + CallFunction("case_when", {cond, ArrayFromJSON(list(int16()), "[[1, 2]]"), + ArrayFromJSON(list(int32()), "[[3, 4]]")})); + ASSERT_RAISES( + NotImplemented, + CallFunction("case_when", + {cond, ArrayFromJSON(fixed_size_list(int32(), 2), "[[1, 2]]"), + ArrayFromJSON(fixed_size_list(int32(), 3), "[[3, 4, 5]]")})); + ASSERT_RAISES( + NotImplemented, + CallFunction("case_when", + {cond, ArrayFromJSON(struct_({field("a", int32())}), R"([{"a": 1}])"), + ArrayFromJSON(struct_({field("a", int64())}), R"([{"a": 2}])")})); +} + template class TestCoalesceNumeric : public ::testing::Test {}; template