[WriteFiles] Support direct writes of large streaming bundles to temp files (#40533)

CHANGES.md

@@ -73,6 +73,8 @@
 * 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)).
 * (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)).
 

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<UserT, DestinationT, OutputT>
     // 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<UserT, DestinationT, OutputT>
           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,24 @@ static class StreamingShardedWriteFactory<UserT, DestinationT, OutputT>
         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 {@code null}.
+     */
+    private @Nullable Long getDirectWriteThreshold(AppliedPTransform<?, ?, ?> transform) {
+      Object original = transform.getTransform();
+      if (original instanceof WriteFiles) {
+        return ((WriteFiles<?, ?, ?>) original).getDirectWriteByteThreshold();
+      }
+      return null;
+    }
+
     @Override
     public Map<PCollection<?>, ReplacementOutput> mapOutputs(
         Map<TupleTag<?>, PCollection<?>> outputs, WriteFilesResult<DestinationT> newOutput) {

runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java

@@ -2149,9 +2149,66 @@ public void testStreamingWriteWithShardingReturnsSameTransform() {
 
     assertNull(replacement.getNumShardsProvider());
     assertNull(replacement.getComputeNumShards());
+    assertNull(replacement.getDirectWriteByteThreshold());
     assertTrue(replacement.getWithAutoSharding());
   }
 
+  @Test
+  public void testStreamingWriteWithDirectWriteByteThreshold() {
+    PipelineOptions options = TestPipeline.testingPipelineOptions();
+
+    TestPipeline p = TestPipeline.fromOptions(options);
+    StreamingShardedWriteFactory<Object, Void, Object> factory =
+        new StreamingShardedWriteFactory<>(p.getOptions());
+    WriteFiles<Object, Void, Object> original =
+        WriteFiles.to(new TestSink(tmpFolder.toString()))
+            .withRunnerDeterminedSharding()
+            .withDirectWriteByteThreshold(1000L);
+    PCollection<Object> objs = (PCollection) p.apply(Create.empty(VoidCoder.of()));
+    AppliedPTransform<PCollection<Object>, WriteFilesResult<Void>, WriteFiles<Object, Void, Object>>
+        originalApplication =
+            AppliedPTransform.of(
+                "writefiles",
+                PValues.expandInput(objs),
+                Collections.emptyMap(),
+                original,
+                ResourceHints.create(),
+                p);
+
+    WriteFiles<Object, Void, Object> replacement =
+        (WriteFiles<Object, Void, Object>)
+            factory.getReplacementTransform(originalApplication).getTransform();
+
+    assertEquals(Long.valueOf(1000L), replacement.getDirectWriteByteThreshold());
+  }
+
+  @Test
+  public void testStreamingWriteDirectWriteThresholdDefaultNull() {
+    PipelineOptions options = TestPipeline.testingPipelineOptions();
+
+    TestPipeline p = TestPipeline.fromOptions(options);
+    StreamingShardedWriteFactory<Object, Void, Object> factory =
+        new StreamingShardedWriteFactory<>(p.getOptions());
+    WriteFiles<Object, Void, Object> original =
+        WriteFiles.to(new TestSink(tmpFolder.toString())).withRunnerDeterminedSharding();
+    PCollection<Object> objs = (PCollection) p.apply(Create.empty(VoidCoder.of()));
+    AppliedPTransform<PCollection<Object>, WriteFilesResult<Void>, WriteFiles<Object, Void, Object>>
+        originalApplication =
+            AppliedPTransform.of(
+                "writefiles",
+                PValues.expandInput(objs),
+                Collections.emptyMap(),
+                original,
+                ResourceHints.create(),
+                p);
+
+    WriteFiles<Object, Void, Object> replacement =
+        (WriteFiles<Object, Void, Object>)
+            factory.getReplacementTransform(originalApplication).getTransform();
+
+    assertNull(replacement.getDirectWriteByteThreshold());
+  }
+
   private void verifyMergingStatefulParDoRejected(PipelineOptions options) throws Exception {
     Pipeline p = Pipeline.create(options);
 

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<BadRecord, ?> getBadRecordErrorHandler();
 
     abstract Builder<DestinationT, UserT> toBuilder();
@@ -1133,6 +1135,9 @@ abstract Builder<DestinationT, UserT> setMaxNumWritersPerBundle(
       abstract Builder<DestinationT, UserT> setBatchMaxBufferingDuration(
           @Nullable Duration batchMaxBufferingDuration);
 
+      abstract Builder<DestinationT, UserT> setDirectWriteByteThreshold(
+          @Nullable Long directWriteByteThreshold);
+
       abstract Builder<DestinationT, UserT> setBadRecordErrorHandler(
           @Nullable ErrorHandler<BadRecord, ?> badRecordErrorHandler);
 
@@ -1436,6 +1441,13 @@ public Write<DestinationT, UserT> withBatchMaxBufferingDuration(
       return toBuilder().setBatchMaxBufferingDuration(batchMaxBufferingDuration).build();
     }
 
+    /** See {@link WriteFiles#withDirectWriteByteThreshold(long)}. */
+    public Write<DestinationT, UserT> withDirectWriteByteThreshold(long thresholdBytes) {
+      checkArgument(
+          thresholdBytes > 0, "thresholdBytes must be positive, but was: %s", thresholdBytes);
+      return toBuilder().setDirectWriteByteThreshold(thresholdBytes).build();
+    }
+
     @VisibleForTesting
     Contextful<Fn<DestinationT, FileNaming>> resolveFileNamingFn() {
       if (getDynamic()) {
@@ -1565,6 +1577,9 @@ public WriteFilesResult<DestinationT> expand(PCollection<UserT> input) {
       if (getBatchMaxBufferingDuration() != null) {
         writeFiles = writeFiles.withBatchMaxBufferingDuration(getBatchMaxBufferingDuration());
       }
+      if (getDirectWriteByteThreshold() != null) {
+        writeFiles = writeFiles.withDirectWriteByteThreshold(getDirectWriteByteThreshold());
+      }
       return input.apply(writeFiles);
     }
 

sdks/java/core/src/main/java/org/apache/beam/sdk/io/TextIO.java

@@ -713,6 +713,9 @@ public abstract static class TypedWrite<UserT, DestinationT>
     /** 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<UserT, DestinationT> setNumShards(
       abstract Builder<UserT, DestinationT> setBatchMaxBufferingDuration(
           @Nullable Duration batchMaxBufferingDuration);
 
+      abstract Builder<UserT, DestinationT> setDirectWriteByteThreshold(
+          @Nullable Long directWriteByteThreshold);
+
       abstract Builder<UserT, DestinationT> setWindowedWrites(boolean windowedWrites);
 
       abstract Builder<UserT, DestinationT> setAutoSharding(boolean windowedWrites);
@@ -926,6 +932,13 @@ public TypedWrite<UserT, DestinationT> withBatchMaxBufferingDuration(
       return toBuilder().setBatchMaxBufferingDuration(batchMaxBufferingDuration).build();
     }
 
+    /** See {@link WriteFiles#withDirectWriteByteThreshold(long)}. */
+    public TypedWrite<UserT, DestinationT> 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<UserT, DestinationT> withTempDirectory(
         ValueProvider<ResourceId> tempDirectory) {
@@ -1195,6 +1208,9 @@ public WriteFilesResult<DestinationT> expand(PCollection<UserT> 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.
      *