Vector Casting

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

@@ -4,16 +4,16 @@
 use std::fmt::Formatter;
 use std::ops::Deref;
 
-use prost::Message;
-use vortex_dtype::DType;
-use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
-use vortex_proto::expr as pb;
-
-use crate::ArrayRef;
 use crate::compute::cast as compute_cast;
 use crate::expr::expression::Expression;
 use crate::expr::{ChildName, ExprId, ExpressionView, StatsCatalog, VTable, VTableExt};
 use crate::stats::Stat;
+use crate::ArrayRef;
+use prost::Message;
+use vortex_dtype::DType;
+use vortex_error::{vortex_bail, vortex_err, VortexExpect, VortexResult};
+use vortex_proto::expr as pb;
+use vortex_vector::Vector;
 
 /// A cast expression that converts values to a target data type.
 pub struct Cast;
@@ -87,6 +87,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>,
@@ -120,6 +130,15 @@ impl VTable for Cast {
     }
 }
 
+/// A trait that defines the compute function for performing a cast.
+///
+/// The `try_cast` method attempts to convert the vector to the specified data type,
+/// returning a `VortexResult` that contains the casted vector or an error if the
+/// cast operation fails.
+trait VectorCast {
+    fn try_cast(self, dtype: &DType) -> VortexResult<Vector>;
+}
+
 /// Creates an expression that casts values to a target data type.
 ///
 /// Converts the input expression's values to the specified target type.
@@ -141,11 +160,11 @@ mod tests {
     use vortex_error::VortexUnwrap as _;
 
     use super::cast;
-    use crate::IntoArray;
     use crate::arrays::StructArray;
     use crate::expr::exprs::get_item::get_item;
     use crate::expr::exprs::root::root;
-    use crate::expr::{Expression, test_harness};
+    use crate::expr::{test_harness, Expression};
+    use crate::IntoArray;
 
     #[test]
     fn dtype() {

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

@@ -5,17 +5,17 @@ pub mod transform;
 
 use std::fmt::{Display, Formatter};
 
+use crate::expr::expression::Expression;
+use crate::expr::field::DisplayFieldNames;
+use crate::expr::{ChildName, ExprId, ExpressionView, VTable, VTableExt};
+use crate::{ArrayRef, IntoArray, ToCanonical};
 use itertools::Itertools;
 use prost::Message;
 use vortex_dtype::{DType, FieldNames};
-use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
+use vortex_error::{vortex_bail, vortex_err, VortexExpect, VortexResult};
 use vortex_proto::expr::select_opts::Opts;
 use vortex_proto::expr::{FieldNames as ProtoFieldNames, SelectOpts};
-
-use crate::expr::expression::Expression;
-use crate::expr::field::DisplayFieldNames;
-use crate::expr::{ChildName, ExprId, ExpressionView, VTable, VTableExt};
-use crate::{ArrayRef, IntoArray, ToCanonical};
+use vortex_vector::Vector;
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub enum FieldSelection {
@@ -144,6 +144,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,45 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use crate::cast::Cast;
+use num_traits::{One, Zero};
+use vortex_buffer::Buffer;
+use vortex_dtype::{match_each_native_ptype, DType};
+use vortex_error::{vortex_bail, VortexResult};
+use vortex_vector::bool::BoolVector;
+use vortex_vector::null::NullVector;
+use vortex_vector::primitive::PVector;
+use vortex_vector::{Vector, VectorOps};
+
+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::{vortex_bail, VortexResult};
+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,21 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use crate::cast::Cast;
+use vortex_dtype::DType;
+use vortex_error::{vortex_bail, VortexResult};
+use vortex_vector::null::NullVector;
+use vortex_vector::{Vector, VectorMut, VectorMutOps, VectorOps};
+
+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().into())
+        } else {
+            vortex_bail!("Cannot cast NullVector to non-nullable type {}", dtype);
+        }
+    }
+}

vortex-compute/src/cast/pvector.rs

@@ -0,0 +1,42 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use crate::cast::Cast;
+use vortex_dtype::{DType, NativePType};
+use vortex_error::{vortex_bail, VortexResult};
+use vortex_vector::null::NullVector;
+use vortex_vector::primitive::{PVector, PrimitiveVector};
+use vortex_vector::{match_each_pvector, Vector, VectorOps};
+
+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;

vortex-vector/src/primitive/generic.rs

@@ -8,11 +8,15 @@ use std::ops::RangeBounds;
 
 use vortex_buffer::Buffer;
 use vortex_dtype::NativePType;
-use vortex_error::{VortexExpect, VortexResult, vortex_ensure};
+use vortex_error::{vortex_ensure, VortexExpect, VortexResult};
 use vortex_mask::Mask;
 
 use crate::primitive::{PScalar, PVectorMut};
+<<<<<<< Updated upstream
 use crate::{Scalar, VectorOps};
+=======
+use crate::VectorOps;
+>>>>>>> Stashed changes
 
 /// An immutable vector of generic primitive values.
 ///