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        // The sort buffer table is only attached when the session enables it at plan time.
106        // Old plans without this table simply keep the pass-through behavior.
107        let sort_buffer_table = self
108            .base
109            .ctx()
110            .session_ctx()
111            .config()
112            .enable_locality_sort_buffer()
113            .then(|| self.build_sort_buffer_catalog(state).to_prost());
114
115        let locality_provider_node = LocalityProviderNode {
116            locality_columns: self.locality_columns().iter().map(|&i| i as u32).collect(),
117            // State table for buffering input data
118            state_table: Some(state_table.to_prost()),
119            // Progress table for tracking backfill progress
120            progress_table: Some(progress_table.to_prost()),
121            // Ephemeral sorted log table for the post-backfill sort buffer
122            sort_buffer_table,
123            rate_limit: self.base.ctx().overwrite_options().backfill_rate_limit,
124        };
125
126        PbNodeBody::LocalityProvider(Box::new(locality_provider_node))
127    }
128}
129
130impl ExprRewritable<Stream> for StreamLocalityProvider {
131    fn has_rewritable_expr(&self) -> bool {
132        false
133    }
134
135    fn rewrite_exprs(&self, _r: &mut dyn ExprRewriter) -> PlanRef {
136        self.clone().into()
137    }
138}
139
140impl ExprVisitable for StreamLocalityProvider {
141    fn visit_exprs(&self, _v: &mut dyn ExprVisitor) {
142        // No expressions to visit
143    }
144}
145
146impl StreamLocalityProvider {
147    /// Build the state table catalog for buffering input data
148    /// Schema: same as input schema (locality handled by primary key ordering)
149    /// Key: `locality_columns` (vnode handled internally by `StateTable`)
150    fn build_state_catalog(&self, state: &mut BuildFragmentGraphState) -> TableCatalog {
151        let mut catalog_builder = TableCatalogBuilder::default();
152        let input = self.input();
153        let input_schema = input.schema();
154
155        // Add all input columns in original order
156        for field in &input_schema.fields {
157            catalog_builder.add_column(field);
158        }
159
160        // Set locality columns as primary key.
161        for locality_col_idx in self.locality_columns() {
162            catalog_builder.add_order_column(*locality_col_idx, OrderType::ascending());
163        }
164        // add streaming key of the input as the rest of the primary key
165        for &key_col_idx in input.expect_stream_key() {
166            catalog_builder.add_order_column(key_col_idx, OrderType::ascending());
167        }
168
169        catalog_builder.set_value_indices((0..input_schema.len()).collect());
170
171        catalog_builder
172            .build(
173                self.input().distribution().dist_column_indices().to_vec(),
174                0,
175            )
176            .with_id(state.gen_table_id_wrapped())
177    }
178
179    /// Build the ephemeral sorted log table catalog for the post-backfill sort buffer.
180    /// Schema: all input columns + `gen` (epoch generation) + `op` + `seq`
181    /// Key: `gen` + `locality_columns` + input stream key + `seq`
182    ///
183    /// Buffered changes of one epoch are written with the current epoch as `gen`, so that a
184    /// prefix scan on `gen` replays exactly the changes of that epoch in locality order. `seq`
185    /// is a monotonically increasing sequence number within the epoch, making every buffered
186    /// change a unique key (append-only writes, per-key order preserved). `op` records the
187    /// change kind. Old generations are cleaned by the state-clean watermark on `gen`.
188    fn build_sort_buffer_catalog(&self, state: &mut BuildFragmentGraphState) -> TableCatalog {
189        let mut catalog_builder = TableCatalogBuilder::default();
190        let input = self.input();
191        let input_schema = input.schema();
192
193        // Add all input columns in original order
194        for field in &input_schema.fields {
195            catalog_builder.add_column(field);
196        }
197
198        let gen_col_idx = catalog_builder.add_column(&Field::with_name(DataType::Int64, "_rw_gen"));
199        catalog_builder.add_column(&Field::with_name(DataType::Int16, "_rw_op"));
200        let seq_col_idx = catalog_builder.add_column(&Field::with_name(DataType::Int64, "_rw_seq"));
201
202        // Primary key: gen + locality columns + stream key + seq
203        catalog_builder.add_order_column(gen_col_idx, OrderType::ascending());
204        for locality_col_idx in self.locality_columns() {
205            catalog_builder.add_order_column(*locality_col_idx, OrderType::ascending());
206        }
207        for &key_col_idx in input.expect_stream_key() {
208            catalog_builder.add_order_column(key_col_idx, OrderType::ascending());
209        }
210        catalog_builder.add_order_column(seq_col_idx, OrderType::ascending());
211
212        // Clean old generations by the state-clean watermark on the `gen` column.
213        catalog_builder.set_clean_watermark_indices(vec![gen_col_idx]);
214
215        let num_of_columns = catalog_builder.columns().len();
216        catalog_builder.set_value_indices((0..num_of_columns).collect());
217
218        catalog_builder
219            .build(
220                self.input().distribution().dist_column_indices().to_vec(),
221                0,
222            )
223            .with_id(state.gen_table_id_wrapped())
224    }
225
226    /// Build the progress table catalog for tracking backfill progress
227    /// Schema: | vnode | pk(locality columns + input stream keys) | `backfill_finished` | `row_count` |
228    /// Key: | vnode | pk(locality columns + input stream keys) |
229    fn build_progress_catalog(&self, state: &mut BuildFragmentGraphState) -> TableCatalog {
230        let mut catalog_builder = TableCatalogBuilder::default();
231        let input = self.input();
232        let input_schema = input.schema();
233
234        // Add vnode column as primary key
235        catalog_builder.add_column(&Field::with_name(VirtualNode::RW_TYPE, "vnode"));
236        catalog_builder.add_order_column(0, OrderType::ascending());
237
238        // Add locality columns as part of primary key
239        for &locality_col_idx in self.locality_columns() {
240            let field = &input_schema.fields[locality_col_idx];
241            catalog_builder.add_column(field);
242        }
243
244        // Add stream key columns as part of primary key (excluding those already added as locality columns)
245        for &key_col_idx in input.expect_stream_key() {
246            let field = &input_schema.fields[key_col_idx];
247            catalog_builder.add_column(field);
248        }
249
250        // Add backfill_finished column
251        catalog_builder.add_column(&Field::with_name(DataType::Boolean, "backfill_finished"));
252
253        // Add row_count column
254        catalog_builder.add_column(&Field::with_name(DataType::Int64, "row_count"));
255
256        // Set vnode column index and distribution key
257        catalog_builder.set_vnode_col_idx(0);
258        catalog_builder.set_dist_key_in_pk(vec![0]);
259
260        let num_of_columns = catalog_builder.columns().len();
261        catalog_builder.set_value_indices((0..num_of_columns).collect_vec());
262
263        catalog_builder
264            .build(vec![0], 1)
265            .with_id(state.gen_table_id_wrapped())
266    }
267}
risingwave_frontend/optimizer/plan_node/stream_locality_provider.rs

risingwave_frontend/optimizer/plan_node/
stream_locality_provider.rs
