Expressions

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

Author: gatesn

Cargo.lock

@@ -39,7 +39,7 @@ pub mod session;
 mod signature;
 mod simplify;
 pub mod stats;
-// pub mod transform;
+pub mod transform;
 pub mod traversal;
 mod vtable;
 

vortex-array/src/expr/mod.rs

@@ -93,7 +93,7 @@ mod tests {
             eq(lit(1), root()),
         );
 
-        let s_expr = (&expr).serialize_proto().unwrap();
+        let s_expr = expr.serialize_proto().unwrap();
         let buf = s_expr.encode_to_vec();
         let s_expr = pb::Expr::decode(buf.as_slice()).unwrap();
         let deser_expr = deserialize_expr_proto(&s_expr, &registry).unwrap();

vortex-array/src/expr/proto.rs

@@ -12,6 +12,7 @@ use vortex_utils::aliases::hash_map::HashMap;
 use crate::expr::Expression;
 use crate::expr::Root;
 use crate::expr::SimplifyCtx;
+use crate::expr::transform::match_between::find_between;
 
 impl Expression {
     /// Simplify the expression, returning a potentially new expression.
@@ -55,7 +56,15 @@ impl Expression {
             dtype_cache: RefCell::new(HashMap::new()),
         };
 
-        inner(self, &cache)?.map_or_else(|| Ok(self.clone()), Ok)
+        let simplified = inner(self, &cache)?.unwrap_or_else(|| self.clone());
+
+        // TODO(ngates): remove the "between" optimization, or rewrite it to not always convert
+        //  to CNF?
+        let simplified = find_between(simplified);
+
+        // TODO(ngates): perform constant folding by executing expressions with all-literal
+        //  children here
+        Ok(simplified)
     }
 }
 

vortex-array/src/expr/simplify.rs

@@ -49,65 +49,69 @@ pub fn find_between(expr: Expression) -> Expression {
 }
 
 fn maybe_match(lhs: &Expression, rhs: &Expression) -> Option<Expression> {
-    let (Some(lhs_e), Some(rhs_e)) = (lhs.as_opt::<Binary>(), rhs.as_opt::<Binary>()) else {
+    let (Some(lhs_op), Some(rhs_op)) = (lhs.as_opt::<Binary>(), rhs.as_opt::<Binary>()) else {
         return None;
     };
 
+    // Extract the grandchildren
+    let lhs_lhs = lhs.child(0);
+    let lhs_rhs = lhs.child(1);
+    let rhs_lhs = rhs.child(0);
+    let rhs_rhs = rhs.child(1);
+
     // Cannot compare to self
-    if lhs_e.lhs().eq(lhs_e.rhs()) || rhs_e.lhs().eq(rhs_e.rhs()) {
+    if lhs_lhs.eq(lhs_rhs) || rhs_lhs.eq(rhs_rhs) {
         return None;
     }
 
     // First, get both halves to have GetItem on the left
-    let lhs = match (lhs_e.lhs().is::<GetItem>(), lhs_e.rhs().is::<GetItem>()) {
+    let lhs = match (lhs_lhs.is::<GetItem>(), lhs_rhs.is::<GetItem>()) {
         (true, false) => lhs.clone(),
-        (false, true) => Binary.new_expr(
-            lhs_e.operator().swap()?,
-            [lhs_e.rhs().clone(), lhs_e.lhs().clone()],
-        ),
+        (false, true) => Binary.new_expr(lhs_op.swap()?, [lhs_rhs.clone(), lhs_lhs.clone()]),
         _ => return None,
     };
-    let lhs_e = lhs.as_::<Binary>();
+    let lhs_op = lhs.as_::<Binary>();
+    let lhs_lhs = lhs.child(0);
 
-    let rhs = match (rhs_e.lhs().is::<GetItem>(), rhs_e.rhs().is::<GetItem>()) {
+    let rhs = match (rhs_lhs.is::<GetItem>(), rhs_rhs.is::<GetItem>()) {
         (true, false) => rhs.clone(),
-        (false, true) => Binary.new_expr(
-            rhs_e.operator().swap()?,
-            [rhs_e.rhs().clone(), rhs_e.lhs().clone()],
-        ),
+        (false, true) => Binary.new_expr(rhs_op.swap()?, [rhs_rhs.clone(), rhs_lhs.clone()]),
         _ => return None,
     };
-    let rhs_e = rhs.as_::<Binary>();
+    let rhs_op = rhs.as_::<Binary>();
+    let rhs_lhs = rhs.child(0);
 
     // Both conjuncts must reference the same GetItem column
-    if !lhs_e.lhs().eq(rhs_e.lhs()) {
+    if !lhs_lhs.eq(rhs_lhs) {
         return None;
     }
 
-    let target = lhs_e.lhs().clone();
+    let target = lhs_lhs.clone();
 
     // Find the lower bound
-    let (lower, upper) = match (lhs_e.operator(), rhs_e.operator()) {
+    let (lower, upper) = match (lhs_op, rhs_op) {
         (Operator::Lt | Operator::Lte, Operator::Gt | Operator::Gte) => (rhs, lhs),
         (Operator::Gt | Operator::Gte, Operator::Lt | Operator::Lte) => (lhs, rhs),
         _ => return None,
     };
-    let lower_e = lower.as_::<Binary>();
-    let upper_e = upper.as_::<Binary>();
+    let lower_op = lower.as_::<Binary>();
+    let lower_rhs = lower.child(1);
+    let upper_op = upper.as_::<Binary>();
+    let upper_rhs = upper.child(1);
 
     // Ensure bounds are literals
-    let _ = lower_e.rhs().as_opt::<Literal>()?;
-    let _ = upper_e.rhs().as_opt::<Literal>()?;
+    let _ = lower_rhs.as_opt::<Literal>()?;
+    let _ = upper_rhs.as_opt::<Literal>()?;
 
-    let lower_strict = is_strict_comparison(lower_e.operator())?;
-    let upper_strict = is_strict_comparison(upper_e.operator())?;
+    let lower_strict = is_strict_comparison(*lower_op)?;
+    let upper_strict = is_strict_comparison(*upper_op)?;
 
     Some(Between.new_expr(
         BetweenOptions {
             lower_strict,
             upper_strict,
         },
-        [target, lower_e.rhs().clone(), upper_e.rhs().clone()],
+        [target, lower_rhs.clone(), upper_rhs.clone()],
     ))
 }
 

vortex-array/src/expr/transform/match_between.rs

@@ -3,13 +3,8 @@
 
 //! A collection of transformations that can be applied to a [`crate::expr::Expression`].
 pub(crate) mod match_between;
-mod optimizer;
 mod partition;
 mod replace;
-pub mod rules;
-mod simplify;
-mod simplify_typed;
 
-pub use optimizer::*;
 pub use partition::*;
 pub use replace::*;

vortex-array/src/expr/transform/mod.rs

@@ -22,7 +22,6 @@ use crate::expr::analysis::descendent_annotations;
 use crate::expr::exprs::get_item::get_item;
 use crate::expr::exprs::pack::pack;
 use crate::expr::exprs::root::root;
-use crate::expr::transform::ExprOptimizer;
 use crate::expr::traversal::NodeExt;
 use crate::expr::traversal::NodeRewriter;
 use crate::expr::traversal::Transformed;
@@ -43,7 +42,6 @@ pub fn partition<A: AnnotationFn>(
     expr: Expression,
     scope: &DType,
     annotate_fn: A,
-    optimizer: &ExprOptimizer,
 ) -> VortexResult<PartitionedExpr<A::Annotation>>
 where
     A::Annotation: Display,
@@ -73,7 +71,7 @@ where
             Nullability::NonNullable,
         );
 
-        let expr = optimizer.optimize_typed(expr.clone(), scope)?;
+        let expr = expr.simplify(scope)?;
         let expr_dtype = expr.return_dtype(scope)?;
 
         partitions.push(expr);
@@ -91,7 +89,7 @@ where
     );
 
     Ok(PartitionedExpr {
-        root: optimizer.optimize_typed(root, &root_scope)?,
+        root: root.simplify(&root_scope)?,
         partitions: partitions.into_boxed_slice(),
         partition_names,
         partition_dtypes: partition_dtypes.into_boxed_slice(),
@@ -222,9 +220,7 @@ mod tests {
     use crate::expr::exprs::pack::pack;
     use crate::expr::exprs::root::root;
     use crate::expr::exprs::select::select;
-    use crate::expr::session::ExprSession;
     use crate::expr::transform::replace::replace_root_fields;
-    use crate::expr::transform::simplify_typed::simplify_typed;
 
     #[fixture]
     fn dtype() -> DType {
@@ -247,48 +243,34 @@ mod tests {
     #[rstest]
     fn test_expr_top_level_ref(dtype: DType) {
         let fields = dtype.as_struct_fields_opt().unwrap();
-        let session = ExprSession::default();
-        let optimizer = ExprOptimizer::new(&session);
 
         let expr = root();
-        let partitioned = partition(
-            expr.clone(),
-            &dtype,
-            annotate_scope_access(fields),
-            &optimizer,
-        )
-        .unwrap();
+        let partitioned = partition(expr.clone(), &dtype, annotate_scope_access(fields)).unwrap();
 
         // An un-expanded root expression is annotated by all fields, but since it is a single node
         assert_eq!(partitioned.partitions.len(), 0);
         assert_eq!(&partitioned.root, &root());
 
         // Instead, callers must expand the root expression themselves.
         let expr = replace_root_fields(expr, fields);
-        let partitioned =
-            partition(expr, &dtype, annotate_scope_access(fields), &optimizer).unwrap();
+        let partitioned = partition(expr, &dtype, annotate_scope_access(fields)).unwrap();
 
         assert_eq!(partitioned.partitions.len(), fields.names().len());
     }
 
     #[rstest]
     fn test_expr_top_level_ref_get_item_and_split(dtype: DType) {
         let fields = dtype.as_struct_fields_opt().unwrap();
-        let session = ExprSession::default();
-        let optimizer = ExprOptimizer::new(&session);
 
         let expr = get_item("y", get_item("a", root()));
 
-        let partitioned =
-            partition(expr, &dtype, annotate_scope_access(fields), &optimizer).unwrap();
+        let partitioned = partition(expr, &dtype, annotate_scope_access(fields)).unwrap();
         assert_eq!(&partitioned.root, &get_item("a_0", get_item("a", root())));
     }
 
     #[rstest]
     fn test_expr_top_level_ref_get_item_and_split_pack(dtype: DType) {
         let fields = dtype.as_struct_fields_opt().unwrap();
-        let session = ExprSession::default();
-        let optimizer = ExprOptimizer::new(&session);
 
         let expr = pack(
             [
@@ -298,17 +280,11 @@ mod tests {
             ],
             NonNullable,
         );
-        let partitioned =
-            partition(expr, &dtype, annotate_scope_access(fields), &optimizer).unwrap();
+        let partitioned = partition(expr, &dtype, annotate_scope_access(fields)).unwrap();
 
         let split_a = partitioned.find_partition(&"a".into()).unwrap();
         assert_eq!(
-            &simplify_typed(
-                split_a.clone(),
-                &dtype,
-                ExprSession::default().rewrite_rules()
-            )
-            .unwrap(),
+            &split_a.simplify(&dtype).unwrap(),
             &pack(
                 [
                     ("a_0", get_item("x", get_item("a", root()))),
@@ -322,12 +298,9 @@ mod tests {
     #[rstest]
     fn test_expr_top_level_ref_get_item_add(dtype: DType) {
         let fields = dtype.as_struct_fields_opt().unwrap();
-        let session = ExprSession::default();
-        let optimizer = ExprOptimizer::new(&session);
 
         let expr = and(get_item("y", get_item("a", root())), lit(1));
-        let partitioned =
-            partition(expr, &dtype, annotate_scope_access(fields), &optimizer).unwrap();
+        let partitioned = partition(expr, &dtype, annotate_scope_access(fields)).unwrap();
 
         // Whole expr is a single split
         assert_eq!(partitioned.partitions.len(), 1);
@@ -336,12 +309,9 @@ mod tests {
     #[rstest]
     fn test_expr_top_level_ref_get_item_add_cannot_split(dtype: DType) {
         let fields = dtype.as_struct_fields_opt().unwrap();
-        let session = ExprSession::default();
-        let optimizer = ExprOptimizer::new(&session);
 
         let expr = and(get_item("y", get_item("a", root())), get_item("b", root()));
-        let partitioned =
-            partition(expr, &dtype, annotate_scope_access(fields), &optimizer).unwrap();
+        let partitioned = partition(expr, &dtype, annotate_scope_access(fields)).unwrap();
 
         // One for id.a and id.b
         assert_eq!(partitioned.partitions.len(), 2);
@@ -351,16 +321,13 @@ mod tests {
     #[rstest]
     fn test_expr_partition_many_occurrences_of_field(dtype: DType) {
         let fields = dtype.as_struct_fields_opt().unwrap();
-        let session = ExprSession::default();
-        let optimizer = ExprOptimizer::new(&session);
 
         let expr = and(
             get_item("y", get_item("a", root())),
             select(["a", "b"], root()),
         );
-        let expr = simplify_typed(expr, &dtype, ExprSession::default().rewrite_rules()).unwrap();
-        let partitioned =
-            partition(expr, &dtype, annotate_scope_access(fields), &optimizer).unwrap();
+        let expr = expr.simplify(&dtype).unwrap();
+        let partitioned = partition(expr, &dtype, annotate_scope_access(fields)).unwrap();
 
         // One for id.a and id.b
         assert_eq!(partitioned.partitions.len(), 2);
@@ -394,13 +361,10 @@ mod tests {
     #[rstest]
     fn test_expr_merge(dtype: DType) {
         let fields = dtype.as_struct_fields_opt().unwrap();
-        let session = ExprSession::default();
-        let optimizer = ExprOptimizer::new(&session);
 
         let expr = merge([col("a"), pack([("b", col("b"))], NonNullable)]);
 
-        let partitioned =
-            partition(expr, &dtype, annotate_scope_access(fields), &optimizer).unwrap();
+        let partitioned = partition(expr, &dtype, annotate_scope_access(fields)).unwrap();
         let expected = pack(
             [
                 ("x", get_item("x", get_item("a_0", col("a")))),