From 337732ae01c3b634d2b6284dc02ccdea0d5a74a5 Mon Sep 17 00:00:00 2001
From: starocean999 <40539150+starocean999@users.noreply.github.com>
Date: Wed, 10 May 2023 17:57:21 +0800
Subject: [PATCH] [fix](nereids) lost exchange before global limit merge node
 sometimes (#19396)

should add exchange node between global and local limit
---
 .../translator/PhysicalPlanTranslator.java    | 22 +++++++++++++------
 .../ChildOutputPropertyDeriver.java           |  4 ++++
 .../limit/sql/useRestrictiveLimitFromSubq.sql |  3 ++-
 .../limit/sql/withGroupByInSubq.sql           |  3 ++-
 .../nereids_p0/limit/sql/withJoinInSubq.sql   |  3 ++-
 .../suites/nereids_p0/limit/sql/withSubq.sql  |  3 ++-
 6 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java
index acb3ebe87b..5d78760c9a 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java
@@ -1602,13 +1602,21 @@ public class PhysicalPlanTranslator extends DefaultPlanVisitor<PlanFragment, Pla
         }
 
         PlanNode child = inputFragment.getPlanRoot();
-
-        // This case means GlobalLimit's child isn't gatherNode, which suggests the child is UNPARTITIONED
-        // When there is valid offset, exchangeNode should be added because other node don't support offset
-        if (physicalLimit.isGlobal() && physicalLimit.hasValidOffset()
-                && !(child instanceof ExchangeNode)) {
-            inputFragment = createParentFragment(inputFragment, DataPartition.UNPARTITIONED, context);
-            child = inputFragment.getPlanRoot();
+        if (physicalLimit.isGlobal()) {
+            if (child instanceof ExchangeNode) {
+                DataPartition outputPartition = DataPartition.UNPARTITIONED;
+                ExchangeNode exchangeNode = (ExchangeNode) inputFragment.getPlanRoot();
+                inputFragment.setOutputPartition(outputPartition);
+                inputFragment.setPlanRoot(exchangeNode.getChild(0));
+                inputFragment.setDestination(exchangeNode);
+                inputFragment = new PlanFragment(context.nextFragmentId(), exchangeNode, DataPartition.UNPARTITIONED);
+                context.addPlanFragment(inputFragment);
+            } else if (physicalLimit.hasValidOffset()) {
+                // This case means GlobalLimit's child isn't gatherNode, which suggests the child is UNPARTITIONED
+                // When there is valid offset, exchangeNode should be added because other node don't support offset
+                inputFragment = createParentFragment(inputFragment, DataPartition.UNPARTITIONED, context);
+                child = inputFragment.getPlanRoot();
+            }
         }
         child.setOffset(physicalLimit.getOffset());
         child.setLimit(physicalLimit.getLimit());
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/ChildOutputPropertyDeriver.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/ChildOutputPropertyDeriver.java
index 7cfffe090e..f4efbe4685 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/ChildOutputPropertyDeriver.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/ChildOutputPropertyDeriver.java
@@ -124,6 +124,10 @@ public class ChildOutputPropertyDeriver extends PlanVisitor<PhysicalProperties,
     public PhysicalProperties visitPhysicalLimit(PhysicalLimit<? extends Plan> limit, PlanContext context) {
         Preconditions.checkState(childrenOutputProperties.size() == 1);
         PhysicalProperties childOutputProperty = childrenOutputProperties.get(0);
+        if (limit.getPhase().isLocal()) {
+            return new PhysicalProperties(childOutputProperty.getDistributionSpec(),
+                    childOutputProperty.getOrderSpec());
+        }
         return new PhysicalProperties(DistributionSpecGather.INSTANCE, childOutputProperty.getOrderSpec());
     }
 
diff --git a/regression-test/suites/nereids_p0/limit/sql/useRestrictiveLimitFromSubq.sql b/regression-test/suites/nereids_p0/limit/sql/useRestrictiveLimitFromSubq.sql
index c7a1a6659a..b9db067d77 100644
--- a/regression-test/suites/nereids_p0/limit/sql/useRestrictiveLimitFromSubq.sql
+++ b/regression-test/suites/nereids_p0/limit/sql/useRestrictiveLimitFromSubq.sql
@@ -1,2 +1,3 @@
 -- database: presto; groups: limit; tables: nation
-SELECT COUNT(*) FROM (SELECT * FROM tpch_tiny_nation LIMIT 2) AS foo LIMIT 5
+SELECT /*+SET_VAR(parallel_fragment_exec_instance_num=2) */
+COUNT(*) FROM (SELECT * FROM tpch_tiny_nation LIMIT 2) AS foo LIMIT 5
diff --git a/regression-test/suites/nereids_p0/limit/sql/withGroupByInSubq.sql b/regression-test/suites/nereids_p0/limit/sql/withGroupByInSubq.sql
index dc01b0ca92..0b53dc0cab 100644
--- a/regression-test/suites/nereids_p0/limit/sql/withGroupByInSubq.sql
+++ b/regression-test/suites/nereids_p0/limit/sql/withGroupByInSubq.sql
@@ -1,4 +1,5 @@
 -- database: presto; groups: limit; tables: partsupp
-SELECT COUNT(*) FROM (
+SELECT /*+SET_VAR(parallel_fragment_exec_instance_num=2) */ 
+COUNT(*) FROM (
     SELECT suppkey, COUNT(*) FROM tpch_tiny_partsupp
     GROUP BY suppkey LIMIT 20) t1
diff --git a/regression-test/suites/nereids_p0/limit/sql/withJoinInSubq.sql b/regression-test/suites/nereids_p0/limit/sql/withJoinInSubq.sql
index 2abb6d48ef..444bccf623 100644
--- a/regression-test/suites/nereids_p0/limit/sql/withJoinInSubq.sql
+++ b/regression-test/suites/nereids_p0/limit/sql/withJoinInSubq.sql
@@ -1 +1,2 @@
-SELECT COUNT(*) FROM (SELECT n1.regionkey, n1.nationkey FROM tpch_tiny_nation n1 JOIN tpch_tiny_nation n2 ON n1.regionkey = n2.regionkey LIMIT 5) foo
+SELECT /*+SET_VAR(parallel_fragment_exec_instance_num=2) */
+COUNT(*) FROM (SELECT n1.regionkey, n1.nationkey FROM tpch_tiny_nation n1 JOIN tpch_tiny_nation n2 ON n1.regionkey = n2.regionkey LIMIT 5) foo
diff --git a/regression-test/suites/nereids_p0/limit/sql/withSubq.sql b/regression-test/suites/nereids_p0/limit/sql/withSubq.sql
index ebb14262df..3cc941a518 100644
--- a/regression-test/suites/nereids_p0/limit/sql/withSubq.sql
+++ b/regression-test/suites/nereids_p0/limit/sql/withSubq.sql
@@ -1,2 +1,3 @@
 -- database: presto; groups: limit; tables: nation
-SELECT COUNT(*) FROM (SELECT * FROM tpch_tiny_nation LIMIT 10) t1
+SELECT /*+SET_VAR(parallel_fragment_exec_instance_num=2) */
+COUNT(*) FROM (SELECT * FROM tpch_tiny_nation LIMIT 10) t1