Merge pull request #9092 from xudong963/topk

feat: optimize topk in cluser mode

Author: mergify[bot]

src/query/sql/src/planner/optimizer/distributed/distributed.rs

@@ -12,23 +12,28 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::cmp;
 use std::sync::Arc;
 
 use common_catalog::table_context::TableContext;
 use common_exception::Result;
 
+use crate::optimizer::distributed::topk::TopK;
 use crate::optimizer::property::require_property;
 use crate::optimizer::Distribution;
 use crate::optimizer::RelExpr;
 use crate::optimizer::RequiredProperty;
 use crate::optimizer::SExpr;
 use crate::plans::Exchange;
+use crate::plans::Limit;
+use crate::plans::RelOperator;
 
 pub fn optimize_distributed_query(ctx: Arc<dyn TableContext>, s_expr: &SExpr) -> Result<SExpr> {
     let required = RequiredProperty {
         distribution: Distribution::Any,
     };
     let mut result = require_property(ctx, &required, s_expr)?;
+    push_down_topk_to_merge(&mut result, None)?;
     let rel_expr = RelExpr::with_s_expr(&result);
     let physical_prop = rel_expr.derive_physical_prop()?;
     let root_required = RequiredProperty {
@@ -41,3 +46,36 @@ pub fn optimize_distributed_query(ctx: Arc<dyn TableContext>, s_expr: &SExpr) ->
 
     Ok(result)
 }
+
+// Traverse the SExpr tree to find top_k, if find, push down it to Exchange::Merge
+fn push_down_topk_to_merge(s_expr: &mut SExpr, mut top_k: Option<TopK>) -> Result<()> {
+    if let RelOperator::Exchange(Exchange::Merge) = s_expr.plan {
+        if let Some(top_k) = top_k {
+            let child = &mut s_expr.children[0];
+            *child = SExpr::create_unary(top_k.sort.into(), child.clone());
+            *child = SExpr::create_unary(top_k.limit.into(), child.clone());
+        }
+        return Ok(());
+    }
+    for child in s_expr.children.iter_mut() {
+        if let RelOperator::Sort(sort) = &child.plan {
+            if let RelOperator::Limit(limit) = &s_expr.plan {
+                // If limit.limit is None, no need to push down.
+                if let Some(mut count) = limit.limit {
+                    count += limit.offset;
+                    top_k = Some(TopK {
+                        sort: sort.clone(),
+                        limit: Limit {
+                            limit: limit
+                                .limit
+                                .map(|origin_limit| cmp::max(origin_limit, count)),
+                            offset: 0,
+                        },
+                    });
+                }
+            }
+        }
+        push_down_topk_to_merge(child, top_k.clone())?;
+    }
+    Ok(())
+}

src/query/sql/src/planner/optimizer/distributed/mod.rs

@@ -14,5 +14,6 @@
 
 #[allow(clippy::module_inception)]
 mod distributed;
+mod topk;
 
 pub use distributed::optimize_distributed_query;

src/query/sql/src/planner/optimizer/distributed/topk.rs

@@ -0,0 +1,22 @@
+// Copyright 2022 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::plans::Limit;
+use crate::plans::Sort;
+
+#[derive(Clone, Debug)]
+pub struct TopK {
+    pub(crate) sort: Sort,
+    pub(crate) limit: Limit,
+}

tests/logictest/suites/mode/cluster/04_0002_explain_v2

@@ -167,6 +167,70 @@ Exchange
 statement ok
 set prefer_broadcast_join = 0;
 
+statement query T
+explain select * from (SELECT number AS a FROM numbers(10)) x  order by x.a limit 3;
+
+----
+Limit
+├── limit: 3
+├── offset: 0
+└── Sort
+    ├── sort keys: [number ASC NULLS LAST]
+    └── Exchange
+        ├── exchange type: Merge
+        └── Limit
+            ├── limit: 3
+            ├── offset: 0
+            └── Sort
+                ├── sort keys: [number ASC NULLS LAST]
+                └── TableScan
+                    ├── table: default.system.numbers
+                    ├── read rows: 10
+                    ├── read bytes: 80
+                    ├── partitions total: 1
+                    ├── partitions scanned: 1
+                    └── push downs: [filters: [], limit: 3]
+
+statement query T
+explain select * from (SELECT number AS a FROM numbers(10)) x right join (SELECT number AS a FROM numbers(5)) y using(a) order by x.a limit 3;
+
+----
+Limit
+├── limit: 3
+├── offset: 0
+└── Sort
+    ├── sort keys: [number ASC NULLS LAST]
+    └── Exchange
+        ├── exchange type: Merge
+        └── Limit
+            ├── limit: 3
+            ├── offset: 0
+            └── Sort
+                ├── sort keys: [number ASC NULLS LAST]
+                └── HashJoin
+                    ├── join type: RIGHT OUTER
+                    ├── build keys: [y.a (#1)]
+                    ├── probe keys: [x.a (#0)]
+                    ├── filters: []
+                    ├── Exchange(Build)
+                    │   ├── exchange type: Hash(y.a (#1))
+                    │   └── TableScan
+                    │       ├── table: default.system.numbers
+                    │       ├── read rows: 5
+                    │       ├── read bytes: 40
+                    │       ├── partitions total: 1
+                    │       ├── partitions scanned: 1
+                    │       └── push downs: [filters: [], limit: NONE]
+                    └── Exchange(Probe)
+                        ├── exchange type: Hash(x.a (#0))
+                        └── TableScan
+                            ├── table: default.system.numbers
+                            ├── read rows: 10
+                            ├── read bytes: 80
+                            ├── partitions total: 1
+                            ├── partitions scanned: 1
+                            └── push downs: [filters: [], limit: NONE]
+
 statement ok
 drop table t1;