(arrow-rs) branch master updated: Add `entries` to `MapBuilder` to return both key and value array builders (#5218)

[tu...@apache.org]

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 <vi...@gmail.com>
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]]