update

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

Author: joseph-isaacs

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

@@ -9,74 +9,33 @@ use crate::expr::session::ExprSession;
 use crate::expr::transform::{simplify, simplify_typed};
 
 /// A unified optimizer for expressions that can work with or without type information.
-///
-/// This provides a convenient API for optimizing expressions, automatically choosing between
-/// typed and untyped optimization based on whether a dtype is provided.
-///
-/// # Examples
-///
-/// ```rust
-/// # use vortex_array::expr::Expression;
-/// # use vortex_array::expr::session::ExprSession;
-/// # use vortex_array::expr::transform::ExprOptimizer;
-/// # use vortex_dtype::DType;
-/// # fn example(expr: Expression, dtype: DType) {
-/// let session = ExprSession::default();
-///
-/// // Untyped optimization
-/// let optimizer = ExprOptimizer::new(&session);
-/// let optimized = optimizer.optimize(expr.clone()).unwrap();
-///
-/// // Typed optimization
-/// let optimizer = ExprOptimizer::with_dtype(&session, dtype);
-/// let optimized = optimizer.optimize(expr).unwrap();
-/// # }
-/// ```
-pub struct ExprOptimizer<'a> {
-    session: &'a ExprSession,
+pub struct ExprOptimizer {
+    session: ExprSession,
     dtype: Option<DType>,
 }
 
-impl<'a> ExprOptimizer<'a> {
+impl ExprOptimizer {
     /// Create a new untyped optimizer.
     ///
     /// This optimizer will use untyped simplification rules only.
-    pub fn new(session: &'a ExprSession) -> Self {
+    pub fn new(session: ExprSession) -> Self {
         Self {
             session,
             dtype: None,
         }
     }
 
-    /// Create a new typed optimizer with the given dtype.
-    ///
-    /// This optimizer will use both typed and untyped simplification rules,
-    /// with access to dtype information.
-    pub fn with_dtype(session: &'a ExprSession, dtype: DType) -> Self {
-        Self {
-            session,
-            dtype: Some(dtype),
-        }
-    }
-
     /// Optimize the given expression.
     ///
     /// If this optimizer was created with a dtype, this will perform typed optimization.
     /// Otherwise, it will perform untyped optimization.
     pub fn optimize(&self, expr: Expression) -> VortexResult<Expression> {
-        match &self.dtype {
-            Some(dtype) => simplify_typed(expr, dtype, self.session),
-            None => simplify(expr, self.session),
-        }
-    }
-
-    /// Get the dtype associated with this optimizer, if any.
-    pub fn dtype(&self) -> Option<&DType> {
-        self.dtype.as_ref()
+        simplify(expr, &self.session)
     }
 
-    /// Get the session associated with this optimizer.
-    pub fn session(&self) -> &ExprSession {
-        self.session
+    /// Apply optimize rules to the expression, with a known dtype. This will also apply rules
+    /// in `optimize`.
+    pub fn optimize_typed(&self, expr: Expression, dtype: &DType) -> VortexResult<Expression> {
+        simplify_typed(expr, dtype, &self.session)
     }
 }

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

@@ -12,11 +12,10 @@ use crate::expr::Expression;
 use crate::expr::exprs::get_item::get_item;
 use crate::expr::exprs::pack::pack;
 use crate::expr::exprs::root::root;
-use crate::expr::session::ExprSession;
+use crate::expr::transform::ExprOptimizer;
 use crate::expr::transform::annotations::{
     Annotation, AnnotationFn, Annotations, descendent_annotations,
 };
-use crate::expr::transform::simplify_typed::simplify_typed;
 use crate::expr::traversal::{NodeExt, NodeRewriter, Transformed, TraversalOrder};
 
 /// Partition an expression into sub-expressions that are uniquely associated with an annotation.
@@ -34,7 +33,7 @@ pub fn partition<A: AnnotationFn>(
     expr: Expression,
     scope: &DType,
     annotate_fn: A,
-    session: &ExprSession,
+    optimizer: &ExprOptimizer,
 ) -> VortexResult<PartitionedExpr<A::Annotation>>
 where
     A::Annotation: Display,
@@ -64,7 +63,7 @@ where
             Nullability::NonNullable,
         );
 
-        let expr = simplify_typed(expr.clone(), scope, session)?;
+        let expr = optimizer.optimize_typed(expr.clone(), scope)?;
         let expr_dtype = expr.return_dtype(scope)?;
 
         partitions.push(expr);
@@ -82,7 +81,7 @@ where
     );
 
     Ok(PartitionedExpr {
-        root: simplify_typed(root, &root_scope, session)?,
+        root: optimizer.optimize_typed(root, &root_scope)?,
         partitions: partitions.into_boxed_slice(),
         partition_names,
         partition_dtypes: partition_dtypes.into_boxed_slice(),

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

@@ -13,7 +13,7 @@ use crate::expr::traversal::{NodeExt, Transformed};
 ///
 /// This applies only untyped rewrite rules registered in the default session.
 /// If the scope dtype is known, see `simplify_typed` for a simplifier which uses dtype.
-pub fn simplify(e: Expression, session: &ExprSession) -> VortexResult<Expression> {
+pub(crate) fn simplify(e: Expression, session: &ExprSession) -> VortexResult<Expression> {
     let ctx = RuleContext;
 
     let e = apply_parent_rules(e, &ctx, session)?;

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

@@ -14,7 +14,7 @@ use crate::expr::traversal::{NodeExt, Transformed};
 ///
 /// NOTE: After typed simplification, returned expressions is "bound" to the scope DType.
 ///     Applying the returned expression to a different DType may produce wrong results.
-pub fn simplify_typed(
+pub(crate) fn simplify_typed(
     expr: Expression,
     dtype: &DType,
     session: &ExprSession,

vortex-layout/src/layouts/struct_/reader.rs

@@ -11,7 +11,7 @@ use vortex_array::arrays::StructArray;
 use vortex_array::expr::session::ExprSessionExt;
 use vortex_array::expr::transform::immediate_access::annotate_scope_access;
 use vortex_array::expr::transform::{
-    PartitionedExpr, partition, replace, replace_root_fields, simplify_typed,
+    ExprOptimizer, PartitionedExpr, partition, replace, replace_root_fields,
 };
 use vortex_array::expr::{ExactExpr, Expression, Merge, Pack, col, root};
 use vortex_array::vtable::ValidityHelper;
@@ -41,6 +41,8 @@ pub struct StructReader {
 
     field_lookup: Option<HashMap<FieldName, usize>>,
     partitioned_expr_cache: DashMap<ExactExpr, Partitioned>,
+
+    optimizer: ExprOptimizer,
 }
 
 impl StructReader {
@@ -95,6 +97,7 @@ impl StructReader {
             lazy_children,
             field_lookup,
             partitioned_expr_cache: Default::default(),
+            optimizer: ExprOptimizer::new(*session.expressions()),
         })
     }
 
@@ -141,7 +144,9 @@ impl StructReader {
                 // First, we expand the root scope into the fields of the struct to ensure
                 // that partitioning works correctly.
                 let expr = replace(expr.clone(), &root(), self.expanded_root_expr.clone());
-                let expr = simplify_typed(expr, self.dtype(), &self.session.expressions())
+                let expr = self
+                    .optimizer
+                    .optimize_typed(expr, self.dtype())
                     .vortex_expect("We should not fail to simplify expression over struct fields");
 
                 // Partition the expression into expressions that can be evaluated over individual fields

vortex-scan/src/scan_builder.rs

@@ -9,8 +9,8 @@ use futures::future::BoxFuture;
 use itertools::Itertools;
 use vortex_array::ArrayRef;
 use vortex_array::expr::session::ExprSessionExt;
+use vortex_array::expr::transform::ExprOptimizer;
 use vortex_array::expr::transform::immediate_access::immediate_scope_access;
-use vortex_array::expr::transform::simplify_typed;
 use vortex_array::expr::{Expression, root};
 use vortex_array::iter::{ArrayIterator, ArrayIteratorAdapter};
 use vortex_array::stats::StatsSet;
@@ -216,15 +216,14 @@ impl<A: 'static + Send> ScanBuilder<A> {
             &self.session,
         ));
 
+        let optimizer = ExprOptimizer::new(*self.session.expressions());
+
         // Normalize and simplify the expressions.
-        let projection = simplify_typed(
-            self.projection,
-            layout_reader.dtype(),
-            &self.session.expressions(),
-        )?;
+        let projection = optimizer.optimize_typed(self.projection, layout_reader.as_ref())?;
+
         let filter = self
             .filter
-            .map(|f| simplify_typed(f, layout_reader.dtype(), &self.session.expressions()))
+            .map(|f| optimizer.optimize_typed(f, layout_reader.dtype()))
             .transpose()?;
 
         // Construct field masks and compute the row splits of the scan.