pitrou commented on code in PR #50347:
URL: https://github.com/apache/arrow/pull/50347#discussion_r3622479022


##########
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;
+  }

Review Comment:
   I don't these cases are valid (`Scalar::Validate` would return an error), so 
no need to check for them (or just as debug checks).



##########
cpp/src/arrow/datum_test.cc:
##########
@@ -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));
+}

Review Comment:
   Can you also test with a chunked array (for example with a dictionary type)?



##########
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()) 
{

Review Comment:
   Same here: not sure it's worth checking for these invalid cases, and it may 
actually be better to let them crash rather than return a seemingly valid value.



##########
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));

Review Comment:
   Unrelated to this PR, but it would may be useful to separately add a 
`VisitScalar` that, similar to `VisitType`, takes a callable object.
   
   This would allow to rewrite this as, for example:
   ```c++
   Result<int64_t> DictionaryIndexValue(const Scalar& index) {
     return VisitScalar(index, [](auto&& typed_index) {
       using TypeClass = typename 
std::decay_t<decltype(typed_index)>::TypeClass;
       if constexpr (is_integer(TypeClass::id)) {
         return static_cast<int64_t>(typed_index.value);
       }
       return Status::TypeError("Invalid dictionary index type");
     });
   }
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to