diff --git a/cpp/src/arrow/compute/kernels/scalar_validity.cc b/cpp/src/arrow/compute/kernels/scalar_validity.cc index c995b8dff82f..0829af59dc1f 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/dict_util_internal.h" #include "arrow/util/float16.h" #include "arrow/util/logging_internal.h" @@ -80,10 +83,42 @@ 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(); + } + 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(*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())); + + const ArrayData& result = *taken.array(); + if (arr.GetNullCount() > 0) { + ::arrow::internal::BitmapOrNot(result.buffers[1]->data(), result.offset, + arr.buffers[0].data, arr.offset, arr.length, + out_offset, out_bitmap); + } else { + CopyBitmap(result.buffers[1]->data(), result.offset, arr.length, out_bitmap, + out_offset); + } + return Status::OK(); +} + // `nan_is_null` can only be true for the is_null kernel since the is_valid and // true_unless_null kernels currently do not take `NullOptions` -Status SetLogicalNullBits(const ArraySpan& span, uint8_t* out_bitmap, int64_t out_offset, - bool set_on_null, bool nan_is_null) { +Status SetLogicalNullBits(KernelContext* ctx, const ArraySpan& span, uint8_t* out_bitmap, + int64_t out_offset, bool set_on_null, bool nan_is_null) { const Type::type t = span.type->id(); if (t == Type::NA) { // Input is all nulls, so all output bits are the same. @@ -98,7 +133,10 @@ Status SetLogicalNullBits(const ArraySpan& span, uint8_t* out_bitmap, int64_t ou // TODO: propagate `nan_is_null` ree_util::SetLogicalNullBits(span, out_bitmap, out_offset, set_on_null); } else if (t == Type::DICTIONARY) { - // TODO: propagate `nan_is_null` + const auto& dict_type = checked_cast(*span.type); + if (nan_is_null && is_floating(dict_type.value_type()->id())) { + return SetNanBitsFromDictionary(ctx, span, out_bitmap, out_offset); + } dict_util::SetLogicalNullBits(span, out_bitmap, out_offset, set_on_null); } else { // Input is a type for which logical and physical nulls are the same, so we can @@ -141,13 +179,15 @@ Status SetLogicalNullBits(const ArraySpan& span, uint8_t* out_bitmap, int64_t ou Status IsValidExec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { ArraySpan* out_span = out->array_span_mutable(); - return SetLogicalNullBits(batch[0].array, out_span->buffers[1].data, out_span->offset, + return SetLogicalNullBits(ctx, batch[0].array, out_span->buffers[1].data, + out_span->offset, /*set_on_null=*/false, /*nan_is_null=*/false); } Status IsNullExec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { ArraySpan* out_span = out->array_span_mutable(); - return SetLogicalNullBits(batch[0].array, out_span->buffers[1].data, out_span->offset, + return SetLogicalNullBits(ctx, batch[0].array, out_span->buffers[1].data, + out_span->offset, /*set_on_null=*/true, /*nan_is_null=*/NanOptionsState::Get(ctx).nan_is_null); } @@ -265,7 +305,8 @@ Status TrueUnlessNullExec(KernelContext* ctx, const ExecSpan& batch, ExecResult* // NullHandling::INTERSECTION and change the validity checks in exec.cc so that // they correctly handle logical nulls, but that would invove significant changes // in exec.cc which might have more side effects - return SetLogicalNullBits(batch[0].array, out_span->buffers[0].data, out_span->offset, + return SetLogicalNullBits(ctx, batch[0].array, out_span->buffers[0].data, + out_span->offset, /*set_on_null=*/false, /*nan_is_null=*/false); } diff --git a/cpp/src/arrow/compute/kernels/scalar_validity_test.cc b/cpp/src/arrow/compute/kernels/scalar_validity_test.cc index 86895cdcec53..0a041e7e47f7 100644 --- a/cpp/src/arrow/compute/kernels/scalar_validity_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_validity_test.cc @@ -208,6 +208,39 @@ 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); + + for (const auto& value_type : {float16(), float32(), float64()}) { + SCOPED_TRACE(value_type->ToString()); + auto dict_ty = dictionary(int32(), value_type); + auto arr = + DictArrayFromJSON(dict_ty, "[0, 1, 2, 3, null, 1]", "[1.5, NaN, -0.0, null]"); + + // Null dictionary values and null indices are always null. + CheckScalarUnary( + "is_null", arr, + ArrayFromJSON(boolean(), "[false, false, false, true, true, false]")); + CheckScalarUnary("is_null", arr, + ArrayFromJSON(boolean(), "[false, false, false, true, true, false]"), + &default_options); + CheckScalarUnary("is_null", arr, + ArrayFromJSON(boolean(), "[false, true, false, true, true, true]"), + &nan_is_null_options); + + auto no_null_indices = + DictArrayFromJSON(dict_ty, "[3, 1, 0, 2]", "[1.5, NaN, -0.0, null]"); + CheckScalarUnary("is_null", no_null_indices, + ArrayFromJSON(boolean(), "[true, true, false, false]"), + &nan_is_null_options); + + auto empty = DictArrayFromJSON(dict_ty, "[]", "[]"); + CheckScalarUnary("is_null", empty, ArrayFromJSON(boolean(), "[]"), + &nan_is_null_options); + } +} + template class TestFloatingPointValidityKernels : public TestValidityKernels { public: