bkietz commented on a change in pull request #10472:
URL: https://github.com/apache/arrow/pull/10472#discussion_r646926844
##########
File path: cpp/src/arrow/compute/kernels/scalar_if_else_test.cc
##########
@@ -271,5 +271,40 @@ TEST_F(TestIfElseKernel, IfElseNull) {
ArrayFromJSON(null(), "[null, null, null, null]"));
}
+TEST_F(TestIfElseKernel, IfElseMultiType) {
+ CheckWithDifferentShapes(ArrayFromJSON(boolean(), "[true, true, true,
false]"),
+ ArrayFromJSON(int32(), "[1, 2, 3, 4]"),
+ ArrayFromJSON(float32(), "[5, 6, 7, 8]"),
+ ArrayFromJSON(float32(), "[1, 2, 3, 8]"));
+}
+
+TEST_F(TestIfElseKernel, IfElseDispatchBest) {
+ std::string name = "if_else";
+ CheckDispatchBest(name, {boolean(), int32(), int32()}, {boolean(), int32(),
int32()});
+ CheckDispatchBest(name, {boolean(), int32(), null()}, {boolean(), int32(),
int32()});
+ CheckDispatchBest(name, {boolean(), null(), int32()}, {boolean(), int32(),
int32()});
+
+ CheckDispatchBest(name, {boolean(), int32(), int8()}, {boolean(), int32(),
int32()});
+ CheckDispatchBest(name, {boolean(), int32(), int16()}, {boolean(), int32(),
int32()});
+ CheckDispatchBest(name, {boolean(), int32(), int32()}, {boolean(), int32(),
int32()});
+ CheckDispatchBest(name, {boolean(), int32(), int64()}, {boolean(), int64(),
int64()});
+
+ CheckDispatchBest(name, {boolean(), int32(), uint8()}, {boolean(), int32(),
int32()});
+ CheckDispatchBest(name, {boolean(), int32(), uint16()}, {boolean(), int32(),
int32()});
+ CheckDispatchBest(name, {boolean(), int32(), uint32()}, {boolean(), int64(),
int64()});
+ CheckDispatchBest(name, {boolean(), int32(), uint64()}, {boolean(), int64(),
int64()});
+
+ CheckDispatchBest(name, {boolean(), uint8(), uint8()}, {boolean(), uint8(),
uint8()});
+ CheckDispatchBest(name, {boolean(), uint8(), uint16()},
+ {boolean(), uint16(), uint16()});
+
Review comment:
It'd be nice to support conditions being `null` in alignment with the
pattern that "anything may be cast to from null"
```suggestion
CheckDispatchBest(name, {null(), uint8(), int8()}, {boolean(), int8(),
int8()});
```
##########
File path: cpp/src/arrow/compute/kernels/scalar_if_else.cc
##########
@@ -543,7 +543,34 @@ struct ResolveIfElseExec {
}
};
-void AddPrimitiveIfElseKernels(const std::shared_ptr<ScalarFunction>&
scalar_function,
+struct IfElseFunction : ScalarFunction {
+ using ScalarFunction::ScalarFunction;
+
+ Result<const Kernel*> DispatchBest(std::vector<ValueDescr>* values) const
override {
+ RETURN_NOT_OK(CheckArity(*values));
+
+ // if-else 0'th descriptor is bool, so skip it
+ std::vector<ValueDescr> values_copy(values->begin() + 1, values->end());
+
+ using arrow::compute::detail::DispatchExactImpl;
+ if (auto kernel = DispatchExactImpl(this, *values)) return kernel;
+
+ internal::EnsureDictionaryDecoded(&values_copy);
+ internal::ReplaceNullWithOtherType(&values_copy);
+
+ if (auto type = internal::CommonNumeric(values_copy)) {
+ internal::ReplaceTypes(type, &values_copy);
+ }
+
+ std::copy(values_copy.begin(), values_copy.end(), values->begin() + 1);
Review comment:
```suggestion
std::move(values_copy.begin(), values_copy.end(), values->begin() + 1);
```
##########
File path: cpp/src/arrow/compute/kernels/scalar_if_else.cc
##########
@@ -543,7 +543,34 @@ struct ResolveIfElseExec {
}
};
-void AddPrimitiveIfElseKernels(const std::shared_ptr<ScalarFunction>&
scalar_function,
+struct IfElseFunction : ScalarFunction {
+ using ScalarFunction::ScalarFunction;
+
+ Result<const Kernel*> DispatchBest(std::vector<ValueDescr>* values) const
override {
+ RETURN_NOT_OK(CheckArity(*values));
+
+ // if-else 0'th descriptor is bool, so skip it
+ std::vector<ValueDescr> values_copy(values->begin() + 1, values->end());
+
+ using arrow::compute::detail::DispatchExactImpl;
+ if (auto kernel = DispatchExactImpl(this, *values)) return kernel;
+
Review comment:
If `values` is an exact match for a kernel, there's no need to construct
`values_copy`
```suggestion
using arrow::compute::detail::DispatchExactImpl;
if (auto kernel = DispatchExactImpl(this, *values)) return kernel;
// if-else 0'th descriptor is bool, so skip it
std::vector<ValueDescr> values_copy(values->begin() + 1, values->end());
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
