diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs index 106abcd88a9..89fc3549222 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs @@ -350,6 +350,7 @@ mod tests { use crate::arrays::UnionArray; use crate::arrays::VarBinViewArray; use crate::arrays::VariantArray; + use crate::arrays::listview::ListViewRebuildMode; use crate::builders::builder_with_capacity; use crate::dtype::DType; use crate::dtype::DecimalDType; @@ -506,9 +507,21 @@ mod tests { let array = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable).into_array(); + // These lists are out of order and leave element 1 unreferenced. `ListViewBuilder` keeps + // the layout it is handed, so the builder round-trip inside + // `materialized_uncompressed_size_in_bytes` retains that element and no longer stands in + // for the logical size. Rebuild to the exact layout, which is what "materialized" means + // here. + let mut ctx = array_session().create_execution_ctx(); + let exact = array + .clone() + .execute::(&mut ctx)? + .rebuild(ListViewRebuildMode::MakeExact, &mut ctx)? + .into_array(); + assert_eq!( aggregate(&array)?, - materialized_uncompressed_size_in_bytes(&array) + materialized_uncompressed_size_in_bytes(&exact) ); Ok(()) } diff --git a/vortex-array/src/builders/child.rs b/vortex-array/src/builders/child.rs new file mode 100644 index 00000000000..ba374194ebd --- /dev/null +++ b/vortex-array/src/builders/child.rs @@ -0,0 +1,561 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_mask::Mask; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::ChunkedArray; +use crate::arrays::MaskedArray; +use crate::builders::ArrayBuilder; +use crate::builders::builder_with_capacity; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::scalar::Scalar; +use crate::validity::Validity; + +/// 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. +/// +/// 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. +/// +/// Nested builders receive values from two sources: individual [`Scalar`]s, which have to be +/// materialized into a canonical builder, and whole arrays, whose encoding the builder has no +/// reason to decode. A `ChildBuilder` keeps the latter as chunks and materializes only the former, +/// stitching everything back together into a [`ChunkedArray`] on [`finish`](Self::finish) when +/// more than one chunk accumulated. +/// +/// This keeps canonical arrays canonical only at the top level, which is all that [`Canonical`] +/// promises: the fields of a `StructArray`, the elements of a list, and the storage of an +/// extension array may all stay compressed. +/// +/// [`Canonical`]: crate::Canonical +pub struct ChildBuilder { + /// The [`DType`] shared by every chunk and by the scalar builder. + dtype: DType, + + /// Completed chunks, in logical order. Never contains an empty chunk. + chunks: Vec, + + /// The summed length of `chunks`. + chunks_len: usize, + + /// Builder holding the scalars appended after the last chunk. + pending: Box, + + /// The shortest appended array that is kept as a chunk rather than copied. + min_chunk_len: usize, +} + +impl ChildBuilder { + /// Creates a new `ChildBuilder` whose scalar builder is pre-allocated for `capacity` values. + pub fn with_capacity(dtype: &DType, capacity: usize) -> Self { + Self { + dtype: dtype.clone(), + 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() + } + + /// Appends every value of `array` to the child. + /// + /// Arrays long enough to earn a chunk are kept in their original encoding; shorter ones are + /// materialized into the scalar builder. + pub fn append_array(&mut self, array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<()> { + vortex_ensure!( + array.dtype() == &self.dtype, + "Cannot append an array of dtype {} to a child builder of dtype {}", + array.dtype(), + self.dtype, + ); + + if array.is_empty() { + return Ok(()); + } + + if array.len() < self.min_chunk_len { + return array.append_to_builder(self.pending.as_mut(), ctx); + } + + self.flush_pending(); + self.chunks_len += array.len(); + self.chunks.push(array.clone()); + + Ok(()) + } + + /// Appends a single [`Scalar`] to the child. + pub fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> { + self.pending.append_scalar(scalar) + } + + /// Appends `n` "zero" values to the child. + /// + /// See [`ArrayBuilder::append_zeros`]. + pub fn append_zeros(&mut self, n: usize) { + self.pending.append_zeros(n) + } + + /// Appends `n` null values to the child. + /// + /// See [`ArrayBuilder::append_nulls`]. + pub fn append_nulls(&mut self, n: usize) { + self.pending.append_nulls(n) + } + + /// Appends `n` default values to the child. + /// + /// See [`ArrayBuilder::append_defaults`]. + pub fn append_defaults(&mut self, n: usize) { + self.pending.append_defaults(n) + } + + /// Allocates space for `additional` more values in the scalar builder. + pub fn reserve_exact(&mut self, additional: usize) { + self.pending.reserve_exact(additional) + } + + /// Overrides the validity of every value appended so far. + /// + /// # Safety + /// + /// `validity` must have the same length as [`self.len()`](Self::len). + /// + /// # Panics + /// + /// Panics if a chunk that was kept in its original encoding contains nulls, since replacing + /// the validity of such a chunk would require decoding it. + pub unsafe fn set_validity_unchecked(&mut self, validity: Mask) { + if !self.dtype.is_nullable() { + return; + } + + if self.chunks.is_empty() { + // Fast path: every value lives in the scalar builder, which owns its null buffer. + unsafe { self.pending.set_validity_unchecked(validity) }; + return; + } + + // The chunks carry their own validity, so the override has to be pushed into each of them. + self.flush_pending(); + let mut offset = 0; + for chunk in &mut self.chunks { + let end = offset + chunk.len(); + *chunk = MaskedArray::try_new( + chunk.clone(), + Validity::from_mask(validity.slice(offset..end), Nullability::Nullable), + ) + .vortex_expect("cannot override the validity of a child chunk that contains nulls") + .into_array(); + offset = end; + } + } + + /// Finishes the child, combining the accumulated chunks into a [`ChunkedArray`] when there is + /// more than one of them. + pub fn finish(&mut self) -> ArrayRef { + if self.chunks.is_empty() { + return self.pending.finish(); + } + + self.flush_pending(); + self.chunks_len = 0; + + let mut chunks = std::mem::take(&mut self.chunks); + if chunks.len() == 1 { + return chunks.remove(0); + } + + ChunkedArray::try_new(chunks, self.dtype.clone()) + .vortex_expect("every child chunk has the child dtype") + .into_array() + } + + /// Moves whatever the scalar builder holds into `chunks`, keeping the chunks in logical order. + fn flush_pending(&mut self) { + if self.pending.is_empty() { + return; + } + self.chunks_len += self.pending.len(); + let pending = self.pending.finish(); + self.chunks.push(pending); + } +} + +#[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::DEFAULT_MIN_CHUNK_LEN as MIN_CHUNK_LEN; + use crate::ArrayRef; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::arrays::Chunked; + 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 { + ConstantArray::new(value, len).into_array() + } + + /// A non-canonical *nullable* array of `len` non-null values, all equal to `value`. + fn nullable_constant(value: i32, len: usize) -> ArrayRef { + ConstantArray::new(Scalar::primitive(value, Nullable), len).into_array() + } + + #[test] + fn test_long_arrays_are_kept_as_chunks() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = ChildBuilder::with_capacity(&DType::from(I32), 0); + + builder.append_array(&constant(1, MIN_CHUNK_LEN), &mut ctx)?; + builder.append_array(&constant(2, MIN_CHUNK_LEN), &mut ctx)?; + assert_eq!(builder.len(), 2 * MIN_CHUNK_LEN); + + let child = builder.finish(); + let chunked = child.as_::(); + assert_eq!(chunked.nchunks(), 2); + // The chunks were never decoded. + assert!(chunked.iter_chunks().all(|c| c.is::())); + + Ok(()) + } + + #[test] + fn test_short_arrays_are_materialized() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = ChildBuilder::with_capacity(&DType::from(I32), 0); + + builder.append_array(&constant(1, MIN_CHUNK_LEN - 1), &mut ctx)?; + assert_eq!(builder.len(), MIN_CHUNK_LEN - 1); + + let child = builder.finish(); + assert!(child.is::()); + + Ok(()) + } + + #[test] + fn test_scalars_interleaved_with_chunks_keep_their_order() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = ChildBuilder::with_capacity(&DType::from(I32), 0); + + builder.append_scalar(&1i32.into())?; + builder.append_array(&constant(2, MIN_CHUNK_LEN), &mut ctx)?; + builder.append_scalar(&3i32.into())?; + + let child = builder.finish(); + assert_eq!(child.len(), MIN_CHUNK_LEN + 2); + + let expected = ChunkedArray::try_new( + vec![ + buffer![1i32].into_array(), + constant(2, MIN_CHUNK_LEN), + buffer![3i32].into_array(), + ], + DType::from(I32), + )? + .into_array(); + assert_arrays_eq!(&child, &expected, &mut ctx); + + Ok(()) + } + + #[test] + fn test_single_chunk_is_not_wrapped() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = ChildBuilder::with_capacity(&DType::from(I32), 0); + + builder.append_array(&constant(7, MIN_CHUNK_LEN), &mut ctx)?; + + let child = builder.finish(); + assert!(child.is::()); + + Ok(()) + } + + #[test] + fn test_empty_arrays_never_become_chunks() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = ChildBuilder::with_capacity(&DType::from(I32), 0); + let empty = constant(1, MIN_CHUNK_LEN).slice(0..0)?; + + builder.append_array(&empty, &mut ctx)?; + builder.append_array(&constant(1, MIN_CHUNK_LEN), &mut ctx)?; + builder.append_array(&empty, &mut ctx)?; + builder.append_array(&constant(2, MIN_CHUNK_LEN), &mut ctx)?; + builder.append_array(&empty, &mut ctx)?; + + assert_eq!(builder.len(), 2 * MIN_CHUNK_LEN); + assert_eq!(builder.finish().as_::().nchunks(), 2); + + Ok(()) + } + + /// An empty child must still finish as an empty array rather than as an empty [`ChunkedArray`]. + #[test] + fn test_empty_child_finishes_without_chunks() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = ChildBuilder::with_capacity(&DType::from(I32), 0); + + builder.append_array(&constant(1, MIN_CHUNK_LEN).slice(0..0)?, &mut ctx)?; + + let child = builder.finish(); + assert!(child.is_empty()); + assert!(child.is::()); + + Ok(()) + } + + /// The dtype check has to run before the length check, so that a mismatched array is rejected + /// whether it would have been chunked or materialized. + #[rstest] + #[case::would_be_materialized(MIN_CHUNK_LEN - 1)] + #[case::would_be_chunked(MIN_CHUNK_LEN)] + fn test_appending_a_mismatched_dtype_is_rejected(#[case] len: usize) { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = ChildBuilder::with_capacity(&DType::from(I32), 0); + + let wrong_dtype = ConstantArray::new(1i64, len).into_array(); + assert!(builder.append_array(&wrong_dtype, &mut ctx).is_err()); + } + + /// Everything the scalar builder can produce has to be flushed ahead of the next chunk. + #[test] + fn test_zeros_and_nulls_around_chunks_keep_their_order() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DType::Primitive(I32, Nullable); + let mut builder = ChildBuilder::with_capacity(&dtype, 0); + + builder.append_array(&nullable_constant(1, MIN_CHUNK_LEN), &mut ctx)?; + builder.append_nulls(2); + builder.append_array(&nullable_constant(2, MIN_CHUNK_LEN), &mut ctx)?; + builder.append_zeros(1); + + let child = builder.finish(); + assert_eq!(child.len(), 2 * MIN_CHUNK_LEN + 3); + assert_eq!(child.as_::().nchunks(), 4); + + let expected = PrimitiveArray::from_option_iter( + std::iter::repeat_n(Some(1i32), MIN_CHUNK_LEN) + .chain([None, None]) + .chain(std::iter::repeat_n(Some(2i32), MIN_CHUNK_LEN)) + .chain([Some(0)]), + ) + .into_array(); + assert_arrays_eq!(&child, &expected, &mut ctx); + + Ok(()) + } + + /// Overriding the validity once chunks exist has to push the override into each chunk, sliced + /// to that chunk's own range. + #[test] + fn test_set_validity_pushes_the_override_into_every_chunk() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DType::Primitive(I32, Nullable); + let mut builder = ChildBuilder::with_capacity(&dtype, 0); + + builder.append_array(&nullable_constant(1, MIN_CHUNK_LEN), &mut ctx)?; + builder.append_array(&nullable_constant(2, MIN_CHUNK_LEN), &mut ctx)?; + + // Straddle the chunk boundary, so an override sliced wrongly cannot pass. + let invalid = [MIN_CHUNK_LEN - 1, MIN_CHUNK_LEN]; + let validity = Mask::from_iter((0..2 * MIN_CHUNK_LEN).map(|i| !invalid.contains(&i))); + unsafe { builder.set_validity_unchecked(validity) }; + + let child = builder.finish(); + let chunked = child.as_::(); + assert_eq!(chunked.nchunks(), 2); + // The override was layered over the chunks rather than decoding them. + assert!(chunked.iter_chunks().all(|chunk| chunk.is::())); + + let expected = + PrimitiveArray::from_option_iter((0..2 * MIN_CHUNK_LEN).map(|i| { + (!invalid.contains(&i)).then_some(if i < MIN_CHUNK_LEN { 1i32 } else { 2 }) + })) + .into_array(); + assert_arrays_eq!(&child, &expected, &mut ctx); + + Ok(()) + } + + /// Values still sitting in the scalar builder are part of the override too. + #[test] + fn test_set_validity_covers_pending_scalars() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DType::Primitive(I32, Nullable); + let mut builder = ChildBuilder::with_capacity(&dtype, 0); + + builder.append_array(&nullable_constant(1, MIN_CHUNK_LEN), &mut ctx)?; + builder.append_scalar(&Scalar::primitive(2i32, Nullable))?; + + let mut validity = vec![true; MIN_CHUNK_LEN + 1]; + validity[MIN_CHUNK_LEN] = false; + unsafe { builder.set_validity_unchecked(Mask::from_iter(validity)) }; + + let child = builder.finish(); + let expected = PrimitiveArray::from_option_iter( + std::iter::repeat_n(Some(1i32), MIN_CHUNK_LEN).chain([None]), + ) + .into_array(); + assert_arrays_eq!(&child, &expected, &mut ctx); + + Ok(()) + } + + /// A non-nullable child cannot carry nulls, so the override is dropped and the chunks are left + /// exactly as they were appended. + #[test] + fn test_set_validity_is_a_noop_for_a_non_nullable_child() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = ChildBuilder::with_capacity(&DType::from(I32), 0); + + builder.append_array(&constant(1, MIN_CHUNK_LEN), &mut ctx)?; + builder.append_array(&constant(2, MIN_CHUNK_LEN), &mut ctx)?; + unsafe { builder.set_validity_unchecked(Mask::new_false(2 * MIN_CHUNK_LEN)) }; + + let child = builder.finish(); + let chunked = child.as_::(); + assert!(chunked.iter_chunks().all(|chunk| chunk.is::())); + + let expected = ChunkedArray::try_new( + vec![constant(1, MIN_CHUNK_LEN), constant(2, MIN_CHUNK_LEN)], + DType::from(I32), + )? + .into_array(); + assert_arrays_eq!(&child, &expected, &mut ctx); + + Ok(()) + } + + /// Replacing the validity of a chunk that already contains nulls would mean decoding it, which + /// is exactly what the chunk exists to avoid. + #[test] + #[should_panic(expected = "cannot override the validity of a child chunk that contains nulls")] + fn test_set_validity_rejects_a_chunk_that_contains_nulls() { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DType::Primitive(I32, Nullable); + let mut builder = ChildBuilder::with_capacity(&dtype, 0); + + let with_nulls = ConstantArray::new(Scalar::null(dtype), MIN_CHUNK_LEN).into_array(); + builder + .append_array(&with_nulls, &mut ctx) + .vortex_expect("append"); + + 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_::(); + assert_eq!(chunked.nchunks(), 2); + assert!(chunked.iter_chunks().all(|chunk| chunk.is::())); + + 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_::().elements().is::()); + + Ok(()) + } + + #[test] + fn test_finish_resets_the_builder() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = ChildBuilder::with_capacity(&DType::from(I32), 0); + + builder.append_array(&constant(1, MIN_CHUNK_LEN), &mut ctx)?; + builder.append_scalar(&2i32.into())?; + assert_eq!(builder.finish().len(), MIN_CHUNK_LEN + 1); + + assert_eq!(builder.len(), 0); + builder.append_scalar(&3i32.into())?; + + let expected = PrimitiveArray::new(buffer![3i32], NonNullable.into()).into_array(); + assert_arrays_eq!(&builder.finish(), &expected, &mut ctx); + + Ok(()) + } +} diff --git a/vortex-array/src/builders/extension.rs b/vortex-array/src/builders/extension.rs index aa8e1b76aa2..1b2d56fec8c 100644 --- a/vortex-array/src/builders/extension.rs +++ b/vortex-array/src/builders/extension.rs @@ -13,8 +13,8 @@ use crate::IntoArray; use crate::arrays::ExtensionArray; use crate::arrays::extension::ExtensionArrayExt; use crate::builders::ArrayBuilder; +use crate::builders::ChildBuilder; use crate::builders::DEFAULT_BUILDER_CAPACITY; -use crate::builders::builder_with_capacity; use crate::canonical::Canonical; use crate::dtype::DType; use crate::dtype::extension::ExtDTypeRef; @@ -24,7 +24,7 @@ use crate::scalar::Scalar; /// The builder for building a [`ExtensionArray`]. pub struct ExtensionBuilder { dtype: DType, - storage: Box, + storage: ChildBuilder, } impl ExtensionBuilder { @@ -36,7 +36,7 @@ impl ExtensionBuilder { /// Creates a new `ExtensionBuilder` with the given `capacity`. pub fn with_capacity(ext_dtype: ExtDTypeRef, capacity: usize) -> Self { Self { - storage: builder_with_capacity(ext_dtype.storage_dtype(), capacity), + storage: ChildBuilder::with_capacity(ext_dtype.storage_dtype(), capacity), dtype: DType::Extension(ext_dtype), } } @@ -53,9 +53,7 @@ impl ExtensionBuilder { array: &ExtensionArray, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - array - .storage_array() - .append_to_builder(self.storage.as_mut(), ctx) + self.storage.append_array(array.storage_array(), ctx) } /// Finishes the builder directly into a [`ExtensionArray`]. @@ -112,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) } diff --git a/vortex-array/src/builders/fixed_size_list.rs b/vortex-array/src/builders/fixed_size_list.rs index 7cd0ac692f0..1c6f001dcf1 100644 --- a/vortex-array/src/builders/fixed_size_list.rs +++ b/vortex-array/src/builders/fixed_size_list.rs @@ -17,9 +17,9 @@ use crate::IntoArray; use crate::arrays::FixedSizeListArray; use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use crate::builders::ArrayBuilder; +use crate::builders::ChildBuilder; use crate::builders::DEFAULT_BUILDER_CAPACITY; use crate::builders::LazyBitBufferBuilder; -use crate::builders::builder_with_capacity; use crate::canonical::Canonical; use crate::dtype::DType; use crate::dtype::Nullability; @@ -34,7 +34,7 @@ pub struct FixedSizeListBuilder { /// The builder for the underlying elements of the [`FixedSizeListArray`]. /// /// This builder will have a capacity equal to the `list_size * capacity`. - elements_builder: Box, + elements_builder: ChildBuilder, /// The null map builder of the [`FixedSizeListArray`]. /// @@ -62,7 +62,7 @@ impl FixedSizeListBuilder { ) -> Self { let elements_capacity = capacity * list_size as usize; - let elements_builder = builder_with_capacity(&element_dtype, elements_capacity); + let elements_builder = ChildBuilder::with_capacity(&element_dtype, elements_capacity); let fsl_dtype = DType::FixedSizeList(element_dtype, list_size, nullability); let nulls = LazyBitBufferBuilder::new(capacity); @@ -98,7 +98,7 @@ impl FixedSizeListBuilder { self.list_size() ); - array.append_to_builder(self.elements_builder.as_mut(), ctx)?; + self.elements_builder.append_array(array, ctx)?; self.nulls.append_non_null(); Ok(()) @@ -115,9 +115,7 @@ impl FixedSizeListBuilder { return Ok(()); } - array - .elements() - .append_to_builder(self.elements_builder.as_mut(), ctx)?; + self.elements_builder.append_array(array.elements(), ctx)?; self.nulls .append_validity_mask(&array.validity()?.execute_mask(array.len(), ctx)?); Ok(()) @@ -257,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); diff --git a/vortex-array/src/builders/list.rs b/vortex-array/src/builders/list.rs index e0b43b98407..38eb0a7b674 100644 --- a/vortex-array/src/builders/list.rs +++ b/vortex-array/src/builders/list.rs @@ -25,10 +25,10 @@ use crate::arrays::PrimitiveArray; use crate::arrays::list::ListArraySlotsExt; use crate::arrays::listview::ListViewArraySlotsExt; use crate::builders::ArrayBuilder; +use crate::builders::ChildBuilder; use crate::builders::DEFAULT_BUILDER_CAPACITY; use crate::builders::LazyBitBufferBuilder; use crate::builders::PrimitiveBuilder; -use crate::builders::builder_with_capacity; use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::dtype::Nullability; @@ -44,7 +44,7 @@ pub struct ListBuilder { dtype: DType, /// The builder for the underlying elements of the [`ListArray`]. - elements_builder: Box, + elements_builder: ChildBuilder, /// The builder for the `offsets` into the `elements` array. offsets_builder: PrimitiveBuilder, @@ -79,7 +79,7 @@ impl ListBuilder { elements_capacity: usize, capacity: usize, ) -> Self { - let elements_builder = builder_with_capacity(value_dtype.as_ref(), elements_capacity); + let elements_builder = ChildBuilder::with_capacity(value_dtype.as_ref(), elements_capacity); let mut offsets_builder = PrimitiveBuilder::::with_capacity(NonNullable, capacity + 1); // The first offset is always 0 and represents an empty list. @@ -111,8 +111,7 @@ impl ListBuilder { self.element_dtype() ); - self.elements_builder.reserve_exact(array.len()); - array.append_to_builder(self.elements_builder.as_mut(), ctx)?; + self.elements_builder.append_array(array, ctx)?; self.nulls.append_non_null(); self.offsets_builder.append_value( O::from_usize(self.elements_builder.len()) @@ -210,7 +209,7 @@ where let list_elements = new_elements .slice(offset..offset + size) .vortex_expect("list builder slice"); - list_elements.append_to_builder(builder.elements_builder.as_mut(), ctx)?; + builder.elements_builder.append_array(&list_elements, ctx)?; curr_offset += size; } @@ -275,6 +274,10 @@ impl ArrayBuilder for ListBuilder { 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); @@ -320,11 +323,8 @@ impl ArrayBuilder for ListBuilder { // in bulk and the offsets rebased onto this builder's elements. let elements_base = self.elements_builder.len(); if last > first { - self.elements_builder.reserve_exact(last - first); - array - .elements() - .slice(first..last)? - .append_to_builder(self.elements_builder.as_mut(), ctx)?; + self.elements_builder + .append_array(&array.elements().slice(first..last)?, ctx)?; } self.offsets_builder.reserve_exact(num_lists); diff --git a/vortex-array/src/builders/listview.rs b/vortex-array/src/builders/listview.rs index a9a234b353a..48413a0b8f5 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -28,13 +28,14 @@ use crate::arrays::ListView; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::arrays::list::ListArraySlotsExt; +use crate::arrays::listview::ListViewArrayExt; use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::listview::ListViewRebuildMode; use crate::builders::ArrayBuilder; +use crate::builders::ChildBuilder; use crate::builders::DEFAULT_BUILDER_CAPACITY; use crate::builders::PrimitiveBuilder; use crate::builders::UninitRange; -use crate::builders::builder_with_capacity; use crate::builders::lazy_null_builder::LazyBitBufferBuilder; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; @@ -58,7 +59,7 @@ pub struct ListViewBuilder { dtype: DType, /// The builder for the underlying elements of the [`ListArray`](crate::arrays::ListArray). - elements_builder: Box, + elements_builder: ChildBuilder, /// The builder for the `offsets` into the `elements` array. offsets_builder: PrimitiveBuilder, @@ -68,6 +69,14 @@ pub struct ListViewBuilder { /// The null map builder of the [`ListViewArray`]. nulls: LazyBitBufferBuilder, + + /// Whether the appends so far leave the result zero-copyable to a [`ListArray`]. + /// + /// Only [`append_listview_array`](ArrayBuilder::append_listview_array) can clear this; every + /// other append writes its lists back to back. + /// + /// [`ListArray`]: crate::arrays::ListArray + zero_copy_to_list: bool, } impl ListViewBuilder { @@ -96,7 +105,7 @@ impl ListViewBuilder { elements_capacity: usize, capacity: usize, ) -> Self { - let elements_builder = builder_with_capacity(&element_dtype, elements_capacity); + let elements_builder = ChildBuilder::with_capacity(&element_dtype, elements_capacity); let offsets_builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, capacity); @@ -111,6 +120,7 @@ impl ListViewBuilder { offsets_builder, sizes_builder, nulls, + zero_copy_to_list: true, } } @@ -142,8 +152,7 @@ impl ListViewBuilder { "appending this list would cause an offset overflow" ); - self.elements_builder.reserve_exact(num_elements); - array.append_to_builder(self.elements_builder.as_mut(), ctx)?; + self.elements_builder.append_array(array, ctx)?; self.nulls.append_non_null(); self.offsets_builder.append_value( @@ -206,6 +215,8 @@ impl ListViewBuilder { let sizes = self.sizes_builder.finish(); let validity = self.nulls.finish_with_nullability(self.dtype.nullability()); + let zero_copy_to_list = std::mem::replace(&mut self.zero_copy_to_list, true); + // SAFETY: // - Both the offsets and the sizes are non-nullable. // - The offsets, sizes, and validity have the same length since we always appended the same @@ -214,11 +225,11 @@ impl ListViewBuilder { // - In every method that adds values to this builder (`append_value`, `append_scalar`, // `append_list_array`, and `append_listview_array`), we checked that `offset + size` // does not overflow. - // - We constructed everything in a way that builds the `ListViewArray` similar to the shape - // of a `ListArray`, so we know the resulting array is zero-copyable to a `ListArray`. + // - Every append writes its lists back to back, so the result is zero-copyable to a + // `ListArray` unless `zero_copy_to_list` recorded an appended layout we left alone. unsafe { ListViewArray::new_unchecked(elements, offsets, sizes, validity) - .with_zero_copy_to_list(true) + .with_zero_copy_to_list(zero_copy_to_list) } } @@ -300,6 +311,10 @@ impl ArrayBuilder for ListViewBuilder { 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); @@ -352,23 +367,29 @@ impl ArrayBuilder for ListViewBuilder { return Ok(()); } - // Normalize to an exact zero-copy-to-list layout and then bulk append. This avoids the - // very expensive scalar_at-per-list path for overlapping / out-of-order list views. + // Trim unreferenced elements, but keep the incoming layout otherwise. Rebasing each offset + // by the builder's element count is correct for any layout, so flattening would only + // discard the source's sharing: a constant list array points every view at one copy of the + // value, and flattening materializes a copy per row. let listview = array .into_owned() - .rebuild(ListViewRebuildMode::MakeExact, ctx)?; - debug_assert!(listview.is_zero_copy_to_list()); + .rebuild(ListViewRebuildMode::TrimElements, ctx)?; + + // The next append starts where this one's elements end, so the combined array stays + // zero-copyable only if these views are packed back to back and reference every element + // they carry. Leading or trailing unreferenced elements would leave an interior gap. + self.zero_copy_to_list &= listview.is_zero_copy_to_list() + && listview.offset_at(0) == 0 + && listview.offset_at(listview.len() - 1) + listview.size_at(listview.len() - 1) + == listview.elements().len(); self.nulls .append_validity_mask(&array.validity()?.execute_mask(array.len(), ctx)?); - // Bulk append the new elements (which should have no gaps or overlaps). + // Bulk append the elements. Gaps and overlaps between the views are preserved. let old_elements_len = self.elements_builder.len(); self.elements_builder - .reserve_exact(listview.elements().len()); - listview - .elements() - .append_to_builder(self.elements_builder.as_mut(), ctx)?; + .append_array(listview.elements(), ctx)?; let new_elements_len = self.elements_builder.len(); // Reserve enough space for the new views. @@ -432,10 +453,9 @@ where ); if last > first { - builder.elements_builder.reserve_exact(last - first); - elements - .slice(first..last)? - .append_to_builder(builder.elements_builder.as_mut(), ctx)?; + builder + .elements_builder + .append_array(&elements.slice(first..last)?, ctx)?; } builder.offsets_builder.reserve_exact(num_lists); @@ -514,6 +534,7 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; + use crate::arrays::ConstantArray; use crate::arrays::ListArray; use crate::arrays::ListViewArray; use crate::arrays::listview::ListViewArrayExt; @@ -834,6 +855,39 @@ mod tests { Ok(()) } + /// A constant list array points every view at a single copy of the value. Flattening it in the + /// builder would materialize a copy per row, undoing the reason to append the array at all + /// instead of the same list in a loop. + #[test] + fn test_constant_list_append_keeps_one_copy_of_the_value() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let element_dtype: Arc = Arc::new(I32.into()); + + const ROWS: usize = 10_000; + let fill = Scalar::list( + Arc::clone(&element_dtype), + vec![1i32.into(), 2i32.into(), 3i32.into()], + NonNullable, + ); + let constant = ConstantArray::new(fill, ROWS).into_array(); + + let mut builder = + ListViewBuilder::::with_capacity(element_dtype, NonNullable, 0, 0); + constant.append_to_builder(&mut builder, &mut ctx)?; + let listview = builder.finish_into_listview(); + + assert_eq!(listview.len(), ROWS); + assert_eq!( + listview.elements().len(), + 3, + "the fill value should be stored once, not once per row", + ); + assert!(!listview.is_zero_copy_to_list()); + assert_arrays_eq!(&listview.into_array(), &constant, &mut ctx); + + Ok(()) + } + #[test] fn test_extend_from_array_overlapping_listview() { let mut ctx = array_session().create_execution_ctx(); @@ -861,7 +915,8 @@ mod tests { let listview = builder.finish_into_listview(); assert_eq!(listview.len(), 3); - assert!(listview.is_zero_copy_to_list()); + // The builder kept the source's overlapping layout, so the result is not zero-copyable. + assert!(!listview.is_zero_copy_to_list()); assert_arrays_eq!( listview.list_elements_at(0).unwrap(), @@ -875,7 +930,8 @@ mod tests { .execute_is_valid(1, &mut ctx) .unwrap() ); - assert_eq!(listview.list_elements_at(1).unwrap().len(), 0); + // List 1 is null, so its size is meaningless; the builder no longer rewrites it to zero. + assert_eq!(listview.size_at(1), source.size_at(1)); assert_arrays_eq!( listview.list_elements_at(2).unwrap(), PrimitiveArray::from_iter([10i32]), diff --git a/vortex-array/src/builders/mod.rs b/vortex-array/src/builders/mod.rs index ac1183177bc..1cadd3c80aa 100644 --- a/vortex-array/src/builders/mod.rs +++ b/vortex-array/src/builders/mod.rs @@ -6,6 +6,12 @@ //! Every logical type in Vortex has a canonical (uncompressed) in-memory encoding. This module //! provides pre-allocated builders to construct new canonical arrays. //! +//! Canonical form is not recursive, and neither are these builders: appending an array to a nested +//! builder keeps the child in the encoding it arrived in instead of decoding it. The fields of a +//! [`StructArray`](crate::arrays::StructArray), the elements of a list, and the storage of an +//! [`ExtensionArray`](crate::arrays::ExtensionArray) may therefore come back compressed, or as a +//! [`ChunkedArray`](crate::arrays::ChunkedArray) when several arrays were appended in turn. +//! //! ## Example: //! //! ``` @@ -53,6 +59,7 @@ mod lazy_null_builder; pub(crate) use lazy_null_builder::LazyBitBufferBuilder; mod bool; +mod child; mod decimal; pub mod dict; mod extension; @@ -65,6 +72,7 @@ mod struct_; mod varbinview; pub use bool::*; +pub(crate) use child::ChildBuilder; pub use decimal::*; pub use extension::*; pub use fixed_size_list::*; @@ -162,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. @@ -182,6 +204,9 @@ pub trait ArrayBuilder: Send { /// Constructs an Array from the builder components. /// + /// The returned array is canonical at the top level only; its children keep whatever encoding + /// they were appended with. + /// /// # Panics /// /// This function may panic if the builder's methods are called with invalid arguments. If only diff --git a/vortex-array/src/builders/struct_.rs b/vortex-array/src/builders/struct_.rs index da08e262760..d2cf1d67b0b 100644 --- a/vortex-array/src/builders/struct_.rs +++ b/vortex-array/src/builders/struct_.rs @@ -17,9 +17,9 @@ use crate::IntoArray; use crate::arrays::StructArray; use crate::arrays::struct_::StructArrayExt; use crate::builders::ArrayBuilder; +use crate::builders::ChildBuilder; use crate::builders::DEFAULT_BUILDER_CAPACITY; use crate::builders::LazyBitBufferBuilder; -use crate::builders::builder_with_capacity; use crate::canonical::Canonical; use crate::dtype::DType; use crate::dtype::Nullability; @@ -30,7 +30,7 @@ use crate::scalar::StructScalar; /// The builder for building a [`StructArray`]. pub struct StructBuilder { dtype: DType, - builders: Vec>, + builders: Vec, nulls: LazyBitBufferBuilder, } @@ -48,7 +48,7 @@ impl StructBuilder { ) -> Self { let builders = struct_dtype .fields() - .map(|dt| builder_with_capacity(&dt, capacity)) + .map(|dt| ChildBuilder::with_capacity(&dt, capacity)) .collect(); Self { @@ -131,7 +131,7 @@ impl StructBuilder { .iter_unmasked_fields() .zip_eq(self.builders.iter_mut()) { - field.append_to_builder(builder.as_mut(), ctx)?; + builder.append_array(field, ctx)?; } self.nulls @@ -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); diff --git a/vortex-array/src/builders/tests.rs b/vortex-array/src/builders/tests.rs index 138e6941def..445a10cb4ef 100644 --- a/vortex-array/src/builders/tests.rs +++ b/vortex-array/src/builders/tests.rs @@ -4,17 +4,46 @@ use std::sync::Arc; use rstest::rstest; +use vortex_buffer::Buffer; +use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; +use crate::ArrayRef; use crate::Canonical; use crate::ExecutionCtx; +use crate::RecursiveCanonical; use crate::VortexSessionExecute; use crate::array::IntoArray; use crate::array_session; +use crate::arrays::Chunked; +use crate::arrays::ChunkedArray; +use crate::arrays::Constant; +use crate::arrays::ConstantArray; +use crate::arrays::Extension; +use crate::arrays::ExtensionArray; +use crate::arrays::FixedSizeList; +use crate::arrays::FixedSizeListArray; +use crate::arrays::List; +use crate::arrays::ListArray; +use crate::arrays::ListView; +use crate::arrays::ListViewArray; +use crate::arrays::Primitive; +use crate::arrays::PrimitiveArray; +use crate::arrays::Struct; +use crate::arrays::StructArray; +use crate::arrays::chunked::ChunkedArrayExt; +use crate::arrays::extension::ExtensionArraySlotsExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use crate::arrays::list::ListArraySlotsExt; +use crate::arrays::listview::ListViewArraySlotsExt; +use crate::arrays::struct_::StructArrayExt; +use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; +use crate::builders::ListBuilder; use crate::builders::builder_with_capacity; +use crate::builders::child::DEFAULT_MIN_CHUNK_LEN as MIN_CHUNK_LEN; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::Nullability; @@ -24,6 +53,7 @@ use crate::dtype::half::f16; use crate::extension::datetime::TimeUnit; use crate::extension::datetime::Timestamp; use crate::scalar::Scalar; +use crate::validity::Validity; /// Test that `append_zeros` produces the same result as manually appending `Scalar::default_value`. /// @@ -897,3 +927,377 @@ fn test_set_validity_noop_when_non_nullable() -> VortexResult<()> { } Ok(()) } + +/// Builders only promise a canonical *top level*, so a child array that is long enough to be worth +/// a chunk must come back out of the builder in the encoding it went in with. +/// +/// Each case appends the same array twice, which additionally checks that the two chunks are +/// stitched back together into a [`ChunkedArray`] rather than being decoded and concatenated. +#[rstest] +#[case::struct_field( + StructArray::try_from_iter([("a", constant_i32())]) + .vortex_expect("struct array") + .into_array(), + |array: &ArrayRef| array.as_::().unmasked_field(0).clone() +)] +#[case::list_elements( + ListViewArray::new( + constant_i32(), + (0..MIN_CHUNK_LEN as u64).collect::>().into_array(), + Buffer::full(1u64, MIN_CHUNK_LEN).into_array(), + Validity::NonNullable, + ) + .into_array(), + |array: &ArrayRef| array.as_::().elements().clone() +)] +#[case::fixed_size_list_elements( + FixedSizeListArray::new( + constant_i32(), + 2, + Validity::NonNullable, + MIN_CHUNK_LEN / 2, + ) + .into_array(), + |array: &ArrayRef| array.as_::().elements().clone() +)] +#[case::extension_storage( + ExtensionArray::new( + Timestamp::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased(), + ConstantArray::new( + Scalar::primitive(0i64, Nullability::NonNullable), + MIN_CHUNK_LEN, + ) + .into_array(), + ) + .into_array(), + |array: &ArrayRef| array.as_::().storage().clone() +)] +fn test_children_are_not_canonicalized( + #[case] array: ArrayRef, + #[case] child_of: fn(&ArrayRef) -> ArrayRef, +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + let mut builder = builder_with_capacity(array.dtype(), 0); + array.append_to_builder(builder.as_mut(), &mut ctx)?; + array.append_to_builder(builder.as_mut(), &mut ctx)?; + let built = builder.finish(); + + let child = child_of(&built); + let chunked = child.as_::(); + assert_eq!( + chunked.nchunks(), + 2, + "expected one chunk per appended array" + ); + assert!( + chunked.iter_chunks().all(|chunk| chunk.is::()), + "the constant-encoded child was decoded by the builder", + ); + + let expected = ChunkedArray::try_new(vec![array.clone(), array], built.dtype().clone())?; + assert_arrays_eq!(&built, &expected, &mut ctx); + + Ok(()) +} + +/// Children short enough that a chunk would cost more than a copy are still materialized, so that +/// scalar-sized appends do not produce a [`ChunkedArray`] with more chunks than values. +#[test] +fn test_short_children_are_materialized() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + let elements = ConstantArray::new(1i32, 2).into_array(); + let array = FixedSizeListArray::new(elements, 2, Validity::NonNullable, 1).into_array(); + + let mut builder = builder_with_capacity(array.dtype(), 0); + for _ in 0..MIN_CHUNK_LEN { + array.append_to_builder(builder.as_mut(), &mut ctx)?; + } + let built = builder.finish(); + + assert!(built.as_::().elements().is::()); + assert_eq!(built.len(), MIN_CHUNK_LEN); + + Ok(()) +} + +/// A builder that mixes appended arrays with scalar appends must keep the two in order. +#[test] +fn test_struct_builder_interleaves_arrays_and_scalars() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + let array = StructArray::try_from_iter([("a", constant_i32())])?.into_array(); + let scalar = Scalar::struct_( + array.dtype().clone(), + vec![Scalar::primitive(1i32, Nullability::NonNullable)], + ); + + let mut builder = builder_with_capacity(array.dtype(), 0); + builder.append_scalar(&scalar)?; + array.append_to_builder(builder.as_mut(), &mut ctx)?; + builder.append_scalar(&scalar)?; + let built = builder.finish(); + + let scalar_array = StructArray::try_from_iter([( + "a", + PrimitiveArray::new(buffer![1i32], Validity::NonNullable), + )])? + .into_array(); + let expected = ChunkedArray::try_new( + vec![scalar_array.clone(), array, scalar_array], + built.dtype().clone(), + )?; + assert_arrays_eq!(&built, &expected, &mut ctx); + + 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_::() + .unmasked_field(0) + .as_::() + .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() +} + +/// Two lists of [`MIN_CHUNK_LEN`] elements each, so that appending them chunks the elements. +fn two_lists_of_chunk_len() -> ListViewArray { + ListViewArray::new( + iota(2 * MIN_CHUNK_LEN), + u64s([0, MIN_CHUNK_LEN]), + u64s([MIN_CHUNK_LEN, MIN_CHUNK_LEN]), + Validity::NonNullable, + ) +} + +/// `0..n` as an `i32` array, so that the values of one chunk are distinguishable from the next. +fn iota(n: usize) -> ArrayRef { + (0..n) + .map(|i| i32::try_from(i).vortex_expect("iota value fits in an i32")) + .collect::>() + .into_array() +} + +/// A `u64` array of list offsets or sizes. +fn u64s(values: impl IntoIterator) -> ArrayRef { + values + .into_iter() + .map(|v| u64::try_from(v).vortex_expect("list offset fits in a u64")) + .collect::>() + .into_array() +} + +/// Once a list builder keeps its elements as chunks, `elements_builder.len()` is a running total +/// across those chunks — every offset appended afterwards has to be rebased onto it. +/// +/// The two cases cover the two bulk paths into the elements builder: appending a `ListViewArray` +/// rebases the view's own offsets, while appending a `ListArray` slices the elements first. +#[rstest] +#[case::from_listview(two_lists_of_chunk_len().into_array())] +#[case::from_list( + ListArray::new( + iota(2 * MIN_CHUNK_LEN), + u64s([0, MIN_CHUNK_LEN, 2 * MIN_CHUNK_LEN]), + Validity::NonNullable, + ) + .into_array() +)] +fn test_list_offsets_are_rebased_across_element_chunks( + #[case] lists: ArrayRef, +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + let mut builder = builder_with_capacity( + &DType::List( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + Nullability::NonNullable, + ), + 0, + ); + lists.append_to_builder(builder.as_mut(), &mut ctx)?; + lists.append_to_builder(builder.as_mut(), &mut ctx)?; + let built = builder.finish(); + + assert!( + built.as_::().elements().is::(), + "the elements should have been kept as chunks", + ); + + let expected = ChunkedArray::try_new(vec![lists.clone(), lists], built.dtype().clone())?; + assert_arrays_eq!(&built, &expected, &mut ctx); + + Ok(()) +} + +/// `ListBuilder` computes each offset from the running element count as well, and reaches the +/// elements builder through `append_array_as_list` rather than a bulk append. +#[test] +fn test_list_builder_offsets_are_rebased_across_element_chunks() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let element_dtype = Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)); + + let mut builder = + ListBuilder::::with_capacity(element_dtype, Nullability::NonNullable, 0, 0); + for value in 0..3i32 { + builder.append_array_as_list( + &ConstantArray::new(value, MIN_CHUNK_LEN).into_array(), + &mut ctx, + )?; + } + let built = builder.finish(); + + assert!(built.as_::().elements().is::()); + + let expected = ListArray::new( + (0..3i32) + .flat_map(|value| std::iter::repeat_n(value, MIN_CHUNK_LEN)) + .collect::>() + .into_array(), + u64s((0..=3).map(|i| i * MIN_CHUNK_LEN)), + Validity::NonNullable, + ); + assert_arrays_eq!(&built, &expected, &mut ctx); + + Ok(()) +} + +/// A nested builder's own validity buffer is independent of its chunked child, so nulls appended +/// alongside chunks must survive. +#[rstest] +#[case::fixed_size_list( + FixedSizeListArray::new( + iota(MIN_CHUNK_LEN), + 4, + Validity::from_iter((0..MIN_CHUNK_LEN / 4).map(|i| i % 3 != 0)), + MIN_CHUNK_LEN / 4, + ) + .into_array(), + |array: &ArrayRef| array.as_::().elements().clone() +)] +#[case::struct_( + StructArray::try_from_iter_with_validity( + [("a", iota(MIN_CHUNK_LEN))], + Validity::from_iter((0..MIN_CHUNK_LEN).map(|i| i % 3 != 0)), + ) + .vortex_expect("struct array") + .into_array(), + |array: &ArrayRef| array.as_::().unmasked_field(0).clone() +)] +fn test_validity_survives_chunked_children( + #[case] array: ArrayRef, + #[case] child_of: fn(&ArrayRef) -> ArrayRef, +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + let mut builder = builder_with_capacity(array.dtype(), 0); + array.append_to_builder(builder.as_mut(), &mut ctx)?; + builder.append_nulls(1); + array.append_to_builder(builder.as_mut(), &mut ctx)?; + let built = builder.finish(); + + assert!(child_of(&built).is::()); + + let mut null = builder_with_capacity(array.dtype(), 1); + null.append_nulls(1); + let expected = ChunkedArray::try_new( + vec![array.clone(), null.finish(), array], + built.dtype().clone(), + )?; + assert_arrays_eq!(&built, &expected, &mut ctx); + + Ok(()) +} + +/// An extension array has no validity of its own — it lives in the storage — so overriding the +/// builder's validity has to reach into the chunked storage. +#[test] +fn test_extension_set_validity_reaches_chunked_storage() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::Nullable).erased(); + + let array = ExtensionArray::new( + ext_dtype.clone(), + ConstantArray::new( + Scalar::primitive(0i64, Nullability::Nullable), + MIN_CHUNK_LEN, + ) + .into_array(), + ) + .into_array(); + + let mut builder = builder_with_capacity(array.dtype(), 0); + array.append_to_builder(builder.as_mut(), &mut ctx)?; + array.append_to_builder(builder.as_mut(), &mut ctx)?; + + let invalid = [0, 2 * MIN_CHUNK_LEN - 1]; + builder.set_validity(Mask::from_iter( + (0..2 * MIN_CHUNK_LEN).map(|i| !invalid.contains(&i)), + )); + let built = builder.finish(); + + assert!(built.as_::().storage().is::()); + + let expected = ExtensionArray::new( + ext_dtype, + PrimitiveArray::from_option_iter( + (0..2 * MIN_CHUNK_LEN).map(|i| (!invalid.contains(&i)).then_some(0i64)), + ) + .into_array(), + ) + .into_array(); + assert_arrays_eq!(&built, &expected, &mut ctx); + + Ok(()) +} + +/// Consumers that genuinely need a fully-decoded tree ask for it, and must still get one. +#[test] +fn test_chunked_children_canonicalize_recursively() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + let array = StructArray::try_from_iter([("a", constant_i32())])?.into_array(); + let mut builder = builder_with_capacity(array.dtype(), 0); + array.append_to_builder(builder.as_mut(), &mut ctx)?; + array.append_to_builder(builder.as_mut(), &mut ctx)?; + let built = builder.finish(); + + let recursive = built.clone().execute::(&mut ctx)?.0; + assert!( + recursive + .clone() + .into_array() + .as_::() + .unmasked_field(0) + .is::() + ); + assert_arrays_eq!(&recursive.into_array(), &built, &mut ctx); + + Ok(()) +}