tejasiyer-dev commented on code in PR #40543:
URL: https://github.com/apache/beam/pull/40543#discussion_r4569336960
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java:
##########
@@ -1114,6 +1202,318 @@ 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<UserT> smallBundlesTag = new TupleTag<>("smallBundles");
+ TupleTag<UserT> largeBundlesTag = new TupleTag<>("largeBundles");
+ TupleTag<FileResult<DestinationT>> directFilesTag = new
TupleTag<>("directFiles");
+ TupleTag<UserT> spilledRecordsTag = new TupleTag<>("spilledRecords");
+
+ // Evaluate the threshold per window, pane and destination, i.e. per
output file.
+ PCollectionTuple lifted =
+ input.apply(
+ "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<UserT> 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<FileResult<DestinationT>> groupedFiles;
+ if (getComputeNumShards() != null || getNumShardsProvider() != null) {
+ groupedFiles =
+ smallBundles
+ .apply(
+ "WriteShardedBundlesToTempFiles",
+ new WriteShardedBundlesToTempFiles(
+ destinationCoder, fileResultCoder, numShardsView))
+ .apply(
+ "DropShardNum",
+ ParDo.of(
+ new DoFn<FileResult<DestinationT>,
FileResult<DestinationT>>() {
+ @ProcessElement
+ public void process(ProcessContext c) {
+ c.output(c.element().withShard(UNKNOWN_SHARDNUM));
+ }
+ }))
+ .setCoder(fileResultCoder);
+ } else {
+ groupedFiles =
+ writeAutoShardedRecordsToTempFiles(
+ smallBundles,
+ destinationCoder,
+ fileResultCoder,
+ getBatchSizeOrDefault(),
+ getBatchSizeBytesOrDefault(),
+ getBatchMaxBufferingDurationOrDefault());
+ }
+
+ // 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<FileResult<DestinationT>> 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. Every file is
stamped with this
+ // firing's pane so files from different firings of the same window get
distinct names.
+ PCollection<FileResult<DestinationT>> allTempFiles =
+ PCollectionList.of(directFiles)
+ .and(groupedFiles)
+ .apply("FlattenTempFiles", Flatten.pCollections())
+ .setCoder(fileResultCoder);
+ return groupTempFilesByDestination(
+ allTempFiles,
+ destinationCoder,
+ fileResultCoder,
+ "GroupTempFilesByDestination",
+ /* stampFiringPane= */ true);
+ }
+ }
+
+ /**
+ * 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<UserT, DestinationT>
+ implements SerializableFunction<UserT, @Nullable Object> {
+ private final DynamicDestinations<UserT, DestinationT, ?>
dynamicDestinations;
+
+ DestinationKeyFn(DynamicDestinations<UserT, DestinationT, ?>
dynamicDestinations) {
+ this.dynamicDestinations = dynamicDestinations;
+ }
+
+ @Override
+ 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<T> extends DoFn<T, T> {
+ 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<T> receiver)
{
+ if (!counted) {
+ groupedBundles.inc();
+ counted = true;
+ }
+ receiver.output(element);
+ }
+ }
+
+ /**
+ * 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 DirectWriteRecordsFn extends DoFn<UserT,
FileResult<DestinationT>> {
+ private final TupleTag<UserT> spilledRecordsTag;
+ private final Coder<UserT> inputCoder;
+
+ private final Counter directBundles = Metrics.counter(WriteFiles.class,
"directWriteBundles");
Review Comment:
I was thinking to keep directWriteBundles because it lives in the existing
write DoFn, so it adds no extra step to the pipeline. Also the tests use it to
verify the direct path actually ran, since the DirectRunner doesn't expose
pcollection metrics.
Happy to remove it if you'd prefer.
--
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]
