Skip to content

Commit 974c3fc

Browse files
zanmato1984zanmato1984
authored
GH-50624: [C++][Compute] Tighten case_when exact dispatch for parameterized types (#50625)
### Rationale for this change `case_when` exact dispatch was too permissive for parameterized value types. When the value arguments were registered by type id only, `DispatchExact` could incorrectly accept incompatible concrete types and return successful but corrupted results instead of failing or falling back to `DispatchBest`. This fixes #50624. ### What changes are included in this PR? - rename the case_when exact-match helper to `AllValueTypesMatchConstraint()` to make its intent explicit - apply the same exact-match constraint across case_when value kernels so exact dispatch only accepts identical value `DataType`s - add regression coverage for mismatched parameterized value types, including `fixed_size_binary`, `list`, `fixed_size_list`, `struct`, and dictionary exact-dispatch mismatches ### Are these changes tested? Yes. This PR extends `TestCaseWhen.DispatchExact` and adds `TestCaseWhen.ParameterizedValueTypeMismatch`. I also ran targeted `arrow-compute-scalar-if-else-test` coverage for: - `TestCaseWhen.DispatchExact` - `TestCaseWhen.DispatchBest` - `TestCaseWhen.ParameterizedValueTypeMismatch` ### Are there any user-facing changes? Yes. `case_when` now rejects incompatible parameterized value types that previously could be incorrectly exact-dispatched, which could lead to corrupted results. **This PR contains a "Critical Fix".** It fixes a bug where `case_when` could return successful but corrupted results for incompatible parameterized value types. * GitHub Issue: #50624 Authored-by: Rossi Sun <zanmato1984@gmail.com> Signed-off-by: Rossi Sun <zanmato1984@gmail.com>
1 parent 9c37a70 commit 974c3fc

2 files changed

Lines changed: 89 additions & 33 deletions

File tree

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

Lines changed: 37 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -1495,13 +1495,11 @@ struct CaseWhenFunction : ScalarFunction {
14951495
return arrow::compute::detail::NoMatchingKernel(this, *types);
14961496
}
14971497

1498-
static std::shared_ptr<MatchConstraint> DecimalMatchConstraint() {
1498+
// For case_when exact dispatch, all value arguments must have identical DataType.
1499+
static std::shared_ptr<MatchConstraint> AllValueTypesMatchConstraint() {
14991500
static auto constraint =
15001501
MatchConstraint::Make([](const std::vector<TypeHolder>& types) -> bool {
15011502
DCHECK_GE(types.size(), 2);
1502-
DCHECK(std::all_of(types.begin() + 1, types.end(), [](const TypeHolder& type) {
1503-
return is_decimal(type.id());
1504-
}));
15051503
return std::all_of(
15061504
types.begin() + 2, types.end(),
15071505
[&types](const TypeHolder& type) { return type == types[1]; });
@@ -2738,10 +2736,10 @@ struct ChooseFunction : ScalarFunction {
27382736

27392737
void AddCaseWhenKernel(const std::shared_ptr<CaseWhenFunction>& scalar_function,
27402738
detail::GetTypeId get_id, ArrayKernelExec exec,
2741-
std::shared_ptr<MatchConstraint> constraint = nullptr) {
2739+
const std::shared_ptr<MatchConstraint>& constraint) {
27422740
ScalarKernel kernel(
27432741
KernelSignature::Make({InputType(Type::STRUCT), InputType(get_id.id)}, LastType,
2744-
/*is_varargs=*/true, std::move(constraint)),
2742+
/*is_varargs=*/true, constraint),
27452743
exec);
27462744
if (is_fixed_width(get_id.id)) {
27472745
kernel.null_handling = NullHandling::COMPUTED_PREALLOCATE;
@@ -2756,38 +2754,42 @@ void AddCaseWhenKernel(const std::shared_ptr<CaseWhenFunction>& scalar_function,
27562754
}
27572755

27582756
void AddPrimitiveCaseWhenKernels(const std::shared_ptr<CaseWhenFunction>& scalar_function,
2759-
const std::vector<std::shared_ptr<DataType>>& types) {
2757+
const std::vector<std::shared_ptr<DataType>>& types,
2758+
const std::shared_ptr<MatchConstraint>& constraint) {
27602759
for (auto&& type : types) {
27612760
auto exec = GenerateTypeAgnosticPrimitive<CaseWhenFunctor>(*type);
2762-
AddCaseWhenKernel(scalar_function, type, std::move(exec));
2761+
AddCaseWhenKernel(scalar_function, type, std::move(exec), constraint);
27632762
}
27642763
}
27652764

27662765
void AddBinaryCaseWhenKernels(const std::shared_ptr<CaseWhenFunction>& scalar_function,
2767-
const std::vector<std::shared_ptr<DataType>>& types) {
2766+
const std::vector<std::shared_ptr<DataType>>& types,
2767+
const std::shared_ptr<MatchConstraint>& constraint) {
27682768
for (auto&& type : types) {
27692769
auto exec = GenerateTypeAgnosticVarBinaryBase<CaseWhenFunctor>(*type);
2770-
AddCaseWhenKernel(scalar_function, type, std::move(exec));
2770+
AddCaseWhenKernel(scalar_function, type, std::move(exec), constraint);
27712771
}
27722772
}
27732773

27742774
template <typename ArrowNestedType>
2775-
void AddNestedCaseWhenKernel(const std::shared_ptr<CaseWhenFunction>& scalar_function) {
2775+
void AddNestedCaseWhenKernel(const std::shared_ptr<CaseWhenFunction>& scalar_function,
2776+
const std::shared_ptr<MatchConstraint>& constraint) {
27762777
AddCaseWhenKernel(scalar_function, ArrowNestedType::type_id,
2777-
CaseWhenFunctor<ArrowNestedType>::Exec);
2778+
CaseWhenFunctor<ArrowNestedType>::Exec, constraint);
27782779
}
27792780

2780-
void AddNestedCaseWhenKernels(const std::shared_ptr<CaseWhenFunction>& scalar_function) {
2781-
AddNestedCaseWhenKernel<FixedSizeListType>(scalar_function);
2782-
AddNestedCaseWhenKernel<ListType>(scalar_function);
2783-
AddNestedCaseWhenKernel<LargeListType>(scalar_function);
2784-
AddNestedCaseWhenKernel<ListViewType>(scalar_function);
2785-
AddNestedCaseWhenKernel<LargeListViewType>(scalar_function);
2786-
AddNestedCaseWhenKernel<MapType>(scalar_function);
2787-
AddNestedCaseWhenKernel<StructType>(scalar_function);
2788-
AddNestedCaseWhenKernel<DenseUnionType>(scalar_function);
2789-
AddNestedCaseWhenKernel<SparseUnionType>(scalar_function);
2790-
AddNestedCaseWhenKernel<DictionaryType>(scalar_function);
2781+
void AddNestedCaseWhenKernels(const std::shared_ptr<CaseWhenFunction>& scalar_function,
2782+
const std::shared_ptr<MatchConstraint>& constraint) {
2783+
AddNestedCaseWhenKernel<FixedSizeListType>(scalar_function, constraint);
2784+
AddNestedCaseWhenKernel<ListType>(scalar_function, constraint);
2785+
AddNestedCaseWhenKernel<LargeListType>(scalar_function, constraint);
2786+
AddNestedCaseWhenKernel<ListViewType>(scalar_function, constraint);
2787+
AddNestedCaseWhenKernel<LargeListViewType>(scalar_function, constraint);
2788+
AddNestedCaseWhenKernel<MapType>(scalar_function, constraint);
2789+
AddNestedCaseWhenKernel<StructType>(scalar_function, constraint);
2790+
AddNestedCaseWhenKernel<DenseUnionType>(scalar_function, constraint);
2791+
AddNestedCaseWhenKernel<SparseUnionType>(scalar_function, constraint);
2792+
AddNestedCaseWhenKernel<DictionaryType>(scalar_function, constraint);
27912793
}
27922794

27932795
void AddCoalesceKernel(const std::shared_ptr<ScalarFunction>& scalar_function,
@@ -2909,19 +2911,21 @@ void RegisterScalarIfElse(FunctionRegistry* registry) {
29092911
{
29102912
auto func = std::make_shared<CaseWhenFunction>(
29112913
"case_when", Arity::VarArgs(/*min_args=*/2), case_when_doc);
2912-
AddPrimitiveCaseWhenKernels(func, NumericTypes());
2913-
AddPrimitiveCaseWhenKernels(func, TemporalTypes());
2914-
AddPrimitiveCaseWhenKernels(func, IntervalTypes());
2915-
AddPrimitiveCaseWhenKernels(func, DurationTypes());
2916-
AddPrimitiveCaseWhenKernels(func, {boolean(), null(), float16()});
2914+
auto all_value_types_match = CaseWhenFunction::AllValueTypesMatchConstraint();
2915+
AddPrimitiveCaseWhenKernels(func, NumericTypes(), all_value_types_match);
2916+
AddPrimitiveCaseWhenKernels(func, TemporalTypes(), all_value_types_match);
2917+
AddPrimitiveCaseWhenKernels(func, IntervalTypes(), all_value_types_match);
2918+
AddPrimitiveCaseWhenKernels(func, DurationTypes(), all_value_types_match);
2919+
AddPrimitiveCaseWhenKernels(func, {boolean(), null(), float16()},
2920+
all_value_types_match);
29172921
AddCaseWhenKernel(func, Type::FIXED_SIZE_BINARY,
2918-
CaseWhenFunctor<FixedSizeBinaryType>::Exec);
2922+
CaseWhenFunctor<FixedSizeBinaryType>::Exec, all_value_types_match);
29192923
AddCaseWhenKernel(func, Type::DECIMAL128, CaseWhenFunctor<FixedSizeBinaryType>::Exec,
2920-
CaseWhenFunction::DecimalMatchConstraint());
2924+
all_value_types_match);
29212925
AddCaseWhenKernel(func, Type::DECIMAL256, CaseWhenFunctor<FixedSizeBinaryType>::Exec,
2922-
CaseWhenFunction::DecimalMatchConstraint());
2923-
AddBinaryCaseWhenKernels(func, BaseBinaryTypes());
2924-
AddNestedCaseWhenKernels(func);
2926+
all_value_types_match);
2927+
AddBinaryCaseWhenKernels(func, BaseBinaryTypes(), all_value_types_match);
2928+
AddNestedCaseWhenKernels(func, all_value_types_match);
29252929
DCHECK_OK(registry->AddFunction(std::move(func)));
29262930
}
29272931
{

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

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2726,6 +2726,36 @@ TEST(TestCaseWhen, UnionBoolStringRandom) {
27262726
}
27272727

27282728
TEST(TestCaseWhen, DispatchExact) {
2729+
// Matching parameterized types should exact-match.
2730+
CheckDispatchExact("case_when", {struct_({field("", boolean())}), fixed_size_binary(4),
2731+
fixed_size_binary(4)});
2732+
CheckDispatchExact("case_when",
2733+
{struct_({field("", boolean())}), list(int32()), list(int32())});
2734+
CheckDispatchExact("case_when",
2735+
{struct_({field("", boolean())}), fixed_size_list(int32(), 2),
2736+
fixed_size_list(int32(), 2)});
2737+
CheckDispatchExact("case_when",
2738+
{struct_({field("", boolean())}), dictionary(int8(), utf8()),
2739+
dictionary(int8(), utf8())});
2740+
2741+
// Mismatched parameterized types should not exact-match.
2742+
CheckDispatchExactFails("case_when", {struct_({field("", boolean())}),
2743+
fixed_size_binary(4), fixed_size_binary(5)});
2744+
CheckDispatchExactFails(
2745+
"case_when", {struct_({field("", boolean())}), list(int16()), list(int32())});
2746+
CheckDispatchExactFails("case_when",
2747+
{struct_({field("", boolean())}), fixed_size_list(int32(), 2),
2748+
fixed_size_list(int32(), 3)});
2749+
CheckDispatchExactFails(
2750+
"case_when", {struct_({field("", boolean())}), struct_({field("a", int32())}),
2751+
struct_({field("a", int64())})});
2752+
CheckDispatchExactFails("case_when",
2753+
{struct_({field("", boolean())}), dictionary(int8(), utf8()),
2754+
dictionary(int8(), large_utf8())});
2755+
CheckDispatchExactFails("case_when",
2756+
{struct_({field("", boolean())}), dictionary(int8(), utf8()),
2757+
dictionary(int16(), utf8())});
2758+
27292759
// Decimal types with same (p, s)
27302760
CheckDispatchExact("case_when", {struct_({field("", boolean())}), decimal128(20, 3),
27312761
decimal128(20, 3)});
@@ -2825,6 +2855,28 @@ TEST(TestCaseWhen, DispatchBest) {
28252855
{struct_({field("", boolean())}), decimal256(23, 3), decimal256(23, 3)});
28262856
}
28272857

2858+
TEST(TestCaseWhen, ParameterizedValueTypeMismatch) {
2859+
auto cond = MakeStruct({ArrayFromJSON(boolean(), "[true]")});
2860+
2861+
ASSERT_RAISES(
2862+
NotImplemented,
2863+
CallFunction("case_when", {cond, ArrayFromJSON(fixed_size_binary(4), R"(["abcd"])"),
2864+
ArrayFromJSON(fixed_size_binary(5), R"(["efghi"])")}));
2865+
ASSERT_RAISES(NotImplemented,
2866+
CallFunction("case_when", {cond, ArrayFromJSON(list(int16()), "[[1, 2]]"),
2867+
ArrayFromJSON(list(int32()), "[[3, 4]]")}));
2868+
ASSERT_RAISES(
2869+
NotImplemented,
2870+
CallFunction("case_when",
2871+
{cond, ArrayFromJSON(fixed_size_list(int32(), 2), "[[1, 2]]"),
2872+
ArrayFromJSON(fixed_size_list(int32(), 3), "[[3, 4, 5]]")}));
2873+
ASSERT_RAISES(
2874+
NotImplemented,
2875+
CallFunction("case_when",
2876+
{cond, ArrayFromJSON(struct_({field("a", int32())}), R"([{"a": 1}])"),
2877+
ArrayFromJSON(struct_({field("a", int64())}), R"([{"a": 2}])")}));
2878+
}
2879+
28282880
template <typename Type>
28292881
class TestCoalesceNumeric : public ::testing::Test {};
28302882
template <typename Type>

0 commit comments

Comments
 (0)