diff --git a/vortex-array/src/arrays/constant/vtable/canonical.rs b/vortex-array/src/arrays/constant/vtable/canonical.rs index 686e59df76e..99eb800de78 100644 --- a/vortex-array/src/arrays/constant/vtable/canonical.rs +++ b/vortex-array/src/arrays/constant/vtable/canonical.rs @@ -20,7 +20,6 @@ use crate::ExecutionCtx; use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::BoolArray; -use crate::arrays::ChunkedArray; use crate::arrays::Constant; use crate::arrays::ConstantArray; use crate::arrays::DecimalArray; @@ -30,6 +29,7 @@ use crate::arrays::ListViewArray; use crate::arrays::MapArray; use crate::arrays::NullArray; use crate::arrays::PrimitiveArray; +use crate::arrays::RepeatedArray; use crate::arrays::StructArray; use crate::arrays::UnionArray; use crate::arrays::VarBinViewArray; @@ -461,7 +461,7 @@ fn constant_canonical_fixed_size_list_array( /// /// A fixed-size list holds its elements back to back, so the rows cannot share one range the way a /// list view's can. Uniform elements stay a [`ConstantArray`] covering the whole run; otherwise the -/// row materializes once and the run chunks that single copy. +/// row materializes once and the run tiles that single copy. fn tile_fixed_size_list_elements( values: &[Scalar], element_dtype: &DType, @@ -490,14 +490,9 @@ fn tile_fixed_size_list_elements( return tile; } - // SAFETY: every chunk is `tile` itself, so they share its dtype and none is empty. - unsafe { - ChunkedArray::new_unchecked( - std::iter::repeat_n(tile, len).collect::>(), - element_dtype.clone(), - ) - } - .into_array() + RepeatedArray::try_new(tile, len) + .vortex_expect("repeated elements length is the run's elements length") + .into_array() } } } @@ -521,7 +516,6 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; use crate::arrays::BoolArray; - use crate::arrays::Chunked; use crate::arrays::Constant; use crate::arrays::ConstantArray; use crate::arrays::DecimalArray; @@ -530,14 +524,15 @@ mod tests { use crate::arrays::ListViewArray; use crate::arrays::NullArray; use crate::arrays::PrimitiveArray; + use crate::arrays::Repeated; use crate::arrays::StructArray; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; - use crate::arrays::chunked::ChunkedArrayExt; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::listview::ListViewRebuildMode; + use crate::arrays::repeated::RepeatedArraySlotsExt; use crate::arrays::struct_::StructArrayExt; use crate::assert_arrays_eq; use crate::dtype::DType; @@ -1072,7 +1067,7 @@ mod tests { } } - /// Mixed elements should materialize once, with every row chunking that one copy. + /// Mixed elements should materialize once, with every row tiling that one copy. #[test] fn test_canonicalize_fixed_size_list_tiles_one_copy_of_mixed_elements() { let mut ctx = SESSION.create_execution_ctx(); @@ -1092,16 +1087,13 @@ mod tests { .unwrap(); assert_eq!(canonical.elements().len(), 2 * LEN); + let repeated = canonical.elements().as_::(); assert_eq!( - canonical.elements().as_::().nchunks(), + repeated.count(), LEN, "every row should reference the same tile rather than hold its own copy", ); - let chunks = canonical.elements().as_::(); - let tile_ptr = chunks.chunk(0).buffer_handles()[0].as_host().as_ptr(); - for chunk in chunks.iter_chunks() { - assert_eq!(chunk.buffer_handles()[0].as_host().as_ptr(), tile_ptr); - } + assert_eq!(repeated.child().len(), 2); for index in [0, 1, LEN - 1] { assert_arrays_eq!( canonical.fixed_size_list_elements_at(index).unwrap(), diff --git a/vortex-array/src/arrays/constant/vtable/mod.rs b/vortex-array/src/arrays/constant/vtable/mod.rs index 3b55071ffb8..804b3591cf2 100644 --- a/vortex-array/src/arrays/constant/vtable/mod.rs +++ b/vortex-array/src/arrays/constant/vtable/mod.rs @@ -424,14 +424,12 @@ mod tests { use crate::ExecutionCtx; 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; @@ -752,8 +750,8 @@ mod tests { 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. + /// A fixed-size list whose elements are all the same scalar repeats a constant array, so the + /// elements collapse into one constant array rather than materializing a value per row. #[test] fn test_uniform_fixed_size_list_constant_append_keeps_elements_constant() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); @@ -777,12 +775,10 @@ mod tests { let elements = result.as_::().elements().clone(); assert!( - elements - .as_::() - .iter_chunks() - .all(|chunk| chunk.is::()), + elements.is::(), "a uniform tile should have stayed constant-encoded", ); + assert_eq!(elements.len(), 2_000); Ok(()) } diff --git a/vortex-array/src/arrays/mod.rs b/vortex-array/src/arrays/mod.rs index 0b5bdcf29df..5f1bf6117fe 100644 --- a/vortex-array/src/arrays/mod.rs +++ b/vortex-array/src/arrays/mod.rs @@ -11,7 +11,8 @@ //! Utility and lazy arrays represent common transformations without immediately materializing //! their result. Examples include [`ChunkedArray`] for concatenation, [`ConstantArray`] for repeated //! values, [`DictArray`] for dictionary encoding, [`FilterArray`] for masked rows, [`SliceArray`] -//! for views, and [`ScalarFnArray`] for deferred scalar-function execution. +//! for views, [`RepeatedArray`] for an array repeated end to end, and [`ScalarFnArray`] for +//! deferred scalar-function execution. //! //! Some public arrays are primarily internal building blocks. Their constructors and extension //! traits document the stable contract; avoid depending on undocumented slot order or metadata @@ -102,6 +103,10 @@ pub mod primitive; pub use primitive::Primitive; pub use primitive::PrimitiveArray; +pub mod repeated; +pub use repeated::Repeated; +pub use repeated::RepeatedArray; + pub mod runend; pub use runend::RunEnd; pub use runend::RunEndArray; diff --git a/vortex-array/src/arrays/repeated/array.rs b/vortex-array/src/arrays/repeated/array.rs new file mode 100644 index 00000000000..cca937a59f3 --- /dev/null +++ b/vortex-array/src/arrays/repeated/array.rs @@ -0,0 +1,65 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt::Display; +use std::fmt::Formatter; + +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::array::ArrayParts; +use crate::array_slots; +use crate::arrays::Repeated; +use crate::arrays::RepeatedArray; + +/// The child slots of a [`RepeatedArray`]. +#[array_slots(Repeated)] +pub struct RepeatedSlots { + /// The array that is repeated. + #[slot(0)] + pub child: ArrayRef, +} + +/// The data of a [`RepeatedArray`]: how many times its child is repeated. +#[derive(Clone, Debug)] +pub struct RepeatedData { + pub(super) count: usize, +} + +impl Display for RepeatedData { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "count: {}", self.count) + } +} + +impl RepeatedData { + /// Returns the number of times the child is repeated. + pub fn count(&self) -> usize { + self.count + } +} + +/// Returns the length of `count` copies of an array of `child_len` rows. +pub(super) fn repeated_len(child_len: usize, count: usize) -> VortexResult { + child_len.checked_mul(count).ok_or_else(|| { + vortex_err!("Repeated length overflows usize: {child_len} rows repeated {count} times") + }) +} + +impl RepeatedArray { + /// Creates an array that repeats `child` end to end `count` times, without reading any values. + /// + /// Returns an error if the length of the result overflows `usize`. + pub fn try_new(child: ArrayRef, count: usize) -> VortexResult { + let len = repeated_len(child.len(), count)?; + let dtype = child.dtype().clone(); + Self::try_from_parts(ArrayParts::new( + Repeated, + dtype, + len, + RepeatedData { count }, + RepeatedSlots { child }.into_slots(), + )) + } +} diff --git a/vortex-array/src/arrays/repeated/mod.rs b/vortex-array/src/arrays/repeated/mod.rs new file mode 100644 index 00000000000..5ce3502d06d --- /dev/null +++ b/vortex-array/src/arrays/repeated/mod.rs @@ -0,0 +1,29 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! An in-memory encoding that repeats a child array end to end. +//! +//! A [`RepeatedArray`] with child `c` and count `n` has `c.len() * n` rows, and row `i` is row +//! `i % c.len()` of `c`. For example, repeating `[1, 2, 3]` twice gives `[1, 2, 3, 1, 2, 3]`. +//! +//! Execution decodes the child once and appends it `n` times, so the child is never decoded per +//! copy. +//! +//! `Repeated` has no wire form, so the default [`ArraySession`] does not register it, and an array +//! that holds a `Repeated` cannot be serialized. +//! +//! [`ArraySession`]: crate::session::ArraySession + +mod array; +mod rules; +mod vtable; + +pub use array::RepeatedArraySlotsExt; +pub use array::RepeatedData; +pub use array::RepeatedSlots; +pub use array::RepeatedSlotsView; +pub use vtable::Repeated; +pub use vtable::RepeatedArray; + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/arrays/repeated/rules.rs b/vortex-array/src/arrays/repeated/rules.rs new file mode 100644 index 00000000000..e4562288974 --- /dev/null +++ b/vortex-array/src/arrays/repeated/rules.rs @@ -0,0 +1,44 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::Repeated; +use crate::arrays::RepeatedArray; +use crate::arrays::repeated::RepeatedArraySlotsExt; +use crate::arrays::slice::SliceReduce; +use crate::arrays::slice::SliceReduceAdaptor; +use crate::optimizer::rules::ParentRuleSet; + +pub(super) const PARENT_RULES: ParentRuleSet = + ParentRuleSet::new(&[ParentRuleSet::lift(&SliceReduceAdaptor(Repeated))]); + +impl SliceReduce for Repeated { + fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { + // The range is non-empty, so the child is too. + let child = array.child(); + let child_len = child.len(); + let first_copy = range.start / child_len; + let offset = first_copy * child_len; + + // A range inside one copy is a slice of the child. + if range.end - offset <= child_len { + return Ok(Some(child.slice(range.start - offset..range.end - offset)?)); + } + + // A range of whole copies is a shorter repetition of the child. + if range.start.is_multiple_of(child_len) && range.end.is_multiple_of(child_len) { + let count = range.len() / child_len; + return Ok(Some( + RepeatedArray::try_new(child.clone(), count)?.into_array(), + )); + } + + Ok(None) + } +} diff --git a/vortex-array/src/arrays/repeated/tests.rs b/vortex-array/src/arrays/repeated/tests.rs new file mode 100644 index 00000000000..722a6b1506a --- /dev/null +++ b/vortex-array/src/arrays/repeated/tests.rs @@ -0,0 +1,232 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; +use std::sync::LazyLock; + +use rstest::rstest; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::arrays::Constant; +use crate::arrays::ConstantArray; +use crate::arrays::Primitive; +use crate::arrays::PrimitiveArray; +use crate::arrays::Repeated; +use crate::arrays::RepeatedArray; +use crate::arrays::StructArray; +use crate::arrays::VarBinViewArray; +use crate::arrays::repeated::RepeatedArraySlotsExt; +use crate::assert_arrays_eq; +use crate::builders::builder_with_capacity_in; +use crate::optimizer::ArrayOptimizer; +use crate::scalar::Scalar; + +static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + +fn repeated(child: impl IntoArray, count: usize) -> VortexResult { + Ok(RepeatedArray::try_new(child.into_array(), count)?.into_array()) +} + +#[test] +fn repeats_primitive() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = repeated(buffer![1i32, 2, 3], 3)?; + + assert_eq!(array.len(), 9); + assert_arrays_eq!( + array, + PrimitiveArray::from_iter([1i32, 2, 3, 1, 2, 3, 1, 2, 3]), + &mut ctx + ); + Ok(()) +} + +#[test] +fn repeats_nulls() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = repeated(PrimitiveArray::from_option_iter([Some(1i64), None]), 2)?; + + assert_arrays_eq!( + array, + PrimitiveArray::from_option_iter([Some(1i64), None, Some(1), None]), + &mut ctx + ); + assert_eq!(array.invalid_count(&mut ctx)?, 2); + Ok(()) +} + +#[test] +fn repeats_strings() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let child = VarBinViewArray::from_iter_nullable_str([Some("a"), None, Some("a long string")]); + let array = repeated(child, 2)?; + + assert_arrays_eq!( + array, + VarBinViewArray::from_iter_nullable_str([ + Some("a"), + None, + Some("a long string"), + Some("a"), + None, + Some("a long string"), + ]), + &mut ctx + ); + Ok(()) +} + +#[test] +fn repeats_struct() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let child = StructArray::try_from_iter([ + ("a", buffer![1u8, 2].into_array()), + ("b", VarBinViewArray::from_iter_str(["x", "y"]).into_array()), + ])?; + let expected = StructArray::try_from_iter([ + ("a", buffer![1u8, 2, 1, 2, 1, 2].into_array()), + ( + "b", + VarBinViewArray::from_iter_str(["x", "y", "x", "y", "x", "y"]).into_array(), + ), + ])?; + + assert_arrays_eq!(repeated(child, 3)?, expected, &mut ctx); + Ok(()) +} + +#[rstest] +#[case::zero_count(buffer![1i32, 2, 3].into_array(), 0)] +#[case::empty_child(PrimitiveArray::from_iter(Vec::::new()).into_array(), 4)] +fn repeats_to_empty(#[case] child: ArrayRef, #[case] count: usize) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = repeated(child, count)?; + + assert!(array.is_empty()); + assert_arrays_eq!( + array, + PrimitiveArray::from_iter(Vec::::new()), + &mut ctx + ); + Ok(()) +} + +#[test] +fn repeats_single_row() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = repeated(PrimitiveArray::from_option_iter([Some(7u16)]), 4)?; + + assert_arrays_eq!( + array, + PrimitiveArray::from_option_iter([Some(7u16); 4]), + &mut ctx + ); + Ok(()) +} + +#[test] +fn scalar_at_wraps() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = repeated(buffer![10u32, 20, 30], 5)?; + + for (index, expected) in [(0, 10u32), (2, 30), (3, 10), (7, 20), (14, 30)] { + assert_eq!( + array.execute_scalar(index, &mut ctx)?, + Scalar::from(expected) + ); + } + Ok(()) +} + +#[test] +fn append_to_builder() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = repeated(buffer![1i8, 2], 3)?; + let mut builder = builder_with_capacity_in(array.dtype(), array.len() + 1, ctx.allocator()); + builder.append_scalar(&Scalar::from(0i8))?; + array.append_to_builder(builder.as_mut(), &mut ctx)?; + + assert_arrays_eq!( + builder.finish(), + PrimitiveArray::from_iter([0i8, 1, 2, 1, 2, 1, 2]), + &mut ctx + ); + Ok(()) +} + +#[rstest] +#[case::within_first_copy(0..2, vec![0, 1])] +#[case::within_later_copy(5..8, vec![1, 2, 3])] +#[case::exactly_one_copy(4..8, vec![0, 1, 2, 3])] +#[case::whole_copies(4..12, vec![0, 1, 2, 3, 0, 1, 2, 3])] +#[case::across_copies(3..6, vec![3, 0, 1])] +#[case::across_many_copies(1..11, vec![1, 2, 3, 0, 1, 2, 3, 0, 1, 2])] +fn slice(#[case] range: Range, #[case] expected: Vec) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = repeated(buffer![0i32, 1, 2, 3], 3)?; + + assert_arrays_eq!( + array.slice(range)?, + PrimitiveArray::from_iter(expected), + &mut ctx + ); + Ok(()) +} + +#[test] +fn slice_of_whole_copies_stays_repeated() -> VortexResult<()> { + let array = repeated(buffer![0i32, 1, 2, 3], 3)?; + let sliced = array.slice(4..12)?; + + let repeated = sliced.as_::(); + assert_eq!(repeated.count(), 2); + assert_eq!(repeated.child().len(), 4); + Ok(()) +} + +#[test] +fn reduce_single_count_to_child() -> VortexResult<()> { + let child = buffer![1i32, 2].into_array(); + let reduced = repeated(child.clone(), 1)?.optimize()?; + + assert!(reduced.is::()); + assert_eq!(reduced.len(), child.len()); + Ok(()) +} + +#[test] +fn reduce_constant_child() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let reduced = repeated(ConstantArray::new(5i32, 3), 4)?.optimize()?; + + assert!(reduced.is::()); + assert_arrays_eq!(reduced, PrimitiveArray::from_iter([5i32; 12]), &mut ctx); + Ok(()) +} + +#[test] +fn reduce_nested_repeats() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let reduced = repeated(repeated(buffer![1i32, 2], 2)?, 3)?.optimize()?; + + let flattened = reduced.as_::(); + assert_eq!(flattened.count(), 6); + assert!(!flattened.child().is::()); + assert_arrays_eq!( + reduced, + PrimitiveArray::from_iter([1i32, 2].repeat(6)), + &mut ctx + ); + Ok(()) +} + +#[test] +fn rejects_overflowing_length() { + let result = RepeatedArray::try_new(buffer![1u8, 2].into_array(), usize::MAX); + assert!(result.is_err()); +} diff --git a/vortex-array/src/arrays/repeated/vtable.rs b/vortex-array/src/arrays/repeated/vtable.rs new file mode 100644 index 00000000000..a5cb2281212 --- /dev/null +++ b/vortex-array/src/arrays/repeated/vtable.rs @@ -0,0 +1,248 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::hash::Hash; +use std::hash::Hasher; +use std::iter; + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::AnyCanonical; +use crate::ArrayEq; +use crate::ArrayHash; +use crate::ArrayRef; +use crate::Canonical; +use crate::EqMode; +use crate::ExecutionCtx; +use crate::ExecutionResult; +use crate::IntoArray; +use crate::array::Array; +use crate::array::ArrayId; +use crate::array::ArrayParts; +use crate::array::ArrayView; +use crate::array::OperationsVTable; +use crate::array::ProbeState; +use crate::array::VTable; +use crate::array::ValidityVTable; +use crate::array::with_empty_buffers; +use crate::arrays::ChunkedArray; +use crate::arrays::Constant; +use crate::arrays::ConstantArray; +use crate::arrays::repeated::RepeatedArraySlotsExt; +use crate::arrays::repeated::RepeatedData; +use crate::arrays::repeated::RepeatedSlots; +use crate::arrays::repeated::array::repeated_len; +use crate::arrays::repeated::rules::PARENT_RULES; +use crate::buffer::BufferHandle; +use crate::builders::ArrayBuilder; +use crate::dtype::DType; +use crate::require_child; +use crate::scalar::Scalar; +use crate::serde::ArrayChildren; +use crate::validity::Validity; + +/// The encoding of a [`RepeatedArray`]. +#[derive(Clone, Debug)] +pub struct Repeated; + +/// A [`Repeated`]-encoded array, which repeats its child end to end. +pub type RepeatedArray = Array; + +impl ArrayHash for RepeatedData { + fn array_hash(&self, state: &mut H, _accuracy: EqMode) { + self.count.hash(state); + } +} + +impl ArrayEq for RepeatedData { + fn array_eq(&self, other: &Self, _accuracy: EqMode) -> bool { + self.count == other.count + } +} + +impl VTable for Repeated { + type TypedArrayData = RepeatedData; + + type OperationsVTable = Self; + type ValidityVTable = Self; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.repeated"); + *ID + } + + fn validate( + &self, + data: &RepeatedData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + vortex_ensure_eq!(slots.len(), RepeatedSlots::COUNT); + let child = slots[RepeatedSlots::CHILD] + .as_ref() + .ok_or_else(|| vortex_err!("Repeated requires a child"))?; + vortex_ensure_eq!( + child.dtype(), + dtype, + "Repeated child dtype does not match outer dtype" + ); + vortex_ensure_eq!( + repeated_len(child.len(), data.count)?, + len, + "Repeated length does not match child length times count" + ); + Ok(()) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + vortex_panic!("Repeated buffer index {idx} out of bounds") + } + + fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option { + None + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + + fn serialize( + _array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + Ok(None) + } + + fn deserialize( + &self, + _dtype: &DType, + _len: usize, + _metadata: &[u8], + _buffers: &[BufferHandle], + _children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + vortex_bail!("Repeated array is not serializable") + } + + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + if array.is_empty() { + return Ok(()); + } + let child = array + .child() + .clone() + .execute::(ctx)? + .into_array(); + for _ in 0..array.count { + child.append_to_builder(builder, ctx)?; + } + Ok(()) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + RepeatedSlots::NAMES[idx].to_string() + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + if array.child().len() == 1 { + let scalar = array.child().execute_scalar(0, ctx)?; + return Ok(ExecutionResult::done( + ConstantArray::new(scalar, array.len()).into_array(), + )); + } + + let array = require_child!(array, array.child(), RepeatedSlots::CHILD => AnyCanonical); + + // Decode the child once, then let `Chunked` append the decoded child once per copy. + let copies = iter::repeat_n(array.child().clone(), array.count); + let chunked = ChunkedArray::try_new(copies, array.dtype().clone())?; + Ok(ExecutionResult::done(chunked.into_array())) + } + + fn reduce(array: ArrayView<'_, Self>) -> VortexResult> { + if array.is_empty() { + return Ok(Some(Canonical::empty(array.dtype()).into_array())); + } + if array.count == 1 { + return Ok(Some(array.child().clone())); + } + if let Some(constant) = array.child().as_opt::() { + return Ok(Some( + ConstantArray::new(constant.scalar().clone(), array.len()).into_array(), + )); + } + if let Some(inner) = array.child().as_opt::() { + let count = inner + .count + .checked_mul(array.count) + .ok_or_else(|| vortex_err!("Repeated count overflows usize"))?; + return Ok(Some( + RepeatedArray::try_new(inner.child().clone(), count)?.into_array(), + )); + } + Ok(None) + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + PARENT_RULES.evaluate(array, parent, child_idx) + } +} + +impl OperationsVTable for Repeated { + type ProbeState = (); + + fn probe_scalar( + state: &mut ProbeState<'_, Repeated>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let child_index = index % state.array().child().len(); + state + .slot(RepeatedSlots::CHILD)? + .ok_or_else(|| vortex_err!("Repeated child slot is missing"))? + .execute_scalar(child_index, ctx) + } + + fn scalar_at( + array: ArrayView<'_, Repeated>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Self::probe_scalar(&mut ProbeState::once(array), index, ctx) + } +} + +impl ValidityVTable for Repeated { + fn validity(array: ArrayView<'_, Repeated>) -> VortexResult { + Ok(match array.child().validity()? { + Validity::Array(validity) => { + Validity::Array(RepeatedArray::try_new(validity, array.count)?.into_array()) + } + validity => validity, + }) + } +} diff --git a/vortex-array/src/builders/child.rs b/vortex-array/src/builders/child.rs index 923b7f76114..415a9a49288 100644 --- a/vortex-array/src/builders/child.rs +++ b/vortex-array/src/builders/child.rs @@ -66,8 +66,8 @@ impl ChildBuilder { /// /// 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 + /// chunks, 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 diff --git a/vortex-array/src/builders/fixed_size_list.rs b/vortex-array/src/builders/fixed_size_list.rs index 0584e138b5b..6a230f99ccd 100644 --- a/vortex-array/src/builders/fixed_size_list.rs +++ b/vortex-array/src/builders/fixed_size_list.rs @@ -14,8 +14,8 @@ use vortex_error::vortex_panic; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; -use crate::arrays::ChunkedArray; use crate::arrays::FixedSizeListArray; +use crate::arrays::RepeatedArray; use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use crate::builders::ArrayBuilder; use crate::builders::ChildBuilder; @@ -24,6 +24,7 @@ use crate::builders::ValidityBuilder; use crate::canonical::Canonical; use crate::dtype::DType; use crate::dtype::Nullability; +use crate::optimizer::ArrayOptimizer; use crate::scalar::ListScalar; use crate::scalar::Scalar; @@ -136,7 +137,7 @@ impl FixedSizeListBuilder { /// 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 + /// go in as a [`RepeatedArray`] of `array` repeated `n` times, 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 @@ -154,15 +155,11 @@ impl FixedSizeListBuilder { 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)?; + // Optimizing collapses a constant tile into one constant run. + let repeated = RepeatedArray::try_new(array.clone(), n)? + .into_array() + .optimize()?; + self.elements_builder.append_array(&repeated, ctx)?; self.nulls.append_n_non_nulls(n); Ok(())