feat(array): let callers choose a nested builder's chunk threshold

`ChildBuilder` copies appended arrays shorter than 64 values instead of
giving them a chunk, which is right for the nested builders that append the
elements of a single list, and wrong for a caller whose appended arrays are
themselves chunks worth preserving.

Add `ArrayBuilder::set_min_chunk_len`, a no-op for builders without array
children, so such a caller can lower the threshold — to zero to keep every
chunk boundary. The threshold is read when an array is appended and applies
transitively to the children of a builder's children.

Signed-off-by: Claude <noreply@anthropic.com>

Author: claude

vortex-array/src/builders/child.rs

@@ -18,13 +18,16 @@ use crate::dtype::Nullability;
 use crate::scalar::Scalar;
 use crate::validity::Validity;
 
-/// The shortest array that a [`ChildBuilder`] keeps as a chunk of its own.
+/// The default shortest array that a [`ChildBuilder`] keeps as a chunk of its own.
 ///
 /// Nested builders routinely append very short arrays (the elements of a single list, for
 /// example), and giving each one its own chunk would produce a [`ChunkedArray`] with more chunks
 /// than values. Below this length, copying the values costs less than the indirection that every
 /// later access to the chunk would pay.
-pub(super) const MIN_CHUNK_LEN: usize = 64;
+///
+/// Callers that would rather keep every chunk boundary, however short, can lower the threshold
+/// with [`ArrayBuilder::set_min_chunk_len`].
+pub(super) const DEFAULT_MIN_CHUNK_LEN: usize = 64;
 
 /// Accumulates the child of a nested [`ArrayBuilder`] without canonicalizing appended arrays.
 ///
@@ -51,6 +54,9 @@ pub struct ChildBuilder {
 
     /// Builder holding the scalars appended after the last chunk.
     pending: Box<dyn ArrayBuilder>,
+
+    /// The shortest appended array that is kept as a chunk rather than copied.
+    min_chunk_len: usize,
 }
 
 impl ChildBuilder {
@@ -61,9 +67,18 @@ impl ChildBuilder {
             chunks: Vec::new(),
             chunks_len: 0,
             pending: builder_with_capacity(dtype, capacity),
+            min_chunk_len: DEFAULT_MIN_CHUNK_LEN,
         }
     }
 
+    /// Sets the shortest appended array that is kept as a chunk rather than copied.
+    ///
+    /// See [`ArrayBuilder::set_min_chunk_len`].
+    pub fn set_min_chunk_len(&mut self, min_chunk_len: usize) {
+        self.min_chunk_len = min_chunk_len;
+        self.pending.set_min_chunk_len(min_chunk_len);
+    }
+
     /// The number of values appended so far.
     pub fn len(&self) -> usize {
         self.chunks_len + self.pending.len()
@@ -85,7 +100,7 @@ impl ChildBuilder {
             return Ok(());
         }
 
-        if array.len() < MIN_CHUNK_LEN {
+        if array.len() < self.min_chunk_len {
             return array.append_to_builder(self.pending.as_mut(), ctx);
         }
 
@@ -196,14 +211,16 @@ impl ChildBuilder {
 
 #[cfg(test)]
 mod tests {
+    use std::sync::Arc;
+
     use rstest::rstest;
     use vortex_buffer::buffer;
     use vortex_error::VortexExpect;
     use vortex_error::VortexResult;
     use vortex_mask::Mask;
 
     use super::ChildBuilder;
-    use super::MIN_CHUNK_LEN;
+    use super::DEFAULT_MIN_CHUNK_LEN as MIN_CHUNK_LEN;
     use crate::ArrayRef;
     use crate::IntoArray;
     use crate::VortexSessionExecute;
@@ -212,16 +229,20 @@ mod tests {
     use crate::arrays::ChunkedArray;
     use crate::arrays::Constant;
     use crate::arrays::ConstantArray;
+    use crate::arrays::ListView;
+    use crate::arrays::ListViewArray;
     use crate::arrays::Masked;
     use crate::arrays::Primitive;
     use crate::arrays::PrimitiveArray;
     use crate::arrays::chunked::ChunkedArrayExt;
+    use crate::arrays::listview::ListViewArraySlotsExt;
     use crate::assert_arrays_eq;
     use crate::dtype::DType;
     use crate::dtype::Nullability::NonNullable;
     use crate::dtype::Nullability::Nullable;
     use crate::dtype::PType::I32;
     use crate::scalar::Scalar;
+    use crate::validity::Validity;
 
     /// A non-canonical array of `len` values, all equal to `value`.
     fn constant(value: i32, len: usize) -> ArrayRef {
@@ -476,6 +497,50 @@ mod tests {
         unsafe { builder.set_validity_unchecked(Mask::new_true(MIN_CHUNK_LEN)) };
     }
 
+    /// A threshold of zero keeps every chunk boundary, however short.
+    #[test]
+    fn test_min_chunk_len_zero_keeps_every_chunk() -> VortexResult<()> {
+        let mut ctx = array_session().create_execution_ctx();
+        let mut builder = ChildBuilder::with_capacity(&DType::from(I32), 0);
+        builder.set_min_chunk_len(0);
+
+        builder.append_array(&constant(1, 1), &mut ctx)?;
+        builder.append_array(&constant(2, 1), &mut ctx)?;
+
+        let child = builder.finish();
+        let chunked = child.as_::<Chunked>();
+        assert_eq!(chunked.nchunks(), 2);
+        assert!(chunked.iter_chunks().all(|chunk| chunk.is::<Constant>()));
+
+        Ok(())
+    }
+
+    /// The threshold reaches the children of a child: an array short enough to be materialized
+    /// lands in the scalar builder, which is itself a nested builder with children of its own.
+    #[test]
+    fn test_min_chunk_len_applies_transitively() -> VortexResult<()> {
+        let mut ctx = array_session().create_execution_ctx();
+        let dtype = DType::List(Arc::new(DType::from(I32)), NonNullable);
+        let mut builder = ChildBuilder::with_capacity(&dtype, 0);
+        builder.set_min_chunk_len(4 * MIN_CHUNK_LEN);
+
+        // Two lists — far too few to be a chunk — holding enough elements between them that those
+        // elements would become a chunk if the threshold stopped at the outer builder.
+        let lists = ListViewArray::new(
+            constant(1, 2 * MIN_CHUNK_LEN),
+            buffer![0u64, MIN_CHUNK_LEN as u64].into_array(),
+            buffer![MIN_CHUNK_LEN as u64; 2].into_array(),
+            Validity::NonNullable,
+        )
+        .into_array();
+        builder.append_array(&lists, &mut ctx)?;
+
+        let child = builder.finish();
+        assert!(child.as_::<ListView>().elements().is::<Primitive>());
+
+        Ok(())
+    }
+
     #[test]
     fn test_finish_resets_the_builder() -> VortexResult<()> {
         let mut ctx = array_session().create_execution_ctx();

vortex-array/src/builders/extension.rs

@@ -110,6 +110,10 @@ impl ArrayBuilder for ExtensionBuilder {
         self.append_value(scalar.as_extension())
     }
 
+    fn set_min_chunk_len(&mut self, min_chunk_len: usize) {
+        self.storage.set_min_chunk_len(min_chunk_len);
+    }
+
     fn reserve_exact(&mut self, capacity: usize) {
         self.storage.reserve_exact(capacity)
     }

vortex-array/src/builders/fixed_size_list.rs

@@ -255,6 +255,10 @@ impl ArrayBuilder for FixedSizeListBuilder {
 
     /// This will increase the capacity if extending with this `array` would go past the original
     /// capacity.
+    fn set_min_chunk_len(&mut self, min_chunk_len: usize) {
+        self.elements_builder.set_min_chunk_len(min_chunk_len);
+    }
+
     fn reserve_exact(&mut self, additional: usize) {
         self.elements_builder
             .reserve_exact(additional * self.list_size() as usize);

vortex-array/src/builders/list.rs

@@ -274,6 +274,10 @@ impl<O: IntegerPType> ArrayBuilder for ListBuilder<O> {
         self.append_value(scalar.as_list())
     }
 
+    fn set_min_chunk_len(&mut self, min_chunk_len: usize) {
+        self.elements_builder.set_min_chunk_len(min_chunk_len);
+    }
+
     fn reserve_exact(&mut self, additional: usize) {
         self.elements_builder.reserve_exact(additional);
         self.offsets_builder.reserve_exact(additional);

vortex-array/src/builders/listview.rs

@@ -311,6 +311,10 @@ impl<O: IntegerPType, S: IntegerPType> ArrayBuilder for ListViewBuilder<O, S> {
         self.append_value(list_scalar)
     }
 
+    fn set_min_chunk_len(&mut self, min_chunk_len: usize) {
+        self.elements_builder.set_min_chunk_len(min_chunk_len);
+    }
+
     fn reserve_exact(&mut self, capacity: usize) {
         self.elements_builder.reserve_exact(capacity * 2);
         self.offsets_builder.reserve_exact(capacity);

vortex-array/src/builders/mod.rs

@@ -170,6 +170,20 @@ pub trait ArrayBuilder: Send {
     /// Allocate space for extra `additional` items
     fn reserve_exact(&mut self, additional: usize);
 
+    /// Sets the shortest array that this builder's children keep as a chunk of their own.
+    ///
+    /// Arrays shorter than this are copied into a canonical child instead, which costs a copy but
+    /// spares every later reader the chunk indirection. Pass `0` to keep every chunk boundary,
+    /// however short — worthwhile when the appended arrays are themselves chunks whose identity is
+    /// the point, as when canonicalizing a [`ChunkedArray`](crate::arrays::ChunkedArray).
+    ///
+    /// The threshold is read when an array is appended, so it only affects subsequent appends. It
+    /// applies transitively to the children of this builder's children. Builders without array
+    /// children ignore it.
+    fn set_min_chunk_len(&mut self, min_chunk_len: usize) {
+        let _ = min_chunk_len;
+    }
+
     /// Override builders validity with the one provided.
     ///
     /// Note that this will have no effect on the final array if the array builder is non-nullable.

vortex-array/src/builders/struct_.rs

@@ -184,6 +184,12 @@ impl ArrayBuilder for StructBuilder {
         self.append_value(scalar.as_struct())
     }
 
+    fn set_min_chunk_len(&mut self, min_chunk_len: usize) {
+        for builder in &mut self.builders {
+            builder.set_min_chunk_len(min_chunk_len);
+        }
+    }
+
     fn reserve_exact(&mut self, capacity: usize) {
         self.builders.iter_mut().for_each(|builder| {
             builder.reserve_exact(capacity);

vortex-array/src/builders/tests.rs

@@ -43,7 +43,7 @@ use crate::assert_arrays_eq;
 use crate::builders::ArrayBuilder;
 use crate::builders::ListBuilder;
 use crate::builders::builder_with_capacity;
-use crate::builders::child::MIN_CHUNK_LEN;
+use crate::builders::child::DEFAULT_MIN_CHUNK_LEN as MIN_CHUNK_LEN;
 use crate::dtype::DType;
 use crate::dtype::DecimalDType;
 use crate::dtype::Nullability;
@@ -1053,6 +1053,35 @@ fn test_struct_builder_interleaves_arrays_and_scalars() -> VortexResult<()> {
     Ok(())
 }
 
+/// Lowering the threshold to zero keeps every chunk boundary a nested builder is handed, however
+/// short — what a caller wants when the appended arrays are themselves chunks worth preserving.
+#[test]
+fn test_min_chunk_len_zero_preserves_short_chunks() -> VortexResult<()> {
+    let mut ctx = array_session().create_execution_ctx();
+
+    let array = StructArray::try_from_iter([("a", buffer![1i32])])?.into_array();
+    let mut builder = builder_with_capacity(array.dtype(), 0);
+    builder.set_min_chunk_len(0);
+    array.append_to_builder(builder.as_mut(), &mut ctx)?;
+    array.append_to_builder(builder.as_mut(), &mut ctx)?;
+    let built = builder.finish();
+
+    assert_eq!(
+        built
+            .as_::<Struct>()
+            .unmasked_field(0)
+            .as_::<Chunked>()
+            .nchunks(),
+        2,
+        "a one-row field should still have earned a chunk",
+    );
+
+    let expected = ChunkedArray::try_new(vec![array.clone(), array], built.dtype().clone())?;
+    assert_arrays_eq!(&built, &expected, &mut ctx);
+
+    Ok(())
+}
+
 /// A non-canonical array of [`MIN_CHUNK_LEN`] `i32` values.
 fn constant_i32() -> ArrayRef {
     ConstantArray::new(0i32, MIN_CHUNK_LEN).into_array()