diff --git a/vortex-array/src/arrays/varbin/builder.rs b/vortex-array/src/arrays/varbin/builder.rs index c798d7ba02c..d7e8945b201 100644 --- a/vortex-array/src/arrays/varbin/builder.rs +++ b/vortex-array/src/arrays/varbin/builder.rs @@ -630,10 +630,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() } @@ -949,7 +945,6 @@ mod tests { 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); 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 index 892fcb0e30b..869bdb44071 100644 --- a/vortex-array/src/builders/child.rs +++ b/vortex-array/src/builders/child.rs @@ -1,22 +1,17 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; -use vortex_mask::Mask; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::ChunkedArray; -use crate::arrays::MaskedArray; use crate::builders::ArrayBuilder; use crate::builders::builder_with_capacity; use crate::dtype::DType; -use crate::dtype::Nullability; use crate::scalar::Scalar; -use crate::validity::Validity; /// Accumulates the child of a nested [`ArrayBuilder`] without canonicalizing appended arrays. /// @@ -121,42 +116,6 @@ impl ChildBuilder { self.pending.reserve_exact(additional) } - /// Overrides the validity of every value appended so far. - /// - /// # Safety - /// - /// `validity` must have the same length as [`self.len()`](Self::len). - /// - /// # Panics - /// - /// Panics if a chunk that was kept in its original encoding contains nulls, since replacing - /// the validity of such a chunk would require decoding it. - pub unsafe fn set_validity_unchecked(&mut self, validity: Mask) { - if !self.dtype.is_nullable() { - return; - } - - if self.chunks.is_empty() { - // Fast path: every value lives in the scalar builder, which owns its null buffer. - unsafe { self.pending.set_validity_unchecked(validity) }; - return; - } - - // The chunks carry their own validity, so the override has to be pushed into each of them. - self.flush_pending(); - let mut offset = 0; - for chunk in &mut self.chunks { - let end = offset + chunk.len(); - *chunk = MaskedArray::try_new( - chunk.clone(), - Validity::from_mask(validity.slice(offset..end), Nullability::Nullable), - ) - .vortex_expect("cannot override the validity of a child chunk that contains nulls") - .into_array(); - offset = end; - } - } - /// Finishes the child, combining the accumulated chunks into a [`ChunkedArray`] when there is /// more than one of them. pub fn finish(&mut self) -> ArrayRef { @@ -172,9 +131,7 @@ impl ChildBuilder { return chunks.remove(0); } - ChunkedArray::try_new(chunks, self.dtype.clone()) - .vortex_expect("every child chunk has the child dtype") - .into_array() + unsafe { ChunkedArray::new_unchecked(chunks, self.dtype.clone()) }.into_array() } /// Moves whatever the scalar builder holds into `chunks`, keeping the chunks in logical order. @@ -192,9 +149,7 @@ impl ChildBuilder { mod tests { use rstest::rstest; use vortex_buffer::buffer; - use vortex_error::VortexExpect; use vortex_error::VortexResult; - use vortex_mask::Mask; use super::ChildBuilder; use crate::ArrayRef; @@ -205,7 +160,6 @@ mod tests { use crate::arrays::ChunkedArray; use crate::arrays::Constant; use crate::arrays::ConstantArray; - use crate::arrays::Masked; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::chunked::ChunkedArrayExt; @@ -380,104 +334,6 @@ mod tests { Ok(()) } - /// Overriding the validity once chunks exist has to push the override into each chunk, sliced - /// to that chunk's own range. - #[test] - fn test_set_validity_pushes_the_override_into_every_chunk() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let dtype = DType::Primitive(I32, Nullable); - let mut builder = ChildBuilder::with_capacity(&dtype, 0); - - builder.append_array(&nullable_constant(1, CHUNK_LEN), &mut ctx)?; - builder.append_array(&nullable_constant(2, CHUNK_LEN), &mut ctx)?; - - // Straddle the chunk boundary, so an override sliced wrongly cannot pass. - let invalid = [CHUNK_LEN - 1, CHUNK_LEN]; - let validity = Mask::from_iter((0..2 * CHUNK_LEN).map(|i| !invalid.contains(&i))); - unsafe { builder.set_validity_unchecked(validity) }; - - let child = builder.finish(); - let chunked = child.as_::(); - assert_eq!(chunked.nchunks(), 2); - // The override was layered over the chunks rather than decoding them. - assert!(chunked.iter_chunks().all(|chunk| chunk.is::())); - - let expected = PrimitiveArray::from_option_iter( - (0..2 * CHUNK_LEN) - .map(|i| (!invalid.contains(&i)).then_some(if i < CHUNK_LEN { 1i32 } else { 2 })), - ) - .into_array(); - assert_arrays_eq!(&child, &expected, &mut ctx); - - Ok(()) - } - - /// Values still sitting in the scalar builder are part of the override too. - #[test] - fn test_set_validity_covers_pending_scalars() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let dtype = DType::Primitive(I32, Nullable); - let mut builder = ChildBuilder::with_capacity(&dtype, 0); - - builder.append_array(&nullable_constant(1, CHUNK_LEN), &mut ctx)?; - builder.append_scalar(&Scalar::primitive(2i32, Nullable))?; - - let mut validity = vec![true; CHUNK_LEN + 1]; - validity[CHUNK_LEN] = false; - unsafe { builder.set_validity_unchecked(Mask::from_iter(validity)) }; - - let child = builder.finish(); - let expected = PrimitiveArray::from_option_iter( - std::iter::repeat_n(Some(1i32), CHUNK_LEN).chain([None]), - ) - .into_array(); - assert_arrays_eq!(&child, &expected, &mut ctx); - - Ok(()) - } - - /// A non-nullable child cannot carry nulls, so the override is dropped and the chunks are left - /// exactly as they were appended. - #[test] - fn test_set_validity_is_a_noop_for_a_non_nullable_child() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let mut builder = ChildBuilder::with_capacity(&DType::from(I32), 0); - - builder.append_array(&constant(1, CHUNK_LEN), &mut ctx)?; - builder.append_array(&constant(2, CHUNK_LEN), &mut ctx)?; - unsafe { builder.set_validity_unchecked(Mask::new_false(2 * CHUNK_LEN)) }; - - let child = builder.finish(); - let chunked = child.as_::(); - assert!(chunked.iter_chunks().all(|chunk| chunk.is::())); - - let expected = ChunkedArray::try_new( - vec![constant(1, CHUNK_LEN), constant(2, CHUNK_LEN)], - DType::from(I32), - )? - .into_array(); - assert_arrays_eq!(&child, &expected, &mut ctx); - - Ok(()) - } - - /// Replacing the validity of a chunk that already contains nulls would mean decoding it, which - /// is exactly what the chunk exists to avoid. - #[test] - #[should_panic(expected = "cannot override the validity of a child chunk that contains nulls")] - fn test_set_validity_rejects_a_chunk_that_contains_nulls() { - let mut ctx = array_session().create_execution_ctx(); - let dtype = DType::Primitive(I32, Nullable); - let mut builder = ChildBuilder::with_capacity(&dtype, 0); - - let with_nulls = ConstantArray::new(Scalar::null(dtype), CHUNK_LEN).into_array(); - builder - .append_array(&with_nulls, &mut ctx) - .vortex_expect("append"); - - unsafe { builder.set_validity_unchecked(Mask::new_true(CHUNK_LEN)) }; - } - #[test] fn test_finish_resets_the_builder() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); 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 e756aef2696..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; @@ -114,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 8dd8a1c3ebc..2dc1749dbe0 100644 --- a/vortex-array/src/builders/fixed_size_list.rs +++ b/vortex-array/src/builders/fixed_size_list.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; @@ -19,7 +18,7 @@ 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::ValidityBuilder; use crate::canonical::Canonical; use crate::dtype::DType; use crate::dtype::Nullability; @@ -39,7 +38,7 @@ pub struct FixedSizeListBuilder { /// 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 { @@ -64,7 +63,7 @@ impl FixedSizeListBuilder { 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, @@ -116,8 +115,7 @@ impl FixedSizeListBuilder { } self.elements_builder.append_array(array.elements(), ctx)?; - self.nulls - .append_validity_mask(&array.validity()?.execute_mask(array.len(), ctx)?); + self.nulls.append_validity(array.validity()?, array.len()); Ok(()) } @@ -261,10 +259,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 4acfb64fecf..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; @@ -28,8 +27,8 @@ 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::ValidityBuilder; use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::dtype::Nullability; @@ -52,7 +51,7 @@ pub struct ListBuilder { offsets_builder: PrimitiveBuilder, /// The null map builder of the [`ListArray`]. - nulls: LazyBitBufferBuilder, + nulls: ValidityBuilder, } impl ListBuilder { @@ -90,7 +89,7 @@ impl ListBuilder { Self { elements_builder, offsets_builder, - nulls: LazyBitBufferBuilder::new(capacity), + nulls: ValidityBuilder::new(capacity), dtype: DType::List(value_dtype, nullability), } } @@ -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)?; @@ -242,8 +240,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()); // 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 @@ -379,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() } diff --git a/vortex-array/src/builders/listview.rs b/vortex-array/src/builders/listview.rs index c2fc1ff70a3..e1dd3af31d2 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -16,7 +16,6 @@ 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; @@ -35,7 +34,7 @@ use crate::builders::ChildBuilder; use crate::builders::DEFAULT_BUILDER_CAPACITY; use crate::builders::PrimitiveBuilder; use crate::builders::UninitRange; -use crate::builders::lazy_null_builder::LazyBitBufferBuilder; +use crate::builders::ValidityBuilder; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::IntegerPType; @@ -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`]. /// @@ -112,7 +111,7 @@ impl ListViewBuilder { 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), @@ -259,8 +258,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| { @@ -304,8 +302,7 @@ impl ListViewBuilder { // 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)?); + self.nulls.append_validity(array.validity()?, array.len()); // Bulk append the trimmed elements; the offsets are rebased onto them below. let old_elements_len = self.elements_builder.len(); @@ -418,10 +415,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() } diff --git a/vortex-array/src/builders/mod.rs b/vortex-array/src/builders/mod.rs index fd5bc8a4765..74019950170 100644 --- a/vortex-array/src/builders/mod.rs +++ b/vortex-array/src/builders/mod.rs @@ -40,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; @@ -66,6 +65,7 @@ mod map; mod null; mod primitive; mod struct_; +mod validity; mod varbinview; pub use bool::*; @@ -79,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; @@ -168,24 +169,6 @@ 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 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 f226d04e4b6..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; @@ -19,7 +18,7 @@ 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::ValidityBuilder; use crate::canonical::Canonical; use crate::dtype::DType; use crate::dtype::Nullability; @@ -31,7 +30,7 @@ use crate::scalar::StructScalar; pub struct StructBuilder { dtype: DType, builders: Vec, - nulls: LazyBitBufferBuilder, + nulls: ValidityBuilder, } impl StructBuilder { @@ -53,7 +52,7 @@ impl StructBuilder { Self { builders, - nulls: LazyBitBufferBuilder::new(capacity), + nulls: ValidityBuilder::new(capacity), dtype: DType::Struct(struct_dtype, nullability), } } @@ -134,8 +133,7 @@ impl StructBuilder { 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 c2bf5ad027e..914cd04e1b1 100644 --- a/vortex-array/src/builders/tests.rs +++ b/vortex-array/src/builders/tests.rs @@ -8,7 +8,6 @@ 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; @@ -866,70 +865,6 @@ fn test_append_scalar_repeated_same_instance() { } } -/// Test that `set_validity` correctly overrides a builder's validity across all mask variants. -/// -/// `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. -#[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::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] -)] -fn test_set_validity_overrides_validity( - #[case] mask: Mask, - #[case] expected: Vec, -) -> VortexResult<()> { - let dtype = DType::Primitive(PType::I32, Nullability::Nullable); - let mut builder = builder_with_capacity(&dtype, mask.len()); - builder.append_zeros(mask.len()); - - builder.set_validity(mask); - - let validity = builder.finish().validity()?; - 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}" - ); - } - Ok(()) -} - -/// Test that `set_validity` is a no-op on a non-nullable builder. -#[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); - - // Providing an all-false mask must not make the non-nullable array invalid. - builder.set_validity(Mask::new_false(4)); - - let validity = builder.finish().validity()?; - 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" - ); - } - Ok(()) -} - /// Builders only promise a canonical *top level*, so a child array that is long enough to be worth /// a chunk must come back out of the builder in the encoding it went in with. /// @@ -1083,6 +1018,37 @@ fn test_struct_builder_interleaves_arrays_and_scalars() -> VortexResult<()> { /// 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(); + + 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() @@ -1234,44 +1200,6 @@ fn test_validity_survives_chunked_children( Ok(()) } -/// An extension array has no validity of its own — it lives in the storage — so overriding the -/// builder's validity has to reach into the chunked storage. -#[test] -fn test_extension_set_validity_reaches_chunked_storage() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::Nullable).erased(); - - let array = ExtensionArray::new( - ext_dtype.clone(), - ConstantArray::new(Scalar::primitive(0i64, Nullability::Nullable), CHUNK_LEN).into_array(), - ) - .into_array(); - - let mut builder = builder_with_capacity(array.dtype(), 0); - array.append_to_builder(builder.as_mut(), &mut ctx)?; - array.append_to_builder(builder.as_mut(), &mut ctx)?; - - let invalid = [0, 2 * CHUNK_LEN - 1]; - builder.set_validity(Mask::from_iter( - (0..2 * CHUNK_LEN).map(|i| !invalid.contains(&i)), - )); - let built = builder.finish(); - - assert!(built.as_::().storage().is::()); - - let expected = ExtensionArray::new( - ext_dtype, - PrimitiveArray::from_option_iter( - (0..2 * CHUNK_LEN).map(|i| (!invalid.contains(&i)).then_some(0i64)), - ) - .into_array(), - ) - .into_array(); - assert_arrays_eq!(&built, &expected, &mut ctx); - - Ok(()) -} - /// Consumers that genuinely need a fully-decoded tree ask for it, and must still get one. #[test] fn test_chunked_children_canonicalize_recursively() -> VortexResult<()> { diff --git a/vortex-array/src/builders/validity.rs b/vortex-array/src/builders/validity.rs new file mode 100644 index 00000000000..5265d363861 --- /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`] 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() }