risingwave_frontend/optimizer/plan_node/
stream_eowc_over_window.rs

1// Copyright 2023 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::HashSet;
16
17use risingwave_common::catalog::Field;
18use risingwave_common::types::DataType;
19use risingwave_common::util::sort_util::OrderType;
20use risingwave_expr::window_function::WindowFuncKind;
21use risingwave_pb::stream_plan::stream_node::PbNodeBody;
22
23use super::generic::{self, PlanWindowFunction};
24use super::stream::prelude::*;
25use super::utils::{TableCatalogBuilder, impl_distill_by_unit};
26use super::{ExprRewritable, PlanBase, PlanTreeNodeUnary, StreamNode, StreamPlanRef as PlanRef};
27use crate::TableCatalog;
28use crate::optimizer::plan_node::expr_visitable::ExprVisitable;
29use crate::optimizer::property::{MonotonicityMap, StreamKind, WatermarkColumns};
30use crate::stream_fragmenter::BuildFragmentGraphState;
31
32#[derive(Debug, Clone, PartialEq, Eq, Hash)]
33pub struct StreamEowcOverWindow {
34    pub base: PlanBase<Stream>,
35    core: generic::OverWindow<PlanRef>,
36}
37
38impl StreamEowcOverWindow {
39    pub fn new(core: generic::OverWindow<PlanRef>) -> Self {
40        assert!(core.funcs_have_same_partition_and_order());
41
42        let input = &core.input;
43        assert!(input.append_only());
44        assert!(input.emit_on_window_close());
45
46        // Should order by a single watermark column.
47        let order_key = &core.window_functions[0].order_by;
48        assert_eq!(order_key.len(), 1);
49        assert_eq!(order_key[0].order_type, OrderType::ascending());
50        let order_key_idx = order_key[0].column_index;
51        let input_watermark_cols = core.input.watermark_columns();
52        assert!(input_watermark_cols.contains(order_key_idx));
53
54        // `EowcOverWindowExecutor` cannot produce any watermark columns, because there may be some
55        // ancient rows in some rarely updated partitions that are emitted at the end of time.
56        let watermark_columns = WatermarkColumns::new();
57
58        let base = PlanBase::new_stream_with_core(
59            &core,
60            input.distribution().clone(),
61            StreamKind::AppendOnly,
62            true,
63            watermark_columns,
64            // we cannot derive monotonicity for any column for the same reason as watermark columns
65            MonotonicityMap::new(),
66        );
67        StreamEowcOverWindow { base, core }
68    }
69
70    fn window_functions(&self) -> &[PlanWindowFunction] {
71        &self.core.window_functions
72    }
73
74    fn partition_key_indices(&self) -> Vec<usize> {
75        self.window_functions()[0]
76            .partition_by
77            .iter()
78            .map(|i| i.index())
79            .collect()
80    }
81
82    fn order_key_index(&self) -> usize {
83        self.window_functions()[0].order_by[0].column_index
84    }
85
86    fn infer_state_table(&self) -> TableCatalog {
87        // The EOWC over window state table has the same schema as the input.
88
89        let in_fields = self.core.input.schema().fields();
90        let mut tbl_builder = TableCatalogBuilder::default();
91        for field in in_fields {
92            tbl_builder.add_column(field);
93        }
94
95        let mut order_cols = HashSet::new();
96        let partition_key_indices = self.partition_key_indices();
97        for idx in &partition_key_indices {
98            if !order_cols.contains(idx) {
99                tbl_builder.add_order_column(*idx, OrderType::ascending());
100                order_cols.insert(*idx);
101            }
102        }
103        let read_prefix_len_hint = tbl_builder.get_current_pk_len();
104        let order_key_index = self.order_key_index();
105        if !order_cols.contains(&order_key_index) {
106            tbl_builder.add_order_column(order_key_index, OrderType::ascending());
107            order_cols.insert(order_key_index);
108        }
109        for idx in self.core.input.expect_stream_key() {
110            if !order_cols.contains(idx) {
111                tbl_builder.add_order_column(*idx, OrderType::ascending());
112                order_cols.insert(*idx);
113            }
114        }
115
116        let in_dist_key = self.core.input.distribution().dist_column_indices();
117        tbl_builder.build(in_dist_key.to_vec(), read_prefix_len_hint)
118    }
119
120    /// Returns true if any window function requires intermediate state persistence.
121    /// This includes numbering functions (`row_number`, `rank`, `dense_rank`), and windows that may
122    /// be closed by watermark, whose output progress is persisted.
123    fn needs_intermediate_state_table(&self) -> bool {
124        self.window_functions().iter().any(|f| {
125            matches!(
126                f.kind,
127                WindowFuncKind::RowNumber | WindowFuncKind::Rank | WindowFuncKind::DenseRank
128            )
129        }) || self.needs_watermark_trigger_table()
130    }
131
132    /// Returns true if the window of any window function may be closed by the watermark of the
133    /// order key.
134    fn needs_watermark_trigger_table(&self) -> bool {
135        self.window_functions()
136            .iter()
137            .any(|f| f.frame.bounds.may_close_by_watermark())
138    }
139
140    /// Map the input distribution key to the given columns of partition key, which start at
141    /// `offset` in an internal table.
142    fn partition_dist_key(&self, offset: usize) -> Vec<usize> {
143        let partition_key_indices = self.partition_key_indices();
144        self.core
145            .input
146            .distribution()
147            .dist_column_indices()
148            .iter()
149            .filter_map(|&idx| {
150                partition_key_indices
151                    .iter()
152                    .position(|&pk_idx| pk_idx == idx)
153            })
154            .map(|pos| pos + offset)
155            .collect()
156    }
157
158    /// Infer the intermediate state table for persisting window function states.
159    ///
160    /// Schema: partition key columns + `state_0..state_{n-1}` (Bytea, one per window function
161    /// call) + `progress` (Bytea, only with the watermark trigger table).
162    /// PK: partition key columns (ascending) only.
163    fn infer_intermediate_state_table(&self) -> TableCatalog {
164        let in_fields = self.core.input.schema().fields();
165        let mut tbl_builder = TableCatalogBuilder::default();
166
167        // Add partition key columns
168        let partition_key_indices = self.partition_key_indices();
169        for &idx in &partition_key_indices {
170            tbl_builder.add_column(&in_fields[idx]);
171        }
172
173        // Add state columns: one per window function call (state_0, state_1, ..., state_{n-1})
174        let num_calls = self.window_functions().len();
175        for i in 0..num_calls {
176            tbl_builder.add_column(&Field::with_name(DataType::Bytea, format!("state_{}", i)));
177        }
178        if self.needs_watermark_trigger_table() {
179            tbl_builder.add_column(&Field::with_name(DataType::Bytea, "progress"));
180        }
181
182        // PK = partition key columns only (ascending)
183        for i in 0..partition_key_indices.len() {
184            tbl_builder.add_order_column(i, OrderType::ascending());
185        }
186        let read_prefix_len_hint = tbl_builder.get_current_pk_len();
187
188        // The partition key columns are added first at indices 0..partition_key_indices.len().
189        tbl_builder.build(self.partition_dist_key(0), read_prefix_len_hint)
190    }
191
192    /// Infer the table that indexes partitions by the watermark of the order key at which to check
193    /// them again, see `WatermarkTriggers` of the executor.
194    ///
195    /// Schema: `trigger_watermark` (order key type) + `trigger_exclusive` (Boolean) + partition key
196    /// columns. PK: all columns (ascending).
197    fn infer_watermark_trigger_table(&self) -> TableCatalog {
198        let in_fields = self.core.input.schema().fields();
199        let mut tbl_builder = TableCatalogBuilder::default();
200
201        let order_key_type = in_fields[self.order_key_index()].data_type();
202        tbl_builder.add_column(&Field::with_name(order_key_type, "trigger_watermark"));
203        tbl_builder.add_column(&Field::with_name(DataType::Boolean, "trigger_exclusive"));
204        let partition_key_indices = self.partition_key_indices();
205        for &idx in &partition_key_indices {
206            tbl_builder.add_column(&in_fields[idx]);
207        }
208        for i in 0..2 + partition_key_indices.len() {
209            tbl_builder.add_order_column(i, OrderType::ascending());
210        }
211
212        // Only range scans by the trigger watermark.
213        tbl_builder.build(self.partition_dist_key(2), 0)
214    }
215}
216
217impl_distill_by_unit!(StreamEowcOverWindow, core, "StreamEowcOverWindow");
218
219impl PlanTreeNodeUnary<Stream> for StreamEowcOverWindow {
220    fn input(&self) -> PlanRef {
221        self.core.input.clone()
222    }
223
224    fn clone_with_input(&self, input: PlanRef) -> Self {
225        let mut core = self.core.clone();
226        core.input = input;
227        Self::new(core)
228    }
229}
230impl_plan_tree_node_for_unary! { Stream, StreamEowcOverWindow }
231
232impl StreamNode for StreamEowcOverWindow {
233    fn to_stream_prost_body(&self, state: &mut BuildFragmentGraphState) -> PbNodeBody {
234        use risingwave_pb::stream_plan::*;
235
236        let calls = self
237            .window_functions()
238            .iter()
239            .map(PlanWindowFunction::to_protobuf)
240            .collect();
241        let partition_by = self.window_functions()[0]
242            .partition_by
243            .iter()
244            .map(|i| i.index() as _)
245            .collect();
246        let order_by = self.window_functions()[0]
247            .order_by
248            .iter()
249            .map(|o| o.to_protobuf())
250            .collect();
251        let state_table = self
252            .infer_state_table()
253            .with_id(state.gen_table_id_wrapped())
254            .to_internal_table_prost();
255
256        // Build intermediate state table if needed
257        let intermediate_state_table = if self.needs_intermediate_state_table() {
258            Some(
259                self.infer_intermediate_state_table()
260                    .with_id(state.gen_table_id_wrapped())
261                    .to_internal_table_prost(),
262            )
263        } else {
264            None
265        };
266
267        let watermark_trigger_table = self.needs_watermark_trigger_table().then(|| {
268            self.infer_watermark_trigger_table()
269                .with_id(state.gen_table_id_wrapped())
270                .to_internal_table_prost()
271        });
272
273        PbNodeBody::EowcOverWindow(Box::new(EowcOverWindowNode {
274            calls,
275            partition_by,
276            order_by,
277            state_table: Some(state_table),
278            intermediate_state_table,
279            watermark_trigger_table,
280            version: PbEowcOverWindowNodeVersion::LATEST as _,
281        }))
282    }
283}
284
285impl ExprRewritable<Stream> for StreamEowcOverWindow {}
286
287impl ExprVisitable for StreamEowcOverWindow {}
risingwave_frontend/optimizer/plan_node/stream_eowc_over_window.rs

risingwave_frontend/optimizer/plan_node/
stream_eowc_over_window.rs
