Vector cast (#5461)

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

Author: gatesn

vortex-array/src/expr/exprs/cast.rs vortex-array/src/expr/exprs/cast/mod.rsvortex-array/src/expr/exprs/cast.rs renamed to vortex-array/src/expr/exprs/cast/mod.rs

@@ -8,6 +8,7 @@ use prost::Message;
 use vortex_dtype::DType;
 use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
 use vortex_proto::expr as pb;
+use vortex_vector::Vector;
 
 use crate::ArrayRef;
 use crate::compute::cast as compute_cast;
@@ -87,6 +88,16 @@ impl VTable for Cast {
         })
     }
 
+    fn execute(
+        &self,
+        expr: &ExpressionView<Self>,
+        vector: &Vector,
+        dtype: &DType,
+    ) -> VortexResult<Vector> {
+        let input = expr.child(0).execute(vector, dtype)?;
+        vortex_compute::cast::Cast::cast(&input, dtype)
+    }
+
     fn stat_expression(
         &self,
         expr: &ExpressionView<Self>,

vortex-array/src/expr/exprs/select/mod.rs

@@ -11,6 +11,7 @@ use vortex_dtype::{DType, FieldNames};
 use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
 use vortex_proto::expr::select_opts::Opts;
 use vortex_proto::expr::{FieldNames as ProtoFieldNames, SelectOpts};
+use vortex_vector::Vector;
 
 use crate::expr::expression::Expression;
 use crate::expr::field::DisplayFieldNames;
@@ -144,6 +145,17 @@ impl VTable for Select {
         }?
         .into_array())
     }
+
+    fn execute(
+        &self,
+        _expr: &ExpressionView<Self>,
+        _vector: &Vector,
+        _dtype: &DType,
+    ) -> VortexResult<Vector> {
+        vortex_bail!(
+            "Select expressions cannot be executed. They must be removed during optimization."
+        )
+    }
 }
 
 /// Creates an expression that selects (includes) specific fields from an array.

vortex-buffer/src/trusted_len.rs

@@ -162,5 +162,6 @@ where
 }
 
 // Arrow bit iterators
+unsafe impl<'a> TrustedLen for crate::bit::BitIterator<'a> {}
 unsafe impl<'a> TrustedLen for crate::bit::BitChunkIterator<'a> {}
 unsafe impl<'a> TrustedLen for crate::bit::UnalignedBitChunkIterator<'a> {}

vortex-compute/src/cast/bool.rs

@@ -0,0 +1,46 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use num_traits::{One, Zero};
+use vortex_buffer::Buffer;
+use vortex_dtype::{DType, match_each_native_ptype};
+use vortex_error::{VortexResult, vortex_bail};
+use vortex_vector::bool::BoolVector;
+use vortex_vector::null::NullVector;
+use vortex_vector::primitive::PVector;
+use vortex_vector::{Vector, VectorOps};
+
+use crate::cast::Cast;
+
+impl Cast for BoolVector {
+    fn cast(&self, dtype: &DType) -> VortexResult<Vector> {
+        match dtype {
+            DType::Null if self.validity().all_false() => {
+                // Can cast an all-null BoolVector to NullVector.
+                Ok(NullVector::new(self.len()).into())
+            }
+            DType::Bool(n) if n.is_nullable() || self.validity().all_true() => {
+                // If the target dtype is nullable, or if the source BoolVector has no nulls,
+                // we can cast directly to BoolVector.
+                Ok(self.clone().into())
+            }
+            DType::Primitive(ptype, _) => {
+                match_each_native_ptype!(ptype, |T| {
+                    Ok(PVector::<T>::new(
+                        Buffer::<T>::from_trusted_len_iter(
+                            self.bits()
+                                .iter()
+                                .map(|b| if b { T::one() } else { T::zero() }),
+                        ),
+                        self.validity().clone(),
+                    )
+                    .into())
+                })
+            }
+            DType::Extension(ext_dtype) => self.cast(ext_dtype.storage_dtype()),
+            _ => {
+                vortex_bail!("Cannot cast BoolVector to type {}", dtype);
+            }
+        }
+    }
+}

vortex-compute/src/cast/mod.rs

@@ -0,0 +1,33 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Lossless casting of Vortex vectors for different logical data types.
+
+mod bool;
+mod null;
+mod pvector;
+
+use vortex_dtype::DType;
+use vortex_error::{VortexResult, vortex_bail};
+use vortex_vector::Vector;
+
+/// Trait for casting vectors to different data types.
+pub trait Cast {
+    /// Cast the vector to the specified data type.
+    fn cast(&self, dtype: &DType) -> VortexResult<Vector>;
+}
+
+impl Cast for Vector {
+    fn cast(&self, dtype: &DType) -> VortexResult<Vector> {
+        // Switch to macro once all vector types implement Cast
+        // match_each_vector!(self, |v| { Cast::cast(v, dtype) })
+        match self {
+            Vector::Null(v) => Cast::cast(v, dtype),
+            Vector::Bool(v) => Cast::cast(v, dtype),
+            Vector::Primitive(v) => Cast::cast(v, dtype),
+            _ => {
+                vortex_bail!("Casting not implemented for vector type {:?}", self);
+            }
+        }
+    }
+}

vortex-compute/src/cast/null.rs

@@ -0,0 +1,22 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_dtype::DType;
+use vortex_error::{VortexResult, vortex_bail};
+use vortex_vector::null::NullVector;
+use vortex_vector::{Vector, VectorMut, VectorMutOps, VectorOps};
+
+use crate::cast::Cast;
+
+impl Cast for NullVector {
+    fn cast(&self, dtype: &DType) -> VortexResult<Vector> {
+        if dtype.is_nullable() {
+            // We can create an all-null vector of _any_ type.
+            let mut vec = VectorMut::with_capacity(dtype, self.len());
+            vec.append_nulls(self.len());
+            Ok(vec.freeze())
+        } else {
+            vortex_bail!("Cannot cast NullVector to non-nullable type {}", dtype);
+        }
+    }
+}

vortex-compute/src/cast/pvector.rs

@@ -0,0 +1,43 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_dtype::{DType, NativePType};
+use vortex_error::{VortexResult, vortex_bail};
+use vortex_vector::null::NullVector;
+use vortex_vector::primitive::{PVector, PrimitiveVector};
+use vortex_vector::{Vector, VectorOps, match_each_pvector};
+
+use crate::cast::Cast;
+
+impl Cast for PrimitiveVector {
+    fn cast(&self, dtype: &DType) -> VortexResult<Vector> {
+        match_each_pvector!(self, |v| { Cast::cast(v, dtype) })
+    }
+}
+
+impl<T: NativePType> Cast for PVector<T> {
+    fn cast(&self, dtype: &DType) -> VortexResult<Vector> {
+        match dtype {
+            // Can cast an all-null PVector to NullVector.
+            DType::Null if self.validity().all_false() => Ok(NullVector::new(self.len()).into()),
+            // We're already the correct PType, and we have compatible nullability.
+            DType::Primitive(target_ptype, n)
+                if *target_ptype == T::PTYPE && (n.is_nullable() || self.validity().all_true()) =>
+            {
+                Ok(self.clone().into())
+            }
+            // We're not the correct PType, but we do have compatible nullability.
+            DType::Primitive(target_ptype, n) if n.is_nullable() || self.validity().all_true() => {
+                vortex_bail!(
+                    "Casting from PType {} to PType {} not yet implemented",
+                    T::PTYPE,
+                    target_ptype
+                );
+            }
+            DType::Extension(ext_dtype) => self.cast(ext_dtype.storage_dtype()),
+            _ => {
+                vortex_bail!("Cannot cast {:?} to dtype {}", self, dtype);
+            }
+        }
+    }
+}

vortex-compute/src/lib.rs

@@ -10,6 +10,7 @@
 pub mod arithmetic;
 #[cfg(feature = "arrow")]
 pub mod arrow;
+pub mod cast;
 pub mod comparison;
 pub mod expand;
 pub mod filter;