risingwave_frontend/optimizer/plan_node/
stream_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::util::sort_util::{ColumnOrder, OrderType};
18use risingwave_expr::window_function::FrameBounds;
19use risingwave_pb::stream_plan::stream_node::PbNodeBody;
20
21use super::generic::{GenericPlanNode, PlanWindowFunction};
22use super::stream::prelude::*;
23use super::utils::{TableCatalogBuilder, impl_distill_by_unit};
24use super::{
25    ExprRewritable, PlanBase, PlanTreeNodeUnary, StreamNode, StreamPlanRef as PlanRef, generic,
26};
27use crate::TableCatalog;
28use crate::optimizer::plan_node::expr_visitable::ExprVisitable;
29use crate::optimizer::property::{MonotonicityMap, WatermarkColumns};
30use crate::stream_fragmenter::BuildFragmentGraphState;
31
32#[derive(Debug, Clone, PartialEq, Eq, Hash)]
33pub struct StreamOverWindow {
34    pub base: PlanBase<Stream>,
35    core: generic::OverWindow<PlanRef>,
36}
37
38impl StreamOverWindow {
39    pub fn new(core: generic::OverWindow<PlanRef>) -> Result<Self> {
40        assert!(core.funcs_have_same_partition_and_order());
41        reject_upsert_input!(core.input);
42
43        let input = &core.input;
44        let watermark_columns = WatermarkColumns::new();
45
46        let base = PlanBase::new_stream_with_core(
47            &core,
48            input.distribution().clone(),
49            StreamKind::Retract, // general over window cannot be append-only
50            false,
51            watermark_columns,
52            MonotonicityMap::new(), // TODO: derive monotonicity
53        );
54
55        Ok(StreamOverWindow { base, core })
56    }
57
58    fn infer_state_table(&self) -> TableCatalog {
59        let mut tbl_builder = TableCatalogBuilder::default();
60
61        let out_schema = self.core.schema();
62        for field in out_schema.fields() {
63            tbl_builder.add_column(field);
64        }
65
66        let mut order_cols = HashSet::new();
67        for idx in self.core.partition_key_indices() {
68            if order_cols.insert(idx) {
69                tbl_builder.add_order_column(idx, OrderType::ascending());
70            }
71        }
72        let read_prefix_len_hint = tbl_builder.get_current_pk_len();
73        for o in self.core.order_key() {
74            if order_cols.insert(o.column_index) {
75                tbl_builder.add_order_column(o.column_index, o.order_type);
76            }
77        }
78        for &idx in self.core.input.expect_stream_key() {
79            if order_cols.insert(idx) {
80                tbl_builder.add_order_column(idx, OrderType::ascending());
81            }
82        }
83
84        let in_dist_key = self.core.input.distribution().dist_column_indices();
85        tbl_builder.build(in_dist_key.to_vec(), read_prefix_len_hint)
86    }
87
88    /// Whether the executor is allowed to clean up state rows below the watermark of the first
89    /// `ORDER BY` column. See `OverWindowExecutor` for the cleaning strategy.
90    ///
91    /// The cleaning is only correct when:
92    /// - the input is append-only, so no existing row will ever be updated or deleted;
93    /// - all window frames are bounded `ROWS` frames, so a row can only affect (and be affected
94    ///   by) a bounded number of neighboring rows;
95    /// - the first `ORDER BY` column is a watermark column with NULLs ordered as the largest
96    ///   values, so that once a watermark is received, rows below it can never get new neighbors
97    ///   on the "smaller" side, and all new rows (including NULLs) land on the "larger" side.
98    fn state_cleaning_enabled(&self) -> bool {
99        let input = &self.core.input;
100        let Some(first_order_key) = self.core.order_key().first() else {
101            return false;
102        };
103        input.append_only()
104            && self.core.window_functions().iter().all(|func| {
105                matches!(&func.frame.bounds, FrameBounds::Rows(bounds)
106                    if !bounds.start.is_unbounded_preceding() && !bounds.end.is_unbounded_following())
107            })
108            && first_order_key.order_type.nulls_are_largest()
109            && input.watermark_columns().contains(first_order_key.column_index)
110            // the first order key column must be part of the state table sub-PK following the
111            // partition key, so that the executor can scan stale rows with it
112            && !self
113                .core
114                .partition_key_indices()
115                .contains(&first_order_key.column_index)
116    }
117}
118
119impl_distill_by_unit!(StreamOverWindow, core, "StreamOverWindow");
120
121impl PlanTreeNodeUnary<Stream> for StreamOverWindow {
122    fn input(&self) -> PlanRef {
123        self.core.input.clone()
124    }
125
126    fn clone_with_input(&self, input: PlanRef) -> Self {
127        let mut core = self.core.clone();
128        core.input = input;
129        Self::new(core).unwrap()
130    }
131}
132impl_plan_tree_node_for_unary! { Stream, StreamOverWindow }
133
134impl StreamNode for StreamOverWindow {
135    fn to_stream_prost_body(&self, state: &mut BuildFragmentGraphState) -> PbNodeBody {
136        use risingwave_pb::stream_plan::*;
137
138        let calls = self
139            .core
140            .window_functions()
141            .iter()
142            .map(PlanWindowFunction::to_protobuf)
143            .collect();
144        let partition_by = self
145            .core
146            .partition_key_indices()
147            .into_iter()
148            .map(|idx| idx as _)
149            .collect();
150        let order_by = self
151            .core
152            .order_key()
153            .iter()
154            .copied()
155            .map(ColumnOrder::to_protobuf)
156            .collect();
157        let state_table = self
158            .infer_state_table()
159            .with_id(state.gen_table_id_wrapped())
160            .to_internal_table_prost();
161
162        PbNodeBody::OverWindow(Box::new(OverWindowNode {
163            calls,
164            partition_by,
165            order_by,
166            state_table: Some(state_table),
167
168            // Cache policy should now be read from per-job config override.
169            #[expect(deprecated)]
170            cache_policy: PbOverWindowCachePolicy::Unspecified as _,
171
172            enable_state_cleaning: self.state_cleaning_enabled(),
173        }))
174    }
175}
176
177impl ExprRewritable<Stream> for StreamOverWindow {}
178
179impl ExprVisitable for StreamOverWindow {}
risingwave_frontend/optimizer/plan_node/stream_over_window.rs

risingwave_frontend/optimizer/plan_node/
stream_over_window.rs
