WriteFiles extends PTransform<PCollection<UserT>, WriteFilesResult<DestinationT>> { public static <UserT, DestinationT, OutputT> WriteFiles<UserT, DestinationT, OutputT> to( FileBasedSink<UserT, DestinationT, OutputT> sink) { checkArgument(sink != null, "sink can not be null"); return new AutoValue_WriteFiles.Builder<UserT, DestinationT, OutputT>() .setSink(sink) .setComputeNumShards(null) .setNumShardsProvider(null) .setWindowedWrites(false) .setMaxNumWritersPerBundle(DEFAULT_MAX_NUM_WRITERS_PER_BUNDLE) .setSideInputs(sink.getDynamicDestinations().getSideInputs()) .build(); }  static WriteFiles<UserT, DestinationT, OutputT> to(
      FileBasedSink<UserT, DestinationT, OutputT> sink); abstract FileBasedSink<UserT, DestinationT, OutputT> getSink(); abstract @Nullable PTransform<PCollection<UserT>, PCollectionView<Integer>> getComputeNumShards(); abstract @Nullable ValueProvider<Integer> getNumShardsProvider(); abstract boolean getWindowedWrites(); abstract @Nullable ShardingFunction<UserT, DestinationT> getShardingFunction(); @Override Map<TupleTag<?>, PValue> getAdditionalInputs(); WriteFiles<UserT, DestinationT, OutputT> withNumShards(int numShards); WriteFiles<UserT, DestinationT, OutputT> withNumShards(
      ValueProvider<Integer> numShardsProvider); WriteFiles<UserT, DestinationT, OutputT> withMaxNumWritersPerBundle(
      int maxNumWritersPerBundle); WriteFiles<UserT, DestinationT, OutputT> withSideInputs(
      List<PCollectionView<?>> sideInputs); WriteFiles<UserT, DestinationT, OutputT> withSharding(
      PTransform<PCollection<UserT>, PCollectionView<Integer>> sharding); WriteFiles<UserT, DestinationT, OutputT> withRunnerDeterminedSharding(); WriteFiles<UserT, DestinationT, OutputT> withShardingFunction(
      ShardingFunction<UserT, DestinationT> shardingFunction); WriteFiles<UserT, DestinationT, OutputT> withWindowedWrites(); WriteFiles<UserT, DestinationT, OutputT> withNoSpilling(); @Override void validate(PipelineOptions options); @Override WriteFilesResult<DestinationT> expand(PCollection<UserT> input); @Override void populateDisplayData(DisplayData.Builder builder); @Internal
  static final Class<? extends WriteFiles> CONCRETE_CLASS; }
@Test @Category(NeedsRunner.class) public void testWrite() throws IOException { List<String> inputs = Arrays.asList( "Critical canary", "Apprehensive eagle", "Intimidating pigeon", "Pedantic gull", "Frisky finch"); runWrite(inputs, IDENTITY_MAP, getBaseOutputFilename(), WriteFiles.to(makeSimpleSink())); }
@Test @Category(NeedsRunner.class) public void testEmptyWrite() throws IOException { runWrite( Collections.emptyList(), IDENTITY_MAP, getBaseOutputFilename(), WriteFiles.to(makeSimpleSink())); checkFileContents( getBaseOutputFilename(), Collections.emptyList(), Optional.of(1), true ); }
@Test @Category(NeedsRunner.class) public void testShardedWrite() throws IOException { runShardedWrite( Arrays.asList("one", "two", "three", "four", "five", "six"), IDENTITY_MAP, getBaseOutputFilename(), WriteFiles.to(makeSimpleSink())); }
@Test @Category(NeedsRunner.class) public void testWriteWithEmptyPCollection() throws IOException { List<String> inputs = new ArrayList<>(); runWrite(inputs, IDENTITY_MAP, getBaseOutputFilename(), WriteFiles.to(makeSimpleSink())); }
@Test @Category(NeedsRunner.class) public void testWriteWindowed() throws IOException { List<String> inputs = Arrays.asList( "Critical canary", "Apprehensive eagle", "Intimidating pigeon", "Pedantic gull", "Frisky finch"); runWrite( inputs, new WindowAndReshuffle<>(Window.into(FixedWindows.of(Duration.millis(2)))), getBaseOutputFilename(), WriteFiles.to(makeSimpleSink())); }
@Test @Category(NeedsRunner.class) public void testWriteWithSessions() throws IOException { List<String> inputs = Arrays.asList( "Critical canary", "Apprehensive eagle", "Intimidating pigeon", "Pedantic gull", "Frisky finch"); runWrite( inputs, new WindowAndReshuffle<>(Window.into(Sessions.withGapDuration(Duration.millis(1)))), getBaseOutputFilename(), WriteFiles.to(makeSimpleSink())); }
@Test public void testDisplayData() { DynamicDestinations<String, Void, String> dynamicDestinations = DynamicFileDestinations.constant( DefaultFilenamePolicy.fromParams( new Params() .withBaseFilename( getBaseOutputDirectory() .resolve("file", StandardResolveOptions.RESOLVE_FILE)) .withShardTemplate("-SS-of-NN"))); SimpleSink<Void> sink = new SimpleSink<Void>( getBaseOutputDirectory(), dynamicDestinations, Compression.UNCOMPRESSED) { @Override public void populateDisplayData(DisplayData.Builder builder) { builder.add(DisplayData.item("foo", "bar")); } }; WriteFiles<String, ?, String> write = WriteFiles.to(sink); DisplayData displayData = DisplayData.from(write); assertThat(displayData, hasDisplayItem("sink", sink.getClass())); assertThat(displayData, includesDisplayDataFor("sink", sink)); }
@Test @Category(NeedsRunner.class) public void testUnboundedNeedsWindowed() { thrown.expect(IllegalArgumentException.class); thrown.expectMessage( "Must use windowed writes when applying WriteFiles to an unbounded PCollection"); SimpleSink<Void> sink = makeSimpleSink(); p.apply(Create.of("foo")).setIsBoundedInternal(IsBounded.UNBOUNDED).apply(WriteFiles.to(sink)); p.run(); }