This is an automated email from the ASF dual-hosted git repository.
zanmato1984 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/main by this push:
new 974c3fc1a77 GH-50624: [C++][Compute] Tighten case_when exact dispatch
for parameterized types (#50625)
974c3fc1a77 is described below
commit 974c3fc1a772eccc85bd3500bad8573a57ff5089
Author: Rossi Sun <[email protected]>
AuthorDate: Mon Jul 27 11:31:51 2026 -0700
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 <[email protected]>
Signed-off-by: Rossi Sun <[email protected]>
---
cpp/src/arrow/compute/kernels/scalar_if_else.cc | 70 ++++++++++++----------
.../arrow/compute/kernels/scalar_if_else_test.cc | 52 ++++++++++++++++
2 files changed, 89 insertions(+), 33 deletions(-)
diff --git a/cpp/src/arrow/compute/kernels/scalar_if_else.cc
b/cpp/src/arrow/compute/kernels/scalar_if_else.cc
index 26f880fc0c6..1510dd9fc83 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<MatchConstraint> DecimalMatchConstraint() {
+ // For case_when exact dispatch, all value arguments must have identical
DataType.
+ static std::shared_ptr<MatchConstraint> AllValueTypesMatchConstraint() {
static auto constraint =
MatchConstraint::Make([](const std::vector<TypeHolder>& 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,10 +2736,10 @@ struct ChooseFunction : ScalarFunction {
void AddCaseWhenKernel(const std::shared_ptr<CaseWhenFunction>&
scalar_function,
detail::GetTypeId get_id, ArrayKernelExec exec,
- std::shared_ptr<MatchConstraint> constraint = nullptr) {
+ const std::shared_ptr<MatchConstraint>& constraint) {
ScalarKernel kernel(
KernelSignature::Make({InputType(Type::STRUCT), InputType(get_id.id)},
LastType,
- /*is_varargs=*/true, std::move(constraint)),
+ /*is_varargs=*/true, constraint),
exec);
if (is_fixed_width(get_id.id)) {
kernel.null_handling = NullHandling::COMPUTED_PREALLOCATE;
@@ -2756,38 +2754,42 @@ void AddCaseWhenKernel(const
std::shared_ptr<CaseWhenFunction>& scalar_function,
}
void AddPrimitiveCaseWhenKernels(const std::shared_ptr<CaseWhenFunction>&
scalar_function,
- const std::vector<std::shared_ptr<DataType>>&
types) {
+ const std::vector<std::shared_ptr<DataType>>&
types,
+ const std::shared_ptr<MatchConstraint>&
constraint) {
for (auto&& type : types) {
auto exec = GenerateTypeAgnosticPrimitive<CaseWhenFunctor>(*type);
- AddCaseWhenKernel(scalar_function, type, std::move(exec));
+ AddCaseWhenKernel(scalar_function, type, std::move(exec), constraint);
}
}
void AddBinaryCaseWhenKernels(const std::shared_ptr<CaseWhenFunction>&
scalar_function,
- const std::vector<std::shared_ptr<DataType>>&
types) {
+ const std::vector<std::shared_ptr<DataType>>&
types,
+ const std::shared_ptr<MatchConstraint>&
constraint) {
for (auto&& type : types) {
auto exec = GenerateTypeAgnosticVarBinaryBase<CaseWhenFunctor>(*type);
- AddCaseWhenKernel(scalar_function, type, std::move(exec));
+ AddCaseWhenKernel(scalar_function, type, std::move(exec), constraint);
}
}
template <typename ArrowNestedType>
-void AddNestedCaseWhenKernel(const std::shared_ptr<CaseWhenFunction>&
scalar_function) {
+void AddNestedCaseWhenKernel(const std::shared_ptr<CaseWhenFunction>&
scalar_function,
+ const std::shared_ptr<MatchConstraint>&
constraint) {
AddCaseWhenKernel(scalar_function, ArrowNestedType::type_id,
- CaseWhenFunctor<ArrowNestedType>::Exec);
+ CaseWhenFunctor<ArrowNestedType>::Exec, constraint);
}
-void AddNestedCaseWhenKernels(const std::shared_ptr<CaseWhenFunction>&
scalar_function) {
- AddNestedCaseWhenKernel<FixedSizeListType>(scalar_function);
- AddNestedCaseWhenKernel<ListType>(scalar_function);
- AddNestedCaseWhenKernel<LargeListType>(scalar_function);
- AddNestedCaseWhenKernel<ListViewType>(scalar_function);
- AddNestedCaseWhenKernel<LargeListViewType>(scalar_function);
- AddNestedCaseWhenKernel<MapType>(scalar_function);
- AddNestedCaseWhenKernel<StructType>(scalar_function);
- AddNestedCaseWhenKernel<DenseUnionType>(scalar_function);
- AddNestedCaseWhenKernel<SparseUnionType>(scalar_function);
- AddNestedCaseWhenKernel<DictionaryType>(scalar_function);
+void AddNestedCaseWhenKernels(const std::shared_ptr<CaseWhenFunction>&
scalar_function,
+ const std::shared_ptr<MatchConstraint>&
constraint) {
+ AddNestedCaseWhenKernel<FixedSizeListType>(scalar_function, constraint);
+ AddNestedCaseWhenKernel<ListType>(scalar_function, constraint);
+ AddNestedCaseWhenKernel<LargeListType>(scalar_function, constraint);
+ AddNestedCaseWhenKernel<ListViewType>(scalar_function, constraint);
+ AddNestedCaseWhenKernel<LargeListViewType>(scalar_function, constraint);
+ AddNestedCaseWhenKernel<MapType>(scalar_function, constraint);
+ AddNestedCaseWhenKernel<StructType>(scalar_function, constraint);
+ AddNestedCaseWhenKernel<DenseUnionType>(scalar_function, constraint);
+ AddNestedCaseWhenKernel<SparseUnionType>(scalar_function, constraint);
+ AddNestedCaseWhenKernel<DictionaryType>(scalar_function, constraint);
}
void AddCoalesceKernel(const std::shared_ptr<ScalarFunction>& scalar_function,
@@ -2909,19 +2911,21 @@ void RegisterScalarIfElse(FunctionRegistry* registry) {
{
auto func = std::make_shared<CaseWhenFunction>(
"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<FixedSizeBinaryType>::Exec);
+ CaseWhenFunctor<FixedSizeBinaryType>::Exec,
all_value_types_match);
AddCaseWhenKernel(func, Type::DECIMAL128,
CaseWhenFunctor<FixedSizeBinaryType>::Exec,
- CaseWhenFunction::DecimalMatchConstraint());
+ all_value_types_match);
AddCaseWhenKernel(func, Type::DECIMAL256,
CaseWhenFunctor<FixedSizeBinaryType>::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 <typename Type>
class TestCoalesceNumeric : public ::testing::Test {};
template <typename Type>
