From 0a49326c69c789dbd95417178a2217a96a498616 Mon Sep 17 00:00:00 2001 From: Tejas Iyer Date: Mon, 12 Oct 2026 18:38:28 +0000 Subject: [PATCH 1/6] [WriteFiles] Support direct writes of large streaming bundles to temp files --- CHANGES.md | 1 + .../beam/runners/dataflow/DataflowRunner.java | 21 + .../runners/dataflow/DataflowRunnerTest.java | 94 +++ .../java/org/apache/beam/sdk/io/FileIO.java | 15 + .../java/org/apache/beam/sdk/io/TextIO.java | 21 + .../org/apache/beam/sdk/io/WriteFiles.java | 649 ++++++++++++++---- .../org/apache/beam/sdk/io/FileIOTest.java | 7 + .../apache/beam/sdk/io/TextIOWriteTest.java | 6 + .../apache/beam/sdk/io/WriteFilesTest.java | 335 +++++++++ 9 files changed, 1025 insertions(+), 124 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index cb5eced42f7c..e7dc4f5bcbff 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -73,6 +73,7 @@ * The `disableCounterMetrics`, `disableStringSetMetrics` and `disableBoundedTrieMetrics` experiments are now honored by the Python SDK, as they already were in Java (Python) ([#38746](https://github.com/apache/beam/issues/38746)). * (Go) Pipeline-level `--resource_hints` are now forwarded to expansion services, so they apply to cross-language transforms as they do in the Java and Python SDKs. The `max_active_bundles_per_worker` hint name is now also accepted by `--resource_hints` ([#23893](https://github.com/apache/beam/issues/23893)). * ReadFromBigQuery now supports Lakehouse runtime catalog (BigLake metastore) tables with `method=DIRECT_READ`, using 4-part `project.catalog.namespace.table` identifiers. Previously `project:catalog.namespace.table` was silently mis-parsed and tables that report no `numBytes` failed to split (Python) ([#39597](https://github.com/apache/beam/issues/39597)). +* (Java) `TextIO`, `FileIO`, and `WriteFiles`: streaming writes can skip the shuffle for large bundles by writing them directly to temp files (`withDirectWriteByteThreshold` or experiment `streaming_file_direct_write_threshold_bytes`). Off by default ([#40533](https://github.com/apache/beam/issues/40533)). ## Breaking Changes diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java index f1b0c89090d3..b066b12790ee 100644 --- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java +++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java @@ -180,6 +180,7 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hashing; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.ByteStreams; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.Files; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.DateTimeUtils; import org.joda.time.DateTimeZone; import org.joda.time.format.DateTimeFormat; @@ -2645,6 +2646,7 @@ static class StreamingShardedWriteFactory // We pick 10 as a default, as it works well with the default number of workers started // by Dataflow. static final int DEFAULT_NUM_SHARDS = 10; + DataflowPipelineWorkerPoolOptions options; StreamingShardedWriteFactory(PipelineOptions options) { @@ -2684,6 +2686,11 @@ static class StreamingShardedWriteFactory replacement = replacement.withWindowedWrites(); } + Long directWriteThreshold = getDirectWriteThreshold(transform); + if (directWriteThreshold != null) { + replacement = replacement.withDirectWriteByteThreshold(directWriteThreshold); + } + if (WriteFilesTranslation.isAutoSharded(transform)) { replacement = replacement.withAutoSharding(); return PTransformReplacement.of( @@ -2693,11 +2700,25 @@ static class StreamingShardedWriteFactory return PTransformReplacement.of( PTransformReplacements.getSingletonMainInput(transform), replacement.withNumShards(numShards)); + } catch (RuntimeException e) { + throw e; } catch (Exception e) { throw new RuntimeException(e); } } + /** + * Returns the direct write threshold set on the original {@link WriteFiles}, else the value of + * the {@link WriteFiles#DIRECT_WRITE_THRESHOLD_EXPERIMENT} experiment, else {@code null}. + */ + private @Nullable Long getDirectWriteThreshold(AppliedPTransform transform) { + Object original = transform.getTransform(); + if (original instanceof WriteFiles) { + return ((WriteFiles) original).resolveDirectWriteThreshold(options); + } + return null; + } + @Override public Map, ReplacementOutput> mapOutputs( Map, PCollection> outputs, WriteFilesResult newOutput) { diff --git a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java index b196aa1714c6..a2ffeca62785 100644 --- a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java +++ b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java @@ -2149,9 +2149,103 @@ public void testStreamingWriteWithShardingReturnsSameTransform() { assertNull(replacement.getNumShardsProvider()); assertNull(replacement.getComputeNumShards()); + assertNull(replacement.getDirectWriteByteThreshold()); assertTrue(replacement.getWithAutoSharding()); } + @Test + public void testStreamingWriteDirectWriteThresholdExperiment() { + PipelineOptions options = TestPipeline.testingPipelineOptions(); + ExperimentalOptions.addExperiment( + options.as(ExperimentalOptions.class), "streaming_file_direct_write_threshold_bytes=1000"); + + TestPipeline p = TestPipeline.fromOptions(options); + StreamingShardedWriteFactory factory = + new StreamingShardedWriteFactory<>(p.getOptions()); + WriteFiles original = + WriteFiles.to(new TestSink(tmpFolder.toString())).withRunnerDeterminedSharding(); + PCollection objs = (PCollection) p.apply(Create.empty(VoidCoder.of())); + AppliedPTransform, WriteFilesResult, WriteFiles> + originalApplication = + AppliedPTransform.of( + "writefiles", + PValues.expandInput(objs), + Collections.emptyMap(), + original, + ResourceHints.create(), + p); + + WriteFiles replacement = + (WriteFiles) + factory.getReplacementTransform(originalApplication).getTransform(); + + assertEquals(Long.valueOf(1000L), replacement.getDirectWriteByteThreshold()); + } + + @Test + public void testStreamingWriteDirectWriteThresholdDefaultNull() { + PipelineOptions options = TestPipeline.testingPipelineOptions(); + + TestPipeline p = TestPipeline.fromOptions(options); + StreamingShardedWriteFactory factory = + new StreamingShardedWriteFactory<>(p.getOptions()); + WriteFiles original = + WriteFiles.to(new TestSink(tmpFolder.toString())).withRunnerDeterminedSharding(); + PCollection objs = (PCollection) p.apply(Create.empty(VoidCoder.of())); + AppliedPTransform, WriteFilesResult, WriteFiles> + originalApplication = + AppliedPTransform.of( + "writefiles", + PValues.expandInput(objs), + Collections.emptyMap(), + original, + ResourceHints.create(), + p); + + WriteFiles replacement = + (WriteFiles) + factory.getReplacementTransform(originalApplication).getTransform(); + + assertNull(replacement.getDirectWriteByteThreshold()); + } + + @Test + public void testStreamingWriteDirectWriteThresholdInvalidExperiment() { + for (String badValue : ImmutableList.of("abc", "0", "-1")) { + PipelineOptions options = TestPipeline.testingPipelineOptions(); + ExperimentalOptions.addExperiment( + options.as(ExperimentalOptions.class), + "streaming_file_direct_write_threshold_bytes=" + badValue); + + TestPipeline p = TestPipeline.fromOptions(options); + StreamingShardedWriteFactory factory = + new StreamingShardedWriteFactory<>(p.getOptions()); + WriteFiles original = + WriteFiles.to(new TestSink(tmpFolder.toString())).withRunnerDeterminedSharding(); + PCollection objs = (PCollection) p.apply(Create.empty(VoidCoder.of())); + AppliedPTransform< + PCollection, WriteFilesResult, WriteFiles> + originalApplication = + AppliedPTransform.of( + "writefiles", + PValues.expandInput(objs), + Collections.emptyMap(), + original, + ResourceHints.create(), + p); + + IllegalArgumentException thrown = + assertThrows( + IllegalArgumentException.class, + () -> factory.getReplacementTransform(originalApplication)); + assertThat( + thrown.getMessage(), + containsString( + "Experiment streaming_file_direct_write_threshold_bytes must be a positive integer, got: " + + badValue)); + } + } + private void verifyMergingStatefulParDoRejected(PipelineOptions options) throws Exception { Pipeline p = Pipeline.create(options); 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 3f4289003614..683c167b1a12 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 @@ -1071,6 +1071,8 @@ public static FileNaming relativeFileNaming( abstract @Nullable Duration getBatchMaxBufferingDuration(); + abstract @Nullable Long getDirectWriteByteThreshold(); + abstract @Nullable ErrorHandler getBadRecordErrorHandler(); abstract Builder toBuilder(); @@ -1133,6 +1135,9 @@ abstract Builder setMaxNumWritersPerBundle( abstract Builder setBatchMaxBufferingDuration( @Nullable Duration batchMaxBufferingDuration); + abstract Builder setDirectWriteByteThreshold( + @Nullable Long directWriteByteThreshold); + abstract Builder setBadRecordErrorHandler( @Nullable ErrorHandler badRecordErrorHandler); @@ -1436,6 +1441,13 @@ public Write withBatchMaxBufferingDuration( return toBuilder().setBatchMaxBufferingDuration(batchMaxBufferingDuration).build(); } + /** See {@link WriteFiles#withDirectWriteByteThreshold(long)}. */ + public Write withDirectWriteByteThreshold(long thresholdBytes) { + checkArgument( + thresholdBytes > 0, "thresholdBytes must be positive, but was: %s", thresholdBytes); + return toBuilder().setDirectWriteByteThreshold(thresholdBytes).build(); + } + @VisibleForTesting Contextful> resolveFileNamingFn() { if (getDynamic()) { @@ -1565,6 +1577,9 @@ public WriteFilesResult expand(PCollection input) { if (getBatchMaxBufferingDuration() != null) { writeFiles = writeFiles.withBatchMaxBufferingDuration(getBatchMaxBufferingDuration()); } + if (getDirectWriteByteThreshold() != null) { + writeFiles = writeFiles.withDirectWriteByteThreshold(getDirectWriteByteThreshold()); + } return input.apply(writeFiles); } 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 b96465e5cab7..dc7dd1c5d49f 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 @@ -713,6 +713,9 @@ public abstract static class TypedWrite /** Batch max buffering duration for input records. */ abstract @Nullable Duration getBatchMaxBufferingDuration(); + /** Byte threshold above which a streaming bundle is written directly to temporary files. */ + abstract @Nullable Long getDirectWriteByteThreshold(); + /** Whether to write windowed output files. */ abstract boolean getWindowedWrites(); @@ -780,6 +783,9 @@ abstract Builder setNumShards( abstract Builder setBatchMaxBufferingDuration( @Nullable Duration batchMaxBufferingDuration); + abstract Builder setDirectWriteByteThreshold( + @Nullable Long directWriteByteThreshold); + abstract Builder setWindowedWrites(boolean windowedWrites); abstract Builder setAutoSharding(boolean windowedWrites); @@ -926,6 +932,13 @@ public TypedWrite withBatchMaxBufferingDuration( return toBuilder().setBatchMaxBufferingDuration(batchMaxBufferingDuration).build(); } + /** See {@link WriteFiles#withDirectWriteByteThreshold(long)}. */ + public TypedWrite withDirectWriteByteThreshold(long thresholdBytes) { + checkArgument( + thresholdBytes > 0, "thresholdBytes must be positive, but was: %s", thresholdBytes); + return toBuilder().setDirectWriteByteThreshold(thresholdBytes).build(); + } + /** Set the base directory used to generate temporary files. */ public TypedWrite withTempDirectory( ValueProvider tempDirectory) { @@ -1195,6 +1208,9 @@ public WriteFilesResult expand(PCollection input) { if (getBatchMaxBufferingDuration() != null) { write = write.withBatchMaxBufferingDuration(getBatchMaxBufferingDuration()); } + if (getDirectWriteByteThreshold() != null) { + write = write.withDirectWriteByteThreshold(getDirectWriteByteThreshold()); + } return input.apply("WriteFiles", write); } @@ -1388,6 +1404,11 @@ public Write withBatchMaxBufferingDuration(@Nullable Duration batchMaxBufferingD return new Write(inner.withBatchMaxBufferingDuration(batchMaxBufferingDuration)); } + /** See {@link TypedWrite#withDirectWriteByteThreshold(long)}. */ + public Write withDirectWriteByteThreshold(long thresholdBytes) { + return new Write(inner.withDirectWriteByteThreshold(thresholdBytes)); + } + /** * 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 2fd77a1ad0c6..05128530bd0c 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 @@ -49,6 +49,9 @@ import org.apache.beam.sdk.io.FileBasedSink.WriteOperation; import org.apache.beam.sdk.io.FileBasedSink.Writer; import org.apache.beam.sdk.io.fs.ResourceId; +import org.apache.beam.sdk.metrics.Counter; +import org.apache.beam.sdk.metrics.Metrics; +import org.apache.beam.sdk.options.ExperimentalOptions; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.ValueProvider; import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider; @@ -79,6 +82,7 @@ import org.apache.beam.sdk.transforms.windowing.Window; import org.apache.beam.sdk.util.CoderUtils; import org.apache.beam.sdk.util.MoreFutures; +import org.apache.beam.sdk.util.common.ElementByteSizeObserver; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.PCollection.IsBounded; @@ -158,6 +162,14 @@ public abstract class WriteFiles public static final Duration FILE_TRIGGERING_RECORD_BUFFERING_DURATION = Duration.standardSeconds(5); + /** + * Experiment name for the byte threshold above which a streaming bundle is written directly to + * temporary files. + */ + @Internal + public static final String DIRECT_WRITE_THRESHOLD_EXPERIMENT = + "streaming_file_direct_write_threshold_bytes"; + static final int UNKNOWN_SHARDNUM = -1; static final int DUMMY_SHARDNUM = 0; private @Nullable WriteOperation writeOperation; @@ -210,6 +222,13 @@ public static WriteFiles> getSideInputs(); public abstract @Nullable ShardingFunction getShardingFunction(); @@ -251,6 +270,9 @@ abstract Builder setBatchSizeBytes( abstract Builder setBatchMaxBufferingDuration( @Nullable Duration batchMaxBufferingDuration); + abstract Builder setDirectWriteByteThreshold( + @Nullable Long directWriteByteThreshold); + abstract Builder setSideInputs( List> sideInputs); @@ -445,6 +467,64 @@ public WriteFiles withBatchMaxBufferingDuration( return toBuilder().setBatchMaxBufferingDuration(batchMaxBufferingDuration).build(); } + /** + * Returns a new {@link WriteFiles} that writes large input bundles directly to temporary files, + * skipping the shuffle before temporary file creation. + * + *

Applies only to unbounded inputs, and is ignored when dynamic destinations use side inputs. + * Records of each bundle are buffered until their encoded size reaches {@code thresholdBytes}; + * bundles that reach it are written straight to temporary files. Smaller bundles keep the + * existing grouped path so output files stay large. Shard numbers are assigned at finalization, + * which may override {@link #withNumShards(int)} (the number of files per window can differ from + * the configured number of shards). + * + *

Can also be set via the {@link #DIRECT_WRITE_THRESHOLD_EXPERIMENT} pipeline experiment. + */ + public WriteFiles withDirectWriteByteThreshold( + long thresholdBytes) { + checkArgument(thresholdBytes > 0, "thresholdBytes must be positive, got %s", thresholdBytes); + return toBuilder().setDirectWriteByteThreshold(thresholdBytes).build(); + } + + /** + * Returns the direct write threshold set via {@link #withDirectWriteByteThreshold(long)}, else + * the value of the {@link #DIRECT_WRITE_THRESHOLD_EXPERIMENT} experiment in {@code options} when + * no explicit shard count is configured, else {@code null}. + */ + @Internal + public @Nullable Long resolveDirectWriteThreshold(PipelineOptions options) { + if (getDirectWriteByteThreshold() != null) { + return getDirectWriteByteThreshold(); + } + if (getComputeNumShards() != null || getNumShardsProvider() != null) { + return null; + } + String fromExperiment = + ExperimentalOptions.getExperimentValue(options, DIRECT_WRITE_THRESHOLD_EXPERIMENT); + if (fromExperiment == null) { + return null; + } + final long threshold; + try { + threshold = Long.parseLong(fromExperiment); + } catch (NumberFormatException e) { + throw new IllegalArgumentException( + "Experiment " + + DIRECT_WRITE_THRESHOLD_EXPERIMENT + + " must be a positive integer, got: " + + fromExperiment, + e); + } + if (threshold <= 0) { + throw new IllegalArgumentException( + "Experiment " + + DIRECT_WRITE_THRESHOLD_EXPERIMENT + + " must be a positive integer, got: " + + fromExperiment); + } + return threshold; + } + public WriteFiles withSideInputs( List> sideInputs) { return toBuilder().setSideInputs(sideInputs).build(); @@ -583,6 +663,22 @@ public WriteFilesResult expand(PCollection input) { (getComputeNumShards() == null) ? null : input.apply(getComputeNumShards()); boolean fixedSharding = getComputeNumShards() != null || getNumShardsProvider() != null; + Long directWriteThreshold = resolveDirectWriteThreshold(input.getPipeline().getOptions()); + if (directWriteThreshold != null + && getSideInputs().isEmpty() + && input.isBounded() == IsBounded.UNBOUNDED) { + // Shard numbers are assigned in finalize across both the direct and grouped files, so any + // fixed shard count is not enforced there. + return input + .apply( + "WriteShardedBundlesWithDirectWrite", + new WriteShardedBundlesWithDirectWrite( + destinationCoder, fileResultCoder, numShardsView, directWriteThreshold)) + .apply( + "FinalizeTempFileBundles", + new FinalizeTempFileBundles( + null, destinationCoder, /* assignShardNumbersInFinalize= */ true)); + } PCollection>> tempFileResults; if (fixedSharding) { tempFileResults = @@ -614,7 +710,9 @@ public WriteFilesResult expand(PCollection input) { } return tempFileResults.apply( - "FinalizeTempFileBundles", new FinalizeTempFileBundles(numShardsView, destinationCoder)); + "FinalizeTempFileBundles", + new FinalizeTempFileBundles( + numShardsView, destinationCoder, /* assignShardNumbersInFinalize= */ false)); } @Override @@ -817,7 +915,6 @@ public void processElement( ProcessContext c, BoundedWindow window, MultiOutputReceiver outputReceiver) throws Exception { getDynamicDestinations().setSideInputAccessorFromProcessContext(c); - PaneInfo paneInfo = c.pane(); // If we are doing windowed writes, we need to ensure that we have separate files for // data in different windows/panes. Similar for dynamic writes, make sure that different // destinations go to different writers. @@ -835,10 +932,10 @@ public void processElement( writer = writers.get(key); if (writer == null) { if (getMaxNumWritersPerBundle() < 0 || writers.size() <= getMaxNumWritersPerBundle()) { - writer = openAndRegisterWriter(key, window, paneInfo, destination); + writer = openAndRegisterWriter(writers, key); } else if (getEvictWritersWhenFull()) { evictOldestWriter(); - writer = openAndRegisterWriter(key, window, paneInfo, destination); + writer = openAndRegisterWriter(writers, key); } else { if (spilledShardNum == UNKNOWN_SHARDNUM) { // Cache the random value so we only call ThreadLocalRandom once per DoFn instance. @@ -865,29 +962,6 @@ public void processElement( writeOrClose(writer, formattedRecord); } - 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 = @@ -927,30 +1001,56 @@ public void finishBundle(FinishBundleContext c) throws Exception { 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); - } + closeAllWritersAndOutputFileResults(writers, c); writers.clear(); evictedFileResults.clear(); } } } + private Writer openAndRegisterWriter( + Map, Writer> writers, + WriterKey key) + throws Exception { + String uuid = UUID.randomUUID().toString(); + LOG.info( + "Opening writer {} for window {} pane {} destination {}", + uuid, + key.window, + key.paneInfo, + key.destination); + Writer writer = writeOperation.createWriter(); + writer.setDestination(key.destination); + writer.open(uuid); + writers.put(key, writer); + LOG.debug("Done opening writer"); + return writer; + } + + private static void closeAllWritersAndOutputFileResults( + Map, Writer> writers, + DoFn>.FinishBundleContext c) + throws Exception { + 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); + } + } + private static void writeOrClose( Writer writer, OutputT t) throws Exception { try { @@ -1118,6 +1218,370 @@ public PCollection> expand(PCollection input) { } } + /** + * Streaming write that sends large bundles straight to temporary files. + * + *

Small bundles, and records that would exceed the per-bundle writer limit, take the usual + * grouped path ({@link WriteShardedBundlesToTempFiles} when a shard count is set, or the {@link + * GroupIntoBatches} auto-sharded path otherwise). All files then drop their shard number and are + * grouped per destination, so finalization assigns shard numbers across both paths. + */ + private class WriteShardedBundlesWithDirectWrite + extends PTransform, PCollection>>> { + private final Coder destinationCoder; + private final Coder> fileResultCoder; + private final @Nullable PCollectionView numShardsView; + private final long thresholdBytes; + + private WriteShardedBundlesWithDirectWrite( + Coder destinationCoder, + Coder> fileResultCoder, + @Nullable PCollectionView numShardsView, + long thresholdBytes) { + this.destinationCoder = destinationCoder; + this.fileResultCoder = fileResultCoder; + this.numShardsView = numShardsView; + this.thresholdBytes = thresholdBytes; + } + + @Override + public PCollection>> expand(PCollection input) { + TupleTag> directFilesTag = new TupleTag<>("directFiles"); + TupleTag smallBundlesTag = new TupleTag<>("smallBundles"); + + PCollectionTuple routed = + input.apply( + "DirectWriteLargeBundles", + ParDo.of( + new DirectWriteLargeBundlesFn( + smallBundlesTag, thresholdBytes, input.getCoder())) + .withSideInputs(getSideInputs()) + .withOutputTags( + directFilesTag, + TupleTagList.of(ImmutableList.of(smallBundlesTag, BAD_RECORD_TAG)))); + addErrorCollection(routed); + + PCollection smallBundles = routed.get(smallBundlesTag).setCoder(input.getCoder()); + PCollection> groupedFiles; + if (getComputeNumShards() != null || getNumShardsProvider() != null) { + groupedFiles = + smallBundles + .apply( + "WriteShardedBundlesToTempFiles", + new WriteShardedBundlesToTempFiles( + destinationCoder, fileResultCoder, numShardsView)) + .apply( + "DropShardNum", + ParDo.of( + new DoFn, FileResult>() { + @ProcessElement + public void process(ProcessContext c) { + c.output(c.element().withShard(UNKNOWN_SHARDNUM)); + } + })) + .setCoder(fileResultCoder); + } else { + groupedFiles = + writeAutoShardedRecordsToTempFiles( + smallBundles, + destinationCoder, + fileResultCoder, + getBatchSize() != null ? getBatchSize() : FILE_TRIGGERING_RECORD_COUNT, + getBatchSizeBytes() != null ? getBatchSizeBytes() : FILE_TRIGGERING_BYTE_COUNT, + getBatchMaxBufferingDuration() != null + ? getBatchMaxBufferingDuration() + : FILE_TRIGGERING_RECORD_BUFFERING_DURATION); + } + + // When the small path goes through GroupIntoShards (GBK), its trigger transitions to the + // continuation trigger. Set the direct path to the same windowing strategy so + // FlattenTempFiles does not reject mismatched triggers when user pipelines use non-default + // triggers. + PCollection> directFiles = + routed + .get(directFilesTag) + .setCoder(fileResultCoder) + .setWindowingStrategyInternal(groupedFiles.getWindowingStrategy()); + + // Group all temp files per destination so each window/pane firing is finalized together + // and shard numbers are assigned once across both paths. + return PCollectionList.of(directFiles) + .and(groupedFiles) + .apply("FlattenTempFiles", Flatten.pCollections()) + .setCoder(fileResultCoder) + .apply( + "KeyedByDestination", + WithKeys.of( + new SimpleFunction, DestinationT>() { + @Override + public DestinationT apply(FileResult input) { + return input.getDestination(); + } + })) + .setCoder(KvCoder.of(destinationCoder, fileResultCoder)) + .apply("GroupTempFilesByDestination", GroupByKey.create()) + .apply( + "ExtractValuesToList", + ParDo.of( + new DoFn< + KV>>, + List>>() { + @ProcessElement + public void processElement(ProcessContext c) { + // Stamp every file with this firing's pane so files from different + // firings of the same window get distinct final names. + List> result = new ArrayList<>(); + for (FileResult e : c.element().getValue()) { + result.add( + new FileResult<>( + e.getTempFilename(), + UNKNOWN_SHARDNUM, + e.getWindow(), + c.pane(), + e.getDestination())); + } + c.output(result); + } + })) + .setCoder(ListCoder.of(fileResultCoder)); + } + } + + private static class ByteSizeObserver extends ElementByteSizeObserver { + private long byteSize; + + @Override + protected void reportElementSize(long elementByteSize) { + this.byteSize = elementByteSize; + } + } + + /** + * Buffers each bundle's records until their encoded size reaches {@code thresholdBytes}. Bundles + * that reach it are written directly to temporary files. Records of smaller bundles, and records + * that would exceed the per-bundle writer limit, are emitted to {@code smallBundlesTag} at the + * end of the bundle for the sharded {@link GroupByKey} path. + * + *

Adapted from the IcebergIO bundle lifter (apache/beam#36720), but buffers at most {@code + * thresholdBytes} rather than the whole bundle. + */ + private class DirectWriteLargeBundlesFn extends DoFn> { + private final TupleTag smallBundlesTag; + private final long thresholdBytes; + private final Coder inputCoder; + + private final Counter directBundles = Metrics.counter(WriteFiles.class, "directWriteBundles"); + private final Counter groupedBundles = Metrics.counter(WriteFiles.class, "groupedWriteBundles"); + + // Per-bundle state, reset in startBundle. + private transient List> buffer; + private transient long bufferedBytes; + private transient boolean directMode; + private transient Map, Writer> writers; + private transient List> spilled; + + DirectWriteLargeBundlesFn( + TupleTag smallBundlesTag, long thresholdBytes, Coder inputCoder) { + this.smallBundlesTag = smallBundlesTag; + this.thresholdBytes = thresholdBytes; + this.inputCoder = inputCoder; + } + + @StartBundle + public void startBundle() { + buffer = new ArrayList<>(); + bufferedBytes = 0; + directMode = false; + writers = new LinkedHashMap<>(); + spilled = new ArrayList<>(); + } + + @ProcessElement + public void processElement( + ProcessContext c, BoundedWindow window, MultiOutputReceiver outputReceiver) + throws Exception { + getDynamicDestinations().setSideInputAccessorFromProcessContext(c); + BufferedRecord record = + new BufferedRecord<>(c.element(), c.timestamp(), window, c.pane()); + if (directMode) { + write(record, outputReceiver); + return; + } + buffer.add(record); + bufferedBytes += estimateSize(c.element()); + if (bufferedBytes >= thresholdBytes) { + directMode = true; + // Direct writes are disabled when destinations use side inputs (see + // expand()), so resolving buffered records' destinations here is safe. + for (BufferedRecord buffered : buffer) { + write(buffered, outputReceiver); + } + buffer.clear(); + } + } + + private long estimateSize(UserT value) throws Exception { + if (inputCoder.isRegisterByteSizeObserverCheap(value)) { + ByteSizeObserver observer = new ByteSizeObserver(); + inputCoder.registerByteSizeObserver(value, observer); + observer.advance(); + return observer.byteSize; + } + return CoderUtils.encodeToByteArray(inputCoder, value).length; + } + + private void write(BufferedRecord record, MultiOutputReceiver outputReceiver) + throws Exception { + MaybeDestination maybeDestination = + getDestinationWithErrorHandling(record.value, outputReceiver, inputCoder); + if (!maybeDestination.isValid) { + return; + } + DestinationT destination = maybeDestination.destination; + WriterKey key = new WriterKey<>(record.window, record.pane, destination); + Writer writer = writers.get(key); + if (writer == null) { + int maxWriters = + getMaxNumWritersPerBundle() < 0 + ? DEFAULT_MAX_NUM_WRITERS_PER_BUNDLE + : getMaxNumWritersPerBundle(); + if (writers.size() >= maxWriters) { + // Too many open writers in this bundle; send the record to the grouped path. + spilled.add(record); + return; + } + writer = openAndRegisterWriter(writers, key); + } + OutputT formattedRecord = + formatRecordWithErrorHandling(record.value, outputReceiver, inputCoder); + if (formattedRecord == null) { + return; + } + writeOrClose(writer, formattedRecord); + } + + @FinishBundle + public void finishBundle(FinishBundleContext c) throws Exception { + if (directMode) { + directBundles.inc(); + } else if (!buffer.isEmpty()) { + groupedBundles.inc(); + } + closeAllWritersAndOutputFileResults(writers, c); + for (BufferedRecord record : buffer) { + c.output(smallBundlesTag, record.value, record.timestamp, record.window); + } + for (BufferedRecord record : spilled) { + c.output(smallBundlesTag, record.value, record.timestamp, record.window); + } + buffer.clear(); + spilled.clear(); + writers.clear(); + bufferedBytes = 0; + directMode = false; + } + } + + /** A record held by {@link DirectWriteLargeBundlesFn} with the metadata needed to re-emit it. */ + private static class BufferedRecord { + private final T value; + private final Instant timestamp; + private final BoundedWindow window; + private final PaneInfo pane; + + BufferedRecord(T value, Instant timestamp, BoundedWindow window, PaneInfo pane) { + this.value = value; + this.timestamp = timestamp; + this.window = window; + this.pane = pane; + } + } + + private PCollection> writeAutoShardedRecordsToTempFiles( + PCollection input, + Coder destinationCoder, + Coder> fileResultCoder, + int batchSize, + int batchSizeBytes, + Duration maxBufferingDuration) { + // Auto-sharding is achieved via GroupIntoBatches.WithShardedKey which shards, groups and at + // the same time batches the input records. The sharding behavior depends on runners. The + // batching is per window and we also emit the batches if there are a certain number of + // records buffered or they have been buffered for a certain time, controlled by + // FILE_TRIGGERING_RECORD_COUNT and BUFFERING_DURATION respectively. + // + // TODO(https://github.com/apache/beam/issues/20928): The implementation doesn't currently + // work with merging windows. + TupleTag> shardTag = new TupleTag<>("shardTag"); + + PCollectionTuple shardedElements = + input.apply( + "KeyedByDestinationHash", + ParDo.of(new KeyByDestinationHash(input.getCoder(), destinationCoder)) + .withOutputTags(shardTag, TupleTagList.of(BAD_RECORD_TAG))); + addErrorCollection(shardedElements); + + PCollection, Iterable>> shardedInput = + shardedElements + .get(shardTag) + .setCoder(KvCoder.of(VarIntCoder.of(), input.getCoder())) + .apply( + "ShardAndBatch", + GroupIntoBatches.ofSize(batchSize) + .withByteSize(batchSizeBytes) + .withMaxBufferingDuration(maxBufferingDuration) + .withShardedKey()) + .setCoder( + KvCoder.of( + org.apache.beam.sdk.util.ShardedKey.Coder.of(VarIntCoder.of()), + IterableCoder.of(input.getCoder()))); + + TupleTag> writtenRecordsTag = new TupleTag<>("writtenRecords"); + // Write grouped elements to temp files. + PCollectionTuple writtenFiles = + shardedInput + .apply( + "AddDummyShard", + MapElements.via( + new SimpleFunction< + KV, Iterable>, + KV, Iterable>>() { + @Override + public KV, Iterable> apply( + KV, Iterable> input) { + // Add dummy shard since it is required by WriteShardsIntoTempFilesFn. It + // will be dropped after we generate the temp files. + return KV.of( + ShardedKey.of(input.getKey().getKey(), DUMMY_SHARDNUM), + input.getValue()); + } + })) + .setCoder( + KvCoder.of( + ShardedKeyCoder.of(VarIntCoder.of()), IterableCoder.of(input.getCoder()))) + .apply( + "WriteShardsIntoTempFiles", + ParDo.of(new WriteShardsIntoTempFilesFn(input.getCoder())) + .withSideInputs(getSideInputs()) + .withOutputTags(writtenRecordsTag, TupleTagList.of(BAD_RECORD_TAG))); + + addErrorCollection(writtenFiles); + + return writtenFiles + .get(writtenRecordsTag) + .setCoder(fileResultCoder) + .apply( + "DropShardNum", + ParDo.of( + new DoFn, FileResult>() { + @ProcessElement + public void process(ProcessContext c) { + c.output(c.element().withShard(UNKNOWN_SHARDNUM)); + } + })) + .setCoder(fileResultCoder); + } + private class WriteAutoShardedBundlesToTempFiles extends PTransform, PCollection>>> { private final Coder destinationCoder; @@ -1144,84 +1608,14 @@ private WriteAutoShardedBundlesToTempFiles( @Override public PCollection>> expand(PCollection input) { - // Auto-sharding is achieved via GroupIntoBatches.WithShardedKey which shards, groups and at - // the same time batches the input records. The sharding behavior depends on runners. The - // batching is per window and we also emit the batches if there are a certain number of - // records buffered or they have been buffered for a certain time, controlled by - // FILE_TRIGGERING_RECORD_COUNT and BUFFERING_DURATION respectively. - // - // TODO(https://github.com/apache/beam/issues/20928): The implementation doesn't currently - // work with merging windows. - TupleTag> shardTag = new TupleTag<>("shardTag"); - - PCollectionTuple shardedElements = - input.apply( - "KeyedByDestinationHash", - ParDo.of(new KeyByDestinationHash(input.getCoder(), destinationCoder)) - .withOutputTags(shardTag, TupleTagList.of(BAD_RECORD_TAG))); - addErrorCollection(shardedElements); - - PCollection, Iterable>> shardedInput = - shardedElements - .get(shardTag) - .setCoder(KvCoder.of(VarIntCoder.of(), input.getCoder())) - .apply( - "ShardAndBatch", - GroupIntoBatches.ofSize(batchSize) - .withByteSize(batchSizeBytes) - .withMaxBufferingDuration(maxBufferingDuration) - .withShardedKey()) - .setCoder( - KvCoder.of( - org.apache.beam.sdk.util.ShardedKey.Coder.of(VarIntCoder.of()), - IterableCoder.of(input.getCoder()))); - - TupleTag> writtenRecordsTag = new TupleTag<>("writtenRecords"); - // Write grouped elements to temp files. - PCollectionTuple writtenFiles = - shardedInput - .apply( - "AddDummyShard", - MapElements.via( - new SimpleFunction< - KV, Iterable>, - KV, Iterable>>() { - @Override - public KV, Iterable> apply( - KV, Iterable> - input) { - // Add dummy shard since it is required by WriteShardsIntoTempFilesFn. It - // will be dropped after we generate the temp files. - return KV.of( - ShardedKey.of(input.getKey().getKey(), DUMMY_SHARDNUM), - input.getValue()); - } - })) - .setCoder( - KvCoder.of( - ShardedKeyCoder.of(VarIntCoder.of()), IterableCoder.of(input.getCoder()))) - .apply( - "WriteShardsIntoTempFiles", - ParDo.of(new WriteShardsIntoTempFilesFn(input.getCoder())) - .withSideInputs(getSideInputs()) - .withOutputTags(writtenRecordsTag, TupleTagList.of(BAD_RECORD_TAG))); - - addErrorCollection(writtenFiles); - PCollection> tempFiles = - writtenFiles - .get(writtenRecordsTag) - .setCoder(fileResultCoder) - .apply( - "DropShardNum", - ParDo.of( - new DoFn, FileResult>() { - @ProcessElement - public void process(ProcessContext c) { - c.output(c.element().withShard(UNKNOWN_SHARDNUM)); - } - })) - .setCoder(fileResultCoder); + writeAutoShardedRecordsToTempFiles( + input, + destinationCoder, + fileResultCoder, + batchSize, + batchSizeBytes, + maxBufferingDuration); // Group temp file results by destinations again to gather all the results in the same window. // This is needed since we don't have shard idx associated with each temp file so have to rely @@ -1484,11 +1878,16 @@ private class FinalizeTempFileBundles PCollection>>, WriteFilesResult> { private final @Nullable PCollectionView numShardsView; private final Coder destinationCoder; + // When true, ignore the fixed shard count and let finalize number the files 0..K-1. + private final boolean assignShardNumbersInFinalize; private FinalizeTempFileBundles( - @Nullable PCollectionView numShardsView, Coder destinationCoder) { + @Nullable PCollectionView numShardsView, + Coder destinationCoder, + boolean assignShardNumbersInFinalize) { this.numShardsView = numShardsView; this.destinationCoder = destinationCoder; + this.assignShardNumbersInFinalize = assignShardNumbersInFinalize; } @Override @@ -1519,7 +1918,9 @@ private class FinalizeFn public void process(ProcessContext c) throws Exception { getDynamicDestinations().setSideInputAccessorFromProcessContext(c); @Nullable Integer fixedNumShards; - if (numShardsView != null) { + if (assignShardNumbersInFinalize) { + fixedNumShards = null; + } else if (numShardsView != null) { fixedNumShards = c.sideInput(numShardsView); } else if (getNumShardsProvider() != null) { fixedNumShards = getNumShardsProvider().get(); diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/FileIOTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/FileIOTest.java index 7a0dfb4822b3..eab7195d0918 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/FileIOTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/FileIOTest.java @@ -723,6 +723,13 @@ public void testWriteUnboundedWithCustomBatchSizeBytes() throws IOException { assertTrue(metadata.size() >= 2); } + @Test + public void testWriteWithDirectWriteByteThreshold() { + FileIO.Write write = + FileIO.write().to("/tmp/x").via(TextIO.sink()).withDirectWriteByteThreshold(1000); + assertEquals(Long.valueOf(1000), write.getDirectWriteByteThreshold()); + } + static void checkFileContents(File rootDir, String prefix, List inputs) throws IOException { List outputFiles = Lists.newArrayList(); diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/TextIOWriteTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/TextIOWriteTest.java index f58906e464be..50b1672c6453 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/TextIOWriteTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/TextIOWriteTest.java @@ -782,4 +782,10 @@ public void testSink() throws Exception { assertEquals(Arrays.asList("header", "a", "b", "c", "footer"), readLinesFromFile(f)); } + + @Test + public void testWriteWithDirectWriteByteThreshold() { + TextIO.Write write = TextIO.write().to("/tmp/x").withDirectWriteByteThreshold(1000); + assertEquals(Long.valueOf(1000), write.inner.getDirectWriteByteThreshold()); + } } 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 09f3f04d5b09..35ddb839dfdf 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 @@ -54,6 +54,7 @@ import java.util.regex.Matcher; import java.util.regex.Pattern; import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.io.DefaultFilenamePolicy.Params; import org.apache.beam.sdk.io.FileBasedSink.DynamicDestinations; @@ -64,7 +65,12 @@ import org.apache.beam.sdk.io.fs.MatchResult.Metadata; import org.apache.beam.sdk.io.fs.ResolveOptions.StandardResolveOptions; import org.apache.beam.sdk.io.fs.ResourceId; +import org.apache.beam.sdk.metrics.MetricNameFilter; +import org.apache.beam.sdk.metrics.MetricQueryResults; +import org.apache.beam.sdk.metrics.MetricResult; +import org.apache.beam.sdk.metrics.MetricsFilter; import org.apache.beam.sdk.options.Description; +import org.apache.beam.sdk.options.ExperimentalOptions; import org.apache.beam.sdk.options.PipelineOptionsFactoryTest.TestPipelineOptions; import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider; import org.apache.beam.sdk.testing.BeamParallelJunit4Runner; @@ -89,10 +95,12 @@ import org.apache.beam.sdk.transforms.display.DisplayData; import org.apache.beam.sdk.transforms.errorhandling.ErrorHandler.BadRecordErrorHandler; import org.apache.beam.sdk.transforms.errorhandling.ErrorHandlingTestUtils.ErrorSinkTransform; +import org.apache.beam.sdk.transforms.windowing.AfterPane; import org.apache.beam.sdk.transforms.windowing.BoundedWindow; import org.apache.beam.sdk.transforms.windowing.FixedWindows; import org.apache.beam.sdk.transforms.windowing.IntervalWindow; import org.apache.beam.sdk.transforms.windowing.PaneInfo; +import org.apache.beam.sdk.transforms.windowing.Repeatedly; import org.apache.beam.sdk.transforms.windowing.Sessions; import org.apache.beam.sdk.transforms.windowing.Window; import org.apache.beam.sdk.util.MimeTypes; @@ -106,6 +114,7 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; import org.apache.commons.compress.utils.Sets; +import org.checkerframework.checker.nullness.qual.Nullable; import org.hamcrest.Matchers; import org.joda.time.Duration; import org.joda.time.Instant; @@ -422,6 +431,288 @@ public void testWithRunnerDeterminedShardingTestStream() throws IOException { getBaseOutputFilename(), elements, Optional.absent(), !write.getWindowedWrites(), null); } + @Test + @Category({NeedsRunner.class, UsesUnboundedPCollections.class, UsesTestStream.class}) + public void testDirectWriteAllBundlesLarge() throws IOException { + runDirectWriteTestStream(1L); + } + + @Test + @Category({NeedsRunner.class, UsesUnboundedPCollections.class, UsesTestStream.class}) + public void testDirectWriteAllBundlesSmall() throws IOException { + runDirectWriteTestStream(Long.MAX_VALUE); + } + + @Test + @Category({NeedsRunner.class, UsesUnboundedPCollections.class, UsesTestStream.class}) + public void testDirectWriteUnshardedViaExperiment() throws IOException { + List elements = Lists.newArrayList(); + for (int i = 0; i < 30; ++i) { + elements.add("number: " + i); + } + Instant startInstant = new Instant(0L); + // Two 10-element bundles (~100 bytes each), one 8-element bundle (~88 bytes), and one + // 2-element bundle (~22 bytes, < 50) so both direct and auto-sharded paths run. + TestStream testStream = + TestStream.create(StringUtf8Coder.of()) + .advanceWatermarkTo(startInstant) + .addElements(elements.get(0), Iterables.toArray(elements.subList(1, 10), String.class)) + .advanceWatermarkTo(startInstant.plus(Duration.standardSeconds(5))) + .addElements( + elements.get(10), Iterables.toArray(elements.subList(11, 20), String.class)) + .advanceWatermarkTo(startInstant.plus(Duration.standardSeconds(10))) + .addElements( + elements.get(20), Iterables.toArray(elements.subList(21, 28), String.class)) + .addElements(elements.get(28), elements.get(29)) + .advanceWatermarkToInfinity(); + + ExperimentalOptions.addExperiment( + p.getOptions().as(ExperimentalOptions.class), + WriteFiles.DIRECT_WRITE_THRESHOLD_EXPERIMENT + "=50"); + + WriteFiles write = + WriteFiles.to(makeSimpleSink()).withWindowedWrites().withRunnerDeterminedSharding(); + p.apply(testStream) + .apply(Window.into(FixedWindows.of(Duration.standardSeconds(10)))) + .apply(write) + .getPerDestinationOutputFilenames() + .apply(new VerifyFilesExist<>()); + PipelineResult result = p.run(); + + checkFileContents(getBaseOutputFilename(), elements, Optional.absent(), false, null); + assertThat(counterValue(result, "directWriteBundles"), Matchers.greaterThan(0L)); + assertThat(counterValue(result, "groupedWriteBundles"), Matchers.greaterThan(0L)); + } + + @Test + @Category({NeedsRunner.class, UsesUnboundedPCollections.class, UsesTestStream.class}) + public void testDirectWriteExperimentIgnoredWithExplicitNumShards() throws IOException { + List elements = Lists.newArrayList(); + for (int i = 0; i < 30; ++i) { + elements.add("number: " + i); + } + Instant startInstant = new Instant(0L); + TestStream testStream = + TestStream.create(StringUtf8Coder.of()) + .advanceWatermarkTo(startInstant) + .addElements(elements.get(0), Iterables.toArray(elements.subList(1, 10), String.class)) + .advanceWatermarkTo(startInstant.plus(Duration.standardSeconds(5))) + .addElements( + elements.get(10), Iterables.toArray(elements.subList(11, 20), String.class)) + .advanceWatermarkTo(startInstant.plus(Duration.standardSeconds(10))) + .addElements( + elements.get(20), Iterables.toArray(elements.subList(21, 30), String.class)) + .advanceWatermarkToInfinity(); + + ExperimentalOptions.addExperiment( + p.getOptions().as(ExperimentalOptions.class), + WriteFiles.DIRECT_WRITE_THRESHOLD_EXPERIMENT + "=1"); + + WriteFiles write = + WriteFiles.to(makeSimpleSink()).withWindowedWrites().withNumShards(3); + p.apply(testStream) + .apply(Window.into(FixedWindows.of(Duration.standardSeconds(10)))) + .apply(write) + .getPerDestinationOutputFilenames() + .apply(new VerifyFilesExist<>()); + PipelineResult result = p.run(); + + checkFileContents(getBaseOutputFilename(), elements, Optional.absent(), false, null); + assertEquals(0L, counterValue(result, "directWriteBundles")); + } + + @Test + @Category({NeedsRunner.class, UsesUnboundedPCollections.class, UsesTestStream.class}) + public void testDirectWriteNonDefaultTriggerLargeBundles() throws IOException { + runDirectWriteNonDefaultTrigger(1L); + } + + @Test + @Category({NeedsRunner.class, UsesUnboundedPCollections.class, UsesTestStream.class}) + public void testDirectWriteNonDefaultTriggerSmallBundles() throws IOException { + runDirectWriteNonDefaultTrigger(Long.MAX_VALUE); + } + + private void runDirectWriteNonDefaultTrigger(long thresholdBytes) throws IOException { + List elements = Lists.newArrayList(); + for (int i = 0; i < 30; ++i) { + elements.add("number: " + i); + } + Instant startInstant = new Instant(0L); + TestStream testStream = + TestStream.create(StringUtf8Coder.of()) + .advanceWatermarkTo(startInstant) + .addElements(elements.get(0), Iterables.toArray(elements.subList(1, 10), String.class)) + .advanceWatermarkTo(startInstant.plus(Duration.standardSeconds(5))) + .addElements( + elements.get(10), Iterables.toArray(elements.subList(11, 20), String.class)) + .advanceWatermarkTo(startInstant.plus(Duration.standardSeconds(10))) + .addElements( + elements.get(20), Iterables.toArray(elements.subList(21, 30), String.class)) + .advanceWatermarkToInfinity(); + + FilenamePolicy paneAwarePolicy = + new PaneAwarePerWindowFiles( + getBaseOutputDirectory().resolve("file", StandardResolveOptions.RESOLVE_FILE), + "simple"); + + WriteFiles write = + WriteFiles.to(SimpleSink.makeSimpleSink(getBaseOutputDirectory(), paneAwarePolicy)) + .withWindowedWrites() + .withNumShards(3) + .withDirectWriteByteThreshold(thresholdBytes); + p.apply(testStream) + .apply( + Window.into(FixedWindows.of(Duration.standardSeconds(10))) + .triggering(Repeatedly.forever(AfterPane.elementCountAtLeast(5))) + .discardingFiredPanes() + .withAllowedLateness(Duration.standardMinutes(1))) + .apply(write) + .getPerDestinationOutputFilenames() + .apply(new VerifyFilesExist<>()); + PipelineResult result = p.run(); + + checkFileContents(getBaseOutputFilename(), elements, Optional.absent(), false, null); + assertDirectWriteMetrics(result, thresholdBytes); + } + + private static class SideInputDestinations extends DynamicDestinations { + private final ResourceId baseOutputDirectory; + private final PCollectionView sideInputView; + + SideInputDestinations(ResourceId baseOutputDirectory, PCollectionView sideInputView) { + this.baseOutputDirectory = baseOutputDirectory; + this.sideInputView = sideInputView; + } + + @Override + public List> getSideInputs() { + return Collections.singletonList(sideInputView); + } + + @Override + public String formatRecord(String record) { + return record; + } + + @Override + public String getDestination(String element) { + return sideInput(sideInputView); + } + + @Override + public String getDefaultDestination() { + return "file"; + } + + @Override + public FilenamePolicy getFilenamePolicy(String destination) { + return new PerWindowFiles( + baseOutputDirectory.resolve(destination, StandardResolveOptions.RESOLVE_FILE), "simple"); + } + } + + @Test + @Category({NeedsRunner.class, UsesUnboundedPCollections.class, UsesTestStream.class}) + public void testDirectWriteDisabledWithSideInputs() throws IOException { + List elements = Lists.newArrayList(); + for (int i = 0; i < 30; ++i) { + elements.add("number: " + i); + } + Instant startInstant = new Instant(0L); + TestStream testStream = + TestStream.create(StringUtf8Coder.of()) + .advanceWatermarkTo(startInstant) + .addElements(elements.get(0), Iterables.toArray(elements.subList(1, 10), String.class)) + .advanceWatermarkTo(startInstant.plus(Duration.standardSeconds(5))) + .addElements( + elements.get(10), Iterables.toArray(elements.subList(11, 20), String.class)) + .advanceWatermarkTo(startInstant.plus(Duration.standardSeconds(10))) + .addElements( + elements.get(20), Iterables.toArray(elements.subList(21, 30), String.class)) + .advanceWatermarkToInfinity(); + + PCollectionView sideInputView = + p.apply("SideInput", Create.of("file")).apply(View.asSingleton()); + SideInputDestinations destinations = + new SideInputDestinations(getBaseOutputDirectory(), sideInputView); + SimpleSink sink = + new SimpleSink<>(getBaseOutputDirectory(), destinations, Compression.UNCOMPRESSED); + + WriteFiles write = + WriteFiles.to(sink).withWindowedWrites().withNumShards(3).withDirectWriteByteThreshold(1L); + p.apply(testStream) + .apply(Window.into(FixedWindows.of(Duration.standardSeconds(10)))) + .apply(write) + .getPerDestinationOutputFilenames() + .apply(new VerifyFilesExist<>()); + PipelineResult result = p.run(); + + checkFileContents(getBaseOutputFilename(), elements, Optional.absent(), false, null); + assertEquals(0L, counterValue(result, "directWriteBundles")); + } + + private void runDirectWriteTestStream(long thresholdBytes) throws IOException { + List elements = Lists.newArrayList(); + for (int i = 0; i < 30; ++i) { + elements.add("number: " + i); + } + Instant startInstant = new Instant(0L); + TestStream testStream = + TestStream.create(StringUtf8Coder.of()) + .advanceWatermarkTo(startInstant) + .addElements(elements.get(0), Iterables.toArray(elements.subList(1, 10), String.class)) + .advanceWatermarkTo(startInstant.plus(Duration.standardSeconds(5))) + .addElements( + elements.get(10), Iterables.toArray(elements.subList(11, 20), String.class)) + .advanceWatermarkTo(startInstant.plus(Duration.standardSeconds(10))) + .addElements( + elements.get(20), Iterables.toArray(elements.subList(21, 30), String.class)) + .advanceWatermarkToInfinity(); + + WriteFiles write = + WriteFiles.to(makeSimpleSink()) + .withWindowedWrites() + .withNumShards(3) + .withDirectWriteByteThreshold(thresholdBytes); + p.apply(testStream) + .apply(Window.into(FixedWindows.of(Duration.standardSeconds(10)))) + .apply(write) + .getPerDestinationOutputFilenames() + .apply(new VerifyFilesExist<>()); + PipelineResult result = p.run(); + + // Shard count is not fixed on the direct-write path, so don't check it. + checkFileContents(getBaseOutputFilename(), elements, Optional.absent(), false, null); + assertDirectWriteMetrics(result, thresholdBytes); + } + + private static void assertDirectWriteMetrics(PipelineResult result, long thresholdBytes) { + long directBundles = counterValue(result, "directWriteBundles"); + long groupedBundles = counterValue(result, "groupedWriteBundles"); + if (thresholdBytes == 1L) { + assertThat(directBundles, Matchers.greaterThan(0L)); + } else if (thresholdBytes == Long.MAX_VALUE) { + assertEquals(0L, directBundles); + assertThat(groupedBundles, Matchers.greaterThan(0L)); + } + } + + private static long counterValue(PipelineResult result, String name) { + MetricQueryResults metrics = + result + .metrics() + .queryMetrics( + MetricsFilter.builder() + .addNameFilter(MetricNameFilter.named(WriteFiles.class, name)) + .build()); + long total = 0; + for (MetricResult counter : metrics.getCounters()) { + total += counter.getAttempted(); + } + return total; + } + /** Test a WriteFiles transform with an empty PCollection. */ @Test @Category(NeedsRunner.class) @@ -1109,6 +1400,50 @@ public ResourceId unwindowedFilename( } } + private static class PaneAwarePerWindowFiles extends FilenamePolicy { + private static final DateTimeFormatter FORMATTER = ISODateTimeFormat.hourMinuteSecondMillis(); + private final ResourceId baseFilename; + private final String suffix; + + public PaneAwarePerWindowFiles(ResourceId baseFilename, String suffix) { + this.baseFilename = baseFilename; + this.suffix = suffix; + } + + @Override + public ResourceId windowedFilename( + int shardNumber, + int numShards, + BoundedWindow window, + PaneInfo paneInfo, + OutputFileHints outputFileHints) { + DecimalFormat df = new DecimalFormat("0000"); + IntervalWindow intervalWindow = (IntervalWindow) window; + String prefix = + baseFilename.isDirectory() ? "" : firstNonNull(baseFilename.getFilename(), ""); + String filename = + String.format( + "%s%s-%s-pane-%d-%s-of-%s%s%s", + prefix, + FORMATTER.print(intervalWindow.start()), + FORMATTER.print(intervalWindow.end()), + paneInfo.getIndex(), + df.format(shardNumber), + df.format(numShards), + outputFileHints.getSuggestedFilenameSuffix(), + suffix); + return baseFilename + .getCurrentDirectory() + .resolve(filename, StandardResolveOptions.RESOLVE_FILE); + } + + @Override + public @Nullable ResourceId unwindowedFilename( + int shardNumber, int numShards, OutputFileHints outputFileHints) { + throw new UnsupportedOperationException(); + } + } + /** * Same as {@link #runShardedWrite(List, PTransform, String, WriteFiles, BiFunction, boolean)} but * without shard content check. This means content will be checked only globally, that shards From d9d8af4137a032d5374da624b7395bceace02e4c Mon Sep 17 00:00:00 2001 From: Tejas Iyer Date: Mon, 12 Oct 2026 19:16:35 +0000 Subject: [PATCH 2/6] Remove experiment flag and configure direct write solely via withDirectWriteByteThreshold --- .../beam/runners/dataflow/DataflowRunner.java | 5 +- .../runners/dataflow/DataflowRunnerTest.java | 45 ++------------- .../org/apache/beam/sdk/io/WriteFiles.java | 57 +------------------ .../apache/beam/sdk/io/WriteFilesTest.java | 49 ++-------------- 4 files changed, 13 insertions(+), 143 deletions(-) diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java index b066b12790ee..d3d6855ec47a 100644 --- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java +++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java @@ -2708,13 +2708,12 @@ static class StreamingShardedWriteFactory } /** - * Returns the direct write threshold set on the original {@link WriteFiles}, else the value of - * the {@link WriteFiles#DIRECT_WRITE_THRESHOLD_EXPERIMENT} experiment, else {@code null}. + * Returns the direct write threshold set on the original {@link WriteFiles}, else {@code null}. */ private @Nullable Long getDirectWriteThreshold(AppliedPTransform transform) { Object original = transform.getTransform(); if (original instanceof WriteFiles) { - return ((WriteFiles) original).resolveDirectWriteThreshold(options); + return ((WriteFiles) original).getDirectWriteByteThreshold(); } return null; } diff --git a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java index a2ffeca62785..3d05c2486c1e 100644 --- a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java +++ b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java @@ -2154,16 +2154,16 @@ public void testStreamingWriteWithShardingReturnsSameTransform() { } @Test - public void testStreamingWriteDirectWriteThresholdExperiment() { + public void testStreamingWriteWithDirectWriteByteThreshold() { PipelineOptions options = TestPipeline.testingPipelineOptions(); - ExperimentalOptions.addExperiment( - options.as(ExperimentalOptions.class), "streaming_file_direct_write_threshold_bytes=1000"); TestPipeline p = TestPipeline.fromOptions(options); StreamingShardedWriteFactory factory = new StreamingShardedWriteFactory<>(p.getOptions()); WriteFiles original = - WriteFiles.to(new TestSink(tmpFolder.toString())).withRunnerDeterminedSharding(); + WriteFiles.to(new TestSink(tmpFolder.toString())) + .withRunnerDeterminedSharding() + .withDirectWriteByteThreshold(1000L); PCollection objs = (PCollection) p.apply(Create.empty(VoidCoder.of())); AppliedPTransform, WriteFilesResult, WriteFiles> originalApplication = @@ -2209,43 +2209,6 @@ public void testStreamingWriteDirectWriteThresholdDefaultNull() { assertNull(replacement.getDirectWriteByteThreshold()); } - @Test - public void testStreamingWriteDirectWriteThresholdInvalidExperiment() { - for (String badValue : ImmutableList.of("abc", "0", "-1")) { - PipelineOptions options = TestPipeline.testingPipelineOptions(); - ExperimentalOptions.addExperiment( - options.as(ExperimentalOptions.class), - "streaming_file_direct_write_threshold_bytes=" + badValue); - - TestPipeline p = TestPipeline.fromOptions(options); - StreamingShardedWriteFactory factory = - new StreamingShardedWriteFactory<>(p.getOptions()); - WriteFiles original = - WriteFiles.to(new TestSink(tmpFolder.toString())).withRunnerDeterminedSharding(); - PCollection objs = (PCollection) p.apply(Create.empty(VoidCoder.of())); - AppliedPTransform< - PCollection, WriteFilesResult, WriteFiles> - originalApplication = - AppliedPTransform.of( - "writefiles", - PValues.expandInput(objs), - Collections.emptyMap(), - original, - ResourceHints.create(), - p); - - IllegalArgumentException thrown = - assertThrows( - IllegalArgumentException.class, - () -> factory.getReplacementTransform(originalApplication)); - assertThat( - thrown.getMessage(), - containsString( - "Experiment streaming_file_direct_write_threshold_bytes must be a positive integer, got: " - + badValue)); - } - } - private void verifyMergingStatefulParDoRejected(PipelineOptions options) throws Exception { Pipeline p = Pipeline.create(options); 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 05128530bd0c..6f68591024f2 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 @@ -51,7 +51,6 @@ import org.apache.beam.sdk.io.fs.ResourceId; import org.apache.beam.sdk.metrics.Counter; import org.apache.beam.sdk.metrics.Metrics; -import org.apache.beam.sdk.options.ExperimentalOptions; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.ValueProvider; import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider; @@ -162,14 +161,6 @@ public abstract class WriteFiles public static final Duration FILE_TRIGGERING_RECORD_BUFFERING_DURATION = Duration.standardSeconds(5); - /** - * Experiment name for the byte threshold above which a streaming bundle is written directly to - * temporary files. - */ - @Internal - public static final String DIRECT_WRITE_THRESHOLD_EXPERIMENT = - "streaming_file_direct_write_threshold_bytes"; - static final int UNKNOWN_SHARDNUM = -1; static final int DUMMY_SHARDNUM = 0; private @Nullable WriteOperation writeOperation; @@ -477,8 +468,6 @@ public WriteFiles withBatchMaxBufferingDuration( * existing grouped path so output files stay large. Shard numbers are assigned at finalization, * which may override {@link #withNumShards(int)} (the number of files per window can differ from * the configured number of shards). - * - *

Can also be set via the {@link #DIRECT_WRITE_THRESHOLD_EXPERIMENT} pipeline experiment. */ public WriteFiles withDirectWriteByteThreshold( long thresholdBytes) { @@ -486,45 +475,6 @@ public WriteFiles withDirectWriteByteThreshold( return toBuilder().setDirectWriteByteThreshold(thresholdBytes).build(); } - /** - * Returns the direct write threshold set via {@link #withDirectWriteByteThreshold(long)}, else - * the value of the {@link #DIRECT_WRITE_THRESHOLD_EXPERIMENT} experiment in {@code options} when - * no explicit shard count is configured, else {@code null}. - */ - @Internal - public @Nullable Long resolveDirectWriteThreshold(PipelineOptions options) { - if (getDirectWriteByteThreshold() != null) { - return getDirectWriteByteThreshold(); - } - if (getComputeNumShards() != null || getNumShardsProvider() != null) { - return null; - } - String fromExperiment = - ExperimentalOptions.getExperimentValue(options, DIRECT_WRITE_THRESHOLD_EXPERIMENT); - if (fromExperiment == null) { - return null; - } - final long threshold; - try { - threshold = Long.parseLong(fromExperiment); - } catch (NumberFormatException e) { - throw new IllegalArgumentException( - "Experiment " - + DIRECT_WRITE_THRESHOLD_EXPERIMENT - + " must be a positive integer, got: " - + fromExperiment, - e); - } - if (threshold <= 0) { - throw new IllegalArgumentException( - "Experiment " - + DIRECT_WRITE_THRESHOLD_EXPERIMENT - + " must be a positive integer, got: " - + fromExperiment); - } - return threshold; - } - public WriteFiles withSideInputs( List> sideInputs) { return toBuilder().setSideInputs(sideInputs).build(); @@ -663,7 +613,7 @@ public WriteFilesResult expand(PCollection input) { (getComputeNumShards() == null) ? null : input.apply(getComputeNumShards()); boolean fixedSharding = getComputeNumShards() != null || getNumShardsProvider() != null; - Long directWriteThreshold = resolveDirectWriteThreshold(input.getPipeline().getOptions()); + Long directWriteThreshold = getDirectWriteByteThreshold(); if (directWriteThreshold != null && getSideInputs().isEmpty() && input.isBounded() == IsBounded.UNBOUNDED) { @@ -1360,10 +1310,7 @@ protected void reportElementSize(long elementByteSize) { * Buffers each bundle's records until their encoded size reaches {@code thresholdBytes}. Bundles * that reach it are written directly to temporary files. Records of smaller bundles, and records * that would exceed the per-bundle writer limit, are emitted to {@code smallBundlesTag} at the - * end of the bundle for the sharded {@link GroupByKey} path. - * - *

Adapted from the IcebergIO bundle lifter (apache/beam#36720), but buffers at most {@code - * thresholdBytes} rather than the whole bundle. + * end of the bundle for the grouped path. */ private class DirectWriteLargeBundlesFn extends DoFn> { private final TupleTag smallBundlesTag; 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 35ddb839dfdf..40d6efb4a671 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 @@ -70,7 +70,6 @@ import org.apache.beam.sdk.metrics.MetricResult; import org.apache.beam.sdk.metrics.MetricsFilter; import org.apache.beam.sdk.options.Description; -import org.apache.beam.sdk.options.ExperimentalOptions; import org.apache.beam.sdk.options.PipelineOptionsFactoryTest.TestPipelineOptions; import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider; import org.apache.beam.sdk.testing.BeamParallelJunit4Runner; @@ -445,7 +444,7 @@ public void testDirectWriteAllBundlesSmall() throws IOException { @Test @Category({NeedsRunner.class, UsesUnboundedPCollections.class, UsesTestStream.class}) - public void testDirectWriteUnshardedViaExperiment() throws IOException { + public void testDirectWriteUnshardedMixedBundles() throws IOException { List elements = Lists.newArrayList(); for (int i = 0; i < 30; ++i) { elements.add("number: " + i); @@ -466,12 +465,11 @@ public void testDirectWriteUnshardedViaExperiment() throws IOException { .addElements(elements.get(28), elements.get(29)) .advanceWatermarkToInfinity(); - ExperimentalOptions.addExperiment( - p.getOptions().as(ExperimentalOptions.class), - WriteFiles.DIRECT_WRITE_THRESHOLD_EXPERIMENT + "=50"); - WriteFiles write = - WriteFiles.to(makeSimpleSink()).withWindowedWrites().withRunnerDeterminedSharding(); + WriteFiles.to(makeSimpleSink()) + .withWindowedWrites() + .withRunnerDeterminedSharding() + .withDirectWriteByteThreshold(50L); p.apply(testStream) .apply(Window.into(FixedWindows.of(Duration.standardSeconds(10)))) .apply(write) @@ -484,43 +482,6 @@ public void testDirectWriteUnshardedViaExperiment() throws IOException { assertThat(counterValue(result, "groupedWriteBundles"), Matchers.greaterThan(0L)); } - @Test - @Category({NeedsRunner.class, UsesUnboundedPCollections.class, UsesTestStream.class}) - public void testDirectWriteExperimentIgnoredWithExplicitNumShards() throws IOException { - List elements = Lists.newArrayList(); - for (int i = 0; i < 30; ++i) { - elements.add("number: " + i); - } - Instant startInstant = new Instant(0L); - TestStream testStream = - TestStream.create(StringUtf8Coder.of()) - .advanceWatermarkTo(startInstant) - .addElements(elements.get(0), Iterables.toArray(elements.subList(1, 10), String.class)) - .advanceWatermarkTo(startInstant.plus(Duration.standardSeconds(5))) - .addElements( - elements.get(10), Iterables.toArray(elements.subList(11, 20), String.class)) - .advanceWatermarkTo(startInstant.plus(Duration.standardSeconds(10))) - .addElements( - elements.get(20), Iterables.toArray(elements.subList(21, 30), String.class)) - .advanceWatermarkToInfinity(); - - ExperimentalOptions.addExperiment( - p.getOptions().as(ExperimentalOptions.class), - WriteFiles.DIRECT_WRITE_THRESHOLD_EXPERIMENT + "=1"); - - WriteFiles write = - WriteFiles.to(makeSimpleSink()).withWindowedWrites().withNumShards(3); - p.apply(testStream) - .apply(Window.into(FixedWindows.of(Duration.standardSeconds(10)))) - .apply(write) - .getPerDestinationOutputFilenames() - .apply(new VerifyFilesExist<>()); - PipelineResult result = p.run(); - - checkFileContents(getBaseOutputFilename(), elements, Optional.absent(), false, null); - assertEquals(0L, counterValue(result, "directWriteBundles")); - } - @Test @Category({NeedsRunner.class, UsesUnboundedPCollections.class, UsesTestStream.class}) public void testDirectWriteNonDefaultTriggerLargeBundles() throws IOException { From 2f21a802135a752b5a1c315dfb8ebd6f1773ceb0 Mon Sep 17 00:00:00 2001 From: Tejas Iyer Date: Tue, 3 Nov 2026 00:07:50 +0000 Subject: [PATCH 3/6] Clarify finishBundle branches and remove stale lines from CHANGES.md --- CHANGES.md | 3 +-- .../org/apache/beam/sdk/io/WriteFiles.java | 18 +++++++++++------- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index aae1d14e820b..4592a2197541 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -73,8 +73,7 @@ * The `disableCounterMetrics`, `disableStringSetMetrics` and `disableBoundedTrieMetrics` experiments are now honored by the Python SDK, as they already were in Java (Python) ([#38746](https://github.com/apache/beam/issues/38746)). * (Go) Pipeline-level `--resource_hints` are now forwarded to expansion services, so they apply to cross-language transforms as they do in the Java and Python SDKs. The `max_active_bundles_per_worker` hint name is now also accepted by `--resource_hints` ([#23893](https://github.com/apache/beam/issues/23893)). * ReadFromBigQuery now supports Lakehouse runtime catalog (BigLake metastore) tables with `method=DIRECT_READ`, using 4-part `project.catalog.namespace.table` identifiers. Previously `project:catalog.namespace.table` was silently mis-parsed and tables that report no `numBytes` failed to split (Python) ([#39597](https://github.com/apache/beam/issues/39597)). -* (Java) `TextIO`, `FileIO`, and `WriteFiles`: streaming writes can skip the shuffle for large bundles by writing them directly to temp files (`withDirectWriteByteThreshold` or experiment `streaming_file_direct_write_threshold_bytes`). Off by default ([#40533](https://github.com/apache/beam/issues/40533)). -* (Python) Added `AwsOptions` with `--aws_*` pipeline options (e.g. `--aws_access_key_id`, `--aws_region_name`, `--aws_endpoint_url`) for configuring AWS service clients such as S3 ([#40528](https://github.com/apache/beam/issues/40528)). +* (Java) `TextIO`, `FileIO`, and `WriteFiles`: streaming writes can skip the shuffle for large bundles by writing them directly to temp files (`withDirectWriteByteThreshold`). Off by default ([#40533](https://github.com/apache/beam/issues/40533)). * (Python) Added `AwsOptions` with `--aws_*` pipeline options (e.g. `--aws_access_key_id`, `--aws_region_name`, `--aws_endpoint_url`) that set defaults for all AWS service clients. The existing `--s3_*` options in `S3Options` override these defaults for the S3 client only ([#40528](https://github.com/apache/beam/issues/40528)). * (Python) Added `--aws_credentials_secret_manager` so AWS credentials (`--aws_access_key_id`, `--aws_secret_access_key`, `--aws_session_token` and their `--s3_*` overrides) can be secret specs resolved through a secret manager such as `GoogleCloudSecretManager`, instead of plain-text values ([#40528](https://github.com/apache/beam/issues/40528)). 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 6f68591024f2..0780dcd4e474 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 @@ -1409,17 +1409,21 @@ private void write(BufferedRecord record, MultiOutputReceiver outputRecei @FinishBundle public void finishBundle(FinishBundleContext c) throws Exception { + closeAllWritersAndOutputFileResults(writers, c); + // Exactly one of these cases applies, and records are emitted in arrival order: + // - Direct mode: the buffer was flushed when the threshold was crossed, so only records + // that exceeded the per-bundle writer limit remain (in `spilled`). + // - Grouped mode: nothing was written directly, so every record is still in `buffer`. if (directMode) { directBundles.inc(); + for (BufferedRecord record : spilled) { + c.output(smallBundlesTag, record.value, record.timestamp, record.window); + } } else if (!buffer.isEmpty()) { groupedBundles.inc(); - } - closeAllWritersAndOutputFileResults(writers, c); - for (BufferedRecord record : buffer) { - c.output(smallBundlesTag, record.value, record.timestamp, record.window); - } - for (BufferedRecord record : spilled) { - c.output(smallBundlesTag, record.value, record.timestamp, record.window); + for (BufferedRecord record : buffer) { + c.output(smallBundlesTag, record.value, record.timestamp, record.window); + } } buffer.clear(); spilled.clear(); From 14904e46cb33e02167b331da5239e257d89a46a4 Mon Sep 17 00:00:00 2001 From: Tejas Iyer Date: Tue, 3 Nov 2026 01:40:31 +0000 Subject: [PATCH 4/6] Share grouping and batch-default logic between direct-write and auto-sharded paths --- .../org/apache/beam/sdk/io/WriteFiles.java | 191 ++++++++++-------- 1 file changed, 105 insertions(+), 86 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 0780dcd4e474..78afc8a1a79b 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 @@ -653,9 +653,9 @@ && getSideInputs().isEmpty() new WriteAutoShardedBundlesToTempFiles( destinationCoder, fileResultCoder, - getBatchSize(), - getBatchSizeBytes(), - getBatchMaxBufferingDuration())); + getBatchSizeOrDefault(), + getBatchSizeBytesOrDefault(), + getBatchMaxBufferingDurationOrDefault())); } } @@ -665,6 +665,23 @@ && getSideInputs().isEmpty() numShardsView, destinationCoder, /* assignShardNumbersInFinalize= */ false)); } + private int getBatchSizeOrDefault() { + Integer batchSize = getBatchSize(); + return batchSize != null ? batchSize : FILE_TRIGGERING_RECORD_COUNT; + } + + private int getBatchSizeBytesOrDefault() { + Integer batchSizeBytes = getBatchSizeBytes(); + return batchSizeBytes != null ? batchSizeBytes : FILE_TRIGGERING_BYTE_COUNT; + } + + private Duration getBatchMaxBufferingDurationOrDefault() { + Duration maxBufferingDuration = getBatchMaxBufferingDuration(); + return maxBufferingDuration != null + ? maxBufferingDuration + : FILE_TRIGGERING_RECORD_BUFFERING_DURATION; + } + @Override public void populateDisplayData(DisplayData.Builder builder) { super.populateDisplayData(builder); @@ -1236,11 +1253,9 @@ public void process(ProcessContext c) { smallBundles, destinationCoder, fileResultCoder, - getBatchSize() != null ? getBatchSize() : FILE_TRIGGERING_RECORD_COUNT, - getBatchSizeBytes() != null ? getBatchSizeBytes() : FILE_TRIGGERING_BYTE_COUNT, - getBatchMaxBufferingDuration() != null - ? getBatchMaxBufferingDuration() - : FILE_TRIGGERING_RECORD_BUFFERING_DURATION); + getBatchSizeOrDefault(), + getBatchSizeBytesOrDefault(), + getBatchMaxBufferingDurationOrDefault()); } // When the small path goes through GroupIntoShards (GBK), its trigger transitions to the @@ -1254,46 +1269,19 @@ public void process(ProcessContext c) { .setWindowingStrategyInternal(groupedFiles.getWindowingStrategy()); // Group all temp files per destination so each window/pane firing is finalized together - // and shard numbers are assigned once across both paths. - return PCollectionList.of(directFiles) - .and(groupedFiles) - .apply("FlattenTempFiles", Flatten.pCollections()) - .setCoder(fileResultCoder) - .apply( - "KeyedByDestination", - WithKeys.of( - new SimpleFunction, DestinationT>() { - @Override - public DestinationT apply(FileResult input) { - return input.getDestination(); - } - })) - .setCoder(KvCoder.of(destinationCoder, fileResultCoder)) - .apply("GroupTempFilesByDestination", GroupByKey.create()) - .apply( - "ExtractValuesToList", - ParDo.of( - new DoFn< - KV>>, - List>>() { - @ProcessElement - public void processElement(ProcessContext c) { - // Stamp every file with this firing's pane so files from different - // firings of the same window get distinct final names. - List> result = new ArrayList<>(); - for (FileResult e : c.element().getValue()) { - result.add( - new FileResult<>( - e.getTempFilename(), - UNKNOWN_SHARDNUM, - e.getWindow(), - c.pane(), - e.getDestination())); - } - c.output(result); - } - })) - .setCoder(ListCoder.of(fileResultCoder)); + // and shard numbers are assigned once across both paths. Every file is stamped with this + // firing's pane so files from different firings of the same window get distinct names. + PCollection> allTempFiles = + PCollectionList.of(directFiles) + .and(groupedFiles) + .apply("FlattenTempFiles", Flatten.pCollections()) + .setCoder(fileResultCoder); + return groupTempFilesByDestination( + allTempFiles, + destinationCoder, + fileResultCoder, + "GroupTempFilesByDestination", + /* stampFiringPane= */ true); } } @@ -1544,17 +1532,14 @@ private class WriteAutoShardedBundlesToTempFiles private WriteAutoShardedBundlesToTempFiles( Coder destinationCoder, Coder> fileResultCoder, - Integer batchSize, - Integer batchSizeBytes, + int batchSize, + int batchSizeBytes, Duration maxBufferingDuration) { this.destinationCoder = destinationCoder; this.fileResultCoder = fileResultCoder; - this.batchSize = batchSize != null ? batchSize : FILE_TRIGGERING_RECORD_COUNT; - this.batchSizeBytes = batchSizeBytes != null ? batchSizeBytes : FILE_TRIGGERING_BYTE_COUNT; - this.maxBufferingDuration = - maxBufferingDuration != null - ? maxBufferingDuration - : FILE_TRIGGERING_RECORD_BUFFERING_DURATION; + this.batchSize = batchSize; + this.batchSizeBytes = batchSizeBytes; + this.maxBufferingDuration = maxBufferingDuration; } @Override @@ -1571,36 +1556,70 @@ public PCollection>> expand(PCollection inp // Group temp file results by destinations again to gather all the results in the same window. // This is needed since we don't have shard idx associated with each temp file so have to rely // on the indexing within a bundle. - return tempFiles - .apply( - "KeyedByDestination", - WithKeys.of( - new SimpleFunction, DestinationT>() { - @Override - public DestinationT apply(FileResult input) { - return input.getDestination(); - } - })) - .setCoder(KvCoder.of(destinationCoder, fileResultCoder)) - .apply(GroupByKey.create()) - .apply( - "ExtractValuesToList", - ParDo.of( - new DoFn< - KV>>, - List>>() { - @ProcessElement - public void processElement( - @Element KV>> element, - ProcessContext c) { - List> result = new ArrayList<>(); - for (FileResult e : element.getValue()) { - result.add(e); - } - c.output(result); - } - })) - .setCoder(ListCoder.of(fileResultCoder)); + // "GroupByKey" matches the default name this step has always had (update compatibility). + return groupTempFilesByDestination( + tempFiles, destinationCoder, fileResultCoder, "GroupByKey", /* stampFiringPane= */ false); + } + } + + /** + * Groups temp file results by destination into one list per destination and window firing. Shared + * by the auto-sharded write and the direct-write path. Produces the steps {@code + * KeyedByDestination}, {@code groupByKeyName} and {@code ExtractValuesToList}. + * + * @param stampFiringPane if true, each file is re-stamped with the firing's pane and an unknown + * shard number so shard numbers can be assigned in finalize across all files of the firing. + */ + private PCollection>> groupTempFilesByDestination( + PCollection> tempFiles, + Coder destinationCoder, + Coder> fileResultCoder, + String groupByKeyName, + boolean stampFiringPane) { + return tempFiles + .apply( + "KeyedByDestination", + WithKeys.of( + new SimpleFunction, DestinationT>() { + @Override + public DestinationT apply(FileResult input) { + return input.getDestination(); + } + })) + .setCoder(KvCoder.of(destinationCoder, fileResultCoder)) + .apply(groupByKeyName, GroupByKey.create()) + .apply("ExtractValuesToList", ParDo.of(new ExtractValuesToListFn<>(stampFiringPane))) + .setCoder(ListCoder.of(fileResultCoder)); + } + + private static class ExtractValuesToListFn + extends DoFn< + KV>>, List>> { + private final boolean stampFiringPane; + + ExtractValuesToListFn(boolean stampFiringPane) { + this.stampFiringPane = stampFiringPane; + } + + @ProcessElement + public void processElement(ProcessContext c) { + List> result = new ArrayList<>(); + for (FileResult e : c.element().getValue()) { + if (stampFiringPane) { + // Stamp every file with this firing's pane so files from different firings of the + // same window get distinct final names. + result.add( + new FileResult<>( + e.getTempFilename(), + UNKNOWN_SHARDNUM, + e.getWindow(), + c.pane(), + e.getDestination())); + } else { + result.add(e); + } + } + c.output(result); } } From b1fe445521135fe623ef2d69ae1dd11361dbf1e4 Mon Sep 17 00:00:00 2001 From: Tejas Iyer Date: Tue, 10 Nov 2026 21:55:23 +0000 Subject: [PATCH 5/6] Stream elements to the large output once BundleLifter's threshold is reached. Evaluate the direct-write threshold per WriterKey instead of per bundle. --- .../org/apache/beam/sdk/io/WriteFiles.java | 185 +++++----- .../apache/beam/sdk/util/BundleLifter.java | 318 +++++++++++++++--- .../apache/beam/sdk/io/WriteFilesTest.java | 101 ++++++ .../beam/sdk/util/BundleLifterTest.java | 159 +++++++++ 4 files changed, 613 insertions(+), 150 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 3242139f6c1c..e1a74333daba 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 @@ -64,6 +64,7 @@ import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.transforms.Reify; import org.apache.beam.sdk.transforms.Reshuffle; +import org.apache.beam.sdk.transforms.SerializableFunction; import org.apache.beam.sdk.transforms.SimpleFunction; import org.apache.beam.sdk.transforms.Values; import org.apache.beam.sdk.transforms.View; @@ -79,9 +80,9 @@ import org.apache.beam.sdk.transforms.windowing.GlobalWindows; import org.apache.beam.sdk.transforms.windowing.PaneInfo; import org.apache.beam.sdk.transforms.windowing.Window; +import org.apache.beam.sdk.util.BundleLifter; import org.apache.beam.sdk.util.CoderUtils; import org.apache.beam.sdk.util.MoreFutures; -import org.apache.beam.sdk.util.common.ElementByteSizeObserver; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.PCollection.IsBounded; @@ -1209,22 +1210,37 @@ private WriteShardedBundlesWithDirectWrite( @Override public PCollection>> expand(PCollection input) { - TupleTag> directFilesTag = new TupleTag<>("directFiles"); TupleTag smallBundlesTag = new TupleTag<>("smallBundles"); + TupleTag largeBundlesTag = new TupleTag<>("largeBundles"); + TupleTag> directFilesTag = new TupleTag<>("directFiles"); + TupleTag spilledRecordsTag = new TupleTag<>("spilledRecords"); - PCollectionTuple routed = + // Evaluate the threshold per window, pane and destination, i.e. per output file. + PCollectionTuple lifted = input.apply( - "DirectWriteLargeBundles", - ParDo.of( - new DirectWriteLargeBundlesFn( - smallBundlesTag, thresholdBytes, input.getCoder())) - .withSideInputs(getSideInputs()) - .withOutputTags( - directFilesTag, - TupleTagList.of(ImmutableList.of(smallBundlesTag, BAD_RECORD_TAG)))); + "LiftLargeBundles", + BundleLifter.of(smallBundlesTag, largeBundlesTag, thresholdBytes, input.getCoder()) + .withKeyFn(new DestinationKeyFn<>(getDynamicDestinations()))); + + PCollectionTuple routed = + lifted + .get(largeBundlesTag) + .setCoder(input.getCoder()) + .apply( + "DirectWriteLargeBundles", + ParDo.of(new DirectWriteRecordsFn(spilledRecordsTag, input.getCoder())) + .withSideInputs(getSideInputs()) + .withOutputTags( + directFilesTag, + TupleTagList.of(ImmutableList.of(spilledRecordsTag, BAD_RECORD_TAG)))); addErrorCollection(routed); - PCollection smallBundles = routed.get(smallBundlesTag).setCoder(input.getCoder()); + PCollection smallBundles = + PCollectionList.of(lifted.get(smallBundlesTag).setCoder(input.getCoder())) + .and(routed.get(spilledRecordsTag).setCoder(input.getCoder())) + .apply("FlattenSmallBundles", Flatten.pCollections()) + .apply("CountGroupedBundles", ParDo.of(new CountGroupedBundlesFn<>())) + .setCoder(input.getCoder()); PCollection> groupedFiles; if (getComputeNumShards() != null || getNumShardsProvider() != null) { groupedFiles = @@ -1281,50 +1297,73 @@ public void process(ProcessContext c) { } } - private static class ByteSizeObserver extends ElementByteSizeObserver { - private long byteSize; + /** + * Key function for {@link BundleLifter}: the record's destination, so that together with the + * window and pane the threshold is evaluated per output file. Returns {@code null} if the + * destination cannot be computed, which sends the record to the grouped path where the usual + * error handling applies. + */ + private static class DestinationKeyFn + implements SerializableFunction { + private final DynamicDestinations dynamicDestinations; + + DestinationKeyFn(DynamicDestinations dynamicDestinations) { + this.dynamicDestinations = dynamicDestinations; + } @Override - protected void reportElementSize(long elementByteSize) { - this.byteSize = elementByteSize; + public @Nullable Object apply(UserT record) { + try { + // Wrapped so that a null destination (e.g. Void) is still a valid key. + return java.util.Optional.ofNullable(dynamicDestinations.getDestination(record)); + } catch (Exception e) { + return null; + } + } + } + + /** Passes records through, counting the bundles that take the grouped path. */ + private static class CountGroupedBundlesFn extends DoFn { + private final Counter groupedBundles = Metrics.counter(WriteFiles.class, "groupedWriteBundles"); + private transient boolean counted; + + @StartBundle + public void startBundle() { + counted = false; + } + + @ProcessElement + public void processElement(@Element T element, OutputReceiver receiver) { + if (!counted) { + groupedBundles.inc(); + counted = true; + } + receiver.output(element); } } /** - * Buffers each bundle's records until their encoded size reaches {@code thresholdBytes}. Bundles - * that reach it are written directly to temporary files. Records of smaller bundles, and records - * that would exceed the per-bundle writer limit, are emitted to {@code smallBundlesTag} at the - * end of the bundle for the grouped path. + * Writes the records lifted by {@link BundleLifter} directly to temporary files, one writer per + * window, pane and destination. Records that would exceed the per-bundle writer limit are emitted + * to {@code spilledRecordsTag} for the grouped path. */ - private class DirectWriteLargeBundlesFn extends DoFn> { - private final TupleTag smallBundlesTag; - private final long thresholdBytes; + private class DirectWriteRecordsFn extends DoFn> { + private final TupleTag spilledRecordsTag; private final Coder inputCoder; private final Counter directBundles = Metrics.counter(WriteFiles.class, "directWriteBundles"); - private final Counter groupedBundles = Metrics.counter(WriteFiles.class, "groupedWriteBundles"); // Per-bundle state, reset in startBundle. - private transient List> buffer; - private transient long bufferedBytes; - private transient boolean directMode; private transient Map, Writer> writers; - private transient List> spilled; - DirectWriteLargeBundlesFn( - TupleTag smallBundlesTag, long thresholdBytes, Coder inputCoder) { - this.smallBundlesTag = smallBundlesTag; - this.thresholdBytes = thresholdBytes; + DirectWriteRecordsFn(TupleTag spilledRecordsTag, Coder inputCoder) { + this.spilledRecordsTag = spilledRecordsTag; this.inputCoder = inputCoder; } @StartBundle public void startBundle() { - buffer = new ArrayList<>(); - bufferedBytes = 0; - directMode = false; writers = new LinkedHashMap<>(); - spilled = new ArrayList<>(); } @ProcessElement @@ -1332,44 +1371,13 @@ public void processElement( ProcessContext c, BoundedWindow window, MultiOutputReceiver outputReceiver) throws Exception { getDynamicDestinations().setSideInputAccessorFromProcessContext(c); - BufferedRecord record = - new BufferedRecord<>(c.element(), c.timestamp(), window, c.pane()); - if (directMode) { - write(record, outputReceiver); - return; - } - buffer.add(record); - bufferedBytes += estimateSize(c.element()); - if (bufferedBytes >= thresholdBytes) { - directMode = true; - // Direct writes are disabled when destinations use side inputs (see - // expand()), so resolving buffered records' destinations here is safe. - for (BufferedRecord buffered : buffer) { - write(buffered, outputReceiver); - } - buffer.clear(); - } - } - - private long estimateSize(UserT value) throws Exception { - if (inputCoder.isRegisterByteSizeObserverCheap(value)) { - ByteSizeObserver observer = new ByteSizeObserver(); - inputCoder.registerByteSizeObserver(value, observer); - observer.advance(); - return observer.byteSize; - } - return CoderUtils.encodeToByteArray(inputCoder, value).length; - } - - private void write(BufferedRecord record, MultiOutputReceiver outputReceiver) - throws Exception { MaybeDestination maybeDestination = - getDestinationWithErrorHandling(record.value, outputReceiver, inputCoder); + getDestinationWithErrorHandling(c.element(), outputReceiver, inputCoder); if (!maybeDestination.isValid) { return; } DestinationT destination = maybeDestination.destination; - WriterKey key = new WriterKey<>(record.window, record.pane, destination); + WriterKey key = new WriterKey<>(window, c.pane(), destination); Writer writer = writers.get(key); if (writer == null) { int maxWriters = @@ -1378,13 +1386,13 @@ private void write(BufferedRecord record, MultiOutputReceiver outputRecei : getMaxNumWritersPerBundle(); if (writers.size() >= maxWriters) { // Too many open writers in this bundle; send the record to the grouped path. - spilled.add(record); + outputReceiver.get(spilledRecordsTag).output(c.element()); return; } writer = openAndRegisterWriter(writers, key); } OutputT formattedRecord = - formatRecordWithErrorHandling(record.value, outputReceiver, inputCoder); + formatRecordWithErrorHandling(c.element(), outputReceiver, inputCoder); if (formattedRecord == null) { return; } @@ -1393,42 +1401,11 @@ private void write(BufferedRecord record, MultiOutputReceiver outputRecei @FinishBundle public void finishBundle(FinishBundleContext c) throws Exception { - closeAllWritersAndOutputFileResults(writers, c); - // Exactly one of these cases applies, and records are emitted in arrival order: - // - Direct mode: the buffer was flushed when the threshold was crossed, so only records - // that exceeded the per-bundle writer limit remain (in `spilled`). - // - Grouped mode: nothing was written directly, so every record is still in `buffer`. - if (directMode) { + if (!writers.isEmpty()) { directBundles.inc(); - for (BufferedRecord record : spilled) { - c.output(smallBundlesTag, record.value, record.timestamp, record.window); - } - } else if (!buffer.isEmpty()) { - groupedBundles.inc(); - for (BufferedRecord record : buffer) { - c.output(smallBundlesTag, record.value, record.timestamp, record.window); - } } - buffer.clear(); - spilled.clear(); + closeAllWritersAndOutputFileResults(writers, c); writers.clear(); - bufferedBytes = 0; - directMode = false; - } - } - - /** A record held by {@link DirectWriteLargeBundlesFn} with the metadata needed to re-emit it. */ - private static class BufferedRecord { - private final T value; - private final Instant timestamp; - private final BoundedWindow window; - private final PaneInfo pane; - - BufferedRecord(T value, Instant timestamp, BoundedWindow window, PaneInfo pane) { - this.value = value; - this.timestamp = timestamp; - this.window = window; - this.pane = pane; } } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/BundleLifter.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/BundleLifter.java index 539ef147f173..0edd929d0a57 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/BundleLifter.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/BundleLifter.java @@ -17,30 +17,55 @@ */ package org.apache.beam.sdk.util; -import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.LinkedHashMap; import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; import org.apache.beam.sdk.annotations.Internal; +import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.transforms.DoFn; -import org.apache.beam.sdk.transforms.DoFn.MultiOutputReceiver; -import org.apache.beam.sdk.transforms.DoFn.OutputReceiver; import org.apache.beam.sdk.transforms.PTransform; import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.transforms.SerializableFunction; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.PaneInfo; +import org.apache.beam.sdk.util.common.ElementByteSizeObserver; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.PCollectionTuple; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.TupleTagList; import org.checkerframework.checker.nullness.qual.MonotonicNonNull; import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Duration; +import org.joda.time.Instant; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** - * A PTransform that buffers elements and outputs them to one of two TupleTags based on the total - * size of the bundle in finish_bundle. + * A PTransform that routes elements to one of two TupleTags depending on how many bytes of them + * arrive together within a bundle. + * + *

Elements are buffered per key. Once the buffered elements of a key reach the threshold, they + * are emitted to the large tag, and all later elements of that key in the same bundle are emitted + * to the large tag directly without buffering. Elements of keys that never reach the threshold are + * emitted to the small tag at the end of the bundle. + * + *

Without a key function, all elements of a bundle share a single key, so the threshold applies + * to the whole bundle. With {@link #withKeyFn}, the key is the element's window, pane and the value + * returned by the key function. Elements for which the key function returns {@code null} are + * emitted to the small tag immediately. + * + *

To bound memory, at most twice the threshold is buffered in total. When that is exceeded, the + * key with the fewest buffered bytes is emitted to the small tag. + * + *

Every element keeps its own timestamp and window. * *

For internal use only; no backwards-compatibility guarantees. * @@ -57,89 +82,272 @@ public class BundleLifter extends PTransform, PCollectionTuple final TupleTag smallBatchTag; final TupleTag largeBatchTag; - final int threshold; + final long threshold; final SerializableFunction elementSizer; + final @Nullable SerializableFunction keyFn; /** - * A DoFn that buffers elements within a bundle and outputs them to different tags in - * finish_bundle based on the total bundle size. + * A DoFn that buffers elements per key within a bundle and outputs them to the large tag once a + * key reaches the threshold, or to the small tag otherwise. * * @param The type of elements being processed. */ static class BundleLiftDoFn extends DoFn { private static final Logger LOG = LoggerFactory.getLogger(BundleLiftDoFn.class); + /** The key shared by all elements when no key function is set. */ + private static final Object UNKEYED = new Object(); + final TupleTag smallBatchTag; final TupleTag largeBatchTag; - final int threshold; + final long threshold; + final long maxBufferedBytes; final SerializableFunction elementSizer; + final @Nullable SerializableFunction keyFn; - private transient @MonotonicNonNull List buffer; - private transient long bundleSizeBytes; - private transient @Nullable MultiOutputReceiver receiver; + // Per-bundle state, reset in startBundle. + private transient @MonotonicNonNull Map> buffers; + private transient @MonotonicNonNull Set liftedKeys; + private transient long bufferedBytes; BundleLiftDoFn( TupleTag smallBatchTag, TupleTag largeBatchTag, - int threshold, + long threshold, SerializableFunction elementSizer) { + this(smallBatchTag, largeBatchTag, threshold, elementSizer, null); + } + + BundleLiftDoFn( + TupleTag smallBatchTag, + TupleTag largeBatchTag, + long threshold, + SerializableFunction elementSizer, + @Nullable SerializableFunction keyFn) { this.smallBatchTag = smallBatchTag; this.largeBatchTag = largeBatchTag; this.threshold = threshold; + this.maxBufferedBytes = threshold > Long.MAX_VALUE / 2 ? Long.MAX_VALUE : 2 * threshold; this.elementSizer = elementSizer; + this.keyFn = keyFn; } @StartBundle public void startBundle() { - buffer = new ArrayList<>(); - receiver = null; - bundleSizeBytes = 0L; + buffers = new LinkedHashMap<>(); + liftedKeys = new HashSet<>(); + bufferedBytes = 0L; } @ProcessElement - public void processElement(@Element T element, MultiOutputReceiver mor) { - if (receiver == null) { - receiver = mor; + public void processElement( + @Element T element, + @Timestamp Instant timestamp, + BoundedWindow window, + PaneInfo pane, + MultiOutputReceiver receiver) { + Map> buffers = + checkStateNotNull(this.buffers, "Buffers should be set by startBundle."); + Set liftedKeys = + checkStateNotNull(this.liftedKeys, "Lifted keys should be set by startBundle."); + + Object key; + if (keyFn == null) { + key = UNKEYED; + } else { + @Nullable Object userKey = keyFn.apply(element); + if (userKey == null) { + receiver.get(smallBatchTag).output(element); + return; + } + key = new WindowedKey(window, pane, userKey); + } + + if (liftedKeys.contains(key)) { + // This key already reached the threshold in this bundle; stream it through. + receiver.get(largeBatchTag).output(element); + return; + } + + KeyBuffer buffer = buffers.get(key); + if (buffer == null) { + buffer = new KeyBuffer<>(); + buffers.put(key, buffer); + } + long elementBytes = elementSizer.apply(element); + buffer.add(new BufferedElement<>(element, timestamp, window, pane), elementBytes); + bufferedBytes += elementBytes; + + if (buffer.sizeBytes >= threshold) { + buffers.remove(key); + bufferedBytes -= buffer.sizeBytes; + liftedKeys.add(key); + LOG.debug( + "Key reached {} estimated bytes; emitting {} buffered elements to tag: '{}'", + buffer.sizeBytes, + buffer.elements.size(), + largeBatchTag.getId()); + emit(buffer, receiver.get(largeBatchTag)); + return; + } + + while (bufferedBytes > maxBufferedBytes && !buffers.isEmpty()) { + evictSmallestKey(buffers, receiver.get(smallBatchTag)); } - checkArgumentNotNull(buffer, "Buffer should be set by startBundle."); - buffer.add(element); - bundleSizeBytes += elementSizer.apply(element); } @FinishBundle - public void finishBundle() { - checkArgumentNotNull(buffer, "Buffer should be set by startBundle."); - if (buffer.isEmpty()) { + public void finishBundle(FinishBundleContext context) { + Map> buffers = + checkStateNotNull(this.buffers, "Buffers should be set by startBundle."); + for (KeyBuffer buffer : buffers.values()) { + LOG.debug( + "Emitting {} elements of {} estimated bytes to tag: '{}'", + buffer.elements.size(), + buffer.sizeBytes, + smallBatchTag.getId()); + for (BufferedElement e : buffer.elements) { + context.output(smallBatchTag, e.value, e.timestamp, e.window); + } + } + buffers.clear(); + checkStateNotNull(this.liftedKeys).clear(); + bufferedBytes = 0L; + } + + /** + * Buffered elements are re-emitted from {@link ProcessElement} with their own timestamps, which + * may be earlier than the current element's. This is safe because they belong to the same + * bundle, so the input watermark is still held at or before them. + */ + @Override + public Duration getAllowedTimestampSkew() { + return Duration.millis(Long.MAX_VALUE); + } + + private void evictSmallestKey(Map> buffers, OutputReceiver small) { + Object smallestKey = null; + KeyBuffer smallest = null; + for (Map.Entry> entry : buffers.entrySet()) { + if (smallest == null || entry.getValue().sizeBytes < smallest.sizeBytes) { + smallestKey = entry.getKey(); + smallest = entry.getValue(); + } + } + if (smallestKey == null || smallest == null) { return; } + buffers.remove(smallestKey); + bufferedBytes -= smallest.sizeBytes; + emit(smallest, small); + } - // Select the target tag based on the bundle size - TupleTag targetTag; - targetTag = (bundleSizeBytes < threshold) ? smallBatchTag : largeBatchTag; - LOG.debug( - "Emitting {} elements of {} estimated bytes to tag: '{}'", - buffer.size(), - bundleSizeBytes, - targetTag.getId()); + private static void emit(KeyBuffer buffer, OutputReceiver receiver) { + for (BufferedElement e : buffer.elements) { + receiver.outputWindowedValue( + e.value, e.timestamp, Collections.singletonList(e.window), e.pane); + } + } + } - checkArgumentNotNull(receiver, "Receiver should be set by startBundle."); - OutputReceiver taggedOutput = receiver.get(targetTag); + /** The buffered elements of one key, in arrival order. */ + private static class KeyBuffer { + private final List> elements = new ArrayList<>(); + private long sizeBytes; - for (T element : buffer) { - taggedOutput.output(element); + void add(BufferedElement element, long elementBytes) { + elements.add(element); + sizeBytes += elementBytes; + } + } + + /** A buffered element with the metadata needed to re-emit it. */ + private static class BufferedElement { + private final T value; + private final Instant timestamp; + private final BoundedWindow window; + private final PaneInfo pane; + + BufferedElement(T value, Instant timestamp, BoundedWindow window, PaneInfo pane) { + this.value = value; + this.timestamp = timestamp; + this.window = window; + this.pane = pane; + } + } + + /** A user key scoped to a window and pane. */ + private static class WindowedKey { + private final BoundedWindow window; + private final PaneInfo pane; + private final Object key; + + WindowedKey(BoundedWindow window, PaneInfo pane, Object key) { + this.window = window; + this.pane = pane; + this.key = key; + } + + @Override + public boolean equals(@Nullable Object o) { + if (this == o) { + return true; + } + if (!(o instanceof WindowedKey)) { + return false; + } + WindowedKey other = (WindowedKey) o; + return window.equals(other.window) && pane.equals(other.pane) && key.equals(other.key); + } + + @Override + public int hashCode() { + return Objects.hash(window, pane, key); + } + } + + /** Sizes elements by their encoded size under the given coder. */ + private static class CoderSizer implements SerializableFunction { + private final Coder coder; + + CoderSizer(Coder coder) { + this.coder = coder; + } + + @Override + public Integer apply(T element) { + try { + long bytes; + if (coder.isRegisterByteSizeObserverCheap(element)) { + SizeObserver observer = new SizeObserver(); + coder.registerByteSizeObserver(element, observer); + observer.advance(); + bytes = observer.bytes; + } else { + bytes = CoderUtils.encodeToByteArray(coder, element).length; + } + return (int) Math.min(bytes, Integer.MAX_VALUE); + } catch (Exception e) { + throw new RuntimeException("Failed to estimate the encoded size of an element", e); } } } - private BundleLifter(TupleTag smallBatchTag, TupleTag largeBatchTag, int threshold) { - this(smallBatchTag, largeBatchTag, threshold, x -> 1); + private static class SizeObserver extends ElementByteSizeObserver { + private long bytes; + + @Override + protected void reportElementSize(long elementByteSize) { + bytes = elementByteSize; + } } private BundleLifter( TupleTag smallBatchTag, TupleTag largeBatchTag, - int threshold, - SerializableFunction elementSizer) { + long threshold, + SerializableFunction elementSizer, + @Nullable SerializableFunction keyFn) { if (smallBatchTag == null || largeBatchTag == null) { throw new IllegalArgumentException("smallBatchTag and largeBatchTag must not be null"); } @@ -154,19 +362,37 @@ private BundleLifter( this.largeBatchTag = largeBatchTag; this.threshold = threshold; this.elementSizer = elementSizer; + this.keyFn = keyFn; } + /** Routes by element count: each element counts as one toward the threshold. */ public static BundleLifter of( - TupleTag smallBatchTag, TupleTag largeBatchTag, int threshold) { - return new BundleLifter<>(smallBatchTag, largeBatchTag, threshold); + TupleTag smallBatchTag, TupleTag largeBatchTag, long threshold) { + return new BundleLifter<>(smallBatchTag, largeBatchTag, threshold, x -> 1, null); } + /** Routes by the sizes returned by {@code elementSizer}. */ public static BundleLifter of( TupleTag smallBatchTag, TupleTag largeBatchTag, - int threshold, + long threshold, SerializableFunction elementSizer) { - return new BundleLifter<>(smallBatchTag, largeBatchTag, threshold, elementSizer); + return new BundleLifter<>(smallBatchTag, largeBatchTag, threshold, elementSizer, null); + } + + /** Routes by each element's encoded size under {@code coder}. */ + public static BundleLifter of( + TupleTag smallBatchTag, TupleTag largeBatchTag, long threshold, Coder coder) { + return new BundleLifter<>( + smallBatchTag, largeBatchTag, threshold, new CoderSizer<>(coder), null); + } + + /** + * Returns a copy that evaluates the threshold separately for each window, pane and key returned + * by {@code keyFn}. Elements for which {@code keyFn} returns {@code null} go to the small tag. + */ + public BundleLifter withKeyFn(SerializableFunction keyFn) { + return new BundleLifter<>(smallBatchTag, largeBatchTag, threshold, elementSizer, keyFn); } @Override @@ -175,7 +401,7 @@ public PCollectionTuple expand(PCollection input) { return input.apply( "BundleLiftDoFn", - ParDo.of(new BundleLiftDoFn<>(smallBatchTag, largeBatchTag, threshold, elementSizer)) + ParDo.of(new BundleLiftDoFn<>(smallBatchTag, largeBatchTag, threshold, elementSizer, keyFn)) .withOutputTags(mainOutputTag, TupleTagList.of(smallBatchTag).and(largeBatchTag))); } } 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 3843efa0038d..29a9b5de3c6f 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 @@ -87,6 +87,7 @@ import org.apache.beam.sdk.transforms.MapElements; import org.apache.beam.sdk.transforms.PTransform; import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.SerializableFunction; import org.apache.beam.sdk.transforms.SimpleFunction; import org.apache.beam.sdk.transforms.Top; import org.apache.beam.sdk.transforms.Values; @@ -613,6 +614,106 @@ public void testDirectWriteDisabledWithSideInputs() throws IOException { assertEquals(0L, counterValue(result, "directWriteBundles")); } + /** Routes each element to a destination chosen by {@code destinationOf}, without side inputs. */ + private static class PerElementDestinations extends DynamicDestinations { + private final ResourceId baseOutputDirectory; + private final SerializableFunction destinationOf; + + PerElementDestinations( + ResourceId baseOutputDirectory, SerializableFunction destinationOf) { + this.baseOutputDirectory = baseOutputDirectory; + this.destinationOf = destinationOf; + } + + @Override + public String formatRecord(String record) { + return record; + } + + @Override + public String getDestination(String element) { + return destinationOf.apply(element); + } + + @Override + public String getDefaultDestination() { + return "file"; + } + + @Override + public FilenamePolicy getFilenamePolicy(String destination) { + return new PerWindowFiles( + baseOutputDirectory.resolve(destination, StandardResolveOptions.RESOLVE_FILE), "simple"); + } + } + + private PipelineResult runPerDestinationDirectWrite( + List elements, SerializableFunction destinationOf, long threshold) + throws IOException { + TestStream testStream = + TestStream.create(StringUtf8Coder.of()) + .advanceWatermarkTo(new Instant(0L)) + // A single bundle containing every element. + .addElements( + elements.get(0), + Iterables.toArray(elements.subList(1, elements.size()), String.class)) + .advanceWatermarkToInfinity(); + + SimpleSink sink = + new SimpleSink<>( + getBaseOutputDirectory(), + new PerElementDestinations(getBaseOutputDirectory(), destinationOf), + Compression.UNCOMPRESSED); + WriteFiles write = + WriteFiles.to(sink) + .withWindowedWrites() + .withNumShards(3) + .withDirectWriteByteThreshold(threshold); + p.apply(testStream) + .apply(Window.into(FixedWindows.of(Duration.standardSeconds(10)))) + .apply(write) + .getPerDestinationOutputFilenames() + .apply(new VerifyFilesExist<>()); + PipelineResult result = p.run(); + + checkFileContents(getBaseOutputFilename(), elements, Optional.absent(), false, null); + return result; + } + + @Test + @Category({NeedsRunner.class, UsesUnboundedPCollections.class, UsesTestStream.class}) + public void testDirectWriteThresholdIsPerDestination() throws IOException { + List elements = Lists.newArrayList(); + for (int i = 0; i < 30; ++i) { + elements.add("number: " + i); + } + // ~300 bytes in one bundle, but split over 10 destinations of ~30 bytes each, so no output + // file reaches the 50-byte threshold and everything takes the grouped path. + PipelineResult result = + runPerDestinationDirectWrite( + elements, e -> "file" + (Integer.parseInt(e.substring(8)) % 10), 50L); + + assertEquals(0L, counterValue(result, "directWriteBundles")); + assertThat(counterValue(result, "groupedWriteBundles"), Matchers.greaterThan(0L)); + } + + @Test + @Category({NeedsRunner.class, UsesUnboundedPCollections.class, UsesTestStream.class}) + public void testDirectWriteLargeAndSmallDestinationsInOneBundle() throws IOException { + List elements = Lists.newArrayList(); + for (int i = 0; i < 23; ++i) { + elements.add("number: " + i); + } + // 20 elements (~220 bytes) go to one destination and 3 (~33 bytes) to another. With a + // 100-byte threshold, the large destination is written directly and the small one is grouped. + PipelineResult result = + runPerDestinationDirectWrite( + elements, e -> Integer.parseInt(e.substring(8)) < 20 ? "fileLarge" : "fileSmall", 100L); + + assertThat(counterValue(result, "directWriteBundles"), Matchers.greaterThan(0L)); + assertThat(counterValue(result, "groupedWriteBundles"), Matchers.greaterThan(0L)); + } + private void runDirectWriteTestStream(long thresholdBytes) throws IOException { List elements = Lists.newArrayList(); for (int i = 0; i < 30; ++i) { diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/BundleLifterTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/BundleLifterTest.java index 0467f1383bcf..e4d3db9fecb4 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/BundleLifterTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/BundleLifterTest.java @@ -18,12 +18,21 @@ package org.apache.beam.sdk.util; import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.contains; import static org.hamcrest.Matchers.containsInAnyOrder; import static org.hamcrest.Matchers.empty; +import static org.hamcrest.Matchers.greaterThanOrEqualTo; +import static org.hamcrest.Matchers.hasSize; +import java.util.ArrayList; +import java.util.List; +import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.transforms.DoFnTester; +import org.apache.beam.sdk.transforms.windowing.IntervalWindow; import org.apache.beam.sdk.util.BundleLifter.BundleLiftDoFn; import org.apache.beam.sdk.values.TupleTag; +import org.apache.beam.sdk.values.ValueInSingleWindow; +import org.joda.time.Instant; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -96,4 +105,154 @@ public void testLargeBundleWithSizer() throws Exception { assertThat( tester.peekOutputElements(STRING_LARGE), containsInAnyOrder("123", "456", "789", "0")); } + + @Test + public void testStreamsElementsThroughOnceThresholdIsReached() throws Exception { + DoFnTester tester = + DoFnTester.of(new BundleLiftDoFn<>(INTEGER_SMALL, INTEGER_LARGE, 2, x -> 1)); + + tester.startBundle(); + tester.processElement(1); + assertThat(tester.peekOutputElements(INTEGER_LARGE), empty()); + tester.processElement(2); + // The buffered elements are flushed as soon as the threshold is reached... + assertThat(tester.peekOutputElements(INTEGER_LARGE), contains(1, 2)); + tester.processElement(3); + // ...and later elements are output directly, before the bundle finishes. + assertThat(tester.peekOutputElements(INTEGER_LARGE), contains(1, 2, 3)); + tester.finishBundle(); + + assertThat(tester.peekOutputElements(INTEGER_SMALL), empty()); + assertThat(tester.peekOutputElements(INTEGER_LARGE), contains(1, 2, 3)); + } + + @Test + public void testKeyedThresholdIsEvaluatedPerKey() throws Exception { + DoFnTester tester = + DoFnTester.of(new BundleLiftDoFn<>(INTEGER_SMALL, INTEGER_LARGE, 3, x -> 1, x -> x % 2)); + + tester.startBundle(); + for (int i = 0; i < 5; i++) { + tester.processElement(i); + } + tester.finishBundle(); + + // Even key has 3 elements and reaches the threshold; odd key has 2 and does not. + assertThat(tester.peekOutputElements(INTEGER_LARGE), containsInAnyOrder(0, 2, 4)); + assertThat(tester.peekOutputElements(INTEGER_SMALL), containsInAnyOrder(1, 3)); + } + + @Test + public void testKeyedThresholdIsEvaluatedPerWindow() throws Exception { + IntervalWindow w1 = new IntervalWindow(new Instant(0), new Instant(10)); + IntervalWindow w2 = new IntervalWindow(new Instant(10), new Instant(20)); + DoFnTester tester = + DoFnTester.of(new BundleLiftDoFn<>(INTEGER_SMALL, INTEGER_LARGE, 3, x -> 1, x -> "key")); + + tester.startBundle(); + tester.processWindowedElement(1, new Instant(1), w1); + tester.processWindowedElement(2, new Instant(2), w1); + tester.processWindowedElement(3, new Instant(11), w2); + tester.processWindowedElement(4, new Instant(12), w2); + tester.finishBundle(); + + // 4 elements share the key, but each window only has 2. + assertThat(tester.peekOutputElements(INTEGER_LARGE), empty()); + assertThat(tester.peekOutputElements(INTEGER_SMALL), containsInAnyOrder(1, 2, 3, 4)); + } + + @Test + public void testUnkeyedThresholdIsEvaluatedPerBundle() throws Exception { + IntervalWindow w1 = new IntervalWindow(new Instant(0), new Instant(10)); + IntervalWindow w2 = new IntervalWindow(new Instant(10), new Instant(20)); + DoFnTester tester = + DoFnTester.of(new BundleLiftDoFn<>(INTEGER_SMALL, INTEGER_LARGE, 3, x -> 1)); + + tester.startBundle(); + tester.processWindowedElement(1, new Instant(1), w1); + tester.processWindowedElement(2, new Instant(2), w1); + tester.processWindowedElement(3, new Instant(11), w2); + tester.processWindowedElement(4, new Instant(12), w2); + tester.finishBundle(); + + // Without a key function the whole bundle counts toward the threshold. + assertThat(tester.peekOutputElements(INTEGER_LARGE), containsInAnyOrder(1, 2, 3, 4)); + assertThat(tester.peekOutputElements(INTEGER_SMALL), empty()); + } + + @Test + public void testNullKeyGoesToSmallOutput() throws Exception { + DoFnTester tester = + DoFnTester.of( + new BundleLiftDoFn<>( + INTEGER_SMALL, INTEGER_LARGE, 1, x -> 1, x -> x < 0 ? null : "key")); + + tester.startBundle(); + tester.processElement(-1); + tester.processElement(5); + tester.finishBundle(); + + assertThat(tester.peekOutputElements(INTEGER_SMALL), contains(-1)); + assertThat(tester.peekOutputElements(INTEGER_LARGE), contains(5)); + } + + @Test + public void testMemoryCapEvictsSmallestKey() throws Exception { + // Threshold 4 caps total buffered bytes at 8. + DoFnTester tester = + DoFnTester.of(new BundleLiftDoFn<>(INTEGER_SMALL, INTEGER_LARGE, 4, x -> 1, x -> x)); + + tester.startBundle(); + for (int i = 0; i < 9; i++) { + tester.processElement(i); + } + // The 9th distinct key pushes the total over the cap, so one key is evicted early. + assertThat(tester.peekOutputElements(INTEGER_SMALL), hasSize(1)); + tester.finishBundle(); + + assertThat( + tester.peekOutputElements(INTEGER_SMALL), containsInAnyOrder(0, 1, 2, 3, 4, 5, 6, 7, 8)); + assertThat(tester.peekOutputElements(INTEGER_LARGE), empty()); + } + + @Test + public void testPreservesEachElementsWindowAndTimestamp() throws Exception { + IntervalWindow w1 = new IntervalWindow(new Instant(0), new Instant(10)); + IntervalWindow w2 = new IntervalWindow(new Instant(10), new Instant(20)); + + DoFnTester large = + DoFnTester.of(new BundleLiftDoFn<>(INTEGER_SMALL, INTEGER_LARGE, 2, x -> 1)); + large.startBundle(); + large.processWindowedElement(1, new Instant(5), w1); + large.processWindowedElement(2, new Instant(15), w2); + large.finishBundle(); + assertThat( + describe(large.getMutableOutput(INTEGER_LARGE)), + contains("1@" + new Instant(5) + " in " + w1, "2@" + new Instant(15) + " in " + w2)); + + DoFnTester small = + DoFnTester.of(new BundleLiftDoFn<>(INTEGER_SMALL, INTEGER_LARGE, 3, x -> 1)); + small.startBundle(); + small.processWindowedElement(1, new Instant(5), w1); + small.processWindowedElement(2, new Instant(15), w2); + small.finishBundle(); + assertThat( + describe(small.getMutableOutput(INTEGER_SMALL)), + contains("1@" + new Instant(5) + " in " + w1, "2@" + new Instant(15) + " in " + w2)); + } + + private static List describe(List> outputs) { + List described = new ArrayList<>(); + for (ValueInSingleWindow output : outputs) { + described.add(output.getValue() + "@" + output.getTimestamp() + " in " + output.getWindow()); + } + return described; + } + + @Test + public void testCoderBasedSizing() { + BundleLifter lifter = + BundleLifter.of(STRING_SMALL, STRING_LARGE, 10L, StringUtf8Coder.of()); + assertThat(lifter.elementSizer.apply("hello"), greaterThanOrEqualTo(5)); + } } From fec8a2e0c674b2e310d28b948d3338e270e46bf3 Mon Sep 17 00:00:00 2001 From: Tejas Iyer Date: Fri, 13 Nov 2026 07:13:25 +0000 Subject: [PATCH 6/6] Disabled Direct Write for merging windows. Fall back to default path with log warning if directWriteThreshold is set but direct write can't be used. Also evict BundleLifter keys in one pass when over the memory cap. --- .../org/apache/beam/sdk/io/WriteFiles.java | 60 ++++++++++++------- .../apache/beam/sdk/util/BundleLifter.java | 39 ++++++------ .../apache/beam/sdk/io/WriteFilesTest.java | 38 ++++++++++++ .../beam/sdk/util/BundleLifterTest.java | 5 +- 4 files changed, 103 insertions(+), 39 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 e1a74333daba..94f79893a55f 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 @@ -459,12 +459,15 @@ public WriteFiles withBatchMaxBufferingDuration( * Returns a new {@link WriteFiles} that writes large input bundles directly to temporary files, * skipping the shuffle before temporary file creation. * - *

Applies only to unbounded inputs, and is ignored when dynamic destinations use side inputs. - * Records of each bundle are buffered until their encoded size reaches {@code thresholdBytes}; - * bundles that reach it are written straight to temporary files. Smaller bundles keep the - * existing grouped path so output files stay large. Shard numbers are assigned at finalization, - * which may override {@link #withNumShards(int)} (the number of files per window can differ from - * the configured number of shards). + *

Records are buffered per window, pane and destination until their encoded size reaches + * {@code thresholdBytes}; those that reach it are written straight to temporary files. Smaller + * groups keep the existing grouped path so output files stay large. Shard numbers are assigned at + * finalization, which may override {@link #withNumShards(int)} (the number of files per window + * can differ from the configured number of shards). + * + *

Direct writes are only used for unbounded inputs with non-merging windows whose dynamic + * destinations don't use side inputs. Otherwise the threshold is ignored, a warning is logged, + * and the grouped write path is used. */ public WriteFiles withDirectWriteByteThreshold( long thresholdBytes) { @@ -611,20 +614,37 @@ public WriteFilesResult expand(PCollection input) { boolean fixedSharding = getComputeNumShards() != null || getNumShardsProvider() != null; Long directWriteThreshold = getDirectWriteByteThreshold(); - if (directWriteThreshold != null - && getSideInputs().isEmpty() - && input.isBounded() == IsBounded.UNBOUNDED) { - // Shard numbers are assigned in finalize across both the direct and grouped files, so any - // fixed shard count is not enforced there. - return input - .apply( - "WriteShardedBundlesWithDirectWrite", - new WriteShardedBundlesWithDirectWrite( - destinationCoder, fileResultCoder, numShardsView, directWriteThreshold)) - .apply( - "FinalizeTempFileBundles", - new FinalizeTempFileBundles( - null, destinationCoder, /* assignShardNumbersInFinalize= */ true)); + if (directWriteThreshold != null) { + List disabledReasons = new ArrayList<>(); + if (input.isBounded() == IsBounded.BOUNDED) { + disabledReasons.add("the input is bounded"); + } + if (!getSideInputs().isEmpty()) { + disabledReasons.add("dynamic destinations use side inputs"); + } + // Merging windows are only final after the GroupByKey, so direct writes would use + // unmerged windows. + if (input.getWindowingStrategy().needsMerge()) { + disabledReasons.add("the windowing strategy uses merging windows"); + } + if (disabledReasons.isEmpty()) { + // Shard numbers are assigned in finalize across both the direct and grouped files, so any + // fixed shard count is not enforced there. + return input + .apply( + "WriteShardedBundlesWithDirectWrite", + new WriteShardedBundlesWithDirectWrite( + destinationCoder, fileResultCoder, numShardsView, directWriteThreshold)) + .apply( + "FinalizeTempFileBundles", + new FinalizeTempFileBundles( + null, destinationCoder, /* assignShardNumbersInFinalize= */ true)); + } + LOG.warn( + "Direct write is disabled (directWriteByteThreshold={}) because {}. Falling back to the" + + " grouped write path.", + directWriteThreshold, + String.join(" and ", disabledReasons)); } PCollection>> tempFileResults; if (fixedSharding) { diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/BundleLifter.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/BundleLifter.java index 0edd929d0a57..3de2936673e1 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/BundleLifter.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/BundleLifter.java @@ -22,6 +22,7 @@ import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; import java.util.ArrayList; import java.util.Collections; +import java.util.Comparator; import java.util.HashSet; import java.util.LinkedHashMap; import java.util.List; @@ -63,7 +64,8 @@ * emitted to the small tag immediately. * *

To bound memory, at most twice the threshold is buffered in total. When that is exceeded, the - * key with the fewest buffered bytes is emitted to the small tag. + * keys with the fewest buffered bytes are emitted to the small tag until at most the threshold + * remains buffered. * *

Every element keeps its own timestamp and window. * @@ -108,7 +110,8 @@ static class BundleLiftDoFn extends DoFn { // Per-bundle state, reset in startBundle. private transient @MonotonicNonNull Map> buffers; private transient @MonotonicNonNull Set liftedKeys; - private transient long bufferedBytes; + // Reset in startBundle; not transient so SpotBugs doesn't flag it as unrestored state. + private long bufferedBytes; BundleLiftDoFn( TupleTag smallBatchTag, @@ -191,8 +194,8 @@ public void processElement( return; } - while (bufferedBytes > maxBufferedBytes && !buffers.isEmpty()) { - evictSmallestKey(buffers, receiver.get(smallBatchTag)); + if (bufferedBytes > maxBufferedBytes) { + evictSmallestKeys(buffers, receiver.get(smallBatchTag)); } } @@ -225,21 +228,23 @@ public Duration getAllowedTimestampSkew() { return Duration.millis(Long.MAX_VALUE); } - private void evictSmallestKey(Map> buffers, OutputReceiver small) { - Object smallestKey = null; - KeyBuffer smallest = null; - for (Map.Entry> entry : buffers.entrySet()) { - if (smallest == null || entry.getValue().sizeBytes < smallest.sizeBytes) { - smallestKey = entry.getKey(); - smallest = entry.getValue(); + /** + * Emits the smallest keys to the small output until at most {@code threshold} bytes remain + * buffered. Each call frees more than {@code threshold} bytes, so at least that many new bytes + * must be buffered before it runs again, which keeps the sort's cost amortized. + */ + private void evictSmallestKeys(Map> buffers, OutputReceiver small) { + List>> bySize = new ArrayList<>(buffers.entrySet()); + bySize.sort(Comparator.comparingLong(entry -> entry.getValue().sizeBytes)); + for (Map.Entry> entry : bySize) { + if (bufferedBytes <= threshold) { + break; } + KeyBuffer buffer = entry.getValue(); + buffers.remove(entry.getKey()); + bufferedBytes -= buffer.sizeBytes; + emit(buffer, small); } - if (smallestKey == null || smallest == null) { - return; - } - buffers.remove(smallestKey); - bufferedBytes -= smallest.sizeBytes; - emit(smallest, small); } private static void emit(KeyBuffer buffer, OutputReceiver receiver) { 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 29a9b5de3c6f..319fd3fd33f4 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 @@ -614,6 +614,44 @@ public void testDirectWriteDisabledWithSideInputs() throws IOException { assertEquals(0L, counterValue(result, "directWriteBundles")); } + @Test + @Category({NeedsRunner.class, UsesUnboundedPCollections.class, UsesTestStream.class}) + public void testDirectWriteDisabledWithMergingWindows() throws IOException { + List elements = Lists.newArrayList(); + for (int i = 0; i < 30; ++i) { + elements.add("number: " + i); + } + Instant startInstant = new Instant(0L); + TestStream testStream = + TestStream.create(StringUtf8Coder.of()) + .advanceWatermarkTo(startInstant) + .addElements(elements.get(0), Iterables.toArray(elements.subList(1, 10), String.class)) + .advanceWatermarkTo(startInstant.plus(Duration.standardSeconds(5))) + .addElements( + elements.get(10), Iterables.toArray(elements.subList(11, 20), String.class)) + .advanceWatermarkTo(startInstant.plus(Duration.standardSeconds(10))) + .addElements( + elements.get(20), Iterables.toArray(elements.subList(21, 30), String.class)) + .advanceWatermarkToInfinity(); + + // Merging windows require an explicit shard count, and a 1-byte threshold would otherwise + // write every bundle directly. + WriteFiles write = + WriteFiles.to(makeSimpleSink()) + .withWindowedWrites() + .withNumShards(3) + .withDirectWriteByteThreshold(1L); + p.apply(testStream) + .apply(Window.into(Sessions.withGapDuration(Duration.standardSeconds(1)))) + .apply(write) + .getPerDestinationOutputFilenames() + .apply(new VerifyFilesExist<>()); + PipelineResult result = p.run(); + + checkFileContents(getBaseOutputFilename(), elements, Optional.absent(), false, null); + assertEquals(0L, counterValue(result, "directWriteBundles")); + } + /** Routes each element to a destination chosen by {@code destinationOf}, without side inputs. */ private static class PerElementDestinations extends DynamicDestinations { private final ResourceId baseOutputDirectory; diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/BundleLifterTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/BundleLifterTest.java index e4d3db9fecb4..49772865b7ae 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/BundleLifterTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/BundleLifterTest.java @@ -206,8 +206,9 @@ public void testMemoryCapEvictsSmallestKey() throws Exception { for (int i = 0; i < 9; i++) { tester.processElement(i); } - // The 9th distinct key pushes the total over the cap, so one key is evicted early. - assertThat(tester.peekOutputElements(INTEGER_SMALL), hasSize(1)); + // The 9th distinct key pushes the total (9) over the cap (8), so the smallest keys are evicted + // in one pass until at most the threshold (4) remains buffered: 5 keys go out early. + assertThat(tester.peekOutputElements(INTEGER_SMALL), hasSize(5)); tester.finishBundle(); assertThat(