From 523009697e3b619dd96d3ebf9bdf297dd9620f03 Mon Sep 17 00:00:00 2001
From: HappenLee <happenlee@hotmail.com>
Date: Fri, 31 Jul 2020 19:33:41 +0800
Subject: [PATCH] [ColocateJoin] ColocateJoin support table join itself 
 (#4230)

if left table and right table is same table, they are naturally colocate relationship.
---
 .../doris/planner/DistributedPlanner.java     | 28 +++++++++++--------
 .../apache/doris/planner/QueryPlanTest.java   |  7 +++++
 2 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/fe/fe-core/src/main/java/org/apache/doris/planner/DistributedPlanner.java b/fe/fe-core/src/main/java/org/apache/doris/planner/DistributedPlanner.java
index 351a7cff559af1..429076fea68156 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/planner/DistributedPlanner.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/planner/DistributedPlanner.java
@@ -494,19 +494,23 @@ private boolean canColocateJoin(HashJoinNode node, PlanNode leftRoot, PlanNode r
         OlapTable leftTable = ((OlapScanNode) leftRoot).getOlapTable();
         OlapTable rightTable = ((OlapScanNode) rightRoot).getOlapTable();
 
-        ColocateTableIndex colocateIndex = Catalog.getCurrentColocateIndex();
-
-        //1 the table must be colocate
-        if (!colocateIndex.isSameGroup(leftTable.getId(), rightTable.getId())) {
-            cannotReason.add("table not in same group");
-            return false;
-        }
+        // if left table and right table is same table, they are naturally colocate relationship
+        // no need to check colocate group
+        if (leftTable.getId() != rightTable.getId()) {
+            ColocateTableIndex colocateIndex = Catalog.getCurrentColocateIndex();
+
+            //1 the table must be colocate
+            if (!colocateIndex.isSameGroup(leftTable.getId(), rightTable.getId())) {
+                cannotReason.add("table not in the same group");
+                return false;
+            }
 
-        //2 the colocate group must be stable
-        GroupId groupId = colocateIndex.getGroup(leftTable.getId());
-        if (colocateIndex.isGroupUnstable(groupId)) {
-            cannotReason.add("group is not stable");
-            return false;
+            //2 the colocate group must be stable
+            GroupId groupId = colocateIndex.getGroup(leftTable.getId());
+            if (colocateIndex.isGroupUnstable(groupId)) {
+                cannotReason.add("group is not stable");
+                return false;
+            }
         }
 
         DistributionInfo leftDistribution = leftTable.getDefaultDistributionInfo();
diff --git a/fe/fe-core/src/test/java/org/apache/doris/planner/QueryPlanTest.java b/fe/fe-core/src/test/java/org/apache/doris/planner/QueryPlanTest.java
index 0134de9cd0e214..57ebc3c619a4de 100644
--- a/fe/fe-core/src/test/java/org/apache/doris/planner/QueryPlanTest.java
+++ b/fe/fe-core/src/test/java/org/apache/doris/planner/QueryPlanTest.java
@@ -895,6 +895,13 @@ public void testColocateJoin() throws Exception {
         Assert.assertTrue(explainString.contains("colocate: false"));
     }
 
+    @Test
+    public void testSelfColocateJoin() throws Exception {
+        String queryStr = "explain select * from test.jointest t1, test.jointest t2 where t1.k1 = t2.k1";
+        String explainString = UtFrameUtils.getSQLPlanOrErrorMsg(connectContext, queryStr);
+        Assert.assertTrue(explainString.contains("colocate: true"));
+    }
+
     @Test
     public void testJoinWithMysqlTable() throws Exception {
         connectContext.setDatabase("default_cluster:test");