diff --git a/encodings/fastlanes/benches/canonicalize_bench.rs b/encodings/fastlanes/benches/canonicalize_bench.rs index 7d8d7abb659..9edf223b8c6 100644 --- a/encodings/fastlanes/benches/canonicalize_bench.rs +++ b/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::::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::::with_capacity_in( chunked.dtype().nullability(), diff --git a/vortex-array/benches/take_chunked.rs b/vortex-array/benches/take_chunked.rs index 3350030533e..665a4aa372b 100644 --- a/vortex-array/benches/take_chunked.rs +++ b/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::>(); - 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 { diff --git a/vortex-array/src/aggregate_fn/accumulator_grouped.rs b/vortex-array/src/aggregate_fn/accumulator_grouped.rs index 631549bd036..295daa90671 100644 --- a/vortex-array/src/aggregate_fn/accumulator_grouped.rs +++ b/vortex-array/src/aggregate_fn/accumulator_grouped.rs @@ -254,11 +254,14 @@ impl DynGroupedAccumulator for GroupedAccumulator { } fn flush(&mut self) -> VortexResult { - 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 { diff --git a/vortex-array/src/arrays/chunked/array.rs b/vortex-array/src/arrays/chunked/array.rs index 807d526c17b..69bd15dd05b 100644 --- a/vortex-array/src/arrays/chunked/array.rs +++ b/vortex-array/src/arrays/chunked/array.rs @@ -170,9 +170,17 @@ impl Array { ) -> VortexResult> { 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::(chunks, dtype, lower) + } + + fn parts_from_chunks_with_capacity( + chunks: impl IntoIterator, + dtype: DType, + capacity: usize, + ) -> VortexResult> { + 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 { && !chunks_to_combine.is_empty() { let canonical = unsafe { - Array::::new_unchecked(chunks_to_combine, self.dtype().clone()) + Array::::new_unchecked( + chunks_to_combine.drain(..), + self.dtype().clone(), + ) } .into_array() .execute::(ctx)? @@ -243,7 +254,6 @@ impl Array { 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 { /// /// All chunks must have exactly the same [`DType`] as the provided `dtype`. pub unsafe fn new_unchecked(chunks: impl IntoIterator, dtype: DType) -> Self { - let parts = Self::parts_from_chunks::(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, + dtype: DType, + expected_nchunks: usize, + ) -> Self { + let parts = Self::parts_from_chunks_with_capacity::(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 for Array { fn from_iter>(iter: T) -> Self { - let chunks: Vec = 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::::try_new(chunks, dtype) @@ -293,8 +323,11 @@ impl FromIterator for Array { #[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::>()); + 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(); diff --git a/vortex-array/src/arrays/chunked/compute/cast.rs b/vortex-array/src/arrays/chunked/compute/cast.rs index 25e312de748..fb5dfde0121 100644 --- a/vortex-array/src/arrays/chunked/compute/cast.rs +++ b/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> { - 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() + }) + }) } } diff --git a/vortex-array/src/arrays/chunked/compute/fill_null.rs b/vortex-array/src/arrays/chunked/compute/fill_null.rs index b207a152ee0..0cafb1c0c6d 100644 --- a/vortex-array/src/arrays/chunked/compute/fill_null.rs +++ b/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> { - let new_chunks = array + let chunks = array .iter_chunks() - .map(|c| c.fill_null(fill_value.clone())) - .collect::>>()?; - - // 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() + }) + }) } } diff --git a/vortex-array/src/arrays/chunked/compute/mask.rs b/vortex-array/src/arrays/chunked/compute/mask.rs index f05988de05c..d708c899c16 100644 --- a/vortex-array/src/arrays/chunked/compute/mask.rs +++ b/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,7 +21,7 @@ impl MaskKernel for Chunked { _ctx: &mut ExecutionCtx, ) -> VortexResult> { let chunk_offsets = array.chunk_offset_values(); - let new_chunks: Vec = array + let chunks = array .iter_chunks() .enumerate() .map(|(i, chunk)| { @@ -28,12 +29,18 @@ impl MaskKernel for Chunked { 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::>()?; - - Ok(Some( - ChunkedArray::try_new(new_chunks, array.dtype().as_nullable())?.into_array(), - )) + }) } } diff --git a/vortex-array/src/arrays/chunked/compute/rules.rs b/vortex-array/src/arrays/chunked/compute/rules.rs index 712578255bb..90dbc9cd67e 100644 --- a/vortex-array/src/arrays/chunked/compute/rules.rs +++ b/vortex-array/src/arrays/chunked/compute/rules.rs @@ -45,18 +45,20 @@ impl ArrayParentReduceRule 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 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 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(), - )) + }) } } diff --git a/vortex-array/src/arrays/chunked/compute/slice.rs b/vortex-array/src/arrays/chunked/compute/slice.rs index 14919301bd1..73882537140 100644 --- a/vortex-array/src/arrays/chunked/compute/slice.rs +++ b/vortex-array/src/arrays/chunked/compute/slice.rs @@ -3,7 +3,6 @@ use std::ops::Range; -use itertools::Itertools; use vortex_error::VortexResult; use crate::ArrayRef; @@ -32,7 +31,7 @@ impl SliceKernel for Chunked { // SAFETY: empty chunked array trivially satisfies all validations unsafe { return Ok(Some( - ChunkedArray::new_unchecked(vec![], array.dtype().clone()).into_array(), + ChunkedArray::new_unchecked([], array.dtype().clone()).into_array(), )); } } @@ -47,18 +46,15 @@ impl SliceKernel for Chunked { )); } - let mut chunks = (offset_chunk..length_chunk + 1) - .map(|i| array.chunk(i).clone()) - .collect_vec(); - if let Some(c) = chunks.first_mut() { - *c = c.slice(offset_in_first_chunk..c.len())?; - } - - if length_in_last_chunk == 0 { - chunks.pop(); - } else if let Some(c) = chunks.last_mut() { - *c = c.slice(0..length_in_last_chunk)?; - } + let first = array.chunk(offset_chunk); + let first = first.slice(offset_in_first_chunk..first.len())?; + let middle = (offset_chunk + 1..length_chunk).map(|i| array.chunk(i).clone()); + let last = if length_in_last_chunk == 0 { + None + } else { + Some(array.chunk(length_chunk).slice(0..length_in_last_chunk)?) + }; + let chunks = std::iter::once(first).chain(middle).chain(last); // SAFETY: chunks are slices of the original valid chunks, preserving their dtype. // All chunks maintain the same dtype as the original array. diff --git a/vortex-array/src/arrays/chunked/compute/zip.rs b/vortex-array/src/arrays/chunked/compute/zip.rs index dc95c194173..0ef45cb5c55 100644 --- a/vortex-array/src/arrays/chunked/compute/zip.rs +++ b/vortex-array/src/arrays/chunked/compute/zip.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; @@ -30,17 +31,23 @@ impl ZipKernel for Chunked { let dtype = if_true .dtype() .union_nullability(if_false.dtype().nullability()); - let mut out_chunks = Vec::with_capacity(if_true.nchunks() + if_false.nchunks()); - - for pair in if_true.paired_chunks(&if_false) { - let pair = pair?; - let mask_slice = mask.slice(pair.pos)?; - out_chunks.push(mask_slice.zip(pair.left, pair.right)?); - } - - // SAFETY: chunks originate from zipping slices of inputs that share dtype/nullability. - let chunked = unsafe { ChunkedArray::new_unchecked(out_chunks, dtype) }; - Ok(Some(chunked.into_array())) + let chunks = if_true + .paired_chunks(&if_false) + .map(|pair| { + let pair = pair?; + mask.slice(pair.pos)?.zip(pair.left, pair.right) + }); + chunks.process_results(|chunks| { + // SAFETY: chunks originate from zipping slices of inputs that share dtype/nullability. + Some(unsafe { + ChunkedArray::new_unchecked_sized( + chunks, + dtype, + if_true.nchunks() + if_false.nchunks(), + ) + .into_array() + }) + }) } } diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index a521fd0e6cc..89dc01aa849 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -120,9 +120,6 @@ fn pack_variant_chunks( } Some(first_shredded) => { let shredded_dtype = first_shredded.dtype().clone(); - let mut shredded_chunks = Vec::with_capacity(variant_chunks.len()); - shredded_chunks.push(first_shredded.clone()); - for chunk in &variant_chunks[1..] { let shredded = chunk.shredded().ok_or_else(|| { vortex_err!( @@ -135,9 +132,14 @@ fn pack_variant_chunks( shredded_dtype, shredded.dtype() ); - shredded_chunks.push(shredded.clone()); } + let shredded_chunks = variant_chunks.iter().map(|chunk| { + chunk + .shredded() + .vortex_expect("validated shredded presence") + .clone() + }); Some(ChunkedArray::try_new(shredded_chunks, shredded_dtype)?.into_array()) } }; @@ -170,7 +172,6 @@ fn swizzle_list_chunks( // Since each list array in `chunks` has offsets local to each array, we can reuse the existing // array's child `elements` as the chunks and recompute offsets. - let mut list_elements_chunks = Vec::with_capacity(chunks.len()); let mut num_elements = 0; // TODO(connor)[ListView]: We could potentially choose a smaller type here, but that would make @@ -184,49 +185,49 @@ fn swizzle_list_chunks( let sizes_slice_out = sizes.as_mut_slice(); let mut next_list = 0usize; - for chunk in chunks { - let chunk_array = chunk.clone().execute::(ctx)?; - // By rebuilding as zero-copy to `List` and trimming all elements (to prevent gaps), we make - // the final output `ListView` also zero-copyable to `List`. - let chunk_array = chunk_array.rebuild(ListViewRebuildMode::MakeExact, ctx)?; - - // Add the `elements` of the current array as a new chunk. - list_elements_chunks.push(chunk_array.elements().clone()); - - // Cast offsets and sizes to `u64`. - let offsets_arr = chunk_array - .offsets() - .clone() - .cast(DType::Primitive(PType::U64, Nullability::NonNullable)) - .vortex_expect("Must be able to fit array offsets in u64") - .execute::(ctx)?; - - let sizes_arr = chunk_array - .sizes() - .clone() - .cast(DType::Primitive(PType::U64, Nullability::NonNullable)) - .vortex_expect("Must be able to fit array offsets in u64") - .execute::(ctx)?; - - let offsets_slice = offsets_arr.as_slice::(); - let sizes_slice = sizes_arr.as_slice::(); - - // Append offsets and sizes, adjusting offsets to point into the combined array. - for (&offset, &size) in offsets_slice.iter().zip(sizes_slice.iter()) { - offsets_out[next_list] = offset + num_elements; - sizes_slice_out[next_list] = size; - next_list += 1; - } + let element_chunks = chunks + .iter() + .map(|chunk| -> VortexResult<_> { + let chunk_array = chunk.clone().execute::(ctx)?; + // By rebuilding as zero-copy to `List` and trimming all elements (to prevent gaps), we make + // the final output `ListView` also zero-copyable to `List`. + let chunk_array = chunk_array.rebuild(ListViewRebuildMode::MakeExact, ctx)?; + + // Cast offsets and sizes to `u64`. + let offsets_arr = chunk_array + .offsets() + .clone() + .cast(DType::Primitive(PType::U64, Nullability::NonNullable)) + .vortex_expect("Must be able to fit array offsets in u64") + .execute::(ctx)?; + + let sizes_arr = chunk_array + .sizes() + .clone() + .cast(DType::Primitive(PType::U64, Nullability::NonNullable)) + .vortex_expect("Must be able to fit array offsets in u64") + .execute::(ctx)?; + + let offsets_slice = offsets_arr.as_slice::(); + let sizes_slice = sizes_arr.as_slice::(); + + // Append offsets and sizes, adjusting offsets to point into the combined array. + for (&offset, &size) in offsets_slice.iter().zip(sizes_slice.iter()) { + offsets_out[next_list] = offset + num_elements; + sizes_slice_out[next_list] = size; + next_list += 1; + } - num_elements += chunk_array.elements().len() as u64; - } + num_elements += chunk_array.elements().len() as u64; + Ok(chunk_array.elements().clone()) + }); + let chunked_elements = element_chunks.process_results(|elements| { + // SAFETY: elements come from valid ListView arrays with the same element dtype. + unsafe { ChunkedArray::new_unchecked_sized(elements, elem_dtype.clone(), chunks.len()) } + .into_array() + })?; debug_assert_eq!(next_list, len); - // SAFETY: elements are sliced from valid `ListViewArray`s (from `to_listview()`). - let chunked_elements = - unsafe { ChunkedArray::new_unchecked(list_elements_chunks, elem_dtype.clone()) } - .into_array(); - let offsets = PrimitiveArray::new(offsets.freeze(), Validity::NonNullable).into_array(); let sizes = PrimitiveArray::new(sizes.freeze(), Validity::NonNullable).into_array(); @@ -260,16 +261,20 @@ fn swizzle_fixed_size_list_chunks( ) -> VortexResult { let len: usize = chunks.iter().map(|c| c.len()).sum(); - let mut element_chunks = Vec::with_capacity(chunks.len()); - for chunk in chunks { - let chunk_array = chunk.clone().execute::(ctx)?; - // A canonical `FixedSizeListArray` keeps its `elements` child trimmed to exactly - // `list_size * chunk.len()` starting at the first list, so the children concatenate - // cleanly into the combined `elements` array. - element_chunks.push(chunk_array.elements().clone()); - } - - let chunked_elements = ChunkedArray::try_new(element_chunks, elem_dtype.clone())?.into_array(); + let element_chunks = chunks + .iter() + .map(|chunk| -> VortexResult<_> { + let chunk_array = chunk.clone().execute::(ctx)?; + // A canonical `FixedSizeListArray` keeps its `elements` child trimmed to exactly + // `list_size * chunk.len()` starting at the first list, so the children concatenate + // cleanly into the combined `elements` array. + Ok(chunk_array.elements().clone()) + }); + let chunked_elements = element_chunks.process_results(|elements| { + // SAFETY: every fixed-size-list chunk has the same element dtype. + unsafe { ChunkedArray::new_unchecked_sized(elements, elem_dtype.clone(), chunks.len()) } + .into_array() + })?; FixedSizeListArray::try_new(chunked_elements, list_size, validity, len) } diff --git a/vortex-array/src/arrays/constant/vtable/canonical.rs b/vortex-array/src/arrays/constant/vtable/canonical.rs index e1f9ad02f6c..d9edd51c655 100644 --- a/vortex-array/src/arrays/constant/vtable/canonical.rs +++ b/vortex-array/src/arrays/constant/vtable/canonical.rs @@ -386,7 +386,7 @@ fn tile_fixed_size_list_elements( // SAFETY: every chunk is `tile` itself, so they share its dtype and none is empty. unsafe { ChunkedArray::new_unchecked( - std::iter::repeat_n(tile, len).collect::>(), + std::iter::repeat_n(tile, len), element_dtype.clone(), ) } diff --git a/vortex-array/src/arrays/dict/compute/rules.rs b/vortex-array/src/arrays/dict/compute/rules.rs index d43c6deeb75..eac588a7ea5 100644 --- a/vortex-array/src/arrays/dict/compute/rules.rs +++ b/vortex-array/src/arrays/dict/compute/rules.rs @@ -64,7 +64,6 @@ impl ArrayParentReduceRule for DictionaryChunkedValuesPullUpRule { ) -> VortexResult> { let values = array.values(); let codes_dtype = array.codes().dtype().clone(); - let mut code_chunks = Vec::with_capacity(parent.nchunks()); let mut all_values_referenced = array.has_all_values_referenced(); for chunk in parent.iter_chunks() { @@ -78,9 +77,11 @@ impl ArrayParentReduceRule for DictionaryChunkedValuesPullUpRule { return Ok(None); } all_values_referenced |= dict.has_all_values_referenced(); - code_chunks.push(dict.codes().clone()); } + let code_chunks = parent + .iter_chunks() + .map(|chunk| chunk.as_::().codes().clone()); let codes = ChunkedArray::try_new(code_chunks, codes_dtype)?.into_array(); let dict = DictArray::try_new(codes, values.clone())?; let dict = if all_values_referenced { diff --git a/vortex-array/src/arrays/listview/compute/zip.rs b/vortex-array/src/arrays/listview/compute/zip.rs index d5e2a6bf406..db10e4b3732 100644 --- a/vortex-array/src/arrays/listview/compute/zip.rs +++ b/vortex-array/src/arrays/listview/compute/zip.rs @@ -6,6 +6,7 @@ use std::ops::BitAnd; use std::ops::BitOr; use std::ops::Not; +use itertools::Either; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; @@ -75,9 +76,7 @@ impl ZipKernel for ListView { // Concatenate the two `elements` arrays without copying. If either side is already a // `ChunkedArray` (e.g. the result of a previous list-view zip), splice its chunks in // directly rather than nesting chunked arrays. - let mut chunks = Vec::with_capacity(2); - push_element_chunks(true_elements, &mut chunks); - push_element_chunks(false_elements, &mut chunks); + let chunks = element_chunks(&true_elements).chain(element_chunks(&false_elements)); let elements = ChunkedArray::try_new(chunks, result_elements_dtype)?.into_array(); let true_offsets = to_u64(if_true.offsets(), ctx)?; @@ -189,12 +188,12 @@ fn select_column( } } -/// Appends `array`'s element chunks to `chunks`, flattening a top-level [`ChunkedArray`] so the +/// Iterates over `array`'s element chunks, flattening a top-level [`ChunkedArray`] so the /// concatenated elements never nest chunked arrays. -fn push_element_chunks(array: ArrayRef, chunks: &mut Vec) { +fn element_chunks(array: &ArrayRef) -> impl Iterator + '_ { match array.as_opt::() { - Some(chunked) => chunks.extend(chunked.iter_chunks().cloned()), - None => chunks.push(array), + Some(chunked) => Either::Left((0..chunked.nchunks()).map(move |i| chunked.chunk(i).clone())), + None => Either::Right(std::iter::once(array.clone())), } } diff --git a/vortex-array/src/arrays/variant/vtable/kernel.rs b/vortex-array/src/arrays/variant/vtable/kernel.rs index 22d52ecd6a7..cbc3543885a 100644 --- a/vortex-array/src/arrays/variant/vtable/kernel.rs +++ b/vortex-array/src/arrays/variant/vtable/kernel.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use itertools::Itertools; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; @@ -177,20 +178,22 @@ fn merge_typed_as_variant( let dtype = DType::Variant(Nullability::Nullable); // TODO(variant): replace this with a Variant builder once one exists. // Chunked canonicalizes to VariantArray, so this row-wise fallback is safe. - let mut chunks = Vec::with_capacity(typed.len()); - - for idx in 0..typed.len() { - let typed_scalar = typed.execute_scalar(idx, ctx)?; - let fallback_scalar = fallback - .as_ref() - .map(|fallback| fallback.execute_scalar(idx, ctx)) - .transpose()?; - let scalar = merge_typed_scalar_as_variant(typed_scalar, fallback_scalar, &dtype)?; - - chunks.push(ConstantArray::new(scalar, 1).into_array()); - } - - let core_storage = ChunkedArray::try_new(chunks, dtype)?.into_array(); + let chunks = (0..typed.len()) + .map(|idx| -> VortexResult<_> { + let typed_scalar = typed.execute_scalar(idx, ctx)?; + let fallback_scalar = fallback + .as_ref() + .map(|fallback| fallback.execute_scalar(idx, ctx)) + .transpose()?; + let scalar = merge_typed_scalar_as_variant(typed_scalar, fallback_scalar, &dtype)?; + + Ok(ConstantArray::new(scalar, 1).into_array()) + }); + let core_storage = chunks.process_results(|chunks| { + // SAFETY: each output scalar is constructed with the requested variant dtype. + unsafe { ChunkedArray::new_unchecked_sized(chunks, dtype.clone(), typed.len()) } + .into_array() + })?; VariantArray::try_new(core_storage, None).map(|array| array.into_array()) } diff --git a/vortex-array/src/builders/child.rs b/vortex-array/src/builders/child.rs index a9f5b46fc84..b0b028f9402 100644 --- a/vortex-array/src/builders/child.rs +++ b/vortex-array/src/builders/child.rs @@ -2,8 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_buffer::BufferAllocatorRef; -use vortex_error::VortexResult; use vortex_error::vortex_ensure; +use vortex_error::{VortexExpect, VortexResult}; use crate::ArrayRef; use crate::ExecutionCtx; @@ -134,12 +134,13 @@ impl ChildBuilder { self.flush_pending(); self.chunks_len = 0; - let mut chunks = std::mem::take(&mut self.chunks); - if chunks.len() == 1 { - return chunks.remove(0); + if self.chunks.len() == 1 { + return self.chunks.pop().vortex_expect("single chunk"); } - unsafe { ChunkedArray::new_unchecked(chunks, self.dtype.clone()) }.into_array() + // SAFETY: every accumulated chunk has the builder's dtype. + unsafe { ChunkedArray::new_unchecked(self.chunks.drain(..), self.dtype.clone()) } + .into_array() } /// Moves whatever the scalar builder holds into `chunks`, keeping the chunks in logical order. diff --git a/vortex-array/src/builders/fixed_size_list.rs b/vortex-array/src/builders/fixed_size_list.rs index e9a66b62d83..b5994336fbe 100644 --- a/vortex-array/src/builders/fixed_size_list.rs +++ b/vortex-array/src/builders/fixed_size_list.rs @@ -177,7 +177,7 @@ impl FixedSizeListBuilder { // SAFETY: every chunk is `array` itself, so they share its dtype and none is empty. let tiled = unsafe { ChunkedArray::new_unchecked( - std::iter::repeat_n(array.clone(), n).collect::>(), + std::iter::repeat_n(array.clone(), n), self.element_dtype().clone(), ) }; diff --git a/vortex-array/src/iter.rs b/vortex-array/src/iter.rs index b4cd01d875d..2a44a4b8df0 100644 --- a/vortex-array/src/iter.rs +++ b/vortex-array/src/iter.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use itertools::Itertools; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use crate::ArrayRef; use crate::IntoArray; @@ -53,6 +54,10 @@ where fn next(&mut self) -> Option { self.inner.next() } + + fn size_hint(&self) -> (usize, Option) { + self.inner.size_hint() + } } impl ArrayIterator for ArrayIteratorAdapter @@ -81,12 +86,26 @@ pub trait ArrayIteratorExt: ArrayIterator { Self: Sized, { let dtype = self.dtype().clone(); - let mut chunks: Vec = self.try_collect()?; - if chunks.len() == 1 { - Ok(chunks.remove(0)) - } else { - Ok(ChunkedArray::try_new(chunks, dtype)?.into_array()) + let mut chunks = self.peekable(); + let Some(first) = chunks.next().transpose()? else { + return Ok(ChunkedArray::try_new([], dtype)?.into_array()); + }; + if chunks.peek().is_none() { + return Ok(first); } + let chunks = std::iter::once(Ok(first)) + .chain(chunks) + .map(|chunk| -> VortexResult<_> { + let chunk = chunk?; + vortex_ensure!(chunk.dtype() == &dtype, MismatchedTypes: &dtype, chunk.dtype()); + Ok(chunk) + }); + let expected_nchunks = chunks.size_hint().0; + chunks.process_results(|chunks| { + // SAFETY: the iterator validates the dtype of every successfully yielded chunk. + unsafe { ChunkedArray::new_unchecked_sized(chunks, dtype.clone(), expected_nchunks) } + .into_array() + }) } } diff --git a/vortex-array/src/scalar_fn/fns/variant_get/mod.rs b/vortex-array/src/scalar_fn/fns/variant_get/mod.rs index db3f795f5a4..a20ff1d4510 100644 --- a/vortex-array/src/scalar_fn/fns/variant_get/mod.rs +++ b/vortex-array/src/scalar_fn/fns/variant_get/mod.rs @@ -5,6 +5,7 @@ use std::fmt; use std::fmt::Display; use std::fmt::Formatter; +use itertools::Itertools; use prost::Message; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -156,15 +157,17 @@ impl ScalarFnVTable for VariantGet { // TODO(variant): replace this with a Variant builder once one exists. // Chunked canonicalizes to VariantArray, so this row-wise fallback is safe. - let mut chunks = Vec::with_capacity(input.len()); - - for idx in 0..input.len() { - let scalar = input.execute_scalar(idx, ctx)?; - let output = variant_get_scalar(&scalar, options, &dtype)?; - chunks.push(ConstantArray::new(output, 1).into_array()); - } - - let array = ChunkedArray::try_new(chunks, dtype)?.into_array(); + let chunks = (0..input.len()) + .map(|idx| -> VortexResult<_> { + let scalar = input.execute_scalar(idx, ctx)?; + let output = variant_get_scalar(&scalar, options, &dtype)?; + Ok(ConstantArray::new(output, 1).into_array()) + }); + let array = chunks.process_results(|chunks| { + // SAFETY: each output scalar is constructed with the requested variant dtype. + unsafe { ChunkedArray::new_unchecked_sized(chunks, dtype.clone(), input.len()) } + .into_array() + })?; VariantArray::try_new(array, None).map(|array| array.into_array()) } diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index 6b4ed871f2e..b3081d6be49 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -10,6 +10,7 @@ use arrow_array::RecordBatch; use arrow_select::concat::concat_batches; use futures::StreamExt; use futures::TryStreamExt; +use itertools::Itertools; use parquet::arrow::AsyncArrowWriter; use parquet::arrow::ParquetRecordBatchStreamBuilder; use parquet::arrow::async_reader::ParquetRecordBatchStream; @@ -142,13 +143,18 @@ pub async fn parquet_to_vortex_chunks_with_batch_size( .ok_or_else(|| anyhow::anyhow!("cannot convert an empty Parquet file"))?; let combined = concat_batches(&schema, &batches)?; - let mut chunks = Vec::with_capacity(combined.num_rows().div_ceil(batch_size)); - for start in (0..combined.num_rows()).step_by(batch_size) { - let len = batch_size.min(combined.num_rows() - start); - chunks.push(record_batch_to_vortex(combined.slice(start, len))?); - } - - Ok(ChunkedArray::from_iter(chunks)) + let dtype = SESSION.arrow().from_arrow_schema(schema.as_ref())?; + let chunks = (0..combined.num_rows()) + .step_by(batch_size) + .map(|start| { + let len = batch_size.min(combined.num_rows() - start); + record_batch_to_vortex(combined.slice(start, len)) + }); + let expected_nchunks = combined.num_rows().div_ceil(batch_size); + Ok(chunks.process_results(|chunks| { + // SAFETY: every batch is a slice of `combined` and is converted using the same schema. + unsafe { ChunkedArray::new_unchecked_sized(chunks, dtype, expected_nchunks) } + })?) } /// Convert one Arrow [`RecordBatch`] into a canonical Vortex array. diff --git a/vortex-bench/src/vector_dataset/convert.rs b/vortex-bench/src/vector_dataset/convert.rs index a0125bd923b..136486cfb06 100644 --- a/vortex-bench/src/vector_dataset/convert.rs +++ b/vortex-bench/src/vector_dataset/convert.rs @@ -3,6 +3,7 @@ // TODO(connor): Should we re-export this through `conversions.rs`? +use itertools::Itertools; use vortex::array::ArrayRef; use vortex::array::EmptyMetadata; use vortex::array::IntoArray; @@ -24,6 +25,7 @@ use vortex::dtype::extension::ExtDType; use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_bail; +use vortex::error::vortex_ensure; use vortex::error::vortex_err; use vortex_tensor::vector::Vector; @@ -62,17 +64,27 @@ use crate::SESSION; /// - The input has zero rows (the dimension cannot be inferred from empty input). pub fn list_to_vector_ext(input: ArrayRef) -> VortexResult { if let Some(chunked) = input.as_opt::() { - let converted: Vec = chunked + let mut converted = chunked .iter_chunks() - .map(|chunk| list_to_vector_ext(chunk.clone())) - .collect::>()?; + .map(|chunk| list_to_vector_ext(chunk.clone())); - let Some(first) = converted.first() else { + let Some(first) = converted.next().transpose()? else { vortex_bail!("list_to_vector_ext: chunked input has no chunks"); }; let dtype = first.dtype().clone(); - return Ok(ChunkedArray::try_new(converted, dtype)?.into_array()); + let chunks = std::iter::once(Ok(first)) + .chain(converted) + .map(|chunk| -> VortexResult<_> { + let chunk = chunk?; + vortex_ensure!(chunk.dtype() == &dtype, MismatchedTypes: &dtype, chunk.dtype()); + Ok(chunk) + }); + return chunks.process_results(|chunks| { + // SAFETY: the iterator checks that all chunks have the same vector dtype and dimension. + unsafe { ChunkedArray::new_unchecked_sized(chunks, dtype.clone(), chunked.nchunks()) } + .into_array() + }); } // `parquet_to_vortex_chunks` produces `ListView` arrays for list columns by default; diff --git a/vortex-compressor/src/compressor/sample.rs b/vortex-compressor/src/compressor/sample.rs index ba1271d4a6d..d3318cf186e 100644 --- a/vortex-compressor/src/compressor/sample.rs +++ b/vortex-compressor/src/compressor/sample.rs @@ -50,14 +50,11 @@ pub(crate) fn sample(input: &ArrayRef, sample_size: u32, sample_count: u32) -> A ); // For every slice, grab the relevant slice and repack into a new PrimitiveArray. - let chunks: Vec<_> = slices - .into_iter() - .map(|(start, end)| { - input - .slice(start..end) - .vortex_expect("slice should succeed") - }) - .collect(); + let chunks = slices.into_iter().map(|(start, end)| { + input + .slice(start..end) + .vortex_expect("slice should succeed") + }); // SAFETY: all chunks are slices of `input`, so they share its dtype. unsafe { ChunkedArray::new_unchecked(chunks, input.dtype().clone()) }.into_array() } diff --git a/vortex-file/benches/split_collection.rs b/vortex-file/benches/split_collection.rs index 2c4ccd0053e..c629a7d2a79 100644 --- a/vortex-file/benches/split_collection.rs +++ b/vortex-file/benches/split_collection.rs @@ -77,8 +77,7 @@ fn make_file(columns: usize, chunks: usize) -> VortexFile { }) .collect::>(); StructArray::from_fields(&fields).unwrap().into_array() - }) - .collect::>(); + }); let array = ChunkedArray::from_iter(struct_chunks).into_array(); let strategy = vortex_file::WriteStrategyBuilder::default() diff --git a/vortex-python/src/arrays/from_arrow.rs b/vortex-python/src/arrays/from_arrow.rs index 1077d6d3aaa..56cf3dcdf39 100644 --- a/vortex-python/src/arrays/from_arrow.rs +++ b/vortex-python/src/arrays/from_arrow.rs @@ -7,13 +7,15 @@ use arrow_array::make_array; use arrow_data::ArrayData as ArrowArrayData; use arrow_schema::DataType; use arrow_schema::Field; +use itertools::Itertools; use pyo3::exceptions::PyValueError; use pyo3::intern; use pyo3::prelude::*; use vortex::array::IntoArray; use vortex::array::arrays::ChunkedArray; use vortex::error::VortexError; -use vortex::error::VortexResult; +use vortex::error::vortex_ensure; +use vortex::error::vortex_err; use vortex_arrow::ArrowSessionExt; use crate::arrays::PyArrayRef; @@ -41,16 +43,6 @@ pub(super) fn from_arrow(obj: &Borrowed<'_, '_, PyAny>) -> PyVortexResult> = obj.getattr(intern!(py, "chunks"))?.extract()?; - let encoded_chunks = chunks - .iter() - .map(|a| { - let arrow_array = ArrowArrayData::from_pyarrow(&a.as_borrowed()).map(make_array)?; - session() - .arrow() - .from_arrow_array(arrow_array, false) - .map_err(PyVortexError::from) - }) - .collect::>>()?; let arrow_dtype = obj .getattr(intern!(py, "type")) .and_then(|v| DataType::from_pyarrow(&v.as_borrowed()))?; @@ -58,27 +50,47 @@ pub(super) fn from_arrow(obj: &Borrowed<'_, '_, PyAny>) -> PyVortexResult PyVortexResult<_> { + let arrow_array = ArrowArrayData::from_pyarrow(&a.as_borrowed()).map(make_array)?; + let encoded = session() + .arrow() + .from_arrow_array(arrow_array, false) + .map_err(PyVortexError::from)?; + if encoded.dtype() != &dtype { + return Err(vortex_err!(MismatchedTypes: &dtype, encoded.dtype()).into()); + } + Ok(encoded) + }); + let expected_nchunks = encoded_chunks.size_hint().0; + let array = encoded_chunks.process_results(|chunks| { + // SAFETY: the iterator validates each converted chunk against the declared dtype. + unsafe { ChunkedArray::new_unchecked_sized(chunks, dtype.clone(), expected_nchunks) } + })?; + Ok(PyArrayRef::from(array.into_array())) } else if obj.is_instance(table)? { let array_stream = ArrowArrayStreamReader::from_pyarrow(&obj.as_borrowed())?; let dtype = session() .arrow() .from_arrow_schema(array_stream.schema().as_ref()) .map_err(|e| PyValueError::new_err(e.to_string()))?; - let chunks = array_stream + let encoded_chunks = array_stream .into_iter() .map(|b| { b.map_err(VortexError::from).and_then(|b| { let schema = b.schema(); - session().arrow().from_arrow_record_batch(b, &schema) + let encoded = session().arrow().from_arrow_record_batch(b, &schema)?; + vortex_ensure!(encoded.dtype() == &dtype, MismatchedTypes: &dtype, encoded.dtype()); + Ok(encoded) }) - }) - .collect::>>()?; - Ok(PyArrayRef::from( - ChunkedArray::try_new(chunks, dtype)?.into_array(), - )) + }); + let expected_nchunks = encoded_chunks.size_hint().0; + let array = encoded_chunks.process_results(|chunks| { + // SAFETY: the iterator validates each converted batch against the stream's dtype. + unsafe { ChunkedArray::new_unchecked_sized(chunks, dtype.clone(), expected_nchunks) } + })?; + Ok(PyArrayRef::from(array.into_array())) } else { Err(PyValueError::new_err("Cannot convert object to Vortex array").into()) }