pitrou commented on code in PR #51000:
URL: https://github.com/apache/arrow/pull/51000#discussion_r4122466614
##########
cpp/src/arrow/compute/kernels/scalar_validity.cc:
##########
@@ -80,10 +83,43 @@ static void SetNanBits(const ArraySpan& arr, uint8_t*
out_bitmap, int64_t out_of
}
}
+// Maps `is_null` over dictionary values and then through the indices, so
+// both NaN and null dictionary entries are reported, whatever the index type.
+static Status SetNanBitsFromDictionary(KernelContext* ctx, const ArraySpan&
arr,
+ uint8_t* out_bitmap, int64_t
out_offset) {
+ if (arr.length == 0) {
+ return Status::OK();
+ }
+ if (arr.GetNullCount() > 0) {
+ InvertBitmap(arr.buffers[0].data, arr.offset, arr.length, out_bitmap,
out_offset);
+ } else {
+ bit_util::SetBitsTo(out_bitmap, out_offset, arr.length, false);
+ }
+ NullOptions nan_is_null_options(/*nan_is_null=*/true);
+ ARROW_ASSIGN_OR_RAISE(Datum dict_is_null,
+ CallFunction("is_null",
{arr.dictionary().ToArrayData()},
+ &nan_is_null_options,
ctx->exec_context()));
+
+ const auto& dict_type = checked_cast<const DictionaryType&>(*arr.type);
+ auto indices = ArrayData::Make(dict_type.index_type(), arr.length,
+ {arr.GetBuffer(0), arr.GetBuffer(1)},
arr.GetNullCount(),
+ arr.offset);
+ ARROW_ASSIGN_OR_RAISE(Datum taken,
+ Take(dict_is_null, Datum(std::move(indices)),
+ TakeOptions::NoBoundsCheck(),
ctx->exec_context()));
+
+ // Null index slots are already set from the input validity bitmap, so the
+ // values bitmap can be OR'ed in without masking null slots out first.
+ const ArrayData& result = *taken.array();
+ ::arrow::internal::BitmapOr(out_bitmap, out_offset,
result.buffers[1]->data(),
Review Comment:
Oops, I forgot one point here. You're calling `Invert` then `BitmapOr`, but
`BitmapOrNot` would do faster in just one step.
Also, if `arr.GetNullCount() == 0`, no logical operation is required and we
can just copy the bitmap returned by `Take`.
--
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]
