risingwave_frontend/optimizer/rule/
top_n_project_transpose_rule.rs

1// Copyright 2026 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 risingwave_common::util::sort_util::ColumnOrder;
16
17use super::prelude::{PlanRef, *};
18use crate::optimizer::plan_node::{LogicalProject, LogicalTopN, PlanTreeNodeUnary};
19use crate::optimizer::property::Order;
20
21/// Transpose `LogicalTopN` and `LogicalProject` when the project is a pure projection.
22pub struct TopNProjectTransposeRule {}
23
24impl Rule<Logical> for TopNProjectTransposeRule {
25    fn apply(&self, plan: PlanRef) -> Option<PlanRef> {
26        let top_n: &LogicalTopN = plan.as_logical_top_n()?;
27        let project: LogicalProject = top_n.input().as_logical_project()?.to_owned();
28        let projection = project.try_as_projection()?;
29
30        let new_order = Order {
31            column_orders: top_n
32                .topn_order()
33                .column_orders
34                .iter()
35                .map(|order| {
36                    let mapped_idx = projection.get(order.column_index).copied()?;
37                    Some(ColumnOrder::new(mapped_idx, order.order_type))
38                })
39                .collect::<Option<Vec<_>>>()?,
40        };
41        let new_group_key = top_n
42            .group_key()
43            .iter()
44            .map(|idx| projection.get(*idx).copied())
45            .collect::<Option<Vec<_>>>()?;
46
47        let limit_attr = top_n.limit_attr();
48        let new_top_n = LogicalTopN::new(
49            project.input(),
50            limit_attr.limit(),
51            top_n.offset(),
52            limit_attr.with_ties(),
53            new_order,
54            new_group_key,
55        );
56        Some(project.clone_with_input(new_top_n.into()).into())
57    }
58}
59
60impl TopNProjectTransposeRule {
61    pub fn create() -> BoxedRule {
62        Box::new(TopNProjectTransposeRule {})
63    }
64}