This is an automated email from the ASF dual-hosted git repository.
viirya pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/master by this push:
new 1883bb691 Add unary_dict_mut (#3804)
1883bb691 is described below
commit 1883bb691a33c39ce13e355e7f0a82414fc74010
Author: Liang-Chi Hsieh <[email protected]>
AuthorDate: Thu Mar 9 00:33:21 2023 -0800
Add unary_dict_mut (#3804)
* Add unary_dict_mut
* Remove unused function
---
arrow-arith/src/arity.rs | 21 ++++++++++++++++-
arrow-array/src/array/dictionary_array.rs | 38 +++++++++++++++++++++++++++++++
2 files changed, 58 insertions(+), 1 deletion(-)
diff --git a/arrow-arith/src/arity.rs b/arrow-arith/src/arity.rs
index 5d4973714..74edd654b 100644
--- a/arrow-arith/src/arity.rs
+++ b/arrow-arith/src/arity.rs
@@ -61,7 +61,7 @@ where
pub fn unary_mut<I, F>(
array: PrimitiveArray<I>,
op: F,
-) -> std::result::Result<PrimitiveArray<I>, PrimitiveArray<I>>
+) -> Result<PrimitiveArray<I>, PrimitiveArray<I>>
where
I: ArrowPrimitiveType,
F: Fn(I::Native) -> I::Native,
@@ -647,4 +647,23 @@ mod tests {
.unwrap()
.expect_err("should got error");
}
+
+ #[test]
+ fn test_unary_dict_mut() {
+ let values = Int32Array::from(vec![Some(10), Some(20), None]);
+ let keys = Int8Array::from_iter_values([0, 0, 1, 2]);
+ let dictionary = DictionaryArray::<Int8Type>::try_new(&keys,
&values).unwrap();
+
+ drop(keys);
+ drop(values);
+
+ let updated = dictionary.unary_mut::<_, Int32Type>(|x| x + 1).unwrap();
+ let typed = updated.downcast_dict::<Int32Array>().unwrap();
+ assert_eq!(typed.value(0), 11);
+ assert_eq!(typed.value(1), 11);
+ assert_eq!(typed.value(2), 21);
+
+ let values = updated.values();
+ assert!(values.is_null(2));
+ }
}
diff --git a/arrow-array/src/array/dictionary_array.rs
b/arrow-array/src/array/dictionary_array.rs
index 74b5cec19..f9a40c6f3 100644
--- a/arrow-array/src/array/dictionary_array.rs
+++ b/arrow-array/src/array/dictionary_array.rs
@@ -433,6 +433,44 @@ impl<K: ArrowDictionaryKeyType> DictionaryArray<K> {
}
}
}
+
+ /// Applies an unary and infallible function to a mutable dictionary array.
+ /// Mutable dictionary array means that the buffers are not shared with
other arrays.
+ /// As a result, this mutates the buffers directly without allocating new
buffers.
+ ///
+ /// # Implementation
+ ///
+ /// This will apply the function for all dictionary values, including
those on null slots.
+ /// This implies that the operation must be infallible for any value of
the corresponding type
+ /// or this function may panic.
+ /// # Example
+ /// ```
+ /// use arrow_array::{Array, ArrayAccessor, DictionaryArray, StringArray,
types::{Int8Type, Int32Type}};
+ /// use arrow_array::{Int8Array, Int32Array};
+ /// let values = Int32Array::from(vec![Some(10), Some(20), None]);
+ /// let keys = Int8Array::from_iter_values([0, 0, 1, 2]);
+ /// let dictionary = DictionaryArray::<Int8Type>::try_new(&keys,
&values).unwrap();
+ /// drop(keys);
+ /// drop(values);
+ /// let c = dictionary.unary_mut::<_, Int32Type>(|x| x + 1).unwrap();
+ /// let typed = c.downcast_dict::<Int32Array>().unwrap();
+ /// assert_eq!(typed.value(0), 11);
+ /// assert_eq!(typed.value(1), 11);
+ /// assert_eq!(typed.value(2), 21);
+ /// ```
+ pub fn unary_mut<F, V>(self, op: F) -> Result<DictionaryArray<K>,
DictionaryArray<K>>
+ where
+ V: ArrowPrimitiveType,
+ F: Fn(V::Native) -> V::Native,
+ {
+ let mut builder: PrimitiveDictionaryBuilder<K, V> =
+ self.into_primitive_dict_builder()?;
+ builder
+ .values_slice_mut()
+ .iter_mut()
+ .for_each(|v| *v = op(*v));
+ Ok(builder.finish())
+ }
}
/// Constructs a `DictionaryArray` from an array data reference.
