[fix](colocate) fix colocate join while multi tables (#37729)

apache · Jul 16, 2024 · 7cd326e · 7cd326e
1 parent 8e42871
commit 7cd326e
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 6 deletions.
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
@@ -310,12 +310,9 @@ public static boolean couldColocateJoin(DistributionSpecHash leftHashSpec, Distr
 
  SlotReference leftSlot = (SlotReference) leftChild;
  SlotReference rightSlot = (SlotReference) rightChild;
- Integer leftIndex = null;
- Integer rightIndex = null;
- if (leftSlot.getTable().isPresent() && leftSlot.getTable().get().getId() == leftHashSpec.getTableId()) {
- leftIndex = leftHashSpec.getExprIdToEquivalenceSet().get(leftSlot.getExprId());
- rightIndex = rightHashSpec.getExprIdToEquivalenceSet().get(rightSlot.getExprId());
- } else {
+ Integer leftIndex = leftHashSpec.getExprIdToEquivalenceSet().get(leftSlot.getExprId());
+ Integer rightIndex = rightHashSpec.getExprIdToEquivalenceSet().get(rightSlot.getExprId());
+ if (leftIndex == null) {
  leftIndex = rightHashSpec.getExprIdToEquivalenceSet().get(leftSlot.getExprId());
  rightIndex = leftHashSpec.getExprIdToEquivalenceSet().get(rightSlot.getExprId());
  }

diff --git a/regression-test/suites/correctness_p0/test_colocate_join_of_column_order.groovy b/regression-test/suites/correctness_p0/test_colocate_join_of_column_order.groovy
@@ -79,4 +79,33 @@ suite("test_colocate_join_of_column_order") {
 
  sql """ DROP TABLE IF EXISTS `test_colocate_join_of_column_order_t1`; """
  sql """ DROP TABLE IF EXISTS `test_colocate_join_of_column_order_t2`; """
+
+ // multi tables
+ sql """ DROP TABLE IF EXISTS `test_colocate_join_of_column_order_ta`; """
+ sql """
+ CREATE TABLE IF NOT EXISTS `test_colocate_join_of_column_order_ta` ( `c1` bigint NULL, `c2` bigint NULL)
+ DISTRIBUTED BY HASH(c1) PROPERTIES ( "replication_num" = "1", "colocate_with" = "group_column_order3");
+ """
+ sql """ DROP TABLE IF EXISTS `test_colocate_join_of_column_order_tb`; """
+ sql """
+ CREATE TABLE IF NOT EXISTS `test_colocate_join_of_column_order_tb` ( `c1` bigint NULL, `c2` bigint NULL)
+ DISTRIBUTED BY HASH(c1) PROPERTIES ( "replication_num" = "1", "colocate_with" = "group_column_order3");
+ """
+ sql """ DROP TABLE IF EXISTS `test_colocate_join_of_column_order_tc`; """
+ sql """
+ CREATE TABLE IF NOT EXISTS `test_colocate_join_of_column_order_tc` ( `c1` bigint NULL, `c2` bigint NULL)
+ DISTRIBUTED BY HASH(c1) PROPERTIES ( "replication_num" = "1", "colocate_with" = "group_column_order3");
+ """
+ sql """insert into test_colocate_join_of_column_order_ta values(1,1);"""
+ sql """insert into test_colocate_join_of_column_order_tb values(1,1);"""
+ sql """insert into test_colocate_join_of_column_order_tc values(1,1);"""
+
+ explain {
+ sql("""select /*+ set_var(disable_join_reorder=true) */ * from test_colocate_join_of_column_order_ta join [shuffle] (select cast((c2 + 1) as bigint) c2 from test_colocate_join_of_column_order_tb) test_colocate_join_of_column_order_tb on test_colocate_join_of_column_order_ta.c1 = test_colocate_join_of_column_order_tb.c2 join [shuffle] test_colocate_join_of_column_order_tc on test_colocate_join_of_column_order_tb.c2 = test_colocate_join_of_column_order_tc.c1;""");
+ contains "COLOCATE"
+ }
+
+ sql """ DROP TABLE IF EXISTS `test_colocate_join_of_column_order_ta`; """
+ sql """ DROP TABLE IF EXISTS `test_colocate_join_of_column_order_tb`; """
+ sql """ DROP TABLE IF EXISTS `test_colocate_join_of_column_order_tc`; """
 }