Improve performance of Merge for MarkDistinct

Author: hqbhoho

core/trino-main/src/main/java/io/trino/sql/planner/QueryPlanner.java

@@ -751,9 +751,9 @@ public MergeWriterNode plan(Merge merge)
                 .process(merge.getTarget());
 
         // Assign a unique id to every target table row
-        Symbol uniqueIdSymbol = symbolAllocator.newSymbol("unique_id", BIGINT);
+        Symbol targetUniqueIdSymbol = symbolAllocator.newSymbol("target_unique_id", BIGINT);
         RelationPlan planWithUniqueId = new RelationPlan(
-                new AssignUniqueId(idAllocator.getNextId(), targetTablePlan.getRoot(), uniqueIdSymbol),
+                new AssignUniqueId(idAllocator.getNextId(), targetTablePlan.getRoot(), targetUniqueIdSymbol),
                 mergeAnalysis.getTargetTableScope(),
                 targetTablePlan.getFieldMappings(),
                 outerContext);
@@ -774,8 +774,16 @@ public MergeWriterNode plan(Merge merge)
         RelationPlan source = new RelationPlanner(analysis, symbolAllocator, idAllocator, lambdaDeclarationToSymbolMap, plannerContext, outerContext, session, recursiveSubqueries)
                 .process(merge.getSource());
 
+        // Assign a unique id to every source table row
+        Symbol sourceUniqueIdSymbol = symbolAllocator.newSymbol("source_unique_id", BIGINT);
+        RelationPlan sourcePlanWithUniqueId = new RelationPlan(
+                new AssignUniqueId(idAllocator.getNextId(), source.getRoot(), sourceUniqueIdSymbol),
+                source.getScope(),
+                source.getFieldMappings(),
+                outerContext);
+
         RelationPlan joinPlan = new RelationPlanner(analysis, symbolAllocator, idAllocator, lambdaDeclarationToSymbolMap, plannerContext, outerContext, session, recursiveSubqueries)
-                .planJoin(merge.getPredicate(), Join.Type.RIGHT, mergeAnalysis.getJoinScope(), planWithPresentColumn, source, analysis.getSubqueries(merge)); // TODO: ir
+                .planJoin(merge.getPredicate(), Join.Type.RIGHT, mergeAnalysis.getJoinScope(), planWithPresentColumn, sourcePlanWithUniqueId, analysis.getSubqueries(merge)); // TODO: ir
 
         PlanBuilder subPlan = newPlanBuilder(joinPlan, analysis, lambdaDeclarationToSymbolMap, session, plannerContext);
 
@@ -864,10 +872,10 @@ public MergeWriterNode plan(Merge merge)
 
             List<io.trino.sql.tree.Expression> constraints = analysis.getCheckConstraints(mergeAnalysis.getTargetTable());
             if (!constraints.isEmpty()) {
-                assignments.putIdentity(uniqueIdSymbol);
+                assignments.putIdentity(targetUniqueIdSymbol);
                 assignments.putIdentity(presentColumn);
                 assignments.putIdentity(rowIdSymbol);
-                assignments.putIdentities(source.getFieldMappings());
+                assignments.putIdentities(sourcePlanWithUniqueId.getFieldMappings());
                 subPlan = subPlan.withNewRoot(new ProjectNode(
                         idAllocator.getNextId(),
                         subPlan.getRoot(),
@@ -887,6 +895,7 @@ public MergeWriterNode plan(Merge merge)
                         .build()),
                         null));
 
+        Symbol uniqueIdSymbol = symbolAllocator.newSymbol("unique_id", BIGINT);
         Symbol mergeRowSymbol = symbolAllocator.newSymbol("merge_row", mergeAnalysis.getMergeRowType());
         Symbol caseNumberSymbol = symbolAllocator.newSymbol("case_number", INTEGER);
 
@@ -897,7 +906,9 @@ public MergeWriterNode plan(Merge merge)
             Symbol symbol = planWithPresentColumn.getFieldMappings().get(fieldIndex);
             projectionAssignmentsBuilder.putIdentity(symbol);
         }
-        projectionAssignmentsBuilder.putIdentity(uniqueIdSymbol);
+
+        Expression uniqueIdExpression = new Coalesce(targetUniqueIdSymbol.toSymbolReference(), sourceUniqueIdSymbol.toSymbolReference());
+        projectionAssignmentsBuilder.put(uniqueIdSymbol, uniqueIdExpression);
         projectionAssignmentsBuilder.putIdentity(rowIdSymbol);
         projectionAssignmentsBuilder.put(mergeRowSymbol, caseExpression);