Refactor tests for clarity and reusability

Extract shared helpers in push_down_filter_regressions.rs
and push_down_filter.rs to reduce code duplication.
Consolidate optimizer-delta test assertions and create
specific plan builders for common expressions.

Add a utility in utils.rs to evaluate predicates under
different null-restriction modes, streamlining
mode-comparison tests and enhancing maintainability.

Author: kosiew

datafusion/core/tests/sql/push_down_filter_regressions.rs

@@ -33,6 +33,9 @@ const WINDOW_SCALAR_SUBQUERY_SQL: &str = r#"
     )
 "#;
 
+const WINDOW_SCALAR_SUBQUERY_EXPECTED: &[&str] =
+    &["+----+", "| rn |", "+----+", "| 1  |", "+----+"];
+
 fn sqllogictest_style_ctx(push_down_filter_enabled: bool) -> SessionContext {
     let ctx =
         SessionContext::new_with_config(SessionConfig::new().with_target_partitions(4));
@@ -56,30 +59,20 @@ async fn capture_window_scalar_subquery_plans(
     ))
 }
 
-#[tokio::test]
-async fn window_scalar_subquery_regression() -> Result<()> {
-    let ctx = SessionContext::new();
+async fn assert_window_scalar_subquery(ctx: SessionContext) -> Result<()> {
     let results = ctx.sql(WINDOW_SCALAR_SUBQUERY_SQL).await?.collect().await?;
-
-    assert_batches_eq!(
-        &["+----+", "| rn |", "+----+", "| 1  |", "+----+",],
-        &results
-    );
-
+    assert_batches_eq!(WINDOW_SCALAR_SUBQUERY_EXPECTED, &results);
     Ok(())
 }
 
 #[tokio::test]
-async fn window_scalar_subquery_sqllogictest_style_regression() -> Result<()> {
-    let ctx = sqllogictest_style_ctx(true);
-    let results = ctx.sql(WINDOW_SCALAR_SUBQUERY_SQL).await?.collect().await?;
-
-    assert_batches_eq!(
-        &["+----+", "| rn |", "+----+", "| 1  |", "+----+",],
-        &results
-    );
+async fn window_scalar_subquery_regression() -> Result<()> {
+    assert_window_scalar_subquery(SessionContext::new()).await
+}
 
-    Ok(())
+#[tokio::test]
+async fn window_scalar_subquery_sqllogictest_style_regression() -> Result<()> {
+    assert_window_scalar_subquery(sqllogictest_style_ctx(true)).await
 }
 
 #[tokio::test]
@@ -212,28 +205,18 @@ async fn window_scalar_subquery_optimizer_delta() -> Result<()> {
     let (disabled_optimized, disabled_physical) =
         capture_window_scalar_subquery_plans(false).await?;
 
+    assert_eq!(enabled_optimized, disabled_optimized);
+    assert_eq!(enabled_physical, disabled_physical);
+
     assert!(
         enabled_optimized
             .contains("Filter: s.acctbal > __scalar_sq_1.avg(suppliers.acctbal)")
     );
     assert!(enabled_optimized.contains("Cross Join:"));
-    assert!(
-        disabled_optimized
-            .contains("Filter: s.acctbal > __scalar_sq_1.avg(suppliers.acctbal)")
-    );
-    assert!(disabled_optimized.contains("Cross Join:"));
-
     assert!(
         enabled_physical.contains("FilterExec: acctbal@1 > avg(suppliers.acctbal)@2")
     );
     assert!(enabled_physical.contains("CrossJoinExec"));
-    assert!(
-        disabled_physical.contains("FilterExec: acctbal@1 > avg(suppliers.acctbal)@2")
-    );
-    assert!(disabled_physical.contains("CrossJoinExec"));
-
-    assert_eq!(enabled_optimized, disabled_optimized);
-    assert_eq!(enabled_physical, disabled_physical);
 
     Ok(())
 }

datafusion/optimizer/src/push_down_filter.rs

@@ -1522,6 +1522,53 @@ mod tests {
 
     use super::*;
 
+    fn scalar_subquery_right_plan() -> Result<LogicalPlan> {
+        LogicalPlanBuilder::from(test_table_scan_with_name("test1")?)
+            .project(vec![col("a").alias("acctbal")])?
+            .aggregate(
+                Vec::<Expr>::new(),
+                vec![avg(col("acctbal")).alias("avg_acctbal")],
+            )?
+            .alias("__scalar_sq_1")?
+            .build()
+    }
+
+    fn row_number_window_expr() -> Expr {
+        Expr::from(WindowFunction::new(
+            WindowFunctionDefinition::WindowUDF(
+                datafusion_functions_window::row_number::row_number_udwf(),
+            ),
+            vec![],
+        ))
+        .partition_by(vec![col("s.nation")])
+        .order_by(vec![col("s.acctbal").sort(false, true)])
+        .build()
+        .unwrap()
+    }
+
+    fn window_over_scalar_subquery_cross_join_plan(
+        with_project_wrapper: bool,
+    ) -> Result<LogicalPlan> {
+        let left = {
+            let builder = LogicalPlanBuilder::from(test_table_scan()?)
+                .project(vec![col("a").alias("nation"), col("b").alias("acctbal")])?
+                .alias("s")?;
+            let builder = if with_project_wrapper {
+                builder.project(vec![col("s.nation"), col("s.acctbal")])?
+            } else {
+                builder
+            };
+            builder.build()?
+        };
+
+        LogicalPlanBuilder::from(left)
+            .cross_join(scalar_subquery_right_plan()?)?
+            .filter(col("s.acctbal").gt(col("__scalar_sq_1.avg_acctbal")))?
+            .project(vec![col("s.nation"), col("s.acctbal")])?
+            .window(vec![row_number_window_expr()])?
+            .build()
+    }
+
     fn observe(_plan: &LogicalPlan, _rule: &dyn OptimizerRule) {}
 
     macro_rules! assert_optimized_plan_equal {
@@ -2443,36 +2490,7 @@ mod tests {
 
     #[test]
     fn window_over_scalar_subquery_cross_join_keeps_filter_above_join() -> Result<()> {
-        let left = LogicalPlanBuilder::from(test_table_scan()?)
-            .project(vec![col("a").alias("nation"), col("b").alias("acctbal")])?
-            .alias("s")?
-            .build()?;
-        let right = LogicalPlanBuilder::from(test_table_scan_with_name("test1")?)
-            .project(vec![col("a").alias("acctbal")])?
-            .aggregate(
-                Vec::<Expr>::new(),
-                vec![avg(col("acctbal")).alias("avg_acctbal")],
-            )?
-            .alias("__scalar_sq_1")?
-            .build()?;
-
-        let window = Expr::from(WindowFunction::new(
-            WindowFunctionDefinition::WindowUDF(
-                datafusion_functions_window::row_number::row_number_udwf(),
-            ),
-            vec![],
-        ))
-        .partition_by(vec![col("s.nation")])
-        .order_by(vec![col("s.acctbal").sort(false, true)])
-        .build()
-        .unwrap();
-
-        let plan = LogicalPlanBuilder::from(left)
-            .cross_join(right)?
-            .filter(col("s.acctbal").gt(col("__scalar_sq_1.avg_acctbal")))?
-            .project(vec![col("s.nation"), col("s.acctbal")])?
-            .window(vec![window])?
-            .build()?;
+        let plan = window_over_scalar_subquery_cross_join_plan(false)?;
 
         assert_optimized_plan_equal!(
             plan,
@@ -2495,37 +2513,7 @@ mod tests {
     #[test]
     fn window_over_scalar_subquery_cross_join_with_project_wrapper_keeps_filter_above_join()
     -> Result<()> {
-        let left = LogicalPlanBuilder::from(test_table_scan()?)
-            .project(vec![col("a").alias("nation"), col("b").alias("acctbal")])?
-            .alias("s")?
-            .project(vec![col("s.nation"), col("s.acctbal")])?
-            .build()?;
-        let right = LogicalPlanBuilder::from(test_table_scan_with_name("test1")?)
-            .project(vec![col("a").alias("acctbal")])?
-            .aggregate(
-                Vec::<Expr>::new(),
-                vec![avg(col("acctbal")).alias("avg_acctbal")],
-            )?
-            .alias("__scalar_sq_1")?
-            .build()?;
-
-        let window = Expr::from(WindowFunction::new(
-            WindowFunctionDefinition::WindowUDF(
-                datafusion_functions_window::row_number::row_number_udwf(),
-            ),
-            vec![],
-        ))
-        .partition_by(vec![col("s.nation")])
-        .order_by(vec![col("s.acctbal").sort(false, true)])
-        .build()
-        .unwrap();
-
-        let plan = LogicalPlanBuilder::from(left)
-            .cross_join(right)?
-            .filter(col("s.acctbal").gt(col("__scalar_sq_1.avg_acctbal")))?
-            .project(vec![col("s.nation"), col("s.acctbal")])?
-            .window(vec![window])?
-            .build()?;
+        let plan = window_over_scalar_subquery_cross_join_plan(true)?;
 
         assert_optimized_plan_equal!(
             plan,

datafusion/optimizer/src/utils.rs

@@ -250,6 +250,23 @@ mod tests {
         Operator, binary_expr, case, col, in_list, is_null, lit, when,
     };
 
+    fn restrict_null_predicate_in_modes(
+        predicate: Expr,
+        join_cols: &[Column],
+    ) -> Result<(bool, bool)> {
+        let auto_result = with_null_restriction_eval_mode_for_test(
+            NullRestrictionEvalMode::Auto,
+            || is_restrict_null_predicate(predicate.clone(), join_cols.iter()),
+        )?;
+
+        let authoritative_result = with_null_restriction_eval_mode_for_test(
+            NullRestrictionEvalMode::AuthoritativeOnly,
+            || is_restrict_null_predicate(predicate.clone(), join_cols.iter()),
+        )?;
+
+        Ok((auto_result, authoritative_result))
+    }
+
     #[test]
     fn expr_is_restrict_null_predicate() -> Result<()> {
         let test_cases = vec![
@@ -465,27 +482,13 @@ mod tests {
     #[test]
     fn null_restriction_eval_mode_auto_vs_authoritative_only() -> Result<()> {
         let predicate = binary_expr(col("a"), Operator::Gt, lit(8i64));
-        let join_cols_of_predicate = predicate.column_refs();
-
-        let auto_result = with_null_restriction_eval_mode_for_test(
-            NullRestrictionEvalMode::Auto,
-            || {
-                is_restrict_null_predicate(
-                    predicate.clone(),
-                    join_cols_of_predicate.iter().copied(),
-                )
-            },
-        )?;
-
-        let authoritative_result = with_null_restriction_eval_mode_for_test(
-            NullRestrictionEvalMode::AuthoritativeOnly,
-            || {
-                is_restrict_null_predicate(
-                    predicate.clone(),
-                    join_cols_of_predicate.iter().copied(),
-                )
-            },
-        )?;
+        let join_cols_of_predicate = predicate
+            .column_refs()
+            .into_iter()
+            .cloned()
+            .collect::<Vec<_>>();
+        let (auto_result, authoritative_result) =
+            restrict_null_predicate_in_modes(predicate, &join_cols_of_predicate)?;
 
         assert_eq!(auto_result, authoritative_result);
 
@@ -496,17 +499,9 @@ mod tests {
     fn mixed_reference_predicate_remains_fast_pathed_in_authoritative_mode() -> Result<()>
     {
         let predicate = binary_expr(col("a"), Operator::Gt, col("b"));
-        let column_a = Column::from_name("a");
-
-        let auto_result = with_null_restriction_eval_mode_for_test(
-            NullRestrictionEvalMode::Auto,
-            || is_restrict_null_predicate(predicate.clone(), std::iter::once(&column_a)),
-        )?;
-
-        let authoritative_only_result = with_null_restriction_eval_mode_for_test(
-            NullRestrictionEvalMode::AuthoritativeOnly,
-            || is_restrict_null_predicate(predicate.clone(), std::iter::once(&column_a)),
-        )?;
+        let join_cols = vec![Column::from_name("a")];
+        let (auto_result, authoritative_only_result) =
+            restrict_null_predicate_in_modes(predicate.clone(), &join_cols)?;
 
         assert!(!auto_result, "{predicate}");
         assert!(!authoritative_only_result, "{predicate}");