diff --git a/vortex-array/src/builders/child.rs b/vortex-array/src/builders/child.rs new file mode 100644 index 00000000000..67d6fee77d9 --- /dev/null +++ b/vortex-array/src/builders/child.rs @@ -0,0 +1,496 @@ +// 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 shortest array that a [`ChildBuilder`] keeps as a chunk of its own. +/// +/// Nested builders routinely append very short arrays (the elements of a single list, for +/// example), and giving each one its own chunk would produce a [`ChunkedArray`] with more chunks +/// than values. Below this length, copying the values costs less than the indirection that every +/// later access to the chunk would pay. +pub(super) const MIN_CHUNK_LEN: usize = 64; + +/// 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, +} + +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), + } + } + + /// 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() < 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 rstest::rstest; + use vortex_buffer::buffer; + use vortex_error::VortexExpect; + use vortex_error::VortexResult; + use vortex_mask::Mask; + + use super::ChildBuilder; + use super::MIN_CHUNK_LEN; + use 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::Masked; + use crate::arrays::Primitive; + use crate::arrays::PrimitiveArray; + use crate::arrays::chunked::ChunkedArrayExt; + 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; + + /// 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)) }; + } + + #[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..e756aef2696 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`]. diff --git a/vortex-array/src/builders/fixed_size_list.rs b/vortex-array/src/builders/fixed_size_list.rs index 7cd0ac692f0..8dd8a1c3ebc 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(()) diff --git a/vortex-array/src/builders/list.rs b/vortex-array/src/builders/list.rs index e0b43b98407..01e36923066 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; } @@ -320,11 +319,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..b91e4e45182 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -31,10 +31,10 @@ use crate::arrays::list::ListArraySlotsExt; 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 +58,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, @@ -96,7 +96,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); @@ -142,8 +142,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( @@ -365,10 +364,7 @@ impl ArrayBuilder for ListViewBuilder { // Bulk append the new elements (which should have no gaps or overlaps). 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 +428,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); diff --git a/vortex-array/src/builders/mod.rs b/vortex-array/src/builders/mod.rs index 561efc2711e..3a877ba2e78 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::*; @@ -180,6 +188,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..f226d04e4b6 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 diff --git a/vortex-array/src/builders/tests.rs b/vortex-array/src/builders/tests.rs index 138e6941def..6f125155eb1 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::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,348 @@ 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(()) +} + +/// 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(()) +}