risingwave_frontend/optimizer/plan_node/
col_pruning.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::ShareParentCounter;
19use crate::optimizer::{LogicalPlanRef as PlanRef, PlanVisitor};
20
21/// The trait for column pruning, only logical plan node will use it, though all plan node impl it.
22pub trait ColPrunable {
23    /// Transform the plan node to only output the required columns ordered by index number.
24    ///
25    /// `required_cols` must be a subset of the range `0..self.schema().len()`.
26    ///
27    /// After calling `prune_col` on the children, their output schema may change, so
28    /// the caller may need to transform its [`InputRef`](crate::expr::InputRef) using
29    /// [`ColIndexMapping`](crate::utils::ColIndexMapping).
30    ///
31    /// When implementing this method for a node, it may require its children to produce additional
32    /// columns besides `required_cols`. In this case, it may need to insert a
33    /// [`LogicalProject`](super::LogicalProject) above to have a correct schema.
34    fn prune_col(&self, required_cols: &[usize], ctx: &mut ColumnPruningContext) -> PlanRef;
35}
36
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38enum ColumnPruningPhase {
39    Idle,
40    Collect,
41    Rewrite,
42}
43
44#[derive(Debug, Clone)]
45struct ShareColumnPruning {
46    original_input: PlanRef,
47    required_cols: Vec<usize>,
48}
49
50#[derive(Debug, Clone)]
51pub struct ColumnPruningContext {
52    pending_required_cols: HashMap<ShareId, Vec<Vec<usize>>>,
53    collected_shares: HashMap<ShareId, ShareColumnPruning>,
54    share_parent_counter: ShareParentCounter,
55    share_mappings: HashMap<ShareId, ColIndexMapping>,
56    skipped_shares: HashSet<ShareId>,
57    phase: ColumnPruningPhase,
58}
59
60impl ColumnPruningContext {
61    pub fn new(root: PlanRef) -> Self {
62        let mut share_parent_counter = ShareParentCounter::default();
63        share_parent_counter.visit(root);
64        Self {
65            pending_required_cols: HashMap::new(),
66            collected_shares: HashMap::new(),
67            share_parent_counter,
68            share_mappings: HashMap::new(),
69            skipped_shares: HashSet::new(),
70            phase: ColumnPruningPhase::Idle,
71        }
72    }
73
74    pub(in crate::optimizer) fn is_running(&self) -> bool {
75        self.phase != ColumnPruningPhase::Idle
76    }
77
78    pub(in crate::optimizer) fn is_collecting(&self) -> bool {
79        self.phase == ColumnPruningPhase::Collect
80    }
81
82    pub(in crate::optimizer) fn get_parent_num(&self, share: &LogicalShare) -> usize {
83        self.share_parent_counter.get_parent_num(share)
84    }
85
86    pub(in crate::optimizer) fn add_required_cols(
87        &mut self,
88        share: &LogicalShare,
89        required_cols: Vec<usize>,
90    ) -> Option<Vec<usize>> {
91        let share_id = share.share_id();
92        let parent_num = self.share_parent_counter.get_parent_num_by_id(share_id);
93        let pending = self.pending_required_cols.entry(share_id).or_default();
94        pending.push(required_cols);
95        assert!(
96            pending.len() <= parent_num,
97            "share {share_id:?} received more column requirements than parents"
98        );
99        if pending.len() != parent_num {
100            return None;
101        }
102
103        let merged_required_cols = self
104            .pending_required_cols
105            .remove(&share_id)
106            .expect("column requirements must exist")
107            .into_iter()
108            .flatten()
109            .sorted()
110            .dedup()
111            .collect_vec();
112        self.collected_shares
113            .try_insert(
114                share_id,
115                ShareColumnPruning {
116                    original_input: share.input(),
117                    required_cols: merged_required_cols.clone(),
118                },
119            )
120            .expect("column requirements must be merged once per share");
121        Some(merged_required_cols)
122    }
123
124    pub(in crate::optimizer) fn share_mapping(&self, share: &LogicalShare) -> ColIndexMapping {
125        self.share_mappings
126            .get(&share.share_id())
127            .unwrap_or_else(|| {
128                panic!(
129                    "logical share {:?} has no column-pruning mapping",
130                    share.share_id()
131                )
132            })
133            .clone()
134    }
135
136    /// Rebuilds one shared definition on first use. Nested shares recursively rebuild first, so
137    /// the call stack provides the child-before-parent order without a separate dependency graph.
138    pub(in crate::optimizer) fn ensure_share_rebuilt(&mut self, share: &LogicalShare) {
139        let share_id = share.share_id();
140        if self.share_mappings.contains_key(&share_id) {
141            return;
142        }
143        assert_eq!(self.phase, ColumnPruningPhase::Rewrite);
144
145        let Some(ShareColumnPruning {
146            original_input,
147            required_cols,
148        }) = self.collected_shares.remove(&share_id)
149        else {
150            // The share was skipped during collection (see `run`): it never received a
151            // requirement from every parent, so it must keep its original definition and
152            // full schema. Parents prune with their own projections above the share.
153            assert!(
154                self.skipped_shares.contains(&share_id),
155                "share {share_id:?} has no collected column requirements"
156            );
157            self.share_mappings
158                .try_insert(share_id, ColIndexMapping::identity(share.schema().len()))
159                .expect("a logical share must be rebuilt once");
160            return;
161        };
162        let old_schema_len = original_input.schema().len();
163        let rebuilt_input = original_input.prune_col(&required_cols, self);
164        let mapping = ColIndexMapping::with_remaining_columns(&required_cols, old_schema_len);
165        debug_assert_eq!(mapping.target_size(), rebuilt_input.schema().len());
166
167        share.ctx().update_logical_share(share_id, rebuilt_input);
168        self.share_mappings
169            .try_insert(share_id, mapping)
170            .expect("a logical share must be rebuilt once");
171    }
172
173    pub(in crate::optimizer) fn run(&mut self, root: PlanRef, required_cols: &[usize]) -> PlanRef {
174        self.phase = ColumnPruningPhase::Collect;
175        let collected = root.prune_col_inner(required_cols, self);
176        // `ShareParentCounter` counts parents via the visitor walk, but the transformation
177        // walk is not guaranteed to reach a share from every parent: some `ColPrunable`
178        // impls (e.g. `LogicalVectorSearchLookupJoin`, which never prunes its lookup side)
179        // legitimately stop recursing into an input. Pruning with a subset of the parents'
180        // requirements would drop columns still referenced by the parents that never
181        // contributed, so such shares are skipped and keep their original definition.
182        self.skipped_shares.extend(
183            self.pending_required_cols
184                .drain()
185                .map(|(share_id, _)| share_id),
186        );
187        if self.collected_shares.is_empty() {
188            self.phase = ColumnPruningPhase::Idle;
189            return collected;
190        }
191
192        self.phase = ColumnPruningPhase::Rewrite;
193        let result = root.prune_col_inner(required_cols, self);
194        self.phase = ColumnPruningPhase::Idle;
195        result
196    }
197}
risingwave_frontend/optimizer/plan_node/col_pruning.rs

risingwave_frontend/optimizer/plan_node/
col_pruning.rs
