WorkerCustomSources { public static SourceOperationResponse performSplit( SourceSplitRequest request, PipelineOptions options) throws Exception { return performSplitWithApiLimit( request, options, DEFAULT_NUM_BUNDLES_LIMIT, DATAFLOW_SPLIT_RESPONSE_API_SIZE_LIMIT); }  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 public void testSplitAndReadBundlesBack() throws Exception { com.google.api.services.dataflow.model.Source source = translateIOToCloudSource(CountingSource.upTo(10L), options); List<WindowedValue<Integer>> elems = readElemsFromSource(options, source); assertEquals(10L, elems.size()); for (long i = 0; i < 10L; i++) { assertEquals(valueInGlobalWindow(i), elems.get((int) i)); } SourceSplitResponse response = performSplit( source, options, 16L , null , null ); assertEquals("SOURCE_SPLIT_OUTCOME_SPLITTING_HAPPENED", response.getOutcome()); List<DerivedSource> bundles = response.getBundles(); assertEquals(5, bundles.size()); for (int i = 0; i < 5; ++i) { DerivedSource bundle = bundles.get(i); assertEquals("SOURCE_DERIVATION_MODE_INDEPENDENT", bundle.getDerivationMode()); com.google.api.services.dataflow.model.Source bundleSource = bundle.getSource(); assertTrue(bundleSource.getDoesNotNeedSplitting()); bundleSource.setCodec(source.getCodec()); List<WindowedValue<Integer>> xs = readElemsFromSource(options, bundleSource); assertThat( "Failed on bundle " + i, xs, contains(valueInGlobalWindow(0L + 2 * i), valueInGlobalWindow(1L + 2 * i))); assertTrue(bundle.getSource().getMetadata().getEstimatedSizeBytes() > 0); } }
@Test public void testSplittingProducedInvalidSource() throws Exception { com.google.api.services.dataflow.model.Source cloudSource = translateIOToCloudSource(new SourceProducingInvalidSplits("original", null), options); expectedException.expect(IllegalArgumentException.class); expectedException.expectMessage( allOf( containsString("Splitting a valid source produced an invalid source"), containsString("original"), containsString("badBundle"))); expectedException.expectCause(hasMessage(containsString("intentionally invalid"))); performSplit( cloudSource, options, null , null , null ); }
@Test public void testLargeSerializedSizeResplits() throws Exception { final long apiSizeLimitForTest = 5 * 1024; BoundedSource<Long> justForSizing = CountingSource.upTo(1000000L); long size = DataflowApiUtils.computeSerializedSizeBytes( translateIOToCloudSource(justForSizing, options)); long numberToSplitToExceedLimit = 2 * apiSizeLimitForTest / size; checkState( numberToSplitToExceedLimit < WorkerCustomSources.DEFAULT_NUM_BUNDLES_LIMIT, "This test expects the number of splits to be less than %s " + "to avoid using SplittableOnlyBoundedSource", WorkerCustomSources.DEFAULT_NUM_BUNDLES_LIMIT); com.google.api.services.dataflow.model.Source source = translateIOToCloudSource(CountingSource.upTo(numberToSplitToExceedLimit), options); SourceSplitResponse split = performSplit(source, options, 8L, null , apiSizeLimitForTest); logged.verifyWarn("too large for the Google Cloud Dataflow API"); logged.verifyWarn(String.format("%d bundles", numberToSplitToExceedLimit)); assertThat((long) split.getBundles().size(), lessThan(numberToSplitToExceedLimit)); }
@Test public void testOversplittingDesiredBundleSizeScaledFirst() throws Exception { com.google.api.services.dataflow.model.Source source = translateIOToCloudSource(CountingSource.upTo(8000), options); SourceSplitResponse bundledWithOnlyNumBundlesLimit = performSplit( source, options, 8L, 100 , 10000 * 1024L ); assertEquals(100, bundledWithOnlyNumBundlesLimit.getBundles().size()); assertThat( DataflowApiUtils.computeSerializedSizeBytes(bundledWithOnlyNumBundlesLimit), greaterThan(10 * 1024L)); SourceSplitResponse bundledWithOnlySizeLimit = performSplit( source, options, 8L, 1000000 , 10 * 1024L ); int numBundlesWithOnlySizeLimit = bundledWithOnlySizeLimit.getBundles().size(); assertThat(numBundlesWithOnlySizeLimit, lessThan(100)); SourceSplitResponse bundledWithSizeLimit = performSplit(source, options, 8L, 100, 10 * 1024L); assertEquals(numBundlesWithOnlySizeLimit, bundledWithSizeLimit.getBundles().size()); }
@Test public void testTooLargeSplitResponseFails() throws Exception { com.google.api.services.dataflow.model.Source source = translateIOToCloudSource(CountingSource.upTo(1000), options); expectedException.expectMessage("[0, 1000)"); expectedException.expectMessage("larger than the limit 100"); performSplit(source, options, 8L, 10, 100L); }
@Test public void testSplitAndReadBundlesBack() throws Exception { DataflowPipelineOptions options = PipelineOptionsFactory.as(DataflowPipelineOptions.class); com.google.api.services.dataflow.model.Source source = translateIOToCloudSource(CountingSource.upTo(10L), options); List<WindowedValue<Integer>> elems = readElemsFromSource(options, source); assertEquals(10L, elems.size()); for (long i = 0; i < 10L; i++) { assertEquals(valueInGlobalWindow(i), elems.get((int) i)); } SourceSplitResponse response = performSplit( source, options, 16L , null , null ); assertEquals("SOURCE_SPLIT_OUTCOME_SPLITTING_HAPPENED", response.getOutcome()); List<DerivedSource> bundles = response.getBundles(); assertEquals(5, bundles.size()); for (int i = 0; i < 5; ++i) { DerivedSource bundle = bundles.get(i); assertEquals("SOURCE_DERIVATION_MODE_INDEPENDENT", bundle.getDerivationMode()); com.google.api.services.dataflow.model.Source bundleSource = bundle.getSource(); assertTrue(bundleSource.getDoesNotNeedSplitting()); bundleSource.setCodec(source.getCodec()); List<WindowedValue<Integer>> xs = readElemsFromSource(options, bundleSource); assertThat( "Failed on bundle " + i, xs, contains(valueInGlobalWindow(0L + 2 * i), valueInGlobalWindow(1L + 2 * i))); assertTrue(bundle.getSource().getMetadata().getEstimatedSizeBytes() > 0); } }
@Test public void testSplittingProducedInvalidSource() throws Exception { DataflowPipelineOptions options = PipelineOptionsFactory.as(DataflowPipelineOptions.class); com.google.api.services.dataflow.model.Source cloudSource = translateIOToCloudSource(new SourceProducingInvalidSplits("original", null), options); expectedException.expect(IllegalArgumentException.class); expectedException.expectMessage( allOf( containsString("Splitting a valid source produced an invalid source"), containsString("original"), containsString("badBundle"))); expectedException.expectCause(hasMessage(containsString("intentionally invalid"))); performSplit( cloudSource, options, null , null , null ); }
@Test public void testLargeSerializedSizeResplits() throws Exception { final long apiSizeLimitForTest = 5 * 1024; DataflowPipelineOptions options = PipelineOptionsFactory.as(DataflowPipelineOptions.class); BoundedSource<Long> justForSizing = CountingSource.upTo(1000000L); long size = DataflowApiUtils.computeSerializedSizeBytes( translateIOToCloudSource(justForSizing, options)); long numberToSplitToExceedLimit = 2 * apiSizeLimitForTest / size; checkState( numberToSplitToExceedLimit < WorkerCustomSources.DEFAULT_NUM_BUNDLES_LIMIT, "This test expects the number of splits to be less than %s " + "to avoid using SplittableOnlyBoundedSource", WorkerCustomSources.DEFAULT_NUM_BUNDLES_LIMIT); com.google.api.services.dataflow.model.Source source = translateIOToCloudSource(CountingSource.upTo(numberToSplitToExceedLimit), options); SourceSplitResponse split = performSplit(source, options, 8L, null , apiSizeLimitForTest); logged.verifyWarn("too large for the Google Cloud Dataflow API"); logged.verifyWarn(String.format("%d bundles", numberToSplitToExceedLimit)); assertThat((long) split.getBundles().size(), lessThan(numberToSplitToExceedLimit)); }
@Test public void testOversplittingDesiredBundleSizeScaledFirst() throws Exception { DataflowPipelineOptions options = PipelineOptionsFactory.as(DataflowPipelineOptions.class); com.google.api.services.dataflow.model.Source source = translateIOToCloudSource(CountingSource.upTo(8000), options); SourceSplitResponse bundledWithOnlyNumBundlesLimit = performSplit( source, options, 8L, 100 , 10000 * 1024L ); assertEquals(100, bundledWithOnlyNumBundlesLimit.getBundles().size()); assertThat( DataflowApiUtils.computeSerializedSizeBytes(bundledWithOnlyNumBundlesLimit), greaterThan(10 * 1024L)); SourceSplitResponse bundledWithOnlySizeLimit = performSplit( source, options, 8L, 1000000 , 10 * 1024L ); int numBundlesWithOnlySizeLimit = bundledWithOnlySizeLimit.getBundles().size(); assertThat(numBundlesWithOnlySizeLimit, lessThan(100)); SourceSplitResponse bundledWithSizeLimit = performSplit(source, options, 8L, 100, 10 * 1024L); assertEquals(numBundlesWithOnlySizeLimit, bundledWithSizeLimit.getBundles().size()); }
@Test public void testTooLargeSplitResponseFails() throws Exception { DataflowPipelineOptions options = PipelineOptionsFactory.as(DataflowPipelineOptions.class); com.google.api.services.dataflow.model.Source source = translateIOToCloudSource(CountingSource.upTo(1000), options); expectedException.expectMessage("[0, 1000)"); expectedException.expectMessage("larger than the limit 100"); performSplit(source, options, 8L, 10, 100L); }