GH-51127: [C++] Fix is_null nan_is_null for dictionary-encoded float arrays

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<const DictionaryType&>(*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<const DictionaryType&>(*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();
 }

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 <typename ArrowType>
 class TestFloatingPointValidityKernels : public TestValidityKernels<ArrowType> {
  public: