Implement efficient DictionaryArray::logical_null_count

apache · Nov 16, 2024 · f8ed2cc · f8ed2cc
1 parent 7b2756a
commit f8ed2cc
Showing 1 changed file with 20 additions and 0 deletions.
diff --git a/arrow-array/src/array/dictionary_array.rs b/arrow-array/src/array/dictionary_array.rs
@@ -749,6 +749,26 @@ impl<T: ArrowDictionaryKeyType> Array for DictionaryArray<T> {
         }
     }
 
+    fn logical_null_count(&self) -> usize {
+        match (self.keys.nulls(), self.values.logical_nulls()) {
+            (None, None) => 0,
+            (Some(key_nulls), None) => key_nulls.null_count(),
+            (None, Some(value_nulls)) => self
+                .keys
+                .values()
+                .iter()
+                .filter(|k| value_nulls.is_null(k.as_usize()))
+                .count(),
+            (Some(key_nulls), Some(value_nulls)) => self
+                .keys
+                .values()
+                .iter()
+                .enumerate()
+                .filter(|(idx, k)| key_nulls.is_null(*idx) || value_nulls.is_null(k.as_usize()))
+                .count(),
+        }
+    }
+
     fn is_nullable(&self) -> bool {
         !self.is_empty() && (self.nulls().is_some() || self.values.is_nullable())
     }