feat[vortex-array]: add an expr array

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

Author: joseph-isaacs

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

@@ -17,6 +17,7 @@ pub struct ExprArray {
 }
 
 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())?);
         Ok(unsafe { Self::unchecked_new(child, expr, dtype) })
@@ -37,7 +38,8 @@ impl ExprArray {
         }
     }
 
-    pub fn new_with_root_dtype(child: ArrayRef, expr: Expression) -> VortexResult<Self> {
+    /// Creates a new ExprArray with the dtype inferred from the expression and child.
+    pub fn new_infer_dtype(child: ArrayRef, expr: Expression) -> VortexResult<Self> {
         let dtype = expr.return_dtype(child.dtype())?;
         Ok(unsafe { Self::unchecked_new(child, expr, dtype) })
     }

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

@@ -27,13 +27,16 @@ impl ArrayVTable<ExprVTable> for ExprVTable {
     fn array_hash<H: std::hash::Hasher>(array: &ExprArray, state: &mut H, precision: Precision) {
         array.child.array_hash(state, precision);
         array.dtype.hash(state);
+        // TODO(joe): fixme
+        array.expr.serialize_metadata().hash()
         // Note: Expression doesn't implement Hash, so we skip it
         // This is acceptable since expressions are typically transient
     }
 
     fn array_eq(array: &ExprArray, other: &ExprArray, precision: Precision) -> bool {
-        array.child.array_eq(&other.child, precision) && array.dtype == other.dtype
-        // Note: We don't compare expressions here as they don't implement Eq
-        // This is acceptable since ExprArray is typically transient
+        array.child.array_eq(&other.child, precision)
+            && array.dtype == other.dtype
+            // TODO(joe): fixme
+            && array.expr.serialize_metadata() == other.expr.serialize_metadata()
     }
 }

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

@@ -9,7 +9,6 @@ use crate::vtable::CanonicalVTable;
 
 impl CanonicalVTable<ExprVTable> for ExprVTable {
     fn canonicalize(array: &ExprArray) -> Canonical {
-        // Evaluate the expression on the child array and canonicalize the result
         array
             .expr
             .evaluate(&array.child)
@@ -45,19 +44,9 @@ mod tests {
         let array = expr_array.into_array();
 
         // Test canonicalize - should evaluate the expression
-        let canonical = array.to_canonical();
+        let canonical = array.to_canonical().into_array();
 
-        // The result should be a primitive array with value 15 repeated for each element
-        let canonical_array = canonical.as_ref();
-        assert_eq!(canonical_array.len(), 3);
-
-        println!("a {}", array.display_tree());
-
-        // Extract the primitive array from the Canonical enum
-        if let crate::Canonical::Primitive(primitive) = canonical {
-            assert_eq!(primitive.buffer::<i32>().as_slice(), &[15i32, 15, 15]);
-        } else {
-            panic!("Expected Canonical::Primitive variant");
-        }
+        let expect = (0..3).map(|_| 15i32).collect::<PrimitiveArray>();
+        assert_eq!(expect, canonical);
     }
 }

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

@@ -12,17 +12,6 @@ use crate::expr::transform::ExprOptimizer;
 use crate::vtable::OperatorVTable;
 
 impl OperatorVTable<ExprVTable> for ExprVTable {
-    fn bind(
-        array: &ExprArray,
-        selection: Option<&ArrayRef>,
-        ctx: &mut dyn BindCtx,
-    ) -> VortexResult<BatchKernelRef> {
-        // Canonicalize the ExprArray and bind the result
-        // This evaluates the expression and returns the result
-        let canonical = array.to_canonical();
-        ctx.bind(&ArrayRef::from(canonical), selection)
-    }
-
     fn reduce(array: &ExprArray) -> VortexResult<Option<ArrayRef>> {
         // Get the default expression session
         let session = ExprSession::default();
@@ -66,7 +55,7 @@ mod tests {
 
         let expr = get_item("a", pack([("a", root())], Nullability::NonNullable));
 
-        let expr_array = ExprArray::new_with_root_dtype(array.clone().into_array(), expr.clone())?;
+        let expr_array = ExprArray::new_infer_dtype(array.clone().into_array(), expr.clone())?;
 
         // Call reduce - it should optimize pack(a: $).a to just $
         let reduced = expr_array.reduce()?.vortex_expect("reduce failed");

vortex-array/src/arrays/struct_/vtable/operator.rs

@@ -238,7 +238,7 @@ mod tests {
 
         let expr = gt(get_item("a", root()), lit(5));
         let expr_array =
-            ExprArray::new_with_root_dtype(struct_array.clone().into_array(), expr.clone())?;
+            ExprArray::new_infer_dtype(struct_array.clone().into_array(), expr.clone())?;
 
         let actual = expr_array.to_canonical().into_array();
         let expected = (0..5)
@@ -270,7 +270,7 @@ mod tests {
             gt(get_item("a", root()), lit(5)),
             lt(get_item("a", root()), lit(10)),
         );
-        let expr_array = ExprArray::new_with_root_dtype(struct_array.clone().into_array(), expr)?;
+        let expr_array = ExprArray::new_infer_dtype(struct_array.clone().into_array(), expr)?;
 
         let actual = expr_array.to_canonical().into_array();
         let expected = (0..5)
@@ -304,7 +304,7 @@ mod tests {
             and(gt(col("a"), lit(5)), lt(col("b"), lit(4))),
             gt(col("a"), lit(6)),
         );
-        let expr_array = ExprArray::new_with_root_dtype(struct_array.clone().into_array(), expr)?;
+        let expr_array = ExprArray::new_infer_dtype(struct_array.clone().into_array(), expr)?;
         let parent: ArrayRef = expr_array.into_array();
         let result = struct_array
             .reduce_parent(&parent, 0)?
@@ -383,7 +383,7 @@ mod tests {
 
         let expr = eq(lit(1), lit(0));
         let expr_array =
-            ExprArray::new_with_root_dtype(struct_array.clone().into_array(), expr.clone())?;
+            ExprArray::new_infer_dtype(struct_array.clone().into_array(), expr.clone())?;
 
         let actual = expr_array.to_canonical().into_array();
         let expected = (0..5).map(|_| false).collect::<BoolArray>().into_array();

vortex-array/src/arrays/struct_/vtable/reduce.rs

@@ -152,7 +152,7 @@ pub(crate) fn apply_partitioned_expr(
             )?
             .into_array();
 
-            Ok(ExprArray::new_with_root_dtype(child, partitioned.root.clone())?.into_array())
+            Ok(ExprArray::new_infer_dtype(child, partitioned.root.clone())?.into_array())
         }
     }
 }