This is an automated email from the ASF dual-hosted git repository.

dataroaring pushed a commit to branch branch-3.0
in repository https://gitbox.apache.org/repos/asf/doris.git


The following commit(s) were added to refs/heads/branch-3.0 by this push:
     new 7e4f29227be branch-3.0: [fix](Nereids) self join not always could do 
colocate join #54323 (#54352)
7e4f29227be is described below

commit 7e4f29227be4c0f25ea7291b48d120b4599dd320
Author: morrySnow <[email protected]>
AuthorDate: Tue Aug 12 10:22:27 2025 +0800

    branch-3.0: [fix](Nereids) self join not always could do colocate join 
#54323 (#54352)
    
    picked from #54323
---
 fe/fe-core/pom.xml                                 |  10 +
 .../org/apache/doris/nereids/util/JoinUtils.java   |   7 +-
 .../apache/doris/nereids/util/JoinUtilsTest.java   | 272 +++++++++++++++++++++
 fe/pom.xml                                         |  11 +
 4 files changed, 295 insertions(+), 5 deletions(-)

diff --git a/fe/fe-core/pom.xml b/fe/fe-core/pom.xml
index 2938d032765..54143c8a8f8 100644
--- a/fe/fe-core/pom.xml
+++ b/fe/fe-core/pom.xml
@@ -954,6 +954,16 @@ under the License.
             <artifactId>sdk-core</artifactId>
             <version>${awssdk.version}</version>
         </dependency>
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-core</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-inline</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
     <repositories>
         <!-- for hive-catalog-shade -->
diff --git 
a/fe/fe-core/src/main/java/org/apache/doris/nereids/util/JoinUtils.java 
b/fe/fe-core/src/main/java/org/apache/doris/nereids/util/JoinUtils.java
index 770905bf199..4de6696bca1 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/util/JoinUtils.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/util/JoinUtils.java
@@ -303,11 +303,8 @@ public class JoinUtils {
         boolean noNeedCheckColocateGroup = hitSameIndex && 
(leftTablePartitions.equals(rightTablePartitions))
                 && (leftTablePartitions.size() <= 1);
         ColocateTableIndex colocateIndex = Env.getCurrentColocateIndex();
-        if (noNeedCheckColocateGroup) {
-            return true;
-        }
-        if (!colocateIndex.isSameGroup(leftTableId, rightTableId)
-                || 
colocateIndex.isGroupUnstable(colocateIndex.getGroup(leftTableId))) {
+        if (!noNeedCheckColocateGroup && 
(!colocateIndex.isSameGroup(leftTableId, rightTableId)
+                || 
colocateIndex.isGroupUnstable(colocateIndex.getGroup(leftTableId)))) {
             return false;
         }
 
diff --git 
a/fe/fe-core/src/test/java/org/apache/doris/nereids/util/JoinUtilsTest.java 
b/fe/fe-core/src/test/java/org/apache/doris/nereids/util/JoinUtilsTest.java
new file mode 100644
index 00000000000..c1e9c77252a
--- /dev/null
+++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/util/JoinUtilsTest.java
@@ -0,0 +1,272 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you 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.
+
+package org.apache.doris.nereids.util;
+
+import org.apache.doris.catalog.ColocateTableIndex;
+import org.apache.doris.catalog.ColocateTableIndex.GroupId;
+import org.apache.doris.catalog.Env;
+import org.apache.doris.nereids.properties.DistributionSpecHash;
+import org.apache.doris.nereids.properties.DistributionSpecHash.ShuffleType;
+import org.apache.doris.nereids.trees.expressions.EqualTo;
+import org.apache.doris.nereids.trees.expressions.ExprId;
+import org.apache.doris.nereids.trees.expressions.Expression;
+import org.apache.doris.nereids.trees.expressions.SlotReference;
+import org.apache.doris.nereids.types.TinyIntType;
+import org.apache.doris.qe.ConnectContext;
+
+import com.google.common.collect.Lists;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
+
+import java.util.Collections;
+import java.util.List;
+
+public class JoinUtilsTest {
+
+    @Test
+    public void testCouldColocateJoinForSameTable() {
+        ConnectContext ctx = new ConnectContext();
+        ctx.setThreadLocalInfo();
+
+        DistributionSpecHash left = new 
DistributionSpecHash(Lists.newArrayList(new ExprId(1)), ShuffleType.NATURAL,
+                1L, 1L, Collections.emptySet());
+        DistributionSpecHash right = new 
DistributionSpecHash(Lists.newArrayList(new ExprId(2)), ShuffleType.NATURAL,
+                1L, 1L, Collections.emptySet());
+
+        Expression leftKey1 = new SlotReference(new ExprId(1), "c1",
+                TinyIntType.INSTANCE, false, Lists.newArrayList());
+        Expression rightKey1 = new SlotReference(new ExprId(2), "c1",
+                TinyIntType.INSTANCE, false, Lists.newArrayList());
+        Expression leftKey2 = new SlotReference(new ExprId(3), "c1",
+                TinyIntType.INSTANCE, false, Lists.newArrayList());
+        Expression rightKey2 = new SlotReference(new ExprId(4), "c1",
+                TinyIntType.INSTANCE, false, Lists.newArrayList());
+
+        List<Expression> conjuncts;
+
+        // key same with distribute key
+        conjuncts = Lists.newArrayList(new EqualTo(leftKey1, rightKey1));
+        Assertions.assertTrue(JoinUtils.couldColocateJoin(left, right, 
conjuncts));
+
+        // key contains distribute key, and have distribute key = distribute 
key
+        conjuncts = Lists.newArrayList(new EqualTo(leftKey1, rightKey1), new 
EqualTo(leftKey2, rightKey2));
+        Assertions.assertTrue(JoinUtils.couldColocateJoin(left, right, 
conjuncts));
+
+        // key contains distribute key, and NOT have distribute key = 
distribute key
+        conjuncts = Lists.newArrayList(new EqualTo(leftKey1, rightKey2), new 
EqualTo(leftKey2, rightKey1));
+        Assertions.assertFalse(JoinUtils.couldColocateJoin(left, right, 
conjuncts));
+
+        // key not contains distribute key
+        conjuncts = Lists.newArrayList(new EqualTo(leftKey2, rightKey2));
+        Assertions.assertFalse(JoinUtils.couldColocateJoin(left, right, 
conjuncts));
+    }
+
+    @Test
+    public void testCouldColocateJoinForDiffTableInSameGroupAndGroupIsStable() 
{
+        ConnectContext ctx = new ConnectContext();
+        ctx.setThreadLocalInfo();
+
+        // same group and group is statble
+        try (MockedStatic<Env> mockedEnv = Mockito.mockStatic(Env.class)) {
+            GroupId groupId = new GroupId(1L, 1L);
+            ColocateTableIndex colocateIndex = 
Mockito.mock(ColocateTableIndex.class);
+            Mockito.when(colocateIndex.isSameGroup(1L, 2L)).thenReturn(true);
+            Mockito.when(colocateIndex.getGroup(1L)).thenReturn(groupId);
+            
Mockito.when(colocateIndex.isGroupUnstable(groupId)).thenReturn(false);
+            mockedEnv.when(() -> 
Env.getCurrentColocateIndex()).thenReturn(colocateIndex);
+
+            DistributionSpecHash left = new 
DistributionSpecHash(Lists.newArrayList(new ExprId(1)),
+                    ShuffleType.NATURAL, 1L, 1L, Collections.emptySet());
+            DistributionSpecHash right = new 
DistributionSpecHash(Lists.newArrayList(new ExprId(2)),
+                    ShuffleType.NATURAL, 2L, 2L, Collections.emptySet());
+
+            Expression leftKey1 = new SlotReference(new ExprId(1), "c1",
+                    TinyIntType.INSTANCE, false, Lists.newArrayList());
+            Expression rightKey1 = new SlotReference(new ExprId(2), "c1",
+                    TinyIntType.INSTANCE, false, Lists.newArrayList());
+            Expression leftKey2 = new SlotReference(new ExprId(3), "c1",
+                    TinyIntType.INSTANCE, false, Lists.newArrayList());
+            Expression rightKey2 = new SlotReference(new ExprId(4), "c1",
+                    TinyIntType.INSTANCE, false, Lists.newArrayList());
+
+            List<Expression> conjuncts;
+
+            // key same with distribute key
+            conjuncts = Lists.newArrayList(new EqualTo(leftKey1, rightKey1));
+            Assertions.assertTrue(JoinUtils.couldColocateJoin(left, right, 
conjuncts));
+
+            // key contains distribute key, and have distribute key = 
distribute key
+            conjuncts = Lists.newArrayList(new EqualTo(leftKey1, rightKey1), 
new EqualTo(leftKey2, rightKey2));
+            Assertions.assertTrue(JoinUtils.couldColocateJoin(left, right, 
conjuncts));
+
+            // key contains distribute key, and NOT have distribute key = 
distribute key
+            conjuncts = Lists.newArrayList(new EqualTo(leftKey1, rightKey2), 
new EqualTo(leftKey2, rightKey1));
+            Assertions.assertFalse(JoinUtils.couldColocateJoin(left, right, 
conjuncts));
+
+            // key not contains distribute key
+            conjuncts = Lists.newArrayList(new EqualTo(leftKey2, rightKey2));
+            Assertions.assertFalse(JoinUtils.couldColocateJoin(left, right, 
conjuncts));
+        }
+    }
+
+    @Test
+    public void testCouldColocateJoinForNotNaturalHashDstribution() {
+        ConnectContext ctx = new ConnectContext();
+        ctx.setThreadLocalInfo();
+
+        // same group and group is statble
+        try (MockedStatic<Env> mockedEnv = Mockito.mockStatic(Env.class)) {
+            GroupId groupId = new GroupId(1L, 1L);
+            ColocateTableIndex colocateIndex = 
Mockito.mock(ColocateTableIndex.class);
+            Mockito.when(colocateIndex.isSameGroup(1L, 2L)).thenReturn(true);
+            Mockito.when(colocateIndex.getGroup(1L)).thenReturn(groupId);
+            
Mockito.when(colocateIndex.isGroupUnstable(groupId)).thenReturn(false);
+            mockedEnv.when(() -> 
Env.getCurrentColocateIndex()).thenReturn(colocateIndex);
+
+            DistributionSpecHash left = new 
DistributionSpecHash(Lists.newArrayList(new ExprId(1)),
+                    ShuffleType.NATURAL, 1L, 1L, Collections.emptySet());
+            DistributionSpecHash right = new 
DistributionSpecHash(Lists.newArrayList(new ExprId(2)),
+                    ShuffleType.EXECUTION_BUCKETED, 2L, 2L, 
Collections.emptySet());
+
+            Expression leftKey1 = new SlotReference(new ExprId(1), "c1",
+                    TinyIntType.INSTANCE, false, Lists.newArrayList());
+            Expression rightKey1 = new SlotReference(new ExprId(2), "c1",
+                    TinyIntType.INSTANCE, false, Lists.newArrayList());
+            Expression leftKey2 = new SlotReference(new ExprId(3), "c1",
+                    TinyIntType.INSTANCE, false, Lists.newArrayList());
+            Expression rightKey2 = new SlotReference(new ExprId(4), "c1",
+                    TinyIntType.INSTANCE, false, Lists.newArrayList());
+
+            List<Expression> conjuncts;
+
+            // key same with distribute key
+            conjuncts = Lists.newArrayList(new EqualTo(leftKey1, rightKey1));
+            Assertions.assertFalse(JoinUtils.couldColocateJoin(left, right, 
conjuncts));
+
+            // key contains distribute key, and have distribute key = 
distribute key
+            conjuncts = Lists.newArrayList(new EqualTo(leftKey1, rightKey1), 
new EqualTo(leftKey2, rightKey2));
+            Assertions.assertFalse(JoinUtils.couldColocateJoin(left, right, 
conjuncts));
+
+            // key contains distribute key, and NOT have distribute key = 
distribute key
+            conjuncts = Lists.newArrayList(new EqualTo(leftKey1, rightKey2), 
new EqualTo(leftKey2, rightKey1));
+            Assertions.assertFalse(JoinUtils.couldColocateJoin(left, right, 
conjuncts));
+
+            // key not contains distribute key
+            conjuncts = Lists.newArrayList(new EqualTo(leftKey2, rightKey2));
+            Assertions.assertFalse(JoinUtils.couldColocateJoin(left, right, 
conjuncts));
+        }
+    }
+
+    @Test
+    public void 
testCouldColocateJoinForDiffTableInSameGroupAndGroupIsUnstable() {
+        ConnectContext ctx = new ConnectContext();
+        ctx.setThreadLocalInfo();
+
+        // same group and group is statble
+        try (MockedStatic<Env> mockedEnv = Mockito.mockStatic(Env.class)) {
+            GroupId groupId = new GroupId(1L, 1L);
+            ColocateTableIndex colocateIndex = 
Mockito.mock(ColocateTableIndex.class);
+            Mockito.when(colocateIndex.isSameGroup(1L, 2L)).thenReturn(true);
+            Mockito.when(colocateIndex.getGroup(1L)).thenReturn(groupId);
+            
Mockito.when(colocateIndex.isGroupUnstable(groupId)).thenReturn(true);
+            mockedEnv.when(() -> 
Env.getCurrentColocateIndex()).thenReturn(colocateIndex);
+
+            DistributionSpecHash left = new 
DistributionSpecHash(Lists.newArrayList(new ExprId(1)),
+                    ShuffleType.NATURAL, 1L, 1L, Collections.emptySet());
+            DistributionSpecHash right = new 
DistributionSpecHash(Lists.newArrayList(new ExprId(2)),
+                    ShuffleType.NATURAL, 2L, 2L, Collections.emptySet());
+
+            Expression leftKey1 = new SlotReference(new ExprId(1), "c1",
+                    TinyIntType.INSTANCE, false, Lists.newArrayList());
+            Expression rightKey1 = new SlotReference(new ExprId(2), "c1",
+                    TinyIntType.INSTANCE, false, Lists.newArrayList());
+            Expression leftKey2 = new SlotReference(new ExprId(3), "c1",
+                    TinyIntType.INSTANCE, false, Lists.newArrayList());
+            Expression rightKey2 = new SlotReference(new ExprId(4), "c1",
+                    TinyIntType.INSTANCE, false, Lists.newArrayList());
+
+            List<Expression> conjuncts;
+
+            // key same with distribute key
+            conjuncts = Lists.newArrayList(new EqualTo(leftKey1, rightKey1));
+            Assertions.assertFalse(JoinUtils.couldColocateJoin(left, right, 
conjuncts));
+
+            // key contains distribute key, and have distribute key = 
distribute key
+            conjuncts = Lists.newArrayList(new EqualTo(leftKey1, rightKey1), 
new EqualTo(leftKey2, rightKey2));
+            Assertions.assertFalse(JoinUtils.couldColocateJoin(left, right, 
conjuncts));
+
+            // key contains distribute key, and NOT have distribute key = 
distribute key
+            conjuncts = Lists.newArrayList(new EqualTo(leftKey1, rightKey2), 
new EqualTo(leftKey2, rightKey1));
+            Assertions.assertFalse(JoinUtils.couldColocateJoin(left, right, 
conjuncts));
+
+            // key not contains distribute key
+            conjuncts = Lists.newArrayList(new EqualTo(leftKey2, rightKey2));
+            Assertions.assertFalse(JoinUtils.couldColocateJoin(left, right, 
conjuncts));
+        }
+    }
+
+    @Test
+    public void testCouldColocateJoinForDiffTableNotInSameGroup() {
+        ConnectContext ctx = new ConnectContext();
+        ctx.setThreadLocalInfo();
+
+        // same group and group is statble
+        try (MockedStatic<Env> mockedEnv = Mockito.mockStatic(Env.class)) {
+            GroupId groupId = new GroupId(1L, 1L);
+            ColocateTableIndex colocateIndex = 
Mockito.mock(ColocateTableIndex.class);
+            Mockito.when(colocateIndex.isSameGroup(1L, 2L)).thenReturn(true);
+            Mockito.when(colocateIndex.getGroup(1L)).thenReturn(groupId);
+            
Mockito.when(colocateIndex.isGroupUnstable(groupId)).thenReturn(true);
+            mockedEnv.when(() -> 
Env.getCurrentColocateIndex()).thenReturn(colocateIndex);
+
+            DistributionSpecHash left = new 
DistributionSpecHash(Lists.newArrayList(new ExprId(1)), ShuffleType.NATURAL,
+                    1L, 1L, Collections.emptySet());
+            DistributionSpecHash right = new 
DistributionSpecHash(Lists.newArrayList(new ExprId(2)), ShuffleType.NATURAL,
+                    2L, 2L, Collections.emptySet());
+
+            Expression leftKey1 = new SlotReference(new ExprId(1), "c1",
+                    TinyIntType.INSTANCE, false, Lists.newArrayList());
+            Expression rightKey1 = new SlotReference(new ExprId(2), "c1",
+                    TinyIntType.INSTANCE, false, Lists.newArrayList());
+            Expression leftKey2 = new SlotReference(new ExprId(3), "c1",
+                    TinyIntType.INSTANCE, false, Lists.newArrayList());
+            Expression rightKey2 = new SlotReference(new ExprId(4), "c1",
+                    TinyIntType.INSTANCE, false, Lists.newArrayList());
+
+            List<Expression> conjuncts;
+
+            // key same with distribute key
+            conjuncts = Lists.newArrayList(new EqualTo(leftKey1, rightKey1));
+            Assertions.assertFalse(JoinUtils.couldColocateJoin(left, right, 
conjuncts));
+
+            // key contains distribute key, and have distribute key = 
distribute key
+            conjuncts = Lists.newArrayList(new EqualTo(leftKey1, rightKey1), 
new EqualTo(leftKey2, rightKey2));
+            Assertions.assertFalse(JoinUtils.couldColocateJoin(left, right, 
conjuncts));
+
+            // key contains distribute key, and NOT have distribute key = 
distribute key
+            conjuncts = Lists.newArrayList(new EqualTo(leftKey1, rightKey2), 
new EqualTo(leftKey2, rightKey1));
+            Assertions.assertFalse(JoinUtils.couldColocateJoin(left, right, 
conjuncts));
+
+            // key not contains distribute key
+            conjuncts = Lists.newArrayList(new EqualTo(leftKey2, rightKey2));
+            Assertions.assertFalse(JoinUtils.couldColocateJoin(left, right, 
conjuncts));
+        }
+    }
+}
diff --git a/fe/pom.xml b/fe/pom.xml
index e88c6b383ff..b4ce8194482 100644
--- a/fe/pom.xml
+++ b/fe/pom.xml
@@ -382,6 +382,7 @@ under the License.
         <aliyun-sdk-oss.version>3.15.0</aliyun-sdk-oss.version>
         <awssdk.version>2.29.26</awssdk.version>
         <s3tables.catalog.version>0.1.4</s3tables.catalog.version>
+        <mockito.version>4.11.0</mockito.version>
     </properties>
     <profiles>
         <profile>
@@ -1702,6 +1703,16 @@ under the License.
                 <artifactId>semver4j</artifactId>
                 <version>${semver4j.version}</version>
             </dependency>
+            <dependency>
+                <groupId>org.mockito</groupId>
+                <artifactId>mockito-core</artifactId>
+                <version>${mockito.version}</version>
+            </dependency>
+            <dependency>
+                <groupId>org.mockito</groupId>
+                <artifactId>mockito-inline</artifactId>
+                <version>${mockito.version}</version>
+            </dependency>
         </dependencies>
     </dependencyManagement>
     <dependencies>


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to