GroupByKeyOnlyEvaluatorFactory implements TransformEvaluatorFactory { @Override public <InputT> TransformEvaluator<InputT> forApplication( AppliedPTransform<?, ?, ?> application, CommittedBundle<?> inputBundle) { @SuppressWarnings({"cast", "unchecked", "rawtypes"}) TransformEvaluator<InputT> evaluator = createEvaluator((AppliedPTransform) application); return evaluator; } GroupByKeyOnlyEvaluatorFactory(EvaluationContext evaluationContext); @Override TransformEvaluator<InputT> forApplication(
      AppliedPTransform<?, ?, ?> application, CommittedBundle<?> inputBundle); @Override void cleanup();  }
@Test public void testInMemoryEvaluator() throws Exception { KV<String, Integer> firstFoo = KV.of("foo", -1); KV<String, Integer> secondFoo = KV.of("foo", 1); KV<String, Integer> thirdFoo = KV.of("foo", 3); KV<String, Integer> firstBar = KV.of("bar", 22); KV<String, Integer> secondBar = KV.of("bar", 12); KV<String, Integer> firstBaz = KV.of("baz", Integer.MAX_VALUE); PCollection<KV<String, Integer>> values = p.apply(Create.of(firstFoo, firstBar, secondFoo, firstBaz, secondBar, thirdFoo)); PCollection<KeyedWorkItem<String, Integer>> groupedKvs = values.apply(new DirectGroupByKeyOnly<>()); CommittedBundle<KV<String, Integer>> inputBundle = bundleFactory.createBundle(values).commit(Instant.now()); EvaluationContext evaluationContext = mock(EvaluationContext.class); StructuralKey<String> fooKey = StructuralKey.of("foo", StringUtf8Coder.of()); UncommittedBundle<KeyedWorkItem<String, Integer>> fooBundle = bundleFactory.createKeyedBundle(fooKey, groupedKvs); StructuralKey<String> barKey = StructuralKey.of("bar", StringUtf8Coder.of()); UncommittedBundle<KeyedWorkItem<String, Integer>> barBundle = bundleFactory.createKeyedBundle(barKey, groupedKvs); StructuralKey<String> bazKey = StructuralKey.of("baz", StringUtf8Coder.of()); UncommittedBundle<KeyedWorkItem<String, Integer>> bazBundle = bundleFactory.createKeyedBundle(bazKey, groupedKvs); when(evaluationContext.createKeyedBundle(fooKey, groupedKvs)).thenReturn(fooBundle); when(evaluationContext.createKeyedBundle(barKey, groupedKvs)).thenReturn(barBundle); when(evaluationContext.createKeyedBundle(bazKey, groupedKvs)).thenReturn(bazBundle); @SuppressWarnings("unchecked") Coder<String> keyCoder = ((KvCoder<String, Integer>) values.getCoder()).getKeyCoder(); TransformEvaluator<KV<String, Integer>> evaluator = new GroupByKeyOnlyEvaluatorFactory(evaluationContext) .forApplication(DirectGraphs.getProducer(groupedKvs), inputBundle); evaluator.processElement(WindowedValue.valueInGlobalWindow(firstFoo)); evaluator.processElement(WindowedValue.valueInGlobalWindow(secondFoo)); evaluator.processElement(WindowedValue.valueInGlobalWindow(thirdFoo)); evaluator.processElement(WindowedValue.valueInGlobalWindow(firstBar)); evaluator.processElement(WindowedValue.valueInGlobalWindow(secondBar)); evaluator.processElement(WindowedValue.valueInGlobalWindow(firstBaz)); evaluator.finishBundle(); assertThat( fooBundle.commit(Instant.now()).getElements(), contains( new KeyedWorkItemMatcher<>( KeyedWorkItems.elementsWorkItem( "foo", ImmutableSet.of( WindowedValue.valueInGlobalWindow(-1), WindowedValue.valueInGlobalWindow(1), WindowedValue.valueInGlobalWindow(3))), keyCoder))); assertThat( barBundle.commit(Instant.now()).getElements(), contains( new KeyedWorkItemMatcher<>( KeyedWorkItems.elementsWorkItem( "bar", ImmutableSet.of( WindowedValue.valueInGlobalWindow(12), WindowedValue.valueInGlobalWindow(22))), keyCoder))); assertThat( bazBundle.commit(Instant.now()).getElements(), contains( new KeyedWorkItemMatcher<>( KeyedWorkItems.elementsWorkItem( "baz", ImmutableSet.of(WindowedValue.valueInGlobalWindow(Integer.MAX_VALUE))), keyCoder))); }