From 627bad78fe7439dc226a917511b737b58be122d9 Mon Sep 17 00:00:00 2001 From: Darshan Siddesh Jagaluru Date: Mon, 21 Sep 2026 12:44:21 +0000 Subject: [PATCH 1/5] Implementing evictWritersWhenFull option in WriteFiles --- .../org/apache/beam/sdk/io/WriteFiles.java | 141 +++++++++++++--- .../apache/beam/sdk/io/WriteFilesTest.java | 157 ++++++++++++++++++ 2 files changed, 278 insertions(+), 20 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 85829d7e4f00..ee375418b106 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,10 @@ 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.List; import java.util.Map; import java.util.UUID; @@ -135,9 +137,11 @@ public abstract class WriteFiles public static final Class 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 oldest 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 +175,7 @@ public static WriteFiles()) @@ -194,6 +199,8 @@ public static WriteFiles setNumShardsProvider( abstract Builder setMaxNumWritersPerBundle( int maxNumWritersPerBundle); + abstract Builder setEvictWritersWhenFull( + boolean evictWritersWhenFull); + abstract Builder setSkipIfEmpty(boolean skipIfEmpty); abstract Builder setBatchSize(@Nullable Integer batchSize); @@ -289,7 +299,10 @@ public WriteFiles withNumShards( 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 oldest open writer if {@link #withEvictWritersWhenFull()} is enabled). + */ public WriteFiles withMaxNumWritersPerBundle( int maxNumWritersPerBundle) { checkArgument( @@ -302,6 +315,41 @@ public WriteFiles withMaxNumWritersPerBundle( return toBuilder().setMaxNumWritersPerBundle(maxNumWritersPerBundle).build(); } + /** + * Returns a new {@link WriteFiles} that evicts the oldest open writer in the bundle (FIFO order, + * by flushing and closing it) instead of spilling unwritten records to shuffle when {@link + * #getMaxNumWritersPerBundle()} is reached. + * + *

Warning: 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. + * + *

This option only applies to writes {@link #withRunnerDeterminedSharding()}. + */ + public WriteFiles withEvictWritersWhenFull() { + return withEvictWritersWhenFull(true); + } + + /** + * Set this sink to evict the oldest open writer in the bundle (FIFO order, by flushing and + * closing it) when {@link #getMaxNumWritersPerBundle()} is reached, instead of spilling unwritten + * records to shuffle. + * + *

Warning: 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. + * + *

This option only applies to writes {@link #withRunnerDeterminedSharding()}. + */ + public WriteFiles withEvictWritersWhenFull( + boolean evictWritersWhenFull) { + return toBuilder().setEvictWritersWhenFull(evictWritersWhenFull).build(); + } + /** Set this sink to skip writing any files if the PCollection is empty. */ public WriteFiles withSkipIfEmpty(boolean skipIfEmpty) { return toBuilder().setSkipIfEmpty(skipIfEmpty).build(); @@ -596,7 +644,10 @@ public PCollection> expand(PCollection input) { TupleTag, UserT>> unwrittenRecordsTag = new TupleTag<>("unwrittenRecords"); Coder 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", @@ -676,6 +727,7 @@ private class WriteUnshardedTempFilesFn extends DoFn, Writer> writers; + private @Nullable List> evictedFileResults; private int spilledShardNum = UNKNOWN_SHARDNUM; @@ -691,7 +743,10 @@ private class WriteUnshardedTempFilesFn extends DoFn 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"); + writer = openAndRegisterWriter(key, window, paneInfo, destination); + } else if (getEvictWritersWhenFull()) { + evictOldestWriter(); + writer = openAndRegisterWriter(key, window, paneInfo, destination); } else { if (spilledShardNum == UNKNOWN_SHARDNUM) { // Cache the random value so we only call ThreadLocalRandom once per DoFn instance. @@ -752,8 +799,63 @@ public void processElement( writeOrClose(writer, formattedRecord); } + private Writer openAndRegisterWriter( + WriterKey key, + BoundedWindow window, + PaneInfo paneInfo, + DestinationT destination) + throws Exception { + String uuid = UUID.randomUUID().toString(); + LOG.info( + "Opening writer {} for window {} pane {} destination {}", + uuid, + window, + paneInfo, + destination); + Writer 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 { + Iterator, Writer>> iterator = + writers.entrySet().iterator(); + Map.Entry, Writer> eldestEntry = + iterator.next(); + iterator.remove(); + + WriterKey evictedKey = eldestEntry.getKey(); + Writer evictedWriter = eldestEntry.getValue(); + LOG.info( + "Evicting oldest writer for window {} pane {} destination {}", + evictedKey.window, + evictedKey.paneInfo, + evictedKey.destination); + try { + evictedWriter.close(); + } catch (Exception e) { + // If anything goes wrong, make sure to delete the temporary file. + evictedWriter.cleanup(); + throw e; + } + evictedFileResults.add( + new FileResult<>( + evictedWriter.getOutputFile(), + UNKNOWN_SHARDNUM, + evictedKey.window, + evictedKey.paneInfo, + evictedKey.destination)); + } + @FinishBundle public void finishBundle(FinishBundleContext c) throws Exception { + for (FileResult evictedResult : evictedFileResults) { + BoundedWindow window = evictedResult.getWindow(); + c.output(evictedResult, window.maxTimestamp(), window); + } for (Map.Entry, Writer> entry : writers.entrySet()) { WriterKey key = entry.getKey(); @@ -1202,8 +1304,7 @@ public WriteShardsIntoTempFilesFn(Coder inputCoder) { 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 cc174002bb46..380a8b660543 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,6 +30,7 @@ import static org.hamcrest.Matchers.nullValue; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import java.io.BufferedReader; @@ -44,6 +45,7 @@ import java.util.List; import java.util.Set; import java.util.concurrent.ThreadLocalRandom; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.BiFunction; import java.util.regex.Matcher; import java.util.regex.Pattern; @@ -491,6 +493,157 @@ public void testWriteNoSpilling() throws IOException { .withNoSpilling()); } + @Test + @Category(NeedsRunner.class) + public void testWriteEvictWritersWhenFull() throws IOException { + List 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 { + private final List elementsToEmit; + + EmitElementsInSingleBundleFn(List elementsToEmit) { + this.elementsToEmit = elementsToEmit; + } + + @ProcessElement + public void processElement(ProcessContext c) { + for (String elem : elementsToEmit) { + c.output(elem); + } + } + } + + private static final class FailingCloseOnEvictSink extends SimpleSink { + static final AtomicBoolean CLEANUP_CALLED = new AtomicBoolean(false); + static final AtomicBoolean THREW_ON_CLOSE = new AtomicBoolean(false); + + FailingCloseOnEvictSink( + ResourceId tempDirectory, + DynamicDestinations dynamicDestinations) { + super(tempDirectory, dynamicDestinations, Compression.UNCOMPRESSED); + } + + @Override + public WriteOperation createWriteOperation() { + return new FailingCloseWriteOperation(this); + } + + private static final class FailingCloseWriteOperation extends WriteOperation { + FailingCloseWriteOperation(FailingCloseOnEvictSink sink) { + super(sink); + } + + @Override + public Writer createWriter() { + return new SimpleWriter(new SimpleWriteOperation<>(getSink())) { + @Override + public void close() throws Exception { + super.close(); + if (THREW_ON_CLOSE.compareAndSet(false, true)) { + throw new IOException("Simulated close failure on eviction"); + } + } + + @Override + public void cleanup() throws Exception { + CLEANUP_CALLED.set(true); + super.cleanup(); + } + }; + } + } + } + + @Test + @Category(NeedsRunner.class) + public void testWriteEvictWritersWhenFullFifoOrderAndReopen() throws IOException { + TestDestinations dynamicDestinations = new TestDestinations(getBaseOutputDirectory()); + SimpleSink 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 (FIFO): [0] + // - "1" (dest 1): opens dest 1. Open writers (FIFO): [0, 1] + // - "6" (dest 1): writes to existing open writer for dest 1. Open writers (FIFO): [0, 1] + // - "2" (dest 2): capacity exceeded -> evicts oldest writer (dest 0), opens dest 2. + // Open writers (FIFO): [1, 2] + // - "11" (dest 1): writes to still-open writer for dest 1 (proving dest 0 was evicted first, + // not dest 1). Open writers (FIFO): [1, 2] + // - "5" (dest 0): dest 0 was previously evicted -> evicts oldest writer (dest 1), re-opens a + // second writer for dest 0. Open writers (FIFO): [2, 0] + List bundleElements = Arrays.asList("0", "1", "6", "2", "11", "5"); + + WriteFilesResult 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 was opened twice (evicted once and closed once at finishBundle) -> 2 shards. + ResourceId base0 = + getBaseOutputDirectory().resolve("file_0", StandardResolveOptions.RESOLVE_FILE); + checkFileContents( + base0.toString(), Arrays.asList("record_0", "record_5"), Optional.of(2), true); + + // Destination 1 stayed open for all three of its elements before being evicted -> 1 shard. + ResourceId base1 = + getBaseOutputDirectory().resolve("file_1", StandardResolveOptions.RESOLVE_FILE); + checkFileContents( + base1.toString(), Arrays.asList("record_1", "record_6", "record_11"), Optional.of(1), true); + + // Destination 2 was opened once and closed at finishBundle -> 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.CLEANUP_CALLED.set(false); + 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")); + assertTrue(FailingCloseOnEvictSink.CLEANUP_CALLED.get()); + } + @Test public void testBuildWrite() { SimpleSink sink = makeSimpleSink(); @@ -512,6 +665,10 @@ public void testBuildWrite() { WriteFiles 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 From 7ebd2d0d82d8e9dd2b82389fef358d0791fd170e Mon Sep 17 00:00:00 2001 From: Darshan Siddesh Jagaluru Date: Mon, 21 Sep 2026 12:57:15 +0000 Subject: [PATCH 2/5] Correcting Unit tests --- .../apache/beam/sdk/io/WriteFilesTest.java | 42 ++++++++++++------- 1 file changed, 28 insertions(+), 14 deletions(-) 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 380a8b660543..c6df03ca117b 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,12 +30,15 @@ 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; @@ -46,6 +49,7 @@ 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; @@ -90,6 +94,7 @@ 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; @@ -525,14 +530,15 @@ public void processElement(ProcessContext c) { } } - private static final class FailingCloseOnEvictSink extends SimpleSink { - static final AtomicBoolean CLEANUP_CALLED = new AtomicBoolean(false); + private static final class FailingCloseOnEvictSink + extends FileBasedSink { + static final AtomicReference EVICTED_TEMP_FILE = new AtomicReference<>(); static final AtomicBoolean THREW_ON_CLOSE = new AtomicBoolean(false); FailingCloseOnEvictSink( ResourceId tempDirectory, DynamicDestinations dynamicDestinations) { - super(tempDirectory, dynamicDestinations, Compression.UNCOMPRESSED); + super(StaticValueProvider.of(tempDirectory), dynamicDestinations, Compression.UNCOMPRESSED); } @Override @@ -547,19 +553,26 @@ private static final class FailingCloseWriteOperation extends WriteOperation createWriter() { - return new SimpleWriter(new SimpleWriteOperation<>(getSink())) { + return new Writer(this, MimeTypes.TEXT) { + private WritableByteChannel channel; + @Override - public void close() throws Exception { - super.close(); - if (THREW_ON_CLOSE.compareAndSet(false, true)) { - throw new IOException("Simulated close failure on eviction"); - } + 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 - public void cleanup() throws Exception { - CLEANUP_CALLED.set(true); - super.cleanup(); + 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"); + } } }; } @@ -619,7 +632,7 @@ public void testWriteEvictWritersWhenFullFifoOrderAndReopen() throws IOException @Test @Category(NeedsRunner.class) public void testWriteEvictWritersWhenFullCloseExceptionCleanup() { - FailingCloseOnEvictSink.CLEANUP_CALLED.set(false); + FailingCloseOnEvictSink.EVICTED_TEMP_FILE.set(null); FailingCloseOnEvictSink.THREW_ON_CLOSE.set(false); TestDestinations dynamicDestinations = new TestDestinations(getBaseOutputDirectory()); @@ -641,7 +654,8 @@ public void testWriteEvictWritersWhenFullCloseExceptionCleanup() { assertThat( thrownException.getCause().getMessage(), containsString("Simulated close failure on eviction")); - assertTrue(FailingCloseOnEvictSink.CLEANUP_CALLED.get()); + assertNotNull(FailingCloseOnEvictSink.EVICTED_TEMP_FILE.get()); + assertFalse(new File(FailingCloseOnEvictSink.EVICTED_TEMP_FILE.get().toString()).exists()); } @Test From ed0157f3405108bddcd97eb24c70d5a76aae625f Mon Sep 17 00:00:00 2001 From: Darshan Siddesh Jagaluru Date: Mon, 21 Sep 2026 15:39:58 +0000 Subject: [PATCH 3/5] Changing the eviction algorithm to LRU --- .../org/apache/beam/sdk/io/WriteFiles.java | 209 ++++++++++-------- .../apache/beam/sdk/io/WriteFilesTest.java | 34 +-- 2 files changed, 132 insertions(+), 111 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 ee375418b106..7959877b62e4 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 @@ -27,6 +27,7 @@ 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; @@ -140,8 +141,8 @@ public abstract class WriteFiles // 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 oldest open writer in the bundle to be closed and evicted if withEvictWritersWhenFull() - // is enabled. + // 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; @@ -301,7 +302,7 @@ public WriteFiles withNumShards( /** * Set the maximum number of writers kept open in a bundle before spilling to shuffle (or evicting - * the oldest open writer if {@link #withEvictWritersWhenFull()} is enabled). + * the least recently used open writer if {@link #withEvictWritersWhenFull()} is enabled). */ public WriteFiles withMaxNumWritersPerBundle( int maxNumWritersPerBundle) { @@ -316,9 +317,9 @@ public WriteFiles withMaxNumWritersPerBundle( } /** - * Returns a new {@link WriteFiles} that evicts the oldest open writer in the bundle (FIFO order, - * by flushing and closing it) instead of spilling unwritten records to shuffle when {@link - * #getMaxNumWritersPerBundle()} is reached. + * 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. * *

Warning: 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 @@ -333,9 +334,9 @@ public WriteFiles withEvictWritersWhenFull() { } /** - * Set this sink to evict the oldest open writer in the bundle (FIFO order, by flushing and - * closing it) when {@link #getMaxNumWritersPerBundle()} is reached, instead of spilling unwritten - * records to shuffle. + * 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. * *

Warning: 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 @@ -725,9 +726,16 @@ private class WriteUnshardedTempFilesFn extends DoFn inputCoder; - // Initialized in startBundle() - private @Nullable Map, Writer> writers; - private @Nullable List> evictedFileResults; + // Initial capacity is maxNumWritersPerBundle (or DEFAULT_MAX_NUM_WRITERS_PER_BUNDLE if -1), + // loadFactor=0.85f prevents rehashing, and accessOrder=true enables LRU eviction order. + private final Map, Writer> writers = + new LinkedHashMap<>( + getMaxNumWritersPerBundle() < 0 + ? DEFAULT_MAX_NUM_WRITERS_PER_BUNDLE + : getMaxNumWritersPerBundle(), + 0.85f, + true); + private final List> evictedFileResults = Lists.newArrayList(); private int spilledShardNum = UNKNOWN_SHARDNUM; @@ -743,10 +751,10 @@ private class WriteUnshardedTempFilesFn extends DoFn key = new WriterKey<>(window, c.pane(), destination); - Writer 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 { - 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 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 = @@ -805,74 +816,82 @@ private Writer openAndRegisterWriter( PaneInfo paneInfo, DestinationT destination) throws Exception { - String uuid = UUID.randomUUID().toString(); - LOG.info( - "Opening writer {} for window {} pane {} destination {}", - uuid, - window, - paneInfo, - destination); - Writer 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 { - Iterator, Writer>> iterator = - writers.entrySet().iterator(); - Map.Entry, Writer> eldestEntry = - iterator.next(); - iterator.remove(); - - WriterKey evictedKey = eldestEntry.getKey(); - Writer evictedWriter = eldestEntry.getValue(); - LOG.info( - "Evicting oldest writer for window {} pane {} destination {}", - evictedKey.window, - evictedKey.paneInfo, - evictedKey.destination); - try { - evictedWriter.close(); - } catch (Exception e) { - // If anything goes wrong, make sure to delete the temporary file. - evictedWriter.cleanup(); - throw e; + synchronized (writers) { + String uuid = UUID.randomUUID().toString(); + LOG.info( + "Opening writer {} for window {} pane {} destination {}", + uuid, + window, + paneInfo, + destination); + Writer writer = writeOperation.createWriter(); + writer.setDestination(destination); + writer.open(uuid); + writers.put(key, writer); + LOG.debug("Done opening writer"); + return writer; } - evictedFileResults.add( - new FileResult<>( - evictedWriter.getOutputFile(), - UNKNOWN_SHARDNUM, - evictedKey.window, - evictedKey.paneInfo, - evictedKey.destination)); } - @FinishBundle - public void finishBundle(FinishBundleContext c) throws Exception { - for (FileResult evictedResult : evictedFileResults) { - BoundedWindow window = evictedResult.getWindow(); - c.output(evictedResult, window.maxTimestamp(), window); - } - for (Map.Entry, Writer> entry : - writers.entrySet()) { - WriterKey key = entry.getKey(); - Writer writer = entry.getValue(); + private void evictOldestWriter() throws Exception { + synchronized (writers) { + Iterator, Writer>> iterator = + writers.entrySet().iterator(); + Map.Entry, Writer> eldestEntry = + iterator.next(); + iterator.remove(); + + WriterKey evictedKey = eldestEntry.getKey(); + Writer 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 evictedResult : evictedFileResults) { + BoundedWindow window = evictedResult.getWindow(); + c.output(evictedResult, window.maxTimestamp(), window); + } + for (Map.Entry, Writer> entry : + writers.entrySet()) { + WriterKey key = entry.getKey(); + Writer 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(); } } } 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 c6df03ca117b..78a9120517ef 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 @@ -581,7 +581,7 @@ protected void writeFooter() throws Exception { @Test @Category(NeedsRunner.class) - public void testWriteEvictWritersWhenFullFifoOrderAndReopen() throws IOException { + public void testWriteEvictWritersWhenFullLruOrderAndReopen() throws IOException { TestDestinations dynamicDestinations = new TestDestinations(getBaseOutputDirectory()); SimpleSink sink = new SimpleSink<>(getBaseOutputDirectory(), dynamicDestinations, Compression.UNCOMPRESSED); @@ -592,16 +592,17 @@ public void testWriteEvictWritersWhenFullFifoOrderAndReopen() throws IOException // 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 (FIFO): [0] - // - "1" (dest 1): opens dest 1. Open writers (FIFO): [0, 1] - // - "6" (dest 1): writes to existing open writer for dest 1. Open writers (FIFO): [0, 1] - // - "2" (dest 2): capacity exceeded -> evicts oldest writer (dest 0), opens dest 2. - // Open writers (FIFO): [1, 2] - // - "11" (dest 1): writes to still-open writer for dest 1 (proving dest 0 was evicted first, - // not dest 1). Open writers (FIFO): [1, 2] - // - "5" (dest 0): dest 0 was previously evicted -> evicts oldest writer (dest 1), re-opens a - // second writer for dest 0. Open writers (FIFO): [2, 0] - List bundleElements = Arrays.asList("0", "1", "6", "2", "11", "5"); + // - "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 bundleElements = Arrays.asList("0", "1", "5", "2", "10", "6"); WriteFilesResult res = p.apply(Create.of("trigger")) @@ -610,19 +611,20 @@ public void testWriteEvictWritersWhenFullFifoOrderAndReopen() throws IOException res.getPerDestinationOutputFilenames().apply(new VerifyFilesExist<>()); p.run(); - // Destination 0 was opened twice (evicted once and closed once at finishBundle) -> 2 shards. + // 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"), Optional.of(2), true); + base0.toString(), Arrays.asList("record_0", "record_5", "record_10"), Optional.of(1), true); - // Destination 1 stayed open for all three of its elements before being evicted -> 1 shard. + // 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", "record_11"), Optional.of(1), true); + base1.toString(), Arrays.asList("record_1", "record_6"), Optional.of(2), true); - // Destination 2 was opened once and closed at finishBundle -> 1 shard. + // Destination 2 was opened once and evicted when "6" arrived -> 1 shard. ResourceId base2 = getBaseOutputDirectory().resolve("file_2", StandardResolveOptions.RESOLVE_FILE); checkFileContents( From 39e7aacd4c6906feea16cf15318b166253e833d5 Mon Sep 17 00:00:00 2001 From: Darshan Siddesh Jagaluru Date: Mon, 21 Sep 2026 15:56:31 +0000 Subject: [PATCH 4/5] Added detailed comments on the tradeoff --- .../org/apache/beam/sdk/io/WriteFiles.java | 57 +++++++++++++++++++ 1 file changed, 57 insertions(+) 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 7959877b62e4..d9132bc697a4 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 @@ -303,6 +303,50 @@ public WriteFiles withNumShards( /** * 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). + * + *

Trade-offs: 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). + * + *

Writer Limit & Overflow Trade-off Matrix (for {@link + * #withRunnerDeterminedSharding()}): + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + *
ConfigurationBehavior when {@code maxNumWritersPerBundle} is reachedWorker Memory ConsumptionShuffle CostOutput File Size / Count
Default (Spill to Shuffle)
{@code maxNumWritersPerBundle > 0},
{@code evictWritersWhenFull = false}		Keeps first {@code N} writers open; spills remaining records to a {@link GroupByKey} shuffle stageBounded ({@code <= N} buffers per bundle)High if many records spill across shuffleFewer, larger files
LRU Writer Eviction
{@code maxNumWritersPerBundle > 0},
{@code evictWritersWhenFull = true}		Flushes and closes the least recently used (LRU) open writer to open a new writer inlineBounded ({@code <= N} buffers per bundle)None (no shuffle stage for unwritten records)May produce more/smaller files (minimal if input is ordered by destination, high if random)
No Spilling
{@link #withNoSpilling()} ({@code maxNumWritersPerBundle = -1})		Opens a new writer for every destination in the bundle without limitUnbounded (risk of OOM with many destinations)None (no shuffle stage for unwritten records)Fewer, larger files (1 file per destination per bundle)
+ * + *

Note that value provided here cannot exceed the default value ({@link + * #DEFAULT_MAX_NUM_WRITERS_PER_BUNDLE}). */ public WriteFiles withMaxNumWritersPerBundle( int maxNumWritersPerBundle) { @@ -321,6 +365,12 @@ public WriteFiles withMaxNumWritersPerBundle( * (LRU order, by flushing and closing it) instead of spilling unwritten records to shuffle when * {@link #getMaxNumWritersPerBundle()} is reached. * + *

Trade-offs: 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. + * *

Warning: 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 @@ -338,6 +388,13 @@ public WriteFiles withEvictWritersWhenFull() { * flushing and closing it) when {@link #getMaxNumWritersPerBundle()} is reached, instead of * spilling unwritten records to shuffle. * + *

Trade-offs: 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. + * *

Warning: 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 From 33a44f0ef915be7ef359ec51161e317aed3ac379 Mon Sep 17 00:00:00 2001 From: Darshan Siddesh Jagaluru Date: Mon, 21 Sep 2026 16:06:24 +0000 Subject: [PATCH 5/5] Corrected a comment --- .../src/main/java/org/apache/beam/sdk/io/WriteFiles.java | 8 +++----- 1 file changed, 3 insertions(+), 5 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 d9132bc697a4..2fd77a1ad0c6 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 @@ -783,15 +783,13 @@ private class WriteUnshardedTempFilesFn extends DoFn inputCoder; - // Initial capacity is maxNumWritersPerBundle (or DEFAULT_MAX_NUM_WRITERS_PER_BUNDLE if -1), - // loadFactor=0.85f prevents rehashing, and accessOrder=true enables LRU eviction order. private final Map, Writer> writers = new LinkedHashMap<>( - getMaxNumWritersPerBundle() < 0 + /* initialCapacity= */ getMaxNumWritersPerBundle() < 0 ? DEFAULT_MAX_NUM_WRITERS_PER_BUNDLE : getMaxNumWritersPerBundle(), - 0.85f, - true); + /* loadFactor= */ 0.85f, + /* accessOrder= */ true); // true = LRU order private final List> evictedFileResults = Lists.newArrayList(); private int spilledShardNum = UNKNOWN_SHARDNUM;