[arrow-rs] branch master updated: Support specifying list capacities for MutableArrayData (#1885)

[alamb]

This is an automated email from the ASF dual-hosted git repository.

alamb 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 f159679ae Support specifying list capacities for MutableArrayData 
(#1885)
f159679ae is described below

commit f159679ae1c8dc05e0ed4c25a58a26bb44bbda9f
Author: Jörn Horstmann <g...@jhorstmann.net>
AuthorDate: Thu Jun 16 22:07:36 2022 +0200

    Support specifying list capacities for MutableArrayData (#1885)
---
 arrow/src/array/transform/mod.rs | 44 ++++++++++++++++++++++++++++++++++++++--
 1 file changed, 42 insertions(+), 2 deletions(-)

diff --git a/arrow/src/array/transform/mod.rs b/arrow/src/array/transform/mod.rs
index 2cad8af0d..68ae7f6d4 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::<Int64Builder>::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::<Int64Builder>::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::<Int64Builder, Int64Builder>::new(