pitrou commented on a change in pull request #10987:
URL: https://github.com/apache/arrow/pull/10987#discussion_r695841734
##########
File path: cpp/src/arrow/compute/kernels/scalar_if_else.cc
##########
@@ -1827,9 +1827,131 @@ Status ExecArrayCoalesce(KernelContext* ctx, const
ExecBatch& batch, Datum* out)
return Status::OK();
}
+// Special case: implement 'coalesce' for an array and a scalar for any
+// fixed-width type (a 'fill_null' operation)
+template <typename Type>
+Status ExecArrayScalarCoalesce(KernelContext* ctx, Datum left, Datum right,
+ int64_t length, Datum* out) {
+ ArrayData* output = out->mutable_array();
+ const int64_t out_offset = output->offset;
+ uint8_t* out_valid = output->buffers[0]->mutable_data();
+ uint8_t* out_values = output->buffers[1]->mutable_data();
+
+ const ArrayData& left_arr = *left.array();
+ const uint8_t* left_valid =
+ left_arr.MayHaveNulls() ? left_arr.buffers[0]->data() : nullptr;
+ arrow::internal::OptionalBitBlockCounter bit_counter(left_valid,
left_arr.offset,
Review comment:
Have you tried using `BitRunReader` instead?
(er, perhaps we already had this conversation? :-S)
##########
File path: cpp/src/arrow/compute/kernels/scalar_if_else.cc
##########
@@ -1827,9 +1827,131 @@ Status ExecArrayCoalesce(KernelContext* ctx, const
ExecBatch& batch, Datum* out)
return Status::OK();
}
+// Special case: implement 'coalesce' for an array and a scalar for any
+// fixed-width type (a 'fill_null' operation)
+template <typename Type>
+Status ExecArrayScalarCoalesce(KernelContext* ctx, Datum left, Datum right,
+ int64_t length, Datum* out) {
+ ArrayData* output = out->mutable_array();
+ const int64_t out_offset = output->offset;
+ uint8_t* out_valid = output->buffers[0]->mutable_data();
+ uint8_t* out_values = output->buffers[1]->mutable_data();
+
+ const ArrayData& left_arr = *left.array();
+ const uint8_t* left_valid =
+ left_arr.MayHaveNulls() ? left_arr.buffers[0]->data() : nullptr;
+ arrow::internal::OptionalBitBlockCounter bit_counter(left_valid,
left_arr.offset,
+ left_arr.length);
+ int64_t offset = 0;
+
+ const uint8_t* left_values = left_arr.buffers[1]->data();
+ const Scalar& right_scalar = *right.scalar();
+ while (offset < length) {
+ const auto block = bit_counter.NextWord();
+ if (block.AllSet()) {
+ // All from left
+ CopyFixedWidth<Type>::CopyArray(*left_arr.type, left_values,
+ left_arr.offset + offset, block.length,
out_values,
+ out_offset + offset);
+ } else if (block.NoneSet()) {
+ // All from right
+ CopyFixedWidth<Type>::CopyScalar(right_scalar, block.length, out_values,
+ out_offset + offset);
+ } else if (block.popcount) {
+ // One by one
+ for (int64_t j = 0; j < block.length; ++j) {
+ if (BitUtil::GetBit(left_valid, left_arr.offset + offset + j)) {
+ CopyFixedWidth<Type>::CopyArray(
+ *left_arr.type, left_values, left_arr.offset + offset + j,
+ /*length=*/1, out_values, out_offset + offset + j);
+ } else {
+ CopyFixedWidth<Type>::CopyScalar(right_scalar, /*length=*/1,
out_values,
+ out_offset + offset + j);
+ }
+ }
+ }
+ offset += block.length;
+ }
+
+ if (right_scalar.is_valid || !left_valid) {
+ BitUtil::SetBitsTo(out_valid, out_offset, length, true);
+ } else {
+ arrow::internal::CopyBitmap(left_valid, left_arr.offset, length, out_valid,
+ out_offset);
+ }
+ return Status::OK();
+}
+
+// Special case: implement 'coalesce' for any 2 arguments for any fixed-width
+// type (a 'fill_null' operation)
+template <typename Type>
+Status ExecBinaryCoalesce(KernelContext* ctx, Datum left, Datum right, int64_t
length,
+ Datum* out) {
+ if (left.is_scalar() && right.is_scalar()) {
+ // Both scalar
+ *out = left.scalar()->is_valid ? left : right;
+ return Status::OK();
+ }
+
+ ArrayData* output = out->mutable_array();
+ const int64_t out_offset = output->offset;
+ uint8_t* out_valid = output->buffers[0]->mutable_data();
+ uint8_t* out_values = output->buffers[1]->mutable_data();
+
+ if (left.is_scalar()) {
+ // LHS is scalar
+ CopyValues<Type>(left.scalar()->is_valid ? left : right, /*in_offset=*/0,
length,
+ out_valid, out_values, out_offset);
+ return Status::OK();
+ }
+
+ // Array with nulls
+ const ArrayData& left_arr = *left.array();
+ const uint8_t* left_valid =
+ left_arr.MayHaveNulls() ? left_arr.buffers[0]->data() : nullptr;
+ arrow::internal::OptionalBitBlockCounter bit_counter(left_valid,
left_arr.offset,
+ left_arr.length);
+ int64_t offset = 0;
+
+ if (right.is_scalar()) {
Review comment:
Shouldn't this be moved up before `// Array with nulls`?
##########
File path: cpp/src/arrow/compute/kernels/scalar_if_else.cc
##########
@@ -1827,9 +1827,131 @@ Status ExecArrayCoalesce(KernelContext* ctx, const
ExecBatch& batch, Datum* out)
return Status::OK();
}
+// Special case: implement 'coalesce' for an array and a scalar for any
+// fixed-width type (a 'fill_null' operation)
+template <typename Type>
+Status ExecArrayScalarCoalesce(KernelContext* ctx, Datum left, Datum right,
+ int64_t length, Datum* out) {
+ ArrayData* output = out->mutable_array();
+ const int64_t out_offset = output->offset;
+ uint8_t* out_valid = output->buffers[0]->mutable_data();
+ uint8_t* out_values = output->buffers[1]->mutable_data();
+
+ const ArrayData& left_arr = *left.array();
+ const uint8_t* left_valid =
+ left_arr.MayHaveNulls() ? left_arr.buffers[0]->data() : nullptr;
+ arrow::internal::OptionalBitBlockCounter bit_counter(left_valid,
left_arr.offset,
+ left_arr.length);
+ int64_t offset = 0;
+
+ const uint8_t* left_values = left_arr.buffers[1]->data();
+ const Scalar& right_scalar = *right.scalar();
+ while (offset < length) {
+ const auto block = bit_counter.NextWord();
+ if (block.AllSet()) {
+ // All from left
+ CopyFixedWidth<Type>::CopyArray(*left_arr.type, left_values,
+ left_arr.offset + offset, block.length,
out_values,
+ out_offset + offset);
+ } else if (block.NoneSet()) {
+ // All from right
+ CopyFixedWidth<Type>::CopyScalar(right_scalar, block.length, out_values,
+ out_offset + offset);
+ } else if (block.popcount) {
+ // One by one
+ for (int64_t j = 0; j < block.length; ++j) {
+ if (BitUtil::GetBit(left_valid, left_arr.offset + offset + j)) {
+ CopyFixedWidth<Type>::CopyArray(
+ *left_arr.type, left_values, left_arr.offset + offset + j,
+ /*length=*/1, out_values, out_offset + offset + j);
+ } else {
+ CopyFixedWidth<Type>::CopyScalar(right_scalar, /*length=*/1,
out_values,
+ out_offset + offset + j);
+ }
+ }
+ }
+ offset += block.length;
+ }
+
+ if (right_scalar.is_valid || !left_valid) {
+ BitUtil::SetBitsTo(out_valid, out_offset, length, true);
+ } else {
+ arrow::internal::CopyBitmap(left_valid, left_arr.offset, length, out_valid,
+ out_offset);
+ }
+ return Status::OK();
+}
+
+// Special case: implement 'coalesce' for any 2 arguments for any fixed-width
+// type (a 'fill_null' operation)
+template <typename Type>
+Status ExecBinaryCoalesce(KernelContext* ctx, Datum left, Datum right, int64_t
length,
+ Datum* out) {
+ if (left.is_scalar() && right.is_scalar()) {
+ // Both scalar
+ *out = left.scalar()->is_valid ? left : right;
+ return Status::OK();
+ }
+
+ ArrayData* output = out->mutable_array();
+ const int64_t out_offset = output->offset;
+ uint8_t* out_valid = output->buffers[0]->mutable_data();
+ uint8_t* out_values = output->buffers[1]->mutable_data();
+
+ if (left.is_scalar()) {
+ // LHS is scalar
+ CopyValues<Type>(left.scalar()->is_valid ? left : right, /*in_offset=*/0,
length,
+ out_valid, out_values, out_offset);
+ return Status::OK();
+ }
+
+ // Array with nulls
+ const ArrayData& left_arr = *left.array();
+ const uint8_t* left_valid =
+ left_arr.MayHaveNulls() ? left_arr.buffers[0]->data() : nullptr;
Review comment:
Is it possible for `left_valid` to be null? If there are no nulls in the
left argument, I would expect this to shortcut.
##########
File path: cpp/src/arrow/compute/kernels/scalar_if_else.cc
##########
@@ -1827,9 +1827,131 @@ Status ExecArrayCoalesce(KernelContext* ctx, const
ExecBatch& batch, Datum* out)
return Status::OK();
}
+// Special case: implement 'coalesce' for an array and a scalar for any
+// fixed-width type (a 'fill_null' operation)
+template <typename Type>
+Status ExecArrayScalarCoalesce(KernelContext* ctx, Datum left, Datum right,
+ int64_t length, Datum* out) {
+ ArrayData* output = out->mutable_array();
+ const int64_t out_offset = output->offset;
+ uint8_t* out_valid = output->buffers[0]->mutable_data();
+ uint8_t* out_values = output->buffers[1]->mutable_data();
+
+ const ArrayData& left_arr = *left.array();
+ const uint8_t* left_valid =
+ left_arr.MayHaveNulls() ? left_arr.buffers[0]->data() : nullptr;
Review comment:
Is it possible for `left_valid` to be null?
##########
File path: cpp/src/arrow/compute/kernels/scalar_if_else_benchmark.cc
##########
@@ -386,8 +441,13 @@ BENCHMARK(CaseWhenBenchString)->Args({kFewItems, 99});
BENCHMARK(CaseWhenBenchStringContiguous)->Args({kFewItems, 0});
BENCHMARK(CaseWhenBenchStringContiguous)->Args({kFewItems, 99});
-BENCHMARK(CoalesceBench64)->Args({kNumItems, 0});
-BENCHMARK(CoalesceBench64)->Args({kNumItems, 99});
+BENCHMARK(CoalesceBench64)->Args({kNumItems, 0, 2});
+BENCHMARK(CoalesceBench64)->Args({kNumItems, 0, 4});
+BENCHMARK(CoalesceBench64)->Args({kNumItems, 99, 2});
+BENCHMARK(CoalesceBench64)->Args({kNumItems, 99, 4});
+
+BENCHMARK(CoalesceScalarBench64)->Args({kNumItems, 0});
+BENCHMARK(CoalesceScalarStringBench)->Args({kNumItems, 0});
BENCHMARK(CoalesceNonNullBench64)->Args({kNumItems, 0});
BENCHMARK(CoalesceNonNullBench64)->Args({kNumItems, 99});
Review comment:
It is possible to also parametrize the null_probability used in the
coalesce benchmarks? I think we have some logic to do this (see e.g.
RegressionArgs and the Take benchmarks).
--
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]
