diff --git a/cpp/src/arrow/compute/kernels/scalar_validity.cc b/cpp/src/arrow/compute/kernels/scalar_validity.cc index 5913b756f1c..9abc6af258b 100644 --- a/cpp/src/arrow/compute/kernels/scalar_validity.cc +++ b/cpp/src/arrow/compute/kernels/scalar_validity.cc @@ -21,8 +21,11 @@ #include "arrow/compute/kernels/common_internal.h" #include "arrow/compute/registry_internal.h" +#include "arrow/compute/api_vector.h" +#include "arrow/type.h" #include "arrow/util/bit_util.h" #include "arrow/util/bitmap_ops.h" +#include "arrow/util/checked_cast.h" #include "arrow/util/float16.h" #include "arrow/util/logging_internal.h" @@ -101,6 +104,34 @@ static void SetNanBits(const ArraySpan& arr, uint8_t* out_bitmap, int64_t out_of } } +// Maps `is_null` over the dictionary values and then through the indices, so that +// both NaN and null dictionary entries are reported, whatever the index type. +static Status SetNullBitsFromDictionary(KernelContext* ctx, const ArraySpan& arr, + const NullOptions& options, uint8_t* out_bitmap, + int64_t out_offset) { + if (arr.length == 0) { + return Status::OK(); + } + const auto& dict_type = checked_cast(*arr.type); + ARROW_ASSIGN_OR_RAISE(Datum dict_is_null, + CallFunction("is_null", {arr.dictionary().ToArrayData()}, + &options, ctx->exec_context())); + + 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())); + + // Slots with a null index are already set from the input validity bitmap, so the + // values bitmap can be OR'ed in without masking the nulls out of it first. + const ArrayData& result = *taken.array(); + ::arrow::internal::BitmapOr(out_bitmap, out_offset, result.buffers[1]->data(), + result.offset, arr.length, out_offset, out_bitmap); + return Status::OK(); +} + Status IsNullExec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { const ArraySpan& arr = batch[0].array; ArraySpan* out_span = out->array_span_mutable(); @@ -136,6 +167,12 @@ Status IsNullExec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { return Status::NotImplemented("NaN detection not implemented for type ", arr.type->ToString()); } + } else if (arr.type->id() == Type::DICTIONARY && options.nan_is_null) { + const auto& dict_type = checked_cast(*arr.type); + if (is_floating(dict_type.value_type()->id())) { + RETURN_NOT_OK( + SetNullBitsFromDictionary(ctx, arr, options, out_bitmap, out_span->offset)); + } } return Status::OK(); } diff --git a/cpp/src/arrow/compute/kernels/scalar_validity_test.cc b/cpp/src/arrow/compute/kernels/scalar_validity_test.cc index 4613176b48c..a0872703773 100644 --- a/cpp/src/arrow/compute/kernels/scalar_validity_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_validity_test.cc @@ -152,6 +152,63 @@ 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 only reported through the nan_is_null path, which is + // where the dictionary is inspected at all. + CheckScalarUnary("is_null", arr, + ArrayFromJSON(boolean(), "[false, false, 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, IsNullDictionaryNanIsNullHalfFloat) { + NullOptions nan_is_null_options(/*nan_is_null=*/true); + + auto dict_ty = dictionary(int8(), float16()); + auto arr = DictArrayFromJSON(dict_ty, "[0, 1]", "[1.5, NaN]"); + + CheckScalarUnary("is_null", arr, ArrayFromJSON(boolean(), "[false, true]"), + &nan_is_null_options); +} + template class TestFloatingPointValidityKernels : public TestValidityKernels { public: