update

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

Author: joseph-isaacs

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

@@ -16,5 +16,5 @@ pub use optimizer::*;
 pub use partition::*;
 pub use replace::*;
 pub use rules::*;
-pub use simplify::*;
-pub use simplify_typed::*;
+pub(crate) use simplify::*;
+pub(crate) use simplify_typed::*;

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

@@ -9,33 +9,29 @@ 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.
-pub struct ExprOptimizer {
-    session: ExprSession,
-    dtype: Option<DType>,
+pub struct ExprOptimizer<'a> {
+    session: &'a ExprSession,
 }
 
-impl ExprOptimizer {
+impl<'a> ExprOptimizer<'a> {
     /// Create a new untyped optimizer.
     ///
     /// This optimizer will use untyped simplification rules only.
-    pub fn new(session: ExprSession) -> Self {
-        Self {
-            session,
-            dtype: None,
-        }
+    pub fn new(session: &'a ExprSession) -> Self {
+        Self { session }
     }
 
     /// 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> {
-        simplify(expr, &self.session)
+        simplify(expr, 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)
+        simplify_typed(expr, dtype, self.session)
     }
 }

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

@@ -234,7 +234,8 @@ mod tests {
     #[rstest]
     fn test_expr_top_level_ref(dtype: DType) {
         let fields = dtype.as_struct_fields_opt().unwrap();
-        let optimizer = ExprOptimizer::new(ExprSession::default());
+        let session = ExprSession::default();
+        let optimizer = ExprOptimizer::new(&session);
 
         let expr = root();
         let partitioned = partition(
@@ -260,7 +261,8 @@ mod tests {
     #[rstest]
     fn test_expr_top_level_ref_get_item_and_split(dtype: DType) {
         let fields = dtype.as_struct_fields_opt().unwrap();
-        let optimizer = ExprOptimizer::new(ExprSession::default());
+        let session = ExprSession::default();
+        let optimizer = ExprOptimizer::new(&session);
 
         let expr = get_item("y", get_item("a", root()));
 
@@ -272,7 +274,8 @@ mod tests {
     #[rstest]
     fn test_expr_top_level_ref_get_item_and_split_pack(dtype: DType) {
         let fields = dtype.as_struct_fields_opt().unwrap();
-        let optimizer = ExprOptimizer::new(ExprSession::default());
+        let session = ExprSession::default();
+        let optimizer = ExprOptimizer::new(&session);
 
         let expr = pack(
             [
@@ -301,7 +304,8 @@ mod tests {
     #[rstest]
     fn test_expr_top_level_ref_get_item_add(dtype: DType) {
         let fields = dtype.as_struct_fields_opt().unwrap();
-        let optimizer = ExprOptimizer::new(ExprSession::default());
+        let session = ExprSession::default();
+        let optimizer = ExprOptimizer::new(&session);
 
         let expr = and(get_item("y", get_item("a", root())), lit(1));
         let partitioned =
@@ -314,7 +318,8 @@ 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 optimizer = ExprOptimizer::new(ExprSession::default());
+        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 =
@@ -328,7 +333,8 @@ mod tests {
     #[rstest]
     fn test_expr_partition_many_occurrences_of_field(dtype: DType) {
         let fields = dtype.as_struct_fields_opt().unwrap();
-        let optimizer = ExprOptimizer::new(ExprSession::default());
+        let session = ExprSession::default();
+        let optimizer = ExprOptimizer::new(&session);
 
         let expr = and(
             get_item("y", get_item("a", root())),
@@ -370,7 +376,8 @@ mod tests {
     #[rstest]
     fn test_expr_merge(dtype: DType) {
         let fields = dtype.as_struct_fields_opt().unwrap();
-        let optimizer = ExprOptimizer::new(ExprSession::default());
+        let session = ExprSession::default();
+        let optimizer = ExprOptimizer::new(&session);
 
         let expr = merge([col("a"), pack([("b", col("b"))], NonNullable)]);
 

vortex-layout/src/layouts/row_idx/mod.rs

@@ -14,7 +14,7 @@ use futures::FutureExt;
 use futures::future::BoxFuture;
 use vortex_array::compute::filter;
 use vortex_array::expr::session::ExprSessionExt;
-use vortex_array::expr::transform::{PartitionedExpr, partition, replace};
+use vortex_array::expr::transform::{ExprOptimizer, PartitionedExpr, partition, replace};
 use vortex_array::expr::{ExactExpr, Expression, is_root, root};
 use vortex_array::{ArrayRef, IntoArray, MaskFuture};
 use vortex_dtype::{DType, FieldMask, FieldName, Nullability, PType};
@@ -32,8 +32,8 @@ pub struct RowIdxLayoutReader {
     name: Arc<str>,
     row_offset: u64,
     child: Arc<dyn LayoutReader>,
-    session: VortexSession,
     partition_cache: DashMap<ExactExpr, Partitioning>,
+    session: VortexSession,
 }
 
 impl RowIdxLayoutReader {
@@ -64,7 +64,7 @@ impl RowIdxLayoutReader {
                             vec![]
                         }
                     },
-                    &self.session.expressions(),
+                    &ExprOptimizer::new(&self.session.expressions()),
                 )
                 .vortex_expect("We should not fail to partition expression over struct fields");
 

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

@@ -33,16 +33,14 @@ pub struct StructReader {
     name: Arc<str>,
     lazy_children: LazyReaderChildren,
 
-    session: VortexSession,
-
     /// A `pack` expression that holds each individual field of the root DType. This expansion
     /// ensures we can correctly partition expressions over the fields of the struct.
     expanded_root_expr: Expression,
 
     field_lookup: Option<HashMap<FieldName, usize>>,
     partitioned_expr_cache: DashMap<ExactExpr, Partitioned>,
 
-    optimizer: ExprOptimizer,
+    session: VortexSession,
 }
 
 impl StructReader {
@@ -92,12 +90,11 @@ impl StructReader {
         Ok(Self {
             layout,
             name,
-            session,
             expanded_root_expr,
             lazy_children,
             field_lookup,
             partitioned_expr_cache: Default::default(),
-            optimizer: ExprOptimizer::new(*session.expressions()),
+            session,
         })
     }
 
@@ -144,8 +141,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 = self
-                    .optimizer
+                let exprs = self.session.expressions();
+                let opt = ExprOptimizer::new(&exprs);
+                let expr = opt
                     .optimize_typed(expr, self.dtype())
                     .vortex_expect("We should not fail to simplify expression over struct fields");
 
@@ -158,7 +156,7 @@ impl StructReader {
                             .as_struct_fields_opt()
                             .vortex_expect("We know it's a struct DType"),
                     ),
-                    &self.session.expressions(),
+                    &opt,
                 )
                 .vortex_expect("We should not fail to partition expression over struct fields");
 

vortex-scan/src/scan_builder.rs

@@ -216,10 +216,11 @@ impl<A: 'static + Send> ScanBuilder<A> {
             &self.session,
         ));
 
-        let optimizer = ExprOptimizer::new(*self.session.expressions());
+        let exprs = self.session.expressions();
+        let optimizer = ExprOptimizer::new(&exprs);
 
         // Normalize and simplify the expressions.
-        let projection = optimizer.optimize_typed(self.projection, layout_reader.as_ref())?;
+        let projection = optimizer.optimize_typed(self.projection, layout_reader.dtype())?;
 
         let filter = self
             .filter