Vector Scalars

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

Author: gatesn

vortex-buffer/src/bit/buf.rs

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use std::ops::{BitAnd, BitOr, BitXor, Not, Range};
+use std::ops::{BitAnd, BitOr, BitXor, Not, RangeBounds};
 
 use crate::bit::ops::{bitwise_binary_op, bitwise_unary_op};
 use crate::bit::{
@@ -194,16 +194,24 @@ impl BitBuffer {
     /// for `len` bits.
     ///
     /// Panics if the slice would extend beyond the end of the buffer.
-    pub fn slice(&self, range: Range<usize>) -> Self {
-        assert!(
-            range.len() <= self.len,
-            "slice from {} to {} exceeds len {}",
-            range.start,
-            range.end,
-            range.len()
-        );
-
-        Self::new_with_offset(self.buffer.clone(), range.len(), self.offset + range.start)
+    pub fn slice(&self, range: impl RangeBounds<usize>) -> Self {
+        let start = match range.start_bound() {
+            std::ops::Bound::Included(&s) => s,
+            std::ops::Bound::Excluded(&s) => s + 1,
+            std::ops::Bound::Unbounded => 0,
+        };
+        let end = match range.end_bound() {
+            std::ops::Bound::Included(&e) => e + 1,
+            std::ops::Bound::Excluded(&e) => e,
+            std::ops::Bound::Unbounded => self.len,
+        };
+
+        assert!(start <= end);
+        assert!(start <= self.len);
+        assert!(end <= self.len);
+        let len = end - start;
+
+        Self::new_with_offset(self.buffer.clone(), len, self.offset + start)
     }
 
     /// Slice any full bytes from the buffer, leaving the offset < 8.

vortex-buffer/src/buffer.rs

@@ -246,7 +246,7 @@ impl<T> Buffer<T> {
     /// Requires that `begin <= end` and `end <= self.len()`.
     /// Also requires that both `begin` and `end` are aligned to the buffer's required alignment.
     #[inline(always)]
-    pub fn slice(&self, range: impl RangeBounds<usize>) -> Self {
+    pub fn slice(&self, range: impl RangeBounds<usize> + Clone + Debug) -> Self {
         self.slice_with_alignment(range, self.alignment)
     }
 
@@ -257,7 +257,7 @@ impl<T> Buffer<T> {
     ///
     /// Requires that `begin <= end` and `end <= self.len()`.
     #[inline(always)]
-    pub fn slice_unaligned(&self, range: impl RangeBounds<usize>) -> Self {
+    pub fn slice_unaligned(&self, range: impl RangeBounds<usize> + Clone + Debug) -> Self {
         self.slice_with_alignment(range, Alignment::of::<u8>())
     }
 
@@ -270,7 +270,7 @@ impl<T> Buffer<T> {
     /// Also requires that both `begin` and `end` are aligned to the given alignment.
     pub fn slice_with_alignment(
         &self,
-        range: impl RangeBounds<usize>,
+        range: impl RangeBounds<usize> + Clone,
         alignment: Alignment,
     ) -> Self {
         let len = self.len();

vortex-buffer/src/string.rs

@@ -10,6 +10,7 @@ use crate::ByteBuffer;
 
 /// A wrapper around a [`ByteBuffer`] that guarantees that the buffer contains valid UTF-8.
 #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[repr(transparent)]
 pub struct BufferString(ByteBuffer);
 
 impl BufferString {

vortex-mask/src/lib.rs

@@ -15,7 +15,7 @@ mod tests;
 
 use std::cmp::Ordering;
 use std::fmt::{Debug, Formatter};
-use std::ops::Range;
+use std::ops::RangeBounds;
 use std::sync::{Arc, OnceLock};
 
 use itertools::Itertools;
@@ -393,11 +393,26 @@ impl Mask {
 
     /// Slice the mask.
     #[inline]
-    pub fn slice(&self, range: Range<usize>) -> Self {
-        assert!(range.end <= self.len());
+    pub fn slice(&self, range: impl RangeBounds<usize>) -> Self {
+        let start = match range.start_bound() {
+            std::ops::Bound::Included(&s) => s,
+            std::ops::Bound::Excluded(&s) => s + 1,
+            std::ops::Bound::Unbounded => 0,
+        };
+        let end = match range.end_bound() {
+            std::ops::Bound::Included(&e) => e + 1,
+            std::ops::Bound::Excluded(&e) => e,
+            std::ops::Bound::Unbounded => self.len(),
+        };
+
+        assert!(start <= end);
+        assert!(start <= self.len());
+        assert!(end <= self.len());
+        let len = end - start;
+
         match &self {
-            Self::AllTrue(_) => Self::new_true(range.len()),
-            Self::AllFalse(_) => Self::new_false(range.len()),
+            Self::AllTrue(_) => Self::new_true(len),
+            Self::AllFalse(_) => Self::new_false(len),
             Self::Values(values) => Self::from_buffer(values.buffer.slice(range)),
         }
     }

vortex-vector/src/binaryview/mod.rs

@@ -8,13 +8,15 @@
 //! [specification](https://arrow.apache.org/docs/format/Columnar.html#variable-size-binary-view-layout),
 //! which are inspired by "German" strings.
 
+pub use scalar::*;
 pub use types::*;
 pub use vector::*;
 pub use vector_mut::*;
 use vortex_error::vortex_panic;
 
 use crate::{Vector, VectorMut};
 
+mod scalar;
 mod types;
 mod vector;
 mod vector_mut;
@@ -28,6 +30,10 @@ pub type BinaryVectorMut = BinaryViewVectorMut<BinaryType>;
 pub type StringVector = BinaryViewVector<StringType>;
 /// Type alias for mutable UTF-8 variable-length string vectors.
 pub type StringVectorMut = BinaryViewVectorMut<StringType>;
+/// Type alias for non-utf8 variable-length binary scalars.
+pub type BinaryScalar = BinaryViewScalar<BinaryType>;
+/// Type alias for UTF-8 variable-length string scalars.
+pub type StringScalar = BinaryViewScalar<StringType>;
 
 impl BinaryViewDowncast for Vector {
     type Output<T: BinaryViewType> = BinaryViewVector<T>;

vortex-vector/src/binaryview/scalar.rs

@@ -0,0 +1,56 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use crate::binaryview::{
+    BinaryType, BinaryViewType, BinaryViewTypeUpcast, BinaryViewVectorMut, StringType,
+};
+use crate::{Scalar, ScalarOps, VectorMutOps};
+
+/// A scalar value for types that implement [`BinaryViewType`].
+pub struct BinaryViewScalar<T: BinaryViewType>(Option<T::Scalar>);
+
+impl<T: BinaryViewType> From<Option<T::Scalar>> for BinaryViewScalar<T> {
+    fn from(value: Option<T::Scalar>) -> Self {
+        Self(value)
+    }
+}
+
+impl<T: BinaryViewType> BinaryViewScalar<T> {
+    /// Returns the scalar value as [`T::Scalar`], or `None` if the scalar is null.
+    pub fn value(&self) -> Option<&T::Scalar> {
+        self.0.as_ref()
+    }
+}
+
+impl<T: BinaryViewType> ScalarOps for BinaryViewScalar<T> {
+    fn is_valid(&self) -> bool {
+        self.0.is_some()
+    }
+
+    fn repeat(&self, n: usize) -> crate::VectorMut {
+        let mut vec = BinaryViewVectorMut::<T>::with_capacity(n);
+        match self.value() {
+            None => vec.append_nulls(n),
+            Some(buf) => vec.append_owned_values(buf.clone(), n),
+        }
+        vec.into()
+    }
+}
+
+impl BinaryViewTypeUpcast for Scalar {
+    type Input<T: BinaryViewType> = BinaryViewScalar<T>;
+
+    fn from_binary(input: Self::Input<BinaryType>) -> Self {
+        Scalar::Binary(input)
+    }
+
+    fn from_string(input: Self::Input<StringType>) -> Self {
+        Scalar::String(input)
+    }
+}
+
+impl<T: BinaryViewType> Into<Scalar> for BinaryViewScalar<T> {
+    fn into(self) -> Scalar {
+        T::upcast(self)
+    }
+}

vortex-vector/src/binaryview/types.rs

@@ -5,6 +5,8 @@
 
 use std::fmt::Debug;
 
+use vortex_buffer::{BufferString, ByteBuffer};
+
 use crate::binaryview::{BinaryViewVector, BinaryViewVectorMut};
 use crate::{Vector, VectorMut};
 
@@ -24,6 +26,8 @@ impl<T: BinaryViewType> From<BinaryViewVectorMut<T>> for VectorMut {
 pub trait BinaryViewType: Debug + Sized + private::Sealed {
     /// The slice type for this variable binary type.
     type Slice: ?Sized + AsRef<[u8]>;
+    /// The scalar type for this variable binary type.
+    type Scalar: Sized + Clone + AsRef<Self::Slice> + Into<ByteBuffer>;
 
     /// Validate if a set of bytes conforms to the logical type constraints of the native `Slice`.
     fn validate(bytes: &[u8]) -> bool;
@@ -40,6 +44,14 @@ pub trait BinaryViewType: Debug + Sized + private::Sealed {
     /// vector operations.
     unsafe fn from_bytes_unchecked(bytes: &[u8]) -> &Self::Slice;
 
+    /// Returns the scalar value from a byte buffer.
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure that the buffer conforms to the type requirements of this binary
+    /// type.
+    unsafe fn scalar_from_buffer_unchecked(buffer: ByteBuffer) -> Self::Scalar;
+
     /// Downcast the provided object to a type-specific instance.
     fn downcast<V: BinaryViewDowncast>(visitor: V) -> V::Output<Self>;
 
@@ -52,18 +64,25 @@ pub trait BinaryViewType: Debug + Sized + private::Sealed {
 pub struct StringType;
 impl BinaryViewType for StringType {
     type Slice = str;
+    type Scalar = BufferString;
 
     #[inline(always)]
     fn validate(bytes: &[u8]) -> bool {
         std::str::from_utf8(bytes).is_ok()
     }
 
+    #[inline(always)]
     unsafe fn from_bytes_unchecked(bytes: &[u8]) -> &Self::Slice {
         // SAFETY: vectors should be checked at the boundary for upholding the UTF8 variant,
         //  or only be built from vectors that are known to satisfy the variant.
         unsafe { std::str::from_utf8_unchecked(bytes) }
     }
 
+    #[inline(always)]
+    unsafe fn scalar_from_buffer_unchecked(buffer: ByteBuffer) -> Self::Scalar {
+        unsafe { BufferString::new_unchecked(buffer) }
+    }
+
     fn downcast<V: BinaryViewDowncast>(visitor: V) -> V::Output<Self> {
         visitor.into_string()
     }
@@ -78,16 +97,23 @@ impl BinaryViewType for StringType {
 pub struct BinaryType;
 impl BinaryViewType for BinaryType {
     type Slice = [u8];
+    type Scalar = ByteBuffer;
 
     #[inline(always)]
     fn validate(_bytes: &[u8]) -> bool {
         true
     }
 
+    #[inline(always)]
     unsafe fn from_bytes_unchecked(bytes: &[u8]) -> &Self::Slice {
         bytes
     }
 
+    #[inline(always)]
+    unsafe fn scalar_from_buffer_unchecked(buffer: ByteBuffer) -> Self::Scalar {
+        buffer
+    }
+
     fn downcast<V: BinaryViewDowncast>(visitor: V) -> V::Output<Self> {
         visitor.into_binary()
     }

vortex-vector/src/binaryview/vector.rs

@@ -3,16 +3,18 @@
 
 //! Variable-length binary vector implementation.
 
+use std::fmt::Debug;
+use std::ops::RangeBounds;
 use std::sync::Arc;
 
-use vortex_buffer::{Buffer, ByteBuffer};
+use vortex_buffer::{Alignment, Buffer, ByteBuffer};
 use vortex_error::{VortexExpect, VortexResult, vortex_ensure};
 use vortex_mask::Mask;
 
-use crate::VectorOps;
-use crate::binaryview::BinaryViewType;
 use crate::binaryview::vector_mut::BinaryViewVectorMut;
 use crate::binaryview::view::{BinaryView, validate_views};
+use crate::binaryview::{BinaryViewScalar, BinaryViewType};
+use crate::{Scalar, VectorOps};
 
 /// A variable-length binary vector.
 ///
@@ -108,10 +110,50 @@ impl<T: BinaryViewType> BinaryViewVector<T> {
         (self.views, self.buffers, self.validity)
     }
 
+    /// Get the `index` item from the vector as an owned `Scalar` type with zero-copy.
+    ///
+    /// This function will panic is `index` is out of range for the vector's length.
+    pub fn get(&self, index: usize) -> Option<T::Scalar> {
+        if !self.validity.value(index) {
+            return None;
+        }
+
+        let view = &self.views[index];
+        if view.is_inlined() {
+            let view = view.as_inlined();
+
+            // We find the occurrence of the inlined data in the views buffer.
+            let buffer = self
+                .views
+                .clone()
+                .into_byte_buffer()
+                .aligned(Alignment::none())
+                .slice_ref(&view.data[..view.size as usize]);
+
+            // SAFETY: validation that the string data contained in this vector is performed
+            //  at construction time, either in the constructor for safe construction, or by
+            //  the caller (when using the unchecked constructor).
+            Some(unsafe { T::scalar_from_buffer_unchecked(buffer) })
+        } else {
+            // Get a pointer into the buffer range
+            let view_ref = view.as_view();
+            let buffer = &self.buffers[view_ref.buffer_index as usize];
+
+            let start = view_ref.offset as usize;
+            let length = view_ref.size as usize;
+            let buffer_slice = buffer.slice(start..start + length);
+
+            // SAFETY: validation that the string data contained in this vector is performed
+            //  at construction time, either in the constructor for safe construction, or by
+            //  the caller (when using the unchecked constructor).
+            Some(unsafe { T::scalar_from_buffer_unchecked(buffer_slice) })
+        }
+    }
+
     /// Get the `index` item from the vector as a native `Slice` type.
     ///
     /// This function will panic is `index` is out of range for the vector's length.
-    pub fn get(&self, index: usize) -> Option<&T::Slice> {
+    pub fn get_ref(&self, index: usize) -> Option<&T::Slice> {
         if !self.validity.value(index) {
             return None;
         }
@@ -160,6 +202,15 @@ impl<T: BinaryViewType> VectorOps for BinaryViewVector<T> {
         &self.validity
     }
 
+    fn scalar_at(&self, index: usize) -> Scalar {
+        assert!(index < self.len());
+        BinaryViewScalar::<T>::from(self.get(index)).into()
+    }
+
+    fn slice(&self, _range: impl RangeBounds<usize> + Clone + Debug) -> Self {
+        todo!()
+    }
+
     fn try_into_mut(self) -> Result<BinaryViewVectorMut<T>, Self> {
         let views_mut = match self.views.try_into_mut() {
             Ok(views_mut) => views_mut,