/** * Constructs the sink. */ SequenceFileSink( ValueProvider<ResourceId> baseOutputDirectoryProvider, FilenamePolicy filenamePolicy, Class<K> keyClass, Class<? extends Serialization<? super K>> keySerializationClass, Class<V> valueClass, Class<? extends Serialization<? super V>> valueSerializationClass) { super(baseOutputDirectoryProvider, DynamicFileDestinations.<KV<K,V>>constant(filenamePolicy), Compression.UNCOMPRESSED); this.keyClass = keyClass; this.valueClass = valueClass; Set<String> serializationNameSet = Sets.newHashSet(); serializationNameSet.add(keySerializationClass.getName()); serializationNameSet.add(valueSerializationClass.getName()); serializationNames = serializationNameSet.toArray(new String[serializationNameSet.size()]); }
@VisibleForTesting TFRecordSink( ValueProvider<ResourceId> outputPrefix, @Nullable String shardTemplate, @Nullable String suffix, Compression compression) { super( outputPrefix, DynamicFileDestinations.constant( DefaultFilenamePolicy.fromStandardParameters( outputPrefix, shardTemplate, suffix, false)), compression); }
TestSink(String tmpFolder) { super( ValueProvider.StaticValueProvider.of(FileSystems.matchNewResource(tmpFolder, true)), DynamicFileDestinations.constant(FILENAME_POLICY, SerializableFunctions.identity())); }
DummySink() { super( StaticValueProvider.of(FileSystems.matchNewResource("nowhere", false)), DynamicFileDestinations.constant( new DummyFilenamePolicy(), SerializableFunctions.constant(null))); }
TestSink(String tmpFolder) { super( StaticValueProvider.of(FileSystems.matchNewResource(tmpFolder, true)), DynamicFileDestinations.constant( new FilenamePolicy() { @Override public ResourceId windowedFilename( int shardNumber, int numShards, BoundedWindow window, PaneInfo paneInfo, OutputFileHints outputFileHints) { throw new UnsupportedOperationException("should not be called"); } @Nullable @Override public ResourceId unwindowedFilename( int shardNumber, int numShards, OutputFileHints outputFileHints) { throw new UnsupportedOperationException("should not be called"); } }, SerializableFunctions.identity())); }
public static SimpleSink<Void> makeSimpleSink( ResourceId baseDirectory, String prefix, String shardTemplate, String suffix, WritableByteChannelFactory writableByteChannelFactory) { DynamicDestinations<String, Void, String> dynamicDestinations = DynamicFileDestinations.constant( DefaultFilenamePolicy.fromParams( new Params() .withBaseFilename( baseDirectory.resolve(prefix, StandardResolveOptions.RESOLVE_FILE)) .withShardTemplate(shardTemplate) .withSuffix(suffix))); return new SimpleSink<>(baseDirectory, dynamicDestinations, writableByteChannelFactory); }
@Test public void withNoShardingSpecifiedReturnsNewTransform() { ResourceId outputDirectory = LocalResources.fromString("/foo", true /* isDirectory */); PTransform<PCollection<Object>, WriteFilesResult<Void>> original = WriteFiles.to( new FileBasedSink<Object, Void, Object>( StaticValueProvider.of(outputDirectory), DynamicFileDestinations.constant(new FakeFilenamePolicy())) { @Override public WriteOperation<Void, Object> createWriteOperation() { throw new IllegalArgumentException("Should not be used"); } }); @SuppressWarnings("unchecked") PCollection<Object> objs = (PCollection) p.apply(Create.empty(VoidCoder.of())); AppliedPTransform< PCollection<Object>, WriteFilesResult<Void>, PTransform<PCollection<Object>, WriteFilesResult<Void>>> originalApplication = AppliedPTransform.of("write", objs.expand(), Collections.emptyMap(), original, p); assertThat( factory.getReplacementTransform(originalApplication).getTransform(), not(equalTo((Object) original))); }
private DynamicDestinations<UserT, DestinationT, String> resolveDynamicDestinations() { DynamicDestinations<UserT, DestinationT, String> dynamicDestinations = getDynamicDestinations(); if (dynamicDestinations == null) { if (getDestinationFunction() != null) { // In this case, DestinationT == Params dynamicDestinations = (DynamicDestinations) DynamicFileDestinations.toDefaultPolicies( getDestinationFunction(), getEmptyDestination(), getFormatFunction()); } else { // In this case, DestinationT == Void FilenamePolicy usedFilenamePolicy = getFilenamePolicy(); if (usedFilenamePolicy == null) { usedFilenamePolicy = DefaultFilenamePolicy.fromStandardParameters( getFilenamePrefix(), getShardTemplate(), getFilenameSuffix(), getWindowedWrites()); } dynamicDestinations = (DynamicDestinations) DynamicFileDestinations.constant(usedFilenamePolicy, getFormatFunction()); } } return dynamicDestinations; }
@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)); }
WriteFiles.to( new FileBasedSink<Integer, Void, Integer>( StaticValueProvider.of(outputDirectory), DynamicFileDestinations.constant(policy)) { @Override public WriteOperation<Void, Integer> createWriteOperation() {
@Test public void testShardedDisplayData() { 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).withNumShards(1); DisplayData displayData = DisplayData.from(write); assertThat(displayData, hasDisplayItem("sink", sink.getClass())); assertThat(displayData, includesDisplayDataFor("sink", sink)); assertThat(displayData, hasDisplayItem("numShards", 1)); }
@Test public void testCustomShardStrategyDisplayData() { DynamicDestinations<String, Void, String> dynamicDestinations = DynamicFileDestinations.constant( DefaultFilenamePolicy.fromParams( new Params()