Source:
private SubPlan reassignPartitioningHandleIfNecessaryHelper(Session session, Metadata metadata, SubPlan subPlan, PartitioningHandle newOutputPartitioningHandle) { PlanFragment fragment = subPlan.getFragment(); PartitioningHandleReassigner partitioningHandleReassigner = new PartitioningHandleReassigner(fragment.getPartitioning(), metadata, session); PlanNode newRoot; if (fragment.getPartitioning().isSingleNode()) { newRoot = fragment.getRoot(); } else { newRoot = SimplePlanRewriter.rewriteWith(partitioningHandleReassigner, fragment.getRoot()); } PartitioningScheme outputPartitioningScheme = fragment.getPartitioningScheme(); Partitioning newOutputPartitioning = outputPartitioningScheme.getPartitioning(); if (outputPartitioningScheme.getPartitioning().getHandle().getConnectorId().isPresent()) { newOutputPartitioning = newOutputPartitioning.withAlternativePartitiongingHandle(newOutputPartitioningHandle); } PlanFragment newFragment = new PlanFragment( fragment.getId(), newRoot, fragment.getSymbols(), fragment.getPartitioning(), fragment.getPartitionedSources(), new PartitioningScheme( newOutputPartitioning, outputPartitioningScheme.getOutputLayout(), outputPartitioningScheme.getHashColumn(), outputPartitioningScheme.isReplicateNullsAndAny(), outputPartitioningScheme.getBucketToPartition()), fragment.getStageExecutionStrategy(), fragment.getStatsAndCosts()); ImmutableList.Builder<SubPlan> childrenBuilder = ImmutableList.builder(); for (SubPlan child : subPlan.getChildren()) { childrenBuilder.add(reassignPartitioningHandleIfNecessaryHelper(session, metadata, child, fragment.getPartitioning())); } return new SubPlan(newFragment, childrenBuilder.build()); }

Target:
private SubPlan reassignPartitioningHandleIfNecessaryHelper(Session session, Metadata metadata, SubPlan subPlan, PartitioningHandle newOutputPartitioningHandle) { PlanFragment fragment = subPlan.getFragment(); PlanNode newRoot = fragment.getRoot(); if (!fragment.getPartitioning().isSingleNode()) { PartitioningHandleReassigner partitioningHandleReassigner = new PartitioningHandleReassigner(fragment.getPartitioning(), metadata, session); newRoot = SimplePlanRewriter.rewriteWith(partitioningHandleReassigner, newRoot); } PartitioningScheme outputPartitioningScheme = fragment.getPartitioningScheme(); Partitioning newOutputPartitioning = outputPartitioningScheme.getPartitioning(); if (outputPartitioningScheme.getPartitioning().getHandle().getConnectorId().isPresent()) { newOutputPartitioning = newOutputPartitioning.withAlternativePartitiongingHandle(newOutputPartitioningHandle); } PlanFragment newFragment = new PlanFragment( fragment.getId(), newRoot, fragment.getSymbols(), fragment.getPartitioning(), fragment.getPartitionedSources(), new PartitioningScheme( newOutputPartitioning, outputPartitioningScheme.getOutputLayout(), outputPartitioningScheme.getHashColumn(), outputPartitioningScheme.isReplicateNullsAndAny(), outputPartitioningScheme.getBucketToPartition()), fragment.getStageExecutionStrategy(), fragment.getStatsAndCosts()); ImmutableList.Builder<SubPlan> childrenBuilder = ImmutableList.builder(); for (SubPlan child : subPlan.getChildren()) { childrenBuilder.add(reassignPartitioningHandleIfNecessaryHelper(session, metadata, child, fragment.getPartitioning())); } return new SubPlan(newFragment, childrenBuilder.build()); }