Source:
private PartitionedOutputOperator createPartitionedOutputOperator(boolean shouldReplicate) { PartitionFunction partitionFunction = new LocalPartitionGenerator(new InterpretedHashGenerator(ImmutableList.of(BIGINT), new int[] {0}), PARTITION_COUNT); PagesSerdeFactory serdeFactory = new PagesSerdeFactory(new BlockEncodingManager(new TypeRegistry()), false); OutputBuffers buffers = createInitialEmptyOutputBuffers(PARTITIONED); for (int partition = 0; partition < PARTITION_COUNT; partition++) { buffers = buffers.withBuffer(new OutputBuffers.OutputBufferId(partition), partition); } PartitionedOutputBuffer buffer = createPartitionedBuffer( buffers.withNoMoreBufferIds(), new DataSize(Long.MAX_VALUE, BYTE)); buffer.registerLifespanCompletionCallback(ignore -> {}); PartitionedOutputOperator.PartitionedOutputFactory operatorFactory = null; if (shouldReplicate) { operatorFactory = new PartitionedOutputOperator.PartitionedOutputFactory( partitionFunction, ImmutableList.of(0), ImmutableList.of(Optional.empty()), true, OptionalInt.of(0), buffer, new DataSize(1, GIGABYTE)); } else { operatorFactory = new PartitionedOutputOperator.PartitionedOutputFactory( partitionFunction, ImmutableList.of(0), ImmutableList.of(Optional.empty(), Optional.empty()), false, OptionalInt.empty(), buffer, new DataSize(1, GIGABYTE)); } return (PartitionedOutputOperator) operatorFactory .createOutputOperator(0, new PlanNodeId("plan-node-0"), TYPES, Function.identity(), serdeFactory) .createOperator(createDriverContext()); }

Target:
private static PartitionedOutputOperator createPartitionedOutputOperator(boolean shouldReplicate) { PartitionFunction partitionFunction = new LocalPartitionGenerator(new InterpretedHashGenerator(ImmutableList.of(BIGINT), new int[] {0}), PARTITION_COUNT); PagesSerdeFactory serdeFactory = new PagesSerdeFactory(new BlockEncodingManager(new TypeRegistry()), false); DriverContext driverContext = TestingTaskContext.builder(EXECUTOR, SCHEDULER, TEST_SESSION) .setMemoryPoolSize(MAX_MEMORY) .setQueryMaxTotalMemory(MAX_MEMORY) .build() .addPipelineContext(0, true, true, false) .addDriverContext(); OutputBuffers buffers = createInitialEmptyOutputBuffers(PARTITIONED); for (int partition = 0; partition < PARTITION_COUNT; partition++) { buffers = buffers.withBuffer(new OutputBuffers.OutputBufferId(partition), partition); } PartitionedOutputBuffer buffer = new PartitionedOutputBuffer( "task-instance-id", new StateMachine<>("bufferState", SCHEDULER, OPEN, TERMINAL_BUFFER_STATES), buffers.withNoMoreBufferIds(), new DataSize(Long.MAX_VALUE, BYTE), () -> new SimpleLocalMemoryContext(newSimpleAggregatedMemoryContext(), "test"), SCHEDULER); buffer.registerLifespanCompletionCallback(ignore -> {}); PartitionedOutputOperator.PartitionedOutputFactory operatorFactory; if (shouldReplicate) { operatorFactory = new PartitionedOutputOperator.PartitionedOutputFactory( partitionFunction, ImmutableList.of(0), ImmutableList.of(Optional.empty()), true, OptionalInt.of(0), buffer, PARTITION_MAX_MEMORY); return (PartitionedOutputOperator) operatorFactory .createOutputOperator(0, new PlanNodeId("plan-node-0"), REPLICATION_TYPES, Function.identity(), serdeFactory) .createOperator(driverContext); } else { operatorFactory = new PartitionedOutputOperator.PartitionedOutputFactory( partitionFunction, ImmutableList.of(0), ImmutableList.of(Optional.empty(), Optional.empty()), false, OptionalInt.empty(), buffer, PARTITION_MAX_MEMORY); return (PartitionedOutputOperator) operatorFactory .createOutputOperator(0, new PlanNodeId("plan-node-0"), TYPES, Function.identity(), serdeFactory) .createOperator(driverContext); } }