pitrou commented on code in PR #36825:
URL: https://github.com/apache/arrow/pull/36825#discussion_r1272124249
##########
cpp/src/arrow/compute/kernels/scalar_if_else_test.cc:
##########
@@ -126,75 +126,105 @@ TYPED_TEST(TestIfElsePrimitive, IfElseFixedSizeRand) {
CheckIfElseOutput(cond, left, right, expected_data);
}
-void CheckWithDifferentShapes(const std::shared_ptr<Array>& cond,
- const std::shared_ptr<Array>& left,
- const std::shared_ptr<Array>& right,
- const std::shared_ptr<Array>& expected) {
- // this will check for whole arrays, every scalar at i'th index and slicing
(offset)
- CheckScalar("if_else", {cond, left, right}, expected);
+Datum ArrayOrBroadcastScalar(const Datum& input, int64_t length) {
+ if (input.is_scalar()) {
+ EXPECT_OK_AND_ASSIGN(auto array, MakeArrayFromScalar(*input.scalar(),
length));
+ return array;
+ }
+ EXPECT_TRUE(input.is_array());
+ return input;
+}
- auto len = left->length();
- std::vector<int64_t> array_indices = {-1}; // sentinel for make_input
- std::vector<int64_t> scalar_indices(len);
- std::iota(scalar_indices.begin(), scalar_indices.end(), 0);
- auto make_input = [&](const std::shared_ptr<Array>& array, int64_t index,
Datum* input,
- Datum* input_broadcast, std::string* trace) {
- if (index >= 0) {
- // Use scalar from array[index] as input; broadcast scalar for computing
expected
- // result
- ASSERT_OK_AND_ASSIGN(auto scalar, array->GetScalar(index));
- *trace += "@" + std::to_string(index) + "=" + scalar->ToString();
- *input = std::move(scalar);
- ASSERT_OK_AND_ASSIGN(*input_broadcast,
MakeArrayFromScalar(*input->scalar(), len));
+Result<Datum> ExpectedFromIfElse(const Datum& cond, const Datum& left, const
Datum& right,
+ std::shared_ptr<DataType> type) {
+ if (cond.is_scalar() && left.is_scalar() && right.is_scalar()) {
+ const auto& scalar = cond.scalar_as<BooleanScalar>();
+ Datum expected;
+ if (scalar.is_valid) {
+ expected = scalar.value ? left : right;
} else {
- // Use array as input
- *trace += "=Array";
- *input = *input_broadcast = array;
+ expected = MakeNullScalar(left.type());
}
- };
-
- enum { COND_SCALAR = 1, LEFT_SCALAR = 2, RIGHT_SCALAR = 4 };
- for (int mask = 1; mask <= (COND_SCALAR | LEFT_SCALAR | RIGHT_SCALAR);
++mask) {
- for (int64_t cond_idx : (mask & COND_SCALAR) ? scalar_indices :
array_indices) {
- Datum cond_in, cond_bcast;
- std::string trace_cond = "Cond";
- make_input(cond, cond_idx, &cond_in, &cond_bcast, &trace_cond);
-
- for (int64_t left_idx : (mask & LEFT_SCALAR) ? scalar_indices :
array_indices) {
- Datum left_in, left_bcast;
- std::string trace_left = "Left";
- make_input(left, left_idx, &left_in, &left_bcast, &trace_left);
-
- for (int64_t right_idx : (mask & RIGHT_SCALAR) ? scalar_indices :
array_indices) {
- Datum right_in, right_bcast;
- std::string trace_right = "Right";
- make_input(right, right_idx, &right_in, &right_bcast, &trace_right);
-
- SCOPED_TRACE(trace_right);
- SCOPED_TRACE(trace_left);
- SCOPED_TRACE(trace_cond);
-
- Datum expected;
- ASSERT_OK_AND_ASSIGN(auto actual, IfElse(cond_in, left_in,
right_in));
- if (mask == (COND_SCALAR | LEFT_SCALAR | RIGHT_SCALAR)) {
- const auto& scalar = cond_in.scalar_as<BooleanScalar>();
- if (scalar.is_valid) {
- expected = scalar.value ? left_in : right_in;
- } else {
- expected = MakeNullScalar(left_in.type());
- }
- if (!left_in.type()->Equals(*right_in.type())) {
- ASSERT_OK_AND_ASSIGN(expected,
- Cast(expected,
CastOptions::Safe(actual.type())));
- }
- } else {
- ASSERT_OK_AND_ASSIGN(expected, IfElse(cond_bcast, left_bcast,
right_bcast));
- }
- AssertDatumsEqual(expected, actual, /*verbose=*/true);
+ if (!left.type()->Equals(*right.type())) {
+ return Cast(expected, CastOptions::Safe(std::move(type)));
+ }
+ return expected;
+ }
+ // When at least one of the inputs is an array, we expect the output
+ // to be the same as if all the scalars were broadcast to arrays.
Review Comment:
If all inputs are array, then the expected datum will be trivially equal to
the actual. Should we reuse the original `expected` argument in that case?
--
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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
