feat(optimizer): index accelerating TopN

src/frontend/src/optimizer/mod.rs

@@ -429,6 +429,13 @@ impl PlanRoot {
             ApplyOrder::TopDown,
         );
 
+        plan = self.optimize_by_rules(
+            plan,
+            "Agg on Index".to_string(),
+            vec![AggOnIndexRule::create()],
+            ApplyOrder::TopDown,
+        );
+
         #[cfg(debug_assertions)]
         InputRefValidator.validate(plan.clone());
 

src/frontend/src/optimizer/plan_node/batch_limit.rs

@@ -49,8 +49,15 @@ impl BatchLimit {
         let new_offset = 0;
         let logical_partial_limit = LogicalLimit::new(input, new_limit, new_offset);
         let batch_partial_limit = Self::new(logical_partial_limit);
-        let ensure_single_dist = RequiredDist::single()
-            .enforce_if_not_satisfies(batch_partial_limit.into(), &Order::any())?;
+        let any_order = Order::any();
+        let ensure_single_dist = RequiredDist::single().enforce_if_not_satisfies(
+            batch_partial_limit.into(),
+            if self.order().field_order.is_empty() {
+                &any_order
+            } else {
+                &self.order()
+            },
+        )?;
         let batch_global_limit = self.clone_with_input(ensure_single_dist);
         Ok(batch_global_limit.into())
     }

src/frontend/src/optimizer/plan_node/logical_topn.rs

@@ -122,6 +122,11 @@ impl LogicalTopN {
         &self.core.group_key
     }
 
+    // For Optimizing TopN
+    pub fn get_child(&self) -> PlanRef {
+        self.core.input.clone()
+    }
+
     pub(super) fn fmt_with_name(&self, f: &mut fmt::Formatter<'_>, name: &str) -> fmt::Result {
         let mut builder = f.debug_struct(name);
         let input = self.input();

src/frontend/src/optimizer/rule/agg_on_index_rule.rs

@@ -0,0 +1,60 @@
+use super::{BoxedRule, Rule};
+use crate::optimizer::plan_node::{LogicalLimit, LogicalScan, LogicalTopN};
+use crate::optimizer::property::Direction::Asc;
+use crate::optimizer::property::{FieldOrder, Order};
+use crate::optimizer::PlanRef;
+
+pub struct AggOnIndexRule {}
+
+impl Rule for AggOnIndexRule {
+    fn apply(&self, plan: PlanRef) -> Option<PlanRef> {
+        let logical_topn: &LogicalTopN = plan.as_logical_top_n()?;
+        let logical_scan: LogicalScan = logical_topn.get_child().as_logical_scan()?.to_owned();
+        let order = logical_topn.topn_order();
+        if order.field_order.is_empty() {
+            return None;
+        }
+        let index = logical_scan.indexes().iter().find(|idx| {
+            Order {
+                field_order: idx
+                    .index_item
+                    .iter()
+                    .map(|idx_item| FieldOrder {
+                        index: idx_item.index,
+                        direct: Asc,
+                    })
+                    .collect(),
+            }
+            .satisfies(order)
+        })?;
+
+        let p2s_mapping = index.primary_to_secondary_mapping();
+
+        let index_scan = if logical_scan
+            .required_col_idx()
+            .iter()
+            .all(|x| p2s_mapping.contains_key(x))
+        {
+            Some(logical_scan.to_index_scan(
+                &index.name,
+                index.index_table.table_desc().into(),
+                p2s_mapping,
+            ))
+        } else {
+            None
+        }?;
+
+        let logical_limit = LogicalLimit::create(
+            index_scan.into(),
+            logical_topn.limit(),
+            logical_topn.offset(),
+        );
+        Some(logical_limit)
+    }
+}
+
+impl AggOnIndexRule {
+    pub fn create() -> BoxedRule {
+        Box::new(AggOnIndexRule {})
+    }
+}

src/frontend/src/optimizer/rule/mod.rs

@@ -81,6 +81,8 @@ mod dag_to_tree_rule;
 pub use dag_to_tree_rule::*;
 mod apply_share_eliminate_rule;
 pub use apply_share_eliminate_rule::*;
+mod agg_on_index_rule;
+pub use agg_on_index_rule::*;
 mod stream;
 pub use stream::filter_with_now_to_join_rule::*;
 
@@ -115,6 +117,7 @@ macro_rules! for_all_rules {
             ,{DagToTreeRule}
             ,{AggDedupGroupKeyRule}
             ,{FilterWithNowToJoinRule}
+            ,{AggOnIndexRule}
         }
     };
 }