Migrate Optimizer tests to insta, part7

datafusion/optimizer/src/eliminate_duplicated_expr.rs

@@ -120,6 +120,7 @@ mod tests {
     use super::*;
     use crate::assert_optimized_plan_eq_snapshot;
     use crate::test::*;
+    use crate::OptimizerContext;
     use datafusion_expr::{col, logical_plan::builder::LogicalPlanBuilder};
     use std::sync::Arc;
 
@@ -128,9 +129,11 @@ mod tests {
             $plan:expr,
             @ $expected:literal $(,)?
         ) => {{
-            let rule: Arc<dyn crate::OptimizerRule + Send + Sync> = Arc::new(EliminateDuplicatedExpr::new());
+            let optimizer_ctx = OptimizerContext::new().with_max_passes(1);
+            let rules: Vec<Arc<dyn crate::OptimizerRule + Send + Sync>> = vec![Arc::new(EliminateDuplicatedExpr::new())];
             assert_optimized_plan_eq_snapshot!(
-                rule,
+                optimizer_ctx,
+                rules,
                 $plan,
                 @ $expected,
             )

datafusion/optimizer/src/eliminate_filter.rs

@@ -82,6 +82,7 @@ mod tests {
     use std::sync::Arc;
 
     use crate::assert_optimized_plan_eq_snapshot;
+    use crate::OptimizerContext;
     use datafusion_common::{Result, ScalarValue};
     use datafusion_expr::{col, lit, logical_plan::builder::LogicalPlanBuilder, Expr};
 
@@ -94,9 +95,11 @@ mod tests {
             $plan:expr,
             @ $expected:literal $(,)?
         ) => {{
-            let rule: Arc<dyn crate::OptimizerRule + Send + Sync> = Arc::new(EliminateFilter::new());
+            let optimizer_ctx = OptimizerContext::new().with_max_passes(1);
+            let rules: Vec<Arc<dyn crate::OptimizerRule + Send + Sync>> = vec![Arc::new(EliminateFilter::new())];
             assert_optimized_plan_eq_snapshot!(
-                rule,
+                optimizer_ctx,
+                rules,
                 $plan,
                 @ $expected,
             )

datafusion/optimizer/src/eliminate_group_by_constant.rs

@@ -117,6 +117,7 @@ mod tests {
     use super::*;
     use crate::assert_optimized_plan_eq_snapshot;
     use crate::test::*;
+    use crate::OptimizerContext;
 
     use arrow::datatypes::DataType;
     use datafusion_common::Result;
@@ -135,9 +136,11 @@ mod tests {
             $plan:expr,
             @ $expected:literal $(,)?
         ) => {{
-            let rule: Arc<dyn crate::OptimizerRule + Send + Sync> = Arc::new(EliminateGroupByConstant::new());
+            let optimizer_ctx = OptimizerContext::new().with_max_passes(1);
+            let rules: Vec<Arc<dyn crate::OptimizerRule + Send + Sync>> = vec![Arc::new(EliminateGroupByConstant::new())];
             assert_optimized_plan_eq_snapshot!(
-                rule,
+                optimizer_ctx,
+                rules,
                 $plan,
                 @ $expected,
             )

datafusion/optimizer/src/eliminate_join.rs

@@ -76,6 +76,7 @@ impl OptimizerRule for EliminateJoin {
 mod tests {
     use crate::assert_optimized_plan_eq_snapshot;
     use crate::eliminate_join::EliminateJoin;
+    use crate::OptimizerContext;
     use datafusion_common::Result;
     use datafusion_expr::JoinType::Inner;
     use datafusion_expr::{lit, logical_plan::builder::LogicalPlanBuilder};
@@ -84,11 +85,13 @@ mod tests {
     macro_rules! assert_optimized_plan_equal {
         (
             $plan:expr,
-            @$expected:literal $(,)?
+            @ $expected:literal $(,)?
         ) => {{
-            let rule: Arc<dyn crate::OptimizerRule + Send + Sync> = Arc::new(EliminateJoin::new());
+            let optimizer_ctx = OptimizerContext::new().with_max_passes(1);
+            let rules: Vec<Arc<dyn crate::OptimizerRule + Send + Sync>> = vec![Arc::new(EliminateJoin::new())];
             assert_optimized_plan_eq_snapshot!(
-                rule,
+                optimizer_ctx,
+                rules,
                 $plan,
                 @ $expected,
             )

datafusion/optimizer/src/eliminate_nested_union.rs

@@ -117,6 +117,7 @@ mod tests {
     use crate::analyzer::type_coercion::TypeCoercion;
     use crate::analyzer::Analyzer;
     use crate::assert_optimized_plan_eq_snapshot;
+    use crate::OptimizerContext;
     use arrow::datatypes::{DataType, Field, Schema};
     use datafusion_common::config::ConfigOptions;
     use datafusion_expr::{col, logical_plan::table_scan};
@@ -137,9 +138,11 @@ mod tests {
             let options = ConfigOptions::default();
             let analyzed_plan = Analyzer::with_rules(vec![Arc::new(TypeCoercion::new())])
                 .execute_and_check($plan, &options, |_, _| {})?;
-            let rule: Arc<dyn crate::OptimizerRule + Send + Sync> = Arc::new(EliminateNestedUnion::new());
+            let optimizer_ctx = OptimizerContext::new().with_max_passes(1);
+            let rules: Vec<Arc<dyn crate::OptimizerRule + Send + Sync>> = vec![Arc::new(EliminateNestedUnion::new())];
             assert_optimized_plan_eq_snapshot!(
-                rule,
+                optimizer_ctx,
+                rules,
                 analyzed_plan,
                 @ $expected,
             )

datafusion/optimizer/src/eliminate_outer_join.rs

@@ -306,6 +306,7 @@ mod tests {
     use super::*;
     use crate::assert_optimized_plan_eq_snapshot;
     use crate::test::*;
+    use crate::OptimizerContext;
     use arrow::datatypes::DataType;
     use datafusion_expr::{
         binary_expr, cast, col, lit,
@@ -317,11 +318,13 @@ mod tests {
     macro_rules! assert_optimized_plan_equal {
         (
             $plan:expr,
-            @$expected:literal $(,)?
+            @ $expected:literal $(,)?
         ) => {{
-            let rule: Arc<dyn crate::OptimizerRule + Send + Sync> = Arc::new(EliminateOuterJoin::new());
+            let optimizer_ctx = OptimizerContext::new().with_max_passes(1);
+            let rules: Vec<Arc<dyn crate::OptimizerRule + Send + Sync>> = vec![Arc::new(EliminateOuterJoin::new())];
             assert_optimized_plan_eq_snapshot!(
-                rule,
+                optimizer_ctx,
+                rules,
                 $plan,
                 @ $expected,
             )

datafusion/optimizer/src/filter_null_join_keys.rs

@@ -108,6 +108,7 @@ fn create_not_null_predicate(filters: Vec<Expr>) -> Expr {
 mod tests {
     use super::*;
     use crate::assert_optimized_plan_eq_snapshot;
+    use crate::OptimizerContext;
     use arrow::datatypes::{DataType, Field, Schema};
     use datafusion_common::Column;
     use datafusion_expr::logical_plan::table_scan;
@@ -118,9 +119,11 @@ mod tests {
             $plan:expr,
             @ $expected:literal $(,)?
         ) => {{
-            let rule: Arc<dyn crate::OptimizerRule + Send + Sync> = Arc::new(FilterNullJoinKeys {});
+            let optimizer_ctx = OptimizerContext::new().with_max_passes(1);
+            let rules: Vec<Arc<dyn crate::OptimizerRule + Send + Sync>> = vec![Arc::new(FilterNullJoinKeys {})];
             assert_optimized_plan_eq_snapshot!(
-                rule,
+                optimizer_ctx,
+                rules,
                 $plan,
                 @ $expected,
             )

datafusion/optimizer/src/optimize_projections/mod.rs

@@ -847,9 +847,11 @@ mod tests {
             $plan:expr,
             @ $expected:literal $(,)?
         ) => {{
-            let rule: Arc<dyn crate::OptimizerRule + Send + Sync> = Arc::new(OptimizeProjections::new());
+            let optimizer_ctx = OptimizerContext::new().with_max_passes(1);
+            let rules: Vec<Arc<dyn crate::OptimizerRule + Send + Sync>> = vec![Arc::new(OptimizeProjections::new())];
             assert_optimized_plan_eq_snapshot!(
-                rule,
+                optimizer_ctx,
+                rules,
                 $plan,
                 @ $expected,
             )

datafusion/optimizer/src/propagate_empty_relation.rs

@@ -249,6 +249,7 @@ mod tests {
         assert_optimized_plan_with_rules, test_table_scan, test_table_scan_fields,
         test_table_scan_with_name,
     };
+    use crate::OptimizerContext;
 
     use super::*;
 
@@ -257,9 +258,11 @@ mod tests {
             $plan:expr,
             @ $expected:literal $(,)?
         ) => {{
-            let rule: Arc<dyn crate::OptimizerRule + Send + Sync> = Arc::new(PropagateEmptyRelation::new());
+            let optimizer_ctx = OptimizerContext::new().with_max_passes(1);
+            let rules: Vec<Arc<dyn crate::OptimizerRule + Send + Sync>> = vec![Arc::new(PropagateEmptyRelation::new())];
             assert_optimized_plan_eq_snapshot!(
-                rule,
+                optimizer_ctx,
+                rules,
                 $plan,
                 @ $expected,
             )

datafusion/optimizer/src/push_down_filter.rs

@@ -1418,9 +1418,11 @@ mod tests {
             $plan:expr,
             @ $expected:literal $(,)?
         ) => {{
-            let rule: Arc<dyn crate::OptimizerRule + Send + Sync> = Arc::new(PushDownFilter::new());
+            let optimizer_ctx = OptimizerContext::new().with_max_passes(1);
+            let rules: Vec<Arc<dyn crate::OptimizerRule + Send + Sync>> = vec![Arc::new(PushDownFilter::new())];
             assert_optimized_plan_eq_snapshot!(
-                rule,
+                optimizer_ctx,
+                rules,
                 $plan,
                 @ $expected,
             )

datafusion/optimizer/src/push_down_limit.rs

@@ -279,6 +279,7 @@ mod test {
     use crate::assert_optimized_plan_eq_snapshot;
     use crate::test::*;
 
+    use crate::OptimizerContext;
     use datafusion_common::DFSchemaRef;
     use datafusion_expr::{
         col, exists, logical_plan::builder::LogicalPlanBuilder, Expr, Extension,
@@ -291,9 +292,11 @@ mod test {
             $plan:expr,
             @ $expected:literal $(,)?
         ) => {{
-            let rule: Arc<dyn crate::OptimizerRule + Send + Sync> = Arc::new(PushDownLimit::new());
+            let optimizer_ctx = OptimizerContext::new().with_max_passes(1);
+            let rules: Vec<Arc<dyn crate::OptimizerRule + Send + Sync>> = vec![Arc::new(PushDownLimit::new())];
             assert_optimized_plan_eq_snapshot!(
-                rule,
+                optimizer_ctx,
+                rules,
                 $plan,
                 @ $expected,
             )

datafusion/optimizer/src/replace_distinct_aggregate.rs

@@ -190,6 +190,7 @@ mod tests {
     use crate::replace_distinct_aggregate::ReplaceDistinctWithAggregate;
     use crate::test::*;
 
+    use crate::OptimizerContext;
     use datafusion_common::Result;
     use datafusion_expr::{col, logical_plan::builder::LogicalPlanBuilder, Expr};
     use datafusion_functions_aggregate::sum::sum;
@@ -199,9 +200,11 @@ mod tests {
             $plan:expr,
             @ $expected:literal $(,)?
         ) => {{
-            let rule: Arc<dyn crate::OptimizerRule + Send + Sync> = Arc::new(ReplaceDistinctWithAggregate::new());
+            let optimizer_ctx = OptimizerContext::new().with_max_passes(1);
+            let rules: Vec<Arc<dyn crate::OptimizerRule + Send + Sync>> = vec![Arc::new(ReplaceDistinctWithAggregate::new())];
             assert_optimized_plan_eq_snapshot!(
-                rule,
+                optimizer_ctx,
+                rules,
                 $plan,
                 @ $expected,
             )

datafusion/optimizer/src/test/mod.rs

@@ -147,21 +147,6 @@ macro_rules! assert_optimized_plan_eq_snapshot {
 
     Ok::<(), datafusion_common::DataFusionError>(())
     }};
-
-    (
-        $rule:expr,
-        $plan:expr,
-        @ $expected:literal $(,)?
-    ) => {{
-    // Apply the rule once
-    let opt_context = $crate::OptimizerContext::new().with_max_passes(1);
-
-    let optimizer = $crate::Optimizer::with_rules(vec![Arc::clone(&$rule)]);
-    let optimized_plan = optimizer.optimize($plan, &opt_context, |_, _| {})?;
-    insta::assert_snapshot!(optimized_plan, @ $expected);
-
-    Ok::<(), datafusion_common::DataFusionError>(())
-    }};
 }
 
 fn generate_optimized_plan_with_rules(