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

risingwave_frontend/optimizer/plan_node/
stream_over_window.rs
