tejasiyer-dev commented on code in PR #40543:
URL: https://github.com/apache/beam/pull/40543#discussion_r4450392511
##########
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:
Partly. The small bundles now reuse the existing write steps, and both paths
share the grouping steps and batch-size defaults through new helpers.
The small bundles can't go through the entire existing graph, though. If
they did, the direct-write files and the small-bundle files would each go
through a separate finalize step, and their file names would collide. Instead,
both sets of temp files are merged before finalize, so a single finalize step
assigns unique shard numbers to every file in the window.
--
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]
