FlattenEvaluatorFactory implements TransformEvaluatorFactory { @Override public <InputT> TransformEvaluator<InputT> forApplication( AppliedPTransform<?, ?, ?> application, CommittedBundle<?> inputBundle) { @SuppressWarnings({"cast", "unchecked", "rawtypes"}) TransformEvaluator<InputT> evaluator = (TransformEvaluator<InputT>) createInMemoryEvaluator((AppliedPTransform) application); return evaluator; } FlattenEvaluatorFactory(EvaluationContext evaluationContext); @Override TransformEvaluator<InputT> forApplication(
      AppliedPTransform<?, ?, ?> application, CommittedBundle<?> inputBundle); @Override void cleanup();  }
@Test public void testFlattenInMemoryEvaluator() throws Exception { PCollection<Integer> left = p.apply("left", Create.of(1, 2, 4)); PCollection<Integer> right = p.apply("right", Create.of(-1, 2, -4)); PCollectionList<Integer> list = PCollectionList.of(left).and(right); PCollection<Integer> flattened = list.apply(Flatten.pCollections()); CommittedBundle<Integer> leftBundle = bundleFactory.createBundle(left).commit(Instant.now()); CommittedBundle<Integer> rightBundle = bundleFactory.createBundle(right).commit(Instant.now()); EvaluationContext context = mock(EvaluationContext.class); UncommittedBundle<Integer> flattenedLeftBundle = bundleFactory.createBundle(flattened); UncommittedBundle<Integer> flattenedRightBundle = bundleFactory.createBundle(flattened); when(context.createBundle(flattened)).thenReturn(flattenedLeftBundle, flattenedRightBundle); FlattenEvaluatorFactory factory = new FlattenEvaluatorFactory(context); AppliedPTransform<?, ?, ?> flattenedProducer = DirectGraphs.getProducer(flattened); TransformEvaluator<Integer> leftSideEvaluator = factory.forApplication(flattenedProducer, leftBundle); TransformEvaluator<Integer> rightSideEvaluator = factory.forApplication(flattenedProducer, rightBundle); leftSideEvaluator.processElement(WindowedValue.valueInGlobalWindow(1)); rightSideEvaluator.processElement(WindowedValue.valueInGlobalWindow(-1)); leftSideEvaluator.processElement( WindowedValue.timestampedValueInGlobalWindow(2, new Instant(1024))); leftSideEvaluator.processElement(WindowedValue.valueInGlobalWindow(4, PaneInfo.NO_FIRING)); rightSideEvaluator.processElement( WindowedValue.valueInGlobalWindow(2, PaneInfo.ON_TIME_AND_ONLY_FIRING)); rightSideEvaluator.processElement( WindowedValue.timestampedValueInGlobalWindow(-4, new Instant(-4096))); TransformResult<Integer> rightSideResult = rightSideEvaluator.finishBundle(); TransformResult<Integer> leftSideResult = leftSideEvaluator.finishBundle(); assertThat(rightSideResult.getOutputBundles(), Matchers.contains(flattenedRightBundle)); assertThat( rightSideResult.getTransform(), Matchers.<AppliedPTransform<?, ?, ?>>equalTo(flattenedProducer)); assertThat(leftSideResult.getOutputBundles(), Matchers.contains(flattenedLeftBundle)); assertThat( leftSideResult.getTransform(), Matchers.<AppliedPTransform<?, ?, ?>>equalTo(flattenedProducer)); assertThat( flattenedLeftBundle.commit(Instant.now()).getElements(), containsInAnyOrder( WindowedValue.timestampedValueInGlobalWindow(2, new Instant(1024)), WindowedValue.valueInGlobalWindow(4, PaneInfo.NO_FIRING), WindowedValue.valueInGlobalWindow(1))); assertThat( flattenedRightBundle.commit(Instant.now()).getElements(), containsInAnyOrder( WindowedValue.valueInGlobalWindow(2, PaneInfo.ON_TIME_AND_ONLY_FIRING), WindowedValue.timestampedValueInGlobalWindow(-4, new Instant(-4096)), WindowedValue.valueInGlobalWindow(-1))); }
@Test public void testFlattenInMemoryEvaluatorWithEmptyPCollectionList() throws Exception { PCollectionList<Integer> list = PCollectionList.empty(p); PCollection<Integer> flattened = list.apply(Flatten.pCollections()); flattened.setCoder(VarIntCoder.of()); EvaluationContext evaluationContext = mock(EvaluationContext.class); when(evaluationContext.createBundle(flattened)) .thenReturn(bundleFactory.createBundle(flattened)); FlattenEvaluatorFactory factory = new FlattenEvaluatorFactory(evaluationContext); AppliedPTransform<?, ?, ?> flattendProducer = DirectGraphs.getProducer(flattened); TransformEvaluator<Integer> emptyEvaluator = factory.forApplication( flattendProducer, bundleFactory.createRootBundle().commit(BoundedWindow.TIMESTAMP_MAX_VALUE)); TransformResult<Integer> leftSideResult = emptyEvaluator.finishBundle(); CommittedBundle<?> outputBundle = Iterables.getOnlyElement(leftSideResult.getOutputBundles()).commit(Instant.now()); assertThat(outputBundle.getElements(), emptyIterable()); assertThat( leftSideResult.getTransform(), Matchers.<AppliedPTransform<?, ?, ?>>equalTo(flattendProducer)); }