feat(optimizer): index accelerating TopN

e2e_test/batch/top_n/top_n_on_index.slt

@@ -0,0 +1,47 @@
+statement ok
+SET RW_IMPLICIT_FLUSH TO true;
+
+statement ok
+create table t(x int, y int);
+
+statement ok
+create index idx on t(x);
+
+statement ok
+insert into t values (100, 3), (1, 0), (2, 3), (3, 4), (5, 4);
+
+query II
+select * from t order by x limit 1;
+----
+1 0
+
+query II
+select * from t order by x limit 3;
+----
+1 0
+2 3
+3 4
+
+statement ok
+create table t1(x int primary key, y int);
+
+statement ok
+insert into t1 values (100, 3), (1, 0), (2, 3), (3, 4), (5, 4);
+
+query II
+select * from t1 order by x limit 1;
+----
+1 0
+
+query II
+select * from t1 order by x limit 3;
+----
+1 0
+2 3
+3 4
+
+statement ok
+drop table t;
+
+statement ok
+drop table t1;

proto/batch_plan.proto

@@ -24,6 +24,8 @@ message RowSeqScanNode {
   common.Buffer vnode_bitmap = 4;
   // Whether the order on output columns should be preserved.
   bool ordered = 5;
+
+  uint32 chunk_size = 6;
 }
 
 message SysRowSeqScanNode {

src/batch/src/executor/join/local_lookup_join.rs

@@ -114,6 +114,7 @@ impl<C: BatchTaskContext> InnerSideExecutorBuilder<C> {
             scan_ranges,
             ordered: false,
             vnode_bitmap: Some(vnode_bitmap.finish().to_protobuf()),
+            chunk_size: self.context.get_config().developer.batch_chunk_size as u32,
         });
 
         Ok(row_seq_scan_node)

src/batch/src/executor/row_seq_scan.rs

@@ -241,7 +241,12 @@ impl BoxedExecutorBuilder for RowSeqScanExecutorBuilder {
         let ordered = seq_scan_node.ordered;
 
         let epoch = source.epoch.clone();
-        let chunk_size = source.context.get_config().developer.batch_chunk_size;
+        let chunk_size = source
+            .context
+            .get_config()
+            .developer
+            .batch_chunk_size
+            .min(seq_scan_node.chunk_size as usize);
         let metrics = source.context().task_metrics();
 
         dispatch_state_store!(source.context().state_store(), state_store, {

src/frontend/planner_test/tests/testdata/explain.yaml

@@ -51,7 +51,7 @@
       "stages": {
         "0": {
           "root": {
-            "plan_node_id": 28,
+            "plan_node_id": 29,
             "plan_node_type": "BatchProject",
             "schema": [
               {
@@ -64,7 +64,7 @@
             ],
             "children": [
               {
-                "plan_node_id": 26,
+                "plan_node_id": 27,
                 "plan_node_type": "BatchValues",
                 "schema": [],
                 "children": [],

src/frontend/planner_test/tests/testdata/expr.yaml

@@ -536,7 +536,6 @@
   logical_plan: |
     LogicalProject { exprs: [Array(1:Int32, null:Int32)] }
     └─LogicalValues { rows: [[]], schema: Schema { fields: [] } }
-
 - name: regression (#7641) -fuzzing test failed at Bind error,types Boolean and Varchar cannot be matched
   sql: |
     select false >= 'LN1O0QP1yi' NOT IN (md5('4SeUPZhUbH'))
@@ -552,5 +551,4 @@
     └─BatchProject { exprs: [((1:Int32 + 2:Int32) + t.v1)] }
       └─BatchScan { table: t, columns: [t.v1], distribution: SomeShard }
 - name: const_eval of division by 0 error
-  sql: |
-    select 1 / 0 t1;
+  sql: select 1 / 0 t1;

src/frontend/planner_test/tests/testdata/index_selection.yaml

@@ -578,3 +578,23 @@
   batch_plan: |
     BatchExchange { order: [], dist: Single }
     └─BatchScan { table: idx1, columns: [idx1.a, idx1.b, idx1.c], scan_ranges: [idx1.a = Int32(1)], distribution: UpstreamHashShard(idx1.a) }
+- name: topn on index
+  sql: |
+    create table t1 (a int, b int);
+    create index idx1 on t1(a);
+    select * from t1 order by a limit 1
+  batch_plan: |
+    BatchLimit { limit: 1, offset: 0 }
+    └─BatchExchange { order: [idx1.a ASC], dist: Single }
+      └─BatchLimit { limit: 1, offset: 0 }
+        └─BatchScan { table: idx1, columns: [idx1.a, idx1.b], distribution: UpstreamHashShard(idx1.a) }
+- name: topn on primary key
+  sql: |
+    create table t1 (a int primary key, b int);
+    create index idx1 on t1(a);
+    select * from t1 order by a limit 1
+  batch_plan: |
+    BatchLimit { limit: 1, offset: 0 }
+    └─BatchExchange { order: [t1.a ASC], dist: Single }
+      └─BatchLimit { limit: 1, offset: 0 }
+        └─BatchScan { table: t1, columns: [t1.a, t1.b], distribution: UpstreamHashShard(t1.a) }

src/frontend/planner_test/tests/testdata/share.yaml

@@ -48,7 +48,7 @@
               | └─StreamProject { exprs: [id, _row_id] }
               |   └─StreamFilter { predicate: (initial_bid = 1:Int32) }
               |     └─StreamProject { exprs: [id, initial_bid, _row_id] }
-              |       └─StreamShare { id = 519 }
+              |       └─StreamShare { id = 539 }
               |         └─StreamProject { exprs: [id, initial_bid, _row_id] }
               |           └─StreamFilter { predicate: ((initial_bid = 1:Int32) OR (initial_bid = 2:Int32)) }
               |             └─StreamRowIdGen { row_id_index: 9 }
@@ -57,7 +57,7 @@
                 └─StreamProject { exprs: [id, _row_id] }
                   └─StreamFilter { predicate: (initial_bid = 2:Int32) }
                     └─StreamProject { exprs: [id, initial_bid, _row_id] }
-                      └─StreamShare { id = 519 }
+                      └─StreamShare { id = 539 }
                         └─StreamProject { exprs: [id, initial_bid, _row_id] }
                           └─StreamFilter { predicate: ((initial_bid = 1:Int32) OR (initial_bid = 2:Int32)) }
                             └─StreamRowIdGen { row_id_index: 9 }
@@ -125,7 +125,7 @@
           |     └─StreamExchange { dist: HashShard(auction, window_start) }
           |       └─StreamHopWindow { time_col: date_time, slide: 00:00:02, size: 00:00:10, output: [auction, window_start, _row_id] }
           |         └─StreamProject { exprs: [auction, date_time, _row_id] }
-          |           └─StreamShare { id = 744 }
+          |           └─StreamShare { id = 764 }
           |             └─StreamProject { exprs: [auction, date_time, _row_id] }
           |               └─StreamRowIdGen { row_id_index: 4 }
           |                 └─StreamSource { source: "bid", columns: ["auction", "bidder", "price", "date_time", "_row_id"] }
@@ -137,7 +137,7 @@
                     └─StreamExchange { dist: HashShard(auction, window_start) }
                       └─StreamHopWindow { time_col: date_time, slide: 00:00:02, size: 00:00:10, output: [auction, window_start, _row_id] }
                         └─StreamProject { exprs: [auction, date_time, _row_id] }
-                          └─StreamShare { id = 744 }
+                          └─StreamShare { id = 764 }
                             └─StreamProject { exprs: [auction, date_time, _row_id] }
                               └─StreamRowIdGen { row_id_index: 4 }
                                 └─StreamSource { source: "bid", columns: ["auction", "bidder", "price", "date_time", "_row_id"] }

src/frontend/planner_test/tests/testdata/shared_views.yaml

@@ -26,7 +26,7 @@
     └─StreamHashJoin { type: Inner, predicate: (t1.x + t1.y) = (t1.x * (t1.x + t1.y)), output: [(t1.x + t1.y), (t1.x * (t1.x + t1.y)), (t1.y * (t1.x + t1.y)), t1._row_id, t1._row_id, t1._row_id, t1.x, (t1.x + t1.y)] }
       ├─StreamExchange { dist: HashShard((t1.x + t1.y)) }
       | └─StreamProject { exprs: [(t1.x + t1.y), t1._row_id] }
-      |   └─StreamShare { id = 207 }
+      |   └─StreamShare { id = 212 }
       |     └─StreamProject { exprs: [(t1.x + t1.y), t1._row_id] }
       |       └─StreamFilter { predicate: (t1.y > 0:Int32) }
       |         └─StreamTableScan { table: t1, columns: [t1.x, t1.y, t1._row_id], pk: [t1._row_id], dist: UpstreamHashShard(t1._row_id) }
@@ -37,7 +37,7 @@
             | └─StreamTableScan { table: t1, columns: [t1.x, t1.y, t1._row_id], pk: [t1._row_id], dist: UpstreamHashShard(t1._row_id) }
             └─StreamExchange { dist: HashShard((t1.x + t1.y)) }
               └─StreamProject { exprs: [(t1.x + t1.y), t1._row_id] }
-                └─StreamShare { id = 207 }
+                └─StreamShare { id = 212 }
                   └─StreamProject { exprs: [(t1.x + t1.y), t1._row_id] }
                     └─StreamFilter { predicate: (t1.y > 0:Int32) }
                       └─StreamTableScan { table: t1, columns: [t1.x, t1.y, t1._row_id], pk: [t1._row_id], dist: UpstreamHashShard(t1._row_id) }

src/frontend/planner_test/tests/testdata/tpch.yaml

@@ -2573,7 +2573,7 @@
       |   ├─StreamExchange { dist: HashShard(supplier.s_suppkey) }
       |   | └─StreamTableScan { table: supplier, columns: [supplier.s_suppkey, supplier.s_name, supplier.s_address, supplier.s_phone], pk: [supplier.s_suppkey], dist: UpstreamHashShard(supplier.s_suppkey) }
       |   └─StreamProject { exprs: [lineitem.l_suppkey, sum((lineitem.l_extendedprice * (1:Int32 - lineitem.l_discount)))] }
-      |     └─StreamShare { id = 900 }
+      |     └─StreamShare { id = 921 }
       |       └─StreamProject { exprs: [lineitem.l_suppkey, sum((lineitem.l_extendedprice * (1:Int32 - lineitem.l_discount)))] }
       |         └─StreamHashAgg { group_key: [lineitem.l_suppkey], aggs: [count, sum((lineitem.l_extendedprice * (1:Int32 - lineitem.l_discount)))] }
       |           └─StreamExchange { dist: HashShard(lineitem.l_suppkey) }
@@ -2587,7 +2587,7 @@
               └─StreamHashAgg { group_key: [Vnode(lineitem.l_suppkey)], aggs: [count, max(sum((lineitem.l_extendedprice * (1:Int32 - lineitem.l_discount))))] }
                 └─StreamProject { exprs: [lineitem.l_suppkey, sum((lineitem.l_extendedprice * (1:Int32 - lineitem.l_discount))), Vnode(lineitem.l_suppkey)] }
                   └─StreamProject { exprs: [lineitem.l_suppkey, sum((lineitem.l_extendedprice * (1:Int32 - lineitem.l_discount)))] }
-                    └─StreamShare { id = 900 }
+                    └─StreamShare { id = 921 }
                       └─StreamProject { exprs: [lineitem.l_suppkey, sum((lineitem.l_extendedprice * (1:Int32 - lineitem.l_discount)))] }
                         └─StreamHashAgg { group_key: [lineitem.l_suppkey], aggs: [count, sum((lineitem.l_extendedprice * (1:Int32 - lineitem.l_discount)))] }
                           └─StreamExchange { dist: HashShard(lineitem.l_suppkey) }

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![TopNOnIndexRule::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/batch_seq_scan.rs

@@ -240,6 +240,7 @@ impl ToBatchProst for BatchSeqScan {
                 // To be filled by the scheduler.
                 vnode_bitmap: None,
                 ordered: !self.order().is_any(),
+                chunk_size: self.logical.get_chunk_size() as u32,
             })
         }
     }

src/frontend/src/optimizer/plan_node/generic/scan.rs

@@ -37,6 +37,8 @@ pub struct Scan {
     pub indexes: Vec<Rc<IndexCatalog>>,
     /// The pushed down predicates. It refers to column indexes of the table.
     pub predicate: Condition,
+    /// Help RowSeqScan executor use a better chunk size
+    pub chunk_size: usize,
 }
 
 impl Scan {

src/frontend/src/optimizer/plan_node/logical_scan.rs

@@ -20,7 +20,7 @@ use fixedbitset::FixedBitSet;
 use itertools::Itertools;
 use risingwave_common::catalog::{ColumnDesc, Field, Schema, TableDesc};
 use risingwave_common::error::{ErrorCode, Result, RwError};
-use risingwave_common::util::sort_util::OrderType;
+use risingwave_common::util::sort_util::{OrderPair, OrderType};
 
 use super::generic::{GenericPlanNode, GenericPlanRef};
 use super::{
@@ -86,6 +86,7 @@ impl LogicalScan {
             table_desc,
             indexes,
             predicate,
+            chunk_size: 1024,
         };
 
         let schema = core.schema();
@@ -300,6 +301,20 @@ impl LogicalScan {
         )
     }
 
+    /// used by optimizer (currently `top_n_on_index_rule`) to help reduce useless `chunk_size` at
+    /// executor
+    pub fn set_chunk_size(&mut self, chunk_size: usize) {
+        self.core.chunk_size = chunk_size;
+    }
+
+    pub fn get_chunk_size(&self) -> usize {
+        self.core.chunk_size
+    }
+
+    pub fn get_primary_key(&self) -> Vec<OrderPair> {
+        self.core.table_desc.pk.clone()
+    }
+
     /// a vec of `InputRef` corresponding to `output_col_idx`, which can represent a pulled project.
     fn output_idx_to_input_ref(&self) -> Vec<ExprImpl> {
         let output_idx = self

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 top_n_on_index_rule;
+pub use top_n_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}
+            ,{TopNOnIndexRule}
         }
     };
 }