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 create a `LogicalFilter` for them above the current
31    ///    `PlanNode`. A leaf node can do this directly:
32    ///
33    ///     ```ignore
34    ///     LogicalFilter::create(self.clone().into(), predicate)
35    ///     ```
36    ///
37    ///    A non-leaf node must still recurse into every input with [`Condition::true_cond`] so each
38    ///    parent of a [`LogicalShare`] contributes to predicate collection. Unary nodes should use
39    ///    [`gen_filter_and_pushdown`]:
40    ///
41    ///     ```ignore
42    ///     gen_filter_and_pushdown(self, predicate, Condition::true_cond(), ctx)
43    ///     ```
44    ///
45    /// 2. those can be merged with current `PlanNode` (e.g., `LogicalJoin`). We just merge
46    ///    the predicates with the `Condition` of it.
47    ///
48    /// 3. those can be pushed down. We pass them to current `PlanNode`'s input.
49    fn predicate_pushdown(
50        &self,
51        predicate: Condition,
52        ctx: &mut PredicatePushdownContext,
53    ) -> PlanRef;
54}
55
56#[inline]
57pub fn gen_filter_and_pushdown<T: PlanTreeNodeUnary<Logical> + LogicalPlanNode>(
58    node: &T,
59    filter_predicate: Condition,
60    pushed_predicate: Condition,
61    ctx: &mut PredicatePushdownContext,
62) -> PlanRef {
63    let new_input = node.input().predicate_pushdown(pushed_predicate, ctx);
64    let new_node = node.clone_with_input(new_input);
65    LogicalFilter::create(new_node.into(), filter_predicate)
66}
67
68fn merge_share_predicates(requirements: Vec<Condition>) -> Condition {
69    let merged = requirements
70        .into_iter()
71        .map(|mut condition| Condition {
72            conjunctions: condition
73                .conjunctions
74                .extract_if(.., |expr| {
75                    // Temporal, impure and correlated predicates must remain above each
76                    // parent and cannot participate in an OR below a share.
77                    let mut finder = ExprCorrelatedIdFinder::default();
78                    finder.visit_expr(expr);
79                    expr.count_nows() == 0 && expr.is_pure() && !finder.has_correlated_input_ref()
80                })
81                .collect(),
82        })
83        .reduce(Condition::or)
84        .expect("a shared plan must have at least one parent predicate");
85
86    let mut rewriter = ExpressionSimplifyRewriter {};
87    merged
88        .conjunctions
89        .into_iter()
90        .fold(Condition::true_cond(), |condition, expr| {
91            condition.and(Condition::with_expr(rewriter.rewrite_cond(expr)))
92        })
93}
94
95#[derive(Debug, Clone, Copy, PartialEq, Eq)]
96enum PredicatePushdownPhase {
97    Idle,
98    Collect,
99    Rewrite,
100}
101
102#[derive(Debug, Clone)]
103struct SharePredicatePushdown {
104    original_input: PlanRef,
105    predicate: Condition,
106}
107
108#[derive(Debug, Clone)]
109pub struct PredicatePushdownContext {
110    pending_predicates: HashMap<ShareId, Vec<Condition>>,
111    collected_shares: HashMap<ShareId, SharePredicatePushdown>,
112    share_parent_counter: ShareParentCounter,
113    rebuilt_shares: HashSet<ShareId>,
114    skipped_shares: HashSet<ShareId>,
115    phase: PredicatePushdownPhase,
116}
117
118impl PredicatePushdownContext {
119    pub fn new(root: PlanRef) -> Self {
120        let mut share_parent_counter = ShareParentCounter::default();
121        share_parent_counter.visit(root);
122        Self {
123            pending_predicates: HashMap::new(),
124            collected_shares: HashMap::new(),
125            share_parent_counter,
126            rebuilt_shares: HashSet::new(),
127            skipped_shares: HashSet::new(),
128            phase: PredicatePushdownPhase::Idle,
129        }
130    }
131
132    pub(in crate::optimizer) fn is_running(&self) -> bool {
133        self.phase != PredicatePushdownPhase::Idle
134    }
135
136    pub(in crate::optimizer) fn is_collecting(&self) -> bool {
137        self.phase == PredicatePushdownPhase::Collect
138    }
139
140    pub(in crate::optimizer) fn get_parent_num(&self, share: &LogicalShare) -> usize {
141        self.share_parent_counter.get_parent_num(share)
142    }
143
144    pub(in crate::optimizer) fn add_predicate(
145        &mut self,
146        share: &LogicalShare,
147        predicate: Condition,
148    ) -> Option<Condition> {
149        let share_id = share.share_id();
150        let parent_num = self.share_parent_counter.get_parent_num_by_id(share_id);
151        let pending = self.pending_predicates.entry(share_id).or_default();
152        pending.push(predicate);
153        assert!(
154            pending.len() <= parent_num,
155            "share {share_id:?} received more predicates than parents"
156        );
157        if pending.len() != parent_num {
158            return None;
159        }
160
161        let merged_predicate = merge_share_predicates(
162            self.pending_predicates
163                .remove(&share_id)
164                .expect("share predicates must exist"),
165        );
166        self.collected_shares
167            .try_insert(
168                share_id,
169                SharePredicatePushdown {
170                    original_input: share.input(),
171                    predicate: merged_predicate.clone(),
172                },
173            )
174            .expect("predicates must be merged once per share");
175        Some(merged_predicate)
176    }
177
178    /// Rebuilds one shared definition on first use. Nested shares recursively rebuild first, so
179    /// the call stack provides the child-before-parent order without a separate dependency graph.
180    pub(in crate::optimizer) fn ensure_share_rebuilt(&mut self, share: &LogicalShare) {
181        let share_id = share.share_id();
182        if self.rebuilt_shares.contains(&share_id) {
183            return;
184        }
185        assert_eq!(self.phase, PredicatePushdownPhase::Rewrite);
186
187        let Some(SharePredicatePushdown {
188            original_input,
189            predicate,
190        }) = self.collected_shares.remove(&share_id)
191        else {
192            // The share was skipped during collection (see `run`): it never received a
193            // predicate from every parent, so it must keep its original definition.
194            // Parents re-create their filters above the share during the rewrite walk.
195            assert!(
196                self.skipped_shares.contains(&share_id),
197                "share {share_id:?} has no collected predicates"
198            );
199            return;
200        };
201        let rebuilt_input = original_input.predicate_pushdown(predicate, self);
202        share.ctx().update_logical_share(share_id, rebuilt_input);
203        assert!(self.rebuilt_shares.insert(share_id));
204    }
205
206    pub(in crate::optimizer) fn run(&mut self, root: PlanRef, predicate: Condition) -> PlanRef {
207        self.phase = PredicatePushdownPhase::Collect;
208        let collected = root.predicate_pushdown_inner(predicate.clone(), self);
209        // `ShareParentCounter` counts parents via the visitor walk, so every non-leaf
210        // `PredicatePushdown` impl must recurse into all of its inputs, pushing
211        // `Condition::true_cond()` when nothing can be pushed. Otherwise a share may receive
212        // predicates from only a subset of its parents. Such a partial predicate must not be
213        // pushed below the share because it could drop rows required by parents that did not
214        // contribute. This should never happen: fail loudly in debug builds and degrade to
215        // skipping the share (keeping its original definition) in release builds.
216        debug_assert!(
217            self.pending_predicates.is_empty(),
218            "shares {:?} did not receive predicates from every parent; some \
219             `PredicatePushdown` impl likely failed to recurse into an input — push \
220             `Condition::true_cond()` instead of returning early",
221            self.pending_predicates.keys().collect::<Vec<_>>(),
222        );
223        self.skipped_shares.extend(
224            self.pending_predicates
225                .drain()
226                .map(|(share_id, _)| share_id),
227        );
228        if self.collected_shares.is_empty() {
229            self.phase = PredicatePushdownPhase::Idle;
230            return collected;
231        }
232
233        self.phase = PredicatePushdownPhase::Rewrite;
234        let result = root.predicate_pushdown_inner(predicate, self);
235        self.phase = PredicatePushdownPhase::Idle;
236        result
237    }
238}
risingwave_frontend/optimizer/plan_node/predicate_pushdown.rs

risingwave_frontend/optimizer/plan_node/
predicate_pushdown.rs
