risingwave_frontend/optimizer/plan_node/
stream_locality_provider.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 itertools::Itertools;
16use pretty_xmlish::XmlNode;
17use risingwave_common::catalog::Field;
18use risingwave_common::hash::VirtualNode;
19use risingwave_common::types::DataType;
20use risingwave_common::util::sort_util::OrderType;
21use risingwave_pb::stream_plan::LocalityProviderNode;
22use risingwave_pb::stream_plan::stream_node::PbNodeBody;
23
24use super::stream::prelude::*;
25use super::utils::{Distill, TableCatalogBuilder, childless_record};
26use super::{ExprRewritable, PlanTreeNodeUnary, StreamNode, StreamPlanRef as PlanRef, generic};
27use crate::TableCatalog;
28use crate::catalog::TableId;
29use crate::expr::{ExprRewriter, ExprVisitor};
30use crate::optimizer::plan_node::PlanBase;
31use crate::optimizer::plan_node::expr_visitable::ExprVisitable;
32use crate::optimizer::property::Distribution;
33use crate::stream_fragmenter::BuildFragmentGraphState;
34
35/// `StreamLocalityProvider` implements [`super::LogicalLocalityProvider`]
36#[derive(Debug, Clone, PartialEq, Eq, Hash)]
37pub struct StreamLocalityProvider {
38    pub base: PlanBase<Stream>,
39    core: generic::LocalityProvider<PlanRef>,
40}
41
42impl StreamLocalityProvider {
43    pub fn new(core: generic::LocalityProvider<PlanRef>) -> Self {
44        let input = core.input.clone();
45
46        let dist = match input.distribution() {
47            Distribution::HashShard(keys) => {
48                // If the input is hash-distributed, we make it a UpstreamHashShard distribution
49                // just like a normal table scan. It is used to ensure locality provider is in its own fragment.
50                // This is important to ensure the backfill ordering can recognize and build
51                // the dependency graph among different backfill-needed fragments.
52                Distribution::UpstreamHashShard(keys.clone(), TableId::placeholder())
53            }
54            Distribution::UpstreamHashShard(keys, table_id) => {
55                Distribution::UpstreamHashShard(keys.clone(), *table_id)
56            }
57            _ => {
58                panic!("LocalityProvider input must be hash-distributed");
59            }
60        };
61
62        // LocalityProvider maintains the append-only behavior if input is append-only
63        let base = PlanBase::new_stream_with_core(
64            &core,
65            dist,
66            input.stream_kind(),
67            input.emit_on_window_close(),
68            input.watermark_columns().clone(),
69            input.columns_monotonicity().clone(),
70        );
71        StreamLocalityProvider { base, core }
72    }
73
74    pub fn locality_columns(&self) -> &[usize] {
75        &self.core.locality_columns
76    }
77}
78
79impl PlanTreeNodeUnary<Stream> for StreamLocalityProvider {
80    fn input(&self) -> PlanRef {
81        self.core.input.clone()
82    }
83
84    fn clone_with_input(&self, input: PlanRef) -> Self {
85        let mut core = self.core.clone();
86        core.input = input;
87        Self::new(core)
88    }
89}
90
91impl_plan_tree_node_for_unary! { Stream, StreamLocalityProvider }
92
93impl Distill for StreamLocalityProvider {
94    fn distill<'a>(&self) -> XmlNode<'a> {
95        let vec = self.core.fields_pretty();
96        childless_record("StreamLocalityProvider", vec)
97    }
98}
99
100impl StreamNode for StreamLocalityProvider {
101    fn to_stream_prost_body(&self, state: &mut BuildFragmentGraphState) -> PbNodeBody {
102        let state_table = self.build_state_catalog(state);
103        let progress_table = self.build_progress_catalog(state);
104
105        let locality_provider_node = LocalityProviderNode {
106            locality_columns: self.locality_columns().iter().map(|&i| i as u32).collect(),
107            // State table for buffering input data
108            state_table: Some(state_table.to_prost()),
109            // Progress table for tracking backfill progress
110            progress_table: Some(progress_table.to_prost()),
111            rate_limit: self.base.ctx().overwrite_options().backfill_rate_limit,
112        };
113
114        PbNodeBody::LocalityProvider(Box::new(locality_provider_node))
115    }
116}
117
118impl ExprRewritable<Stream> for StreamLocalityProvider {
119    fn has_rewritable_expr(&self) -> bool {
120        false
121    }
122
123    fn rewrite_exprs(&self, _r: &mut dyn ExprRewriter) -> PlanRef {
124        self.clone().into()
125    }
126}
127
128impl ExprVisitable for StreamLocalityProvider {
129    fn visit_exprs(&self, _v: &mut dyn ExprVisitor) {
130        // No expressions to visit
131    }
132}
133
134impl StreamLocalityProvider {
135    /// Build the state table catalog for buffering input data
136    /// Schema: same as input schema (locality handled by primary key ordering)
137    /// Key: `locality_columns` (vnode handled internally by `StateTable`)
138    fn build_state_catalog(&self, state: &mut BuildFragmentGraphState) -> TableCatalog {
139        let mut catalog_builder = TableCatalogBuilder::default();
140        let input = self.input();
141        let input_schema = input.schema();
142
143        // Add all input columns in original order
144        for field in &input_schema.fields {
145            catalog_builder.add_column(field);
146        }
147
148        // Set locality columns as primary key.
149        for locality_col_idx in self.locality_columns() {
150            catalog_builder.add_order_column(*locality_col_idx, OrderType::ascending());
151        }
152        // add streaming key of the input as the rest of the primary key
153        for &key_col_idx in input.expect_stream_key() {
154            catalog_builder.add_order_column(key_col_idx, OrderType::ascending());
155        }
156
157        catalog_builder.set_value_indices((0..input_schema.len()).collect());
158
159        catalog_builder
160            .build(
161                self.input().distribution().dist_column_indices().to_vec(),
162                0,
163            )
164            .with_id(state.gen_table_id_wrapped())
165    }
166
167    /// Build the progress table catalog for tracking backfill progress
168    /// Schema: | vnode | pk(locality columns + input stream keys) | `backfill_finished` | `row_count` |
169    /// Key: | vnode | pk(locality columns + input stream keys) |
170    fn build_progress_catalog(&self, state: &mut BuildFragmentGraphState) -> TableCatalog {
171        let mut catalog_builder = TableCatalogBuilder::default();
172        let input = self.input();
173        let input_schema = input.schema();
174
175        // Add vnode column as primary key
176        catalog_builder.add_column(&Field::with_name(VirtualNode::RW_TYPE, "vnode"));
177        catalog_builder.add_order_column(0, OrderType::ascending());
178
179        // Add locality columns as part of primary key
180        for &locality_col_idx in self.locality_columns() {
181            let field = &input_schema.fields[locality_col_idx];
182            catalog_builder.add_column(field);
183        }
184
185        // Add stream key columns as part of primary key (excluding those already added as locality columns)
186        for &key_col_idx in input.expect_stream_key() {
187            let field = &input_schema.fields[key_col_idx];
188            catalog_builder.add_column(field);
189        }
190
191        // Add backfill_finished column
192        catalog_builder.add_column(&Field::with_name(DataType::Boolean, "backfill_finished"));
193
194        // Add row_count column
195        catalog_builder.add_column(&Field::with_name(DataType::Int64, "row_count"));
196
197        // Set vnode column index and distribution key
198        catalog_builder.set_vnode_col_idx(0);
199        catalog_builder.set_dist_key_in_pk(vec![0]);
200
201        let num_of_columns = catalog_builder.columns().len();
202        catalog_builder.set_value_indices((0..num_of_columns).collect_vec());
203
204        catalog_builder
205            .build(vec![0], 1)
206            .with_id(state.gen_table_id_wrapped())
207    }
208}
risingwave_frontend/optimizer/plan_node/stream_locality_provider.rs

risingwave_frontend/optimizer/plan_node/
stream_locality_provider.rs
