laminar_sql/planner/predicate_split/
mod.rs

1//! Predicate splitting and pushdown for lookup joins.
2//!
3//! This module classifies WHERE/ON predicates in lookup join queries and
4//! splits them into pushdown vs local evaluation categories. It implements
5//! a DataFusion optimizer rule (`PredicateSplitterRule`) that absorbs
6//! filter nodes above `LookupJoinNode` and assigns each predicate to the
7//! correct execution site.
8//!
9//! ## Key Safety Rules
10//!
11//! - **H10 (LEFT JOIN safety):** WHERE-clause predicates on lookup-only
12//!   columns above a `LeftOuter` join must NOT be pushed down — doing so
13//!   changes the semantics by filtering out NULL-extended rows.
14//! - **C7 (qualified columns):** When aliases are present, `col.relation`
15//!   is checked first for unambiguous resolution before falling back to
16//!   unqualified column name matching.
17//! - **`NotEq`** predicates are classified normally but are never pushed
18//!   down (they cannot use equality indexes on the source).
19
20#[allow(clippy::disallowed_types)] // cold path: query planning
21use std::collections::{HashMap, HashSet};
22use std::sync::Arc;
23
24use datafusion::logical_expr::logical_plan::LogicalPlan;
25use datafusion::logical_expr::{
26    BinaryExpr, Expr, Extension, Filter, Operator as DfOperator, UserDefinedLogicalNodeCore,
27};
28use datafusion_common::tree_node::Transformed;
29use datafusion_common::Result;
30use datafusion_optimizer::optimizer::{ApplyOrder, OptimizerConfig, OptimizerRule};
31
32use crate::datafusion::lookup_join::{LookupJoinNode, LookupJoinType};
33
34// ---------------------------------------------------------------------------
35// Predicate Classification
36// ---------------------------------------------------------------------------
37
38/// Classification of a predicate based on which side(s) it references.
39#[derive(Debug, Clone, Copy, PartialEq, Eq)]
40pub enum PredicateClass {
41    /// References only lookup table columns — candidate for pushdown.
42    LookupOnly,
43    /// References only stream columns — evaluate locally.
44    StreamOnly,
45    /// References columns from both sides — evaluate locally.
46    CrossReference,
47    /// References no columns (constant expression) — evaluate locally.
48    Constant,
49}
50
51/// Classifies predicates based on column membership.
52///
53/// Uses both unqualified column names and qualified `"alias.col"` names
54/// for resolution (audit C7). When a column has a relation qualifier,
55/// the qualified form is checked first.
56#[derive(Debug)]
57pub struct PredicateClassifier {
58    /// Unqualified lookup column names.
59    lookup_columns: HashSet<String>,
60    /// Unqualified stream column names.
61    stream_columns: HashSet<String>,
62    /// Qualified lookup names: `"alias.col"` or `"table.col"`.
63    lookup_qualified: HashSet<String>,
64    /// Qualified stream names: `"alias.col"` or `"table.col"`.
65    stream_qualified: HashSet<String>,
66}
67
68impl PredicateClassifier {
69    /// Creates a new classifier from column sets.
70    ///
71    /// `lookup_alias` / `stream_alias` are the SQL aliases (e.g., `c` for
72    /// `customers c`). When provided, qualified lookups like `c.name` can
73    /// be resolved unambiguously.
74    #[must_use]
75    pub fn new(
76        lookup_columns: HashSet<String>,
77        stream_columns: HashSet<String>,
78        lookup_alias: Option<&str>,
79        stream_alias: Option<&str>,
80    ) -> Self {
81        let mut lookup_qualified = HashSet::new();
82        let mut stream_qualified = HashSet::new();
83
84        if let Some(alias) = lookup_alias {
85            for col in &lookup_columns {
86                lookup_qualified.insert(format!("{alias}.{col}"));
87            }
88        }
89        if let Some(alias) = stream_alias {
90            for col in &stream_columns {
91                stream_qualified.insert(format!("{alias}.{col}"));
92            }
93        }
94
95        Self {
96            lookup_columns,
97            stream_columns,
98            lookup_qualified,
99            stream_qualified,
100        }
101    }
102
103    /// Classify a predicate expression.
104    #[must_use]
105    pub fn classify(&self, expr: &Expr) -> PredicateClass {
106        let mut has_lookup = false;
107        let mut has_stream = false;
108        self.walk_columns(expr, &mut has_lookup, &mut has_stream);
109
110        match (has_lookup, has_stream) {
111            (true, false) => PredicateClass::LookupOnly,
112            (false, true) => PredicateClass::StreamOnly,
113            (true, true) => PredicateClass::CrossReference,
114            (false, false) => PredicateClass::Constant,
115        }
116    }
117
118    /// Recursively walk an expression to find column references.
119    fn walk_columns(&self, expr: &Expr, has_lookup: &mut bool, has_stream: &mut bool) {
120        match expr {
121            Expr::Column(col) => {
122                // C7: check qualified form first
123                if let Some(relation) = &col.relation {
124                    let qualified = format!("{}.{}", relation, col.name);
125                    if self.lookup_qualified.contains(&qualified) {
126                        *has_lookup = true;
127                        return;
128                    }
129                    if self.stream_qualified.contains(&qualified) {
130                        *has_stream = true;
131                        return;
132                    }
133                }
134                // Fall back to unqualified
135                if self.lookup_columns.contains(&col.name) {
136                    *has_lookup = true;
137                }
138                if self.stream_columns.contains(&col.name) {
139                    *has_stream = true;
140                }
141            }
142            Expr::BinaryExpr(BinaryExpr { left, right, .. }) => {
143                self.walk_columns(left, has_lookup, has_stream);
144                self.walk_columns(right, has_lookup, has_stream);
145            }
146            Expr::Not(inner)
147            | Expr::IsNull(inner)
148            | Expr::IsNotNull(inner)
149            | Expr::Negative(inner)
150            | Expr::Cast(datafusion::logical_expr::Cast { expr: inner, .. })
151            | Expr::TryCast(datafusion::logical_expr::TryCast { expr: inner, .. }) => {
152                self.walk_columns(inner, has_lookup, has_stream);
153            }
154            Expr::Between(between) => {
155                self.walk_columns(&between.expr, has_lookup, has_stream);
156                self.walk_columns(&between.low, has_lookup, has_stream);
157                self.walk_columns(&between.high, has_lookup, has_stream);
158            }
159            Expr::InList(in_list) => {
160                self.walk_columns(&in_list.expr, has_lookup, has_stream);
161                for item in &in_list.list {
162                    self.walk_columns(item, has_lookup, has_stream);
163                }
164            }
165            Expr::ScalarFunction(func) => {
166                for arg in &func.args {
167                    self.walk_columns(arg, has_lookup, has_stream);
168                }
169            }
170            Expr::Like(like) => {
171                self.walk_columns(&like.expr, has_lookup, has_stream);
172                self.walk_columns(&like.pattern, has_lookup, has_stream);
173            }
174            Expr::Case(case) => {
175                if let Some(operand) = &case.expr {
176                    self.walk_columns(operand, has_lookup, has_stream);
177                }
178                for (when, then) in &case.when_then_expr {
179                    self.walk_columns(when, has_lookup, has_stream);
180                    self.walk_columns(then, has_lookup, has_stream);
181                }
182                if let Some(else_expr) = &case.else_expr {
183                    self.walk_columns(else_expr, has_lookup, has_stream);
184                }
185            }
186            // Literals, placeholders — no columns
187            Expr::Literal(..) | Expr::Placeholder(_) => {}
188            // Catch-all: conservative — mark both sides
189            _ => {
190                *has_lookup = true;
191                *has_stream = true;
192            }
193        }
194    }
195}
196
197// ---------------------------------------------------------------------------
198// Source Capabilities
199// ---------------------------------------------------------------------------
200
201/// Mode describing how far predicates can be pushed to a source.
202#[derive(Debug, Clone, Copy, PartialEq, Eq)]
203pub enum PlanPushdownMode {
204    /// Full predicate pushdown (eq, range, in, null checks).
205    Full,
206    /// Only key equality predicates.
207    KeyOnly,
208    /// No pushdown at all.
209    None,
210}
211
212/// Describes a source's pushdown capabilities for the optimizer.
213#[derive(Debug, Clone)]
214pub struct PlanSourceCapabilities {
215    /// Overall pushdown mode.
216    pub pushdown_mode: PlanPushdownMode,
217    /// Columns that support equality pushdown.
218    pub eq_columns: HashSet<String>,
219    /// Columns that support range pushdown.
220    pub range_columns: HashSet<String>,
221    /// Columns that support IN-list pushdown.
222    pub in_columns: HashSet<String>,
223    /// Whether the source supports IS NULL / IS NOT NULL checks.
224    pub supports_null_check: bool,
225}
226
227impl Default for PlanSourceCapabilities {
228    fn default() -> Self {
229        Self {
230            pushdown_mode: PlanPushdownMode::None,
231            eq_columns: HashSet::new(),
232            range_columns: HashSet::new(),
233            in_columns: HashSet::new(),
234            supports_null_check: false,
235        }
236    }
237}
238
239/// Registry mapping lookup table names to their source capabilities.
240#[derive(Debug, Default)]
241pub struct SourceCapabilitiesRegistry {
242    capabilities: HashMap<String, PlanSourceCapabilities>,
243}
244
245impl SourceCapabilitiesRegistry {
246    /// Register capabilities for a lookup table.
247    pub fn register(&mut self, table_name: String, caps: PlanSourceCapabilities) {
248        self.capabilities.insert(table_name, caps);
249    }
250
251    /// Get capabilities for a lookup table.
252    #[must_use]
253    pub fn get(&self, table_name: &str) -> Option<&PlanSourceCapabilities> {
254        self.capabilities.get(table_name)
255    }
256}
257
258// ---------------------------------------------------------------------------
259// Conjunction Splitting
260// ---------------------------------------------------------------------------
261
262/// Splits a conjunction (AND chain) into individual predicates.
263///
264/// `A AND B AND C` → `[A, B, C]`.
265/// OR expressions and non-AND binary expressions are kept as single items.
266#[must_use]
267pub fn split_conjunction(expr: &Expr) -> Vec<Expr> {
268    match expr {
269        Expr::BinaryExpr(BinaryExpr {
270            left,
271            op: DfOperator::And,
272            right,
273        }) => {
274            let mut parts = split_conjunction(left);
275            parts.extend(split_conjunction(right));
276            parts
277        }
278        other => vec![other.clone()],
279    }
280}
281
282// ---------------------------------------------------------------------------
283// Optimizer Rule
284// ---------------------------------------------------------------------------
285
286/// DataFusion optimizer rule that splits predicates for lookup joins.
287///
288/// Runs `TopDown` to catch `Filter` nodes above `LookupJoinNode` first.
289///
290/// Two cases:
291/// 1. **Filter above LookupJoinNode** — absorb the filter, classify
292///    each conjunct, and assign to pushdown or local.
293/// 2. **Direct LookupJoinNode** — re-classify existing pushdown predicates
294///    (e.g., after a previous pass added them).
295#[derive(Debug)]
296pub struct PredicateSplitterRule {
297    /// Per-table source capabilities.
298    capabilities: SourceCapabilitiesRegistry,
299}
300
301impl PredicateSplitterRule {
302    /// Creates a new rule with the given capabilities registry.
303    #[must_use]
304    pub fn new(capabilities: SourceCapabilitiesRegistry) -> Self {
305        Self { capabilities }
306    }
307
308    /// Split predicates for a `LookupJoinNode`, given a list of predicates
309    /// that come from an absorbed `Filter` (if any) plus the node's
310    /// existing predicates.
311    fn split_for_node(
312        &self,
313        node: &LookupJoinNode,
314        filter_predicates: &[Expr],
315    ) -> (Vec<Expr>, Vec<Expr>) {
316        // Build column sets from schemas
317        let lookup_columns: HashSet<String> = node
318            .lookup_schema()
319            .fields()
320            .iter()
321            .map(|f| f.name().clone())
322            .collect();
323
324        let input_schema = node.inputs()[0].schema();
325        let stream_columns: HashSet<String> = input_schema
326            .fields()
327            .iter()
328            .map(|f| f.name().clone())
329            .collect();
330
331        let classifier = PredicateClassifier::new(
332            lookup_columns,
333            stream_columns,
334            node.lookup_alias(),
335            node.stream_alias(),
336        );
337
338        let caps = self.capabilities.get(node.lookup_table_name());
339        let pushdown_disabled = caps.is_none_or(|c| c.pushdown_mode == PlanPushdownMode::None);
340
341        let is_left_outer = node.join_type() == LookupJoinType::LeftOuter;
342
343        let mut pushdown = Vec::new();
344        let mut local = Vec::new();
345
346        // Include existing predicates from the node
347        let all_predicates = node
348            .pushdown_predicates()
349            .iter()
350            .chain(node.local_predicates().iter())
351            .chain(filter_predicates.iter())
352            .cloned();
353
354        for pred in all_predicates {
355            let class = classifier.classify(&pred);
356
357            // NotEq predicates never push down
358            let has_not_eq = contains_not_eq(&pred);
359
360            match class {
361                PredicateClass::LookupOnly => {
362                    // H10: LEFT OUTER WHERE-clause lookup-only preds stay local
363                    if is_left_outer || pushdown_disabled || has_not_eq {
364                        local.push(pred);
365                    } else {
366                        pushdown.push(pred);
367                    }
368                }
369                PredicateClass::StreamOnly
370                | PredicateClass::CrossReference
371                | PredicateClass::Constant => {
372                    local.push(pred);
373                }
374            }
375        }
376
377        (pushdown, local)
378    }
379}
380
381/// Check if an expression contains a `NotEq` operator.
382fn contains_not_eq(expr: &Expr) -> bool {
383    match expr {
384        Expr::BinaryExpr(BinaryExpr { left, op, right }) => {
385            *op == DfOperator::NotEq || contains_not_eq(left) || contains_not_eq(right)
386        }
387        Expr::Not(inner) => contains_not_eq(inner),
388        _ => false,
389    }
390}
391
392impl OptimizerRule for PredicateSplitterRule {
393    fn name(&self) -> &'static str {
394        "predicate_splitter"
395    }
396
397    fn apply_order(&self) -> Option<ApplyOrder> {
398        Some(ApplyOrder::TopDown)
399    }
400
401    fn rewrite(
402        &self,
403        plan: LogicalPlan,
404        _config: &dyn OptimizerConfig,
405    ) -> Result<Transformed<LogicalPlan>> {
406        // Case 1: Filter above a LookupJoinNode
407        if let LogicalPlan::Filter(Filter {
408            predicate, input, ..
409        }) = &plan
410        {
411            if let LogicalPlan::Extension(ext) = input.as_ref() {
412                if let Some(node) = ext.node.as_any().downcast_ref::<LookupJoinNode>() {
413                    let filter_preds = split_conjunction(predicate);
414                    let (pushdown, local) = self.split_for_node(node, &filter_preds);
415
416                    let inputs = node.inputs();
417                    let rebuilt = LookupJoinNode::new(
418                        inputs[0].clone(),
419                        node.lookup_table_name().to_string(),
420                        node.lookup_schema().clone(),
421                        node.join_keys().to_vec(),
422                        node.join_type(),
423                        pushdown,
424                        node.required_lookup_columns().clone(),
425                        UserDefinedLogicalNodeCore::schema(node).clone(),
426                        node.metadata().clone(),
427                    )
428                    .with_local_predicates(local)
429                    .with_aliases(
430                        node.lookup_alias().map(String::from),
431                        node.stream_alias().map(String::from),
432                    );
433
434                    return Ok(Transformed::yes(LogicalPlan::Extension(Extension {
435                        node: Arc::new(rebuilt),
436                    })));
437                }
438            }
439        }
440
441        // Case 2: Direct LookupJoinNode (re-classify existing predicates)
442        if let LogicalPlan::Extension(ext) = &plan {
443            if let Some(node) = ext.node.as_any().downcast_ref::<LookupJoinNode>() {
444                // Only re-classify if there are predicates to work with
445                if !node.pushdown_predicates().is_empty() || !node.local_predicates().is_empty() {
446                    let (pushdown, local) = self.split_for_node(node, &[]);
447                    let inputs = node.inputs();
448                    let rebuilt = LookupJoinNode::new(
449                        inputs[0].clone(),
450                        node.lookup_table_name().to_string(),
451                        node.lookup_schema().clone(),
452                        node.join_keys().to_vec(),
453                        node.join_type(),
454                        pushdown,
455                        node.required_lookup_columns().clone(),
456                        UserDefinedLogicalNodeCore::schema(node).clone(),
457                        node.metadata().clone(),
458                    )
459                    .with_local_predicates(local)
460                    .with_aliases(
461                        node.lookup_alias().map(String::from),
462                        node.stream_alias().map(String::from),
463                    );
464
465                    return Ok(Transformed::yes(LogicalPlan::Extension(Extension {
466                        node: Arc::new(rebuilt),
467                    })));
468                }
469            }
470        }
471
472        Ok(Transformed::no(plan))
473    }
474}
475
476#[cfg(test)]
477mod tests;
laminar_sql/planner/predicate_split/mod.rs

laminar_sql/planner/predicate_split/
mod.rs
