risingwave_frontend/optimizer/plan_node/
stream_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 risingwave_pb::stream_plan::ChangeLogNode;
16use risingwave_pb::stream_plan::change_log_node::{Keyed, Mode};
17use risingwave_pb::stream_plan::stream_node::PbNodeBody;
18
19use super::expr_visitable::ExprVisitable;
20use super::stream::StreamPlanNodeMetadata;
21use super::stream::prelude::PhysicalPlanRef;
22use super::utils::impl_distill_by_unit;
23use super::{ExprRewritable, PlanBase, PlanTreeNodeUnary, Stream, StreamNode, generic};
24use crate::optimizer::StreamPlanRef as PlanRef;
25use crate::optimizer::property::{Distribution, MonotonicityMap, StreamKind};
26use crate::stream_fragmenter::BuildFragmentGraphState;
27
28#[derive(Debug, Clone, PartialEq, Eq, Hash)]
29pub struct StreamChangeLog {
30    pub base: PlanBase<Stream>,
31    core: generic::ChangeLog<PlanRef>,
32    distribution_keys: Vec<u32>,
33}
34
35impl StreamChangeLog {
36    pub fn new_with_dist(
37        core: generic::ChangeLog<PlanRef>,
38        dist: Distribution,
39        distribution_keys: Vec<u32>,
40    ) -> Self {
41        let input = core.input.clone();
42        let base = PlanBase::new_stream_with_core(
43            &core,
44            dist,
45            // The changelog will convert all delete/update to insert, so it must be append-only here.
46            StreamKind::AppendOnly,
47            input.emit_on_window_close(),
48            input.watermark_columns().clone(),
49            MonotonicityMap::new(), // TODO: derive monotonicity
50        );
51        StreamChangeLog {
52            base,
53            core,
54            distribution_keys,
55        }
56    }
57}
58
59impl PlanTreeNodeUnary<Stream> for StreamChangeLog {
60    fn input(&self) -> PlanRef {
61        self.core.input.clone()
62    }
63
64    fn clone_with_input(&self, input: PlanRef) -> Self {
65        let mut core = self.core.clone();
66        core.input = input;
67        Self::new_with_dist(
68            core,
69            self.base.distribution().clone(),
70            self.distribution_keys.clone(),
71        )
72    }
73}
74
75impl_plan_tree_node_for_unary! { Stream, StreamChangeLog }
76impl_distill_by_unit!(StreamChangeLog, core, "StreamChangeLog");
77
78impl StreamNode for StreamChangeLog {
79    fn to_stream_prost_body(&self, _state: &mut BuildFragmentGraphState) -> PbNodeBody {
80        let mode = self.core.key_indices.as_ref().map(|key| {
81            Mode::Keyed(Keyed {
82                additional_key_columns: key
83                    .iter()
84                    .map(|&index| index as u32)
85                    .filter(|index| !self.distribution_keys.contains(index))
86                    .collect(),
87            })
88        });
89        PbNodeBody::Changelog(Box::new(ChangeLogNode {
90            need_op: self.core.need_op,
91            distribution_keys: self.distribution_keys.clone(),
92            mode,
93            stream_keys: self
94                .core
95                .input
96                .expect_stream_key()
97                .iter()
98                .map(|&index| index as u32)
99                .collect(),
100        }))
101    }
102}
103
104impl ExprRewritable<Stream> for StreamChangeLog {}
105
106impl ExprVisitable for StreamChangeLog {}
107
108#[cfg(test)]
109mod tests {
110    use risingwave_common::catalog::{Field, Schema};
111    use risingwave_common::types::DataType;
112
113    use super::*;
114    use crate::optimizer::optimizer_context::OptimizerContext;
115    use crate::optimizer::plan_node::{StreamExchange, StreamValues};
116
117    #[test]
118    fn test_changelog_separates_input_distribution_and_additional_keys() {
119        for (key_indices, distribution_keys, additional_key_columns) in [
120            (Some(vec![1, 2]), vec![1], Some(vec![2])),
121            (Some(vec![1, 2]), vec![2, 1], Some(vec![])),
122            (None, vec![1], None),
123        ] {
124            let values = StreamValues::new(generic::Values::new_with_stream_key(
125                vec![],
126                Schema::new(vec![
127                    Field::with_name(DataType::Int32, "id"),
128                    Field::with_name(DataType::Int32, "tenant"),
129                    Field::with_name(DataType::Int32, "business_key"),
130                ]),
131                OptimizerContext::mock(),
132                vec![0],
133            ));
134            let input = StreamExchange::new(
135                values.into(),
136                Distribution::HashShard(distribution_keys.clone()),
137            );
138            let core = generic::ChangeLog::new(
139                input.into(),
140                key_indices,
141                Some(distribution_keys.clone()),
142                true,
143                true,
144            );
145            let distribution_keys = distribution_keys
146                .into_iter()
147                .map(|index| index as u32)
148                .collect::<Vec<_>>();
149            let changelog = StreamChangeLog::new_with_dist(
150                core,
151                Distribution::HashShard(vec![4]),
152                distribution_keys.clone(),
153            );
154            let PbNodeBody::Changelog(node) =
155                changelog.to_stream_prost_body(&mut BuildFragmentGraphState::default())
156            else {
157                panic!("expected a changelog node");
158            };
159
160            assert_eq!(node.stream_keys, vec![0]);
161            assert_eq!(node.distribution_keys, distribution_keys);
162            assert_eq!(
163                node.mode,
164                additional_key_columns.map(|additional_key_columns| Mode::Keyed(Keyed {
165                    additional_key_columns,
166                }))
167            );
168        }
169    }
170}
risingwave_frontend/optimizer/plan_node/stream_changelog.rs

risingwave_frontend/optimizer/plan_node/
stream_changelog.rs
