ReplacePgbkWithPrecombineFunction implements Function<MutableNetwork<Node, Edge>, MutableNetwork<Node, Edge>> { @Override public MutableNetwork<Node, Edge> apply(MutableNetwork<Node, Edge> network) { Networks.replaceDirectedNetworkNodes( network, (Node node) -> { if (!isPrecombinePgbk(node)) { return node; } ParallelInstructionNode castNode = ((ParallelInstructionNode) node); ParallelInstruction parallelInstruction = castNode.getParallelInstruction(); Map<String, Object> cloudUserFnSpec = parallelInstruction.getPartialGroupByKey().getValueCombiningFn(); addString(cloudUserFnSpec, WorkerPropertyNames.PHASE, CombinePhase.ADD); ParDoInstruction newParDoInstruction = new ParDoInstruction(); newParDoInstruction.setUserFn(cloudUserFnSpec); ParallelInstruction newParallelInstruction = parallelInstruction.clone(); newParallelInstruction.setPartialGroupByKey(null); newParallelInstruction.setParDo(newParDoInstruction); return ParallelInstructionNode.create(newParallelInstruction, ExecutionLocation.UNKNOWN); }); return network; }  @Override MutableNetwork<Node, Edge> apply(MutableNetwork<Node, Edge> network);  }
@Test public void testPrecombinePgbkIsReplaced() throws Exception { Map<String, Object> valueCombiningFn = new HashMap<>(); Node out1 = createInstructionOutputNode("out1"); String pgbkName = "precombine_pgbk"; Node precombinePgbk = createPrecombinePgbkNode(pgbkName, valueCombiningFn); Node out2 = createInstructionOutputNode("out2"); MutableNetwork<Node, Edge> network = createEmptyNetwork(); network.addNode(out1); network.addNode(precombinePgbk); network.addNode(out2); network.addEdge(out1, precombinePgbk, DefaultEdge.create()); network.addEdge(precombinePgbk, out2, DefaultEdge.create()); Network<Node, Edge> inputNetwork = ImmutableNetwork.copyOf(network); network = new ReplacePgbkWithPrecombineFunction().apply(network); assertEquals(inputNetwork.nodes().size(), network.nodes().size()); assertEquals(inputNetwork.edges().size(), network.edges().size()); List<List<Node>> oldPaths = Networks.allPathsFromRootsToLeaves(inputNetwork); List<List<Node>> newPaths = Networks.allPathsFromRootsToLeaves(network); assertEquals(oldPaths.size(), newPaths.size()); for (Node node : network.nodes()) { if (node instanceof ParallelInstructionNode) { ParallelInstructionNode createdCombineNode = (ParallelInstructionNode) node; ParallelInstruction parallelInstruction = createdCombineNode.getParallelInstruction(); assertEquals(parallelInstruction.getName(), pgbkName); assertNull(parallelInstruction.getPartialGroupByKey()); assertNotNull(parallelInstruction.getParDo()); ParDoInstruction parDoInstruction = parallelInstruction.getParDo(); assertEquals(parDoInstruction.getUserFn(), valueCombiningFn); break; } } }
@Test public void testNormalPgbkIsNotReplaced() throws Exception { Node out1 = createInstructionOutputNode("out1"); Node pgbk = createPrecombinePgbkNode("pgbk", null); Node out2 = createInstructionOutputNode("out2"); MutableNetwork<Node, Edge> network = createEmptyNetwork(); network.addNode(out1); network.addNode(pgbk); network.addNode(out2); network.addEdge(out1, pgbk, DefaultEdge.create()); network.addEdge(pgbk, out2, DefaultEdge.create()); Network<Node, Edge> inputNetwork = ImmutableNetwork.copyOf(network); network = new ReplacePgbkWithPrecombineFunction().apply(network); List<List<Node>> oldPaths = Networks.allPathsFromRootsToLeaves(inputNetwork); List<List<Node>> newPaths = Networks.allPathsFromRootsToLeaves(network); assertThat(oldPaths, containsInAnyOrder(newPaths.toArray())); for (Node node : network.nodes()) { if (node instanceof ParallelInstructionNode) { ParallelInstructionNode newNode = (ParallelInstructionNode) node; ParallelInstruction parallelInstruction = newNode.getParallelInstruction(); assertEquals( parallelInstruction, ((ParallelInstructionNode) pgbk).getParallelInstruction()); break; } } }