diff --git a/cpp/src/arrow/compute/kernels/scalar_validity.cc b/cpp/src/arrow/compute/kernels/scalar_validity.cc index 5913b756f1c..1b6eac717da 100644 --- a/cpp/src/arrow/compute/kernels/scalar_validity.cc +++ b/cpp/src/arrow/compute/kernels/scalar_validity.cc @@ -21,8 +21,10 @@ #include "arrow/compute/kernels/common_internal.h" #include "arrow/compute/registry_internal.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 +103,64 @@ static void SetNanBits(const ArraySpan& arr, uint8_t* out_bitmap, int64_t out_of } } +template +static void SetNanBitsDictionary(const ArraySpan& arr, const ArraySpan& dict_span, + uint8_t* out_bitmap, int64_t out_offset) { + const IndexType* indices = arr.GetValues(1); + const ValueType* dict_values = dict_span.GetValues(1); + for (int64_t i = 0; i < arr.length; ++i) { + if (arr.IsNull(i)) { + continue; + } + auto dict_index = indices[i]; + bool is_nan; + if constexpr (std::is_same_v) { + is_nan = Float16::FromBits(dict_values[dict_index]).is_nan(); + } else { + is_nan = std::isnan(dict_values[dict_index]); + } + if (is_nan) { + bit_util::SetBit(out_bitmap, i + out_offset); + } + } +} + +template +static void DispatchIndexType(const ArraySpan& arr, const ArraySpan& dict_span, + uint8_t* out_bitmap, int64_t out_offset) { + const auto& dict_type = checked_cast(*arr.type); + switch (dict_type.index_type()->id()) { + case Type::INT8: + SetNanBitsDictionary(arr, dict_span, out_bitmap, out_offset); + break; + case Type::INT16: + SetNanBitsDictionary(arr, dict_span, out_bitmap, out_offset); + break; + case Type::INT32: + SetNanBitsDictionary(arr, dict_span, out_bitmap, out_offset); + break; + case Type::INT64: + SetNanBitsDictionary(arr, dict_span, out_bitmap, out_offset); + break; + case Type::UINT8: + SetNanBitsDictionary(arr, dict_span, out_bitmap, out_offset); + break; + case Type::UINT16: + SetNanBitsDictionary(arr, dict_span, out_bitmap, out_offset); + break; + case Type::UINT32: + SetNanBitsDictionary(arr, dict_span, out_bitmap, out_offset); + break; + case Type::UINT64: + SetNanBitsDictionary(arr, dict_span, out_bitmap, out_offset); + break; + default: + DCHECK(false) << "unreachable: unsupported dictionary index type " + << dict_type.index_type()->ToString(); + break; + } +} + Status IsNullExec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { const ArraySpan& arr = batch[0].array; ArraySpan* out_span = out->array_span_mutable(); @@ -136,6 +196,24 @@ 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())) { + const ArraySpan& dict_span = arr.dictionary(); + switch (dict_type.value_type()->id()) { + case Type::FLOAT: + DispatchIndexType(arr, dict_span, out_bitmap, out_span->offset); + break; + case Type::DOUBLE: + DispatchIndexType(arr, dict_span, out_bitmap, out_span->offset); + break; + case Type::HALF_FLOAT: + DispatchIndexType(arr, dict_span, out_bitmap, out_span->offset); + break; + default: + break; + } + } } 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..f2e5bc874c1 100644 --- a/cpp/src/arrow/compute/kernels/scalar_validity_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_validity_test.cc @@ -152,6 +152,37 @@ 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, 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: