risingwave_frontend/optimizer/plan_node/
predicate_pushdown.rs

1// Copyright 2022 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::{HashMap, HashSet};
16
17use super::*;
18use crate::optimizer::plan_visitor::{ExprCorrelatedIdFinder, ShareParentCounter};
19use crate::optimizer::{
20    ExpressionSimplifyRewriter, LogicalPlanRef as PlanRef, PlanVisitor, ShareId,
21};
22
23/// The trait for predicate pushdown, only logical plan node will use it, though all plan node impl
24/// it.
25pub trait PredicatePushdown {
26    /// Push predicate down for every logical plan node.
27    ///
28    /// There are three kinds of predicates:
29    ///
30    /// 1. those can't be pushed down. We just create a `LogicalFilter` for them above the current
31    ///    `PlanNode`. i.e.,
32    ///
33    ///     ```ignore
34    ///     LogicalFilter::create(self.clone().into(), predicate)
35    ///     ```
36    ///
37    /// 2. those can be merged with current `PlanNode` (e.g., `LogicalJoin`). We just merge
38    ///    the predicates with the `Condition` of it.
39    ///
40    /// 3. those can be pushed down. We pass them to current `PlanNode`'s input.
41    fn predicate_pushdown(
42        &self,
43        predicate: Condition,
44        ctx: &mut PredicatePushdownContext,
45    ) -> PlanRef;
46}
47
48#[inline]
49pub fn gen_filter_and_pushdown<T: PlanTreeNodeUnary<Logical> + LogicalPlanNode>(
50    node: &T,
51    filter_predicate: Condition,
52    pushed_predicate: Condition,
53    ctx: &mut PredicatePushdownContext,
54) -> PlanRef {
55    let new_input = node.input().predicate_pushdown(pushed_predicate, ctx);
56    let new_node = node.clone_with_input(new_input);
57    LogicalFilter::create(new_node.into(), filter_predicate)
58}
59
60fn merge_share_predicates(requirements: Vec<Condition>) -> Condition {
61    let merged = requirements
62        .into_iter()
63        .map(|mut condition| Condition {
64            conjunctions: condition
65                .conjunctions
66                .extract_if(.., |expr| {
67                    // Temporal, impure and correlated predicates must remain above each
68                    // parent and cannot participate in an OR below a share.
69                    let mut finder = ExprCorrelatedIdFinder::default();
70                    finder.visit_expr(expr);
71                    expr.count_nows() == 0 && expr.is_pure() && !finder.has_correlated_input_ref()
72                })
73                .collect(),
74        })
75        .reduce(Condition::or)
76        .expect("a shared plan must have at least one parent predicate");
77
78    let mut rewriter = ExpressionSimplifyRewriter {};
79    merged
80        .conjunctions
81        .into_iter()
82        .fold(Condition::true_cond(), |condition, expr| {
83            condition.and(Condition::with_expr(rewriter.rewrite_cond(expr)))
84        })
85}
86
87#[derive(Debug, Clone, Copy, PartialEq, Eq)]
88enum PredicatePushdownPhase {
89    Idle,
90    Collect,
91    Rewrite,
92}
93
94#[derive(Debug, Clone)]
95struct SharePredicatePushdown {
96    original_input: PlanRef,
97    predicate: Condition,
98}
99
100#[derive(Debug, Clone)]
101pub struct PredicatePushdownContext {
102    pending_predicates: HashMap<ShareId, Vec<Condition>>,
103    collected_shares: HashMap<ShareId, SharePredicatePushdown>,
104    share_parent_counter: ShareParentCounter,
105    rebuilt_shares: HashSet<ShareId>,
106    skipped_shares: HashSet<ShareId>,
107    phase: PredicatePushdownPhase,
108}
109
110impl PredicatePushdownContext {
111    pub fn new(root: PlanRef) -> Self {
112        let mut share_parent_counter = ShareParentCounter::default();
113        share_parent_counter.visit(root);
114        Self {
115            pending_predicates: HashMap::new(),
116            collected_shares: HashMap::new(),
117            share_parent_counter,
118            rebuilt_shares: HashSet::new(),
119            skipped_shares: HashSet::new(),
120            phase: PredicatePushdownPhase::Idle,
121        }
122    }
123
124    pub(in crate::optimizer) fn is_running(&self) -> bool {
125        self.phase != PredicatePushdownPhase::Idle
126    }
127
128    pub(in crate::optimizer) fn is_collecting(&self) -> bool {
129        self.phase == PredicatePushdownPhase::Collect
130    }
131
132    pub(in crate::optimizer) fn get_parent_num(&self, share: &LogicalShare) -> usize {
133        self.share_parent_counter.get_parent_num(share)
134    }
135
136    pub(in crate::optimizer) fn add_predicate(
137        &mut self,
138        share: &LogicalShare,
139        predicate: Condition,
140    ) -> Option<Condition> {
141        let share_id = share.share_id();
142        let parent_num = self.share_parent_counter.get_parent_num_by_id(share_id);
143        let pending = self.pending_predicates.entry(share_id).or_default();
144        pending.push(predicate);
145        assert!(
146            pending.len() <= parent_num,
147            "share {share_id:?} received more predicates than parents"
148        );
149        if pending.len() != parent_num {
150            return None;
151        }
152
153        let merged_predicate = merge_share_predicates(
154            self.pending_predicates
155                .remove(&share_id)
156                .expect("share predicates must exist"),
157        );
158        self.collected_shares
159            .try_insert(
160                share_id,
161                SharePredicatePushdown {
162                    original_input: share.input(),
163                    predicate: merged_predicate.clone(),
164                },
165            )
166            .expect("predicates must be merged once per share");
167        Some(merged_predicate)
168    }
169
170    /// Rebuilds one shared definition on first use. Nested shares recursively rebuild first, so
171    /// the call stack provides the child-before-parent order without a separate dependency graph.
172    pub(in crate::optimizer) fn ensure_share_rebuilt(&mut self, share: &LogicalShare) {
173        let share_id = share.share_id();
174        if self.rebuilt_shares.contains(&share_id) {
175            return;
176        }
177        assert_eq!(self.phase, PredicatePushdownPhase::Rewrite);
178
179        let Some(SharePredicatePushdown {
180            original_input,
181            predicate,
182        }) = self.collected_shares.remove(&share_id)
183        else {
184            // The share was skipped during collection (see `run`): it never received a
185            // predicate from every parent, so it must keep its original definition.
186            // Parents re-create their filters above the share during the rewrite walk.
187            assert!(
188                self.skipped_shares.contains(&share_id),
189                "share {share_id:?} has no collected predicates"
190            );
191            return;
192        };
193        let rebuilt_input = original_input.predicate_pushdown(predicate, self);
194        share.ctx().update_logical_share(share_id, rebuilt_input);
195        assert!(self.rebuilt_shares.insert(share_id));
196    }
197
198    pub(in crate::optimizer) fn run(&mut self, root: PlanRef, predicate: Condition) -> PlanRef {
199        self.phase = PredicatePushdownPhase::Collect;
200        let collected = root.predicate_pushdown_inner(predicate.clone(), self);
201        // `ShareParentCounter` counts parents via the visitor walk, but the transformation
202        // walk is not guaranteed to reach a share from every parent: some `PredicatePushdown`
203        // impls (e.g. `LogicalOverWindow`, `LogicalGapFill`) legitimately stop recursing into
204        // their inputs. A predicate merged from a subset of parents must not be pushed below
205        // the share — it would drop rows required by the parents that never contributed — so
206        // such shares are skipped and keep their original definition, like on the base plan.
207        //
208        // Every non-leaf `PredicatePushdown` impl is expected to recurse into all of its
209        // inputs, pushing `Condition::true_cond()` when nothing can be pushed (like
210        // `LogicalExpand`), so this should never happen: fail loudly in debug builds and
211        // degrade to skipping the share in release builds.
212        debug_assert!(
213            self.pending_predicates.is_empty(),
214            "shares {:?} did not receive predicates from every parent; some \
215             `PredicatePushdown` impl likely failed to recurse into an input — push \
216             `Condition::true_cond()` instead of returning early",
217            self.pending_predicates.keys().collect::<Vec<_>>(),
218        );
219        self.skipped_shares.extend(
220            self.pending_predicates
221                .drain()
222                .map(|(share_id, _)| share_id),
223        );
224        if self.collected_shares.is_empty() {
225            self.phase = PredicatePushdownPhase::Idle;
226            return collected;
227        }
228
229        self.phase = PredicatePushdownPhase::Rewrite;
230        let result = root.predicate_pushdown_inner(predicate, self);
231        self.phase = PredicatePushdownPhase::Idle;
232        result
233    }
234}
risingwave_frontend/optimizer/plan_node/predicate_pushdown.rs

risingwave_frontend/optimizer/plan_node/
predicate_pushdown.rs
