[GitHub] [arrow-rs] alamb commented on a diff in pull request #3715: Add `into_primitive_dict_builder` to `DictionaryArray`

["alamb (via GitHub)" <gi...@apache.org>]

alamb commented on code in PR #3715:
URL: https://github.com/apache/arrow-rs/pull/3715#discussion_r1123221345


##########
arrow-array/src/array/dictionary_array.rs:
##########
@@ -941,4 +982,62 @@ mod tests {
         let a = DictionaryArray::<Int32Type>::from_iter(["32"]);
         let _ = DictionaryArray::<Int64Type>::from(a.into_data());
     }
+
+    #[test]
+    fn test_into_primitive_dict_builder() {
+        let values = Int32Array::from_iter_values([10_i32, 12, 15]);
+        let keys = Int8Array::from_iter_values([1_i8, 0, 2, 0]);
+
+        let dict_array = DictionaryArray::<Int8Type>::try_new(&keys, &values).unwrap();
+
+        let boxed: ArrayRef = Arc::new(dict_array);
+        let col: DictionaryArray<Int8Type> = as_dictionary_array(&boxed).clone();
+
+        drop(boxed);
+        drop(keys);
+        drop(values);
+
+        let mut builder = col.into_primitive_dict_builder::<Int32Type>().unwrap();
+
+        let slice = builder.values_slice_mut();
+        assert_eq!(slice, &[10, 12, 15]);
+
+        slice[0] = 4;
+        slice[1] = 2;
+        slice[2] = 1;
+
+        let values = Int32Array::from_iter_values([4_i32, 2, 1]);
+        let keys = Int8Array::from_iter_values([1_i8, 0, 2, 0]);
+
+        let expected = DictionaryArray::<Int8Type>::try_new(&keys, &values).unwrap();
+
+        let new_array = builder.finish();
+        assert_eq!(expected, new_array);
+    }
+
+    #[test]
+    fn test_into_primitive_dict_builder_cloned_array() {
+        let values = Int32Array::from_iter_values([10_i32, 12, 15]);
+        let keys = Int8Array::from_iter_values([1_i8, 0, 2, 0]);
+
+        let dict_array = DictionaryArray::<Int8Type>::try_new(&keys, &values).unwrap();
+
+        let boxed: ArrayRef = Arc::new(dict_array);
+
+        let col: DictionaryArray<Int8Type> =
+            DictionaryArray::<Int8Type>::from(boxed.data().clone());
+        let err = col.into_primitive_dict_builder::<Int32Type>();
+
+        match err {
+            Ok(_) => panic!("Should not get builder from cloned array"),
+            Err(returned) => {

Review Comment:
   You could also write this like `let returned = err.unwrap_err()`



##########
arrow-array/src/array/dictionary_array.rs:
##########
@@ -402,6 +403,44 @@ impl<K: ArrowPrimitiveType> DictionaryArray<K> {
         // Offsets were valid before and verified length is greater than or equal
         Self::from(unsafe { builder.build_unchecked() })
     }
+
+    /// Returns `PrimitiveDictionaryBuilder` of this dictionary array for mutating
+    /// its keys and values if the underlying data buffer is not shared by others.
+    pub fn into_primitive_dict_builder<V>(
+        self,
+    ) -> Result<PrimitiveDictionaryBuilder<K, V>, Self>
+    where
+        V: ArrowPrimitiveType,
+    {
+        if !self.value_type().is_primitive() {
+            return Err(self);
+        }
+
+        let key_array = as_primitive_array::<K>(self.keys()).clone();
+        let value_array = as_primitive_array::<V>(self.values()).clone();
+
+        drop(self.data);

Review Comment:
   is this needed to drop the ref counts to be able to call `into_builder()` without copying data (as self has a ref count)?



-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org