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 MapBuilder { &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::::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::() + .unwrap(); + let values = first + .column(1) + .as_any() + .downcast_ref::() + .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::() + .unwrap(); + let values = second + .column(1) + .as_any() + .downcast_ref::() + .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]]