From 204a95389aac30c61d3cf84825861e86fb264a91 Mon Sep 17 00:00:00 2001
From: praveenkrishna <praveenkrishna@tutanota.com>
Date: Wed, 16 Jun 2021 13:31:07 +0530
Subject: [PATCH] Specify stream partitioning for AssignUniqueId in
 PropertyDerivations

---
 .../optimizations/PropertyDerivations.java    |  2 +-
 .../trino/sql/planner/TestLogicalPlanner.java | 20 +++++++------------
 2 files changed, 8 insertions(+), 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
index 61b453fa0a10..1ee622230347 100644
--- 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
index 3d369255ef1e..fa6f27e4ab83 100644
--- 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