BoundedReadEvaluatorFactory implements TransformEvaluatorFactory { @SuppressWarnings({"unchecked", "rawtypes"}) @Override public @Nullable <InputT> TransformEvaluator<InputT> forApplication( AppliedPTransform<?, ?, ?> application, CommittedBundle<?> inputBundle) throws IOException { return createEvaluator((AppliedPTransform) application); } BoundedReadEvaluatorFactory(EvaluationContext evaluationContext, PipelineOptions options); @VisibleForTesting  BoundedReadEvaluatorFactory(
      EvaluationContext evaluationContext, PipelineOptions options, long minimumDynamicSplitSize); @SuppressWarnings({"unchecked", "rawtypes"}) @Override @Nullable TransformEvaluator<InputT> forApplication(
      AppliedPTransform<?, ?, ?> application, CommittedBundle<?> inputBundle); @Override void cleanup();  }
@Test public void boundedSourceInMemoryTransformEvaluatorProducesElements() throws Exception { when(context.createRootBundle()).thenReturn(bundleFactory.createRootBundle()); UncommittedBundle<Long> outputBundle = bundleFactory.createBundle(longs); when(context.createBundle(longs)).thenReturn(outputBundle); Collection<CommittedBundle<?>> initialInputs = new BoundedReadEvaluatorFactory.InputProvider(context, options) .getInitialInputs(longsProducer, 1); List<WindowedValue<?>> outputs = new ArrayList<>(); for (CommittedBundle<?> shardBundle : initialInputs) { TransformEvaluator<?> evaluator = factory.forApplication(longsProducer, null); for (WindowedValue<?> shard : shardBundle.getElements()) { evaluator.processElement((WindowedValue) shard); } TransformResult<?> result = evaluator.finishBundle(); assertThat(result.getWatermarkHold(), equalTo(BoundedWindow.TIMESTAMP_MAX_VALUE)); assertThat( Iterables.size(result.getOutputBundles()), equalTo(Iterables.size(shardBundle.getElements()))); for (UncommittedBundle<?> output : result.getOutputBundles()) { CommittedBundle<?> committed = output.commit(BoundedWindow.TIMESTAMP_MAX_VALUE); for (WindowedValue<?> val : committed.getElements()) { outputs.add(val); } } } assertThat( outputs, containsInAnyOrder( gw(1L), gw(2L), gw(4L), gw(8L), gw(9L), gw(7L), gw(6L), gw(5L), gw(3L), gw(0L))); }
@Test public void boundedSourceEvaluatorProducesDynamicSplits() throws Exception { BoundedReadEvaluatorFactory factory = new BoundedReadEvaluatorFactory(context, options, 0L); when(context.createRootBundle()).thenReturn(bundleFactory.createRootBundle()); int numElements = 10; Long[] elems = new Long[numElements]; for (int i = 0; i < numElements; i++) { elems[i] = (long) i; } PCollection<Long> read = p.apply(Read.from(new TestSource<>(VarLongCoder.of(), 5, elems))); SplittableParDo.convertReadBasedSplittableDoFnsToPrimitiveReads(p); AppliedPTransform<?, ?, ?> transform = DirectGraphs.getProducer(read); Collection<CommittedBundle<?>> unreadInputs = new BoundedReadEvaluatorFactory.InputProvider(context, options) .getInitialInputs(transform, 1); Collection<WindowedValue<?>> outputs = new ArrayList<>(); int numIterations = 0; while (!unreadInputs.isEmpty()) { numIterations++; UncommittedBundle<Long> outputBundle = bundleFactory.createBundle(read); when(context.createBundle(read)).thenReturn(outputBundle); Collection<CommittedBundle<?>> newUnreadInputs = new ArrayList<>(); for (CommittedBundle<?> shardBundle : unreadInputs) { TransformEvaluator<Long> evaluator = factory.forApplication(transform, null); for (WindowedValue<?> shard : shardBundle.getElements()) { evaluator.processElement((WindowedValue) shard); } TransformResult<Long> result = evaluator.finishBundle(); assertThat(result.getWatermarkHold(), equalTo(BoundedWindow.TIMESTAMP_MAX_VALUE)); assertThat( Iterables.size(result.getOutputBundles()), equalTo(Iterables.size(shardBundle.getElements()))); for (UncommittedBundle<?> output : result.getOutputBundles()) { CommittedBundle<?> committed = output.commit(BoundedWindow.TIMESTAMP_MAX_VALUE); for (WindowedValue<?> val : committed.getElements()) { outputs.add(val); } } if (!Iterables.isEmpty(result.getUnprocessedElements())) { newUnreadInputs.add(shardBundle.withElements((Iterable) result.getUnprocessedElements())); } } unreadInputs = newUnreadInputs; } assertThat(numIterations, greaterThan(1)); WindowedValue[] expectedValues = new WindowedValue[numElements]; for (long i = 0L; i < numElements; i++) { expectedValues[(int) i] = gw(i); } assertThat(outputs, Matchers.<WindowedValue<?>>containsInAnyOrder(expectedValues)); }
@Test public void boundedSourceEvaluatorDynamicSplitsUnsplittable() throws Exception { BoundedReadEvaluatorFactory factory = new BoundedReadEvaluatorFactory(context, options, 0L); PCollection<Long> read = p.apply(Read.from(SourceTestUtils.toUnsplittableSource(CountingSource.upTo(10L)))); SplittableParDo.convertReadBasedSplittableDoFnsToPrimitiveReads(p); AppliedPTransform<?, ?, ?> transform = DirectGraphs.getProducer(read); when(context.createRootBundle()).thenReturn(bundleFactory.createRootBundle()); when(context.createRootBundle()).thenReturn(bundleFactory.createRootBundle()); Collection<CommittedBundle<?>> initialInputs = new BoundedReadEvaluatorFactory.InputProvider(context, options) .getInitialInputs(transform, 1); UncommittedBundle<Long> outputBundle = bundleFactory.createBundle(read); when(context.createBundle(read)).thenReturn(outputBundle); List<WindowedValue<?>> outputs = new ArrayList<>(); for (CommittedBundle<?> shardBundle : initialInputs) { TransformEvaluator<?> evaluator = factory.forApplication(transform, null); for (WindowedValue<?> shard : shardBundle.getElements()) { evaluator.processElement((WindowedValue) shard); } TransformResult<?> result = evaluator.finishBundle(); assertThat(result.getWatermarkHold(), equalTo(BoundedWindow.TIMESTAMP_MAX_VALUE)); assertThat( Iterables.size(result.getOutputBundles()), equalTo(Iterables.size(shardBundle.getElements()))); for (UncommittedBundle<?> output : result.getOutputBundles()) { CommittedBundle<?> committed = output.commit(BoundedWindow.TIMESTAMP_MAX_VALUE); for (WindowedValue<?> val : committed.getElements()) { outputs.add(val); } } } assertThat( outputs, containsInAnyOrder( gw(1L), gw(2L), gw(4L), gw(8L), gw(9L), gw(7L), gw(6L), gw(5L), gw(3L), gw(0L))); }
@Test public void boundedSourceInMemoryTransformEvaluatorShardsOfSource() throws Exception { PipelineOptions options = PipelineOptionsFactory.create(); List<? extends BoundedSource<Long>> splits = source.split(source.getEstimatedSizeBytes(options) / 2, options); UncommittedBundle<BoundedSourceShard<Long>> rootBundle = bundleFactory.createRootBundle(); for (BoundedSource<Long> split : splits) { BoundedSourceShard<Long> shard = BoundedSourceShard.of(split); rootBundle.add(WindowedValue.valueInGlobalWindow(shard)); } CommittedBundle<BoundedSourceShard<Long>> shards = rootBundle.commit(Instant.now()); TransformEvaluator<BoundedSourceShard<Long>> evaluator = factory.forApplication(longsProducer, shards); for (WindowedValue<BoundedSourceShard<Long>> shard : shards.getElements()) { UncommittedBundle<Long> outputBundle = bundleFactory.createBundle(longs); when(context.createBundle(longs)).thenReturn(outputBundle); evaluator.processElement(shard); } TransformResult<?> result = evaluator.finishBundle(); assertThat(Iterables.size(result.getOutputBundles()), equalTo(splits.size())); List<WindowedValue<?>> outputElems = new ArrayList<>(); for (UncommittedBundle<?> outputBundle : result.getOutputBundles()) { CommittedBundle<?> outputs = outputBundle.commit(Instant.now()); for (WindowedValue<?> outputElem : outputs.getElements()) { outputElems.add(outputElem); } } assertThat( outputElems, containsInAnyOrder( gw(1L), gw(2L), gw(4L), gw(8L), gw(9L), gw(7L), gw(6L), gw(5L), gw(3L), gw(0L))); }
@Test public void boundedSourceEvaluatorClosesReader() throws Exception { TestSource<Long> source = new TestSource<>(BigEndianLongCoder.of(), 1L, 2L, 3L); PCollection<Long> pcollection = p.apply(Read.from(source)); AppliedPTransform<?, ?, ?> sourceTransform = DirectGraphs.getProducer(pcollection); UncommittedBundle<Long> output = bundleFactory.createBundle(pcollection); when(context.createBundle(pcollection)).thenReturn(output); TransformEvaluator<BoundedSourceShard<Long>> evaluator = factory.forApplication( sourceTransform, bundleFactory.createRootBundle().commit(Instant.now())); evaluator.processElement(WindowedValue.valueInGlobalWindow(BoundedSourceShard.of(source))); evaluator.finishBundle(); CommittedBundle<Long> committed = output.commit(Instant.now()); assertThat(committed.getElements(), containsInAnyOrder(gw(2L), gw(3L), gw(1L))); assertThat(TestSource.readerClosed, is(true)); }
@Test public void boundedSourceEvaluatorNoElementsClosesReader() throws Exception { TestSource<Long> source = new TestSource<>(BigEndianLongCoder.of()); PCollection<Long> pcollection = p.apply(Read.from(source)); AppliedPTransform<?, ?, ?> sourceTransform = DirectGraphs.getProducer(pcollection); UncommittedBundle<Long> output = bundleFactory.createBundle(pcollection); when(context.createBundle(pcollection)).thenReturn(output); TransformEvaluator<BoundedSourceShard<Long>> evaluator = factory.forApplication( sourceTransform, bundleFactory.createRootBundle().commit(Instant.now())); evaluator.processElement(WindowedValue.valueInGlobalWindow(BoundedSourceShard.of(source))); evaluator.finishBundle(); CommittedBundle<Long> committed = output.commit(Instant.now()); assertThat(committed.getElements(), emptyIterable()); assertThat(TestSource.readerClosed, is(true)); }