GH-50338: [C++] Add ComputeLogicalNullCount to Datum

cpp/src/arrow/datum.cc

@@ -142,7 +142,22 @@ int64_t Datum::null_count() const {
     const auto& val = *std::get<std::shared_ptr<Scalar>>(this->value);
     return val.is_valid ? 0 : 1;
   } else {
-    DCHECK(false) << "This function only valid for array-like values";
+    DCHECK(false) << "This function only valid for scalar or array-like values";
+    return 0;
+  }
+}
+
+int64_t Datum::ComputeLogicalNullCount() const {
+  if (this->kind() == Datum::ARRAY) {
+    return std::get<std::shared_ptr<ArrayData>>(this->value)->ComputeLogicalNullCount();
+  } else if (this->kind() == Datum::CHUNKED_ARRAY) {
+    return std::get<std::shared_ptr<ChunkedArray>>(this->value)
+        ->ComputeLogicalNullCount();
+  } else if (this->kind() == Datum::SCALAR) {
+    const auto& val = *std::get<std::shared_ptr<Scalar>>(this->value);
+    return val.IsLogicalNull() ? 1 : 0;
+  } else {
+    DCHECK(false) << "This function only valid for scalar or array-like values";
     return 0;
   }
 }

cpp/src/arrow/datum.h

@@ -276,6 +276,18 @@ struct ARROW_EXPORT Datum {
   /// Only valid for scalar and array-like data.
   int64_t null_count() const;
 
+  /// \brief Compute the logical null count.
+  ///
+  /// Only valid for scalar and array-like data. Unlike null_count(), this
+  /// accounts for types whose logical nulls are not captured by the top-level
+  /// validity bitmap, such as union, run-end encoded and dictionary types; for
+  /// those types the count is recomputed on every call.
+  ///
+  /// \see Scalar::IsLogicalNull
+  /// \see ArrayData::ComputeLogicalNullCount
+  /// \see ChunkedArray::ComputeLogicalNullCount
+  int64_t ComputeLogicalNullCount() const;
+
   /// \brief The value type of the variant, if any
   ///
   /// \return nullptr if no type

cpp/src/arrow/datum_test.cc

@@ -23,6 +23,8 @@
 
 #include "arrow/array/array_base.h"
 #include "arrow/array/array_binary.h"
+#include "arrow/array/builder_primitive.h"
+#include "arrow/array/builder_run_end.h"
 #include "arrow/chunked_array.h"
 #include "arrow/datum.h"
 #include "arrow/record_batch.h"
@@ -139,6 +141,89 @@ TEST(Datum, NullCount) {
   ASSERT_EQ(3, val3.null_count());
 }
 
+TEST(Datum, ComputeLogicalNullCount) {
+  // For most scalars, is_valid already reflects logical validity.
+  Datum valid_scalar(std::make_shared<Int8Scalar>(1));
+  ASSERT_EQ(0, valid_scalar.ComputeLogicalNullCount());
+
+  Datum null_scalar(MakeNullScalar(int8()));
+  ASSERT_EQ(1, null_scalar.ComputeLogicalNullCount());
+
+  // Arrays and scalars of type null() are entirely null.
+  Datum null_type_arr(ArrayFromJSON(null(), "[null, null, null]"));
+  ASSERT_EQ(3, null_type_arr.null_count());
+  ASSERT_EQ(3, null_type_arr.ComputeLogicalNullCount());
+
+  Datum null_type_scalar(MakeNullScalar(null()));
+  ASSERT_EQ(1, null_type_scalar.null_count());
+  ASSERT_EQ(1, null_type_scalar.ComputeLogicalNullCount());
+
+  // For arrays with a validity bitmap, the logical null count matches
+  // null_count().
+  Datum int8_arr(ArrayFromJSON(int8(), "[1, null, null, null]"));
+  ASSERT_EQ(3, int8_arr.null_count());
+  ASSERT_EQ(3, int8_arr.ComputeLogicalNullCount());
+
+  // Union arrays carry logical nulls in their children without a top-level
+  // validity bitmap, so null_count() is 0 while the logical null count is not.
+  auto union_type = sparse_union({field("a", int8()), field("b", boolean())});
+  auto union_arr =
+      ArrayFromJSON(union_type, R"([[0, null], [1, true], [0, 5], [1, null]])");
+  Datum union_datum(union_arr);
+  ASSERT_EQ(0, union_datum.null_count());
+  ASSERT_EQ(2, union_datum.ComputeLogicalNullCount());
+
+  // Scalars extracted from an array preserve its logical validity.
+  auto scalar_logical_null_count = [](const Array& arr,
+                                      int64_t index) -> Result<int64_t> {
+    ARROW_ASSIGN_OR_RAISE(auto scalar, arr.GetScalar(index));
+    return Datum(scalar).ComputeLogicalNullCount();
+  };
+  ASSERT_OK_AND_EQ(1, scalar_logical_null_count(*union_arr, 0));
+  ASSERT_OK_AND_EQ(0, scalar_logical_null_count(*union_arr, 1));
+
+  // Chunked arrays sum the logical null count over the chunks.
+  auto union_chunk = ArrayFromJSON(union_type, R"([[0, 1], [1, null]])");
+  ASSERT_OK_AND_ASSIGN(auto chunked, ChunkedArray::Make({union_arr, union_chunk}));
+  Datum chunked_datum(chunked);
+  ASSERT_EQ(0, chunked_datum.null_count());
+  ASSERT_EQ(3, chunked_datum.ComputeLogicalNullCount());
+
+  // Run-end encoded arrays carry logical nulls in their values child, also
+  // without a top-level validity bitmap.
+  auto pool = default_memory_pool();
+  auto ree_type = run_end_encoded(int32(), int32());
+  RunEndEncodedBuilder ree_builder(pool, std::make_shared<Int32Builder>(pool),
+                                   std::make_shared<Int32Builder>(pool), ree_type);
+  ASSERT_OK(ree_builder.AppendScalar(*MakeScalar<int32_t>(7), 2));
+  ASSERT_OK(ree_builder.AppendNulls(3));
+  ASSERT_OK_AND_ASSIGN(auto ree_arr, ree_builder.Finish());
+  Datum ree_datum(ree_arr);
+  ASSERT_EQ(0, ree_datum.null_count());
+  ASSERT_EQ(3, ree_datum.ComputeLogicalNullCount());
+  ASSERT_OK_AND_EQ(0, scalar_logical_null_count(*ree_arr, 0));
+  ASSERT_OK_AND_EQ(1, scalar_logical_null_count(*ree_arr, 2));
+
+  // Dictionary arrays have a validity bitmap on the indices, but a valid
+  // index referencing a null dictionary value is also a logical null.
+  auto dict_type = dictionary(int32(), utf8());
+  auto dict_arr = DictArrayFromJSON(dict_type, /*indices=*/"[0, 1, null, 1]",
+                                    /*dictionary=*/R"([null, "a"])");
+  Datum dict_datum(dict_arr);
+  ASSERT_EQ(1, dict_datum.null_count());
+  ASSERT_EQ(2, dict_datum.ComputeLogicalNullCount());
+
+  // A DictionaryScalar's is_valid only reflects index validity, but the
+  // logical null count also accounts for a valid index referencing a null
+  // dictionary value, consistently with the array path.
+  ASSERT_OK_AND_ASSIGN(auto dict_scalar, dict_arr->GetScalar(0));
+  Datum dict_scalar_datum(dict_scalar);
+  ASSERT_EQ(0, dict_scalar_datum.null_count());
+  ASSERT_EQ(1, dict_scalar_datum.ComputeLogicalNullCount());
+  ASSERT_OK_AND_EQ(0, scalar_logical_null_count(*dict_arr, 1));
+  ASSERT_OK_AND_EQ(1, scalar_logical_null_count(*dict_arr, 2));
+}
+
 TEST(Datum, MutableArray) {
   auto arr = ArrayFromJSON(int8(), "[1, 2, 3, 4]");
 

cpp/src/arrow/scalar.cc

@@ -770,54 +770,56 @@ DictionaryScalar::DictionaryScalar(std::shared_ptr<DataType> type)
                             0)
                 .ValueOrDie()} {}
 
+namespace {
+
+struct DictionaryIndexValueImpl {
+  int64_t value = 0;
+
+  Status Visit(const Scalar&) {
+    return Status::TypeError("Not implemented dictionary index type");
+  }
+
+  template <typename ScalarType, typename Type = typename ScalarType::TypeClass>
+  enable_if_integer<Type, Status> Visit(const ScalarType& scalar) {
+    value = static_cast<int64_t>(scalar.value);
+    return Status::OK();
+  }
+};
+
+Result<int64_t> DictionaryIndexValue(const Scalar& index) {
+  DictionaryIndexValueImpl impl;
+  RETURN_NOT_OK(VisitScalarInline(index, &impl));
+  return impl.value;
+}
+
+}  // namespace
+
 Result<std::shared_ptr<Scalar>> DictionaryScalar::GetEncodedValue() const {
   const auto& dict_type = checked_cast<DictionaryType&>(*type);
 
   if (!is_valid) {
     return MakeNullScalar(dict_type.value_type());
   }
 
-  int64_t index_value = 0;
-  switch (dict_type.index_type()->id()) {
-    case Type::UINT8:
-      index_value =
-          static_cast<int64_t>(checked_cast<const UInt8Scalar&>(*value.index).value);
-      break;
-    case Type::INT8:
-      index_value =
-          static_cast<int64_t>(checked_cast<const Int8Scalar&>(*value.index).value);
-      break;
-    case Type::UINT16:
-      index_value =
-          static_cast<int64_t>(checked_cast<const UInt16Scalar&>(*value.index).value);
-      break;
-    case Type::INT16:
-      index_value =
-          static_cast<int64_t>(checked_cast<const Int16Scalar&>(*value.index).value);
-      break;
-    case Type::UINT32:
-      index_value =
-          static_cast<int64_t>(checked_cast<const UInt32Scalar&>(*value.index).value);
-      break;
-    case Type::INT32:
-      index_value =
-          static_cast<int64_t>(checked_cast<const Int32Scalar&>(*value.index).value);
-      break;
-    case Type::UINT64:
-      index_value =
-          static_cast<int64_t>(checked_cast<const UInt64Scalar&>(*value.index).value);
-      break;
-    case Type::INT64:
-      index_value =
-          static_cast<int64_t>(checked_cast<const Int64Scalar&>(*value.index).value);
-      break;
-    default:
-      return Status::TypeError("Not implemented dictionary index type");
-      break;
-  }
+  ARROW_ASSIGN_OR_RAISE(int64_t index_value, DictionaryIndexValue(*value.index));
   return value.dictionary->GetScalar(index_value);
 }
 
+bool DictionaryScalar::IsLogicalNull() const {
+  if (!is_valid) {
+    return true;
+  }
+  const auto& dict = value.dictionary;
+  if (value.index == nullptr || dict == nullptr) {
+    return false;
+  }
+  auto index_value = DictionaryIndexValue(*value.index);
+  if (!index_value.ok() || *index_value < 0 || *index_value >= dict->length()) {
+    return false;
+  }
+  return dict->IsNull(*index_value);
+}
+
 std::shared_ptr<DictionaryScalar> DictionaryScalar::Make(std::shared_ptr<Scalar> index,
                                                          std::shared_ptr<Array> dict) {
   auto type = dictionary(index->type, dict->type());

cpp/src/arrow/scalar.h

@@ -59,8 +59,27 @@ struct ARROW_EXPORT Scalar : public std::enable_shared_from_this<Scalar>,
   std::shared_ptr<DataType> type;
 
   /// \brief Whether the value is valid (not null) or not
+  ///
+  /// Note that this may not reflect logical nullness for all types; see
+  /// IsLogicalNull().
   bool is_valid = false;
 
+  /// \brief Whether the value is logically null
+  ///
+  /// For most types this is simply !is_valid. For dictionary scalars, however,
+  /// is_valid only reflects the validity of the index: a valid index can still
+  /// refer to a null dictionary value, in which case the scalar is logically
+  /// null. For an ill-formed dictionary scalar whose dictionary value cannot
+  /// be resolved (for example an out-of-bounds index), this falls back to
+  /// !is_valid.
+  ///
+  /// Like the array-level logical null computation, this does not recurse into
+  /// nested values: a dictionary value wrapped in a union, run-end encoded or
+  /// extension scalar is reported according to the wrapper's own is_valid.
+  ///
+  /// \see ArraySpan::ComputeLogicalNullCount
+  virtual bool IsLogicalNull() const { return !is_valid; }
+
   bool Equals(const Scalar& other,
               const EqualOptions& options = EqualOptions::Defaults()) const;
 
@@ -860,7 +879,8 @@ struct ARROW_EXPORT RunEndEncodedScalar
 /// \brief A Scalar value for DictionaryType
 ///
 /// `is_valid` denotes the validity of the `index`, regardless of
-/// the corresponding value in the `dictionary`.
+/// the corresponding value in the `dictionary`; use IsLogicalNull()
+/// to also account for the referenced dictionary value.
 struct ARROW_EXPORT DictionaryScalar : public internal::PrimitiveScalarBase {
   using TypeClass = DictionaryType;
   struct ValueType {
@@ -879,6 +899,8 @@ struct ARROW_EXPORT DictionaryScalar : public internal::PrimitiveScalarBase {
 
   Result<std::shared_ptr<Scalar>> GetEncodedValue() const;
 
+  bool IsLogicalNull() const override;
+
   const void* data() const override {
     return internal::checked_cast<internal::PrimitiveScalarBase&>(*value.index).data();
   }

cpp/src/arrow/scalar_test.cc

@@ -1685,6 +1685,12 @@ TEST(TestDictionaryScalar, Basics) {
     ASSERT_OK(encoded_gamma->ValidateFull());
     ASSERT_TRUE(encoded_gamma->Equals(*MakeScalar("gamma")));
 
+    ASSERT_FALSE(scalar_alpha.IsLogicalNull());
+    ASSERT_FALSE(scalar_gamma.IsLogicalNull());
+    ASSERT_TRUE(scalar_null->IsLogicalNull());
+    // The index is valid but refers to a null dictionary value
+    ASSERT_TRUE(scalar_null_value.IsLogicalNull());
+
     // test Array.GetScalar
     DictionaryArray arr(ty, ArrayFromJSON(index_ty, "[2, 0, 1, null]"), dict);
     ASSERT_OK_AND_ASSIGN(auto first, arr.GetScalar(0));
@@ -1763,6 +1769,9 @@ TEST(TestDictionaryScalar, ValidateErrors) {
     scalar = DictionaryScalar(invalid, dict_ty);
     ASSERT_OK(scalar.Validate());
     ASSERT_RAISES(Invalid, scalar.ValidateFull());
+    // The dictionary value cannot be resolved, so IsLogicalNull falls back
+    // to !is_valid
+    ASSERT_FALSE(scalar.IsLogicalNull());
   }
 }