diff --git a/vortex-vector/src/macros.rs b/vortex-vector/src/macros.rs index c86b81fa342..59f11c89b20 100644 --- a/vortex-vector/src/macros.rs +++ b/vortex-vector/src/macros.rs @@ -39,36 +39,16 @@ macro_rules! match_each_vector { ($self:expr, | $vec:ident | $body:block) => {{ match $self { - $crate::Vector::Null(v) => { - let $vec = v; - $body - } - $crate::Vector::Bool(v) => { - let $vec = v; - $body - } - $crate::Vector::Primitive(v) => { - 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 - } + $crate::Vector::Null($vec) => $body, + $crate::Vector::Bool($vec) => $body, + $crate::Vector::Primitive($vec) => $body, + $crate::Vector::String($vec) => $body, + $crate::Vector::Binary($vec) => $body, + $crate::Vector::Struct($vec) => $body, } }}; } -pub(crate) use match_each_vector; - /// Matches on all variants of [`VectorMut`] and executes the same code for each variant branch. /// /// This macro eliminates repetitive match statements when implementing operations that need to work @@ -104,32 +84,96 @@ pub(crate) use match_each_vector; macro_rules! match_each_vector_mut { ($self:expr, | $vec:ident | $body:block) => {{ match $self { - $crate::VectorMut::Null(v) => { - let $vec = v; - $body - } - $crate::VectorMut::Bool(v) => { - let $vec = v; - $body - } - $crate::VectorMut::Primitive(v) => { - 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 - } + $crate::VectorMut::Null($vec) => $body, + $crate::VectorMut::Bool($vec) => $body, + $crate::VectorMut::Primitive($vec) => $body, + $crate::VectorMut::String($vec) => $body, + $crate::VectorMut::Binary($vec) => $body, + $crate::VectorMut::Struct($vec) => $body, + } + }}; +} + +/// Internal macro to generate match arms for vector pairs. +#[doc(hidden)] +#[macro_export] +macro_rules! __match_vector_pair_arms { + ( + $left:expr, + $right:expr, + $enum_left:ident, + $enum_right:ident, + $a:ident, + $b:ident, + $body:expr + ) => {{ + match ($left, $right) { + ($crate::$enum_left::Null($a), $crate::$enum_right::Null($b)) => $body, + ($crate::$enum_left::Bool($a), $crate::$enum_right::Bool($b)) => $body, + ($crate::$enum_left::Primitive($a), $crate::$enum_right::Primitive($b)) => $body, + ($crate::$enum_left::String($a), $crate::$enum_right::String($b)) => $body, + ($crate::$enum_left::Binary($a), $crate::$enum_right::Binary($b)) => $body, + ($crate::$enum_left::Struct($a), $crate::$enum_right::Struct($b)) => $body, + _ => ::vortex_error::vortex_panic!("Mismatched vector types"), } }}; } -pub(crate) use match_each_vector_mut; +/// Matches on pairs of vector variants and executes the same code for matching variant pairs. +/// +/// This macro eliminates repetitive match statements when implementing operations that need to work +/// with pairs of vectors where the variants must match. +/// +/// Specify the types of the left and right vectors (either `Vector` or `VectorMut`) and the macro +/// generates the appropriate match arms. +/// +/// The macro binds the matched inner values to identifiers in the closure that can be used in the +/// body expression. +/// +/// # Examples +/// +/// ``` +/// use vortex_vector::{ +/// BoolVector, BoolVectorMut, Vector, VectorMut, VectorMutOps, match_vector_pair +/// }; +/// +/// fn extend_vector(left: &mut VectorMut, right: &Vector) { +/// match_vector_pair!(left, right, |a: VectorMut, b: Vector| { +/// a.extend_from_vector(b); +/// }) +/// } +/// +/// let mut mut_vec: VectorMut = BoolVectorMut::from_iter([true, false, true]).into(); +/// let vec: Vector = BoolVectorMut::from_iter([false, true]).freeze().into(); +/// +/// extend_vector(&mut mut_vec, &vec); +/// assert_eq!(mut_vec.len(), 5); +/// ``` +/// +/// Note that the vectors can also be owned: +/// +/// ``` +/// use vortex_vector::{ +/// BoolVector, BoolVectorMut, Vector, VectorMut, VectorMutOps, match_vector_pair +/// }; +/// +/// fn extend_vector_owned(mut dest: VectorMut, src: Vector) -> VectorMut { +/// match_vector_pair!(&mut dest, src, |a: VectorMut, b: Vector| { +/// a.extend_from_vector(&b); +/// dest +/// }) +/// } +/// +/// let mut_vec: VectorMut = BoolVectorMut::from_iter([true, false, true]).into(); +/// let vec: Vector = BoolVectorMut::from_iter([false, true]).freeze().into(); +/// +/// let new_bool_mut = extend_vector_owned(mut_vec, vec); +/// assert_eq!(new_bool_mut.len(), 5); +/// ``` +#[macro_export] // DO NOT ADD `#[rustfmt::skip]`!!! https://github.com/rust-lang/rust/pull/52234#issuecomment-903419099 +macro_rules! match_vector_pair { + ($left:expr, $right:expr, | $a:ident : Vector, $b:ident : Vector | $body:expr) => {{ $crate::__match_vector_pair_arms!($left, $right, Vector, Vector, $a, $b, $body) }}; + ($left:expr, $right:expr, | $a:ident : Vector, $b:ident : VectorMut | $body:expr) => {{ $crate::__match_vector_pair_arms!($left, $right, Vector, VectorMut, $a, $b, $body) }}; + ($left:expr, $right:expr, | $a:ident : VectorMut, $b:ident : Vector | $body:expr) => {{ $crate::__match_vector_pair_arms!($left, $right, VectorMut, Vector, $a, $b, $body) }}; + ($left:expr, $right:expr, | $a:ident : VectorMut, $b:ident : VectorMut | $body:expr) => {{ $crate::__match_vector_pair_arms!($left, $right, VectorMut, VectorMut, $a, $b, $body) }}; +} diff --git a/vortex-vector/src/primitive/macros.rs b/vortex-vector/src/primitive/macros.rs index 04d8ae6b840..19ce4819a11 100644 --- a/vortex-vector/src/primitive/macros.rs +++ b/vortex-vector/src/primitive/macros.rs @@ -44,56 +44,21 @@ macro_rules! match_each_pvector { ($self:expr, | $vec:ident | $body:block) => {{ match $self { - $crate::PrimitiveVector::U8(v) => { - let $vec = v; - $body - } - $crate::PrimitiveVector::U16(v) => { - let $vec = v; - $body - } - $crate::PrimitiveVector::U32(v) => { - let $vec = v; - $body - } - $crate::PrimitiveVector::U64(v) => { - let $vec = v; - $body - } - $crate::PrimitiveVector::I8(v) => { - let $vec = v; - $body - } - $crate::PrimitiveVector::I16(v) => { - let $vec = v; - $body - } - $crate::PrimitiveVector::I32(v) => { - let $vec = v; - $body - } - $crate::PrimitiveVector::I64(v) => { - let $vec = v; - $body - } - $crate::PrimitiveVector::F16(v) => { - let $vec = v; - $body - } - $crate::PrimitiveVector::F32(v) => { - let $vec = v; - $body - } - $crate::PrimitiveVector::F64(v) => { - let $vec = v; - $body - } + $crate::PrimitiveVector::U8($vec) => $body, + $crate::PrimitiveVector::U16($vec) => $body, + $crate::PrimitiveVector::U32($vec) => $body, + $crate::PrimitiveVector::U64($vec) => $body, + $crate::PrimitiveVector::I8($vec) => $body, + $crate::PrimitiveVector::I16($vec) => $body, + $crate::PrimitiveVector::I32($vec) => $body, + $crate::PrimitiveVector::I64($vec) => $body, + $crate::PrimitiveVector::F16($vec) => $body, + $crate::PrimitiveVector::F32($vec) => $body, + $crate::PrimitiveVector::F64($vec) => $body, } }}; } -pub(crate) use match_each_pvector; - /// Matches on all primitive type variants of [`PrimitiveVectorMut`] and executes the same code /// for each variant branch. /// @@ -129,52 +94,17 @@ pub(crate) use match_each_pvector; macro_rules! match_each_pvector_mut { ($self:expr, | $vec:ident | $body:block) => {{ match $self { - $crate::PrimitiveVectorMut::U8(v) => { - let $vec = v; - $body - } - $crate::PrimitiveVectorMut::U16(v) => { - let $vec = v; - $body - } - $crate::PrimitiveVectorMut::U32(v) => { - let $vec = v; - $body - } - $crate::PrimitiveVectorMut::U64(v) => { - let $vec = v; - $body - } - $crate::PrimitiveVectorMut::I8(v) => { - let $vec = v; - $body - } - $crate::PrimitiveVectorMut::I16(v) => { - let $vec = v; - $body - } - $crate::PrimitiveVectorMut::I32(v) => { - let $vec = v; - $body - } - $crate::PrimitiveVectorMut::I64(v) => { - let $vec = v; - $body - } - $crate::PrimitiveVectorMut::F16(v) => { - let $vec = v; - $body - } - $crate::PrimitiveVectorMut::F32(v) => { - let $vec = v; - $body - } - $crate::PrimitiveVectorMut::F64(v) => { - let $vec = v; - $body - } + $crate::PrimitiveVectorMut::U8($vec) => $body, + $crate::PrimitiveVectorMut::U16($vec) => $body, + $crate::PrimitiveVectorMut::U32($vec) => $body, + $crate::PrimitiveVectorMut::U64($vec) => $body, + $crate::PrimitiveVectorMut::I8($vec) => $body, + $crate::PrimitiveVectorMut::I16($vec) => $body, + $crate::PrimitiveVectorMut::I32($vec) => $body, + $crate::PrimitiveVectorMut::I64($vec) => $body, + $crate::PrimitiveVectorMut::F16($vec) => $body, + $crate::PrimitiveVectorMut::F32($vec) => $body, + $crate::PrimitiveVectorMut::F64($vec) => $body, } }}; } - -pub(crate) use match_each_pvector_mut; diff --git a/vortex-vector/src/primitive/vector.rs b/vortex-vector/src/primitive/vector.rs index 81f0e2b898b..835adaa8754 100644 --- a/vortex-vector/src/primitive/vector.rs +++ b/vortex-vector/src/primitive/vector.rs @@ -7,8 +7,7 @@ use vortex_dtype::half::f16; use vortex_dtype::{NativePType, PType, PTypeDowncast, PTypeUpcast}; use vortex_error::vortex_panic; -use super::macros::match_each_pvector; -use crate::{PVector, PrimitiveVectorMut, VectorOps}; +use crate::{PVector, PrimitiveVectorMut, VectorOps, match_each_pvector}; /// An immutable vector of primitive values. /// diff --git a/vortex-vector/src/primitive/vector_mut.rs b/vortex-vector/src/primitive/vector_mut.rs index ba6f70a6649..a77bba9c954 100644 --- a/vortex-vector/src/primitive/vector_mut.rs +++ b/vortex-vector/src/primitive/vector_mut.rs @@ -7,8 +7,7 @@ use vortex_dtype::half::f16; use vortex_dtype::{NativePType, PType, PTypeDowncast, PTypeUpcast}; use vortex_error::vortex_panic; -use super::macros::match_each_pvector_mut; -use crate::{PVectorMut, PrimitiveVector, VectorMutOps}; +use crate::{PVectorMut, PrimitiveVector, VectorMutOps, match_each_pvector_mut}; /// A mutable vector of primitive values. /// diff --git a/vortex-vector/src/struct_/vector.rs b/vortex-vector/src/struct_/vector.rs index 6900da67e2b..b69b3daef1b 100644 --- a/vortex-vector/src/struct_/vector.rs +++ b/vortex-vector/src/struct_/vector.rs @@ -33,8 +33,7 @@ pub struct StructVector { /// The length of the vector (which is the same as all field vectors). /// - /// This is stored here as a convenience, and also helps in the case that the `StructVector` has - /// no fields. + /// This is stored here as a convenience, as the validity also tracks this information. pub(super) len: usize, } @@ -111,14 +110,14 @@ impl StructVector { } } - /// Decomposes the struct vector into its constituent parts (fields, validity, and length). + /// Decomposes the struct vector into its constituent parts (fields and validity). pub fn into_parts(self) -> (Arc>, Mask) { (self.fields, self.validity) } /// Returns the fields of the `StructVector`, each stored column-wise as a [`Vector`]. - pub fn fields(&self) -> &[Vector] { - self.fields.as_ref() + pub fn fields(&self) -> &Arc> { + &self.fields } } diff --git a/vortex-vector/src/struct_/vector_mut.rs b/vortex-vector/src/struct_/vector_mut.rs index 5ac6dceba43..f4c64b2c9a2 100644 --- a/vortex-vector/src/struct_/vector_mut.rs +++ b/vortex-vector/src/struct_/vector_mut.rs @@ -5,10 +5,10 @@ use std::sync::Arc; -use vortex_error::{VortexExpect, VortexResult, vortex_ensure, vortex_panic}; +use vortex_error::{VortexExpect, VortexResult, vortex_ensure}; use vortex_mask::MaskMut; -use crate::{StructVector, Vector, VectorMut, VectorMutOps, VectorOps}; +use crate::{StructVector, Vector, VectorMut, VectorMutOps, VectorOps, match_vector_pair}; /// A mutable vector of struct values (values with named fields). /// @@ -237,17 +237,11 @@ impl VectorMutOps for StructVectorMut { ); // Extend each field vector. - let pairs = self.fields.iter_mut().zip(other.fields()); + let pairs = self.fields.iter_mut().zip(other.fields().as_ref()); for (self_mut_vector, other_vec) in pairs { - match (self_mut_vector, other_vec) { - (VectorMut::Null(a), Vector::Null(b)) => a.extend_from_vector(b), - (VectorMut::Bool(a), Vector::Bool(b)) => a.extend_from_vector(b), - (VectorMut::Primitive(a), Vector::Primitive(b)) => a.extend_from_vector(b), - (VectorMut::Struct(a), Vector::Struct(b)) => a.extend_from_vector(b), - _ => { - vortex_panic!("Mismatched field types in `StructVectorMut::extend_from_vector`") - } - } + match_vector_pair!(self_mut_vector, other_vec, |a: VectorMut, b: Vector| { + a.extend_from_vector(b) + }) } // Extend the validity mask. @@ -320,15 +314,11 @@ impl VectorMutOps for StructVectorMut { // Unsplit each field vector. let pairs = self.fields.iter_mut().zip(other.fields); for (self_mut_vector, other_mut_vec) in pairs { - match (self_mut_vector, other_mut_vec) { - (VectorMut::Null(a), VectorMut::Null(b)) => a.unsplit(b), - (VectorMut::Bool(a), VectorMut::Bool(b)) => a.unsplit(b), - (VectorMut::Primitive(a), VectorMut::Primitive(b)) => a.unsplit(b), - (VectorMut::Struct(a), VectorMut::Struct(b)) => a.unsplit(b), - _ => { - vortex_panic!("Mismatched field types in `StructVectorMut::unsplit`") - } - } + match_vector_pair!( + self_mut_vector, + other_mut_vec, + |a: VectorMut, b: VectorMut| a.unsplit(b) + ) } self.validity.unsplit(other.validity); diff --git a/vortex-vector/src/vector.rs b/vortex-vector/src/vector.rs index 649de47d4f9..5d6700a209d 100644 --- a/vortex-vector/src/vector.rs +++ b/vortex-vector/src/vector.rs @@ -8,9 +8,10 @@ use vortex_error::vortex_panic; -use crate::macros::match_each_vector; use crate::varbin::{BinaryVector, StringVector}; -use crate::{BoolVector, NullVector, PrimitiveVector, StructVector, VectorMut, VectorOps}; +use crate::{ + BoolVector, NullVector, PrimitiveVector, StructVector, VectorMut, VectorOps, match_each_vector, +}; /// An enum over all kinds of immutable vectors, which represent fully decompressed (canonical) /// array data. diff --git a/vortex-vector/src/vector_mut.rs b/vortex-vector/src/vector_mut.rs index 9774c07adf2..22ec47ba583 100644 --- a/vortex-vector/src/vector_mut.rs +++ b/vortex-vector/src/vector_mut.rs @@ -10,10 +10,10 @@ use vortex_dtype::DType; use vortex_error::vortex_panic; use vortex_mask::MaskMut; -use super::macros::match_each_vector_mut; use crate::varbin::{BinaryVectorMut, StringVectorMut}; use crate::{ BoolVectorMut, NullVectorMut, PrimitiveVectorMut, StructVectorMut, Vector, VectorMutOps, + match_each_vector_mut, match_vector_pair, }; /// An enum over all kinds of mutable vectors, which represent fully decompressed (canonical) array @@ -94,13 +94,9 @@ impl VectorMutOps for VectorMut { } fn extend_from_vector(&mut self, other: &Self::Immutable) { - match (self, other) { - (VectorMut::Null(a), Vector::Null(b)) => a.extend_from_vector(b), - (VectorMut::Bool(a), Vector::Bool(b)) => a.extend_from_vector(b), - (VectorMut::Primitive(a), Vector::Primitive(b)) => a.extend_from_vector(b), - (VectorMut::Struct(a), Vector::Struct(b)) => a.extend_from_vector(b), - _ => vortex_panic!("Mismatched vector types"), - } + match_vector_pair!(self, other, |a: VectorMut, b: Vector| { + a.extend_from_vector(b) + }) } fn append_nulls(&mut self, n: usize) { @@ -116,15 +112,7 @@ impl VectorMutOps for VectorMut { } fn unsplit(&mut self, other: Self) { - match (self, other) { - (VectorMut::Null(a), VectorMut::Null(b)) => a.unsplit(b), - (VectorMut::Bool(a), VectorMut::Bool(b)) => a.unsplit(b), - (VectorMut::Primitive(a), VectorMut::Primitive(b)) => a.unsplit(b), - (VectorMut::String(a), VectorMut::String(b)) => a.unsplit(b), - (VectorMut::Binary(a), VectorMut::Binary(b)) => a.unsplit(b), - (VectorMut::Struct(a), VectorMut::Struct(b)) => a.unsplit(b), - _ => vortex_panic!("Mismatched vector types"), - } + match_vector_pair!(self, other, |a: VectorMut, b: VectorMut| a.unsplit(b)) } }