risingwave_frontend/optimizer/property/
order.rs

1// Copyright 2025 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::fmt;
16
17use itertools::Itertools;
18use pretty_xmlish::Pretty;
19use risingwave_common::catalog::Schema;
20use risingwave_common::util::sort_util::{ColumnOrder, ColumnOrderDisplay};
21use risingwave_pb::common::PbColumnOrder;
22
23use super::super::plan_node::*;
24use crate::error::Result;
25
26#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)]
27pub struct Order {
28    pub column_orders: Vec<ColumnOrder>,
29}
30
31impl Order {
32    pub const fn new(column_orders: Vec<ColumnOrder>) -> Self {
33        Self { column_orders }
34    }
35
36    pub fn to_protobuf(&self) -> Vec<PbColumnOrder> {
37        self.column_orders
38            .iter()
39            .copied()
40            .map(ColumnOrder::to_protobuf)
41            .collect_vec()
42    }
43
44    #[allow(clippy::len_without_is_empty)]
45    pub fn len(&self) -> usize {
46        self.column_orders.len()
47    }
48
49    pub fn concat(self, other: Self) -> Self {
50        Self {
51            column_orders: self
52                .column_orders
53                .into_iter()
54                .chain(other.column_orders)
55                .collect(),
56        }
57    }
58}
59
60impl fmt::Display for Order {
61    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
62        write!(f, "[")?;
63        for (i, column_order) in self.column_orders.iter().enumerate() {
64            if i > 0 {
65                write!(f, ", ")?;
66            }
67            write!(f, "{}", column_order)?;
68        }
69        write!(f, "]")
70    }
71}
72
73pub struct OrderDisplay<'a> {
74    pub order: &'a Order,
75    pub input_schema: &'a Schema,
76}
77
78impl OrderDisplay<'_> {
79    pub fn distill<'a>(self) -> Pretty<'a> {
80        let iter = self.order.column_orders.iter();
81        let vec = iter.map(|column_order| {
82            Pretty::display(&ColumnOrderDisplay {
83                column_order,
84                input_schema: self.input_schema,
85            })
86        });
87        Pretty::Array(vec.collect())
88    }
89}
90
91const ANY_ORDER: Order = Order {
92    column_orders: vec![],
93};
94
95impl Order {
96    pub fn enforce_if_not_satisfies(&self, plan: BatchPlanRef) -> Result<BatchPlanRef> {
97        use crate::optimizer::plan_node::batch::prelude::*;
98
99        if !plan.order().satisfies(self) {
100            Ok(self.enforce(plan))
101        } else {
102            Ok(plan)
103        }
104    }
105
106    fn enforce(&self, plan: BatchPlanRef) -> BatchPlanRef {
107        BatchSort::new(plan, self.clone()).into()
108    }
109
110    pub fn satisfies(&self, other: &Order) -> bool {
111        if self.column_orders.len() < other.column_orders.len() {
112            return false;
113        }
114        #[expect(clippy::disallowed_methods)]
115        for (order, other_order) in self.column_orders.iter().zip(other.column_orders.iter()) {
116            if order != other_order {
117                return false;
118            }
119        }
120        true
121    }
122
123    #[inline(always)]
124    pub const fn any() -> Self {
125        ANY_ORDER
126    }
127
128    #[inline(always)]
129    pub fn is_any(&self) -> bool {
130        self.column_orders.is_empty()
131    }
132}
133
134#[cfg(test)]
135mod tests {
136    use risingwave_common::util::sort_util::{ColumnOrder, OrderType};
137
138    use super::Order;
139
140    #[test]
141    fn test_order_satisfy() {
142        let o1 = Order {
143            column_orders: vec![
144                ColumnOrder {
145                    column_index: 0,
146                    order_type: OrderType::ascending(),
147                },
148                ColumnOrder {
149                    column_index: 1,
150                    order_type: OrderType::descending(),
151                },
152                ColumnOrder {
153                    column_index: 2,
154                    order_type: OrderType::ascending(),
155                },
156            ],
157        };
158        let o2 = Order {
159            column_orders: vec![
160                ColumnOrder {
161                    column_index: 0,
162                    order_type: OrderType::ascending(),
163                },
164                ColumnOrder {
165                    column_index: 1,
166                    order_type: OrderType::descending(),
167                },
168            ],
169        };
170        let o3 = Order {
171            column_orders: vec![
172                ColumnOrder {
173                    column_index: 0,
174                    order_type: OrderType::ascending(),
175                },
176                ColumnOrder {
177                    column_index: 1,
178                    order_type: OrderType::ascending(),
179                },
180            ],
181        };
182
183        assert!(o1.satisfies(&o2));
184        assert!(!o2.satisfies(&o1));
185        assert!(o1.satisfies(&o1));
186
187        assert!(!o2.satisfies(&o3));
188        assert!(!o3.satisfies(&o2));
189    }
190}
risingwave_frontend/optimizer/property/order.rs

risingwave_frontend/optimizer/property/
order.rs
