ahmedabu98 commented on code in PR #40543: URL: https://github.com/apache/beam/pull/40543#discussion_r4330999923
########## 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)). Review Comment: Not seeing the`streaming_file_direct_write_threshold_bytes` flag in the code here. is it in the Dataflow service? ########## 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)). Review Comment: Remove irrelevant line? Maybe it's getting accidentally picked up from master ########## sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java: ########## @@ -1118,6 +1168,367 @@ public PCollection<FileResult<DestinationT>> expand(PCollection<UserT> input) { } } + /** + * Streaming write that sends large bundles straight to temporary files. + * + * <p>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<UserT>, PCollection<List<FileResult<DestinationT>>>> { + private final Coder<DestinationT> destinationCoder; + private final Coder<FileResult<DestinationT>> fileResultCoder; + private final @Nullable PCollectionView<Integer> numShardsView; + private final long thresholdBytes; + + private WriteShardedBundlesWithDirectWrite( + Coder<DestinationT> destinationCoder, + Coder<FileResult<DestinationT>> fileResultCoder, + @Nullable PCollectionView<Integer> numShardsView, + long thresholdBytes) { + this.destinationCoder = destinationCoder; + this.fileResultCoder = fileResultCoder; + this.numShardsView = numShardsView; + this.thresholdBytes = thresholdBytes; + } + + @Override + public PCollection<List<FileResult<DestinationT>>> expand(PCollection<UserT> input) { + TupleTag<FileResult<DestinationT>> directFilesTag = new TupleTag<>("directFiles"); + TupleTag<UserT> 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<UserT> smallBundles = routed.get(smallBundlesTag).setCoder(input.getCoder()); Review Comment: I see you did a good amount for that using `writeAutoShardedRecordsToTempFiles(...)`. Wondering if other parts can be similarly modularized ########## sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java: ########## @@ -1118,6 +1168,367 @@ public PCollection<FileResult<DestinationT>> expand(PCollection<UserT> input) { } } + /** + * Streaming write that sends large bundles straight to temporary files. + * + * <p>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<UserT>, PCollection<List<FileResult<DestinationT>>>> { + private final Coder<DestinationT> destinationCoder; + private final Coder<FileResult<DestinationT>> fileResultCoder; + private final @Nullable PCollectionView<Integer> numShardsView; + private final long thresholdBytes; + + private WriteShardedBundlesWithDirectWrite( + Coder<DestinationT> destinationCoder, + Coder<FileResult<DestinationT>> fileResultCoder, + @Nullable PCollectionView<Integer> numShardsView, + long thresholdBytes) { + this.destinationCoder = destinationCoder; + this.fileResultCoder = fileResultCoder; + this.numShardsView = numShardsView; + this.thresholdBytes = thresholdBytes; + } + + @Override + public PCollection<List<FileResult<DestinationT>>> expand(PCollection<UserT> input) { + TupleTag<FileResult<DestinationT>> directFilesTag = new TupleTag<>("directFiles"); + TupleTag<UserT> 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<UserT> smallBundles = routed.get(smallBundlesTag).setCoder(input.getCoder()); Review Comment: Is there a way to reuse the existing pipeline graph (built in `WriteFilesResult`'s expand) and pass `smallBundles` through to it? Instead of attempting to recreate the same steps here? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
