update

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

Author: joseph-isaacs

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

@@ -208,6 +208,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::immediate_access::annotate_scope_access;
     use crate::expr::transform::replace::replace_root_fields;
     use crate::expr::transform::simplify_typed::simplify_typed;
@@ -233,13 +234,14 @@ 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 expr = root();
         let partitioned = partition(
             expr.clone(),
             &dtype,
             annotate_scope_access(fields),
-            &ExprSession::default(),
+            &optimizer,
         )
         .unwrap();
 
@@ -249,36 +251,28 @@ mod tests {
 
         // Instead, callers must expand the root expression themselves.
         let expr = replace_root_fields(expr, fields);
-        let partitioned = partition(
-            expr,
-            &dtype,
-            annotate_scope_access(fields),
-            &ExprSession::default(),
-        )
-        .unwrap();
+        let partitioned =
+            partition(expr, &dtype, annotate_scope_access(fields), &optimizer).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 optimizer = ExprOptimizer::new(ExprSession::default());
 
         let expr = get_item("y", get_item("a", root()));
 
-        let partitioned = partition(
-            expr,
-            &dtype,
-            annotate_scope_access(fields),
-            &ExprSession::default(),
-        )
-        .unwrap();
+        let partitioned =
+            partition(expr, &dtype, annotate_scope_access(fields), &optimizer).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 optimizer = ExprOptimizer::new(ExprSession::default());
 
         let expr = pack(
             [
@@ -288,17 +282,12 @@ mod tests {
             ],
             NonNullable,
         );
-        let partitioned = partition(
-            expr,
-            &dtype,
-            annotate_scope_access(fields),
-            &ExprSession::default(),
-        )
-        .unwrap();
+        let partitioned =
+            partition(expr, &dtype, annotate_scope_access(fields), &optimizer).unwrap();
 
         let split_a = partitioned.find_partition(&"a".into()).unwrap();
         assert_eq!(
-            &simplify_typed(split_a.clone(), &dtype, &ExprSession::default(),).unwrap(),
+            &simplify_typed(split_a.clone(), &dtype, &ExprSession::default()).unwrap(),
             &pack(
                 [
                     ("a_0", get_item("x", get_item("a", root()))),
@@ -312,15 +301,11 @@ 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 expr = and(get_item("y", get_item("a", root())), lit(1));
-        let partitioned = partition(
-            expr,
-            &dtype,
-            annotate_scope_access(fields),
-            &ExprSession::default(),
-        )
-        .unwrap();
+        let partitioned =
+            partition(expr, &dtype, annotate_scope_access(fields), &optimizer).unwrap();
 
         // Whole expr is a single split
         assert_eq!(partitioned.partitions.len(), 1);
@@ -329,15 +314,11 @@ 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 expr = and(get_item("y", get_item("a", root())), get_item("b", root()));
-        let partitioned = partition(
-            expr,
-            &dtype,
-            annotate_scope_access(fields),
-            &ExprSession::default(),
-        )
-        .unwrap();
+        let partitioned =
+            partition(expr, &dtype, annotate_scope_access(fields), &optimizer).unwrap();
 
         // One for id.a and id.b
         assert_eq!(partitioned.partitions.len(), 2);
@@ -347,19 +328,15 @@ 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 expr = and(
             get_item("y", get_item("a", root())),
             select(["a", "b"], root()),
         );
         let expr = simplify_typed(expr, &dtype, &ExprSession::default()).unwrap();
-        let partitioned = partition(
-            expr,
-            &dtype,
-            annotate_scope_access(fields),
-            &ExprSession::default(),
-        )
-        .unwrap();
+        let partitioned =
+            partition(expr, &dtype, annotate_scope_access(fields), &optimizer).unwrap();
 
         // One for id.a and id.b
         assert_eq!(partitioned.partitions.len(), 2);
@@ -393,16 +370,12 @@ mod tests {
     #[rstest]
     fn test_expr_merge(dtype: DType) {
         let fields = dtype.as_struct_fields_opt().unwrap();
+        let optimizer = ExprOptimizer::new(ExprSession::default());
 
         let expr = merge([col("a"), pack([("b", col("b"))], NonNullable)]);
 
-        let partitioned = partition(
-            expr,
-            &dtype,
-            annotate_scope_access(fields),
-            &ExprSession::default(),
-        )
-        .unwrap();
+        let partitioned =
+            partition(expr, &dtype, annotate_scope_access(fields), &optimizer).unwrap();
         let expected = pack(
             [
                 ("x", get_item("x", get_item("a_0", col("a")))),