Specify stream partitioning for AssignUniqueId in PropertyDerivations

trinodb · Oct 8, 2021 · 204a953 · 204a953
1 parent 33bb270
commit 204a953
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 14 deletions.
diff --git a/core/trino-main/src/main/java/io/trino/sql/planner/optimizations/PropertyDerivations.java b/core/trino-main/src/main/java/io/trino/sql/planner/optimizations/PropertyDerivations.java
@@ -230,7 +230,7 @@ public ActualProperties visitAssignUniqueId(AssignUniqueId node, List<ActualProp
             }
 
             return ActualProperties.builderFrom(properties)
-                    .global(partitionedOn(ARBITRARY_DISTRIBUTION, ImmutableList.of(node.getIdColumn()), Optional.empty()))
+                    .global(partitionedOn(ARBITRARY_DISTRIBUTION, ImmutableList.of(node.getIdColumn()), Optional.of(ImmutableList.of(node.getIdColumn()))))
                     .local(newLocalProperties.build())
                     .build();
         }

diff --git a/core/trino-main/src/test/java/io/trino/sql/planner/TestLogicalPlanner.java b/core/trino-main/src/test/java/io/trino/sql/planner/TestLogicalPlanner.java
@@ -777,19 +777,13 @@ public void testCorrelatedScalarSubqueryInSelect()
                 anyTree(
                         filter(format("CASE \"is_distinct\" WHEN true THEN true ELSE CAST(fail(%s, 'Scalar sub-query has returned multiple rows') AS boolean) END", SUBQUERY_MULTIPLE_ROWS.toErrorCode().getCode()),
                                 project(
-                                        markDistinct("is_distinct", ImmutableList.of("n_name", "n_regionkey", "unique"), "hash",
-                                                exchange(LOCAL, GATHER,
-                                                        exchange(REMOTE, REPARTITION,
-                                                                project(
-                                                                        ImmutableMap.of(
-                                                                                "hash",
-                                                                                expression("combine_hash(combine_hash(combine_hash(bigint '0', COALESCE(\"$operator$hash_code\"(\"n_name\"), 0)), COALESCE(\"$operator$hash_code\"(\"n_regionkey\"), 0)), COALESCE(\"$operator$hash_code\"(\"unique\"), 0))")),
-                                                                        join(LEFT, ImmutableList.of(equiJoinClause("n_regionkey", "r_regionkey")),
-                                                                                project(
-                                                                                        assignUniqueId("unique",
-                                                                                                tableScan("nation", ImmutableMap.of("n_regionkey", "regionkey", "n_name", "name")))),
-                                                                                anyTree(
-                                                                                        tableScan("region", ImmutableMap.of("r_regionkey", "regionkey"))))))))))));
+                                        markDistinct("is_distinct", ImmutableList.of("unique"),
+                                                join(LEFT, ImmutableList.of(equiJoinClause("n_regionkey", "r_regionkey")),
+                                                        project(
+                                                                assignUniqueId("unique",
+                                                                        tableScan("nation", ImmutableMap.of("n_regionkey", "regionkey", "n_name", "name")))),
+                                                        anyTree(
+                                                                tableScan("region", ImmutableMap.of("r_regionkey", "regionkey")))))))));
     }
 
     @Test