risingwave_frontend/optimizer/plan_node/
stream_sort.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::{Pretty, XmlNode};
18use risingwave_common::catalog::FieldDisplay;
19use risingwave_common::util::sort_util::OrderType;
20use risingwave_pb::stream_plan::stream_node::PbNodeBody;
21
22use super::stream::prelude::*;
23use super::utils::{Distill, TableCatalogBuilder, childless_record};
24use super::{ExprRewritable, PlanBase, PlanTreeNodeUnary, StreamNode, StreamPlanRef as PlanRef};
25use crate::TableCatalog;
26use crate::optimizer::plan_node::expr_visitable::ExprVisitable;
27use crate::optimizer::property::{Monotonicity, MonotonicityMap, WatermarkColumns};
28use crate::stream_fragmenter::BuildFragmentGraphState;
29
30#[derive(Debug, Clone, PartialEq, Eq, Hash)]
31pub struct StreamEowcSort {
32    pub base: PlanBase<Stream>,
33
34    input: PlanRef,
35    sort_column_index: usize,
36    /// See [`Self::with_secondary_order`]. Empty for plain watermark sorting.
37    secondary_order_columns: Vec<usize>,
38}
39
40impl Distill for StreamEowcSort {
41    fn distill<'a>(&self) -> XmlNode<'a> {
42        let mut fields = vec![(
43            "sort_column",
44            Pretty::display(&FieldDisplay(&self.input.schema()[self.sort_column_index])),
45        )];
46        // EXPLAIN must show the full order the sort actually enforces, not just the leading
47        // watermark column.
48        if !self.secondary_order_columns.is_empty() {
49            fields.push((
50                "secondary_order_columns",
51                Pretty::Array(
52                    self.secondary_order_columns
53                        .iter()
54                        .map(|&i| Pretty::display(&FieldDisplay(&self.input.schema()[i])))
55                        .collect(),
56                ),
57            ));
58        }
59        childless_record("StreamEowcSort", fields)
60    }
61}
62
63impl StreamEowcSort {
64    pub fn new(input: PlanRef, sort_column_index: usize) -> Self {
65        Self::with_secondary_order(input, sort_column_index, vec![])
66    }
67
68    /// Like [`Self::new`], with additional order columns appended to the buffer table's key right
69    /// after the sort column. The executor emits rows in `(sort_column, buffer-table key)` order,
70    /// so this makes the emission order `(sort_column, secondary columns, ...)` — what an
71    /// order-sensitive consumer with a multi-column ORDER BY (e.g. `MATCH_RECOGNIZE`) requires.
72    /// No executor or proto change: the order is carried entirely by the inferred table key.
73    pub fn with_secondary_order(
74        input: PlanRef,
75        sort_column_index: usize,
76        secondary_order_columns: Vec<usize>,
77    ) -> Self {
78        assert!(input.watermark_columns().contains(sort_column_index));
79
80        let schema = input.schema().clone();
81        let stream_key = input.stream_key().map(|v| v.to_vec());
82        let fd_set = input.functional_dependency().clone();
83        let dist = input.distribution().clone();
84
85        let mut watermark_columns = WatermarkColumns::new();
86        watermark_columns.insert(
87            sort_column_index,
88            // `StreamSort` operator will propagate input watermark as it is,
89            // so we can assign the same watermark group.
90            input
91                .watermark_columns()
92                .get_group(sort_column_index)
93                .unwrap(),
94        );
95
96        // StreamEowcSort makes the sorting watermark column non-decreasing
97        let mut columns_monotonicity = MonotonicityMap::new();
98        columns_monotonicity.insert(sort_column_index, Monotonicity::NonDecreasing);
99
100        let base = PlanBase::new_stream(
101            input.ctx(),
102            schema,
103            stream_key,
104            fd_set,
105            dist,
106            StreamKind::AppendOnly,
107            true,
108            watermark_columns,
109            columns_monotonicity,
110        );
111        Self {
112            base,
113            input,
114            sort_column_index,
115            secondary_order_columns,
116        }
117    }
118
119    fn infer_state_table(&self) -> TableCatalog {
120        // The sort state table has the same schema as the input.
121
122        let in_fields = self.input.schema().fields();
123        let mut tbl_builder = TableCatalogBuilder::default();
124        for field in in_fields {
125            tbl_builder.add_column(field);
126        }
127
128        let mut order_cols = HashSet::new();
129        tbl_builder.add_order_column(self.sort_column_index, OrderType::ascending());
130        order_cols.insert(self.sort_column_index);
131
132        // Secondary order columns go right after the sort column and before the distribution and
133        // stream keys: the executor emits in `(sort_column, table key)` order, so this is what
134        // makes the emission order the caller's full ORDER BY (see `with_secondary_order`).
135        for idx in &self.secondary_order_columns {
136            if !order_cols.contains(idx) {
137                tbl_builder.add_order_column(*idx, OrderType::ascending());
138                order_cols.insert(*idx);
139            }
140        }
141
142        let dist_key = self.base.distribution().dist_column_indices().to_vec();
143        for idx in &dist_key {
144            if !order_cols.contains(idx) {
145                tbl_builder.add_order_column(*idx, OrderType::ascending());
146                order_cols.insert(*idx);
147            }
148        }
149
150        for idx in self.input.expect_stream_key() {
151            if !order_cols.contains(idx) {
152                tbl_builder.add_order_column(*idx, OrderType::ascending());
153                order_cols.insert(*idx);
154            }
155        }
156
157        let read_prefix_len_hint = 0;
158        tbl_builder.build(dist_key, read_prefix_len_hint)
159    }
160}
161
162impl PlanTreeNodeUnary<Stream> for StreamEowcSort {
163    fn input(&self) -> PlanRef {
164        self.input.clone()
165    }
166
167    fn clone_with_input(&self, input: PlanRef) -> Self {
168        Self::with_secondary_order(
169            input,
170            self.sort_column_index,
171            self.secondary_order_columns.clone(),
172        )
173    }
174}
175
176impl_plan_tree_node_for_unary! { Stream, StreamEowcSort }
177
178impl StreamNode for StreamEowcSort {
179    fn to_stream_prost_body(&self, state: &mut BuildFragmentGraphState) -> PbNodeBody {
180        use risingwave_pb::stream_plan::*;
181        PbNodeBody::Sort(Box::new(SortNode {
182            state_table: Some(
183                self.infer_state_table()
184                    .with_id(state.gen_table_id_wrapped())
185                    .to_internal_table_prost(),
186            ),
187            sort_column_index: self.sort_column_index as _,
188        }))
189    }
190}
191
192impl ExprRewritable<Stream> for StreamEowcSort {}
193
194impl ExprVisitable for StreamEowcSort {}
risingwave_frontend/optimizer/plan_node/stream_sort.rs

risingwave_frontend/optimizer/plan_node/
stream_sort.rs
