risingwave_frontend/optimizer/rule/
top_n_project_transpose_rule.rs1use 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
21pub 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}