Avoid allocating Vec of chunks just to copy it in ChunkedArray constructor

encodings/fastlanes/benches/canonicalize_bench.rs

@@ -64,7 +64,7 @@ fn into_canonical_non_nullable(
     bencher
         .with_inputs(|| {
             (
-                ChunkedArray::from_iter(chunks.clone()).into_array(),
+                ChunkedArray::from_iter(chunks.iter().cloned()).into_array(),
                 SESSION.create_execution_ctx(),
             )
         })
@@ -94,7 +94,7 @@ fn canonical_into_non_nullable(
 
     bencher
         .with_inputs(|| {
-            let chunked = ChunkedArray::from_iter(chunks.clone()).into_array();
+            let chunked = ChunkedArray::from_iter(chunks.iter().cloned()).into_array();
             let ctx = SESSION.create_execution_ctx();
             let primitive_builder = PrimitiveBuilder::<i32>::with_capacity_in(
                 chunked.dtype().nullability(),
@@ -145,7 +145,7 @@ fn into_canonical_nullable(
     bencher
         .with_inputs(|| {
             (
-                ChunkedArray::from_iter(chunks.clone()).into_array(),
+                ChunkedArray::from_iter(chunks.iter().cloned()).into_array(),
                 SESSION.create_execution_ctx(),
             )
         })
@@ -175,7 +175,7 @@ fn canonical_into_nullable(
 
     bencher
         .with_inputs(|| {
-            let chunked = ChunkedArray::from_iter(chunks.clone()).into_array();
+            let chunked = ChunkedArray::from_iter(chunks.iter().cloned()).into_array();
             let ctx = SESSION.create_execution_ctx();
             let primitive_builder = PrimitiveBuilder::<i32>::with_capacity_in(
                 chunked.dtype().nullability(),

vortex-array/benches/take_chunked.rs

@@ -237,10 +237,8 @@ fn chunked_values(case: Case) -> ArrayRef {
             let start = chunk_idx * chunk_len;
             let end = start + chunk_len;
             chunk_values(case.value_kind, start, end)
-        })
-        .collect::<Vec<_>>();
-    let dtype = chunks[0].dtype().clone();
-    ChunkedArray::try_new(chunks, dtype).unwrap().into_array()
+        });
+    ChunkedArray::from_iter(chunks).into_array()
 }
 
 fn advance_random(state: &mut u64) -> u64 {

vortex-array/src/aggregate_fn/accumulator_grouped.rs

@@ -254,11 +254,14 @@ impl<V: AggregateFnVTable> DynGroupedAccumulator for GroupedAccumulator<V> {
     }
 
     fn flush(&mut self) -> VortexResult<ArrayRef> {
-        let mut states = std::mem::take(&mut self.partials);
-        if states.len() == 1 {
-            return Ok(states.pop().vortex_expect("checked one partial"));
+        if self.partials.len() == 1 {
+            return Ok(self.partials.pop().vortex_expect("checked one partial"));
         }
-        Ok(ChunkedArray::try_new(states, self.dtypes.partial_dtype.clone())?.into_array())
+
+        Ok(
+            ChunkedArray::try_new(self.partials.drain(..), self.dtypes.partial_dtype.clone())?
+                .into_array(),
+        )
     }
 
     fn finish(&mut self) -> VortexResult<ArrayRef> {

vortex-array/src/arrays/chunked/array.rs

@@ -170,9 +170,17 @@ impl Array<Chunked> {
     ) -> VortexResult<ArrayParts<Chunked>> {
         let chunks = chunks.into_iter();
         let (lower, _) = chunks.size_hint();
-        let mut slots = ArraySlots::with_capacity(ChunkedSlots::CHUNKS_OFFSET + lower);
+        Self::parts_from_chunks_with_capacity::<VALIDATE>(chunks, dtype, lower)
+    }
+
+    fn parts_from_chunks_with_capacity<const VALIDATE: bool>(
+        chunks: impl IntoIterator<Item = ArrayRef>,
+        dtype: DType,
+        capacity: usize,
+    ) -> VortexResult<ArrayParts<Chunked>> {
+        let mut slots = ArraySlots::with_capacity(ChunkedSlots::CHUNKS_OFFSET + capacity);
         slots.push(None);
-        let mut chunk_offsets = Vec::with_capacity(lower + 1);
+        let mut chunk_offsets = Vec::with_capacity(capacity + 1);
         chunk_offsets.push(0);
         let mut len = 0usize;
 
@@ -234,7 +242,10 @@ impl Array<Chunked> {
                 && !chunks_to_combine.is_empty()
             {
                 let canonical = unsafe {
-                    Array::<Chunked>::new_unchecked(chunks_to_combine, self.dtype().clone())
+                    Array::<Chunked>::new_unchecked(
+                        chunks_to_combine.drain(..),
+                        self.dtype().clone(),
+                    )
                 }
                 .into_array()
                 .execute::<Canonical>(ctx)?
@@ -243,7 +254,6 @@ impl Array<Chunked> {
 
                 new_chunk_n_bytes = 0;
                 new_chunk_n_elements = 0;
-                chunks_to_combine = Vec::new();
             }
 
             if n_bytes > target_bytesize || n_elements > target_rowsize {
@@ -273,17 +283,37 @@ impl Array<Chunked> {
     ///
     /// All chunks must have exactly the same [`DType`] as the provided `dtype`.
     pub unsafe fn new_unchecked(chunks: impl IntoIterator<Item = ArrayRef>, dtype: DType) -> Self {
-        let parts = Self::parts_from_chunks::<false>(chunks, dtype)
+        let chunks = chunks.into_iter();
+        let (expected_nchunks, _) = chunks.size_hint();
+        // SAFETY: the caller guarantees the dtype of every chunk.
+        unsafe { Self::new_unchecked_sized(chunks, dtype, expected_nchunks) }
+    }
+
+    /// Creates a chunked array without validation, reserving space for `expected_nchunks` chunks.
+    ///
+    /// The expected count is only an allocation hint: fewer or more chunks are allowed.
+    /// This is useful when iterator adapters such as `process_results` lose the source size hint.
+    ///
+    /// # Safety
+    ///
+    /// All chunks must have exactly the same [`DType`] as the provided `dtype`.
+    pub unsafe fn new_unchecked_sized(
+        chunks: impl IntoIterator<Item = ArrayRef>,
+        dtype: DType,
+        expected_nchunks: usize,
+    ) -> Self {
+        let parts = Self::parts_from_chunks_with_capacity::<false>(chunks, dtype, expected_nchunks)
             .vortex_expect("unchecked chunked construction cannot fail");
+        // SAFETY: the caller guarantees the dtype of every chunk.
         unsafe { Array::from_parts_unchecked(parts) }
     }
 }
 
 impl FromIterator<ArrayRef> for Array<Chunked> {
     fn from_iter<T: IntoIterator<Item = ArrayRef>>(iter: T) -> Self {
-        let chunks: Vec<ArrayRef> = iter.into_iter().collect();
+        let mut chunks = iter.into_iter().peekable();
         let dtype = chunks
-            .first()
+            .peek()
             .map(|c| c.dtype().clone())
             .vortex_expect("Cannot infer DType from an empty iterator");
         Array::<Chunked>::try_new(chunks, dtype)
@@ -293,8 +323,11 @@ impl FromIterator<ArrayRef> for Array<Chunked> {
 
 #[cfg(test)]
 mod test {
+    use itertools::Itertools;
+    use rstest::rstest;
     use vortex_buffer::buffer;
     use vortex_error::VortexResult;
+    use vortex_error::vortex_err;
 
     use crate::IntoArray;
     use crate::VortexSessionExecute;
@@ -308,6 +341,52 @@ mod test {
     use crate::dtype::PType;
     use crate::validity::Validity;
 
+    #[rstest]
+    fn sized_chunks(
+        #[values(0, 1, 17)] nchunks: usize,
+        #[values(0, 1, 17)] expected_nchunks: usize,
+    ) -> VortexResult<()> {
+        let dtype = DType::from(PType::U64);
+        let chunks = (0..nchunks as u64).map(|i| buffer![i].into_array());
+        // SAFETY: every chunk contains non-nullable u64 values.
+        let array = unsafe {
+            ChunkedArray::new_unchecked_sized(chunks, dtype.clone(), expected_nchunks)
+        };
+
+        assert_eq!(array.dtype(), &dtype);
+        assert_eq!(array.nchunks(), nchunks);
+        assert_eq!(array.chunk_offset_values(), (0..=nchunks).collect::<Vec<_>>());
+        let mut ctx = array_session().create_execution_ctx();
+        assert_arrays_eq!(array, PrimitiveArray::from_iter(0..nchunks as u64), &mut ctx);
+        Ok(())
+    }
+
+    #[rstest]
+    fn sized_chunks_with_process_results(#[values(0, 1)] error_at: usize) {
+        let mut visited = 0;
+        let chunks = (0..3).map(|i| {
+            visited += 1;
+            if i == error_at {
+                Err(vortex_err!("chunk source failed"))
+            } else {
+                Ok(buffer![1u64].into_array())
+            }
+        });
+        let dtype = DType::from(PType::U64);
+        let result = chunks.process_results(|chunks| {
+            // SAFETY: every successfully yielded chunk contains non-nullable u64 values.
+            unsafe { ChunkedArray::new_unchecked_sized(chunks, dtype, 3) }
+        });
+
+        assert!(
+            result
+                .unwrap_err()
+                .to_string()
+                .contains("chunk source failed")
+        );
+        assert_eq!(visited, error_at + 1);
+    }
+
     #[test]
     fn test_rechunk_one_chunk() {
         let mut ctx = array_session().create_execution_ctx();

vortex-array/src/arrays/chunked/compute/cast.rs

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use itertools::Itertools;
 use vortex_error::VortexResult;
 
 use crate::ArrayRef;
@@ -15,17 +16,14 @@ use crate::scalar_fn::fns::cast::CastReduce;
 
 impl CastReduce for Chunked {
     fn cast(array: ArrayView<'_, Chunked>, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
-        let mut cast_chunks = Vec::new();
-        for chunk in array.iter_chunks() {
-            cast_chunks.push(chunk.cast(dtype.clone())?);
-        }
-
-        // SAFETY: casting all chunks retains all chunks have same DType
-        unsafe {
-            Ok(Some(
-                ChunkedArray::new_unchecked(cast_chunks, dtype.clone()).into_array(),
-            ))
-        }
+        let chunks = array.iter_chunks().map(|chunk| chunk.cast(dtype.clone()));
+        chunks.process_results(|chunks| {
+            // SAFETY: every chunk is cast to the requested dtype.
+            Some(unsafe {
+                ChunkedArray::new_unchecked_sized(chunks, dtype.clone(), array.nchunks())
+                    .into_array()
+            })
+        })
     }
 }
 

vortex-array/src/arrays/chunked/compute/fill_null.rs

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use itertools::Itertools;
 use vortex_error::VortexResult;
 
 use crate::ArrayRef;
@@ -18,16 +19,20 @@ impl FillNullReduce for Chunked {
         array: ArrayView<'_, Chunked>,
         fill_value: &Scalar,
     ) -> VortexResult<Option<ArrayRef>> {
-        let new_chunks = array
+        let chunks = array
             .iter_chunks()
-            .map(|c| c.fill_null(fill_value.clone()))
-            .collect::<VortexResult<Vec<_>>>()?;
-
-        // SAFETY: wrapping each chunk in ScalarFnArray preserves the same DType across all chunks.
-        Ok(Some(
-            unsafe { ChunkedArray::new_unchecked(new_chunks, fill_value.dtype().clone()) }
-                .into_array(),
-        ))
+            .map(|c| c.fill_null(fill_value.clone()));
+        chunks.process_results(|chunks| {
+            // SAFETY: filling nulls gives every chunk the fill value's dtype.
+            Some(unsafe {
+                ChunkedArray::new_unchecked_sized(
+                    chunks,
+                    fill_value.dtype().clone(),
+                    array.nchunks(),
+                )
+                .into_array()
+            })
+        })
     }
 }
 

vortex-array/src/arrays/chunked/compute/mask.rs

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use itertools::Itertools;
 use vortex_error::VortexResult;
 
 use crate::ArrayRef;
@@ -20,20 +21,26 @@ impl MaskKernel for Chunked {
         _ctx: &mut ExecutionCtx,
     ) -> VortexResult<Option<ArrayRef>> {
         let chunk_offsets = array.chunk_offset_values();
-        let new_chunks: Vec<ArrayRef> = array
+        let chunks = array
             .iter_chunks()
             .enumerate()
             .map(|(i, chunk)| {
                 let start = chunk_offsets[i];
                 let end = chunk_offsets[i + 1];
                 let chunk_mask = mask.slice(start..end)?;
                 MaskExpr::try_new(chunk.clone(), chunk_mask).map(IntoArray::into_array)
+            });
+        chunks.process_results(|chunks| {
+            // SAFETY: masking makes every chunk's dtype nullable without changing its logical type.
+            Some(unsafe {
+                ChunkedArray::new_unchecked_sized(
+                    chunks,
+                    array.dtype().as_nullable(),
+                    array.nchunks(),
+                )
+                .into_array()
             })
-            .collect::<VortexResult<_>>()?;
-
-        Ok(Some(
-            ChunkedArray::try_new(new_chunks, array.dtype().as_nullable())?.into_array(),
-        ))
+        })
     }
 }
 

vortex-array/src/arrays/chunked/compute/rules.rs

@@ -45,18 +45,20 @@ impl ArrayParentReduceRule<Chunked> for ChunkedUnaryScalarFnPushDownRule {
             return Ok(None);
         }
 
-        let new_chunks: Vec<_> = array
+        let chunks = array
             .iter_chunks()
             .map(|chunk| {
                 ScalarFnArray::try_new(parent.scalar_fn().clone(), vec![chunk.clone()])?
                     .into_array()
                     .optimize()
+            });
+        chunks.process_results(|chunks| {
+            // SAFETY: applying the same scalar function gives every chunk the parent's dtype.
+            Some(unsafe {
+                ChunkedArray::new_unchecked_sized(chunks, parent.dtype().clone(), array.nchunks())
+                    .into_array()
             })
-            .try_collect()?;
-
-        Ok(Some(
-            unsafe { ChunkedArray::new_unchecked(new_chunks, parent.dtype().clone()) }.into_array(),
-        ))
+        })
     }
 }
 
@@ -81,7 +83,7 @@ impl ArrayParentReduceRule<Chunked> for ChunkedConstantScalarFnPushDownRule {
             }
         }
 
-        let new_chunks: Vec<_> = array
+        let chunks = array
             .iter_chunks()
             .map(|chunk| {
                 let new_children: Vec<_> = parent
@@ -103,11 +105,13 @@ impl ArrayParentReduceRule<Chunked> for ChunkedConstantScalarFnPushDownRule {
                 ScalarFnArray::try_new(parent.scalar_fn().clone(), new_children)?
                     .into_array()
                     .optimize()
+            });
+        chunks.process_results(|chunks| {
+            // SAFETY: applying the same scalar function gives every chunk the parent's dtype.
+            Some(unsafe {
+                ChunkedArray::new_unchecked_sized(chunks, parent.dtype().clone(), array.nchunks())
+                    .into_array()
             })
-            .try_collect()?;
-
-        Ok(Some(
-            unsafe { ChunkedArray::new_unchecked(new_chunks, parent.dtype().clone()) }.into_array(),
-        ))
+        })
     }
 }