This is an automated email from the ASF dual-hosted git repository.
tustvold 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 9e060dc6c6 Add `entries` to `MapBuilder` to return both key and value
array builders (#5218)
9e060dc6c6 is described below
commit 9e060dc6c68e6aac30d0ae643d1d1854db44a378
Author: Liang-Chi Hsieh <[email protected]>
AuthorDate: Mon Dec 18 01:19:12 2023 -0800
Add `entries` to `MapBuilder` to return both key and value array builders
(#5218)
* Add keys_and_values to MapBuilder to return both key and value array
builders
* Rename to `entries`.
---
arrow-array/src/builder/map_builder.rs | 5 ++++
arrow/tests/array_transform.rs | 51 ++++++++++++++++++++++++++++++++++
2 files changed, 56 insertions(+)
diff --git a/arrow-array/src/builder/map_builder.rs
b/arrow-array/src/builder/map_builder.rs
index 5785bb5faf..edc45c2c23 100644
--- a/arrow-array/src/builder/map_builder.rs
+++ b/arrow-array/src/builder/map_builder.rs
@@ -119,6 +119,11 @@ impl<K: ArrayBuilder, V: ArrayBuilder> MapBuilder<K, V> {
&mut self.value_builder
}
+ /// Returns both the key and value array builders of the map
+ pub fn entries(&mut self) -> (&mut K, &mut V) {
+ (&mut self.key_builder, &mut self.value_builder)
+ }
+
/// Finish the current map array slot
///
/// Returns an error if the key and values builders are in an inconsistent
state.
diff --git a/arrow/tests/array_transform.rs b/arrow/tests/array_transform.rs
index 6f5b245b8e..5a267c876d 100644
--- a/arrow/tests/array_transform.rs
+++ b/arrow/tests/array_transform.rs
@@ -830,6 +830,57 @@ fn test_map_nulls_append() {
assert_eq!(result, expected_list_data);
}
+#[test]
+fn test_map_keys_values_append() {
+ let mut builder = MapBuilder::<Int64Builder, Int64Builder>::new(
+ None,
+ Int64Builder::with_capacity(32),
+ Int64Builder::with_capacity(32),
+ );
+ let (keys, values) = builder.entries();
+ keys.append_slice(&[1, 2, 3]);
+ values.append_slice(&[1, 3, 4]);
+ builder.append(true).unwrap();
+
+ let (keys, values) = builder.entries();
+ keys.append_slice(&[4, 5]);
+ values.append_slice(&[4, 6]);
+ builder.append(true).unwrap();
+
+ builder.append(false).unwrap();
+
+ let map = builder.finish();
+ assert!(map.is_null(2));
+
+ let first = map.value(0);
+ let keys = first
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap();
+ let values = first
+ .column(1)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap();
+ assert_eq!(keys, &Int64Array::from(vec![Some(1), Some(2), Some(3)]));
+ assert_eq!(values, &Int64Array::from(vec![Some(1), Some(3), Some(4)]));
+
+ let second = map.value(1);
+ let keys = second
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap();
+ let values = second
+ .column(1)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap();
+ assert_eq!(keys, &Int64Array::from(vec![Some(4), Some(5)]));
+ assert_eq!(values, &Int64Array::from(vec![Some(4), Some(6)]));
+}
+
#[test]
fn test_list_of_strings_append() {
// [["alpha", "beta", None]]
