jayzhan211 commented on code in PR #25338:
URL: https://github.com/apache/datafusion/pull/25338#discussion_r4062777274


##########
datafusion/optimizer/src/decorrelate.rs:
##########
@@ -184,6 +195,11 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr {
                         .all(|&e| can_pullup_over_aggregation(e));
                 let (mut join_filters, subquery_filters) =
                     find_join_exprs(subquery_filter_exprs)?;
+                for expr in &join_filters {

Review Comment:
   `correlated_filters` comes from `find_join_exprs`, which strips outer refs → 
a subquery column sharing the outer value's qualified name is read as the value 
→ join not null-aware → wrong result (regression vs `main`).
   
   Repro (expected `NULL` for the NULL row; PR gives `false`, `main` gives 
`NULL`; the `NOT IN` filter keeps the NULL row):
   ```sql
   create table o(x int, k int) as values (null, 1), (7, 1), (8, 1);
   create table t(x int, y int not null) as values (1, 7);
   select a.x, a.x in (select a.y from t as a where a.x = b.k) as r
   from o as a, (select 1 as k) as b order by 1;
   select a.x from o as a, (select 1 as k) as b
   where a.x not in (select a.y from t as a where a.x = b.k);
   ```
   
   Fix (tested: repro correct, optimizer tests + 
`subquery_projection`/`null_aware*`/`joins`/`subquery` slt pass): keep the 
outer refs and match per side.
   ```diff
   -                let (mut join_filters, subquery_filters) =
   -                    find_join_exprs(subquery_filter_exprs)?;
   -                for expr in &join_filters {
   -                    if !self.correlated_filters.contains(expr) {
   -                        self.correlated_filters.push(expr.clone());
   +                for expr in &subquery_filter_exprs {
   +                    if expr.contains_outer() && 
!self.correlated_filters.contains(expr) {
   +                        self.correlated_filters.push((*expr).clone());
                        }
                    }
   +                let (mut join_filters, subquery_filters) =
   +                    find_join_exprs(subquery_filter_exprs)?;
   ```
   ```rs
   fn filter_rejects_null(filter: &Expr, key: &Expr, key_is_outer: bool) -> 
bool {
       let is_key = |side: &Expr| {
           let side = strip_casts(side);
           if key_is_outer {
               side.contains_outer()
                   && side.column_refs().is_empty()
                   && &strip_outer_reference(side.clone()) == key
           } else {
               !side.contains_outer() && side == key
           }
       };
       match filter {
           Expr::BinaryExpr(BinaryExpr { left, op, right }) => {
               matches!(
                   op,
                   Operator::Eq
                       | Operator::NotEq
                       | Operator::Lt
                       | Operator::LtEq
                       | Operator::Gt
                       | Operator::GtEq
               ) && (is_key(left) || is_key(right))
           }
           Expr::IsNotNull(expr) => is_key(expr),
           _ => false,
       }
   }
   ```
   Pass `true` for `value_as_written` and `false` for `output_expr` from 
`InValue::may_be_null_in_scope`, and add the repro to `subquery_projection.slt`.



##########
datafusion/optimizer/src/decorrelate_predicate_subquery.rs:
##########
@@ -222,39 +241,89 @@ fn rewrite_inner_subqueries(
     Ok((cur_input, expr_without_subqueries.data))
 }
 
+/// Rewrites an `IN` subquery that gives a value, for example in a SELECT list.
+/// The value follows SQL three-valued logic: TRUE for a match, FALSE for a 
miss
+/// and NULL (UNKNOWN) when the answer depends on a NULL.
+///
+/// There are two paths:
+///
+/// * One mark join. When the mark column is already exact under three-valued
+///   logic (see [`MarkJoin::three_valued_exact`]), the mark column is the
+///   answer and this single join is the full rewrite. This is the usual case.
+/// * Three mark joins. A residual non-equality filter stays on the join in the
+///   other case. The mark column then only tells TRUE from not-TRUE, so the
+///   UNKNOWN cases must be materialized: one more join tells if the subquery
+///   gives a NULL, and one more tells if the subquery gives any row. A `CASE`
+///   expression puts the three marks together. The two extra joins have no
+///   join predicate, so use them only when the first path cannot apply.
 fn in_subquery_value_mark_join(
     left: &LogicalPlan,
     subquery: &LogicalPlan,
     expr: Expr,
     negated: bool,
     alias: &Arc<AliasGenerator>,
 ) -> Result<Option<(LogicalPlan, Expr)>> {
+    // An outer reference in the value belongs to an enclosing subquery. It
+    // cannot be resolved in the joins this builds, and `build_join` would read
+    // it as a constant, so leave the predicate for the enclosing rule.
+    if expr.contains_outer() {
+        return Ok(None);
+    }
+
     let output_expr = subquery
         .head_output_expr()?
         .map_or(plan_err!("single expression required."), Ok)?;
     let in_predicate = Expr::eq(expr.clone(), output_expr.clone());
-    let Some((matched_plan, matched)) =
-        mark_join(left, subquery, Some(&in_predicate), false, alias)?
+    let Some(MarkJoin {
+        plan: matched_plan,
+        mark: matched,
+        three_valued_exact,
+    }) = mark_join(left, subquery, Some(&in_predicate), false, alias)?
     else {
         return Ok(None);
     };
 
-    // SQL IN needs three facts per outer row to distinguish FALSE from 
UNKNOWN.
-    let null_subquery = LogicalPlanBuilder::from(subquery.clone())
-        .filter(output_expr.is_null())?
-        .build()?;
-    let Some((null_plan, subquery_has_null)) =
-        mark_join(&matched_plan, &null_subquery, None, false, alias)?
-    else {
-        return Ok(None);
-    };
-    let Some((final_plan, subquery_non_empty)) =
-        mark_join(&null_plan, subquery, None, false, alias)?
+    // The mark column is the full answer when it is exact. Negation does not
+    // change that, because NOT UNKNOWN is UNKNOWN.
+    if three_valued_exact {
+        return Ok(Some((
+            matched_plan,
+            if negated { not(matched) } else { matched },
+        )));
+    }
+
+    // The value is not matched here, so the answer is UNKNOWN when the
+    // subquery gives a NULL, and also when the value is NULL and the subquery
+    // gives any row at all. One mark join answers both: it keeps a subquery 
row
+    // in the scope of the outer row when the subquery value is NULL, or when
+    // the outer value is NULL and the row is in scope at all.
+    //
+    // For an outer value that is not NULL the mark reads "the subquery gives a
+    // NULL". For an outer value that is NULL it reads "the subquery gives a
+    // row", which is the weaker fact that this case needs, and which the first
+    // reading implies. `IS NULL` is two-valued on both sides, so neither test
+    // adds an UNKNOWN of its own.
+    let unknown_alias = alias.next("__correlated_sq");

Review Comment:
   Commit `50d6e4a3a` (one merged UNKNOWN join) slows the residual fallback 
~1.6x: the NLJ now evaluates `k < k AND (y IS NULL OR x IS NULL)` on every 
pair; before, one NLJ ran bare `k < k` and the other had its right side 
pre-filtered to `y IS NULL`.
   
   `release-nonlto`, 200k x 200k, 3 interleaved runs: main 2.80/3.55/3.78 s · 
`bb043f10e` (parent) 3.05/3.06/3.67 s · PR 4.74/5.14/6.48 s. `EXPLAIN ANALYZE` 
NLJ `elapsed_compute`: main 21.7 s + 2.7 s, PR 44.7 s.
   ```sql
   create table outer_t as select i as id, case when i % 10 = 0 then null else 
i end as x, i % 100 as k from generate_series(1, 200000) t(i);
   create table inner_t as select i as id, case when i % 7 = 0 then null else i 
* 2 end as y, i % 100 as k from generate_series(1, 200000) t(i);
   select id, x in (select y from inner_t where inner_t.k < outer_t.k) as r 
from outer_t;
   ```
   Your Q6 number went the other way, so this is data-dependent. Please drop 
the commit from this PR (the fallback then matches `main`) and revisit it with 
#25336.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to