u

Signed-off-by: Joe Isaacs <[email protected]>

Author: joseph-isaacs

vortex-array/src/arrays/expr/array.rs

@@ -2,24 +2,58 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use vortex_dtype::DType;
-use vortex_error::VortexResult;
+use vortex_error::{VortexResult, vortex_ensure};
 
 use crate::expr::Expression;
 use crate::stats::ArrayStats;
 use crate::{Array, ArrayRef};
 
+/// A array that represents an expression to be evaluated over a child array.
+///
+/// `ExprArray` enables deferred evaluation of expressions by wrapping a child array
+/// with an expression that operates on it. The expression is not evaluated until the
+/// array is canonicalized/executed.
+///
+/// # Examples
+///
+/// ```ignore
+/// // Create an expression that filters an integer array
+/// let data = PrimitiveArray::from_iter([1, 2, 3, 4, 5]);
+/// let expr = gt(root(), lit(3)); // $ > 3
+/// let expr_array = ExprArray::new_infer_dtype(data.into_array(), expr)?;
+///
+/// // The expression is evaluated when canonicalized
+/// let result = expr_array.to_canonical(); // Returns BoolArray([false, false, false, true, true])
+/// ```
+///
+/// # Type Safety
+///
+/// The `dtype` field must match `expr.return_dtype(child.dtype())`. This invariant
+/// is enforced by the safe constructors ([`try_new`](ExprArray::try_new) and
+/// [`new_infer_dtype`](ExprArray::new_infer_dtype)) but can be bypassed
+/// with [`unchecked_new`](ExprArray::unchecked_new) for performance-critical code.
 #[derive(Clone, Debug)]
 pub struct ExprArray {
+    /// The underlying array that the expression will operate on.
     pub(super) child: ArrayRef,
+    /// The expression to evaluate over the child array.
     pub(super) expr: Expression,
+    /// The data type of the result after evaluating the expression.
     pub(super) dtype: DType,
+    /// Statistics about the resulting array (may be computed lazily).
     pub(super) stats: ArrayStats,
 }
 
 impl ExprArray {
     /// Creates a new ExprArray with the dtype validated to match the expression's return type.
     pub fn try_new(child: ArrayRef, expr: Expression, dtype: DType) -> VortexResult<Self> {
-        assert_eq!(dtype, expr.return_dtype(child.dtype())?);
+        let expected_dtype = expr.return_dtype(child.dtype())?;
+        vortex_ensure!(
+            dtype == expected_dtype,
+            "ExprArray dtype mismatch: expected {}, got {}",
+            expected_dtype,
+            dtype
+        );
         Ok(unsafe { Self::unchecked_new(child, expr, dtype) })
     }
 
@@ -34,6 +68,7 @@ impl ExprArray {
             child,
             expr,
             dtype,
+            // TODO(joe): Propagate or compute statistics from the child array and expression.
             stats: ArrayStats::default(),
         }
     }

vortex-array/src/arrays/expr/vtable/mod.rs

@@ -27,7 +27,7 @@ pub struct ExprEncoding;
 impl VTable for ExprVTable {
     type Array = ExprArray;
     type Encoding = ExprEncoding;
-    type Metadata = ExprDisplay;
+    type Metadata = ExprArrayMetadata;
 
     type ArrayVTable = Self;
     type CanonicalVTable = Self;
@@ -47,7 +47,7 @@ impl VTable for ExprVTable {
     }
 
     fn metadata(array: &ExprArray) -> VortexResult<Self::Metadata> {
-        Ok(ExprDisplay(array.expr.clone()))
+        Ok(ExprArrayMetadata((array.expr.clone(), array.dtype.clone())))
     }
 
     fn serialize(_metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
@@ -62,26 +62,27 @@ impl VTable for ExprVTable {
         _encoding: &ExprEncoding,
         dtype: &DType,
         len: usize,
-        metadata: &Self::Metadata,
+        ExprArrayMetadata((expr, root_dtype)): &Self::Metadata,
         buffers: &[ByteBuffer],
         children: &dyn ArrayChildren,
     ) -> VortexResult<ExprArray> {
         if !buffers.is_empty() {
             vortex_bail!("Expected 0 buffers, got {}", buffers.len());
         }
 
-        let Ok(child) = children.get(0, dtype, len) else {
+        let Ok(child) = children.get(0, root_dtype, len) else {
             vortex_bail!("Expected 1 child, got {}", children.len());
         };
 
-        ExprArray::try_new(child, metadata.0.clone(), dtype.clone())
+        ExprArray::try_new(child, expr.clone(), dtype.clone())
     }
 }
 
-pub struct ExprDisplay(Expression);
+pub struct ExprArrayMetadata((Expression, DType));
 
-impl Debug for ExprDisplay {
+impl Debug for ExprArrayMetadata {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        self.0.fmt_sql(f)
+        // Since this is used in display method we can omit the dtype.
+        self.0.0.fmt_sql(f)
     }
 }