VarBinVector

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

Cargo.lock

@@ -24,3 +24,5 @@ vortex-buffer = { workspace = true }
 vortex-dtype = { workspace = true }
 vortex-error = { workspace = true }
 vortex-mask = { workspace = true }
+
+static_assertions = { workspace = true }

vortex-vector/Cargo.toml

@@ -15,11 +15,16 @@ mod bool;
 mod null;
 mod primitive;
 mod struct_;
+mod varbin;
 
 pub use bool::{BoolVector, BoolVectorMut};
 pub use null::{NullVector, NullVectorMut};
 pub use primitive::{PVector, PVectorMut, PrimitiveVector, PrimitiveVectorMut};
 pub use struct_::{StructVector, StructVectorMut};
+pub use varbin::{
+    BinaryVector, BinaryVectorMut, StringVector, StringVectorMut, VarBinType, VarBinVector,
+    VarBinVectorMut,
+};
 
 mod ops;
 mod vector;

vortex-vector/src/lib.rs

@@ -49,6 +49,14 @@ macro_rules! match_each_vector {
                 let $vec = v;
                 $body
             }
+            $crate::Vector::String(v) => {
+                let $vec = v;
+                $body
+            }
+            $crate::Vector::Binary(v) => {
+                let $vec = v;
+                $body
+            }
             $crate::Vector::Struct(v) => {
                 let $vec = v;
                 $body
@@ -104,6 +112,14 @@ macro_rules! match_each_vector_mut {
                 let $vec = v;
                 $body
             }
+            $crate::VectorMut::String(v) => {
+                let $vec = v;
+                $body
+            }
+            $crate::VectorMut::Binary(v) => {
+                let $vec = v;
+                $body
+            }
             $crate::VectorMut::Struct(v) => {
                 let $vec = v;
                 $body

vortex-vector/src/macros.rs

@@ -29,5 +29,8 @@ impl Sealed for PrimitiveVectorMut {}
 impl<T: NativePType> Sealed for PVector<T> {}
 impl<T: NativePType> Sealed for PVectorMut<T> {}
 
+impl<T: VarBinType> Sealed for VarBinVector<T> {}
+impl<T: VarBinType> Sealed for VarBinVectorMut<T> {}
+
 impl Sealed for StructVector {}
 impl Sealed for StructVectorMut {}

vortex-vector/src/private.rs

@@ -0,0 +1,157 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use crate::{Vector, VectorMut};
+use std::fmt::Debug;
+
+mod vector;
+pub use vector::VarBinVector;
+
+mod vector_mut;
+pub use vector_mut::VarBinVectorMut;
+use vortex_error::vortex_panic;
+
+mod view;
+
+/// Type alias for non-utf8 variable-length binary vectors.
+pub type BinaryVector = VarBinVector<BinaryType>;
+/// Type alias for mutable non-utf8 variable-length binary vectors.
+pub type BinaryVectorMut = VarBinVectorMut<BinaryType>;
+/// Type alias for UTF-8 variable-length string vectors.
+pub type StringVector = VarBinVector<StringType>;
+/// Type alias for mutable UTF-8 variable-length string vectors.
+pub type StringVectorMut = VarBinVectorMut<StringType>;
+
+impl<T: VarBinType> From<VarBinVector<T>> for Vector {
+    fn from(value: VarBinVector<T>) -> Self {
+        T::upcast(value)
+    }
+}
+
+impl<T: VarBinType> From<VarBinVectorMut<T>> for VectorMut {
+    fn from(value: VarBinVectorMut<T>) -> Self {
+        T::upcast(value)
+    }
+}
+
+/// Trait to mark supported binary view types.
+pub trait VarBinType: Debug + Sized + private::Sealed {
+    /// The slice type for this variable binary type.
+    type Slice: ?Sized + AsRef<[u8]>;
+
+    /// Downcast the provided object to a type-specific instance.
+    fn downcast<V: VarBinTypeDowncast>(visitor: V) -> V::Output<Self>;
+
+    /// Upcast a type-specific instance to a generic instance.
+    fn upcast<V: VarBinTypeUpcast>(input: V::Input<Self>) -> V;
+}
+
+/// [`BinaryType`] for UTF-8 strings.
+#[derive(Clone, Debug)]
+pub struct StringType;
+impl VarBinType for StringType {
+    type Slice = str;
+
+    fn downcast<V: VarBinTypeDowncast>(visitor: V) -> V::Output<Self> {
+        visitor.into_string()
+    }
+
+    fn upcast<V: VarBinTypeUpcast>(input: V::Input<Self>) -> V {
+        V::from_string(input)
+    }
+}
+
+/// [`BinaryType`] for raw binary data.
+#[derive(Clone, Debug)]
+pub struct BinaryType;
+impl VarBinType for BinaryType {
+    type Slice = [u8];
+
+    fn downcast<V: VarBinTypeDowncast>(visitor: V) -> V::Output<Self> {
+        visitor.into_binary()
+    }
+
+    fn upcast<V: VarBinTypeUpcast>(input: V::Input<Self>) -> V {
+        V::from_binary(input)
+    }
+}
+
+pub trait VarBinTypeDowncast {
+    type Output<T: VarBinType>;
+
+    fn into_binary(self) -> Self::Output<BinaryType>;
+    fn into_string(self) -> Self::Output<StringType>;
+}
+
+pub trait VarBinTypeUpcast {
+    type Input<T: VarBinType>;
+
+    fn from_binary(input: Self::Input<BinaryType>) -> Self;
+    fn from_string(input: Self::Input<StringType>) -> Self;
+}
+
+impl VarBinTypeDowncast for Vector {
+    type Output<T: VarBinType> = VarBinVector<T>;
+
+    fn into_binary(self) -> Self::Output<BinaryType> {
+        if let Vector::Binary(v) = self {
+            return v;
+        }
+        vortex_panic!("Expected BinaryVector, got {self:?}");
+    }
+
+    fn into_string(self) -> Self::Output<StringType> {
+        if let Vector::String(v) = self {
+            return v;
+        }
+        vortex_panic!("Expected StringVector, got {self:?}");
+    }
+}
+
+impl VarBinTypeUpcast for Vector {
+    type Input<T: VarBinType> = VarBinVector<T>;
+
+    fn from_binary(input: Self::Input<BinaryType>) -> Self {
+        Vector::Binary(input)
+    }
+
+    fn from_string(input: Self::Input<StringType>) -> Self {
+        Vector::String(input)
+    }
+}
+
+impl VarBinTypeDowncast for VectorMut {
+    type Output<T: VarBinType> = VarBinVectorMut<T>;
+
+    fn into_binary(self) -> Self::Output<BinaryType> {
+        if let VectorMut::Binary(v) = self {
+            return v;
+        }
+        vortex_panic!("Expected BinaryVector, got {self:?}");
+    }
+
+    fn into_string(self) -> Self::Output<StringType> {
+        if let VectorMut::String(v) = self {
+            return v;
+        }
+        vortex_panic!("Expected StringVector, got {self:?}");
+    }
+}
+
+impl VarBinTypeUpcast for VectorMut {
+    type Input<T: VarBinType> = VarBinVectorMut<T>;
+
+    fn from_binary(input: Self::Input<BinaryType>) -> Self {
+        VectorMut::Binary(input)
+    }
+
+    fn from_string(input: Self::Input<StringType>) -> Self {
+        VectorMut::String(input)
+    }
+}
+
+mod private {
+    pub trait Sealed {}
+    impl Sealed for super::StringType {}
+    impl Sealed for super::BinaryType {}
+}

vortex-vector/src/varbin/mod.rs

@@ -0,0 +1,63 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use crate::varbin::vector_mut::VarBinVectorMut;
+use crate::varbin::view::BinaryView;
+use crate::varbin::VarBinType;
+use crate::VectorOps;
+use std::sync::Arc;
+use vortex_buffer::{Buffer, ByteBuffer};
+use vortex_mask::Mask;
+
+/// A variable-length binary vector.
+#[derive(Debug, Clone)]
+pub struct VarBinVector<T: VarBinType> {
+    views: Buffer<BinaryView>,
+    validity: Mask,
+    buffers: Arc<Box<[ByteBuffer]>>,
+    _marker: std::marker::PhantomData<T>,
+}
+
+impl<T: VarBinType> VarBinVector<T> {
+    /// Creates a new [`VarBinVector`] from the provided components.
+    ///
+    /// # Safety
+    ///
+    /// This function is unsafe because it does not validate the consistency of the provided
+    /// components.
+    ///
+    /// The caller must ensure that:
+    /// - The length of the `validity` mask matches the length of the `views` buffer.
+    /// - The `views` buffer correctly references the data in the `buffers`.
+    pub unsafe fn new_unchecked(
+        views: Buffer<BinaryView>,
+        validity: Mask,
+        buffers: Arc<Box<[ByteBuffer]>>,
+    ) -> Self {
+        Self {
+            views,
+            validity,
+            buffers,
+            _marker: std::marker::PhantomData,
+        }
+    }
+}
+
+impl<T: VarBinType> VectorOps for VarBinVector<T> {
+    type Mutable = VarBinVectorMut<T>;
+
+    fn len(&self) -> usize {
+        self.views.len()
+    }
+
+    fn validity(&self) -> &Mask {
+        &self.validity
+    }
+
+    fn try_into_mut(self) -> Result<Self::Mutable, Self>
+    where
+        Self: Sized,
+    {
+        todo!()
+    }
+}

vortex-vector/src/varbin/vector.rs

@@ -0,0 +1,74 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use crate::varbin::vector::VarBinVector;
+use crate::varbin::view::BinaryView;
+use crate::varbin::VarBinType;
+use crate::VectorMutOps;
+use vortex_buffer::{BufferMut, ByteBuffer};
+use vortex_mask::MaskMut;
+
+/// Mutable variable-length binary vector.
+#[derive(Clone, Debug)]
+pub struct VarBinVectorMut<T: VarBinType> {
+    views: BufferMut<BinaryView>,
+    validity: MaskMut,
+
+    buffers: Vec<ByteBuffer>,
+    open_buffer: Option<ByteBuffer>,
+
+    _marker: std::marker::PhantomData<T>,
+}
+
+impl<T: VarBinType> VarBinVectorMut<T> {
+    pub(super) fn new(
+        views: BufferMut<BinaryView>,
+        validity: MaskMut,
+        buffers: Vec<ByteBuffer>,
+    ) -> Self {
+        Self {
+            views,
+            validity,
+            buffers,
+            open_buffer: None,
+            _marker: std::marker::PhantomData,
+        }
+    }
+}
+
+impl<T: VarBinType> VectorMutOps for VarBinVectorMut<T> {
+    type Immutable = VarBinVector<T>;
+
+    fn len(&self) -> usize {
+        self.views.len()
+    }
+
+    fn capacity(&self) -> usize {
+        self.views.capacity()
+    }
+
+    fn reserve(&mut self, additional: usize) {
+        self.views.reserve(additional);
+    }
+
+    fn extend_from_vector(&mut self, other: &Self::Immutable) {
+        todo!()
+    }
+
+    fn append_nulls(&mut self, n: usize) {
+        self.views.push_n(BinaryView::empty_view(), n);
+        self.validity.append_n(false, n);
+    }
+
+    fn freeze(self) -> Self::Immutable {
+        todo!()
+    }
+
+    fn split_off(&mut self, at: usize) -> Self {
+        todo!()
+    }
+
+    fn unsplit(&mut self, other: Self) {
+        todo!()
+    }
+}