diff --git a/encodings/sparse/src/canonical.rs b/encodings/sparse/src/canonical.rs index 9f6a3ee91bb..ea76a7e88c5 100644 --- a/encodings/sparse/src/canonical.rs +++ b/encodings/sparse/src/canonical.rs @@ -21,6 +21,7 @@ use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinView; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; +use vortex_array::arrays::listview::ListViewArrayExt; use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::arrays::listview::ListViewRebuildMode; use vortex_array::arrays::primitive::PrimitiveArrayExt; @@ -31,6 +32,7 @@ use vortex_array::builders::ArrayBuilder; use vortex_array::builders::DecimalBuilder; use vortex_array::builders::FixedSizeListBuilder; use vortex_array::builders::ListViewBuilder; +use vortex_array::builders::builder_with_capacity; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::DecimalType; @@ -47,6 +49,7 @@ use vortex_array::match_each_unsigned_integer_ptype; use vortex_array::match_smallest_list_offset_type; use vortex_array::patches::Patches; use vortex_array::scalar::DecimalScalar; +use vortex_array::scalar::ListScalar; use vortex_array::scalar::Scalar; use vortex_array::scalar::StructScalar; use vortex_array::validity::Validity; @@ -232,51 +235,89 @@ fn execute_sparse_lists_inner( total_canonical_values, len, ); - let patch_values = patch_values.into_array(); + // The fill's elements become an array once, up front. Every gap then appends that same array, + // so the fill's elements are stored once for the whole result however many gaps reference them. + let fill_elements = list_scalar_elements_array(fill_value.as_list()); + + // One mask for the whole patch array rather than a validity lookup per patch. + let patch_validity = patch_values + .listview_validity() + .execute_mask(patch_values.len(), ctx) + .vortex_expect("sparse list validity mask failed to execute"); let mut next_index = 0; - let mut patch_idx = 0; - - while patch_idx < patch_indices.len() { - let sparse_idx = sparse_index_at(patch_indices, patch_idx); - append_fill(&mut builder, fill_value, sparse_idx - next_index, ctx); - - // Patches landing on consecutive rows go in as one slice of the patch array rather than a - // list at a time, so the builder sees a run per gap instead of a run per patch. - let run_end = consecutive_run_end(patch_indices, patch_idx); - patch_values - .slice(patch_idx..run_end) - .vortex_expect("patch run is in bounds") - .append_to_builder(&mut builder, ctx) - .vortex_expect("Failed to append sparse values"); - - next_index = sparse_index_at(patch_indices, run_end - 1) + 1; - patch_idx = run_end; + + for ((patch_idx, sparse_idx), patch_valid) in + patch_indices.iter().enumerate().zip(patch_validity.iter()) + { + let sparse_idx = sparse_idx + .to_usize() + .vortex_expect("patch index must fit in usize"); + + append_list_fill( + &mut builder, + fill_elements.as_ref(), + sparse_idx - next_index, + ctx, + ); + + // Take each patch's elements rather than slicing the patch array itself: slicing a + // `ListView` slices its offsets, its sizes and its elements, and every one of those slices + // pays an optimizer pass, where this pays one for the elements alone. + if patch_valid { + let patch_list = patch_values + .list_elements_at(patch_idx) + .vortex_expect("list_elements_at"); + builder + .append_array_as_list(&patch_list, ctx) + .vortex_expect("Failed to append sparse value"); + } else { + builder.append_null(); + } + + next_index = sparse_idx + 1; } - append_fill(&mut builder, fill_value, len - next_index, ctx); + append_list_fill(&mut builder, fill_elements.as_ref(), len - next_index, ctx); builder.finish() } -/// The sparse row that the patch at `patch_idx` occupies. -fn sparse_index_at(patch_indices: &[I], patch_idx: usize) -> usize { - patch_indices[patch_idx] - .to_usize() - .vortex_expect("patch index must fit in usize") +/// Materializes a list scalar's elements into an array, or `None` if the scalar is null. +fn list_scalar_elements_array(list: ListScalar) -> Option { + list.elements().map(|elements| { + let mut builder = builder_with_capacity(list.element_dtype(), elements.len()); + for element in elements { + builder + .append_scalar(&element) + .vortex_expect("list element scalar was invalid"); + } + builder.finish() + }) } -/// The end of the run of patches starting at `start` that occupy consecutive sparse rows. +/// Appends the run of `count` fill lists that covers the gap before the next patch. /// -/// Patch indices are strictly increasing, so a run is any stretch over which they step by one. -fn consecutive_run_end(patch_indices: &[I], start: usize) -> usize { - let mut end = start + 1; - while end < patch_indices.len() - && sparse_index_at(patch_indices, end) == sparse_index_at(patch_indices, end - 1) + 1 - { - end += 1; +/// The whole run goes in as one append that points `count` views at a single copy of +/// `fill_elements`, so a gap costs nothing per row it covers. +fn append_list_fill( + builder: &mut ListViewBuilder, + fill_elements: Option<&ArrayRef>, + count: usize, + ctx: &mut ExecutionCtx, +) { + if count == 0 { + return; + } + + match fill_elements { + Some(fill_elements) => builder + .append_array_as_repeated_list(fill_elements, count, ctx) + .vortex_expect("Failed to append sparse fill value"), + // A null fill has no elements to share, and the builder can record the nulls without + // going through an array at all. + None => builder.append_nulls(count), } - end } /// Canonicalize a sparse [`FixedSizeListArray`] by expanding it into a dense representation. @@ -328,45 +369,89 @@ fn execute_sparse_fixed_size_list_inner( nullability, array_len, ); - let values = values.into_array(); + // The fill's elements become an array once, up front, so that a gap does not rebuild them. + // They are tiled per row rather than shared - a fixed-size list holds its elements back to + // back - unless they are all the same scalar, in which case the tile stays constant-encoded + // and the tiling costs nothing. + let fill_elements = fixed_size_list_fill_tile(fill_value.as_list(), list_size); + + // One mask for the whole patch array rather than a validity lookup per patch. + let patch_validity = values + .validity() + .vortex_expect("sparse fixed-size-list validity should be derivable") + .execute_mask(values.len(), ctx) + .vortex_expect("sparse fixed-size-list validity mask failed to execute"); let mut next_index = 0; - let mut patch_idx = 0; - while patch_idx < patch_indices.len() { + for ((patch_idx, sparse_idx), patch_valid) in + patch_indices.iter().enumerate().zip(patch_validity.iter()) + { // Fill gap before this patch with fill values. - let sparse_idx = sparse_index_at(patch_indices, patch_idx); - append_fill(&mut builder, fill_value, sparse_idx - next_index, ctx); - - // Patches landing on consecutive rows go in as one slice of the patch array, whose - // elements and validity are both appended in bulk. A null patch carries its own - // placeholder elements along, which is what the builder would have appended for it anyway. - let run_end = consecutive_run_end(patch_indices, patch_idx); - values - .slice(patch_idx..run_end) - .vortex_expect("patch run is in bounds") - .append_to_builder(&mut builder, ctx) - .vortex_expect("Failed to append sparse fixed-size-list values"); - - next_index = sparse_index_at(patch_indices, run_end - 1) + 1; - patch_idx = run_end; + let sparse_idx = sparse_idx + .to_usize() + .vortex_expect("patch index must fit in usize"); + append_fixed_size_list_fill( + &mut builder, + fill_elements.as_ref(), + sparse_idx - next_index, + ctx, + ); + + // Take each patch's elements rather than slicing the patch array itself: slicing a + // `FixedSizeList` slices its elements and its validity, and every one of those slices pays + // an optimizer pass, where this pays one for the elements alone. + if patch_valid { + let patch_list = values + .fixed_size_list_elements_at(patch_idx) + .vortex_expect("fixed_size_list_elements_at"); + builder + .append_array_as_list(&patch_list, ctx) + .vortex_expect("Failed to append sparse fixed-size-list value"); + } else { + builder.append_null(); + } + + next_index = sparse_idx + 1; } // Fill remaining positions after last patch. - append_fill(&mut builder, fill_value, array_len - next_index, ctx); + append_fixed_size_list_fill( + &mut builder, + fill_elements.as_ref(), + array_len - next_index, + ctx, + ); builder.finish_into_fixed_size_list() } -/// Appends the run of `count` fill values that covers the gap before the next patch. +/// Materializes the elements a fixed-size-list fill value covers each of its rows with, or `None` +/// if the fill is null. /// -/// The run goes in as a single [`ConstantArray`] rather than one appended value per row. For a -/// list dtype that is the difference between storing the fill value once per gap and once per row: -/// canonicalizing a constant list array points every view at one copy of the value, and the -/// builder keeps that layout. -fn append_fill( - builder: &mut dyn ArrayBuilder, - fill_value: &Scalar, +/// Elements that are all the same scalar stay a constant array, so tiling them over a gap costs +/// nothing however many rows it covers. +fn fixed_size_list_fill_tile(fill: ListScalar, list_size: u32) -> Option { + let elements = fill.elements()?; + + Some(match elements.iter().all_equal_value() { + Ok(uniform) => ConstantArray::new(uniform.clone(), list_size as usize).into_array(), + Err(_) => { + let mut builder = builder_with_capacity(fill.element_dtype(), elements.len()); + for element in &elements { + builder + .append_scalar(element) + .vortex_expect("fixed-size-list element scalar was invalid"); + } + builder.finish() + } + }) +} + +/// Appends the run of `count` fill lists that covers the gap before the next patch. +fn append_fixed_size_list_fill( + builder: &mut FixedSizeListBuilder, + fill_elements: Option<&ArrayRef>, count: usize, ctx: &mut ExecutionCtx, ) { @@ -374,17 +459,13 @@ fn append_fill( return; } - if fill_value.is_null() { - // A null fill has no elements to share, and the builder can record the nulls without - // going through an array at all. - builder.append_nulls(count); - return; + match fill_elements { + Some(fill_elements) => builder + .append_array_as_repeated_list(fill_elements, count, ctx) + .vortex_expect("Failed to append sparse fixed-size-list fill value"), + // A null fill has no elements of its own, only the placeholders the builder writes. + None => builder.append_nulls(count), } - - ConstantArray::new(fill_value.clone(), count) - .into_array() - .append_to_builder(builder, ctx) - .vortex_expect("Failed to append sparse fill value"); } fn execute_sparse_bools( @@ -1342,10 +1423,11 @@ mod test { } /// Nested builders chunk a child on the boundaries it is appended on, so the number of appends - /// canonicalization makes is now visible in the elements child. A run of consecutive patches - /// and the gap after it should cost one chunk each, not one chunk per row. + /// canonicalization makes is visible in the elements child. Patches go in one at a time, so + /// they cost a chunk each; a gap covers all its rows with a single append, so it costs one + /// chunk however many rows it fills. #[test] - fn test_sparse_list_chunks_elements_per_run_not_per_patch() -> VortexResult<()> { + fn test_sparse_list_chunks_elements_per_patch_and_once_per_gap() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); const PATCHES: usize = 100; @@ -1369,8 +1451,8 @@ mod test { let actual = sparse.execute::(&mut ctx)?; assert_eq!( actual.elements().as_::().nchunks(), - 2, - "expected one chunk for the patch run and one for the gap", + PATCHES + 1, + "expected one chunk per patch and a single chunk for the whole gap", ); let expected_lists = (0..patches_i32) diff --git a/vortex-array/src/arrays/constant/vtable/mod.rs b/vortex-array/src/arrays/constant/vtable/mod.rs index 2e4c982a8ea..f409236e1d1 100644 --- a/vortex-array/src/arrays/constant/vtable/mod.rs +++ b/vortex-array/src/arrays/constant/vtable/mod.rs @@ -33,12 +33,16 @@ use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::BoolBuilder; use crate::builders::DecimalBuilder; +use crate::builders::ExtensionBuilder; +use crate::builders::FixedSizeListBuilder; use crate::builders::NullBuilder; use crate::builders::PrimitiveBuilder; +use crate::builders::StructBuilder; use crate::builders::VarBinViewBuilder; use crate::canonical::Canonical; use crate::dtype::DType; use crate::match_each_decimal_value; +use crate::match_each_list_builder; use crate::match_each_native_ptype; use crate::match_each_varbin_builder; use crate::scalar::DecimalValue; @@ -252,21 +256,66 @@ impl VTable for Constant { }); } } - // TODO: add fast paths for DType::Struct, DType::List, DType::FixedSizeList, DType::Extension. - _ => { - let canonical = array - .array() - .clone() - .execute::(ctx)? - .into_array(); - canonical.append_to_builder(builder, ctx)?; + DType::List(..) => append_constant_list_run(array, n, builder, ctx)?, + DType::Struct(..) => match builder.as_any_mut().downcast_mut::() { + Some(b) => b.append_constant(scalar.as_struct(), n, ctx)?, + None => append_via_canonical(array, builder, ctx)?, + }, + DType::Extension(..) => match builder.as_any_mut().downcast_mut::() { + Some(b) => b.append_constant(scalar.as_extension(), n, ctx)?, + None => append_via_canonical(array, builder, ctx)?, + }, + DType::FixedSizeList(..) => { + match builder.as_any_mut().downcast_mut::() { + Some(b) => b.append_constant(scalar.as_list(), n, ctx)?, + None => append_via_canonical(array, builder, ctx)?, + } } + // TODO: add a fast path for DType::Union once it has a builder. + _ => append_via_canonical(array, builder, ctx)?, } Ok(()) } } +/// Appends the constant list `array` as one repeated run per list builder. +/// +/// Only the concrete list builders can record a run from the scalar; any other builder for a list +/// dtype has to go the long way around. +// The complexity comes from the expansion of `match_each_list_builder!`. +#[expect(clippy::cognitive_complexity)] +fn append_constant_list_run( + array: ArrayView<'_, Constant>, + n: usize, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let scalar = array.scalar(); + match match_each_list_builder!(builder, |b| b.append_constant_list( + scalar.as_list(), + n, + ctx + )) { + Some(result) => result, + None => append_via_canonical(array, builder, ctx), + } +} + +/// Appends `array` by canonicalizing it first, for the dtypes with no fast path of their own. +fn append_via_canonical( + array: ArrayView<'_, Constant>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let canonical = array + .array() + .clone() + .execute::(ctx)? + .into_array(); + canonical.append_to_builder(builder, ctx) +} + /// Downcasts `builder` to `B`, then either appends `n` nulls or calls `fill` with the typed /// builder depending on `is_null`. /// @@ -291,19 +340,30 @@ fn append_value_or_nulls( #[cfg(test)] mod tests { + use std::sync::Arc; + use rstest::rstest; use vortex_error::VortexResult; use crate::IntoArray; use crate::VortexSessionExecute; + use crate::arrays::Constant; use crate::arrays::ConstantArray; + use crate::arrays::Extension; + use crate::arrays::FixedSizeList; + use crate::arrays::Struct; use crate::arrays::constant::vtable::canonical::constant_canonicalize; + use crate::arrays::extension::ExtensionArrayExt; + use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; + use crate::arrays::struct_::StructArrayExt; use crate::assert_arrays_eq; use crate::builders::builder_with_capacity; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::StructFields; + use crate::extension::datetime::Date; + use crate::extension::datetime::TimeUnit; use crate::scalar::Scalar; /// Appends `array` into a fresh builder and asserts the result matches `constant_canonicalize`. @@ -438,4 +498,115 @@ mod tests { let dtype = DType::Struct(fields, Nullability::Nullable); assert_append_matches_canonical(ConstantArray::new(Scalar::null(dtype), 4)) } + + /// A run of identical structs should leave each field constant-encoded rather than materialize + /// a value per row per field. + #[test] + fn test_struct_constant_append_keeps_fields_constant() -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let fields = StructFields::new( + ["x", "y"].into(), + vec![ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Utf8(Nullability::NonNullable), + ], + ); + let scalar = Scalar::struct_( + DType::Struct(fields, Nullability::NonNullable), + [ + Scalar::primitive(42i32, Nullability::NonNullable), + Scalar::utf8("hi", Nullability::NonNullable), + ], + ); + let array = ConstantArray::new(scalar, 1_000); + + let mut builder = builder_with_capacity(array.dtype(), array.len()); + array + .into_array() + .append_to_builder(builder.as_mut(), &mut ctx)?; + let result = builder.finish(); + + let struct_array = result.as_::(); + for field in 0..2 { + assert!( + struct_array.unmasked_field(field).is::(), + "field {field} should have stayed constant-encoded", + ); + } + Ok(()) + } + + #[rstest] + #[case::non_uniform(vec![Scalar::from(1i32), Scalar::from(2i32)])] + #[case::uniform(vec![Scalar::from(7i32), Scalar::from(7i32)])] + fn test_fixed_size_list_constant_append(#[case] elements: Vec) -> VortexResult<()> { + let scalar = Scalar::fixed_size_list( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + elements, + Nullability::NonNullable, + ); + assert_append_matches_canonical(ConstantArray::new(scalar, 4)) + } + + #[test] + fn test_null_fixed_size_list_constant_append() -> VortexResult<()> { + let dtype = DType::FixedSizeList( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + 2, + Nullability::Nullable, + ); + assert_append_matches_canonical(ConstantArray::new(Scalar::null(dtype), 3)) + } + + /// A fixed-size list whose elements are all the same scalar tiles to a constant array, so the + /// elements child should not be materialized at all. + #[test] + fn test_uniform_fixed_size_list_constant_append_keeps_elements_constant() -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let scalar = Scalar::fixed_size_list( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + vec![Scalar::from(7i32), Scalar::from(7i32)], + Nullability::NonNullable, + ); + let array = ConstantArray::new(scalar, 1_000); + + let mut builder = builder_with_capacity(array.dtype(), array.len()); + array + .into_array() + .append_to_builder(builder.as_mut(), &mut ctx)?; + let result = builder.finish(); + + assert!( + result.as_::().elements().is::(), + "a uniform tile should have stayed constant-encoded", + ); + Ok(()) + } + + #[test] + fn test_extension_constant_append() -> VortexResult<()> { + let scalar = Scalar::extension::(TimeUnit::Days, Scalar::from(Some(42i32))); + assert_append_matches_canonical(ConstantArray::new(scalar, 5)) + } + + /// An extension array is its storage wearing a dtype, so a run of identical values should leave + /// the storage constant-encoded. + #[test] + fn test_extension_constant_append_keeps_storage_constant() -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let scalar = Scalar::extension::(TimeUnit::Days, Scalar::from(Some(42i32))); + let array = ConstantArray::new(scalar, 1_000); + + let mut builder = builder_with_capacity(array.dtype(), array.len()); + array + .into_array() + .append_to_builder(builder.as_mut(), &mut ctx)?; + let result = builder.finish(); + + assert!( + result.as_::().storage_array().is::(), + "the storage should have stayed constant-encoded", + ); + Ok(()) + } } diff --git a/vortex-array/src/builders/child.rs b/vortex-array/src/builders/child.rs index 869bdb44071..d1ff5f2954e 100644 --- a/vortex-array/src/builders/child.rs +++ b/vortex-array/src/builders/child.rs @@ -63,6 +63,12 @@ impl ChildBuilder { /// at a boundary only the caller can see. Callers that would rather have the values copied /// should append them as scalars. /// + /// An appended [`ChunkedArray`] stays one chunk rather than giving up its own. Unpacking it + /// would spill a chunk per row into this child for a caller expressing a repeated run as + /// chunks - a tiled fixed-size list, say - and the child's chunk list is what every later + /// append and the final [`finish`](Self::finish) walk. Nesting costs one level, not one per + /// appended array. + /// /// Nothing is decoded here, so `_ctx` goes unused; it stays in the signature so that the /// nested builders forwarding their [`ExecutionCtx`] here do not have to explain why they /// don't. @@ -80,6 +86,7 @@ impl ChildBuilder { self.flush_pending(); self.chunks_len += array.len(); + self.chunks.push(array.clone()); Ok(()) diff --git a/vortex-array/src/builders/extension.rs b/vortex-array/src/builders/extension.rs index c1a91f202d0..d246850a7bb 100644 --- a/vortex-array/src/builders/extension.rs +++ b/vortex-array/src/builders/extension.rs @@ -9,6 +9,7 @@ use vortex_error::vortex_ensure; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; +use crate::arrays::ConstantArray; use crate::arrays::ExtensionArray; use crate::arrays::extension::ExtensionArrayExt; use crate::builders::ArrayBuilder; @@ -55,6 +56,24 @@ impl ExtensionBuilder { self.storage.append_array(array.storage_array(), ctx) } + /// Appends the same extension `value` `n` times. + /// + /// An extension array is its storage array wearing a dtype, so a run of identical values is a + /// constant storage array, which goes in as a single chunk and stays constant-encoded. + pub(crate) fn append_constant( + &mut self, + value: ExtScalar, + n: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if n == 0 { + return Ok(()); + } + + let storage = ConstantArray::new(value.to_storage_scalar(), n).into_array(); + self.storage.append_array(&storage, ctx) + } + /// Finishes the builder directly into a [`ExtensionArray`]. pub fn finish_into_extension(&mut self) -> ExtensionArray { let storage = self.storage.finish(); diff --git a/vortex-array/src/builders/fixed_size_list.rs b/vortex-array/src/builders/fixed_size_list.rs index 2dc1749dbe0..cade5c09fbe 100644 --- a/vortex-array/src/builders/fixed_size_list.rs +++ b/vortex-array/src/builders/fixed_size_list.rs @@ -4,6 +4,7 @@ use std::any::Any; use std::sync::Arc; +use itertools::Itertools as _; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -13,12 +14,15 @@ use vortex_error::vortex_panic; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; +use crate::arrays::ChunkedArray; +use crate::arrays::ConstantArray; 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::ValidityBuilder; +use crate::builders::builder_with_capacity; use crate::canonical::Canonical; use crate::dtype::DType; use crate::dtype::Nullability; @@ -103,6 +107,110 @@ impl FixedSizeListBuilder { Ok(()) } + /// Appends `array` as `n` identical non-null lists. + /// + /// A fixed-size list array holds its elements back to back, so `n` identical lists are the + /// array's elements tiled `n` times - there is no layout that lets the rows share one range of + /// elements the way a list view's can. The tiling costs nothing to build even so: the elements + /// go in as a [`ChunkedArray`] of `n` clones of the same array, so the tile's values are stored + /// once however many rows reference them, and the child holds the whole run as one chunk. A + /// constant tile does better still, collapsing to a single constant chunk. + /// + /// A caller with a run of appends to make should hand over the same `array` each time rather + /// than rebuild it, which is the whole reason this takes an array instead of a scalar. + pub fn append_array_as_repeated_list( + &mut self, + array: &ArrayRef, + n: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + vortex_ensure!( + array.dtype() == self.element_dtype(), + "Array dtype {:?} does not match list element dtype {:?}", + array.dtype(), + self.element_dtype() + ); + vortex_ensure!( + array.len() == self.list_size() as usize, + "Array length {} does not match fixed list size {}", + array.len(), + self.list_size() + ); + + if n == 0 { + return Ok(()); + } + + let tiled_len = n * array.len(); + match array.as_constant() { + Some(element) => { + let tiled = ConstantArray::new(element, tiled_len).into_array(); + self.elements_builder.append_array(&tiled, ctx)?; + } + None => { + // SAFETY: every chunk is `array` itself, so they share its dtype and none is empty. + let tiled = unsafe { + ChunkedArray::new_unchecked( + std::iter::repeat_n(array.clone(), n).collect::>(), + self.element_dtype().clone(), + ) + }; + self.elements_builder + .append_array(&tiled.into_array(), ctx)?; + } + } + self.nulls.append_n_non_nulls(n); + + Ok(()) + } + + /// Appends the same fixed-size list `value` `n` times. + /// + /// This materializes the value's elements and hands them to + /// [`append_array_as_repeated_list`](Self::append_array_as_repeated_list) - as a constant array + /// when they are all the same scalar, so that the tiling costs nothing. A caller with a run of + /// appends to make should materialize the elements once itself and call that directly. + pub(crate) fn append_constant( + &mut self, + value: ListScalar, + n: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if n == 0 { + return Ok(()); + } + + let Some(elements) = value.elements() else { + vortex_ensure!( + self.dtype.is_nullable(), + "Cannot append a null fixed-size list to a non-nullable builder" + ); + self.append_nulls(n); + return Ok(()); + }; + + let list_size = self.list_size() as usize; + vortex_ensure!( + elements.len() == list_size, + "Scalar list length {} does not match fixed list size {}", + elements.len(), + list_size + ); + + let tile = match elements.iter().all_equal_value() { + Ok(uniform) => ConstantArray::new(uniform.clone(), list_size).into_array(), + Err(_) => { + let mut tile_builder = builder_with_capacity(self.element_dtype(), list_size); + for element in &elements { + tile_builder.append_scalar(element)?; + } + tile_builder.finish() + } + }; + + self.append_array_as_repeated_list(&tile, n, ctx) + } + /// Appends the values of a canonical [`FixedSizeListArray`] to the builder, recursing into the /// elements builder. pub(crate) fn append_fixed_size_list_array( diff --git a/vortex-array/src/builders/list.rs b/vortex-array/src/builders/list.rs index ca74b7b9328..032464790e8 100644 --- a/vortex-array/src/builders/list.rs +++ b/vortex-array/src/builders/list.rs @@ -149,6 +149,47 @@ impl ListBuilder { Ok(()) } + /// Appends the same list `value` `n` times. + /// + /// A `ListArray`'s offsets can only describe contiguous, in-order lists, so the elements go in + /// once per row - there is nothing to share, unlike + /// [`ListViewBuilder::append_constant_list`](crate::builders::ListViewBuilder::append_constant_list). + /// What this does save is the array machinery: canonicalizing the run and appending it back + /// costs a `ListViewArray` construction, a rebuild and two offset casts that the values + /// themselves do not need. + pub fn append_constant_list( + &mut self, + value: ListScalar, + n: usize, + _ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if n == 0 { + return Ok(()); + } + + let Some(elements) = value.elements() else { + if self.dtype.nullability() == NonNullable { + vortex_bail!("Cannot append null value to non-nullable list"); + } + self.append_nulls(n); + return Ok(()); + }; + + self.offsets_builder.reserve_exact(n); + for _ in 0..n { + for scalar in &elements { + self.elements_builder.append_scalar(scalar)?; + } + self.offsets_builder.append_value( + O::from_usize(self.elements_builder.len()) + .vortex_expect("Failed to convert from usize to O"), + ); + } + self.nulls.append_n_non_nulls(n); + + Ok(()) + } + /// Finishes the builder directly into a [`ListArray`]. pub fn finish_into_list(&mut self) -> ListArray { assert_eq!( diff --git a/vortex-array/src/builders/listview.rs b/vortex-array/src/builders/listview.rs index e1dd3af31d2..7cf13a725f5 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -12,6 +12,7 @@ use std::sync::Arc; +use num_traits::ToPrimitive; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -27,15 +28,15 @@ 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::ValidityBuilder; -use crate::builtins::ArrayBuiltins; +use crate::builders::builder_with_capacity; use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::dtype::Nullability; @@ -204,6 +205,90 @@ impl ListViewBuilder { Ok(()) } + /// Appends `array` as `n` identical non-null lists, storing its elements once. + /// + /// A `ListViewArray` can point many views at one range of elements, so a repeated list costs + /// its elements once however many rows it covers. The elements go in as one appended array, so + /// a caller that hands over the same `array` on every call - a sparse array filling the gaps + /// between its patches, say - stores those elements once for the whole result. + /// + /// The views share their elements, so the result is no longer zero-copyable to a + /// [`ListArray`](crate::arrays::ListArray) unless it covers a single row or empty lists. + pub fn append_array_as_repeated_list( + &mut self, + array: &ArrayRef, + n: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + vortex_ensure!( + array.dtype() == self.element_dtype(), + "Array dtype {:?} does not match list element dtype {:?}", + array.dtype(), + self.element_dtype() + ); + + if n == 0 { + return Ok(()); + } + + let curr_offset = self.elements_builder.len(); + let num_elements = array.len(); + + // We must assert this even in release mode to ensure that the safety comment in + // `finish_into_listview` is correct. + assert!( + ((curr_offset + num_elements) as u64) < O::max_value_as_u64(), + "appending this list would cause an offset overflow" + ); + + self.elements_builder.append_array(array, ctx)?; + + let offset = + O::from_usize(curr_offset).vortex_expect("Failed to convert from usize to `O`"); + let size = S::from_usize(num_elements).vortex_expect("Failed to convert from usize to `S`"); + self.offsets_builder.append_n_values(offset, n); + self.sizes_builder.append_n_values(size, n); + self.nulls.append_n_non_nulls(n); + + if n > 1 && num_elements > 0 { + self.zero_copy_to_list = false; + } + + Ok(()) + } + + /// Appends the same list `value` `n` times, storing its elements once. + /// + /// This materializes the scalar's elements and hands them to + /// [`append_array_as_repeated_list`](Self::append_array_as_repeated_list). A caller with a run + /// of appends to make should materialize the elements once itself and call that directly. + pub fn append_constant_list( + &mut self, + value: ListScalar, + n: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if n == 0 { + return Ok(()); + } + + let Some(elements) = value.elements() else { + vortex_ensure!( + self.dtype.is_nullable(), + "Cannot append null value to non-nullable list builder" + ); + self.append_nulls(n); + return Ok(()); + }; + + let mut elements_builder = builder_with_capacity(self.element_dtype(), elements.len()); + for scalar in elements { + elements_builder.append_scalar(&scalar)?; + } + + self.append_array_as_repeated_list(&elements_builder.finish(), n, ctx) + } + /// Finishes the builder directly into a [`ListViewArray`]. pub fn finish_into_listview(&mut self) -> ListViewArray { debug_assert_eq!(self.offsets_builder.len(), self.sizes_builder.len()); @@ -291,52 +376,85 @@ impl ListViewBuilder { return Ok(()); } - // Drop leading and trailing unreferenced elements so we do not copy them in, but keep the - // layout otherwise: rebasing the offsets is correct whatever it is, and flattening would - // throw away the source's sharing - a constant list array points every view at one copy. - let listview = array - .into_owned() - .rebuild(ListViewRebuildMode::TrimElements, ctx)?; + let len = array.len(); - // A trimmed zero-copy-to-list source references every element it carries, back to back, so - // it lands flush against the elements already in the builder. Any other layout does not. - self.zero_copy_to_list &= listview.is_zero_copy_to_list(); + // Materialize the metadata once and do the trimming and the rebase by hand. Going through + // `rebuild(ListViewRebuildMode::TrimElements)` would subtract the window start from every + // offset with a compute kernel, only for the rebase below to add this builder's elements + // base straight back on - two passes, one of them through the compute stack, for one + // addition per offset. Casting the sizes to the builder's type is another kernel for what + // is a copy. + let offsets = array.offsets().clone().execute::(ctx)?; + let sizes = array.sizes().clone().execute::(ctx)?; + + // The window of `elements` that the views actually reference. Everything outside it is + // unreachable and must not be appended. An exact source covers its window back to back and + // in order, so the first and last view bound it; any other layout has to be searched. + let (start, end) = if array.is_zero_copy_to_list() { + let last = len - 1; + ( + metadata_at(&offsets, 0), + metadata_at(&offsets, last) + metadata_at(&sizes, last), + ) + } else { + array.into_owned().referenced_element_bounds(ctx)? + }; - self.nulls.append_validity(array.validity()?, array.len()); + // An exact source references every element it carries, back to back, so it lands flush + // against the elements already in the builder. Any other layout does not. + self.zero_copy_to_list &= array.is_zero_copy_to_list(); + + self.nulls.append_validity(array.validity()?, len); + + // Bulk append the referenced elements; the offsets are rebased onto them below. + let elements_base = self.elements_builder.len(); + + // We must assert this even in release mode to ensure that the safety comment in + // `finish_into_listview` is correct. + assert!( + ((elements_base + (end - start)) as u64) < O::max_value_as_u64(), + "appending this list would cause an offset overflow" + ); + + if end > start { + self.elements_builder + .append_array(&array.elements().slice(start..end)?, ctx)?; + } + + // Every view lies inside `start..end`, so rebasing it onto `elements_base` keeps it inside + // the elements this builder now holds - which is what lets `finish_into_listview` build the + // array unchecked. + assert_eq!( + self.elements_builder.len(), + elements_base + (end - start), + "appending the referenced elements did not extend the child by the window's length" + ); - // Bulk append the trimmed elements; the offsets are rebased onto them below. - let old_elements_len = self.elements_builder.len(); - self.elements_builder - .append_array(listview.elements(), ctx)?; - let new_elements_len = self.elements_builder.len(); - - // Reserve enough space for the new views. - let extend_length = listview.len(); - self.sizes_builder.reserve_exact(extend_length); - self.offsets_builder.reserve_exact(extend_length); - - // The incoming sizes might have a different type than the builder, so we need to cast. - let cast_sizes = listview - .sizes() - .clone() - .cast(self.sizes_builder.dtype().clone())?; - cast_sizes.append_to_builder(&mut self.sizes_builder, ctx)?; - - // Now we need to adjust all of the offsets by adding the current number of elements in the - // builder. - let uninit_range = self.offsets_builder.uninit_range(extend_length); - - // This should be cheap because trimming rebases the offsets, it does not compress them. - let new_offsets = listview.offsets().clone().execute::(ctx)?; - - match_each_integer_ptype!(new_offsets.ptype(), |A| { - adjust_and_extend_offsets::( - uninit_range, - new_offsets, - old_elements_len, - new_elements_len, + self.offsets_builder.reserve_exact(len); + let offsets_range = self.offsets_builder.uninit_range(len); + match_each_integer_ptype!(offsets.ptype(), |A| { + extend_rebased_offsets::( + offsets_range, + offsets.as_slice::(), + start, + elements_base, ); }); + + self.sizes_builder.reserve_exact(len); + let mut sizes_range = self.sizes_builder.uninit_range(len); + if sizes.ptype() == S::PTYPE { + // The sizes already have the builder's type, so there is nothing to convert. + sizes_range.copy_from_slice(0, sizes.as_slice::()); + // SAFETY: `copy_from_slice` initialized all `len` values, and the sizes builder is + // non-nullable. + unsafe { sizes_range.finish() }; + } else { + match_each_integer_ptype!(sizes.ptype(), |A| { + extend_converted_sizes::(sizes_range, sizes.as_slice::()); + }); + } + Ok(()) } } @@ -483,44 +601,67 @@ where Ok(()) } -/// Given new offsets, adds them to the `UninitRange` after adding the `old_elements_len` to each -/// offset. -fn adjust_and_extend_offsets( - mut uninit_range: UninitRange, - new_offsets: PrimitiveArray, - old_elements_len: usize, - new_elements_len: usize, -) { - let new_offsets_slice = new_offsets.as_slice::(); - let old_elements_len = O::from_usize(old_elements_len) - .vortex_expect("the old elements length did not fit into the offset type (impossible)"); - let new_elements_len = O::from_usize(new_elements_len) - .vortex_expect("the current elements length did not fit into the offset type (impossible)"); - - for i in 0..uninit_range.len() { - let new_offset = O::from_usize( - new_offsets_slice[i] - .to_usize() - .vortex_expect("Offsets must always fit in usize"), - ) - .vortex_expect("New offset somehow did not fit into the builder's offset type"); +/// Reads one non-nullable integer value of list view metadata as a `usize`. +fn metadata_at(metadata: &PrimitiveArray, index: usize) -> usize { + match_each_integer_ptype!(metadata.ptype(), |A| { + metadata.as_slice::()[index] + .to_usize() + .vortex_expect("list view metadata must fit in a usize") + }) +} - // We have to check this even in release mode to ensure the final `new_unchecked` - // construction in `finish_into_listview` is valid. - let adjusted_new_offset = new_offset + old_elements_len; - assert!( - adjusted_new_offset <= new_elements_len, - "[{i}/{}]: {new_offset} + {old_elements_len} \ - = {adjusted_new_offset} <= {new_elements_len} failed", - uninit_range.len() +/// Writes `offsets` into `range`, moved off the source's element window and onto the elements the +/// builder already holds. +/// +/// `window_start` is the first element the source's views reference, so every offset is at least +/// `window_start` and the subtraction cannot underflow. +fn extend_rebased_offsets( + mut range: UninitRange, + offsets: &[A], + window_start: usize, + elements_base: usize, +) { + debug_assert_eq!(range.len(), offsets.len()); + + for (i, &offset) in offsets.iter().enumerate() { + let offset = offset + .to_usize() + .vortex_expect("offsets must always fit in usize"); + debug_assert!( + offset >= window_start, + "offset {offset} precedes the referenced window at {window_start}" ); - - uninit_range.set_value(i, adjusted_new_offset); + let rebased = O::from_usize(offset - window_start + elements_base) + .vortex_expect("rebased offset did not fit into the builder's offset type"); + range.set_value(i, rebased); } // SAFETY: We have set all the values in the range, and since `offsets` are non-nullable, we are // done. - unsafe { uninit_range.finish() }; + unsafe { range.finish() }; +} + +/// Writes `sizes` into `range`, converting them to the builder's size type. +/// +/// Sizes that already have the builder's type are copied in bulk by the caller instead. +fn extend_converted_sizes( + mut range: UninitRange, + sizes: &[A], +) { + debug_assert_eq!(range.len(), sizes.len()); + + for (i, &size) in sizes.iter().enumerate() { + let size = S::from_usize( + size.to_usize() + .vortex_expect("sizes must always fit in usize"), + ) + .vortex_expect("size did not fit into the builder's size type"); + range.set_value(i, size); + } + + // SAFETY: We have set all the values in the range, and since `sizes` are non-nullable, we are + // done. + unsafe { range.finish() }; } #[cfg(test)] @@ -888,6 +1029,48 @@ mod tests { Ok(()) } + /// Only the elements the views reference land in the builder: appending a slice out of the + /// middle of a list array leaves the elements on either side of it behind. + #[test] + fn test_append_listview_array_trims_unreferenced_elements() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype: Arc = Arc::new(I32.into()); + + // Five lists of two elements each, of which we append only the middle three. + let source = ListArray::from_iter_slow::( + (0..5).map(|i| vec![2 * i, 2 * i + 1]), + Arc::clone(&dtype), + )? + .into_array() + .execute::(&mut ctx)?; + let middle = source.slice(1..4)?.execute::(&mut ctx)?; + + let mut builder = + ListViewBuilder::::with_capacity(Arc::clone(&dtype), NonNullable, 0, 0); + builder.append_listview_array(middle.as_view(), &mut ctx)?; + // A second append has to rebase onto the elements already in the builder. + builder.append_listview_array(middle.as_view(), &mut ctx)?; + let listview = builder.finish_into_listview(); + + assert_eq!( + listview.elements().len(), + 12, + "only the six referenced elements of each append should have landed", + ); + assert!( + listview.is_zero_copy_to_list(), + "trimming an exact source keeps the result exact", + ); + + let expected = ListArray::from_iter_slow::( + (1..4).chain(1..4).map(|i| vec![2 * i, 2 * i + 1]), + dtype, + )?; + assert_arrays_eq!(listview, expected, &mut ctx); + + Ok(()) + } + #[test] fn test_extend_from_array_overlapping_listview() { let mut ctx = array_session().create_execution_ctx(); diff --git a/vortex-array/src/builders/struct_.rs b/vortex-array/src/builders/struct_.rs index 66a2c66358d..28a6cf1bd6b 100644 --- a/vortex-array/src/builders/struct_.rs +++ b/vortex-array/src/builders/struct_.rs @@ -13,6 +13,7 @@ use vortex_error::vortex_panic; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; +use crate::arrays::ConstantArray; use crate::arrays::StructArray; use crate::arrays::struct_::StructArrayExt; use crate::builders::ArrayBuilder; @@ -84,6 +85,57 @@ impl StructBuilder { Ok(()) } + /// Appends the same struct `value` `n` times. + /// + /// Each field takes a constant array of its own field value as a single chunk, so a run of + /// identical structs costs one chunk per field rather than any work per row. The fields stay + /// constant-encoded, which is all [`Canonical`] asks of them. + pub(crate) fn append_constant( + &mut self, + value: StructScalar, + n: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if !self.dtype.is_nullable() && value.is_null() { + vortex_bail!("Tried to append a null `StructScalar` to a non-nullable struct builder",); + } + + if value.struct_fields() != self.struct_fields() { + vortex_bail!( + "Tried to append a `StructScalar` with fields {} to a \ + struct builder with fields {}", + value.struct_fields(), + self.struct_fields() + ); + } + + if n == 0 { + return Ok(()); + } + + match value.fields_iter() { + Some(fields) => { + for (builder, field) in self.builders.iter_mut().zip_eq(fields) { + builder.append_array(&ConstantArray::new(field, n).into_array(), ctx)?; + } + self.nulls.append_n_non_nulls(n); + } + None => { + // The struct is null, so the fields only need placeholder values of the right + // dtype. `default_value` gives a zero for a non-nullable field and a null for a + // nullable one, preserving each field's nullability. + let field_dtypes: Vec<_> = self.struct_fields().fields().collect(); + for (builder, field_dtype) in self.builders.iter_mut().zip_eq(field_dtypes) { + let placeholder = Scalar::default_value(&field_dtype); + builder.append_array(&ConstantArray::new(placeholder, n).into_array(), ctx)?; + } + self.nulls.append_n_nulls(n); + } + } + + Ok(()) + } + /// Finishes the builder directly into a [`StructArray`]. pub fn finish_into_struct(&mut self) -> StructArray { let len = self.len();