Use putIdentity method to reduce code duplication

trinodb · Jul 20, 2023 · b9f4848 · b9f4848
1 parent 7a1df27
commit b9f4848
Show file tree

Hide file tree

Showing 7 changed files with 16 additions and 16 deletions.
diff --git a/core/trino-main/src/main/java/io/trino/sql/planner/QueryPlanner.java b/core/trino-main/src/main/java/io/trino/sql/planner/QueryPlanner.java
@@ -547,7 +547,7 @@ public PlanNode plan(Delete node)
             Symbol symbol = relationPlan.getFieldMappings().get(fieldIndex);
             columnSymbolsBuilder.add(symbol);
             if (mergeAnalysis.getRedistributionColumnHandles().contains(columnHandle)) {
-                assignmentsBuilder.put(symbol, symbol.toSymbolReference());
+                assignmentsBuilder.putIdentity(symbol);
             }
             else {
                 assignmentsBuilder.put(symbol, new NullLiteral());
@@ -720,11 +720,11 @@ public PlanNode plan(Update node)
         for (ColumnHandle column : mergeAnalysis.getRedistributionColumnHandles()) {
             int fieldIndex = requireNonNull(mergeAnalysis.getColumnHandleFieldNumbers().get(column), "Could not find fieldIndex for redistribution column");
             Symbol symbol = relationPlan.getFieldMappings().get(fieldIndex);
-            projectionAssignmentsBuilder.put(symbol, symbol.toSymbolReference());
+            projectionAssignmentsBuilder.putIdentity(symbol);
         }
 
         // Add the rest of the page columns: rowId, merge row, case number and is_distinct
-        projectionAssignmentsBuilder.put(rowIdSymbol, rowIdSymbol.toSymbolReference());
+        projectionAssignmentsBuilder.putIdentity(rowIdSymbol);
         projectionAssignmentsBuilder.put(mergeRowSymbol, mergeRow);
         projectionAssignmentsBuilder.put(caseNumberSymbol, new GenericLiteral("INTEGER", "0"));
         projectionAssignmentsBuilder.put(isDistinctSymbol, TRUE_LITERAL);
@@ -858,10 +858,10 @@ public MergeWriterNode plan(Merge merge)
         for (ColumnHandle column : mergeAnalysis.getRedistributionColumnHandles()) {
             int fieldIndex = requireNonNull(mergeAnalysis.getColumnHandleFieldNumbers().get(column), "Could not find fieldIndex for redistribution column");
             Symbol symbol = planWithPresentColumn.getFieldMappings().get(fieldIndex);
-            projectionAssignmentsBuilder.put(symbol, symbol.toSymbolReference());
+            projectionAssignmentsBuilder.putIdentity(symbol);
         }
-        projectionAssignmentsBuilder.put(uniqueIdSymbol, uniqueIdSymbol.toSymbolReference());
-        projectionAssignmentsBuilder.put(rowIdSymbol, rowIdSymbol.toSymbolReference());
+        projectionAssignmentsBuilder.putIdentity(uniqueIdSymbol);
+        projectionAssignmentsBuilder.putIdentity(rowIdSymbol);
         projectionAssignmentsBuilder.put(mergeRowSymbol, caseExpression);
 
         ProjectNode subPlanProject = new ProjectNode(

diff --git a/core/trino-main/src/main/java/io/trino/sql/planner/RelationPlanner.java b/core/trino-main/src/main/java/io/trino/sql/planner/RelationPlanner.java
@@ -975,13 +975,13 @@ If casts are redundant (due to column type and common type being equal),
         for (int field : joinAnalysis.getOtherLeftFields()) {
             Symbol symbol = left.getFieldMappings().get(field);
             outputs.add(symbol);
-            assignments.put(symbol, symbol.toSymbolReference());
+            assignments.putIdentity(symbol);
         }
 
         for (int field : joinAnalysis.getOtherRightFields()) {
             Symbol symbol = right.getFieldMappings().get(field);
             outputs.add(symbol);
-            assignments.put(symbol, symbol.toSymbolReference());
+            assignments.putIdentity(symbol);
         }
 
         return new RelationPlan(

diff --git a/...in/src/main/java/io/trino/sql/planner/iterative/rule/PushAggregationThroughOuterJoin.java b/...in/src/main/java/io/trino/sql/planner/iterative/rule/PushAggregationThroughOuterJoin.java
@@ -264,7 +264,7 @@ private Optional<PlanNode> coalesceWithNullAggregation(AggregationNode aggregati
                 assignmentsBuilder.put(symbol, new CoalesceExpression(symbol.toSymbolReference(), sourceAggregationToOverNullMapping.get(symbol).toSymbolReference()));
             }
             else {
-                assignmentsBuilder.put(symbol, symbol.toSymbolReference());
+                assignmentsBuilder.putIdentity(symbol);
             }
         }
         return Optional.of(new ProjectNode(idAllocator.getNextId(), crossJoin, assignmentsBuilder.build()));

diff --git a/...main/src/main/java/io/trino/sql/planner/iterative/rule/PushProjectionThroughExchange.java b/...main/src/main/java/io/trino/sql/planner/iterative/rule/PushProjectionThroughExchange.java
@@ -97,15 +97,15 @@ public Result apply(ProjectNode project, Captures captures, Context context)
             partitioningColumns.stream()
                     .map(outputToInputMap::get)
                     .forEach(inputSymbol -> {
-                        projections.put(inputSymbol, inputSymbol.toSymbolReference());
+                        projections.putIdentity(inputSymbol);
                         inputs.add(inputSymbol);
                     });
 
             // Need to retain the hash symbol for the exchange
             exchange.getPartitioningScheme().getHashColumn()
                     .map(outputToInputMap::get)
                     .ifPresent(inputSymbol -> {
-                        projections.put(inputSymbol, inputSymbol.toSymbolReference());
+                        projections.putIdentity(inputSymbol);
                         inputs.add(inputSymbol);
                     });
 
@@ -116,7 +116,7 @@ public Result apply(ProjectNode project, Captures captures, Context context)
                         .filter(symbol -> !partitioningColumns.contains(symbol))
                         .map(outputToInputMap::get)
                         .forEach(inputSymbol -> {
-                            projections.put(inputSymbol, inputSymbol.toSymbolReference());
+                            projections.putIdentity(inputSymbol);
                             inputs.add(inputSymbol);
                         });
             }

diff --git a/.../trino-main/src/main/java/io/trino/sql/planner/optimizations/HashGenerationOptimizer.java b/.../trino-main/src/main/java/io/trino/sql/planner/optimizations/HashGenerationOptimizer.java
@@ -759,7 +759,7 @@ private PlanWithProperties enforce(PlanWithProperties planWithProperties, HashCo
             for (Symbol symbol : planWithProperties.getNode().getOutputSymbols()) {
                 HashComputation partitionSymbols = resultHashSymbols.get(symbol);
                 if (partitionSymbols == null || requiredHashes.getHashes().contains(partitionSymbols)) {
-                    assignments.put(symbol, symbol.toSymbolReference());
+                    assignments.putIdentity(symbol);
 
                     if (partitionSymbols != null) {
                         outputHashSymbols.put(partitionSymbols, symbol);

diff --git a/.../src/test/java/io/trino/sql/planner/iterative/rule/TestPushProjectionThroughExchange.java b/.../src/test/java/io/trino/sql/planner/iterative/rule/TestPushProjectionThroughExchange.java
@@ -61,8 +61,8 @@ public void testDoesNotFireNarrowingProjection()
 
                     return p.project(
                             Assignments.builder()
-                                    .put(a, a.toSymbolReference())
-                                    .put(b, b.toSymbolReference())
+                                    .putIdentity(a)
+                                    .putIdentity(b)
                                     .build(),
                             p.exchange(e -> e
                                     .addSource(p.values(a, b, c))

diff --git a/...no-main/src/test/java/io/trino/sql/planner/iterative/rule/TestPushTopNThroughProject.java b/...no-main/src/test/java/io/trino/sql/planner/iterative/rule/TestPushTopNThroughProject.java
@@ -173,7 +173,7 @@ public void testPushTopNThroughOverlappingDereferences()
                                     Assignments.builder()
                                             .put(p.symbol("b"), new SubscriptExpression(a.toSymbolReference(), new LongLiteral("1")))
                                             .put(p.symbol("c", rowType), a.toSymbolReference())
-                                            .put(d, d.toSymbolReference())
+                                            .putIdentity(d)
                                             .build(),
                                     p.values(a, d)));
                 })