Aligns MutableDictionaryArray's with MutablePrimitiveArrays with TryPush

src/array/dictionary/mutable.rs

@@ -3,9 +3,8 @@ use std::{collections::hash_map::DefaultHasher, sync::Arc};
 
 use hash_hasher::HashedMap;
 
-use crate::array::TryExtend;
 use crate::{
-    array::{primitive::MutablePrimitiveArray, Array, MutableArray},
+    array::{primitive::MutablePrimitiveArray, Array, MutableArray, TryExtend, TryPush},
     bitmap::MutableBitmap,
     datatypes::DataType,
     error::{ArrowError, Result},
@@ -14,6 +13,19 @@ use crate::{
 use super::{DictionaryArray, DictionaryKey};
 
 /// A mutable, strong-typed version of [`DictionaryArray`].
+///
+/// # Example
+/// Building a UTF8 dictionary with `i32` keys.
+/// ```
+/// # use arrow2::array::{MutableDictionaryArray, MutableUtf8Array, TryPush};
+/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
+/// let mut array: MutableDictionaryArray<i32, MutableUtf8Array<i32>> = MutableDictionaryArray::new();
+/// array.try_push(Some("A"))?;
+/// array.try_push(Some("B"))?;
+/// array.try_push(None)?;
+/// array.try_push(Some("C"))?;
+/// # }
+/// ```
 #[derive(Debug)]
 pub struct MutableDictionaryArray<K: DictionaryKey, M: MutableArray> {
     data_type: DataType,
@@ -68,7 +80,7 @@ impl<K: DictionaryKey, M: MutableArray + Default> Default for MutableDictionaryA
 
 impl<K: DictionaryKey, M: MutableArray> MutableDictionaryArray<K, M> {
     /// Returns whether the value should be pushed to the values or not
-    pub fn try_push_valid<T: Hash>(&mut self, value: &T) -> Result<bool> {
+    fn try_push_valid<T: Hash>(&mut self, value: &T) -> Result<bool> {
         let mut hasher = DefaultHasher::new();
         value.hash(&mut hasher);
         let hash = hasher.finish();
@@ -181,3 +193,63 @@ where
         Ok(())
     }
 }
+
+impl<K, M, T> TryPush<Option<T>> for MutableDictionaryArray<K, M>
+where
+    K: DictionaryKey,
+    M: MutableArray + TryPush<Option<T>>,
+    T: Hash,
+{
+    fn try_push(&mut self, item: Option<T>) -> Result<()> {
+        if let Some(value) = item {
+            if self.try_push_valid(&value)? {
+                self.values.try_push(Some(value))
+            } else {
+                Ok(())
+            }
+        } else {
+            self.push_null();
+            Ok(())
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::array::{MutableUtf8Array, Utf8Array};
+
+    #[test]
+    fn test_push_dictionary() {
+        let mut new: MutableDictionaryArray<i32, MutableUtf8Array<i32>> =
+            MutableDictionaryArray::new();
+
+        for value in [Some("A"), Some("B"), None, Some("C"), Some("A"), Some("B")] {
+            new.try_push(value).unwrap();
+        }
+
+        let arc = new.values.as_arc();
+        let values_array: &Utf8Array<i32> = arc.as_any().downcast_ref().unwrap();
+
+        assert_eq!(*values_array, Utf8Array::<i32>::from_slice(["A", "B", "C"]));
+
+        let mut expected_keys = MutablePrimitiveArray::from_slice(&[0, 1]);
+        expected_keys.push(None);
+        expected_keys.push(Some(2));
+        expected_keys.push(Some(0));
+        expected_keys.push(Some(1));
+        assert_eq!(new.keys, expected_keys);
+
+        let expected_map = ["A", "B", "C"]
+            .iter()
+            .enumerate()
+            .map(|(index, value)| {
+                let mut hasher = DefaultHasher::new();
+                value.hash(&mut hasher);
+                let hash = hasher.finish();
+                (hash, index as i32)
+            })
+            .collect::<HashedMap<_, _>>();
+        assert_eq!(new.map, expected_map);
+    }
+}