WorkerCustomSources { public static NativeReader<WindowedValue<?>> create( final CloudObject spec, final PipelineOptions options, DataflowExecutionContext executionContext) throws Exception { @SuppressWarnings("unchecked") final Source<Object> source = (Source<Object>) deserializeFromCloudSource(spec); if (source instanceof BoundedSource) { @SuppressWarnings({"unchecked", "rawtypes"}) NativeReader<WindowedValue<?>> reader = (NativeReader) new NativeReader<WindowedValue<Object>>() { @Override public NativeReaderIterator<WindowedValue<Object>> iterator() throws IOException { return new BoundedReaderIterator<>( ((BoundedSource<Object>) source).createReader(options)); } }; return reader; } else if (source instanceof UnboundedSource) { @SuppressWarnings({"unchecked", "rawtypes"}) NativeReader<WindowedValue<?>> reader = (NativeReader) new UnboundedReader<Object>( options, spec, (StreamingModeExecutionContext) executionContext); return reader; } else { throw new IllegalArgumentException("Unexpected source kind: " + source.getClass()); } }  static DynamicSourceSplit toSourceSplit(BoundedSourceSplit<?> sourceSplitResult); static SourceOperationResponse performSplit(
      SourceSplitRequest request, PipelineOptions options); static NativeReader<WindowedValue<?>> create(
      final CloudObject spec,
      final PipelineOptions options,
      DataflowExecutionContext executionContext); static boolean isFirstUnboundedSourceSplit(ByteString splitKey); static final long DATAFLOW_SPLIT_RESPONSE_API_SIZE_LIMIT; }
@Test @SuppressWarnings("unchecked") public void testProgressAndSourceSplitTranslation() throws Exception { NativeReader<WindowedValue<Integer>> reader = (NativeReader<WindowedValue<Integer>>) ReaderRegistry.defaultRegistry() .create( translateIOToCloudSource(CountingSource.upTo(10), options), options, null, TestOperationContext.create()); try (NativeReader.NativeReaderIterator<WindowedValue<Integer>> iterator = reader.iterator()) { assertTrue(iterator.start()); assertEquals(valueInGlobalWindow(0L), iterator.getCurrent()); assertEquals( 0.0, readerProgressToCloudProgress(iterator.getProgress()).getFractionConsumed().doubleValue(), 1e-6); assertTrue(iterator.advance()); assertEquals(valueInGlobalWindow(1L), iterator.getCurrent()); assertEquals( 0.1, readerProgressToCloudProgress(iterator.getProgress()).getFractionConsumed().doubleValue(), 1e-6); assertTrue(iterator.advance()); assertEquals(valueInGlobalWindow(2L), iterator.getCurrent()); assertNull(iterator.requestDynamicSplit(ReaderTestUtils.splitRequestAtFraction(0))); assertNull(iterator.requestDynamicSplit(ReaderTestUtils.splitRequestAtFraction(0.1f))); WorkerCustomSources.BoundedSourceSplit<Integer> sourceSplit = (WorkerCustomSources.BoundedSourceSplit<Integer>) iterator.requestDynamicSplit(ReaderTestUtils.splitRequestAtFraction(0.5f)); assertNotNull(sourceSplit); assertThat(readFromSource(sourceSplit.primary, options), contains(0L, 1L, 2L, 3L, 4L)); assertThat(readFromSource(sourceSplit.residual, options), contains(5L, 6L, 7L, 8L, 9L)); sourceSplit = (WorkerCustomSources.BoundedSourceSplit<Integer>) iterator.requestDynamicSplit(ReaderTestUtils.splitRequestAtFraction(0.8f)); assertNotNull(sourceSplit); assertThat(readFromSource(sourceSplit.primary, options), contains(0L, 1L, 2L, 3L)); assertThat(readFromSource(sourceSplit.residual, options), contains(4L)); assertTrue(iterator.advance()); assertEquals(valueInGlobalWindow(3L), iterator.getCurrent()); assertFalse(iterator.advance()); } }
@Test public void testReadUnboundedReader() throws Exception { CounterSet counterSet = new CounterSet(); StreamingModeExecutionStateRegistry executionStateRegistry = new StreamingModeExecutionStateRegistry(null); StreamingModeExecutionContext context = new StreamingModeExecutionContext( counterSet, "computationId", new ReaderCache(), ImmutableMap.of(), null, StreamingStepMetricsContainer.createRegistry(), new DataflowExecutionStateTracker( ExecutionStateSampler.newForTest(), executionStateRegistry.getState( NameContext.forStage("stageName"), "other", null, NoopProfileScope.NOOP), counterSet, PipelineOptionsFactory.create(), "test-work-item-id"), executionStateRegistry, Long.MAX_VALUE); options.setNumWorkers(5); ByteString state = ByteString.EMPTY; for (int i = 0; i < 10 * WorkerCustomSources.maxUnboundedBundleSize; ) { context.start( "key", Windmill.WorkItem.newBuilder() .setKey(ByteString.copyFromUtf8("0000000000000001")) .setWorkToken(0) .setSourceState( Windmill.SourceState.newBuilder().setState(state).build()) .build(), new Instant(0), null, null, null, null, Windmill.WorkItemCommitRequest.newBuilder()); @SuppressWarnings({"unchecked", "rawtypes"}) NativeReader<WindowedValue<ValueWithRecordId<KV<Integer, Integer>>>> reader = (NativeReader) WorkerCustomSources.create( (CloudObject) serializeToCloudSource(new TestCountingSource(Integer.MAX_VALUE), options) .getSpec(), options, context); Instant beforeReading = Instant.now(); int numReadOnThisIteration = 0; for (WindowedValue<ValueWithRecordId<KV<Integer, Integer>>> value : ReaderUtils.readAllFromReader(reader)) { assertEquals(KV.of(0, i), value.getValue().getValue()); assertArrayEquals( encodeToByteArray(KvCoder.of(VarIntCoder.of(), VarIntCoder.of()), KV.of(0, i)), value.getValue().getId()); assertThat(value.getWindows(), contains((BoundedWindow) GlobalWindow.INSTANCE)); assertEquals(i, value.getTimestamp().getMillis()); i++; numReadOnThisIteration++; } Instant afterReading = Instant.now(); assertThat( new Duration(beforeReading, afterReading).getStandardSeconds(), lessThanOrEqualTo( WorkerCustomSources.MAX_UNBOUNDED_BUNDLE_READ_TIME.getStandardSeconds() + 1)); assertThat( numReadOnThisIteration, lessThanOrEqualTo(WorkerCustomSources.maxUnboundedBundleSize)); context.flushState(); state = context.getOutputBuilder().getSourceStateUpdates().getState(); assertEquals( TimeUnit.MILLISECONDS.toMicros(i), context.getOutputBuilder().getSourceWatermark()); assertEquals( 1, context.getOutputBuilder().getSourceStateUpdates().getFinalizeIdsList().size()); assertNotNull(context.getCachedReader()); assertEquals(7L, context.getBacklogBytes()); } }
@Test @SuppressWarnings("unchecked") public void testProgressAndSourceSplitTranslation() throws Exception { DataflowPipelineOptions options = PipelineOptionsFactory.create().as(DataflowPipelineOptions.class); NativeReader<WindowedValue<Integer>> reader = (NativeReader<WindowedValue<Integer>>) ReaderRegistry.defaultRegistry() .create( translateIOToCloudSource(CountingSource.upTo(10), options), options, null, TestOperationContext.create()); try (NativeReader.NativeReaderIterator<WindowedValue<Integer>> iterator = reader.iterator()) { assertTrue(iterator.start()); assertEquals(valueInGlobalWindow(0L), iterator.getCurrent()); assertEquals( 0.0, readerProgressToCloudProgress(iterator.getProgress()).getFractionConsumed().doubleValue(), 1e-6); assertTrue(iterator.advance()); assertEquals(valueInGlobalWindow(1L), iterator.getCurrent()); assertEquals( 0.1, readerProgressToCloudProgress(iterator.getProgress()).getFractionConsumed().doubleValue(), 1e-6); assertTrue(iterator.advance()); assertEquals(valueInGlobalWindow(2L), iterator.getCurrent()); assertNull(iterator.requestDynamicSplit(ReaderTestUtils.splitRequestAtFraction(0))); assertNull(iterator.requestDynamicSplit(ReaderTestUtils.splitRequestAtFraction(0.1f))); WorkerCustomSources.BoundedSourceSplit<Integer> sourceSplit = (WorkerCustomSources.BoundedSourceSplit<Integer>) iterator.requestDynamicSplit(ReaderTestUtils.splitRequestAtFraction(0.5f)); assertNotNull(sourceSplit); assertThat(readFromSource(sourceSplit.primary, options), contains(0L, 1L, 2L, 3L, 4L)); assertThat(readFromSource(sourceSplit.residual, options), contains(5L, 6L, 7L, 8L, 9L)); sourceSplit = (WorkerCustomSources.BoundedSourceSplit<Integer>) iterator.requestDynamicSplit(ReaderTestUtils.splitRequestAtFraction(0.8f)); assertNotNull(sourceSplit); assertThat(readFromSource(sourceSplit.primary, options), contains(0L, 1L, 2L, 3L)); assertThat(readFromSource(sourceSplit.residual, options), contains(4L)); assertTrue(iterator.advance()); assertEquals(valueInGlobalWindow(3L), iterator.getCurrent()); assertFalse(iterator.advance()); } }
@Test public void testReadUnboundedReader() throws Exception { CounterSet counterSet = new CounterSet(); StreamingModeExecutionStateRegistry executionStateRegistry = new StreamingModeExecutionStateRegistry(null); StreamingModeExecutionContext context = new StreamingModeExecutionContext( counterSet, "computationId", new ReaderCache(), ImmutableMap.of(), null, StreamingStepMetricsContainer.createRegistry(), new DataflowExecutionStateTracker( ExecutionStateSampler.newForTest(), executionStateRegistry.getState( NameContext.forStage("stageName"), "other", null, NoopProfileScope.NOOP), counterSet, PipelineOptionsFactory.create(), "test-work-item-id"), executionStateRegistry, Long.MAX_VALUE); DataflowPipelineOptions options = PipelineOptionsFactory.as(DataflowPipelineOptions.class); options.setNumWorkers(5); ByteString state = ByteString.EMPTY; for (int i = 0; i < 10 * WorkerCustomSources.maxUnboundedBundleSize; ) { context.start( "key", Windmill.WorkItem.newBuilder() .setKey(ByteString.copyFromUtf8("0000000000000001")) .setWorkToken(0) .setSourceState( Windmill.SourceState.newBuilder().setState(state).build()) .build(), new Instant(0), null, null, null, null, Windmill.WorkItemCommitRequest.newBuilder()); @SuppressWarnings({"unchecked", "rawtypes"}) NativeReader<WindowedValue<ValueWithRecordId<KV<Integer, Integer>>>> reader = (NativeReader) WorkerCustomSources.create( (CloudObject) serializeToCloudSource(new TestCountingSource(Integer.MAX_VALUE), options) .getSpec(), options, context); Instant beforeReading = Instant.now(); int numReadOnThisIteration = 0; for (WindowedValue<ValueWithRecordId<KV<Integer, Integer>>> value : ReaderUtils.readAllFromReader(reader)) { assertEquals(KV.of(0, i), value.getValue().getValue()); assertArrayEquals( encodeToByteArray(KvCoder.of(VarIntCoder.of(), VarIntCoder.of()), KV.of(0, i)), value.getValue().getId()); assertThat(value.getWindows(), contains((BoundedWindow) GlobalWindow.INSTANCE)); assertEquals(i, value.getTimestamp().getMillis()); i++; numReadOnThisIteration++; } Instant afterReading = Instant.now(); assertThat( new Duration(beforeReading, afterReading).getStandardSeconds(), lessThanOrEqualTo( WorkerCustomSources.MAX_UNBOUNDED_BUNDLE_READ_TIME.getStandardSeconds() + 1)); assertThat( numReadOnThisIteration, lessThanOrEqualTo(WorkerCustomSources.maxUnboundedBundleSize)); context.flushState(); state = context.getOutputBuilder().getSourceStateUpdates().getState(); assertEquals( TimeUnit.MILLISECONDS.toMicros(i), context.getOutputBuilder().getSourceWatermark()); assertEquals( 1, context.getOutputBuilder().getSourceStateUpdates().getFinalizeIdsList().size()); assertNotNull(context.getCachedReader()); assertEquals(7L, context.getBacklogBytes()); } }