risingwave_frontend/optimizer/plan_node/
logical_changelog.rs

1// Copyright 2024 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 itertools::Itertools;
16
17use super::expr_visitable::ExprVisitable;
18use super::generic::{_CHANGELOG_ROW_ID, CHANGELOG_OP, GenericPlanRef};
19use super::utils::impl_distill_by_unit;
20use super::{
21    BatchPlanRef, ColPrunable, ColumnPruningContext, ExprRewritable, Logical,
22    LogicalPlanRef as PlanRef, LogicalProject, PlanBase, PlanTreeNodeUnary, PredicatePushdown,
23    RewriteStreamContext, StreamChangeLog, StreamPlanRef, ToBatch, ToStream, ToStreamContext,
24    gen_filter_and_pushdown, generic,
25};
26use crate::error::ErrorCode::BindError;
27use crate::error::Result;
28use crate::optimizer::plan_node::generic::PhysicalPlanRef;
29use crate::optimizer::property::{Distribution, RequiredDist};
30use crate::utils::{ColIndexMapping, Condition};
31
32#[derive(Debug, Clone, PartialEq, Eq, Hash)]
33pub struct LogicalChangeLog {
34    pub base: PlanBase<Logical>,
35    core: generic::ChangeLog<PlanRef>,
36}
37
38impl LogicalChangeLog {
39    pub fn create(input: PlanRef, key_indices: Option<Vec<usize>>) -> PlanRef {
40        Self::new(input, key_indices, true, true).into()
41    }
42
43    pub fn new(
44        input: PlanRef,
45        key_indices: Option<Vec<usize>>,
46        need_op: bool,
47        need_changelog_row_id: bool,
48    ) -> Self {
49        let core = generic::ChangeLog::new(input, key_indices, need_op, need_changelog_row_id);
50        Self::with_core(core)
51    }
52
53    pub fn with_core(core: generic::ChangeLog<PlanRef>) -> Self {
54        let base = PlanBase::new_logical_with_core(&core);
55        LogicalChangeLog { base, core }
56    }
57}
58
59impl PlanTreeNodeUnary<Logical> for LogicalChangeLog {
60    fn input(&self) -> PlanRef {
61        self.core.input.clone()
62    }
63
64    fn clone_with_input(&self, input: PlanRef) -> Self {
65        let core = self.core.clone_with_input(input);
66
67        Self::with_core(core)
68    }
69
70    fn rewrite_with_input(
71        &self,
72        input: PlanRef,
73        input_col_change: ColIndexMapping,
74    ) -> (Self, ColIndexMapping) {
75        let key_indices = self
76            .core
77            .key_indices
78            .as_ref()
79            .map(|key| input_col_change.map_all(key));
80        let changelog = Self::new(input, key_indices, self.core.need_op, true);
81
82        let out_col_change = if self.core.need_op {
83            let (mut output_vec, len) = input_col_change.into_parts();
84            output_vec.push(Some(len));
85            ColIndexMapping::new(output_vec, len + 1)
86        } else {
87            input_col_change
88        };
89
90        let (mut output_vec, len) = out_col_change.into_parts();
91        let out_col_change = if self.core.need_changelog_row_id {
92            output_vec.push(Some(len));
93            ColIndexMapping::new(output_vec, len + 1)
94        } else {
95            ColIndexMapping::new(output_vec, len + 1)
96        };
97
98        (changelog, out_col_change)
99    }
100}
101
102impl_plan_tree_node_for_unary! { Logical, LogicalChangeLog}
103impl_distill_by_unit!(LogicalChangeLog, core, "LogicalChangeLog");
104
105impl ExprRewritable<Logical> for LogicalChangeLog {}
106
107impl ExprVisitable for LogicalChangeLog {}
108
109impl PredicatePushdown for LogicalChangeLog {
110    fn predicate_pushdown(
111        &self,
112        predicate: Condition,
113        ctx: &mut super::PredicatePushdownContext,
114    ) -> PlanRef {
115        gen_filter_and_pushdown(self, predicate, Condition::true_cond(), ctx)
116    }
117}
118
119impl ColPrunable for LogicalChangeLog {
120    fn prune_col(&self, required_cols: &[usize], ctx: &mut ColumnPruningContext) -> PlanRef {
121        let fields = self.schema().fields();
122        let mut need_op = false;
123        let mut need_changelog_row_id = false;
124        // columns required as input from children,
125        // generated columns like op and event_id are removed
126        let mut input_required_cols = required_cols
127            .iter()
128            .filter_map(|a| {
129                if let Some(f) = fields.get(*a) {
130                    if f.name == CHANGELOG_OP {
131                        need_op = true;
132                        None
133                    } else if f.name == _CHANGELOG_ROW_ID {
134                        need_changelog_row_id = true;
135                        None
136                    } else {
137                        Some(*a)
138                    }
139                } else {
140                    Some(*a)
141                }
142            })
143            .collect_vec();
144
145        // `StreamChangeLog` must see the input stream key to co-locate changes for the same row.
146        // Keep key columns internally even when they are not selected from the changelog CTE, and
147        // project them away above the operator afterwards.
148        if let Some(stream_key) = self.input().stream_key() {
149            for &key in stream_key {
150                if !input_required_cols.contains(&key) {
151                    input_required_cols.push(key);
152                }
153            }
154        }
155
156        // add declared business keys to input request
157        if let Some(key) = &self.core.key_indices {
158            for &index in key {
159                if !input_required_cols.contains(&index) {
160                    input_required_cols.push(index);
161                }
162            }
163        }
164
165        let new_input = self.input().prune_col(&input_required_cols, ctx);
166        let input_mapping = ColIndexMapping::with_remaining_columns(
167            &input_required_cols,
168            self.input().schema().len(),
169        );
170        let key_indices = self
171            .core
172            .key_indices
173            .as_ref()
174            .map(|key| input_mapping.map_all(key));
175
176        // updated plan with new inputs
177        let changelog: PlanRef =
178            Self::new(new_input, key_indices, need_op, need_changelog_row_id).into();
179
180        // output now contains business keys as well, we need to hide it if not requested by the parent
181        let (mut output_mapping, new_output_len) = input_mapping.into_parts();
182
183        if self.core.need_op {
184            output_mapping.push(need_op.then_some(new_output_len));
185        }
186
187        if self.core.need_changelog_row_id {
188            output_mapping
189                .push(need_changelog_row_id.then_some(new_output_len + usize::from(need_op)));
190        }
191
192        let output_len = changelog.schema().len();
193        let output_mapping = ColIndexMapping::new(output_mapping, output_len);
194        let output_required_cols = output_mapping.map_all(required_cols);
195
196        if output_required_cols.iter().copied().eq(0..output_len) {
197            changelog
198        } else {
199            let output_mapping =
200                ColIndexMapping::with_remaining_columns(&output_required_cols, output_len);
201
202            LogicalProject::with_mapping(changelog, output_mapping).into()
203        }
204    }
205}
206
207impl ToBatch for LogicalChangeLog {
208    fn to_batch(&self) -> Result<BatchPlanRef> {
209        Err(BindError("With changelog cte only support with create mv/sink".to_owned()).into())
210    }
211}
212
213impl ToStream for LogicalChangeLog {
214    fn to_stream(&self, ctx: &mut ToStreamContext) -> Result<StreamPlanRef> {
215        let input = self.input().to_stream(ctx)?;
216
217        let input = if let Some(key) = &self.core.key_indices {
218            RequiredDist::hash_shard(key).streaming_enforce_if_not_satisfies(input)?
219        } else if matches!(input.distribution(), Distribution::SomeShard) {
220            RequiredDist::hash_shard(input.expect_stream_key())
221                .streaming_enforce_if_not_satisfies(input)?
222        } else {
223            input
224        };
225
226        let dist = input.distribution();
227        let distribution_keys = match dist {
228            Distribution::HashShard(distribution_keys)
229            | Distribution::UpstreamHashShard(distribution_keys, _) => distribution_keys.clone(),
230            Distribution::Single => {
231                vec![]
232            }
233            _ => {
234                return Err(BindError(format!(
235                    "ChangeLog requires input to be hash distributed, single, but got {:?}",
236                    dist
237                ))
238                .into());
239            }
240        };
241        let core = self.core.clone_with_input(input);
242        let row_id_index = self.schema().fields().len() - 1;
243        let plan = StreamChangeLog::new_with_dist(
244            core,
245            Distribution::HashShard(vec![row_id_index]),
246            distribution_keys.into_iter().map(|k| k as u32).collect(),
247        )
248        .into();
249
250        Ok(plan)
251    }
252
253    fn logical_rewrite_for_stream(
254        &self,
255        ctx: &mut RewriteStreamContext,
256    ) -> Result<(PlanRef, ColIndexMapping)> {
257        let (input, input_col_change) = self.input().logical_rewrite_for_stream(ctx)?;
258        let (changelog, out_col_change) = self.rewrite_with_input(input, input_col_change);
259        Ok((changelog.into(), out_col_change))
260    }
261}
262
263#[cfg(test)]
264mod tests {
265    use std::rc::Rc;
266
267    use risingwave_common::catalog::{CdcTableDesc, ColumnDesc, ColumnId, TableId};
268    use risingwave_common::id::SourceId;
269    use risingwave_common::types::DataType;
270    use risingwave_common::util::sort_util::{ColumnOrder, OrderType};
271    use risingwave_connector::source::cdc::CdcScanOptions;
272
273    use super::*;
274    use crate::optimizer::optimizer_context::OptimizerContext;
275    use crate::optimizer::plan_node::{BackfillType, LogicalCdcScan, PlanTreeNodeUnary};
276
277    #[test]
278    fn test_changelog_on_cdc_scan_enforces_hash_distribution() {
279        let desc = CdcTableDesc {
280            table_id: TableId::new(1),
281            source_id: SourceId::new(2),
282            external_table_name: "mydb.orders".to_owned(),
283            pk: vec![ColumnOrder::new(0, OrderType::ascending())],
284            pk_comparisons: vec![risingwave_common::catalog::CdcKeyComparison::Native],
285            columns: vec![
286                ColumnDesc::named("id", ColumnId::new(1), DataType::Int32),
287                ColumnDesc::named("payload", ColumnId::new(2), DataType::Varchar),
288            ],
289            stream_key: vec![0],
290            ..Default::default()
291        };
292        let scan = LogicalCdcScan::create(
293            "orders".to_owned(),
294            Rc::new(desc),
295            OptimizerContext::mock(),
296            CdcScanOptions {
297                disable_backfill: true,
298                ..Default::default()
299            },
300        );
301        let changelog = LogicalChangeLog::create(scan.clone().into(), None);
302        let stream = changelog
303            .to_stream(&mut ToStreamContext::new_with_backfill_type(
304                false,
305                BackfillType::Replicated,
306            ))
307            .unwrap();
308
309        let input = stream.as_stream_change_log().unwrap().input();
310        assert!(input.as_stream_exchange().is_some());
311        assert_eq!(input.distribution(), &Distribution::HashShard(vec![0]));
312
313        // The primary key is retained internally even when it is not selected from the changelog.
314        let changelog = LogicalChangeLog::create(scan.into(), None);
315        let pruned = changelog.prune_col(
316            &[1, 2, 3],
317            &mut ColumnPruningContext::new(changelog.clone()),
318        );
319        let stream = pruned
320            .to_stream(&mut ToStreamContext::new_with_backfill_type(
321                false,
322                BackfillType::Replicated,
323            ))
324            .unwrap();
325        let changelog = stream.as_stream_project().unwrap().input();
326        let input = changelog.as_stream_change_log().unwrap().input();
327        assert!(input.as_stream_exchange().is_some());
328        assert_eq!(input.distribution(), &Distribution::HashShard(vec![1]));
329    }
330}
risingwave_frontend/optimizer/plan_node/logical_changelog.rs

risingwave_frontend/optimizer/plan_node/
logical_changelog.rs
