You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by pa...@apache.org on 2019/03/16 06:05:19 UTC
[beam] branch master updated: [BEAM-6735] Add noSpilling option to
WriteFiles.
This is an automated email from the ASF dual-hosted git repository.
pabloem pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 41c9110 [BEAM-6735] Add noSpilling option to WriteFiles.
new 66040ee Merge pull request #7929 from kyle-winkelman/writefiles
41c9110 is described below
commit 41c911002e30b080dd54700dffe7a933de5fc0e1
Author: Kyle Winkelman <ky...@optum.com>
AuthorDate: Fri Feb 22 13:57:40 2019 -0600
[BEAM-6735] Add noSpilling option to WriteFiles.
---
.../main/java/org/apache/beam/sdk/io/AvroIO.java | 15 +++++++++-
.../main/java/org/apache/beam/sdk/io/FileIO.java | 20 +++++++++++++
.../java/org/apache/beam/sdk/io/TFRecordIO.java | 14 +++++++++
.../main/java/org/apache/beam/sdk/io/TextIO.java | 19 ++++++++++++
.../java/org/apache/beam/sdk/io/WriteFiles.java | 34 ++++++++++++++++------
.../org/apache/beam/sdk/io/WriteFilesTest.java | 18 ++++++++++--
6 files changed, 108 insertions(+), 12 deletions(-)
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/AvroIO.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/AvroIO.java
index 2830a04..e0bf7b4 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/AvroIO.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/AvroIO.java
@@ -431,7 +431,8 @@ public class AvroIO {
.setNumShards(0)
.setCodec(TypedWrite.DEFAULT_SERIALIZABLE_CODEC)
.setMetadata(ImmutableMap.of())
- .setWindowedWrites(false);
+ .setWindowedWrites(false)
+ .setNoSpilling(false);
}
private static <T> PCollection<T> setBeamSchema(
@@ -954,6 +955,8 @@ public class AvroIO {
abstract boolean getWindowedWrites();
+ abstract boolean getNoSpilling();
+
@Nullable
abstract FilenamePolicy getFilenamePolicy();
@@ -995,6 +998,8 @@ public class AvroIO {
abstract Builder<UserT, DestinationT, OutputT> setWindowedWrites(boolean windowedWrites);
+ abstract Builder<UserT, DestinationT, OutputT> setNoSpilling(boolean noSpilling);
+
abstract Builder<UserT, DestinationT, OutputT> setFilenamePolicy(
FilenamePolicy filenamePolicy);
@@ -1189,6 +1194,11 @@ public class AvroIO {
return toBuilder().setWindowedWrites(true).build();
}
+ /** See {@link WriteFiles#withNoSpilling()}. */
+ public TypedWrite<UserT, DestinationT, OutputT> withNoSpilling() {
+ return toBuilder().setNoSpilling(true).build();
+ }
+
/** Writes to Avro file(s) compressed using specified codec. */
public TypedWrite<UserT, DestinationT, OutputT> withCodec(CodecFactory codec) {
return toBuilder().setCodec(new SerializableAvroCodecFactory(codec)).build();
@@ -1275,6 +1285,9 @@ public class AvroIO {
if (getWindowedWrites()) {
write = write.withWindowedWrites();
}
+ if (getNoSpilling()) {
+ write = write.withNoSpilling();
+ }
return input.apply("Write", write);
}
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileIO.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileIO.java
index 8d48a28..c663e52 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileIO.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileIO.java
@@ -360,6 +360,7 @@ public class FileIO {
.setDynamic(false)
.setCompression(Compression.UNCOMPRESSED)
.setIgnoreWindowing(false)
+ .setNoSpilling(false)
.build();
}
@@ -372,6 +373,7 @@ public class FileIO {
.setDynamic(true)
.setCompression(Compression.UNCOMPRESSED)
.setIgnoreWindowing(false)
+ .setNoSpilling(false)
.build();
}
@@ -798,6 +800,11 @@ public class FileIO {
return defaultNaming(StaticValueProvider.of(prefix), StaticValueProvider.of(suffix));
}
+ /**
+ * Defines a default {@link FileNaming} which will use the prefix and suffix supplied to create
+ * a name based on the window, pane, number of shards, shard index, and compression. Removes
+ * window when in the {@link GlobalWindow} and pane info when it is the only firing of the pane.
+ */
public static FileNaming defaultNaming(
final ValueProvider<String> prefix, final ValueProvider<String> suffix) {
return (window, pane, numShards, shardIndex, compression) -> {
@@ -894,6 +901,8 @@ public class FileIO {
abstract boolean getIgnoreWindowing();
+ abstract boolean getNoSpilling();
+
abstract Builder<DestinationT, UserT> toBuilder();
@AutoValue.Builder
@@ -938,6 +947,8 @@ public class FileIO {
abstract Builder<DestinationT, UserT> setIgnoreWindowing(boolean ignoreWindowing);
+ abstract Builder<DestinationT, UserT> setNoSpilling(boolean noSpilling);
+
abstract Write<DestinationT, UserT> build();
}
@@ -1159,6 +1170,11 @@ public class FileIO {
return toBuilder().setIgnoreWindowing(true).build();
}
+ /** See {@link WriteFiles#withNoSpilling()}. */
+ public Write<DestinationT, UserT> withNoSpilling() {
+ return toBuilder().setNoSpilling(true).build();
+ }
+
@VisibleForTesting
Contextful<Fn<DestinationT, FileNaming>> resolveFileNamingFn() {
if (getDynamic()) {
@@ -1243,6 +1259,7 @@ public class FileIO {
resolvedSpec.setNumShards(getNumShards());
resolvedSpec.setSharding(getSharding());
resolvedSpec.setIgnoreWindowing(getIgnoreWindowing());
+ resolvedSpec.setNoSpilling(getNoSpilling());
Write<DestinationT, UserT> resolved = resolvedSpec.build();
WriteFiles<UserT, DestinationT, ?> writeFiles =
@@ -1258,6 +1275,9 @@ public class FileIO {
if (!getIgnoreWindowing()) {
writeFiles = writeFiles.withWindowedWrites();
}
+ if (getNoSpilling()) {
+ writeFiles = writeFiles.withNoSpilling();
+ }
return input.apply(writeFiles);
}
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TFRecordIO.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TFRecordIO.java
index 09e50ed..96a753a 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TFRecordIO.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TFRecordIO.java
@@ -84,6 +84,7 @@ public class TFRecordIO {
.setFilenameSuffix(null)
.setNumShards(0)
.setCompression(Compression.UNCOMPRESSED)
+ .setNoSpilling(false)
.build();
}
@@ -231,6 +232,9 @@ public class TFRecordIO {
/** Option to indicate the output sink's compression type. Default is NONE. */
abstract Compression getCompression();
+ /** Whether to skip the spilling of data caused by having maxNumWritersPerBundle. */
+ abstract boolean getNoSpilling();
+
abstract Builder toBuilder();
@AutoValue.Builder
@@ -245,6 +249,8 @@ public class TFRecordIO {
abstract Builder setCompression(Compression compression);
+ abstract Builder setNoSpilling(boolean noSpilling);
+
abstract Write build();
}
@@ -340,6 +346,11 @@ public class TFRecordIO {
return toBuilder().setCompression(compression).build();
}
+ /** See {@link WriteFiles#withNoSpilling()}. */
+ public Write withNoSpilling() {
+ return toBuilder().setNoSpilling(true).build();
+ }
+
@Override
public PDone expand(PCollection<byte[]> input) {
checkState(
@@ -352,6 +363,9 @@ public class TFRecordIO {
if (getNumShards() > 0) {
write = write.withNumShards(getNumShards());
}
+ if (getNoSpilling()) {
+ write = write.withNoSpilling();
+ }
input.apply("Write", write);
return PDone.in(input.getPipeline());
}
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TextIO.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TextIO.java
index 11bc3e4..f3cf940 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TextIO.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TextIO.java
@@ -252,6 +252,7 @@ public class TextIO {
.setWritableByteChannelFactory(FileBasedSink.CompressionType.UNCOMPRESSED)
.setWindowedWrites(false)
.setNumShards(0)
+ .setNoSpilling(false)
.build();
}
@@ -629,6 +630,9 @@ public class TextIO {
/** Whether to write windowed output files. */
abstract boolean getWindowedWrites();
+ /** Whether to skip the spilling of data caused by having maxNumWritersPerBundle. */
+ abstract boolean getNoSpilling();
+
/**
* The {@link WritableByteChannelFactory} to be used by the {@link FileBasedSink}. Default is
* {@link FileBasedSink.CompressionType#UNCOMPRESSED}.
@@ -673,6 +677,8 @@ public class TextIO {
abstract Builder<UserT, DestinationT> setWindowedWrites(boolean windowedWrites);
+ abstract Builder<UserT, DestinationT> setNoSpilling(boolean noSpilling);
+
abstract Builder<UserT, DestinationT> setWritableByteChannelFactory(
WritableByteChannelFactory writableByteChannelFactory);
@@ -899,6 +905,11 @@ public class TextIO {
return toBuilder().setWindowedWrites(true).build();
}
+ /** See {@link WriteFiles#withNoSpilling()}. */
+ public TypedWrite<UserT, DestinationT> withNoSpilling() {
+ return toBuilder().setNoSpilling(true).build();
+ }
+
private DynamicDestinations<UserT, DestinationT, String> resolveDynamicDestinations() {
DynamicDestinations<UserT, DestinationT, String> dynamicDestinations =
getDynamicDestinations();
@@ -981,6 +992,9 @@ public class TextIO {
if (getWindowedWrites()) {
write = write.withWindowedWrites();
}
+ if (getNoSpilling()) {
+ write = write.withNoSpilling();
+ }
return input.apply("WriteFiles", write);
}
@@ -1145,6 +1159,11 @@ public class TextIO {
return new Write(inner.withWindowedWrites());
}
+ /** See {@link TypedWrite#withNoSpilling}. */
+ public Write withNoSpilling() {
+ return new Write(inner.withNoSpilling());
+ }
+
/**
* Specify that output filenames are wanted.
*
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java
index bafcfa7..bbcf109 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java
@@ -271,6 +271,17 @@ public abstract class WriteFiles<UserT, DestinationT, OutputT>
return toBuilder().setWindowedWrites(true).build();
}
+ /**
+ * Returns a new {@link WriteFiles} that writes all data without spilling, simplifying the
+ * pipeline. This option should not be used with {@link #withMaxNumWritersPerBundle(int)} and it
+ * will eliminate this limit possibly causing many writers to be opened. Use with caution.
+ *
+ * <p>This option only applies to writes {@link #withRunnerDeterminedSharding()}.
+ */
+ public WriteFiles<UserT, DestinationT, OutputT> withNoSpilling() {
+ return toBuilder().setMaxNumWritersPerBundle(-1).build();
+ }
+
@Override
public void validate(PipelineOptions options) {
getSink().validate(options);
@@ -403,15 +414,21 @@ public abstract class WriteFiles<UserT, DestinationT, OutputT>
@Override
public PCollection<FileResult<DestinationT>> expand(PCollection<UserT> input) {
+ if (getMaxNumWritersPerBundle() < 0) {
+ return input
+ .apply(
+ "WritedUnshardedBundles",
+ ParDo.of(new WriteUnshardedTempFilesFn(null, destinationCoder))
+ .withSideInputs(getSideInputs()))
+ .setCoder(fileResultCoder);
+ }
TupleTag<FileResult<DestinationT>> writtenRecordsTag = new TupleTag<>("writtenRecords");
TupleTag<KV<ShardedKey<Integer>, UserT>> unwrittenRecordsTag =
new TupleTag<>("unwrittenRecords");
PCollectionTuple writeTuple =
input.apply(
"WriteUnshardedBundles",
- ParDo.of(
- new WriteUnshardedTempFilesWithSpillingFn(
- unwrittenRecordsTag, destinationCoder))
+ ParDo.of(new WriteUnshardedTempFilesFn(unwrittenRecordsTag, destinationCoder))
.withSideInputs(getSideInputs())
.withOutputTags(writtenRecordsTag, TupleTagList.of(unwrittenRecordsTag)));
PCollection<FileResult<DestinationT>> writtenBundleFiles =
@@ -453,9 +470,8 @@ public abstract class WriteFiles<UserT, DestinationT, OutputT>
* Writes all the elements in a bundle using a {@link Writer} produced by the {@link
* WriteOperation} associated with the {@link FileBasedSink}.
*/
- private class WriteUnshardedTempFilesWithSpillingFn
- extends DoFn<UserT, FileResult<DestinationT>> {
- private final TupleTag<KV<ShardedKey<Integer>, UserT>> unwrittenRecordsTag;
+ private class WriteUnshardedTempFilesFn extends DoFn<UserT, FileResult<DestinationT>> {
+ @Nullable private final TupleTag<KV<ShardedKey<Integer>, UserT>> unwrittenRecordsTag;
private final Coder<DestinationT> destinationCoder;
// Initialized in startBundle()
@@ -463,8 +479,8 @@ public abstract class WriteFiles<UserT, DestinationT, OutputT>
private int spilledShardNum = UNKNOWN_SHARDNUM;
- WriteUnshardedTempFilesWithSpillingFn(
- TupleTag<KV<ShardedKey<Integer>, UserT>> unwrittenRecordsTag,
+ WriteUnshardedTempFilesFn(
+ @Nullable TupleTag<KV<ShardedKey<Integer>, UserT>> unwrittenRecordsTag,
Coder<DestinationT> destinationCoder) {
this.unwrittenRecordsTag = unwrittenRecordsTag;
this.destinationCoder = destinationCoder;
@@ -489,7 +505,7 @@ public abstract class WriteFiles<UserT, DestinationT, OutputT>
WriterKey<DestinationT> key = new WriterKey<>(window, c.pane(), destination);
Writer<DestinationT, OutputT> writer = writers.get(key);
if (writer == null) {
- if (writers.size() <= getMaxNumWritersPerBundle()) {
+ if (getMaxNumWritersPerBundle() < 0 || writers.size() <= getMaxNumWritersPerBundle()) {
String uuid = UUID.randomUUID().toString();
LOG.info(
"Opening writer {} for window {} pane {} destination {}",
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/WriteFilesTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/WriteFilesTest.java
index 82e708a..4109dbb 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/WriteFilesTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/WriteFilesTest.java
@@ -332,12 +332,26 @@ public class WriteFilesTest {
}
runWrite(
inputs,
- Window.into(FixedWindows.of(Duration.millis(2))),
+ Window.into(FixedWindows.of(Duration.millis(1))),
+ getBaseOutputFilename(),
+ WriteFiles.to(makeSimpleSink()).withMaxNumWritersPerBundle(2).withWindowedWrites());
+ }
+
+ @Test
+ @Category(NeedsRunner.class)
+ public void testWriteNoSpilling() throws IOException {
+ List<String> inputs = Lists.newArrayList();
+ for (int i = 0; i < 100; ++i) {
+ inputs.add("mambo_number_" + i);
+ }
+ runWrite(
+ inputs,
+ Window.into(FixedWindows.of(Duration.millis(1))),
getBaseOutputFilename(),
WriteFiles.to(makeSimpleSink())
.withMaxNumWritersPerBundle(2)
.withWindowedWrites()
- .withNumShards(1));
+ .withNoSpilling());
}
@Test