risingwave_frontend/optimizer/plan_node/
predicate_pushdown.rs1use std::collections::HashMap;
16
17use super::*;
18use crate::optimizer::plan_visitor::ShareParentCounter;
19use crate::optimizer::{LogicalPlanRef as PlanRef, PlanVisitor};
20
21pub trait PredicatePushdown {
24 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}