From 7179c37cf895295ed2be4486511a5a1215970475 Mon Sep 17 00:00:00 2001 From: darshan-sj Date: Mon, 21 Sep 2026 20:20:33 +0000 Subject: [PATCH] [FileWrites] Implementing evictWritersWhenFull option in WriteFiles (#40200) --- .../org/apache/beam/sdk/io/WriteFiles.java | 277 ++++++++++++++---- .../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 85829d7e4f00..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 @@ -23,8 +23,11 @@ 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 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 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 static WriteFiles()) @@ -194,6 +200,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 +300,54 @@ 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 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) { checkArgument( @@ -302,6 +360,54 @@ public WriteFiles withMaxNumWritersPerBundle( 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. + * + *

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

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 + * 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 +702,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", @@ -674,8 +783,14 @@ private class WriteUnshardedTempFilesFn extends DoFn inputCoder; - // Initialized in startBundle() - private @Nullable Map, Writer> writers; + private final Map, Writer> writers = + new LinkedHashMap<>( + /* initialCapacity= */ getMaxNumWritersPerBundle() < 0 + ? DEFAULT_MAX_NUM_WRITERS_PER_BUNDLE + : getMaxNumWritersPerBundle(), + /* loadFactor= */ 0.85f, + /* accessOrder= */ true); // true = LRU order + private final List> evictedFileResults = Lists.newArrayList(); private int spilledShardNum = UNKNOWN_SHARDNUM; @@ -691,7 +806,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()) { - 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 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 void processElement( writeOrClose(writer, formattedRecord); } - @FinishBundle - public void finishBundle(FinishBundleContext c) throws Exception { - for (Map.Entry, Writer> entry : - writers.entrySet()) { - WriterKey key = entry.getKey(); - Writer writer = entry.getValue(); + private Writer openAndRegisterWriter( + WriterKey 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 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, 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(); } } } @@ -1202,8 +1378,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..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 @@ -30,11 +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; @@ -44,6 +48,8 @@ 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.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 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 FileBasedSink { + static final AtomicReference EVICTED_TEMP_FILE = new AtomicReference<>(); + static final AtomicBoolean THREW_ON_CLOSE = new AtomicBoolean(false); + + FailingCloseOnEvictSink( + ResourceId tempDirectory, + DynamicDestinations dynamicDestinations) { + super(StaticValueProvider.of(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 Writer(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 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 bundleElements = Arrays.asList("0", "1", "5", "2", "10", "6"); + + 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 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 sink = makeSimpleSink(); @@ -512,6 +681,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