This is an automated email from the ASF dual-hosted git repository.
Jefffrey pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new bab2817d31 Support map/struct capacities in MutableArrayData (#10801)
bab2817d31 is described below
commit bab2817d31706b6164c148434efd49cf98a06f3e
Author: Emily Matheys <[email protected]>
AuthorDate: Mon Aug 24 08:03:44 2026 +0300
Support map/struct capacities in MutableArrayData (#10801)
The tests seems to completely skip over this and not even use a
MutableArrayData, causing them to pass, but all 3 added tests would
panic before adding the Map and Struct cases to the match statement.
With this change we can now use MutableArrayData on maps properly with
preallocation(meaning using Capacities::List and Capacities::Struct).
Not actually a breaking change, although it does not panic anymore,
which is a change I guess.
---
arrow-data/src/transform/mod.rs | 52 +++++++++++++++++++++++++++++++++--
arrow/tests/array_transform.rs | 60 ++++++++++++++++++++++++++++-------------
2 files changed, 91 insertions(+), 21 deletions(-)
diff --git a/arrow-data/src/transform/mod.rs b/arrow-data/src/transform/mod.rs
index 7558d1e887..989e502e09 100644
--- a/arrow-data/src/transform/mod.rs
+++ b/arrow-data/src/transform/mod.rs
@@ -371,11 +371,14 @@ pub enum Capacities {
/// * the capacity of the array offsets
/// * the capacity of the binary/ str buffer
Binary(usize, Option<usize>),
- /// List and LargeList data types
+ /// List, LargeList and Map data types
///
/// Defines
/// * the capacity of the array offsets
/// * the capacity of the child data
+ ///
+ /// For Map the child data is the entries [`DataType::Struct`], so the
child
+ /// capacity is a [`Capacities::Struct`] holding the key and value
capacities.
List(usize, Option<Box<Capacities>>),
/// Struct type
///
@@ -492,12 +495,17 @@ impl<'a> MutableArrayData<'a> {
| DataType::LargeList(_)
| DataType::ListView(_)
| DataType::LargeListView(_)
- | DataType::FixedSizeList(_, _),
+ | DataType::FixedSizeList(_, _)
+ | DataType::Map(_, _),
Capacities::List(capacity, _),
) => {
array_capacity = *capacity;
new_buffers(data_type, *capacity)
}
+ (DataType::Struct(_), Capacities::Struct(capacity, _)) => {
+ array_capacity = *capacity;
+ new_buffers(data_type, *capacity)
+ }
_ => panic!("Capacities: {capacities:?} not yet supported"),
};
@@ -1029,4 +1037,44 @@ mod test {
assert_eq!(mutable.data.buffer1.capacity(), 64);
assert_eq!(mutable.data.child_data[0].data.buffer1.capacity(), 192);
}
+
+ #[test]
+ fn test_map_append_with_capacities() {
+ let entries = Arc::new(Field::new(
+ "entries",
+ DataType::Struct(
+ vec![
+ Field::new("keys", DataType::Int64, false),
+ Field::new("values", DataType::Int64, true),
+ ]
+ .into(),
+ ),
+ false,
+ ));
+ let array = ArrayData::new_empty(&DataType::Map(entries, false));
+
+ let mutable = MutableArrayData::with_capacities(
+ vec![&array],
+ false,
+ Capacities::List(
+ 6,
+ Some(Box::new(Capacities::Struct(
+ 17,
+ Some(vec![Capacities::Array(17), Capacities::Array(17)]),
+ ))),
+ ),
+ );
+
+ // capacities are rounded up to multiples of 64 by MutableBuffer
+ // the map offsets buffer holds `1 + 6` i32s
+ assert_eq!(mutable.data.buffer1.capacity(), 64);
+
+ // the entries struct itself has no buffers of its own
+ let entries = &mutable.data.child_data[0];
+ assert_eq!(entries.data.buffer1.capacity(), 0);
+
+ // both key and value buffers hold 17 i64s
+ assert_eq!(entries.data.child_data[0].data.buffer1.capacity(), 192);
+ assert_eq!(entries.data.child_data[1].data.buffer1.capacity(), 192);
+ }
}
diff --git a/arrow/tests/array_transform.rs b/arrow/tests/array_transform.rs
index 1f82ca4927..ff005b8630 100644
--- a/arrow/tests/array_transform.rs
+++ b/arrow/tests/array_transform.rs
@@ -25,7 +25,7 @@ use arrow::datatypes::{Int16Type, IntervalMonthDayNanoType};
use arrow_array::StringViewArray;
use arrow_buffer::{Buffer, ScalarBuffer};
use arrow_data::ArrayData;
-use arrow_data::transform::MutableArrayData;
+use arrow_data::transform::{Capacities, MutableArrayData};
use arrow_schema::{DataType, Field, Fields, UnionFields};
use std::sync::Arc;
@@ -449,14 +449,6 @@ fn test_struct_many() {
let array = StructArray::try_from(vec![("f1", strings.clone()), ("f2",
ints.clone())])
.unwrap()
.into_data();
- let arrays = vec![&array, &array];
- let mut mutable = MutableArrayData::new(arrays, false, 0);
-
- mutable.try_extend(0, 1, 3).unwrap();
- mutable.try_extend(1, 0, 2).unwrap();
- let data = mutable.freeze();
- let array = StructArray::from(data);
-
let expected_string =
Arc::new(StringArray::from(vec![None, None, Some("joe"), None])) as
ArrayRef;
let expected_int =
@@ -464,7 +456,24 @@ fn test_struct_many() {
let expected =
StructArray::try_from(vec![("f1", expected_string), ("f2",
expected_int)]).unwrap();
- assert_eq!(array, expected)
+
+ // exact capacities must produce the same result as letting the buffers
grow
+ for capacities in [
+ Capacities::Array(0),
+ Capacities::Struct(
+ 4,
+ // 4 slots per field, and "joe" is the only string kept
+ Some(vec![Capacities::Binary(4, Some(3)), Capacities::Array(4)]),
+ ),
+ ] {
+ let mut mutable =
+ MutableArrayData::with_capacities(vec![&array, &array], false,
capacities);
+
+ mutable.try_extend(0, 1, 3).unwrap();
+ mutable.try_extend(1, 0, 2).unwrap();
+
+ assert_eq!(StructArray::from(mutable.freeze()), expected);
+ }
}
#[test]
@@ -732,14 +741,6 @@ fn test_map_nulls_append() {
let c = b.slice(1, 2);
let d = b.slice(2, 2);
- let mut mutable = MutableArrayData::new(vec![&a, &b, &c, &d], false, 10);
-
- mutable.try_extend(0, 0, a.len()).unwrap();
- mutable.try_extend(1, 0, b.len()).unwrap();
- mutable.try_extend(2, 0, c.len()).unwrap();
- mutable.try_extend(3, 0, d.len()).unwrap();
- let result = mutable.freeze();
-
let expected_key_array = Int64Array::from(vec![
Some(1),
Some(2),
@@ -836,7 +837,28 @@ fn test_map_nulls_append() {
vec![expected_entry_array.into_data()],
)
.unwrap();
- assert_eq!(result, expected_list_data);
+
+ // exact capacities must produce the same result as letting the buffers
grow
+ for capacities in [
+ Capacities::Array(10),
+ Capacities::List(
+ 12,
+ Some(Box::new(Capacities::Struct(
+ 23,
+ Some(vec![Capacities::Array(23), Capacities::Array(23)]),
+ ))),
+ ),
+ ] {
+ let mut mutable =
+ MutableArrayData::with_capacities(vec![&a, &b, &c, &d], false,
capacities);
+
+ mutable.try_extend(0, 0, a.len()).unwrap();
+ mutable.try_extend(1, 0, b.len()).unwrap();
+ mutable.try_extend(2, 0, c.len()).unwrap();
+ mutable.try_extend(3, 0, d.len()).unwrap();
+
+ assert_eq!(mutable.freeze(), expected_list_data);
+ }
}
#[test]
