diff --git a/arrow/src/array/transform/mod.rs b/arrow/src/array/transform/mod.rs index 2cad8af0d9c..68ae7f6d4d0 100644 --- a/arrow/src/array/transform/mod.rs +++ b/arrow/src/array/transform/mod.rs @@ -395,6 +395,13 @@ impl<'a> MutableArrayData<'a> { array_capacity = *capacity; new_buffers(data_type, *capacity) } + ( + DataType::List(_) | DataType::LargeList(_), + Capacities::List(capacity, _), + ) => { + array_capacity = *capacity; + new_buffers(data_type, *capacity) + } _ => panic!("Capacities: {:?} not yet supported", capacities), }; @@ -434,11 +441,10 @@ impl<'a> MutableArrayData<'a> { let capacities = if let Capacities::List(capacity, ref child_capacities) = capacities { - array_capacity = capacity; child_capacities .clone() .map(|c| *c) - .unwrap_or(Capacities::Array(array_capacity)) + .unwrap_or(Capacities::Array(capacity)) } else { Capacities::Array(array_capacity) }; @@ -1318,6 +1324,40 @@ mod tests { Ok(()) } + #[test] + fn test_list_append_with_capacities() -> Result<()> { + let mut builder = ListBuilder::::new(Int64Builder::new(24)); + builder.values().append_slice(&[1, 2, 3])?; + builder.append(true)?; + builder.values().append_slice(&[4, 5])?; + builder.append(true)?; + builder.values().append_slice(&[6, 7, 8])?; + builder.values().append_slice(&[9, 10, 11])?; + builder.append(true)?; + let a = builder.finish(); + + let a_builder = Int64Builder::new(24); + let mut a_builder = ListBuilder::::new(a_builder); + a_builder.values().append_slice(&[12, 13])?; + a_builder.append(true)?; + a_builder.append(true)?; + a_builder.values().append_slice(&[14, 15, 16, 17])?; + a_builder.append(true)?; + let b = a_builder.finish(); + + let mutable = MutableArrayData::with_capacities( + vec![a.data(), b.data()], + false, + Capacities::List(6, Some(Box::new(Capacities::Array(17)))), + ); + + // capacities are rounded up to multiples of 64 by MutableBuffer + assert_eq!(mutable.data.buffer1.capacity(), 64); + assert_eq!(mutable.data.child_data[0].data.buffer1.capacity(), 192); + + Ok(()) + } + #[test] fn test_map_nulls_append() -> Result<()> { let mut builder = MapBuilder::::new(