krishiv_plan/optimizer/
predicate_pushdown.rs

1//! Predicate push-down logical optimizer rule.
2
3use crate::{LogicalPlan, NodeOp};
4
5use super::OptimizerRule;
6
7/// Push `Filter` predicates down into `TableScan` nodes.
8///
9/// Walks the logical plan looking for `Filter` nodes and decomposes each
10/// filter's predicate into AND-conjuncts. Conjuncts that reference only
11/// columns present in one scan's output schema are pushed into that scan
12/// node's `filters` list. If all conjuncts are pushed the `Filter` node is
13/// removed; remaining cross-join conjuncts stay in place.
14///
15/// Two patterns are handled:
16/// - **Filter-above-Scan**: filter's direct input is a scan.
17/// - **Filter-above-Join**: filter sits above a join; each conjunct is tested
18///   against the left and right scan inputs independently and pushed as far
19///   down as it can go. Cross-join predicates (referencing both sides) remain
20///   in the filter.
21pub struct PredicatePushdownRule;
22
23impl OptimizerRule for PredicatePushdownRule {
24    fn name(&self) -> &str {
25        "predicate-pushdown"
26    }
27
28    fn apply(&self, plan: &LogicalPlan) -> Option<LogicalPlan> {
29        let nodes = plan.nodes().to_vec();
30        let id_to_idx: std::collections::HashMap<&str, usize> =
31            nodes.iter().enumerate().map(|(i, n)| (n.id(), i)).collect();
32
33        // Collect pushdown candidates: filter nodes whose input is a scan.
34        struct FilterPushdown {
35            filter_idx: usize,
36            scan_pushes: Vec<(usize, Vec<String>)>,
37            remaining: Vec<String>,
38        }
39
40        let mut pushdowns: Vec<FilterPushdown> = Vec::new();
41
42        for (i, node) in nodes.iter().enumerate() {
43            let predicate = match node.op() {
44                Some(NodeOp::Filter { predicate }) => predicate.clone(),
45                _ => continue,
46            };
47
48            // Collect all scan nodes reachable in one or two hops from this
49            // filter. One hop covers Filter-above-Scan; two hops covers
50            // Filter-above-Join-above-Scan so each side of the join can
51            // independently receive the conjuncts that belong to it.
52            let direct_inputs: Vec<usize> = node
53                .inputs()
54                .iter()
55                .filter_map(|input_id| id_to_idx.get(input_id.as_str()).copied())
56                .collect();
57
58            let mut scan_indices: Vec<usize> = direct_inputs
59                .iter()
60                .copied()
61                .filter(|&idx| {
62                    nodes
63                        .get(idx)
64                        .is_some_and(|n| matches!(n.op(), Some(NodeOp::Scan { .. })))
65                })
66                .collect();
67
68            // Filter-above-Join: descend through join nodes to collect
69            // both left and right scan inputs for per-side pushdown.
70            for join_idx in direct_inputs.iter().copied().filter(|&idx| {
71                nodes.get(idx).is_some_and(|n| {
72                    matches!(
73                        n.op(),
74                        Some(NodeOp::Join {
75                            join_type: crate::JoinType::Inner
76                        })
77                    )
78                })
79            }) {
80                let join_inputs: Vec<String> = nodes
81                    .get(join_idx)
82                    .map(|n| n.inputs().to_vec())
83                    .unwrap_or_default();
84                for child_id in &join_inputs {
85                    if let Some(&child_idx) = id_to_idx.get(child_id.as_str())
86                        && nodes
87                            .get(child_idx)
88                            .is_some_and(|n| matches!(n.op(), Some(NodeOp::Scan { .. })))
89                    {
90                        scan_indices.push(child_idx);
91                    }
92                }
93            }
94            scan_indices.sort_unstable();
95            scan_indices.dedup();
96
97            if scan_indices.is_empty() {
98                continue;
99            }
100
101            // C5: Use sqlparser to split predicate conjuncts properly
102            // instead of naively splitting on the literal string " AND ".
103            let conjuncts = split_predicate_conjuncts(&predicate);
104
105            if conjuncts.is_empty() {
106                continue;
107            }
108
109            let scan_contracts = scan_indices
110                .iter()
111                .filter_map(|&scan_idx| {
112                    let scan_node = nodes.get(scan_idx)?;
113                    let columns = scan_node
114                        .output_schema()
115                        .fields()
116                        .iter()
117                        .map(|field| field.name())
118                        .collect::<Vec<_>>();
119                    let table = match scan_node.op() {
120                        Some(NodeOp::Scan { table, .. }) => table.as_str(),
121                        _ => "",
122                    };
123                    Some((scan_idx, table, columns))
124                })
125                .collect::<Vec<_>>();
126            let mut scan_pushes = std::collections::HashMap::<usize, Vec<String>>::new();
127            let mut remaining = Vec::new();
128
129            for conjunct in conjuncts {
130                let columns = extract_column_refs(&conjunct);
131                let matching_scans = scan_contracts
132                    .iter()
133                    .filter_map(|(scan_idx, table, scan_columns)| {
134                        (!columns.is_empty()
135                            && columns
136                                .iter()
137                                .all(|column| column_belongs_to_scan(column, table, scan_columns)))
138                        .then_some(*scan_idx)
139                    })
140                    .collect::<Vec<_>>();
141                if let [scan_idx] = matching_scans.as_slice() {
142                    scan_pushes.entry(*scan_idx).or_default().push(conjunct);
143                } else {
144                    remaining.push(conjunct);
145                }
146            }
147
148            if !scan_pushes.is_empty() {
149                let mut scan_pushes = scan_pushes.into_iter().collect::<Vec<_>>();
150                scan_pushes.sort_by_key(|(scan_idx, _)| *scan_idx);
151                pushdowns.push(FilterPushdown {
152                    filter_idx: i,
153                    scan_pushes,
154                    remaining,
155                });
156            }
157        }
158
159        if pushdowns.is_empty() {
160            return None;
161        }
162
163        let mut new_nodes = nodes.clone();
164        let mut to_remove: Vec<usize> = Vec::new();
165
166        for pd in &pushdowns {
167            for (scan_idx, pushable) in &pd.scan_pushes {
168                if let Some(node) = new_nodes.get(*scan_idx)
169                    && let Some(NodeOp::Scan { table, filters }) = node.op()
170                {
171                    let table = table.clone();
172                    let mut new_filters = filters.clone();
173                    new_filters.extend(pushable.iter().cloned());
174                    if let Some(n) = new_nodes.get_mut(*scan_idx) {
175                        *n = n.clone().with_op(NodeOp::Scan {
176                            table,
177                            filters: new_filters,
178                        });
179                    }
180                }
181            }
182
183            if pd.remaining.is_empty() {
184                to_remove.push(pd.filter_idx);
185            } else if let Some(n) = new_nodes.get(pd.filter_idx) {
186                let updated = n.clone().with_op(NodeOp::Filter {
187                    predicate: pd.remaining.join(" AND "),
188                });
189                if let Some(slot) = new_nodes.get_mut(pd.filter_idx) {
190                    *slot = updated;
191                }
192            }
193        }
194
195        // Remove filter nodes and rewire downstream node inputs.
196        for &idx in to_remove.iter().rev() {
197            let (filter_id, filter_inputs) = new_nodes
198                .get(idx)
199                .map(|n| (n.id().to_string(), n.inputs().to_vec()))
200                .unwrap_or_default();
201            new_nodes.remove(idx);
202
203            for node in &mut new_nodes {
204                let inputs: Vec<String> = node.inputs().to_vec();
205                if inputs.contains(&filter_id) {
206                    let new_inputs: Vec<String> = inputs
207                        .iter()
208                        .flat_map(|input| {
209                            if input == &filter_id {
210                                filter_inputs.clone()
211                            } else {
212                                vec![input.clone()]
213                            }
214                        })
215                        .collect();
216                    *node = node.clone().with_inputs(new_inputs);
217                }
218            }
219        }
220
221        let mut out = LogicalPlan::new(plan.name(), plan.kind());
222        for node in new_nodes {
223            out.add_node(node);
224        }
225        Some(out)
226    }
227}
228
229/// Extract likely column-name identifiers from a predicate expression string.
230///
231/// Skips string literals and function names, retaining unquoted and quoted
232/// identifier paths such as `column` and `table.column`.
233pub(super) fn extract_column_refs(predicate: &str) -> Vec<String> {
234    const SQL_KEYWORDS: &[&str] = &[
235        "AND", "OR", "NOT", "IN", "IS", "NULL", "TRUE", "FALSE", "WHERE", "SELECT", "FROM", "AS",
236        "ON", "BETWEEN", "LIKE", "EXISTS", "HAVING", "GROUP", "ORDER", "BY", "ASC", "DESC",
237        "LIMIT", "OFFSET", "DISTINCT", "ALL", "ANY", "SOME", "CASE", "WHEN", "THEN", "ELSE", "END",
238        "CAST",
239    ];
240
241    let chars = predicate.char_indices().collect::<Vec<_>>();
242    let mut refs = Vec::new();
243    let mut cursor = 0usize;
244    while cursor < chars.len() {
245        let Some(&(_, ch)) = chars.get(cursor) else {
246            break;
247        };
248        if ch == '\'' {
249            cursor += 1;
250            while cursor < chars.len() {
251                if chars.get(cursor).is_some_and(|(_, c)| *c == '\'') {
252                    if cursor + 1 < chars.len()
253                        && chars.get(cursor + 1).is_some_and(|(_, c)| *c == '\'')
254                    {
255                        cursor += 2;
256                        continue;
257                    }
258                    cursor += 1;
259                    break;
260                }
261                cursor += 1;
262            }
263            continue;
264        }
265        if ch == '"' || ch == '`' {
266            let quote = ch;
267            let start = chars
268                .get(cursor)
269                .map_or(predicate.len(), |(o, _)| *o + ch.len_utf8());
270            cursor += 1;
271            while cursor < chars.len() && chars.get(cursor).is_none_or(|(_, c)| *c != quote) {
272                cursor += 1;
273            }
274            let end = chars
275                .get(cursor)
276                .map_or(predicate.len(), |(offset, _)| *offset);
277            if end > start {
278                refs.push(predicate.get(start..end).unwrap_or("").to_string());
279            }
280            cursor = cursor.saturating_add(1);
281            continue;
282        }
283        if ch.is_ascii_alphabetic() || ch == '_' {
284            let start = chars.get(cursor).map_or(predicate.len(), |(o, _)| *o);
285            cursor += 1;
286            while cursor < chars.len()
287                && chars
288                    .get(cursor)
289                    .is_some_and(|(_, c)| c.is_ascii_alphanumeric() || *c == '_' || *c == '.')
290            {
291                cursor += 1;
292            }
293            let end = chars
294                .get(cursor)
295                .map_or(predicate.len(), |(offset, _)| *offset);
296            let token = predicate.get(start..end).unwrap_or("");
297            let next_non_whitespace = chars
298                .get(cursor..)
299                .unwrap_or(&[])
300                .iter()
301                .find_map(|(_, next)| (!next.is_whitespace()).then_some(*next));
302            if next_non_whitespace != Some('(')
303                && !SQL_KEYWORDS.contains(&token.to_uppercase().as_str())
304                && !refs.iter().any(|existing| existing == token)
305            {
306                refs.push(token.to_string());
307            }
308            continue;
309        }
310        cursor += 1;
311    }
312    refs
313}
314
315/// C5: Split a SQL predicate string into conjuncts using sqlparser for correct
316/// AND splitting.  Respects quoted strings, nested expressions, etc.
317pub(super) fn split_predicate_conjuncts(predicate: &str) -> Vec<String> {
318    use sqlparser::dialect::GenericDialect;
319    use sqlparser::parser::Parser;
320
321    let dialect = GenericDialect {};
322    let expression = predicate
323        .strip_prefix("WHERE ")
324        .or_else(|| predicate.strip_prefix("where "))
325        .unwrap_or(predicate);
326    let statement = format!("SELECT * FROM __krishiv_predicate WHERE {expression}");
327    let Ok(mut stmts) = Parser::parse_sql(&dialect, &statement) else {
328        return Vec::new();
329    };
330    let Some(stmt) = stmts.pop() else {
331        return vec![predicate.to_string()];
332    };
333    // Extract the expression and split on top-level AND.
334    let sqlparser::ast::Statement::Query(query) = stmt else {
335        return vec![predicate.to_string()];
336    };
337    let Some(select_body) = query.body.as_select() else {
338        return vec![predicate.to_string()];
339    };
340    let Some(selection) = &select_body.selection else {
341        return vec![predicate.to_string()];
342    };
343    collect_binary_conjuncts(selection, "AND")
344}
345
346/// Recursively collect top-level conjuncts from a binary expression tree.
347pub(super) fn collect_binary_conjuncts(expr: &sqlparser::ast::Expr, op: &str) -> Vec<String> {
348    match expr {
349        sqlparser::ast::Expr::BinaryOp {
350            left,
351            op: bin_op,
352            right,
353        } if bin_op.to_string().to_uppercase() == op => {
354            let mut left_conjuncts = collect_binary_conjuncts(left, op);
355            let right_conjuncts = collect_binary_conjuncts(right, op);
356            left_conjuncts.extend(right_conjuncts);
357            left_conjuncts
358        }
359        other => vec![other.to_string()],
360    }
361}
362
363/// Check whether `col` (possibly qualified like `"t.id"`) belongs to `scan_table`
364/// with the given column names.  C5: When a column reference has an explicit
365/// qualifier, require an exact case-insensitive table match. Aliases are not
366/// represented in `PlanNode`, so guessing them would permit unsafe pushdown.
367pub(super) fn column_belongs_to_scan(col: &str, scan_table: &str, scan_columns: &[&str]) -> bool {
368    if let Some(dot_pos) = col.rfind('.') {
369        let qualifier = &col[..dot_pos];
370        let unqualified = &col[dot_pos + 1..];
371        if !qualifier.is_empty() {
372            let scan_lower = scan_table.to_ascii_lowercase();
373            let qual_lower = qualifier.to_ascii_lowercase();
374            if qual_lower == scan_lower {
375                return scan_columns.contains(&unqualified);
376            }
377            // Reject qualification that doesn't match this table at all.
378            return false;
379        }
380        return scan_columns.contains(&unqualified);
381    }
382    scan_columns.contains(&col)
383}
krishiv_plan/optimizer/predicate_pushdown.rs

krishiv_plan/optimizer/
predicate_pushdown.rs
