Add repeated array that repeats a child array n times

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::<Vec<_>>(),
-                    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_::<Repeated>();
         assert_eq!(
-            canonical.elements().as_::<Chunked>().nchunks(),
+            repeated.count(),
             LEN,
             "every row should reference the same tile rather than hold its own copy",
         );
-        let chunks = canonical.elements().as_::<Chunked>();
-        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(),

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_::<FixedSizeList>().elements().clone();
         assert!(
-            elements
-                .as_::<Chunked>()
-                .iter_chunks()
-                .all(|chunk| chunk.is::<Constant>()),
+            elements.is::<Constant>(),
             "a uniform tile should have stayed constant-encoded",
         );
+        assert_eq!(elements.len(), 2_000);
         Ok(())
     }
 

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;

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<usize> {
+    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<Self> {
+        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(),
+        ))
+    }
+}

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;

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<Repeated> =
+    ParentRuleSet::new(&[ParentRuleSet::lift(&SliceReduceAdaptor(Repeated))]);
+
+impl SliceReduce for Repeated {
+    fn slice(array: ArrayView<'_, Self>, range: Range<usize>) -> VortexResult<Option<ArrayRef>> {
+        // 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)
+    }
+}