sql: Simplify quantified comparisons in JOIN ON clauses

src/adapter/src/optimize.rs

@@ -226,6 +226,9 @@ impl From<&OptimizerConfig> for mz_sql::plan::HirToMirConfig {
             enable_variadic_left_join_lowering: config.features.enable_variadic_left_join_lowering,
             enable_guard_subquery_tablefunc: config.features.enable_guard_subquery_tablefunc,
             enable_cast_elimination: config.features.enable_cast_elimination,
+            enable_simplify_quantified_comparisons: config
+                .features
+                .enable_simplify_quantified_comparisons,
         }
     }
 }

src/repr/src/optimize.rs

@@ -134,6 +134,8 @@ optimizer_feature_flags!({
     enable_cast_elimination: bool,
     // See the feature flag of the same name.
     enable_case_literal_transform: bool,
+    // See the feature flag of the same name.
+    enable_simplify_quantified_comparisons: bool,
 });
 
 /// A trait used to implement layered config construction.

src/sql/src/plan/lowering.rs

@@ -141,6 +141,7 @@ pub struct Config {
     pub enable_variadic_left_join_lowering: bool,
     pub enable_guard_subquery_tablefunc: bool,
     pub enable_cast_elimination: bool,
+    pub enable_simplify_quantified_comparisons: bool,
 }
 
 impl Default for Config {
@@ -150,6 +151,7 @@ impl Default for Config {
             enable_variadic_left_join_lowering: false,
             enable_guard_subquery_tablefunc: false,
             enable_cast_elimination: false,
+            enable_simplify_quantified_comparisons: false,
         }
     }
 }
@@ -161,6 +163,7 @@ impl From<&SystemVars> for Config {
             enable_variadic_left_join_lowering: vars.enable_variadic_left_join_lowering(),
             enable_guard_subquery_tablefunc: vars.enable_guard_subquery_tablefunc(),
             enable_cast_elimination: vars.enable_cast_elimination(),
+            enable_simplify_quantified_comparisons: vars.enable_simplify_quantified_comparisons(),
         }
     }
 }
@@ -203,7 +206,10 @@ impl HirRelationExpr {
             mut other => {
                 let mut id_gen = mz_ore::id_gen::IdGen::default();
                 transform_hir::split_subquery_predicates(&mut other)?;
-                transform_hir::try_simplify_quantified_comparisons(&mut other)?;
+                transform_hir::try_simplify_quantified_comparisons(
+                    &mut other,
+                    context.config.enable_simplify_quantified_comparisons,
+                )?;
                 transform_hir::fuse_window_functions(&mut other, &context)?;
                 MirRelationExpr::constant(vec![vec![]], ReprRelationType::new(vec![])).let_in(
                     &mut id_gen,

src/sql/src/plan/statement/ddl.rs

@@ -5168,6 +5168,7 @@ pub fn unplan_create_cluster(
                 enable_fast_path_plan_insights: _,
                 enable_cast_elimination: _,
                 enable_case_literal_transform: _,
+                enable_simplify_quantified_comparisons: _,
             } = optimizer_feature_overrides;
             // The ones from above that don't occur below are not wired up to cluster features.
             let features_extracted = ClusterFeatureExtracted {

src/sql/src/plan/statement/dml.rs

@@ -650,6 +650,7 @@ impl TryFrom<ExplainPlanOptionExtracted> for ExplainConfig {
                 enable_fast_path_plan_insights: Default::default(),
                 enable_cast_elimination: Default::default(),
                 enable_case_literal_transform: Default::default(),
+                enable_simplify_quantified_comparisons: Default::default(),
             },
         })
     }

src/sql/src/plan/transform_hir.rs

@@ -177,18 +177,20 @@ pub fn split_subquery_predicates(expr: &mut HirRelationExpr) -> Result<(), Recur
 /// M. Elhemali, et al.
 pub fn try_simplify_quantified_comparisons(
     expr: &mut HirRelationExpr,
+    simplify_join_on: bool,
 ) -> Result<(), RecursionLimitError> {
     fn walk_relation(
         expr: &mut HirRelationExpr,
         outers: &[SqlRelationType],
+        simplify_join_on: bool,
     ) -> Result<(), RecursionLimitError> {
         match expr {
             HirRelationExpr::Map { scalars, input } => {
-                walk_relation(input, outers)?;
+                walk_relation(input, outers, simplify_join_on)?;
                 let mut outers = outers.to_vec();
                 outers.insert(0, input.typ(&outers, &NO_PARAMS));
                 for scalar in scalars {
-                    walk_scalar(scalar, &outers, false)?;
+                    walk_scalar(scalar, &outers, false, simplify_join_on)?;
                     let (inner, outers) = outers
                         .split_first_mut()
                         .expect("outers known to have at least one element");
@@ -197,30 +199,42 @@ pub fn try_simplify_quantified_comparisons(
                 }
             }
             HirRelationExpr::Filter { predicates, input } => {
-                walk_relation(input, outers)?;
+                walk_relation(input, outers, simplify_join_on)?;
                 let mut outers = outers.to_vec();
                 outers.insert(0, input.typ(&outers, &NO_PARAMS));
                 for pred in predicates {
-                    walk_scalar(pred, &outers, true)?;
+                    walk_scalar(pred, &outers, true, simplify_join_on)?;
                 }
             }
             HirRelationExpr::CallTable { exprs, .. } => {
                 let mut outers = outers.to_vec();
                 outers.insert(0, SqlRelationType::empty());
                 for scalar in exprs {
-                    walk_scalar(scalar, &outers, false)?;
+                    walk_scalar(scalar, &outers, false, simplify_join_on)?;
                 }
             }
-            HirRelationExpr::Join { left, right, .. } => {
-                walk_relation(left, outers)?;
+            HirRelationExpr::Join {
+                left, right, on, ..
+            } => {
+                walk_relation(left, outers, simplify_join_on)?;
+                let left_type = left.typ(outers, &NO_PARAMS);
                 let mut outers = outers.to_vec();
-                outers.insert(0, left.typ(&outers, &NO_PARAMS));
-                walk_relation(right, &outers)?;
+                outers.insert(0, left_type);
+                walk_relation(right, &outers, simplify_join_on)?;
+                if simplify_join_on {
+                    // Build outers with the full join output type, since the
+                    // ON clause can reference columns from both sides.
+                    let right_type = right.typ(&outers, &NO_PARAMS);
+                    let mut join_columns = outers[0].column_types.clone();
+                    join_columns.extend(right_type.column_types);
+                    outers[0] = SqlRelationType::new(join_columns);
+                    walk_scalar(on, &outers, true, simplify_join_on)?;
+                }
             }
             expr => {
                 #[allow(deprecated)]
                 let _ = expr.visit1_mut(0, &mut |expr, _| -> Result<(), RecursionLimitError> {
-                    walk_relation(expr, outers)
+                    walk_relation(expr, outers, simplify_join_on)
                 });
             }
         }
@@ -231,12 +245,15 @@ pub fn try_simplify_quantified_comparisons(
         expr: &mut HirScalarExpr,
         outers: &[SqlRelationType],
         mut in_filter: bool,
+        simplify_join_on: bool,
     ) -> Result<(), RecursionLimitError> {
         expr.try_visit_mut_pre(&mut |e| {
             match e {
-                HirScalarExpr::Exists(input, _name) => walk_relation(input, outers)?,
+                HirScalarExpr::Exists(input, _name) => {
+                    walk_relation(input, outers, simplify_join_on)?
+                }
                 HirScalarExpr::Select(input, _name) => {
-                    walk_relation(input, outers)?;
+                    walk_relation(input, outers, simplify_join_on)?;
 
                     // We're inside a `(SELECT ...)` subquery. Now let's see if
                     // it has the form `(SELECT <any|all>(...) FROM <input>)`.
@@ -295,7 +312,7 @@ pub fn try_simplify_quantified_comparisons(
         })
     }
 
-    walk_relation(expr, &[])
+    walk_relation(expr, &[], simplify_join_on)
 }
 
 /// An empty parameter type map.

src/sql/src/session/vars/definitions.rs

@@ -2262,6 +2262,12 @@ feature_flags!(
         default: false,
         enable_for_item_parsing: false,
     },
+    {
+        name: enable_simplify_quantified_comparisons,
+        desc: "Allow the optimizer to simplify quantified comparisons in JOIN ON clauses into semi/anti-join EXISTS form during HIR-to-MIR lowering.",
+        default: true,
+        enable_for_item_parsing: false,
+    },
 );
 
 impl From<&super::SystemVars> for OptimizerFeatures {
@@ -2286,6 +2292,7 @@ impl From<&super::SystemVars> for OptimizerFeatures {
             enable_fast_path_plan_insights: vars.enable_fast_path_plan_insights(),
             enable_cast_elimination: vars.enable_cast_elimination(),
             enable_case_literal_transform: vars.enable_case_literal_transform(),
+            enable_simplify_quantified_comparisons: vars.enable_simplify_quantified_comparisons(),
         }
     }
 }
@@ -2327,6 +2334,7 @@ mod tests {
             enable_fast_path_plan_insights,
             enable_cast_elimination,
             enable_case_literal_transform,
+            enable_simplify_quantified_comparisons,
         } = false_features;
 
         let mut vars = SystemVars::new();
@@ -2356,6 +2364,7 @@ mod tests {
         set_var!(enable_fast_path_plan_insights);
         set_var!(enable_cast_elimination);
         set_var!(enable_case_literal_transform);
+        set_var!(enable_simplify_quantified_comparisons);
 
         // Enable for item parsing, then ensure we still get the same optimizer features.
         vars.enable_for_item_parsing();

test/sqllogictest/subquery.slt

@@ -708,3 +708,58 @@ Source materialize.public.y
 Target cluster: quickstart
 
 EOF
+
+# Regression test: try_simplify_quantified_comparisons (= ANY -> EXISTS) must
+# also apply inside JOIN ON clauses, not just WHERE. The decorrelated plan
+# should use Filter/Distinct/Map(true) (the EXISTS semijoin pattern), not
+# Reduce with an any() aggregate (the generic scalar subquery path).
+query T multiline
+EXPLAIN DECORRELATED PLAN WITH(humanized expressions) AS VERBOSE TEXT FOR
+SELECT * FROM x JOIN y ON y.b = ANY(SELECT a FROM x)
+----
+With
+  cte l0 =
+    Project (#0{a}, #1{b})
+      CrossJoin
+        CrossJoin
+          Constant
+            - ()
+          Get materialize.public.x
+        CrossJoin
+          Constant
+            - ()
+          Get materialize.public.y
+  cte l1 =
+    Distinct project=[#1{b}]
+      Get l0
+  cte l2 =
+    Map (true)
+      Distinct project=[#0{b}]
+        Filter (#0{b} = #1{a})
+          CrossJoin
+            Get l1
+            Get materialize.public.x
+Return
+  Project (#0{a}, #1{b})
+    Filter #2
+      Project (#0{a}, #1{b}, #3)
+        Join on=(#1{b} = #2{b})
+          Get l0
+          Union
+            Get l2
+            CrossJoin
+              Project (#0{b})
+                Join on=(#0{b} = #1{b})
+                  Union
+                    Negate
+                      Distinct project=[#0{b}]
+                        Get l2
+                    Distinct project=[#0{b}]
+                      Get l1
+                  Get l1
+              Constant
+                - (false)
+
+Target cluster: quickstart
+
+EOF