diff --git a/encodings/sparse/src/canonical.rs b/encodings/sparse/src/canonical.rs index 2cfed7fd25b..ea76a7e88c5 100644 --- a/encodings/sparse/src/canonical.rs +++ b/encodings/sparse/src/canonical.rs @@ -23,6 +23,7 @@ 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; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::arrays::varbinview::build_views::BinaryView; @@ -180,18 +181,31 @@ fn execute_sparse_lists( // instead of 8. `O` is already unsigned (from `match_smallest_list_offset_type`). let indices = resolved.indices().as_::().into_owned(); let indices = indices.reinterpret_cast(indices.ptype().to_unsigned()); - let values = resolved.values().as_::().into_owned(); let fill_list = fill_value.as_list(); + // Flatten the patch views up front so that runs of them can be appended as slices below. An + // exact layout also trims `elements` to exactly what the patches reference, so the patch half + // of the count below is what the builder will actually hold rather than an over-estimate. + let values = resolved + .values() + .clone() + .downcast::() + .rebuild(ListViewRebuildMode::MakeExact, ctx)?; + + // Each gap between patches is appended as one constant array, whose canonical form points + // every view at a single copy of the fill value, so a gap costs the fill value's elements once + // however many rows it covers. Bound the number of gaps: there is at most one on either side of + // each patch, and each one covers at least one row. let n_filled = len - resolved.num_patches(); - let total_canonical_values = values.elements().len() + fill_list.len() * n_filled; + let n_fill_runs = (resolved.num_patches() + 1).min(n_filled); + let total_canonical_values = values.elements().len() + fill_list.len() * n_fill_runs; Ok(match_each_unsigned_integer_ptype!(indices.ptype(), |I| { match_smallest_list_offset_type!(total_canonical_values, |O| { execute_sparse_lists_inner::( indices.as_slice(), values, - fill_list, + fill_value, values_dtype, len, total_canonical_values, @@ -206,7 +220,7 @@ fn execute_sparse_lists( fn execute_sparse_lists_inner( patch_indices: &[I], patch_values: ListViewArray, - fill_scalar: ListScalar, + fill_value: &Scalar, values_dtype: Arc, len: usize, total_canonical_values: usize, @@ -221,18 +235,20 @@ fn execute_sparse_lists_inner( total_canonical_values, len, ); - let fill_elements = list_scalar_elements_array(fill_scalar); - let patch_values_validity = patch_values + // 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; - for ((patch_idx, sparse_idx), patch_valid) in patch_indices - .iter() - .enumerate() - .zip(patch_values_validity.iter()) + for ((patch_idx, sparse_idx), patch_valid) in + patch_indices.iter().enumerate().zip(patch_validity.iter()) { let sparse_idx = sparse_idx .to_usize() @@ -245,6 +261,9 @@ fn execute_sparse_lists_inner( 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) @@ -264,6 +283,43 @@ fn execute_sparse_lists_inner( builder.finish() } +/// 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() + }) +} + +/// Appends the run of `count` fill lists that covers the gap before the next patch. +/// +/// 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), + } +} + /// Canonicalize a sparse [`FixedSizeListArray`] by expanding it into a dense representation. fn execute_sparse_fixed_size_list( resolved: &Patches, @@ -274,13 +330,12 @@ fn execute_sparse_fixed_size_list( ) -> VortexResult { let indices = resolved.indices().as_::().into_owned(); let values = resolved.values().as_::().into_owned(); - let fill_scalar = fill_value.as_list(); Ok(match_each_integer_ptype!(indices.ptype(), |I| { execute_sparse_fixed_size_list_inner::( indices.as_slice(), values, - fill_scalar, + fill_value, len, nullability, ctx, @@ -296,9 +351,9 @@ fn execute_sparse_fixed_size_list( /// elements (or defaults if null). Since all lists have the same size, we can directly append /// elements without tracking offsets. fn execute_sparse_fixed_size_list_inner( - indices: &[I], + patch_indices: &[I], values: FixedSizeListArray, - fill_scalar: ListScalar, + fill_value: &Scalar, array_len: usize, nullability: Nullability, ctx: &mut ExecutionCtx, @@ -314,20 +369,28 @@ fn execute_sparse_fixed_size_list_inner( nullability, array_len, ); - let fill_elements = list_scalar_elements_array(fill_scalar); - let values_validity = values + // 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 indices = indices - .iter() - .map(|x| (*x).to_usize().vortex_expect("index must fit in usize")); - for ((patch_idx, sparse_idx), patch_valid) in indices.enumerate().zip(values_validity.iter()) { + 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_idx + .to_usize() + .vortex_expect("patch index must fit in usize"); append_fixed_size_list_fill( &mut builder, fill_elements.as_ref(), @@ -335,7 +398,9 @@ fn execute_sparse_fixed_size_list_inner( ctx, ); - // Append the patch value, handling null patches by appending defaults. + // 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) @@ -361,49 +426,45 @@ fn execute_sparse_fixed_size_list_inner( builder.finish_into_fixed_size_list() } -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"); +/// Materializes the elements a fixed-size-list fill value covers each of its rows with, or `None` +/// if the fill is null. +/// +/// 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() } - builder.finish() }) } -fn append_list_fill( - builder: &mut ListViewBuilder, - fill_elements: Option<&ArrayRef>, - count: usize, - ctx: &mut ExecutionCtx, -) { - if let Some(fill_elements) = fill_elements { - for _ in 0..count { - builder - .append_array_as_list(fill_elements, ctx) - .vortex_expect("Failed to append sparse fill value"); - } - } else { - builder.append_nulls(count); - } -} - +/// 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, ) { - if let Some(fill_elements) = fill_elements { - for _ in 0..count { - builder - .append_array_as_list(fill_elements, ctx) - .vortex_expect("Failed to append sparse fixed-size-list fill value"); - } - } else { - builder.append_nulls(count); + 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 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), } } @@ -606,6 +667,7 @@ mod test { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; + use vortex_array::arrays::Chunked; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::ListArray; @@ -614,6 +676,7 @@ mod test { use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::VarBinViewArray; + use vortex_array::arrays::chunked::ChunkedArrayExt; use vortex_array::arrays::listview::ListViewArrayExt; use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::assert_arrays_eq; @@ -1305,6 +1368,149 @@ mod test { Ok(()) } + /// Each gap between patches is appended as a single constant array, so the fill value's + /// elements are stored once per gap however many rows the gap covers. + #[test] + fn test_sparse_list_fill_stores_one_copy_per_gap() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + + // Two single-element patch lists: [1] and [2]. + let lists = unsafe { + ListViewArray::new_unchecked( + buffer![1i32, 2].into_array(), + buffer![0u32, 1].into_array(), + buffer![1u32, 1].into_array(), + Validity::AllValid, + ) + .with_zero_copy_to_list(true) + } + .into_array(); + + // Patches at 10 and 20 of 10,000 rows, so the fill covers three gaps. + let indices = buffer![10u32, 20].into_array(); + let fill = vec![7i32, 8, 9]; + let sparse = + Sparse::try_new(indices, lists, 10_000, Scalar::from(Some(fill.clone())))?.into_array(); + + let actual = sparse.execute::(&mut ctx)?; + assert_eq!(actual.len(), 10_000); + assert_eq!( + actual.elements().len(), + 2 + 3 * fill.len(), + "the fill value should be stored once per gap, not once per row", + ); + + let fill_elements = PrimitiveArray::from_iter(fill); + for index in [0, 9, 11, 19, 21, 9_999] { + assert_arrays_eq!( + actual.list_elements_at(index).vortex_expect("fill list"), + fill_elements, + &mut ctx + ); + } + assert_arrays_eq!( + actual.list_elements_at(10).vortex_expect("patch list"), + PrimitiveArray::from_iter([1i32]), + &mut ctx + ); + assert_arrays_eq!( + actual.list_elements_at(20).vortex_expect("patch list"), + PrimitiveArray::from_iter([2i32]), + &mut ctx + ); + + Ok(()) + } + + /// Nested builders chunk a child on the boundaries it is appended on, so the number of appends + /// 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_patch_and_once_per_gap() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + + const PATCHES: usize = 100; + let patches_u32 = u32::try_from(PATCHES).vortex_expect("fits in u32"); + let patches_i32 = i32::try_from(PATCHES).vortex_expect("fits in i32"); + + // `PATCHES` single-element lists, patched onto rows 0..PATCHES of a 2 * PATCHES-row array, + // so there is exactly one patch run followed by exactly one gap. + let patch_values = ListViewArray::new( + PrimitiveArray::from_iter(0..patches_i32).into_array(), + PrimitiveArray::from_iter(0..patches_u32).into_array(), + PrimitiveArray::from_iter(std::iter::repeat_n(1u32, PATCHES)).into_array(), + Validity::AllValid, + ) + .into_array(); + + let indices = PrimitiveArray::from_iter(0..patches_u32).into_array(); + let fill = Scalar::from(Some(vec![-1i32])); + let sparse = Sparse::try_new(indices, patch_values, 2 * PATCHES, fill)?.into_array(); + + let actual = sparse.execute::(&mut ctx)?; + assert_eq!( + actual.elements().as_::().nchunks(), + PATCHES + 1, + "expected one chunk per patch and a single chunk for the whole gap", + ); + + let expected_lists = (0..patches_i32) + .map(|i| Some(vec![i])) + .chain(std::iter::repeat_n(Some(vec![-1i32]), PATCHES)); + let expected = ListArray::from_iter_opt_slow::( + expected_lists, + Arc::new(PType::I32.into()), + )?; + assert_arrays_eq!(actual, expected, &mut ctx); + + Ok(()) + } + + /// Patches on consecutive rows are appended as one slice of the patch array, so this covers the + /// run arithmetic together with everything that has to survive it: a null patch inside a run, + /// patch views that overlap and are out of order, runs separated by gaps, and a trailing gap. + #[test] + fn test_sparse_list_appends_consecutive_patches_as_one_run() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + + // Overlapping, out-of-order patch views over three elements: + // - patch 0: [20, 30] + // - patch 1: [10] (null, so its elements are never read) + // - patch 2: [10, 20, 30] + // - patch 3: [10, 20] + let patches = unsafe { + ListViewArray::new_unchecked( + buffer![10i32, 20, 30].into_array(), + buffer![1u32, 0, 0, 0].into_array(), + buffer![2u32, 1, 3, 2].into_array(), + Validity::from_iter([true, false, true, true]), + ) + }; + assert!(!patches.is_zero_copy_to_list()); + + // Rows 0..3 are one run of patches, row 5 is another, and rows 3-4 and 6 are gaps. + let indices = buffer![0u8, 1, 2, 5].into_array(); + let fill = Scalar::from(Some(vec![7i32, 8])); + let sparse = Sparse::try_new(indices, patches.into_array(), 7, fill)?.into_array(); + + let actual = sparse.execute::(&mut ctx)?; + + // The seven elements the patches reference, plus one copy of the two-element fill for each + // of the two gaps. The run of patches on rows 0..3 has no gap inside it to pay for. + assert_eq!(actual.elements().len(), 7 + 2 * 2); + + let expected = ListViewArray::new( + buffer![20i32, 30, 10, 20, 30, 7, 8, 7, 8, 10, 20, 7, 8].into_array(), + buffer![0u8, 2, 2, 5, 7, 9, 11].into_array(), + buffer![2u8, 0, 3, 2, 2, 2, 2].into_array(), + Validity::from_iter([true, false, true, true, true, true, true]), + ); + assert_arrays_eq!(actual, expected, &mut ctx); + + Ok(()) + } + #[test] fn test_sparse_binary_varbin_null_fill() { let mut ctx = SESSION.create_execution_ctx(); 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 a10eb816580..bc6dfcbf28a 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 @@ -338,6 +338,7 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; + use crate::RecursiveCanonical; use crate::VortexSessionExecute; use crate::aggregate_fn::Accumulator; use crate::aggregate_fn::AggregateFnVTable; @@ -376,15 +377,26 @@ mod tests { use crate::scalar::ScalarValue; use crate::validity::Validity; + /// The size the array occupies once rebuilt through the canonical builders, which is the + /// layout [`UncompressedSizeInBytes`] is defined against: the builders normalize physical + /// widths that the input array is free to choose differently, picking the smallest decimal + /// value type for a precision and `u64` list-view offsets and sizes. + /// + /// Builders no longer canonicalize their children, so the finished array is only canonical at + /// the top level - recursively canonicalize it before measuring. fn materialized_uncompressed_size_in_bytes(array: &ArrayRef) -> u64 { + let mut ctx = array_session().create_execution_ctx(); let mut builder = builder_with_capacity(array.dtype(), array.len()); array - .append_to_builder( - builder.as_mut(), - &mut array_session().create_execution_ctx(), - ) + .append_to_builder(builder.as_mut(), &mut ctx) .vortex_expect("appended"); - builder.finish().nbytes() + builder + .finish() + .execute::(&mut ctx) + .vortex_expect("recursively canonicalized") + .0 + .into_array() + .nbytes() } fn aggregate(array: &ArrayRef) -> VortexResult { diff --git a/vortex-array/src/arrays/constant/vtable/mod.rs b/vortex-array/src/arrays/constant/vtable/mod.rs index 2e4c982a8ea..6917d979239 100644 --- a/vortex-array/src/arrays/constant/vtable/mod.rs +++ b/vortex-array/src/arrays/constant/vtable/mod.rs @@ -5,6 +5,7 @@ use std::fmt::Debug; use std::hash::Hash; use std::hash::Hasher; +use itertools::Itertools; use vortex_buffer::ByteBufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -26,6 +27,7 @@ use crate::array::ArrayId; use crate::array::ArrayView; use crate::array::VTable; use crate::array::unsupported_buffer_replacement; +use crate::arrays::ExtensionArray; use crate::arrays::constant::ConstantData; use crate::arrays::constant::compute::rules::PARENT_RULES; use crate::arrays::constant::vtable::canonical::constant_canonicalize; @@ -33,15 +35,21 @@ use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::BoolBuilder; use crate::builders::DecimalBuilder; +use crate::builders::FixedSizeListBuilder; +use crate::builders::ListViewBuilder; use crate::builders::NullBuilder; use crate::builders::PrimitiveBuilder; use crate::builders::VarBinViewBuilder; +use crate::builders::builder_with_capacity; use crate::canonical::Canonical; use crate::dtype::DType; +use crate::dtype::OffsetBuilderPType; use crate::match_each_decimal_value; +use crate::match_each_listview_builder; use crate::match_each_native_ptype; use crate::match_each_varbin_builder; use crate::scalar::DecimalValue; +use crate::scalar::ListScalar; use crate::scalar::Scalar; use crate::scalar::ScalarValue; use crate::serde::ArrayChildren; @@ -252,21 +260,135 @@ 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::Extension(ext_dtype) => { + // An extension array is its storage wearing a dtype, so a run of identical values + // is a constant storage array, which stays constant-encoded in the builder. + // Canonicalizing instead would materialize the storage: see the note in + // `constant_canonicalize` about `ExtensionConstantRule`. + let storage = ConstantArray::new(scalar.as_extension().to_storage_scalar(), n); + ExtensionArray::new(ext_dtype.clone(), storage.into_array()) + .into_array() + .append_to_builder(builder, ctx)? } + DType::FixedSizeList(..) => { + append_constant_fixed_size_list_run(array, n, builder, ctx)? + } + // The remaining dtypes canonicalize cheaply: a constant struct canonicalizes to + // constant fields, and a constant map to views sharing one copy of the entries, so + // appending the canonical array preserves the run's economy. + // 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 run sharing a single copy of its elements. +/// +/// The list's elements materialize once, and +/// [`ListViewBuilder::append_array_as_repeated_list`] points the run's `n` views at that one +/// copy. Only a list-view builder has a layout that can share elements; any other builder for a +/// list dtype - a [`ListBuilder`](crate::builders::ListBuilder), whose offsets can only describe +/// contiguous lists - appends the canonical run instead. +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_listview_builder!(builder, |b| append_repeated_list_run( + b, + scalar.as_list(), + n, + ctx + )) { + Some(result) => result, + None => append_via_canonical(array, builder, ctx), + } +} + +/// Appends the list `scalar` to a [`ListViewBuilder`] `n` times, storing its elements once. +fn append_repeated_list_run( + builder: &mut ListViewBuilder, + scalar: ListScalar, + n: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + if n == 0 { + return Ok(()); + } + + let Some(elements) = scalar.elements() else { + // A null run stores no elements at all. + builder.append_nulls(n); + return Ok(()); + }; + + let mut elements_builder = builder_with_capacity(scalar.element_dtype(), elements.len()); + for element in &elements { + elements_builder.append_scalar(element)?; + } + + builder.append_array_as_repeated_list(&elements_builder.finish(), n, ctx) +} + +/// Appends the constant fixed-size-list `array` as its list's elements tiled `n` times. +/// +/// The list's elements materialize into a tile once - a single [`ConstantArray`] when they are +/// all the same scalar, so that the tiling costs nothing - and +/// [`FixedSizeListBuilder::append_array_as_repeated_list`] shares that one tile across the run. +fn append_constant_fixed_size_list_run( + array: ArrayView<'_, Constant>, + n: usize, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let Some(builder) = builder.as_any_mut().downcast_mut::() else { + return append_via_canonical(array, builder, ctx); + }; + + if n == 0 { + return Ok(()); + } + + let scalar = array.scalar().as_list(); + let Some(elements) = scalar.elements() else { + // A null run stores no elements of its own, only the placeholders the builder writes. + builder.append_nulls(n); + return Ok(()); + }; + + let tile = match elements.iter().all_equal_value() { + Ok(uniform) => ConstantArray::new(uniform.clone(), elements.len()).into_array(), + Err(_) => { + let mut tile_builder = builder_with_capacity(builder.element_dtype(), elements.len()); + for element in &elements { + tile_builder.append_scalar(element)?; + } + tile_builder.finish() + } + }; + + builder.append_array_as_repeated_list(&tile, n, 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 +413,36 @@ 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::Chunked; + use crate::arrays::Constant; use crate::arrays::ConstantArray; + use crate::arrays::Extension; + use crate::arrays::FixedSizeList; + use crate::arrays::ListView; + use crate::arrays::Struct; + use crate::arrays::chunked::ChunkedArrayExt; use crate::arrays::constant::vtable::canonical::constant_canonicalize; + use crate::arrays::extension::ExtensionArrayExt; + use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; + 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::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`. @@ -410,6 +549,79 @@ mod tests { )) } + #[rstest] + #[case::non_empty(vec![Scalar::from(1i32), Scalar::from(2i32)], 4)] + #[case::empty(vec![], 3)] + #[case::n_zero(vec![Scalar::from(1i32)], 0)] + fn test_list_constant_append( + #[case] elements: Vec, + #[case] n: usize, + ) -> VortexResult<()> { + let scalar = Scalar::list( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + elements, + Nullability::NonNullable, + ); + assert_append_matches_canonical(ConstantArray::new(scalar, n)) + } + + #[test] + fn test_null_list_constant_append() -> VortexResult<()> { + let dtype = DType::List( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + Nullability::Nullable, + ); + assert_append_matches_canonical(ConstantArray::new(Scalar::null(dtype), 3)) + } + + /// A run of identical lists appended into a list-view builder shares one copy of its elements + /// across the whole run. + #[test] + fn test_list_constant_append_keeps_one_copy_of_the_elements() -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let scalar = Scalar::list( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + vec![Scalar::from(1i32), Scalar::from(2i32), Scalar::from(3i32)], + 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_eq!( + result.as_::().elements().len(), + 3, + "the run's elements should be stored once, not once per row", + ); + Ok(()) + } + + /// A `ListBuilder`'s offsets can only describe contiguous lists, so a constant run cannot + /// share its elements there and takes the canonical path instead. + #[test] + fn test_list_constant_append_into_list_builder() -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let element_dtype: Arc = + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)); + let scalar = Scalar::list( + Arc::clone(&element_dtype), + vec![Scalar::from(1i32), Scalar::from(2i32)], + Nullability::NonNullable, + ); + let array = ConstantArray::new(scalar, 4).into_array(); + + let mut builder = + ListBuilder::::with_capacity(element_dtype, Nullability::NonNullable, 0, 0); + array.append_to_builder(&mut builder, &mut ctx)?; + + assert_arrays_eq!(&builder.finish(), &array, &mut ctx); + Ok(()) + } + #[test] fn test_struct_constant_append() -> VortexResult<()> { let fields = StructFields::new( @@ -438,4 +650,119 @@ 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 a constant array, so the + /// tile's chunks stay constant-encoded rather than materializing a value per row. + #[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(); + + let elements = result.as_::().elements().clone(); + assert!( + elements + .as_::() + .iter_chunks() + .all(|chunk| chunk.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/arrays/varbin/builder.rs b/vortex-array/src/arrays/varbin/builder.rs index c798d7ba02c..5990ecee371 100644 --- a/vortex-array/src/arrays/varbin/builder.rs +++ b/vortex-array/src/arrays/varbin/builder.rs @@ -456,17 +456,6 @@ impl VarBinBuilder { } } - fn replace_validity(&mut self, validity: Mask) { - self.validity = match validity { - Mask::AllTrue(len) => BitBufferMut::new_set(len), - Mask::AllFalse(len) => BitBufferMut::new_unset(len), - values @ Mask::Values(_) => values - .into_bit_buffer() - .try_into_mut() - .unwrap_or_else(|buffer| BitBufferMut::copy_from(&buffer)), - }; - } - /// Appends `count` end offsets derived from `end_offsets`, shifted past the current data end. /// /// `end_offsets` must be monotonically non-decreasing and end at exactly `num_bytes`. Offsets @@ -630,10 +619,6 @@ impl ArrayBuilder for VarBinBuilder { self.validity.reserve(additional); } - unsafe fn set_validity_unchecked(&mut self, validity: Mask) { - self.replace_validity(validity) - } - fn finish(&mut self) -> ArrayRef { self.finish_into_varbin().into_array() } @@ -938,22 +923,19 @@ mod tests { #[case(true)] fn test_array_builder_methods(#[case] large_offsets: bool) -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); - for validity in [ - Mask::new_true(3), - Mask::new_false(3), - Mask::from_iter([true, false, true]), - ] { - let result = with_offsets(large_offsets, DType::Utf8(Nullable), |builder| { - builder.reserve_exact(3); - builder.append_zero(); - builder.append_scalar(&Scalar::utf8("hello", Nullable))?; - builder.append_null(); - assert_eq!(builder.len(), 3); - builder.set_validity(validity.clone()); - Ok(()) - })?; - assert_eq!(result.validity()?.execute_mask(3, &mut ctx)?, validity); - } + let result = with_offsets(large_offsets, DType::Utf8(Nullable), |builder| { + builder.reserve_exact(3); + builder.append_zero(); + builder.append_scalar(&Scalar::utf8("hello", Nullable))?; + builder.append_null(); + assert_eq!(builder.len(), 3); + Ok(()) + })?; + + assert_eq!( + result.validity()?.execute_mask(3, &mut ctx)?, + Mask::from_iter([true, true, false]) + ); Ok(()) } diff --git a/vortex-array/src/builders/bool.rs b/vortex-array/src/builders/bool.rs index 73d2089e221..3587f6fc3ba 100644 --- a/vortex-array/src/builders/bool.rs +++ b/vortex-array/src/builders/bool.rs @@ -7,7 +7,6 @@ use std::mem; use vortex_buffer::BitBufferMut; use vortex_error::VortexResult; use vortex_error::vortex_ensure; -use vortex_mask::Mask; use crate::ArrayRef; use crate::ExecutionCtx; @@ -127,10 +126,6 @@ impl ArrayBuilder for BoolBuilder { self.nulls.reserve_exact(additional); } - unsafe fn set_validity_unchecked(&mut self, validity: Mask) { - self.nulls = LazyBitBufferBuilder::from_validity_mask(validity); - } - fn finish(&mut self) -> ArrayRef { self.finish_into_bool().into_array() } diff --git a/vortex-array/src/builders/child.rs b/vortex-array/src/builders/child.rs new file mode 100644 index 00000000000..d1ff5f2954e --- /dev/null +++ b/vortex-array/src/builders/child.rs @@ -0,0 +1,361 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::ChunkedArray; +use crate::builders::ArrayBuilder; +use crate::builders::builder_with_capacity; +use crate::dtype::DType; +use crate::scalar::Scalar; + +/// 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 as a chunk of its own, keeping its encoding. + /// + /// However short the array, it becomes a chunk: the caller had a whole array to hand, and + /// deciding on its behalf that its values are cheaper copied than referenced would be guessing + /// 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. + 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(()); + } + + 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) + } + + /// 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); + } + + unsafe { ChunkedArray::new_unchecked(chunks, self.dtype.clone()) }.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::VortexResult; + + use super::ChildBuilder; + 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::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; + + /// An arbitrary array length. `ChildBuilder` treats no length specially, so the tests only + /// need a length long enough to tell chunks apart. + const CHUNK_LEN: usize = 64; + + /// 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_appended_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, CHUNK_LEN), &mut ctx)?; + builder.append_array(&constant(2, CHUNK_LEN), &mut ctx)?; + assert_eq!(builder.len(), 2 * 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(()) + } + + /// However short the appended array, its encoding survives: a single-value array is a chunk + /// too. A caller that wants the values copied appends them as scalars instead. + #[test] + fn test_short_arrays_are_kept_as_chunks_too() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = ChildBuilder::with_capacity(&DType::from(I32), 0); + + builder.append_array(&constant(1, 1), &mut ctx)?; + builder.append_array(&constant(2, 1), &mut ctx)?; + assert_eq!(builder.len(), 2); + + let child = builder.finish(); + let chunked = child.as_::(); + assert_eq!(chunked.nchunks(), 2); + assert!(chunked.iter_chunks().all(|c| c.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, CHUNK_LEN), &mut ctx)?; + builder.append_scalar(&3i32.into())?; + + let child = builder.finish(); + assert_eq!(child.len(), CHUNK_LEN + 2); + + let expected = ChunkedArray::try_new( + vec![ + buffer![1i32].into_array(), + constant(2, 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, 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, CHUNK_LEN).slice(0..0)?; + + builder.append_array(&empty, &mut ctx)?; + builder.append_array(&constant(1, CHUNK_LEN), &mut ctx)?; + builder.append_array(&empty, &mut ctx)?; + builder.append_array(&constant(2, CHUNK_LEN), &mut ctx)?; + builder.append_array(&empty, &mut ctx)?; + + assert_eq!(builder.len(), 2 * 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, 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 empty check, so that a mismatched array is rejected + /// whether or not it would have become a chunk. + #[rstest] + #[case::empty(0)] + #[case::non_empty(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, CHUNK_LEN), &mut ctx)?; + builder.append_nulls(2); + builder.append_array(&nullable_constant(2, CHUNK_LEN), &mut ctx)?; + builder.append_zeros(1); + + let child = builder.finish(); + assert_eq!(child.len(), 2 * CHUNK_LEN + 3); + assert_eq!(child.as_::().nchunks(), 4); + + let expected = PrimitiveArray::from_option_iter( + std::iter::repeat_n(Some(1i32), CHUNK_LEN) + .chain([None, None]) + .chain(std::iter::repeat_n(Some(2i32), CHUNK_LEN)) + .chain([Some(0)]), + ) + .into_array(); + assert_arrays_eq!(&child, &expected, &mut ctx); + + 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, CHUNK_LEN), &mut ctx)?; + builder.append_scalar(&2i32.into())?; + assert_eq!(builder.finish().len(), 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/decimal.rs b/vortex-array/src/builders/decimal.rs index c08dfbc4d87..027b7e62c75 100644 --- a/vortex-array/src/builders/decimal.rs +++ b/vortex-array/src/builders/decimal.rs @@ -9,7 +9,6 @@ use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_error::vortex_panic; -use vortex_mask::Mask; use crate::ArrayRef; use crate::ExecutionCtx; @@ -219,10 +218,6 @@ impl ArrayBuilder for DecimalBuilder { self.nulls.reserve_exact(additional); } - unsafe fn set_validity_unchecked(&mut self, validity: Mask) { - self.nulls = LazyBitBufferBuilder::from_validity_mask(validity); - } - fn finish(&mut self) -> ArrayRef { self.finish_into_decimal().into_array() } diff --git a/vortex-array/src/builders/extension.rs b/vortex-array/src/builders/extension.rs index aa8e1b76aa2..c1a91f202d0 100644 --- a/vortex-array/src/builders/extension.rs +++ b/vortex-array/src/builders/extension.rs @@ -5,7 +5,6 @@ use std::any::Any; use vortex_error::VortexResult; use vortex_error::vortex_ensure; -use vortex_mask::Mask; use crate::ArrayRef; use crate::ExecutionCtx; @@ -13,8 +12,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 +23,7 @@ use crate::scalar::Scalar; /// The builder for building a [`ExtensionArray`]. pub struct ExtensionBuilder { dtype: DType, - storage: Box, + storage: ChildBuilder, } impl ExtensionBuilder { @@ -36,7 +35,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 +52,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`]. @@ -116,10 +113,6 @@ impl ArrayBuilder for ExtensionBuilder { self.storage.reserve_exact(capacity) } - unsafe fn set_validity_unchecked(&mut self, validity: Mask) { - unsafe { self.storage.set_validity_unchecked(validity) }; - } - fn finish(&mut self) -> ArrayRef { self.finish_into_extension().into_array() } diff --git a/vortex-array/src/builders/fixed_size_list.rs b/vortex-array/src/builders/fixed_size_list.rs index 7cd0ac692f0..98acd2c468e 100644 --- a/vortex-array/src/builders/fixed_size_list.rs +++ b/vortex-array/src/builders/fixed_size_list.rs @@ -9,17 +9,17 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; -use vortex_mask::Mask; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; +use crate::arrays::ChunkedArray; 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::builders::ValidityBuilder; use crate::canonical::Canonical; use crate::dtype::DType; use crate::dtype::Nullability; @@ -34,12 +34,12 @@ 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`]. /// /// We also use this type to store the length of the final output array. - nulls: LazyBitBufferBuilder, + nulls: ValidityBuilder, } impl FixedSizeListBuilder { @@ -62,9 +62,9 @@ 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); + let nulls = ValidityBuilder::new(capacity); Self { dtype: fsl_dtype, @@ -98,12 +98,59 @@ 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(()) } + /// 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 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(()); + } + + // 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 values of a canonical [`FixedSizeListArray`] to the builder, recursing into the /// elements builder. pub(crate) fn append_fixed_size_list_array( @@ -115,11 +162,8 @@ impl FixedSizeListBuilder { return Ok(()); } - array - .elements() - .append_to_builder(self.elements_builder.as_mut(), ctx)?; - self.nulls - .append_validity_mask(&array.validity()?.execute_mask(array.len(), ctx)?); + self.elements_builder.append_array(array.elements(), ctx)?; + self.nulls.append_validity(array.validity()?, array.len()); Ok(()) } @@ -263,10 +307,6 @@ impl ArrayBuilder for FixedSizeListBuilder { self.nulls.reserve_exact(additional); } - unsafe fn set_validity_unchecked(&mut self, validity: Mask) { - self.nulls = LazyBitBufferBuilder::from_validity_mask(validity); - } - fn finish(&mut self) -> ArrayRef { self.finish_into_fixed_size_list().into_array() } diff --git a/vortex-array/src/builders/lazy_null_builder.rs b/vortex-array/src/builders/lazy_null_builder.rs index 24abe9ba17d..8a9f62d1a03 100644 --- a/vortex-array/src/builders/lazy_null_builder.rs +++ b/vortex-array/src/builders/lazy_null_builder.rs @@ -31,42 +31,6 @@ impl LazyBitBufferBuilder { } } - /// Creates a builder pre-populated from a validity mask, taking ownership of the mask's buffer - /// instead of copying it where possible. - /// - /// This is the counterpart to [`append_validity_mask`](Self::append_validity_mask) for callers - /// that want to *replace* the builder's contents with the mask rather than extend them: because - /// we own the mask, we can move its buffer in instead of copying it. - pub fn from_validity_mask(validity_mask: Mask) -> Self { - match validity_mask { - // An unmaterialized builder already represents `len` non-null values, so an all-valid - // mask stays lazy. - Mask::AllTrue(len) => Self { - inner: None, - len, - capacity: len, - }, - Mask::AllFalse(len) => Self::from_buffer(BitBufferMut::new_unset(len)), - // Take ownership of the underlying buffer; `into_bit_buffer` and `try_into_mut` only - // copy when the buffer is shared, otherwise this is a move. - values @ Mask::Values(_) => Self::from_buffer( - values - .into_bit_buffer() - .try_into_mut() - .unwrap_or_else(|buffer| BitBufferMut::copy_from(&buffer)), - ), - } - } - - /// Creates a builder backed by an already-materialized buffer. - fn from_buffer(inner: BitBufferMut) -> Self { - Self { - inner: Some(inner), - len: 0, - capacity: 0, - } - } - /// Appends `n` non-null values to the builder. #[inline] pub fn append_n_non_nulls(&mut self, n: usize) { diff --git a/vortex-array/src/builders/list.rs b/vortex-array/src/builders/list.rs index 99a826e4471..ca74b7b9328 100644 --- a/vortex-array/src/builders/list.rs +++ b/vortex-array/src/builders/list.rs @@ -10,7 +10,6 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; -use vortex_mask::Mask; use crate::ArrayRef; use crate::Canonical; @@ -24,11 +23,12 @@ use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; 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::LazyBitBufferBuilder; use crate::builders::PrimitiveBuilder; -use crate::builders::builder_with_capacity; +use crate::builders::ValidityBuilder; use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::dtype::Nullability; @@ -45,13 +45,13 @@ 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, /// The null map builder of the [`ListArray`]. - nulls: LazyBitBufferBuilder, + nulls: ValidityBuilder, } impl ListBuilder { @@ -80,7 +80,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. @@ -89,7 +89,7 @@ impl ListBuilder { Self { elements_builder, offsets_builder, - nulls: LazyBitBufferBuilder::new(capacity), + nulls: ValidityBuilder::new(capacity), dtype: DType::List(value_dtype, nullability), } } @@ -112,8 +112,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()) @@ -191,8 +190,7 @@ impl ListBuilder { return Ok(()); } - self.nulls - .append_validity_mask(&array.validity()?.execute_mask(array.len(), ctx)?); + self.nulls.append_validity(array.validity()?, array.len()); let num_lists = array.len(); let offsets = array.offsets().clone().execute::(ctx)?; @@ -205,11 +203,8 @@ impl 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); @@ -233,6 +228,9 @@ impl ListBuilder { /// /// See [`append_list_array`](Self::append_list_array); this is the same hook for the canonical /// [`ListViewArray`] encoding. + /// + /// A `ListArray`'s offsets can only describe contiguous, in-order lists, so views laid out any + /// other way (overlapping, out of order, or with interior gaps) are flattened first. pub fn append_listview_array( &mut self, array: ArrayView<'_, ListView>, @@ -242,8 +240,15 @@ impl ListBuilder { return Ok(()); } - self.nulls - .append_validity_mask(&array.validity()?.execute_mask(array.len(), ctx)?); + self.nulls.append_validity(array.validity()?, array.len()); + + // Flatten the views into the only layout `ListArray` offsets can express. This is a cheap + // clone when they already are laid out that way, and the flattened result keeps the + // original validity, so the null map appended above still describes it. + let array = array + .into_owned() + .rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx)?; + debug_assert!(array.is_zero_copy_to_list()); // Note that `ListViewArray` has `n` offsets and sizes, not `n+1` offsets like `ListArray`. let elements = array.elements(); @@ -265,8 +270,13 @@ impl ListBuilder { } } -/// Appends `ListViewArray`-layout lists (`n` offsets and sizes) into a [`ListBuilder`], converting -/// into the `ListArray` (`n + 1` offsets) layout. +/// Appends the lists of a zero-copy-to-list [`ListViewArray`] (`n` offsets and sizes) into a +/// [`ListBuilder`], converting into the `ListArray` (`n + 1` offsets) layout. +/// +/// The caller must have made `new_offsets` and `new_sizes` zero-copyable to a `ListArray`, so the +/// lists they describe are contiguous and in order — which is the only layout `ListArray` offsets +/// can express. That lets the referenced elements be appended in bulk, with the offsets rebased +/// onto this builder's elements, instead of appending a slice per list. fn extend_from_listview( builder: &mut ListBuilder, new_elements: &ArrayRef, @@ -282,30 +292,26 @@ where let num_lists = new_offsets.len(); debug_assert_eq!(num_lists, new_sizes.len()); - let total_elements: usize = new_sizes.iter().map(|size| size.as_()).sum(); - builder.elements_builder.reserve_exact(total_elements); + // Leading and trailing unreferenced elements are allowed even in a zero-copy-to-list layout, + // so the referenced range is bounded by the first list's start and the last list's end. + let first: usize = new_offsets[0].as_(); + let last: usize = new_offsets[num_lists - 1].as_() + new_sizes[num_lists - 1].as_(); + + let elements_base = builder.elements_builder.len(); + if last > first { + builder + .elements_builder + .append_array(&new_elements.slice(first..last)?, ctx)?; + } - let mut curr_offset = builder.elements_builder.len(); builder.offsets_builder.reserve_exact(num_lists); let mut offsets_range = builder.offsets_builder.uninit_range(num_lists); - - // We need to append each list individually, converting from `ListViewArray` format to - // the `ListArray` format that `ListBuilder` expects. - for i in 0..new_offsets.len() { - let offset: usize = new_offsets[i].as_(); - let size: usize = new_sizes[i].as_(); - - if size > 0 { - 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)?; - curr_offset += size; - } - - let new_offset = O::from_usize(curr_offset).vortex_expect("Failed to convert offset"); - - offsets_range.set_value(i, new_offset); + for i in 0..num_lists { + let end: usize = new_offsets[i].as_() + new_sizes[i].as_(); + offsets_range.set_value( + i, + O::from_usize(end - first + elements_base).vortex_expect("Failed to convert offset"), + ); } // SAFETY: We have initialized all `num_lists` values, and since the `offsets` array is @@ -370,10 +376,6 @@ impl ArrayBuilder for ListBuilder { self.nulls.reserve_exact(additional); } - unsafe fn set_validity_unchecked(&mut self, validity: Mask) { - self.nulls = LazyBitBufferBuilder::from_validity_mask(validity); - } - fn finish(&mut self) -> ArrayRef { self.finish_into_list().into_array() } @@ -666,6 +668,62 @@ mod tests { Ok(()) } + /// A `ListArray`'s offsets can only describe contiguous, in-order lists, so an overlapping + /// source has to be flattened before its elements can be appended in bulk. A sliced source, + /// meanwhile, keeps the layout it has and is appended from wherever its first list starts. + #[test] + fn test_append_listview_array_flattens_overlaps_and_skips_leading_elements() -> VortexResult<()> + { + let mut ctx = array_session().create_execution_ctx(); + let dtype: Arc = Arc::new(I32.into()); + + // Overlapping source, so not zero-copyable to a list: + // - List 0: [10, 20] + // - List 1: null (size is intentionally non-zero in the source metadata) + // - List 2: [10], sharing the elements list 0 already referenced + let overlapping = unsafe { + ListViewArray::new_unchecked( + buffer![10i32, 20, 30].into_array(), + buffer![0u32, 1, 0].into_array(), + buffer![2u8, 2, 1].into_array(), + Validity::from_iter([true, false, true]), + ) + }; + assert!(!overlapping.is_zero_copy_to_list()); + + // Zero-copyable source sliced past its first list, so its elements start at offset 2. + let sliced = unsafe { + ListViewArray::new_unchecked( + buffer![40i32, 50, 60, 70].into_array(), + buffer![0u32, 2].into_array(), + buffer![2u32, 2].into_array(), + Validity::AllValid, + ) + .with_zero_copy_to_list(true) + } + .into_array() + .slice(1..2)? + .execute::(&mut ctx)?; + + let mut builder = ListBuilder::::with_capacity(dtype, Nullable, 0, 0); + builder.append_listview_array(overlapping.as_view(), &mut ctx)?; + builder.append_listview_array(sliced.as_view(), &mut ctx)?; + + let list = builder.finish_into_list(); + assert_arrays_eq!( + list.elements(), + PrimitiveArray::from_iter([10i32, 20, 10, 60, 70]), + &mut ctx + ); + assert_arrays_eq!( + list.offsets(), + PrimitiveArray::from_iter([0u32, 2, 2, 3, 5]), + &mut ctx + ); + + Ok(()) + } + #[test] fn test_extend_builder() { test_extend_builder_gen::(); diff --git a/vortex-array/src/builders/listview.rs b/vortex-array/src/builders/listview.rs index 7b49fad57e6..2efc2d45d74 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -12,11 +12,11 @@ use std::sync::Arc; +use num_traits::ToPrimitive; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; -use vortex_mask::Mask; use crate::ArrayRef; use crate::Canonical; @@ -28,15 +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::builders::ValidityBuilder; use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::dtype::Nullability; @@ -59,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, @@ -68,7 +67,7 @@ pub struct ListViewBuilder { sizes_builder: PrimitiveBuilder, /// The null map builder of the [`ListViewArray`]. - nulls: LazyBitBufferBuilder, + nulls: ValidityBuilder, /// Whether the appends so far leave the result zero-copyable to a [`ListArray`]. /// @@ -105,14 +104,14 @@ 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); let sizes_builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, capacity); - let nulls = LazyBitBufferBuilder::new(capacity); + let nulls = ValidityBuilder::new(capacity); Self { dtype: DType::List(element_dtype, nullability), @@ -152,8 +151,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 +204,58 @@ 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(()) + } + /// 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()); @@ -260,8 +310,7 @@ impl ListViewBuilder { return Ok(()); } - self.nulls - .append_validity_mask(&array.validity()?.execute_mask(array.len(), ctx)?); + self.nulls.append_validity(array.validity()?, array.len()); let offsets = array.offsets().clone().execute::(ctx)?; match_each_integer_ptype!(offsets.ptype(), |OffsetType| { @@ -294,56 +343,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)?; - - // 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(); - - self.nulls - .append_validity_mask(&array.validity()?.execute_mask(array.len(), ctx)?); - - // Bulk append the trimmed elements; the offsets are rebased onto them below. - 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)?; - 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, + let len = array.len(); + + // 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)? + }; + + // 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" + ); + + 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(()) } } @@ -422,10 +500,6 @@ impl ArrayBuilder for ListViewBuil self.nulls.reserve_exact(capacity); } - unsafe fn set_validity_unchecked(&mut self, validity: Mask) { - self.nulls = LazyBitBufferBuilder::from_validity_mask(validity); - } - fn finish(&mut self) -> ArrayRef { self.finish_into_listview().into_array() } @@ -465,10 +539,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); @@ -495,44 +568,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)] @@ -900,6 +996,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/map.rs b/vortex-array/src/builders/map.rs index 362c2211a30..6714cf8f881 100644 --- a/vortex-array/src/builders/map.rs +++ b/vortex-array/src/builders/map.rs @@ -6,7 +6,6 @@ use std::sync::Arc; use vortex_error::VortexResult; use vortex_error::vortex_ensure; -use vortex_mask::Mask; use crate::ArrayRef; use crate::Canonical; @@ -147,10 +146,6 @@ impl ArrayBuilder for MapBuilder ArrayRef { self.finish_into_map().into_array() } diff --git a/vortex-array/src/builders/mod.rs b/vortex-array/src/builders/mod.rs index b50a114ebe6..22c806f74a1 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: //! //! ``` @@ -34,7 +40,6 @@ use std::any::Any; use std::sync::Arc; use vortex_error::VortexResult; -use vortex_mask::Mask; use crate::ArrayRef; use crate::ExecutionCtx; @@ -49,6 +54,7 @@ mod lazy_null_builder; pub(crate) use lazy_null_builder::LazyBitBufferBuilder; mod bool; +mod child; mod decimal; pub mod dict; mod extension; @@ -59,9 +65,11 @@ mod map; mod null; mod primitive; mod struct_; +mod validity; mod varbinview; pub use bool::*; +pub(crate) use child::ChildBuilder; pub use decimal::*; pub use extension::*; pub use fixed_size_list::*; @@ -71,6 +79,7 @@ pub use map::*; pub use null::*; pub use primitive::*; pub use struct_::*; +pub(crate) use validity::ValidityBuilder; pub use varbinview::*; pub use crate::arrays::varbin::builder::VarBinBuilder; @@ -160,26 +169,11 @@ pub trait ArrayBuilder: Send { /// Allocate space for extra `additional` items fn reserve_exact(&mut self, additional: usize); - /// 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. - fn set_validity(&mut self, validity: Mask) { - if !self.dtype().is_nullable() { - return; - } - assert_eq!(self.len(), validity.len()); - unsafe { self.set_validity_unchecked(validity) } - } - - /// override validity with the one provided, without checking lengths - /// - /// # Safety - /// - /// Given validity must have an equal length to [`self.len()`](Self::len). - unsafe fn set_validity_unchecked(&mut self, validity: Mask); - /// 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 @@ -228,6 +222,24 @@ macro_rules! match_each_list_builder { }}; } +/// Matches a `&mut dyn ArrayBuilder` against every concrete [`ListViewBuilder`]`` +/// instantiation over the [`OffsetBuilderPType`](crate::dtype::OffsetBuilderPType) offset/size +/// types (`u32`, `u64`, `i32`, `i64`), and only those. +/// +/// Binds the downcast builder as `$builder` and evaluates `$body` with it, yielding +/// `Some($body)`; yields `None` when the builder is not a list-view builder - including when it +/// is a [`ListBuilder`]. Callers reach for this instead of +/// [`match_each_list_builder!`](crate::match_each_list_builder) when the body needs methods only +/// a list-view builder has, such as +/// [`append_array_as_repeated_list`](ListViewBuilder::append_array_as_repeated_list). +#[macro_export] +macro_rules! match_each_listview_builder { + ($dyn_builder:expr, | $builder:ident | $body:expr) => {{ + let __dyn_builder: &mut dyn $crate::builders::ArrayBuilder = $dyn_builder; + $crate::__match_each_listview_builder!(__dyn_builder, $builder, $body, [u32, u64, i32, i64]) + }}; +} + #[doc(hidden)] #[macro_export] macro_rules! __match_each_list_builder { diff --git a/vortex-array/src/builders/null.rs b/vortex-array/src/builders/null.rs index 28eab14dd7a..541a27c2c89 100644 --- a/vortex-array/src/builders/null.rs +++ b/vortex-array/src/builders/null.rs @@ -5,7 +5,6 @@ use std::any::Any; use vortex_error::VortexResult; use vortex_error::vortex_ensure; -use vortex_mask::Mask; use crate::ArrayRef; use crate::ExecutionCtx; @@ -72,8 +71,6 @@ impl ArrayBuilder for NullBuilder { fn reserve_exact(&mut self, _additional: usize) {} - unsafe fn set_validity_unchecked(&mut self, _validity: Mask) {} - fn finish(&mut self) -> ArrayRef { NullArray::new(self.length).into_array() } diff --git a/vortex-array/src/builders/primitive.rs b/vortex-array/src/builders/primitive.rs index 4ce1aafe1f7..aca2db36286 100644 --- a/vortex-array/src/builders/primitive.rs +++ b/vortex-array/src/builders/primitive.rs @@ -201,10 +201,6 @@ impl ArrayBuilder for PrimitiveBuilder { self.nulls.reserve_exact(additional); } - unsafe fn set_validity_unchecked(&mut self, validity: Mask) { - self.nulls = LazyBitBufferBuilder::from_validity_mask(validity); - } - fn finish(&mut self) -> ArrayRef { self.finish_into_primitive().into_array() } diff --git a/vortex-array/src/builders/struct_.rs b/vortex-array/src/builders/struct_.rs index da08e262760..66a2c66358d 100644 --- a/vortex-array/src/builders/struct_.rs +++ b/vortex-array/src/builders/struct_.rs @@ -9,7 +9,6 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; -use vortex_mask::Mask; use crate::ArrayRef; use crate::ExecutionCtx; @@ -17,9 +16,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::builders::ValidityBuilder; use crate::canonical::Canonical; use crate::dtype::DType; use crate::dtype::Nullability; @@ -30,8 +29,8 @@ use crate::scalar::StructScalar; /// The builder for building a [`StructArray`]. pub struct StructBuilder { dtype: DType, - builders: Vec>, - nulls: LazyBitBufferBuilder, + builders: Vec, + nulls: ValidityBuilder, } impl StructBuilder { @@ -48,12 +47,12 @@ impl StructBuilder { ) -> Self { let builders = struct_dtype .fields() - .map(|dt| builder_with_capacity(&dt, capacity)) + .map(|dt| ChildBuilder::with_capacity(&dt, capacity)) .collect(); Self { builders, - nulls: LazyBitBufferBuilder::new(capacity), + nulls: ValidityBuilder::new(capacity), dtype: DType::Struct(struct_dtype, nullability), } } @@ -131,11 +130,10 @@ 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 - .append_validity_mask(&array.validity()?.execute_mask(array.len(), ctx)?); + self.nulls.append_validity(array.validity()?, array.len()); Ok(()) } } @@ -191,10 +189,6 @@ impl ArrayBuilder for StructBuilder { self.nulls.reserve_exact(capacity); } - unsafe fn set_validity_unchecked(&mut self, validity: Mask) { - self.nulls = LazyBitBufferBuilder::from_validity_mask(validity); - } - fn finish(&mut self) -> ArrayRef { self.finish_into_struct().into_array() } diff --git a/vortex-array/src/builders/tests.rs b/vortex-array/src/builders/tests.rs index a9db688f239..914cd04e1b1 100644 --- a/vortex-array/src/builders/tests.rs +++ b/vortex-array/src/builders/tests.rs @@ -4,16 +4,43 @@ 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::dtype::DType; use crate::dtype::DecimalDType; @@ -24,6 +51,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`. /// @@ -837,66 +865,362 @@ fn test_append_scalar_repeated_same_instance() { } } -/// Test that `set_validity` correctly overrides a builder's validity across all mask variants. +/// 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. /// -/// `set_validity` moves the mask's buffer into the builder rather than copying it, so the -/// `sliced_offset` case is important: slicing a `Mask::Values` at a non-byte-aligned boundary -/// yields a buffer with a non-zero bit offset, which the move path must preserve. +/// 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::all_true(Mask::new_true(8), vec![true; 8])] -#[case::all_false(Mask::new_false(8), vec![false; 8])] -#[case::values( - Mask::from_iter([true, false, true, true, false, false, true, false]), - vec![true, false, true, true, false, false, true, false] +#[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::sliced_offset( - Mask::from_iter([ - false, false, false, // dropped by the slice - true, false, true, true, false, false, true, false, // kept: indices 3..11 - true, true, true, true, true, // dropped by the slice - ]) - .slice(3..11), - vec![true, false, true, true, false, false, true, false] +#[case::list_elements( + ListViewArray::new( + constant_i32(), + (0..CHUNK_LEN as u64).collect::>().into_array(), + Buffer::full(1u64, CHUNK_LEN).into_array(), + Validity::NonNullable, + ) + .into_array(), + |array: &ArrayRef| array.as_::().elements().clone() )] -fn test_set_validity_overrides_validity( - #[case] mask: Mask, - #[case] expected: Vec, +#[case::fixed_size_list_elements( + FixedSizeListArray::new( + constant_i32(), + 2, + Validity::NonNullable, + 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), + 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 dtype = DType::Primitive(PType::I32, Nullability::Nullable); - let mut builder = builder_with_capacity(&dtype, mask.len()); - builder.append_zeros(mask.len()); + 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", + ); - builder.set_validity(mask); + let expected = ChunkedArray::try_new(vec![array.clone(), array], built.dtype().clone())?; + assert_arrays_eq!(&built, &expected, &mut ctx); - let validity = builder.finish().validity()?; + Ok(()) +} + +/// A child is chunked on the boundaries it is appended on, however small the appends. Appending +/// scalars instead is what asks the builder to copy the values into one canonical child. +#[test] +fn test_children_are_chunked_on_the_boundaries_they_are_appended_on() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); - for (i, &valid) in expected.iter().enumerate() { - assert_eq!( - validity.execute_is_valid(i, &mut ctx)?, - valid, - "validity mismatch at index {i}" - ); + + 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..CHUNK_LEN { + array.append_to_builder(builder.as_mut(), &mut ctx)?; } + let built = builder.finish(); + + assert_eq!(built.len(), CHUNK_LEN); + assert_eq!( + built + .as_::() + .elements() + .as_::() + .nchunks(), + CHUNK_LEN, + "one chunk per appended array", + ); + + // The same values appended as scalars land in a single canonical child. + let mut builder = builder_with_capacity(array.dtype(), 0); + let scalar = array.execute_scalar(0, &mut ctx)?; + for _ in 0..CHUNK_LEN { + builder.append_scalar(&scalar)?; + } + let built_from_scalars = builder.finish(); + + assert!( + built_from_scalars + .as_::() + .elements() + .is::() + ); + assert_arrays_eq!(&built_from_scalars, &built, &mut ctx); + Ok(()) } -/// Test that `set_validity` is a no-op on a non-nullable builder. +/// A builder that mixes appended arrays with scalar appends must keep the two in order. #[test] -fn test_set_validity_noop_when_non_nullable() -> VortexResult<()> { - let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut builder = builder_with_capacity(&dtype, 4); - builder.append_zeros(4); +fn test_struct_builder_interleaves_arrays_and_scalars() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); - // Providing an all-false mask must not make the non-nullable array invalid. - builder.set_validity(Mask::new_false(4)); + 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); - let validity = builder.finish().validity()?; + Ok(()) +} + +/// An arbitrary array length. Nested builders treat no length specially, so the tests only need a +/// length long enough to tell chunks apart. +const CHUNK_LEN: usize = 64; + +/// A nested builder's own validity is accumulated the same way its children are: an appended +/// array's validity is kept as it arrived rather than executed into a mask and copied bit by bit. +#[test] +fn test_appended_validity_is_not_materialized() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); - for i in 0..4 { - assert!( - validity.execute_is_valid(i, &mut ctx)?, - "index {i} should remain valid" - ); + + let array = StructArray::try_from_iter_with_validity( + [("a", iota(CHUNK_LEN))], + Validity::from_iter((0..CHUNK_LEN).map(|i| i % 3 != 0)), + )? + .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 Validity::Array(validity) = built.validity()? else { + panic!("expected array-backed validity"); + }; + assert!( + validity.is::(), + "the two appended validities should have been concatenated, not copied into one buffer", + ); + + 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 [`CHUNK_LEN`] `i32` values. +fn constant_i32() -> ArrayRef { + ConstantArray::new(0i32, CHUNK_LEN).into_array() +} + +/// Two lists of [`CHUNK_LEN`] elements each, so that appending them chunks the elements. +fn two_lists_of_chunk_len() -> ListViewArray { + ListViewArray::new( + iota(2 * CHUNK_LEN), + u64s([0, CHUNK_LEN]), + u64s([CHUNK_LEN, 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 * CHUNK_LEN), + u64s([0, CHUNK_LEN, 2 * 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, 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, CHUNK_LEN)) + .collect::>() + .into_array(), + u64s((0..=3).map(|i| i * 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(CHUNK_LEN), + 4, + Validity::from_iter((0..CHUNK_LEN / 4).map(|i| i % 3 != 0)), + CHUNK_LEN / 4, + ) + .into_array(), + |array: &ArrayRef| array.as_::().elements().clone() +)] +#[case::struct_( + StructArray::try_from_iter_with_validity( + [("a", iota(CHUNK_LEN))], + Validity::from_iter((0..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(()) +} + +/// 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(()) } diff --git a/vortex-array/src/builders/validity.rs b/vortex-array/src/builders/validity.rs new file mode 100644 index 00000000000..d5417407b2f --- /dev/null +++ b/vortex-array/src/builders/validity.rs @@ -0,0 +1,256 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexExpect; +use vortex_error::vortex_panic; + +use crate::builders::LazyBitBufferBuilder; +use crate::dtype::Nullability; +use crate::dtype::Nullability::NonNullable; +use crate::dtype::Nullability::Nullable; +use crate::validity::Validity; + +/// Accumulates the validity of a nested [`ArrayBuilder`](crate::builders::ArrayBuilder) without +/// materializing a null buffer for it. +/// +/// A nested builder learns about validity from two sources: one row at a time, as scalars are +/// appended, and a whole array's worth at a time, as arrays are. Only the former needs a null +/// buffer. An appended array already carries its validity in whatever form it was stored in — +/// [`Validity::AllValid`] and [`Validity::AllInvalid`] cost nothing at all, and an array-backed +/// validity is a bool array that is already built — so this builder keeps those as runs and +/// concatenates them at the end, exactly as [`Validity::concat`] does. +/// +/// Materializing them instead would mean executing every appended array's validity into a +/// [`Mask`](vortex_mask::Mask) and copying its bits, which for a builder assembling many chunks is +/// the dominant cost of tracking validity at all. +pub(crate) struct ValidityBuilder { + /// Completed runs, in logical order, with the number of values each covers. Never contains an + /// empty run. + runs: Vec<(Validity, usize)>, + + /// The summed length of `runs`. + runs_len: usize, + + /// Null buffer holding the bits appended since the last run. + pending: LazyBitBufferBuilder, +} + +impl ValidityBuilder { + /// Creates a new `ValidityBuilder` whose null buffer is pre-allocated for `capacity` bits. + pub fn new(capacity: usize) -> Self { + Self { + runs: Vec::new(), + runs_len: 0, + pending: LazyBitBufferBuilder::new(capacity), + } + } + + /// The number of values whose validity has been recorded so far. + pub fn len(&self) -> usize { + self.runs_len + self.pending.len() + } + + /// Records one valid value. + pub fn append_non_null(&mut self) { + self.pending.append_non_null() + } + + /// Records `n` valid values. + pub fn append_n_non_nulls(&mut self, n: usize) { + self.pending.append_n_non_nulls(n) + } + + /// Records `n` null values. + pub fn append_n_nulls(&mut self, n: usize) { + self.pending.append_n_nulls(n) + } + + /// Records the validity of a whole appended array, covering `len` values, as a run of its own. + /// + /// However few values the run covers, it is kept as it arrived rather than executed into a + /// mask, so a builder's validity is split on exactly the boundaries its children are. + pub fn append_validity(&mut self, validity: Validity, len: usize) { + if len == 0 { + return; + } + + self.flush_pending(); + self.runs_len += len; + self.runs.push((validity, len)); + } + + /// Allocates space for `additional` more bits in the null buffer. + pub fn reserve_exact(&mut self, additional: usize) { + self.pending.reserve_exact(additional) + } + + /// Finishes the validity, concatenating the accumulated runs. + /// + /// # Panics + /// + /// Panics if a non-nullable builder recorded a null, matching + /// [`LazyBitBufferBuilder::finish_with_nullability`]. + pub fn finish_with_nullability(&mut self, nullability: Nullability) -> Validity { + if self.runs.is_empty() { + return self.pending.finish_with_nullability(nullability); + } + + self.flush_pending(); + self.runs_len = 0; + let runs = std::mem::take(&mut self.runs); + + // `Validity::concat` treats `NonNullable` and `AllValid` as different kinds and falls back + // to a bool array when both appear, which they do as soon as a non-nullable array is + // appended next to a scalar. Both mean "no nulls", so answer from the nullability instead. + if runs + .iter() + .all(|(validity, _)| validity.definitely_no_nulls()) + { + return nullability.into(); + } + + let validity = Validity::concat(runs).vortex_expect("runs is not empty"); + if nullability == NonNullable { + vortex_panic!("cannot finish a non-nullable builder holding {validity:?} validity"); + } + validity + } + + /// Moves whatever the null buffer holds into `runs`, keeping the runs in logical order. + fn flush_pending(&mut self) { + let len = self.pending.len(); + if len == 0 { + return; + } + // A run is only ever read back through `Validity::concat`, which takes the nullability + // from the runs as a whole, so an all-valid null buffer can stay lazy here. + let validity = self.pending.finish_with_nullability(Nullable); + self.runs_len += len; + self.runs.push((validity, len)); + } +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + use vortex_mask::Mask; + + use super::ValidityBuilder; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::arrays::Chunked; + use crate::dtype::Nullability::NonNullable; + use crate::dtype::Nullability::Nullable; + use crate::validity::Validity; + + /// An arbitrary run length. `ValidityBuilder` treats no length specially, so the tests only + /// need a length long enough to tell runs apart. + const RUN_LEN: usize = 64; + + /// A `Validity` backed by a bool array, which is the case a run avoids executing. + fn array_backed(len: usize) -> Validity { + Validity::from_mask(Mask::from_iter((0..len).map(|i| i % 2 == 0)), Nullable) + } + + #[test] + fn test_whole_array_validity_is_kept_as_a_run() { + let mut builder = ValidityBuilder::new(0); + + builder.append_validity(array_backed(RUN_LEN), RUN_LEN); + builder.append_validity(array_backed(RUN_LEN), RUN_LEN); + assert_eq!(builder.len(), 2 * RUN_LEN); + + let Validity::Array(array) = builder.finish_with_nullability(Nullable) else { + panic!("expected array-backed validity"); + }; + assert!( + array.is::(), + "the runs should have been concatenated, not copied into one buffer", + ); + } + + /// Uniform runs collapse instead of becoming a bool array. + #[test] + fn test_all_valid_runs_stay_lazy() { + let mut builder = ValidityBuilder::new(0); + + builder.append_validity(Validity::AllValid, RUN_LEN); + builder.append_validity(Validity::AllValid, RUN_LEN); + + assert!(matches!( + builder.finish_with_nullability(Nullable), + Validity::AllValid + )); + } + + /// A one-row validity earns a run too. Uniform runs still collapse, so a builder fed an array + /// at a time does not pay a bool array for validity it never had. + #[test] + fn test_short_validity_is_kept_as_a_run_too() { + let mut builder = ValidityBuilder::new(0); + + for _ in 0..RUN_LEN { + builder.append_validity(Validity::AllInvalid, 1); + } + assert_eq!(builder.len(), RUN_LEN); + + assert!(matches!( + builder.finish_with_nullability(Nullable), + Validity::AllInvalid + )); + } + + /// Bits and runs interleave, and have to come back out in the order they went in. + #[test] + fn test_bits_and_runs_keep_their_order() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let mut builder = ValidityBuilder::new(0); + + builder.append_n_nulls(1); + builder.append_validity(Validity::AllValid, RUN_LEN); + builder.append_non_null(); + builder.append_validity(Validity::AllInvalid, RUN_LEN); + + let validity = builder.finish_with_nullability(Nullable); + let mask = validity.execute_mask(2 * RUN_LEN + 2, &mut ctx)?; + + let expected = Mask::from_iter( + [false] + .into_iter() + .chain(std::iter::repeat_n(true, RUN_LEN)) + .chain([true]) + .chain(std::iter::repeat_n(false, RUN_LEN)), + ); + assert_eq!(mask, expected); + + Ok(()) + } + + #[test] + fn test_non_nullable_finishes_non_nullable() { + let mut builder = ValidityBuilder::new(0); + + builder.append_validity(Validity::NonNullable, RUN_LEN); + builder.append_n_non_nulls(1); + + assert!(matches!( + builder.finish_with_nullability(NonNullable), + Validity::NonNullable + )); + } + + #[test] + fn test_finish_resets_the_builder() { + let mut builder = ValidityBuilder::new(0); + + builder.append_validity(Validity::AllInvalid, RUN_LEN); + assert_eq!(builder.finish_with_nullability(Nullable).maybe_len(), None); + + assert_eq!(builder.len(), 0); + builder.append_n_nulls(1); + assert!(matches!( + builder.finish_with_nullability(Nullable), + Validity::AllInvalid + )); + } +} diff --git a/vortex-array/src/builders/varbinview.rs b/vortex-array/src/builders/varbinview.rs index 1ec39138ebb..46b6e13874d 100644 --- a/vortex-array/src/builders/varbinview.rs +++ b/vortex-array/src/builders/varbinview.rs @@ -384,10 +384,6 @@ impl ArrayBuilder for VarBinViewBuilder { self.nulls.reserve_exact(additional); } - unsafe fn set_validity_unchecked(&mut self, validity: Mask) { - self.nulls = LazyBitBufferBuilder::from_validity_mask(validity); - } - fn finish(&mut self) -> ArrayRef { self.finish_into_varbinview().into_array() }