pitrou commented on code in PR #51000:
URL: https://github.com/apache/arrow/pull/51000#discussion_r4120695707
##########
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 checked 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::BoundsCheck(), ctx->exec_context()));
Review Comment:
Bounds checking is not necessary since the dictionary array is presumed
valid, and therefore dict indices are within bounds already.
##########
cpp/src/arrow/compute/kernels/scalar_validity_test.cc:
##########
@@ -208,6 +208,91 @@ TEST(TestValidityKernels, IsNullSetsZeroNullCount) {
ASSERT_EQ(out.array()->null_count, 0);
}
+TEST(TestValidityKernels, IsNullDictionaryNanIsNull) {
+ NullOptions default_options;
+ NullOptions nan_is_null_options(/*nan_is_null=*/true);
+
+ auto dict_ty = dictionary(int32(), float64());
+ auto arr = DictArrayFromJSON(dict_ty, "[0, 1, 2, null, 1]", "[1.5, NaN,
-0.0]");
+
+ // Without nan_is_null, dictionary-encoded NaNs are not treated as null.
+ CheckScalarUnary("is_null", arr,
+ ArrayFromJSON(boolean(), "[false, false, false, true,
false]"));
+ CheckScalarUnary("is_null", arr,
+ ArrayFromJSON(boolean(), "[false, false, false, true,
false]"),
+ &default_options);
+
+ // With nan_is_null, the dictionary entry backing index 1 is NaN, so every
+ // slot referencing it is null; the pre-existing null index stays null.
+ CheckScalarUnary("is_null", arr,
+ ArrayFromJSON(boolean(), "[false, true, false, true,
true]"),
+ &nan_is_null_options);
+}
+
+TEST(TestValidityKernels, IsNullDictionaryNullValues) {
+ NullOptions default_options;
+ NullOptions nan_is_null_options(/*nan_is_null=*/true);
+
+ auto dict_ty = dictionary(int32(), float64());
+ auto arr = DictArrayFromJSON(dict_ty, "[0, 1, 2, null]", "[1.5, null, NaN]");
+
+ // A null dictionary value is reported regardless of nan_is_null: index 1
points at
+ // a null dictionary entry, so it is a logical null even with default
options.
+ CheckScalarUnary("is_null", arr, ArrayFromJSON(boolean(), "[false, true,
false, true]"),
+ &default_options);
+ CheckScalarUnary("is_null", arr, ArrayFromJSON(boolean(), "[false, true,
true, true]"),
+ &nan_is_null_options);
+}
+
+TEST(TestValidityKernels, IsNullDictionaryNanIsNullUnsignedIndices) {
+ NullOptions nan_is_null_options(/*nan_is_null=*/true);
+
+ auto dict_ty = dictionary(uint8(), float32());
+ auto arr = DictArrayFromJSON(dict_ty, "[2, 0, 1]", "[1.5, NaN, 2.5]");
+
+ CheckScalarUnary("is_null", arr, ArrayFromJSON(boolean(), "[false, false,
true]"),
+ &nan_is_null_options);
+}
+
+TEST(TestValidityKernels, IsNullDictionaryNanIsNullBounds) {
Review Comment:
This test is creating an invalid Arrow array. Compute functions are only
specified for valid Arrow arrays. Let's please remove it.
##########
cpp/src/arrow/compute/kernels/scalar_validity_test.cc:
##########
@@ -208,6 +208,91 @@ TEST(TestValidityKernels, IsNullSetsZeroNullCount) {
ASSERT_EQ(out.array()->null_count, 0);
}
+TEST(TestValidityKernels, IsNullDictionaryNanIsNull) {
+ NullOptions default_options;
+ NullOptions nan_is_null_options(/*nan_is_null=*/true);
+
+ auto dict_ty = dictionary(int32(), float64());
+ auto arr = DictArrayFromJSON(dict_ty, "[0, 1, 2, null, 1]", "[1.5, NaN,
-0.0]");
+
+ // Without nan_is_null, dictionary-encoded NaNs are not treated as null.
+ CheckScalarUnary("is_null", arr,
+ ArrayFromJSON(boolean(), "[false, false, false, true,
false]"));
+ CheckScalarUnary("is_null", arr,
+ ArrayFromJSON(boolean(), "[false, false, false, true,
false]"),
+ &default_options);
+
+ // With nan_is_null, the dictionary entry backing index 1 is NaN, so every
+ // slot referencing it is null; the pre-existing null index stays null.
+ CheckScalarUnary("is_null", arr,
+ ArrayFromJSON(boolean(), "[false, true, false, true,
true]"),
+ &nan_is_null_options);
+}
+
+TEST(TestValidityKernels, IsNullDictionaryNullValues) {
Review Comment:
Can you merge this test with the previous one?
##########
cpp/src/arrow/compute/kernels/scalar_validity_test.cc:
##########
@@ -208,6 +208,91 @@ TEST(TestValidityKernels, IsNullSetsZeroNullCount) {
ASSERT_EQ(out.array()->null_count, 0);
}
+TEST(TestValidityKernels, IsNullDictionaryNanIsNull) {
+ NullOptions default_options;
+ NullOptions nan_is_null_options(/*nan_is_null=*/true);
+
+ auto dict_ty = dictionary(int32(), float64());
+ auto arr = DictArrayFromJSON(dict_ty, "[0, 1, 2, null, 1]", "[1.5, NaN,
-0.0]");
+
+ // Without nan_is_null, dictionary-encoded NaNs are not treated as null.
+ CheckScalarUnary("is_null", arr,
+ ArrayFromJSON(boolean(), "[false, false, false, true,
false]"));
+ CheckScalarUnary("is_null", arr,
+ ArrayFromJSON(boolean(), "[false, false, false, true,
false]"),
+ &default_options);
+
+ // With nan_is_null, the dictionary entry backing index 1 is NaN, so every
+ // slot referencing it is null; the pre-existing null index stays null.
+ CheckScalarUnary("is_null", arr,
+ ArrayFromJSON(boolean(), "[false, true, false, true,
true]"),
+ &nan_is_null_options);
+}
+
+TEST(TestValidityKernels, IsNullDictionaryNullValues) {
+ NullOptions default_options;
+ NullOptions nan_is_null_options(/*nan_is_null=*/true);
+
+ auto dict_ty = dictionary(int32(), float64());
+ auto arr = DictArrayFromJSON(dict_ty, "[0, 1, 2, null]", "[1.5, null, NaN]");
+
+ // A null dictionary value is reported regardless of nan_is_null: index 1
points at
+ // a null dictionary entry, so it is a logical null even with default
options.
+ CheckScalarUnary("is_null", arr, ArrayFromJSON(boolean(), "[false, true,
false, true]"),
+ &default_options);
+ CheckScalarUnary("is_null", arr, ArrayFromJSON(boolean(), "[false, true,
true, true]"),
+ &nan_is_null_options);
+}
+
+TEST(TestValidityKernels, IsNullDictionaryNanIsNullUnsignedIndices) {
+ NullOptions nan_is_null_options(/*nan_is_null=*/true);
+
+ auto dict_ty = dictionary(uint8(), float32());
+ auto arr = DictArrayFromJSON(dict_ty, "[2, 0, 1]", "[1.5, NaN, 2.5]");
+
+ CheckScalarUnary("is_null", arr, ArrayFromJSON(boolean(), "[false, false,
true]"),
+ &nan_is_null_options);
+}
+
+TEST(TestValidityKernels, IsNullDictionaryNanIsNullBounds) {
+ NullOptions options(/*nan_is_null=*/true);
+ auto dict_ty = dictionary(int32(), float64());
+ auto values = ArrayFromJSON(float64(), "[1.5, null, NaN]");
+ for (const auto* indices_json :
+ {"[0, -1]", "[0, 3]", "[0, -2147483648]", "[0, 2147483647]"}) {
+ SCOPED_TRACE(indices_json);
+ auto indices = ArrayFromJSON(int32(), indices_json);
+ auto arr = std::make_shared<DictionaryArray>(dict_ty, indices, values);
+ ASSERT_RAISES(IndexError, IsNull(arr, options));
+ }
+
+ auto empty = DictArrayFromJSON(dict_ty, "[]", "[]");
+ CheckScalarUnary("is_null", empty, ArrayFromJSON(boolean(), "[]"), &options);
+}
+
+TEST(TestValidityKernels, IsNullDictionaryNanIsNullSkipsNullIndex) {
Review Comment:
I don't think this test makes sense either.
##########
cpp/src/arrow/compute/kernels/scalar_validity_test.cc:
##########
@@ -208,6 +208,91 @@ TEST(TestValidityKernels, IsNullSetsZeroNullCount) {
ASSERT_EQ(out.array()->null_count, 0);
}
+TEST(TestValidityKernels, IsNullDictionaryNanIsNull) {
+ NullOptions default_options;
+ NullOptions nan_is_null_options(/*nan_is_null=*/true);
+
+ auto dict_ty = dictionary(int32(), float64());
+ auto arr = DictArrayFromJSON(dict_ty, "[0, 1, 2, null, 1]", "[1.5, NaN,
-0.0]");
+
+ // Without nan_is_null, dictionary-encoded NaNs are not treated as null.
+ CheckScalarUnary("is_null", arr,
+ ArrayFromJSON(boolean(), "[false, false, false, true,
false]"));
+ CheckScalarUnary("is_null", arr,
+ ArrayFromJSON(boolean(), "[false, false, false, true,
false]"),
+ &default_options);
+
+ // With nan_is_null, the dictionary entry backing index 1 is NaN, so every
+ // slot referencing it is null; the pre-existing null index stays null.
+ CheckScalarUnary("is_null", arr,
+ ArrayFromJSON(boolean(), "[false, true, false, true,
true]"),
+ &nan_is_null_options);
+}
+
+TEST(TestValidityKernels, IsNullDictionaryNullValues) {
+ NullOptions default_options;
+ NullOptions nan_is_null_options(/*nan_is_null=*/true);
+
+ auto dict_ty = dictionary(int32(), float64());
+ auto arr = DictArrayFromJSON(dict_ty, "[0, 1, 2, null]", "[1.5, null, NaN]");
+
+ // A null dictionary value is reported regardless of nan_is_null: index 1
points at
+ // a null dictionary entry, so it is a logical null even with default
options.
+ CheckScalarUnary("is_null", arr, ArrayFromJSON(boolean(), "[false, true,
false, true]"),
+ &default_options);
+ CheckScalarUnary("is_null", arr, ArrayFromJSON(boolean(), "[false, true,
true, true]"),
+ &nan_is_null_options);
+}
+
+TEST(TestValidityKernels, IsNullDictionaryNanIsNullUnsignedIndices) {
+ NullOptions nan_is_null_options(/*nan_is_null=*/true);
+
+ auto dict_ty = dictionary(uint8(), float32());
+ auto arr = DictArrayFromJSON(dict_ty, "[2, 0, 1]", "[1.5, NaN, 2.5]");
+
+ CheckScalarUnary("is_null", arr, ArrayFromJSON(boolean(), "[false, false,
true]"),
+ &nan_is_null_options);
+}
+
+TEST(TestValidityKernels, IsNullDictionaryNanIsNullBounds) {
+ NullOptions options(/*nan_is_null=*/true);
+ auto dict_ty = dictionary(int32(), float64());
+ auto values = ArrayFromJSON(float64(), "[1.5, null, NaN]");
+ for (const auto* indices_json :
+ {"[0, -1]", "[0, 3]", "[0, -2147483648]", "[0, 2147483647]"}) {
+ SCOPED_TRACE(indices_json);
+ auto indices = ArrayFromJSON(int32(), indices_json);
+ auto arr = std::make_shared<DictionaryArray>(dict_ty, indices, values);
+ ASSERT_RAISES(IndexError, IsNull(arr, options));
+ }
+
+ auto empty = DictArrayFromJSON(dict_ty, "[]", "[]");
+ CheckScalarUnary("is_null", empty, ArrayFromJSON(boolean(), "[]"), &options);
+}
+
+TEST(TestValidityKernels, IsNullDictionaryNanIsNullSkipsNullIndex) {
+ NullOptions options(/*nan_is_null=*/true);
+ auto dict_ty = dictionary(int32(), float64());
+ auto indices = ArrayFromJSON(int32(), "[0, null, 1]");
+ // The physical value in a null slot must not be dereferenced.
+ reinterpret_cast<int32_t*>(indices->data()->buffers[1]->mutable_data())[1] =
-1;
+ auto values = ArrayFromJSON(float64(), "[1.5, NaN]");
+ auto arr = std::make_shared<DictionaryArray>(dict_ty, indices, values);
+ ASSERT_OK(arr->ValidateFull());
+ CheckScalarUnary("is_null", arr, ArrayFromJSON(boolean(), "[false, true,
true]"),
+ &options);
+}
+
+TEST(TestValidityKernels, IsNullDictionaryNanIsNullHalfFloat) {
Review Comment:
Instead of a separate test for half-float, can you loop over all
floating-point types in the main test?
--
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]
