This is an automated email from the ASF dual-hosted git repository.
Abacn 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 bd24aab048f [FileWrites] Implementing evictWritersWhenFull option in
WriteFiles (#40200)
bd24aab048f is described below
commit bd24aab048f76317f4acfd5cc13bf9468b197d2d
Author: darshan-sj <[email protected]>
AuthorDate: Mon Sep 21 20:20:33 2026 +0000
[FileWrites] Implementing evictWritersWhenFull option in WriteFiles (#40200)
---
.../java/org/apache/beam/sdk/io/WriteFiles.java | 277 +++++++++++++++++----
.../org/apache/beam/sdk/io/WriteFilesTest.java | 173 +++++++++++++
2 files changed, 399 insertions(+), 51 deletions(-)
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 85829d7e4f0..2fd77a1ad0c 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
@@ -23,8 +23,11 @@ import static
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Pr
import com.google.auto.value.AutoValue;
import java.io.IOException;
+import java.io.ObjectInputStream;
import java.util.ArrayList;
import java.util.Collection;
+import java.util.Iterator;
+import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
@@ -135,9 +138,11 @@ public abstract class WriteFiles<UserT, DestinationT,
OutputT>
public static final Class<? extends WriteFiles> CONCRETE_CLASS =
AutoValue_WriteFiles.class;
// The maximum number of file writers to keep open in a single bundle at a
time, since file
- // writers default to 64mb buffers. This comes into play when writing
per-window files.
- // The first 20 files from a single WriteFiles transform will write files
inline in the
- // transform. Anything beyond that might be shuffled.
+ // writers default to 64mb buffers. This comes into play when writing
per-window or dynamic
+ // destination files. The first 20 files from a single WriteFiles transform
will write files
+ // inline in the transform. Anything beyond that might be spilled to shuffle
(default) or cause
+ // the least recently used open writer in the bundle to be closed and
evicted if
+ // withEvictWritersWhenFull() is enabled.
// Keep in mind that specific runners may decide to run multiple bundles in
parallel, based on
// their own policy.
private static final int DEFAULT_MAX_NUM_WRITERS_PER_BUNDLE = 20;
@@ -171,6 +176,7 @@ public abstract class WriteFiles<UserT, DestinationT,
OutputT>
.setWindowedWrites(false)
.setWithAutoSharding(false)
.setMaxNumWritersPerBundle(DEFAULT_MAX_NUM_WRITERS_PER_BUNDLE)
+ .setEvictWritersWhenFull(false)
.setSideInputs(sink.getDynamicDestinations().getSideInputs())
.setSkipIfEmpty(false)
.setBadRecordErrorHandler(new DefaultErrorHandler<>())
@@ -194,6 +200,8 @@ public abstract class WriteFiles<UserT, DestinationT,
OutputT>
abstract int getMaxNumWritersPerBundle();
+ abstract boolean getEvictWritersWhenFull();
+
abstract boolean getSkipIfEmpty();
abstract @Nullable Integer getBatchSize();
@@ -230,6 +238,9 @@ public abstract class WriteFiles<UserT, DestinationT,
OutputT>
abstract Builder<UserT, DestinationT, OutputT> setMaxNumWritersPerBundle(
int maxNumWritersPerBundle);
+ abstract Builder<UserT, DestinationT, OutputT> setEvictWritersWhenFull(
+ boolean evictWritersWhenFull);
+
abstract Builder<UserT, DestinationT, OutputT> setSkipIfEmpty(boolean
skipIfEmpty);
abstract Builder<UserT, DestinationT, OutputT> setBatchSize(@Nullable
Integer batchSize);
@@ -289,7 +300,54 @@ public abstract class WriteFiles<UserT, DestinationT,
OutputT>
return toBuilder().setNumShardsProvider(numShardsProvider).build();
}
- /** Set the maximum number of writers created in a bundle before spilling to
shuffle. */
+ /**
+ * Set the maximum number of writers kept open in a bundle before spilling
to shuffle (or evicting
+ * the least recently used open writer if {@link
#withEvictWritersWhenFull()} is enabled).
+ *
+ * <p><b>Trade-offs:</b> A higher value here can cause more worker memory
consumption (since each
+ * open writer maintains an in-memory write buffer), but reduces the cost of
shuffling spilled
+ * records (or reduces how frequently writers are closed and evicted when
{@link
+ * #withEvictWritersWhenFull()} is enabled). A lower value reduces peak
memory consumption per
+ * bundle at the cost of either more records spilled to shuffle or more
frequent writer evictions
+ * (resulting in smaller output files).
+ *
+ * <p><b>Writer Limit & Overflow Trade-off Matrix (for {@link
+ * #withRunnerDeterminedSharding()}):</b>
+ *
+ * <table>
+ * <tr>
+ * <th>Configuration</th>
+ * <th>Behavior when {@code maxNumWritersPerBundle} is reached</th>
+ * <th>Worker Memory Consumption</th>
+ * <th>Shuffle Cost</th>
+ * <th>Output File Size / Count</th>
+ * </tr>
+ * <tr>
+ * <td><b>Default (Spill to Shuffle)</b><br>{@code
maxNumWritersPerBundle > 0},<br>{@code evictWritersWhenFull = false}</td>
+ * <td>Keeps first {@code N} writers open; spills remaining records to a
{@link GroupByKey} shuffle stage</td>
+ * <td>Bounded ({@code <= N} buffers per bundle)</td>
+ * <td>High if many records spill across shuffle</td>
+ * <td>Fewer, larger files</td>
+ * </tr>
+ * <tr>
+ * <td><b>LRU Writer Eviction</b><br>{@code maxNumWritersPerBundle >
0},<br>{@code evictWritersWhenFull = true}</td>
+ * <td>Flushes and closes the least recently used (LRU) open writer to
open a new writer inline</td>
+ * <td>Bounded ({@code <= N} buffers per bundle)</td>
+ * <td>None (no shuffle stage for unwritten records)</td>
+ * <td>May produce more/smaller files (minimal if input is ordered by
destination, high if random)</td>
+ * </tr>
+ * <tr>
+ * <td><b>No Spilling</b><br>{@link #withNoSpilling()} ({@code
maxNumWritersPerBundle = -1})</td>
+ * <td>Opens a new writer for every destination in the bundle without
limit</td>
+ * <td>Unbounded (risk of OOM with many destinations)</td>
+ * <td>None (no shuffle stage for unwritten records)</td>
+ * <td>Fewer, larger files (1 file per destination per bundle)</td>
+ * </tr>
+ * </table>
+ *
+ * <p>Note that value provided here cannot exceed the default value ({@link
+ * #DEFAULT_MAX_NUM_WRITERS_PER_BUNDLE}).
+ */
public WriteFiles<UserT, DestinationT, OutputT> withMaxNumWritersPerBundle(
int maxNumWritersPerBundle) {
checkArgument(
@@ -302,6 +360,54 @@ public abstract class WriteFiles<UserT, DestinationT,
OutputT>
return
toBuilder().setMaxNumWritersPerBundle(maxNumWritersPerBundle).build();
}
+ /**
+ * Returns a new {@link WriteFiles} that evicts the least recently used open
writer in the bundle
+ * (LRU order, by flushing and closing it) instead of spilling unwritten
records to shuffle when
+ * {@link #getMaxNumWritersPerBundle()} is reached.
+ *
+ * <p><b>Trade-offs:</b> Setting this to {@code true} avoids the cost of
shuffling records while
+ * keeping concurrent writer memory consumption bounded by {@link
#getMaxNumWritersPerBundle()},
+ * but may lead to smaller and more numerous output files since evicted
writers are closed before
+ * the end of the bundle. See {@link #withMaxNumWritersPerBundle(int)} for
the full trade-off
+ * matrix.
+ *
+ * <p><b>Warning:</b> This option should only be used when the input {@link
PCollection} elements
+ * within a bundle are already grouped or ordered by writer keys
(destination/window/pane), such
+ * that consecutive records belong to the same destination. If the input
{@link PCollection} rows
+ * arrive in random order across more destinations than {@link
#getMaxNumWritersPerBundle()},
+ * writers will be repeatedly closed and reopened, creating too many small
files.
+ *
+ * <p>This option only applies to writes {@link
#withRunnerDeterminedSharding()}.
+ */
+ public WriteFiles<UserT, DestinationT, OutputT> withEvictWritersWhenFull() {
+ return withEvictWritersWhenFull(true);
+ }
+
+ /**
+ * Set this sink to evict the least recently used open writer in the bundle
(LRU order, by
+ * flushing and closing it) when {@link #getMaxNumWritersPerBundle()} is
reached, instead of
+ * spilling unwritten records to shuffle.
+ *
+ * <p><b>Trade-offs:</b> Setting this to {@code true} avoids the cost of
shuffling records while
+ * keeping concurrent writer memory consumption bounded by {@link
#getMaxNumWritersPerBundle()},
+ * but may lead to smaller and more numerous output files since evicted
writers are closed before
+ * the end of the bundle. Setting this to {@code false} (default) preserves
larger output files by
+ * spilling excess records to a shuffle stage. See {@link
#withMaxNumWritersPerBundle(int)} for
+ * the full trade-off matrix.
+ *
+ * <p><b>Warning:</b> This option should only be used when the input {@link
PCollection} elements
+ * within a bundle are already grouped or ordered by writer keys
(destination/window/pane), such
+ * that consecutive records belong to the same destination. If the input
{@link PCollection} rows
+ * arrive in random order across more destinations than {@link
#getMaxNumWritersPerBundle()},
+ * writers will be repeatedly closed and reopened, creating too many small
files.
+ *
+ * <p>This option only applies to writes {@link
#withRunnerDeterminedSharding()}.
+ */
+ public WriteFiles<UserT, DestinationT, OutputT> withEvictWritersWhenFull(
+ boolean evictWritersWhenFull) {
+ return toBuilder().setEvictWritersWhenFull(evictWritersWhenFull).build();
+ }
+
/** Set this sink to skip writing any files if the PCollection is empty. */
public WriteFiles<UserT, DestinationT, OutputT> withSkipIfEmpty(boolean
skipIfEmpty) {
return toBuilder().setSkipIfEmpty(skipIfEmpty).build();
@@ -596,7 +702,10 @@ public abstract class WriteFiles<UserT, DestinationT,
OutputT>
TupleTag<KV<ShardedKey<Integer>, UserT>> unwrittenRecordsTag =
new TupleTag<>("unwrittenRecords");
Coder<UserT> inputCoder = input.getCoder();
- if (getMaxNumWritersPerBundle() < 0) {
+ // When no spilling is requested (maxNumWritersPerBundle < 0) or when
writer eviction is
+ // enabled (evictWritersWhenFull), all records are written inline in
WriteUnshardedTempFilesFn
+ // and no records are spilled to the GroupUnwritten shuffle stage.
+ if (getMaxNumWritersPerBundle() < 0 || getEvictWritersWhenFull()) {
PCollectionTuple writeTuple =
input.apply(
"WritedUnshardedBundles",
@@ -674,8 +783,14 @@ public abstract class WriteFiles<UserT, DestinationT,
OutputT>
private final Coder<UserT> inputCoder;
- // Initialized in startBundle()
- private @Nullable Map<WriterKey<DestinationT>, Writer<DestinationT,
OutputT>> writers;
+ private final Map<WriterKey<DestinationT>, Writer<DestinationT, OutputT>>
writers =
+ new LinkedHashMap<>(
+ /* initialCapacity= */ getMaxNumWritersPerBundle() < 0
+ ? DEFAULT_MAX_NUM_WRITERS_PER_BUNDLE
+ : getMaxNumWritersPerBundle(),
+ /* loadFactor= */ 0.85f,
+ /* accessOrder= */ true); // true = LRU order
+ private final List<FileResult<DestinationT>> evictedFileResults =
Lists.newArrayList();
private int spilledShardNum = UNKNOWN_SHARDNUM;
@@ -691,7 +806,10 @@ public abstract class WriteFiles<UserT, DestinationT,
OutputT>
@StartBundle
public void startBundle(StartBundleContext unused) {
// Reset state in case of reuse. We need to make sure that each bundle
gets unique writers.
- writers = Maps.newHashMap();
+ synchronized (writers) {
+ writers.clear();
+ evictedFileResults.clear();
+ }
}
@ProcessElement
@@ -712,36 +830,31 @@ public abstract class WriteFiles<UserT, DestinationT,
OutputT>
}
DestinationT destination = maybeDestination.destination;
WriterKey<DestinationT> key = new WriterKey<>(window, c.pane(),
destination);
- Writer<DestinationT, OutputT> writer = writers.get(key);
- if (writer == null) {
- if (getMaxNumWritersPerBundle() < 0 || writers.size() <=
getMaxNumWritersPerBundle()) {
- String uuid = UUID.randomUUID().toString();
- LOG.info(
- "Opening writer {} for window {} pane {} destination {}",
- uuid,
- window,
- paneInfo,
- destination);
- writer = writeOperation.createWriter();
- writer.setDestination(destination);
- writer.open(uuid);
- writers.put(key, writer);
- LOG.debug("Done opening writer");
- } else {
- if (spilledShardNum == UNKNOWN_SHARDNUM) {
- // Cache the random value so we only call ThreadLocalRandom once
per DoFn instance.
- spilledShardNum =
ThreadLocalRandom.current().nextInt(SPILLED_RECORD_SHARDING_FACTOR);
+ Writer<DestinationT, OutputT> writer;
+ synchronized (writers) {
+ writer = writers.get(key);
+ if (writer == null) {
+ if (getMaxNumWritersPerBundle() < 0 || writers.size() <=
getMaxNumWritersPerBundle()) {
+ writer = openAndRegisterWriter(key, window, paneInfo, destination);
+ } else if (getEvictWritersWhenFull()) {
+ evictOldestWriter();
+ writer = openAndRegisterWriter(key, window, paneInfo, destination);
} else {
- spilledShardNum = (spilledShardNum + 1) %
SPILLED_RECORD_SHARDING_FACTOR;
+ if (spilledShardNum == UNKNOWN_SHARDNUM) {
+ // Cache the random value so we only call ThreadLocalRandom once
per DoFn instance.
+ spilledShardNum =
ThreadLocalRandom.current().nextInt(SPILLED_RECORD_SHARDING_FACTOR);
+ } else {
+ spilledShardNum = (spilledShardNum + 1) %
SPILLED_RECORD_SHARDING_FACTOR;
+ }
+ outputReceiver
+ .get(unwrittenRecordsTag)
+ .output(
+ KV.of(
+ ShardedKey.of(
+ hashDestination(destination, destinationCoder),
spilledShardNum),
+ c.element()));
+ return;
}
- outputReceiver
- .get(unwrittenRecordsTag)
- .output(
- KV.of(
- ShardedKey.of(
- hashDestination(destination, destinationCoder),
spilledShardNum),
- c.element()));
- return;
}
}
OutputT formattedRecord =
@@ -752,25 +865,88 @@ public abstract class WriteFiles<UserT, DestinationT,
OutputT>
writeOrClose(writer, formattedRecord);
}
- @FinishBundle
- public void finishBundle(FinishBundleContext c) throws Exception {
- for (Map.Entry<WriterKey<DestinationT>, Writer<DestinationT, OutputT>>
entry :
- writers.entrySet()) {
- WriterKey<DestinationT> key = entry.getKey();
- Writer<DestinationT, OutputT> writer = entry.getValue();
+ private Writer<DestinationT, OutputT> openAndRegisterWriter(
+ WriterKey<DestinationT> key,
+ BoundedWindow window,
+ PaneInfo paneInfo,
+ DestinationT destination)
+ throws Exception {
+ synchronized (writers) {
+ String uuid = UUID.randomUUID().toString();
+ LOG.info(
+ "Opening writer {} for window {} pane {} destination {}",
+ uuid,
+ window,
+ paneInfo,
+ destination);
+ Writer<DestinationT, OutputT> writer = writeOperation.createWriter();
+ writer.setDestination(destination);
+ writer.open(uuid);
+ writers.put(key, writer);
+ LOG.debug("Done opening writer");
+ return writer;
+ }
+ }
+
+ private void evictOldestWriter() throws Exception {
+ synchronized (writers) {
+ Iterator<Map.Entry<WriterKey<DestinationT>, Writer<DestinationT,
OutputT>>> iterator =
+ writers.entrySet().iterator();
+ Map.Entry<WriterKey<DestinationT>, Writer<DestinationT, OutputT>>
eldestEntry =
+ iterator.next();
+ iterator.remove();
+
+ WriterKey<DestinationT> evictedKey = eldestEntry.getKey();
+ Writer<DestinationT, OutputT> evictedWriter = eldestEntry.getValue();
+ LOG.info(
+ "Evicting least recently used writer for window {} pane {}
destination {}",
+ evictedKey.window,
+ evictedKey.paneInfo,
+ evictedKey.destination);
try {
- writer.close();
+ evictedWriter.close();
} catch (Exception e) {
// If anything goes wrong, make sure to delete the temporary file.
- writer.cleanup();
+ evictedWriter.cleanup();
throw e;
}
- BoundedWindow window = key.window;
- c.output(
+ evictedFileResults.add(
new FileResult<>(
- writer.getOutputFile(), UNKNOWN_SHARDNUM, window,
key.paneInfo, key.destination),
- window.maxTimestamp(),
- window);
+ evictedWriter.getOutputFile(),
+ UNKNOWN_SHARDNUM,
+ evictedKey.window,
+ evictedKey.paneInfo,
+ evictedKey.destination));
+ }
+ }
+
+ @FinishBundle
+ public void finishBundle(FinishBundleContext c) throws Exception {
+ synchronized (writers) {
+ for (FileResult<DestinationT> evictedResult : evictedFileResults) {
+ BoundedWindow window = evictedResult.getWindow();
+ c.output(evictedResult, window.maxTimestamp(), window);
+ }
+ for (Map.Entry<WriterKey<DestinationT>, Writer<DestinationT, OutputT>>
entry :
+ writers.entrySet()) {
+ WriterKey<DestinationT> key = entry.getKey();
+ Writer<DestinationT, OutputT> writer = entry.getValue();
+ try {
+ writer.close();
+ } catch (Exception e) {
+ // If anything goes wrong, make sure to delete the temporary file.
+ writer.cleanup();
+ throw e;
+ }
+ BoundedWindow window = key.window;
+ c.output(
+ new FileResult<>(
+ writer.getOutputFile(), UNKNOWN_SHARDNUM, window,
key.paneInfo, key.destination),
+ window.maxTimestamp(),
+ window);
+ }
+ writers.clear();
+ evictedFileResults.clear();
}
}
}
@@ -1202,8 +1378,7 @@ public abstract class WriteFiles<UserT, DestinationT,
OutputT>
new ArrayList<>();
// Ensure that transient fields are initialized.
- private void readObject(java.io.ObjectInputStream in)
- throws IOException, ClassNotFoundException {
+ private void readObject(ObjectInputStream in) throws IOException,
ClassNotFoundException {
in.defaultReadObject();
closeFutures = new ArrayList<>();
deferredOutput = new ArrayList<>();
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 cc174002bb4..78a9120517e 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
@@ -30,11 +30,15 @@ import static org.hamcrest.Matchers.not;
import static org.hamcrest.Matchers.nullValue;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertThrows;
import static org.junit.Assert.assertTrue;
import java.io.BufferedReader;
import java.io.File;
import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.channels.WritableByteChannel;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.text.DecimalFormat;
@@ -44,6 +48,8 @@ import java.util.Collections;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ThreadLocalRandom;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BiFunction;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
@@ -88,6 +94,7 @@ import
org.apache.beam.sdk.transforms.windowing.IntervalWindow;
import org.apache.beam.sdk.transforms.windowing.PaneInfo;
import org.apache.beam.sdk.transforms.windowing.Sessions;
import org.apache.beam.sdk.transforms.windowing.Window;
+import org.apache.beam.sdk.util.MimeTypes;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollection.IsBounded;
@@ -491,6 +498,168 @@ public class WriteFilesTest {
.withNoSpilling());
}
+ @Test
+ @Category(NeedsRunner.class)
+ public void testWriteEvictWritersWhenFull() 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()
+ .withEvictWritersWhenFull());
+ }
+
+ private static class EmitElementsInSingleBundleFn extends DoFn<String,
String> {
+ private final List<String> elementsToEmit;
+
+ EmitElementsInSingleBundleFn(List<String> elementsToEmit) {
+ this.elementsToEmit = elementsToEmit;
+ }
+
+ @ProcessElement
+ public void processElement(ProcessContext c) {
+ for (String elem : elementsToEmit) {
+ c.output(elem);
+ }
+ }
+ }
+
+ private static final class FailingCloseOnEvictSink
+ extends FileBasedSink<String, Integer, String> {
+ static final AtomicReference<ResourceId> EVICTED_TEMP_FILE = new
AtomicReference<>();
+ static final AtomicBoolean THREW_ON_CLOSE = new AtomicBoolean(false);
+
+ FailingCloseOnEvictSink(
+ ResourceId tempDirectory,
+ DynamicDestinations<String, Integer, String> dynamicDestinations) {
+ super(StaticValueProvider.of(tempDirectory), dynamicDestinations,
Compression.UNCOMPRESSED);
+ }
+
+ @Override
+ public WriteOperation<Integer, String> createWriteOperation() {
+ return new FailingCloseWriteOperation(this);
+ }
+
+ private static final class FailingCloseWriteOperation extends
WriteOperation<Integer, String> {
+ FailingCloseWriteOperation(FailingCloseOnEvictSink sink) {
+ super(sink);
+ }
+
+ @Override
+ public Writer<Integer, String> createWriter() {
+ return new Writer<Integer, String>(this, MimeTypes.TEXT) {
+ private WritableByteChannel channel;
+
+ @Override
+ protected void prepareWrite(WritableByteChannel channel) {
+ this.channel = channel;
+ }
+
+ @Override
+ public void write(String value) throws Exception {
+ channel.write(ByteBuffer.wrap((value +
"\n").getBytes(StandardCharsets.UTF_8)));
+ }
+
+ @Override
+ protected void writeFooter() throws Exception {
+ if (THREW_ON_CLOSE.compareAndSet(false, true)) {
+ EVICTED_TEMP_FILE.set(getOutputFile());
+ assertTrue(new File(getOutputFile().toString()).exists());
+ throw new IOException("Simulated close failure on eviction");
+ }
+ }
+ };
+ }
+ }
+ }
+
+ @Test
+ @Category(NeedsRunner.class)
+ public void testWriteEvictWritersWhenFullLruOrderAndReopen() throws
IOException {
+ TestDestinations dynamicDestinations = new
TestDestinations(getBaseOutputDirectory());
+ SimpleSink<Integer> sink =
+ new SimpleSink<>(getBaseOutputDirectory(), dynamicDestinations,
Compression.UNCOMPRESSED);
+
+ WriteOptions options =
TestPipeline.testingPipelineOptions().as(WriteOptions.class);
+ options.setTestFlag("test_value");
+ Pipeline p = TestPipeline.create(options);
+
+ // Emit all elements in a single bundle with maxNumWritersPerBundle = 1
(allows up to 2 open
+ // writers simultaneously before evicting):
+ // - "0" (dest 0): opens dest 0. Open writers (LRU -> MRU): [0]
+ // - "1" (dest 1): opens dest 1. Open writers (LRU -> MRU): [0, 1]
+ // - "5" (dest 0): writes to existing open writer for dest 0, refreshing
its recency.
+ // Open writers (LRU -> MRU): [1, 0]
+ // - "2" (dest 2): capacity exceeded -> evicts least recently used writer
(dest 1, even though
+ // dest 0 was opened first), opens dest 2. Open writers
(LRU -> MRU): [0, 2]
+ // - "10" (dest 0): writes to still-open writer for dest 0 (proving dest 1
was evicted, not
+ // dest 0). Open writers (LRU -> MRU): [2, 0]
+ // - "6" (dest 1): dest 1 was previously evicted -> evicts least recently
used writer (dest 2),
+ // re-opens a second writer for dest 1. Open writers (LRU
-> MRU): [0, 1]
+ List<String> bundleElements = Arrays.asList("0", "1", "5", "2", "10", "6");
+
+ WriteFilesResult<Integer> res =
+ p.apply(Create.of("trigger"))
+ .apply(ParDo.of(new EmitElementsInSingleBundleFn(bundleElements)))
+
.apply(WriteFiles.to(sink).withMaxNumWritersPerBundle(1).withEvictWritersWhenFull());
+ res.getPerDestinationOutputFilenames().apply(new VerifyFilesExist<>());
+ p.run();
+
+ // Destination 0 stayed open across all three of its elements due to LRU
refresh -> 1 shard.
+ ResourceId base0 =
+ getBaseOutputDirectory().resolve("file_0",
StandardResolveOptions.RESOLVE_FILE);
+ checkFileContents(
+ base0.toString(), Arrays.asList("record_0", "record_5", "record_10"),
Optional.of(1), true);
+
+ // Destination 1 became LRU after "5" accessed dest 0, was evicted when
"2" arrived, and was
+ // re-opened when "6" arrived -> 2 shards.
+ ResourceId base1 =
+ getBaseOutputDirectory().resolve("file_1",
StandardResolveOptions.RESOLVE_FILE);
+ checkFileContents(
+ base1.toString(), Arrays.asList("record_1", "record_6"),
Optional.of(2), true);
+
+ // Destination 2 was opened once and evicted when "6" arrived -> 1 shard.
+ ResourceId base2 =
+ getBaseOutputDirectory().resolve("file_2",
StandardResolveOptions.RESOLVE_FILE);
+ checkFileContents(
+ base2.toString(), Collections.singletonList("record_2"),
Optional.of(1), true);
+ }
+
+ @Test
+ @Category(NeedsRunner.class)
+ public void testWriteEvictWritersWhenFullCloseExceptionCleanup() {
+ FailingCloseOnEvictSink.EVICTED_TEMP_FILE.set(null);
+ FailingCloseOnEvictSink.THREW_ON_CLOSE.set(false);
+
+ TestDestinations dynamicDestinations = new
TestDestinations(getBaseOutputDirectory());
+ FailingCloseOnEvictSink sink =
+ new FailingCloseOnEvictSink(getBaseOutputDirectory(),
dynamicDestinations);
+
+ WriteOptions options =
TestPipeline.testingPipelineOptions().as(WriteOptions.class);
+ options.setTestFlag("test_value");
+ Pipeline p = TestPipeline.create(options);
+
+ // Emitting 3 distinct destinations ("0", "1", "2") in a single bundle with
+ // maxNumWritersPerBundle = 1 triggers evictOldestWriter() when "2"
arrives.
+ p.apply(Create.of("trigger"))
+ .apply(ParDo.of(new EmitElementsInSingleBundleFn(Arrays.asList("0",
"1", "2"))))
+
.apply(WriteFiles.to(sink).withMaxNumWritersPerBundle(1).withEvictWritersWhenFull());
+
+ Pipeline.PipelineExecutionException thrownException =
+ assertThrows(Pipeline.PipelineExecutionException.class, p::run);
+ assertThat(
+ thrownException.getCause().getMessage(),
+ containsString("Simulated close failure on eviction"));
+ assertNotNull(FailingCloseOnEvictSink.EVICTED_TEMP_FILE.get());
+ assertFalse(new
File(FailingCloseOnEvictSink.EVICTED_TEMP_FILE.get().toString()).exists());
+ }
+
@Test
public void testBuildWrite() {
SimpleSink<Void> sink = makeSimpleSink();
@@ -512,6 +681,10 @@ public class WriteFilesTest {
WriteFiles<String, ?, ?> writeUnsharded =
write2.withRunnerDeterminedSharding();
assertThat(writeUnsharded.getComputeNumShards(), nullValue());
assertThat(write.getComputeNumShards(), equalTo(originalSharding));
+ assertFalse(write.getEvictWritersWhenFull());
+ assertTrue(write.withEvictWritersWhenFull().getEvictWritersWhenFull());
+ assertTrue(write.withEvictWritersWhenFull(true).getEvictWritersWhenFull());
+
assertFalse(write.withEvictWritersWhenFull(false).getEvictWritersWhenFull());
}
@Test
