UnboundedReadEvaluatorFactory implements TransformEvaluatorFactory { @SuppressWarnings({"unchecked", "rawtypes"}) @Override public @Nullable <InputT> TransformEvaluator<InputT> forApplication( AppliedPTransform<?, ?, ?> application, CommittedBundle<?> inputBundle) { return createEvaluator((AppliedPTransform) application); } UnboundedReadEvaluatorFactory(EvaluationContext evaluationContext, PipelineOptions options); @VisibleForTesting  UnboundedReadEvaluatorFactory(
      EvaluationContext evaluationContext, PipelineOptions options, double readerReuseChance); @SuppressWarnings({"unchecked", "rawtypes"}) @Override @Nullable TransformEvaluator<InputT> forApplication(
      AppliedPTransform<?, ?, ?> application, CommittedBundle<?> inputBundle); @Override void cleanup();  }
@Test public void unboundedSourceInMemoryTransformEvaluatorProducesElements() throws Exception { when(context.createRootBundle()).thenReturn(bundleFactory.createRootBundle()); Collection<CommittedBundle<?>> initialInputs = new UnboundedReadEvaluatorFactory.InputProvider(context, p.getOptions()) .getInitialInputs(graph.getProducer(longs), 1); CommittedBundle<?> inputShards = Iterables.getOnlyElement(initialInputs); UnboundedSourceShard<Long, ?> inputShard = (UnboundedSourceShard<Long, ?>) Iterables.getOnlyElement(inputShards.getElements()).getValue(); TransformEvaluator<? super UnboundedSourceShard<Long, ?>> evaluator = factory.forApplication(graph.getProducer(longs), inputShards); evaluator.processElement((WindowedValue) Iterables.getOnlyElement(inputShards.getElements())); TransformResult<? super UnboundedSourceShard<Long, ?>> result = evaluator.finishBundle(); WindowedValue<? super UnboundedSourceShard<Long, ?>> residual = Iterables.getOnlyElement(result.getUnprocessedElements()); assertThat(residual.getTimestamp(), Matchers.lessThan(DateTime.now().toInstant())); UnboundedSourceShard<Long, ?> residualShard = (UnboundedSourceShard<Long, ?>) residual.getValue(); assertThat(residualShard.getSource(), equalTo(inputShard.getSource())); assertThat(residualShard.getCheckpoint(), not(nullValue())); assertThat( output.commit(Instant.now()).getElements(), containsInAnyOrder( tgw(1L), tgw(2L), tgw(4L), tgw(8L), tgw(9L), tgw(7L), tgw(6L), tgw(5L), tgw(3L), tgw(0L))); }
@Test public void unboundedSourceWithDuplicatesMultipleCalls() throws Exception { Long[] outputs = new Long[20]; for (long i = 0L; i < 20L; i++) { outputs[(int) i] = i % 5L; } TestUnboundedSource<Long> source = new TestUnboundedSource<>(BigEndianLongCoder.of(), outputs); source.dedupes = true; PCollection<Long> pcollection = p.apply(Read.from(source)); SplittableParDo.convertReadBasedSplittableDoFnsToPrimitiveReads(p); AppliedPTransform<?, ?, ?> sourceTransform = getProducer(pcollection); when(context.createRootBundle()).thenReturn(bundleFactory.createRootBundle()); Collection<CommittedBundle<?>> initialInputs = new UnboundedReadEvaluatorFactory.InputProvider(context, p.getOptions()) .getInitialInputs(sourceTransform, 1); UncommittedBundle<Long> output = bundleFactory.createBundle(pcollection); when(context.createBundle(pcollection)).thenReturn(output); CommittedBundle<?> inputBundle = Iterables.getOnlyElement(initialInputs); TransformEvaluator<UnboundedSourceShard<Long, TestCheckpointMark>> evaluator = factory.forApplication(sourceTransform, inputBundle); for (WindowedValue<?> value : inputBundle.getElements()) { evaluator.processElement( (WindowedValue<UnboundedSourceShard<Long, TestCheckpointMark>>) value); } TransformResult<UnboundedSourceShard<Long, TestCheckpointMark>> result = evaluator.finishBundle(); assertThat( output.commit(Instant.now()).getElements(), containsInAnyOrder(tgw(1L), tgw(2L), tgw(4L), tgw(3L), tgw(0L))); UncommittedBundle<Long> secondOutput = bundleFactory.createBundle(longs); when(context.createBundle(longs)).thenReturn(secondOutput); TransformEvaluator<UnboundedSourceShard<Long, TestCheckpointMark>> secondEvaluator = factory.forApplication(sourceTransform, inputBundle); WindowedValue<UnboundedSourceShard<Long, TestCheckpointMark>> residual = (WindowedValue<UnboundedSourceShard<Long, TestCheckpointMark>>) Iterables.getOnlyElement(result.getUnprocessedElements()); secondEvaluator.processElement(residual); secondEvaluator.finishBundle(); assertThat(secondOutput.commit(Instant.now()).getElements(), Matchers.emptyIterable()); }
@Test public void noElementsAvailableReaderIncludedInResidual() throws Exception { PCollection<Long> pcollection = p.apply(Read.from(new TestUnboundedSource<>(VarLongCoder.of(), 1L))); SplittableParDo.convertReadBasedSplittableDoFnsToPrimitiveReads(p); AppliedPTransform<?, ?, ?> sourceTransform = getProducer(pcollection); when(context.createRootBundle()).thenReturn(bundleFactory.createRootBundle()); Collection<CommittedBundle<?>> initialInputs = new UnboundedReadEvaluatorFactory.InputProvider(context, p.getOptions()) .getInitialInputs(sourceTransform, 1); when(context.createBundle(pcollection)).thenReturn(bundleFactory.createBundle(pcollection)); CommittedBundle<?> inputBundle = Iterables.getOnlyElement(initialInputs); TransformEvaluator<UnboundedSourceShard<Long, TestCheckpointMark>> evaluator = factory.forApplication(sourceTransform, inputBundle); for (WindowedValue<?> value : inputBundle.getElements()) { evaluator.processElement( (WindowedValue<UnboundedSourceShard<Long, TestCheckpointMark>>) value); } TransformResult<UnboundedSourceShard<Long, TestCheckpointMark>> result = evaluator.finishBundle(); UncommittedBundle<Long> secondOutput = bundleFactory.createBundle(longs); when(context.createBundle(longs)).thenReturn(secondOutput); TransformEvaluator<UnboundedSourceShard<Long, TestCheckpointMark>> secondEvaluator = factory.forApplication(sourceTransform, inputBundle); WindowedValue<UnboundedSourceShard<Long, TestCheckpointMark>> residual = (WindowedValue<UnboundedSourceShard<Long, TestCheckpointMark>>) Iterables.getOnlyElement(result.getUnprocessedElements()); secondEvaluator.processElement(residual); TransformResult<UnboundedSourceShard<Long, TestCheckpointMark>> secondResult = secondEvaluator.finishBundle(); assertThat(secondOutput.commit(Instant.now()).getElements(), Matchers.emptyIterable()); WindowedValue<UnboundedSourceShard<Long, TestCheckpointMark>> unprocessed = (WindowedValue<UnboundedSourceShard<Long, TestCheckpointMark>>) Iterables.getOnlyElement(secondResult.getUnprocessedElements()); assertThat(unprocessed.getTimestamp(), Matchers.greaterThan(residual.getTimestamp())); assertThat(unprocessed.getValue().getExistingReader(), not(nullValue())); }
@Test public void evaluatorReusesReaderAndClosesAtTheEnd() throws Exception { int numElements = 1000; ContiguousSet<Long> elems = ContiguousSet.create(Range.openClosed(0L, (long) numElements), DiscreteDomain.longs()); TestUnboundedSource<Long> source = new TestUnboundedSource<>(BigEndianLongCoder.of(), elems.toArray(new Long[0])); source.advanceWatermarkToInfinity = true; PCollection<Long> pcollection = p.apply(Read.from(source)); SplittableParDo.convertReadBasedSplittableDoFnsToPrimitiveReads(p); DirectGraph graph = DirectGraphs.getGraph(p); AppliedPTransform<?, ?, ?> sourceTransform = graph.getProducer(pcollection); when(context.createRootBundle()).thenReturn(bundleFactory.createRootBundle()); UncommittedBundle<Long> output = mock(UncommittedBundle.class); when(context.createBundle(pcollection)).thenReturn(output); WindowedValue<UnboundedSourceShard<Long, TestCheckpointMark>> shard = WindowedValue.valueInGlobalWindow( UnboundedSourceShard.unstarted(source, NeverDeduplicator.create())); CommittedBundle<UnboundedSourceShard<Long, TestCheckpointMark>> inputBundle = bundleFactory .<UnboundedSourceShard<Long, TestCheckpointMark>>createRootBundle() .add(shard) .commit(Instant.now()); UnboundedReadEvaluatorFactory factory = new UnboundedReadEvaluatorFactory(context, p.getOptions(), 1.0 ); new UnboundedReadEvaluatorFactory.InputProvider(context, p.getOptions()) .getInitialInputs(sourceTransform, 1); CommittedBundle<UnboundedSourceShard<Long, TestCheckpointMark>> residual = inputBundle; do { TransformEvaluator<UnboundedSourceShard<Long, TestCheckpointMark>> evaluator = factory.forApplication(sourceTransform, residual); evaluator.processElement(Iterables.getOnlyElement(residual.getElements())); TransformResult<UnboundedSourceShard<Long, TestCheckpointMark>> result = evaluator.finishBundle(); residual = inputBundle.withElements( (Iterable<WindowedValue<UnboundedSourceShard<Long, TestCheckpointMark>>>) result.getUnprocessedElements()); } while (!Iterables.isEmpty(residual.getElements())); verify(output, times(numElements)).add(any()); assertThat(TestUnboundedSource.readerCreatedCount, equalTo(1)); assertThat(TestUnboundedSource.readerClosedCount, equalTo(1)); }
@Test public void evaluatorClosesReaderAndResumesFromCheckpoint() throws Exception { ContiguousSet<Long> elems = ContiguousSet.create(Range.closed(0L, 20L), DiscreteDomain.longs()); TestUnboundedSource<Long> source = new TestUnboundedSource<>(BigEndianLongCoder.of(), elems.toArray(new Long[0])); PCollection<Long> pcollection = p.apply(Read.from(source)); AppliedPTransform<?, ?, ?> sourceTransform = DirectGraphs.getGraph(p).getProducer(pcollection); when(context.createRootBundle()).thenReturn(bundleFactory.createRootBundle()); UncommittedBundle<Long> output = bundleFactory.createBundle(pcollection); when(context.createBundle(pcollection)).thenReturn(output); WindowedValue<UnboundedSourceShard<Long, TestCheckpointMark>> shard = WindowedValue.valueInGlobalWindow( UnboundedSourceShard.unstarted(source, NeverDeduplicator.create())); CommittedBundle<UnboundedSourceShard<Long, TestCheckpointMark>> inputBundle = bundleFactory .<UnboundedSourceShard<Long, TestCheckpointMark>>createRootBundle() .add(shard) .commit(Instant.now()); UnboundedReadEvaluatorFactory factory = new UnboundedReadEvaluatorFactory(context, p.getOptions(), 0.0 ); TransformEvaluator<UnboundedSourceShard<Long, TestCheckpointMark>> evaluator = factory.forApplication(sourceTransform, inputBundle); evaluator.processElement(shard); TransformResult<UnboundedSourceShard<Long, TestCheckpointMark>> result = evaluator.finishBundle(); CommittedBundle<UnboundedSourceShard<Long, TestCheckpointMark>> residual = inputBundle.withElements( (Iterable<WindowedValue<UnboundedSourceShard<Long, TestCheckpointMark>>>) result.getUnprocessedElements()); TransformEvaluator<UnboundedSourceShard<Long, TestCheckpointMark>> secondEvaluator = factory.forApplication(sourceTransform, residual); secondEvaluator.processElement(Iterables.getOnlyElement(residual.getElements())); secondEvaluator.finishBundle(); assertThat(TestUnboundedSource.readerClosedCount, equalTo(2)); assertThat( Iterables.getOnlyElement(residual.getElements()).getValue().getCheckpoint().isFinalized(), is(true)); }
@Test public void evaluatorThrowsInCloseRethrows() throws Exception { ContiguousSet<Long> elems = ContiguousSet.create(Range.closed(0L, 20L), DiscreteDomain.longs()); TestUnboundedSource<Long> source = new TestUnboundedSource<>(BigEndianLongCoder.of(), elems.toArray(new Long[0])) .throwsOnClose(); PCollection<Long> pcollection = p.apply(Read.from(source)); AppliedPTransform<?, ?, ?> sourceTransform = DirectGraphs.getGraph(p).getProducer(pcollection); when(context.createRootBundle()).thenReturn(bundleFactory.createRootBundle()); UncommittedBundle<Long> output = bundleFactory.createBundle(pcollection); when(context.createBundle(pcollection)).thenReturn(output); WindowedValue<UnboundedSourceShard<Long, TestCheckpointMark>> shard = WindowedValue.valueInGlobalWindow( UnboundedSourceShard.unstarted(source, NeverDeduplicator.create())); CommittedBundle<UnboundedSourceShard<Long, TestCheckpointMark>> inputBundle = bundleFactory .<UnboundedSourceShard<Long, TestCheckpointMark>>createRootBundle() .add(shard) .commit(Instant.now()); UnboundedReadEvaluatorFactory factory = new UnboundedReadEvaluatorFactory(context, p.getOptions(), 0.0 ); TransformEvaluator<UnboundedSourceShard<Long, TestCheckpointMark>> evaluator = factory.forApplication(sourceTransform, inputBundle); thrown.expect(IOException.class); thrown.expectMessage("throws on close"); evaluator.processElement(shard); }