pitrou commented on code in PR #50647:
URL: https://github.com/apache/arrow/pull/50647#discussion_r3672065745
##########
cpp/src/arrow/compute/kernels/scalar_compare.cc:
##########
@@ -529,108 +492,246 @@ struct ScalarMinMax {
}
}
+ // Fold left/right into out_values/out_valid a word at a time (validity
bitmap
+ // null => all valid). out_values may alias left. Returns the null count
+ static int64_t CombineWordwise(const OutValue* left, const uint8_t*
left_valid,
+ int64_t left_offset, const OutValue* right,
+ const uint8_t* right_valid, int64_t
right_offset,
+ bool skip_nulls, int64_t length, OutValue*
out_values,
+ uint8_t* out_valid) {
+ auto left_reader = ::arrow::internal::BitmapUInt64Reader(
+ left_valid, left_valid ? left_offset : 0, left_valid ? length : 0);
+ auto right_reader = ::arrow::internal::BitmapUInt64Reader(
+ right_valid, right_valid ? right_offset : 0, right_valid ? length : 0);
+
+ int64_t null_count = 0;
+ int64_t i = 0;
+ for (int64_t words = length / 64; words > 0; --words) {
+ const uint64_t left_word = left_valid ? left_reader.NextWord() :
~uint64_t(0);
+ const uint64_t right_word = right_valid ? right_reader.NextWord() :
~uint64_t(0);
+ const uint64_t out_word =
+ skip_nulls ? (left_word | right_word) : (left_word & right_word);
+ // out_valid is allocated at bit offset 0, so store the word directly
+ const uint64_t out_word_le = bit_util::ToLittleEndian(out_word);
+ std::memcpy(out_valid + i / 8, &out_word_le, sizeof(out_word_le));
+ null_count += 64 - std::popcount(out_word);
+ if (out_word == 0) {
+ // All null: the values are never read, so skip them
+ i += 64;
+ } else if (left_word == ~uint64_t(0) && right_word == ~uint64_t(0)) {
+ // All valid: no per-lane validity check
+ for (int j = 0; j < 64; ++j, ++i) {
+ out_values[i] =
+ Op::template Call<OutValue, OutValue, OutValue>(left[i],
right[i]);
+ }
+ } else {
+ for (int j = 0; j < 64; ++j, ++i) {
+ out_values[i] =
+ CombineOne((left_word >> j) & 1, (right_word >> j) & 1, left[i],
right[i]);
+ }
+ }
+ }
+ for (int bit = 0; i < length; ++i, ++bit) {
+ const bool left_bit = !left_valid || bit_util::GetBit(left_valid,
left_offset + i);
+ const bool right_bit =
+ !right_valid || bit_util::GetBit(right_valid, right_offset + i);
+ const bool out_bit = skip_nulls ? (left_bit || right_bit) : (left_bit &&
right_bit);
+ bit_util::SetBitTo(out_valid, i, out_bit);
+ null_count += !out_bit;
+ if (out_bit) out_values[i] = CombineOne(left_bit, right_bit, left[i],
right[i]);
+ }
+ return null_count;
+ }
+
+ // Seed the accumulator from the first two arrays in a single pass
+ static Status CombineArrays(KernelContext* ctx, const ArraySpan& lhs,
+ const ArraySpan& rhs, bool skip_nulls,
ArrayData* output) {
+ const int64_t length = output->length;
+ const OutValue* left = lhs.GetValues<OutValue>(1);
+ const OutValue* right = rhs.GetValues<OutValue>(1);
+ OutValue* out_values = output->GetMutableValues<OutValue>(1);
+ const uint8_t* left_valid = lhs.MayHaveNulls() ? lhs.buffers[0].data :
nullptr;
+ const uint8_t* right_valid = rhs.MayHaveNulls() ? rhs.buffers[0].data :
nullptr;
+
+ if (!left_valid && !right_valid) {
+ for (int64_t i = 0; i < length; ++i) {
+ out_values[i] =
+ Op::template Call<OutValue, OutValue, OutValue>(left[i], right[i]);
+ }
+ output->buffers[0] = nullptr;
+ output->null_count = 0;
+ return Status::OK();
+ }
+
+ ARROW_ASSIGN_OR_RAISE(output->buffers[0], ctx->AllocateBitmap(length));
+ // CombineWordwise skips all-null words, so pre-zero the values to keep
those
+ // slots initialized for later reads
+ std::memset(out_values, 0, static_cast<size_t>(length) * sizeof(OutValue));
+ output->null_count = CombineWordwise(left, left_valid, lhs.offset, right,
right_valid,
+ rhs.offset, skip_nulls, length,
out_values,
+ output->buffers[0]->mutable_data());
+ if (output->null_count == 0) {
+ output->buffers[0] = nullptr;
+ }
+ return Status::OK();
+ }
+
+ // Seed the accumulator from a single array (result == that array)
+ static Status CopyArrayToOutput(KernelContext* ctx, const ArraySpan& arr,
+ ArrayData* output) {
+ const int64_t length = output->length;
+ const OutValue* arr_values = arr.GetValues<OutValue>(1);
+ OutValue* out_values = output->GetMutableValues<OutValue>(1);
+ std::copy(arr_values, arr_values + length, out_values);
+ if (arr.MayHaveNulls()) {
+ ARROW_ASSIGN_OR_RAISE(output->buffers[0], ctx->AllocateBitmap(length));
+ ::arrow::internal::CopyBitmap(arr.buffers[0].data, arr.offset, length,
+ output->buffers[0]->mutable_data(),
+ /*dest_offset=*/0);
+ output->null_count = arr.null_count;
+ } else {
+ output->buffers[0] = nullptr;
+ output->null_count = 0;
+ }
+ return Status::OK();
+ }
+
+ // Fold `arr` into the accumulator already in `output` (buffers[0] null =>
+ // all valid), in place, leaving a result that can be folded again
+ static Status FoldArrayIntoOutput(KernelContext* ctx, const ArraySpan& arr,
+ bool skip_nulls, ArrayData* output) {
+ const int64_t length = output->length;
+ OutValue* acc_values = output->GetMutableValues<OutValue>(1);
+ const OutValue* arr_values = arr.GetValues<OutValue>(1);
+ // Keep the current validity buffer alive while CombineWordwise writes a
new one
+ std::shared_ptr<Buffer> acc_valid_buf = output->buffers[0];
+ const uint8_t* acc_valid = acc_valid_buf ? acc_valid_buf->data() : nullptr;
+ const uint8_t* arr_valid = arr.MayHaveNulls() ? arr.buffers[0].data :
nullptr;
+
+ if (!acc_valid && !arr_valid) {
+ for (int64_t i = 0; i < length; ++i) {
+ acc_values[i] =
+ Op::template Call<OutValue, OutValue, OutValue>(acc_values[i],
arr_values[i]);
+ }
+ output->null_count = 0;
+ return Status::OK();
+ }
+
+ // Accumulator all valid with skip_nulls: the result stays all valid, so no
+ // output bitmap is needed. Where the array is null the accumulator already
+ // holds the right value
+ if (!acc_valid && skip_nulls) {
+ auto reader = ::arrow::internal::BitmapUInt64Reader(arr_valid,
arr.offset, length);
Review Comment:
I'd be tempted to go with `VisitSetBitRuns` under the assumption that most
inputs are either mostly-null or mostly-not-null. What do you think
@zanmato1984 ?
--
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]
