risingwave_frontend/optimizer/plan_node/
stream_group_topn.rs

1// Copyright 2022 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 pretty_xmlish::XmlNode;
16use risingwave_common::util::sort_util::topn_watermark_forwardable_order_key;
17use risingwave_pb::stream_plan::stream_node::PbNodeBody;
18
19use super::generic::{DistillUnit, TopNLimit};
20use super::stream::prelude::*;
21use super::utils::{Distill, plan_node_name, watermark_pretty};
22use super::{
23    ExprRewritable, PlanBase, PlanTreeNodeUnary, StreamNode, StreamPlanRef as PlanRef, generic,
24};
25use crate::error::Result;
26use crate::optimizer::plan_node::expr_visitable::ExprVisitable;
27use crate::optimizer::plan_node::generic::GenericPlanNode;
28use crate::optimizer::property::{MonotonicityMap, Order};
29use crate::stream_fragmenter::BuildFragmentGraphState;
30
31#[derive(Debug, Clone, PartialEq, Eq, Hash)]
32pub struct StreamGroupTopN {
33    pub base: PlanBase<Stream>,
34    core: generic::TopN<PlanRef>,
35    /// an optional column index which is the vnode of each row computed by the input's consistent
36    /// hash distribution
37    vnode_col_idx: Option<usize>,
38}
39
40impl StreamGroupTopN {
41    pub fn new(core: generic::TopN<PlanRef>, vnode_col_idx: Option<usize>) -> Result<Self> {
42        assert!(!core.group_key.is_empty());
43        assert!(core.limit_attr.limit() > 0);
44        reject_upsert_input!(core.input);
45
46        let input = &core.input;
47
48        // The executor forwards watermarks on group key columns, and on the first `ORDER BY`
49        // column if it's ordered `ASC NULLS LAST`. Keep the optimizer in sync so EOWC won't be
50        // enabled on unsupported plans. See `topn_watermark_forwardable_order_key` for the
51        // reasoning.
52        let watermark_columns = {
53            let mut cols = core.group_key.clone();
54            cols.extend(topn_watermark_forwardable_order_key(
55                &core.order.column_orders,
56            ));
57            input.watermark_columns().retain_clone(&cols)
58        };
59
60        let mut stream_key = core
61            .stream_key()
62            .expect("logical node should have stream key here");
63        if let Some(vnode_col_idx) = vnode_col_idx
64            && stream_key.len() > 1
65        {
66            // The output stream key of `GroupTopN` is a union of group key and input stream key,
67            // while vnode is calculated from a subset of input stream key. So we can safely remove
68            // the vnode column from output stream key. While at meanwhile we cannot leave the stream key
69            // as empty, so we only remove it when stream key length is > 1.
70            stream_key.remove(stream_key.iter().position(|i| *i == vnode_col_idx).unwrap());
71        }
72
73        let base = PlanBase::new_stream(
74            core.ctx(),
75            core.schema(),
76            Some(stream_key),
77            core.functional_dependency(),
78            input.distribution().clone(),
79            StreamKind::Retract,
80            // TODO: https://github.com/risingwavelabs/risingwave/issues/8348
81            false,
82            watermark_columns,
83            MonotonicityMap::new(), // TODO: derive monotonicity
84        );
85
86        Ok(StreamGroupTopN {
87            base,
88            core,
89            vnode_col_idx,
90        })
91    }
92
93    pub fn limit_attr(&self) -> TopNLimit {
94        self.core.limit_attr
95    }
96
97    pub fn offset(&self) -> u64 {
98        self.core.offset
99    }
100
101    pub fn topn_order(&self) -> &Order {
102        &self.core.order
103    }
104
105    pub fn group_key(&self) -> &[usize] {
106        &self.core.group_key
107    }
108}
109
110impl StreamNode for StreamGroupTopN {
111    fn to_stream_prost_body(&self, state: &mut BuildFragmentGraphState) -> PbNodeBody {
112        use risingwave_pb::stream_plan::*;
113
114        let input = self.input();
115        let table = self
116            .core
117            .infer_internal_table_catalog(
118                input.schema(),
119                input.ctx(),
120                input.expect_stream_key(),
121                self.vnode_col_idx,
122            )
123            .with_id(state.gen_table_id_wrapped());
124        assert!(!self.group_key().is_empty());
125        let group_topn_node = GroupTopNNode {
126            limit: self.limit_attr().limit(),
127            offset: self.offset(),
128            with_ties: self.limit_attr().with_ties(),
129            group_key: self.group_key().iter().map(|idx| *idx as u32).collect(),
130            table: Some(table.to_internal_table_prost()),
131            order_by: self.topn_order().to_protobuf(),
132        };
133        if self.input().append_only() {
134            PbNodeBody::AppendOnlyGroupTopN(Box::new(group_topn_node))
135        } else {
136            PbNodeBody::GroupTopN(Box::new(group_topn_node))
137        }
138    }
139}
140
141impl Distill for StreamGroupTopN {
142    fn distill<'a>(&self) -> XmlNode<'a> {
143        let name = plan_node_name!("StreamGroupTopN",
144            { "append_only", self.input().append_only() },
145        );
146        let mut node = self.core.distill_with_name(name);
147        if let Some(ow) = watermark_pretty(self.base.watermark_columns(), self.schema()) {
148            node.fields.push(("output_watermarks".into(), ow));
149        }
150        node
151    }
152}
153
154impl_plan_tree_node_for_unary! { Stream, StreamGroupTopN }
155
156impl PlanTreeNodeUnary<Stream> for StreamGroupTopN {
157    fn input(&self) -> PlanRef {
158        self.core.input.clone()
159    }
160
161    fn clone_with_input(&self, input: PlanRef) -> Self {
162        let mut core = self.core.clone();
163        core.input = input;
164        Self::new(core, self.vnode_col_idx).unwrap()
165    }
166}
167
168impl ExprRewritable<Stream> for StreamGroupTopN {}
169
170impl ExprVisitable for StreamGroupTopN {}
risingwave_frontend/optimizer/plan_node/stream_group_topn.rs

risingwave_frontend/optimizer/plan_node/
stream_group_topn.rs
