risingwave_frontend/optimizer/plan_node/
predicate_pushdown.rs

1// Copyright 2025 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;
16
17use super::*;
18use crate::optimizer::plan_visitor::ShareParentCounter;
19use crate::optimizer::{LogicalPlanRef as PlanRef, PlanVisitor};
20
21/// The trait for predicate pushdown, only logical plan node will use it, though all plan node impl
22/// it.
23pub trait PredicatePushdown {
24    /// Push predicate down for every logical plan node.
25    ///
26    /// There are three kinds of predicates:
27    ///
28    /// 1. those can't be pushed down. We just create a `LogicalFilter` for them above the current
29    ///    `PlanNode`. i.e.,
30    ///
31    ///     ```ignore
32    ///     LogicalFilter::create(self.clone().into(), predicate)
33    ///     ```
34    ///
35    /// 2. those can be merged with current `PlanNode` (e.g., `LogicalJoin`). We just merge
36    ///    the predicates with the `Condition` of it.
37    ///
38    /// 3. those can be pushed down. We pass them to current `PlanNode`'s input.
39    fn predicate_pushdown(
40        &self,
41        predicate: Condition,
42        ctx: &mut PredicatePushdownContext,
43    ) -> PlanRef;
44}
45
46#[inline]
47pub fn gen_filter_and_pushdown<T: PlanTreeNodeUnary<Logical> + LogicalPlanNode>(
48    node: &T,
49    filter_predicate: Condition,
50    pushed_predicate: Condition,
51    ctx: &mut PredicatePushdownContext,
52) -> PlanRef {
53    let new_input = node.input().predicate_pushdown(pushed_predicate, ctx);
54    let new_node = node.clone_with_input(new_input);
55    LogicalFilter::create(new_node.into(), filter_predicate)
56}
57
58#[derive(Debug, Clone)]
59pub struct PredicatePushdownContext {
60    share_predicate_map: HashMap<PlanNodeId, Vec<Condition>>,
61    share_parent_counter: ShareParentCounter,
62}
63
64impl PredicatePushdownContext {
65    pub fn new(root: PlanRef) -> Self {
66        let mut share_parent_counter = ShareParentCounter::default();
67        share_parent_counter.visit(root);
68        Self {
69            share_predicate_map: Default::default(),
70            share_parent_counter,
71        }
72    }
73
74    pub fn get_parent_num(&self, share: &LogicalShare) -> usize {
75        self.share_parent_counter.get_parent_num(share)
76    }
77
78    pub fn add_predicate(&mut self, plan_node_id: PlanNodeId, predicate: Condition) -> usize {
79        self.share_predicate_map
80            .entry(plan_node_id)
81            .and_modify(|e| e.push(predicate.clone()))
82            .or_insert_with(|| vec![predicate])
83            .len()
84    }
85
86    pub fn take_predicate(&mut self, plan_node_id: PlanNodeId) -> Option<Vec<Condition>> {
87        self.share_predicate_map.remove(&plan_node_id)
88    }
89}