From 9e060dc6c68e6aac30d0ae643d1d1854db44a378 Mon Sep 17 00:00:00 2001
From: Liang-Chi Hsieh <viirya@gmail.com>
Date: Mon, 18 Dec 2023 01:19:12 -0800
Subject: [PATCH] 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 5785bb5fafb5..edc45c2c23b7 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 6f5b245b8e3b..5a267c876d6a 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]]