add scaffolding for ListVector

Signed-off-by: Connor Tsui <[email protected]>

Author: connortsui20

vortex-compute/src/arrow/list.rs

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use arrow_array::ArrayRef;
+use vortex_error::VortexResult;
+use vortex_vector::listview::ListViewVector;
+
+use crate::arrow::IntoArrow;
+
+impl IntoArrow<ArrayRef> for ListViewVector {
+    fn into_arrow(self) -> VortexResult<ArrayRef> {
+        todo!("Figure out how to do this")
+    }
+}

vortex-compute/src/arrow/mod.rs

@@ -9,6 +9,7 @@ mod binaryview;
 mod bool;
 mod decimal;
 mod fixed_size_list;
+mod list;
 mod mask;
 mod null;
 mod primitive;

vortex-compute/src/mask/mod.rs

@@ -11,6 +11,7 @@ use vortex_vector::binaryview::{BinaryViewType, BinaryViewVector};
 use vortex_vector::bool::BoolVector;
 use vortex_vector::decimal::{DVector, DecimalVector};
 use vortex_vector::fixed_size_list::FixedSizeListVector;
+use vortex_vector::listview::ListViewVector;
 use vortex_vector::null::NullVector;
 use vortex_vector::primitive::{PVector, PrimitiveVector};
 use vortex_vector::struct_::StructVector;
@@ -46,31 +47,31 @@ impl MaskValidity for BoolVector {
     }
 }
 
-impl MaskValidity for DecimalVector {
+impl MaskValidity for PrimitiveVector {
     fn mask_validity(self, mask: &Mask) -> Self {
-        match_each_dvector!(self, |v| { MaskValidity::mask_validity(v, mask).into() })
+        match_each_pvector!(self, |v| { MaskValidity::mask_validity(v, mask).into() })
     }
 }
 
-impl<D: NativeDecimalType> MaskValidity for DVector<D> {
+impl<T: NativePType> MaskValidity for PVector<T> {
     fn mask_validity(self, mask: &Mask) -> Self {
-        let (ps, elements, validity) = self.into_parts();
-        // SAFETY: we are preserving the original elements buffer and only modifying the validity.
-        unsafe { Self::new_unchecked(ps, elements, validity.bitand(mask)) }
+        let (data, validity) = self.into_parts();
+        // SAFETY: we are preserving the original data buffer and only modifying the validity.
+        unsafe { Self::new_unchecked(data, validity.bitand(mask)) }
     }
 }
 
-impl MaskValidity for PrimitiveVector {
+impl MaskValidity for DecimalVector {
     fn mask_validity(self, mask: &Mask) -> Self {
-        match_each_pvector!(self, |v| { MaskValidity::mask_validity(v, mask).into() })
+        match_each_dvector!(self, |v| { MaskValidity::mask_validity(v, mask).into() })
     }
 }
 
-impl<T: NativePType> MaskValidity for PVector<T> {
+impl<D: NativeDecimalType> MaskValidity for DVector<D> {
     fn mask_validity(self, mask: &Mask) -> Self {
-        let (data, validity) = self.into_parts();
-        // SAFETY: we are preserving the original data buffer and only modifying the validity.
-        unsafe { Self::new_unchecked(data, validity.bitand(mask)) }
+        let (ps, elements, validity) = self.into_parts();
+        // SAFETY: we are preserving the original elements buffer and only modifying the validity.
+        unsafe { Self::new_unchecked(ps, elements, validity.bitand(mask)) }
     }
 }
 
@@ -82,6 +83,15 @@ impl<T: BinaryViewType> MaskValidity for BinaryViewVector<T> {
     }
 }
 
+impl MaskValidity for ListViewVector {
+    fn mask_validity(self, mask: &Mask) -> Self {
+        let (elements, offsets, sizes, validity) = self.into_parts();
+        // SAFETY: we are preserving the original elements and `list_size`, only modifying the
+        // validity.
+        unsafe { Self::new_unchecked(elements, offsets, sizes, validity.bitand(mask)) }
+    }
+}
+
 impl MaskValidity for FixedSizeListVector {
     fn mask_validity(self, mask: &Mask) -> Self {
         let (elements, list_size, validity) = self.into_parts();

vortex-vector/src/fixed_size_list/vector.rs

@@ -149,22 +149,24 @@ impl VectorOps for FixedSizeListVector {
         &self.validity
     }
 
-    fn try_into_mut(self) -> Result<Self::Mutable, Self>
+    fn try_into_mut(self) -> Result<FixedSizeListVectorMut, Self>
     where
         Self: Sized,
     {
         let len = self.len;
         let list_size = self.list_size;
 
+        // Try to unwrap the `Arc`.
         let elements = match Arc::try_unwrap(self.elements) {
             Ok(elements) => elements,
-            Err(elements) => return Err(FixedSizeListVector { elements, ..self }),
+            Err(elements) => return Err(Self { elements, ..self }),
         };
 
+        // Try to make validity mutable.
         let validity = match self.validity.try_into_mut() {
             Ok(validity) => validity,
             Err(validity) => {
-                return Err(FixedSizeListVector {
+                return Err(Self {
                     elements: Arc::new(elements),
                     list_size,
                     validity,
@@ -173,14 +175,15 @@ impl VectorOps for FixedSizeListVector {
             }
         };
 
+        // Try to make the elements mutable.
         match elements.try_into_mut() {
             Ok(mutable_elements) => Ok(FixedSizeListVectorMut {
                 elements: Box::new(mutable_elements),
                 list_size,
                 validity,
                 len,
             }),
-            Err(elements) => Err(FixedSizeListVector {
+            Err(elements) => Err(Self {
                 elements: Arc::new(elements),
                 list_size,
                 validity: validity.freeze(),

vortex-vector/src/lib.rs

@@ -14,6 +14,7 @@ pub mod binaryview;
 pub mod bool;
 pub mod decimal;
 pub mod fixed_size_list;
+pub mod listview;
 pub mod null;
 pub mod primitive;
 pub mod struct_;

vortex-vector/src/listview/mod.rs

@@ -0,0 +1,31 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Definition and implementation of [`ListViewVector`] and [`ListViewVectorMut`].
+//!
+//! A [`ListViewVector`] represents a collection of variable-width lists, where each list can
+//! contain a different number of elements.
+//!
+//! The structure uses separate offset and size vectors to track the boundaries of each list
+//! within the flat elements array. This allows for efficient access to individual lists without
+//! copying data.  This is similar to Apache Arrow's `ListView` type.
+
+mod vector;
+pub use vector::ListViewVector;
+
+mod vector_mut;
+pub use vector_mut::ListViewVectorMut;
+
+use crate::{Vector, VectorMut};
+
+impl From<ListViewVector> for Vector {
+    fn from(v: ListViewVector) -> Self {
+        Self::List(v)
+    }
+}
+
+impl From<ListViewVectorMut> for VectorMut {
+    fn from(v: ListViewVectorMut) -> Self {
+        Self::List(v)
+    }
+}

vortex-vector/src/listview/vector.rs

@@ -0,0 +1,212 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Definition and implementation of [`ListViewVector`].
+
+use std::sync::Arc;
+
+use vortex_error::{VortexExpect, VortexResult};
+use vortex_mask::Mask;
+
+use super::ListViewVectorMut;
+use crate::Vector;
+use crate::ops::{VectorMutOps, VectorOps};
+use crate::primitive::PrimitiveVector;
+
+/// A vector of variable-width lists.
+///
+/// Each list is defined by 2 integers: an offset and a size (a "list view"), which point into a
+/// child `elements` vector.
+///
+/// Note that the list views **do not** need to be sorted, nor do they have to be contiguous or
+/// fully cover the `elements` vector. This means that multiple views can be pointing to the same
+/// elements.
+///
+/// # Structure
+///
+/// - `elements`: The child vector of all list elements, stored as an [`Arc<Vector>`].
+/// - `offsets`: A [`PrimitiveVector`] containing the starting offset of each list in the `elements`
+///   vector.
+/// - `sizes`: A [`PrimitiveVector`] containing the size (number of elements) of each list.
+/// - `validity`: A [`Mask`] indicating which lists are null.
+#[derive(Debug, Clone)]
+pub struct ListViewVector {
+    /// The child vector of elements.
+    pub(super) elements: Arc<Vector>,
+
+    /// Offsets for each list into the elements vector.
+    pub(super) offsets: PrimitiveVector,
+
+    /// Sizes (lengths) of each list.
+    pub(super) sizes: PrimitiveVector,
+
+    /// The validity mask (where `true` represents a list is **not** null).
+    ///
+    /// Note that the `elements` vector will have its own internal validity, denoting if individual
+    /// list elements are null.
+    pub(super) validity: Mask,
+
+    /// The length of the vector (which is the same as the length of the validity mask).
+    ///
+    /// This is stored here as a convenience, as the validity also tracks this information.
+    pub(super) len: usize,
+}
+
+impl ListViewVector {
+    /// Creates a new [`ListViewVector`] from its components.
+    ///
+    /// # Panics
+    ///
+    /// TODO
+    pub fn new(
+        elements: Arc<Vector>,
+        offsets: PrimitiveVector,
+        sizes: PrimitiveVector,
+        validity: Mask,
+    ) -> Self {
+        Self::try_new(elements, offsets, sizes, validity)
+            .vortex_expect("Invalid ListViewVector construction")
+    }
+
+    /// Attempts to create a new [`ListViewVector`] from its components.
+    ///
+    /// # Errors
+    ///
+    /// TODO
+    pub fn try_new(
+        _elements: Arc<Vector>,
+        _offsets: PrimitiveVector,
+        _sizes: PrimitiveVector,
+        _validity: Mask,
+    ) -> VortexResult<Self> {
+        todo!()
+    }
+
+    /// Creates a new [`ListViewVector`] without validation.
+    ///
+    /// # Safety
+    ///
+    /// TODO
+    pub unsafe fn new_unchecked(
+        elements: Arc<Vector>,
+        offsets: PrimitiveVector,
+        sizes: PrimitiveVector,
+        validity: Mask,
+    ) -> Self {
+        let len = validity.len();
+
+        if cfg!(debug_assertions) {
+            Self::new(elements, offsets, sizes, validity)
+        } else {
+            Self {
+                elements,
+                offsets,
+                sizes,
+                validity,
+                len,
+            }
+        }
+    }
+
+    /// Decomposes the [`ListViewVector`] into its constituent parts (child elements, offsets,
+    /// sizes, and validity).
+    pub fn into_parts(self) -> (Arc<Vector>, PrimitiveVector, PrimitiveVector, Mask) {
+        (self.elements, self.offsets, self.sizes, self.validity)
+    }
+
+    /// Returns a reference to the `elements` vector.
+    #[inline]
+    pub fn elements(&self) -> &Vector {
+        &self.elements
+    }
+
+    /// Returns a reference to the integer `offsets` vector.
+    #[inline]
+    pub fn offsets(&self) -> &PrimitiveVector {
+        &self.offsets
+    }
+
+    /// Returns a reference to the integer `sizes` vector.
+    #[inline]
+    pub fn sizes(&self) -> &PrimitiveVector {
+        &self.sizes
+    }
+}
+
+impl VectorOps for ListViewVector {
+    type Mutable = ListViewVectorMut;
+
+    fn len(&self) -> usize {
+        self.len
+    }
+
+    fn validity(&self) -> &Mask {
+        &self.validity
+    }
+
+    fn try_into_mut(self) -> Result<ListViewVectorMut, Self> {
+        // Try to unwrap the `Arc`.
+        let elements = match Arc::try_unwrap(self.elements) {
+            Ok(elements) => elements,
+            Err(elements) => return Err(Self { elements, ..self }),
+        };
+
+        // Try to make the validity mutable.
+        let validity = match self.validity.try_into_mut() {
+            Ok(v) => v,
+            Err(validity) => {
+                return Err(Self {
+                    elements: Arc::new(elements),
+                    validity,
+                    ..self
+                });
+            }
+        };
+
+        // Try to make the offsets mutable.
+        let offsets = match self.offsets.try_into_mut() {
+            Ok(mutable_offsets) => mutable_offsets,
+            Err(offsets) => {
+                return Err(Self {
+                    offsets,
+                    sizes: self.sizes,
+                    elements: Arc::new(elements),
+                    validity: validity.freeze(),
+                    len: self.len,
+                });
+            }
+        };
+
+        // Try to make the sizes mutable.
+        let sizes = match self.sizes.try_into_mut() {
+            Ok(mutable_sizes) => mutable_sizes,
+            Err(sizes) => {
+                return Err(Self {
+                    offsets: offsets.freeze(),
+                    sizes,
+                    elements: Arc::new(elements),
+                    validity: validity.freeze(),
+                    len: self.len,
+                });
+            }
+        };
+
+        // Try to make the elements mutable.
+        match elements.try_into_mut() {
+            Ok(mut_elements) => Ok(ListViewVectorMut {
+                offsets,
+                sizes,
+                elements: Box::new(mut_elements),
+                validity,
+                len: self.len,
+            }),
+            Err(elements) => Err(Self {
+                offsets: offsets.freeze(),
+                sizes: sizes.freeze(),
+                elements: Arc::new(elements),
+                validity: validity.freeze(),
+                len: self.len,
+            }),
+        }
+    }
+}