arunpandianp commented on code in PR #40543:
URL: https://github.com/apache/beam/pull/40543#discussion_r4331147305
##########
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());
+ 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,
+ 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<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.
+ return PCollectionList.of(directFiles)
+ .and(groupedFiles)
+ .apply("FlattenTempFiles", Flatten.pCollections())
+ .setCoder(fileResultCoder)
+ .apply(
+ "KeyedByDestination",
+ WithKeys.of(
+ new SimpleFunction<FileResult<DestinationT>, DestinationT>()
{
+ @Override
+ public DestinationT apply(FileResult<DestinationT> input) {
+ return input.getDestination();
+ }
+ }))
+ .setCoder(KvCoder.of(destinationCoder, fileResultCoder))
+ .apply("GroupTempFilesByDestination", GroupByKey.create())
+ .apply(
+ "ExtractValuesToList",
+ ParDo.of(
+ new DoFn<
+ KV<DestinationT, Iterable<FileResult<DestinationT>>>,
+ List<FileResult<DestinationT>>>() {
+ @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<FileResult<DestinationT>> result = new
ArrayList<>();
+ for (FileResult<DestinationT> 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 grouped path.
+ */
+ private class DirectWriteLargeBundlesFn extends DoFn<UserT,
FileResult<DestinationT>> {
+ private final TupleTag<UserT> smallBundlesTag;
+ private final long thresholdBytes;
+ private final Coder<UserT> 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<BufferedRecord<UserT>> buffer;
+ private transient long bufferedBytes;
+ private transient boolean directMode;
+ private transient Map<WriterKey<DestinationT>, Writer<DestinationT,
OutputT>> writers;
+ private transient List<BufferedRecord<UserT>> spilled;
+
+ DirectWriteLargeBundlesFn(
+ TupleTag<UserT> smallBundlesTag, long thresholdBytes, Coder<UserT>
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<UserT> 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<UserT> 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<UserT> record, MultiOutputReceiver
outputReceiver)
+ throws Exception {
+ MaybeDestination<DestinationT> maybeDestination =
+ getDestinationWithErrorHandling(record.value, outputReceiver,
inputCoder);
+ if (!maybeDestination.isValid) {
+ return;
+ }
+ DestinationT destination = maybeDestination.destination;
+ WriterKey<DestinationT> key = new WriterKey<>(record.window,
record.pane, destination);
+ Writer<DestinationT, OutputT> 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<UserT> record : buffer) {
+ c.output(smallBundlesTag, record.value, record.timestamp,
record.window);
+ }
+ for (BufferedRecord<UserT> record : spilled) {
+ c.output(smallBundlesTag, record.value, record.timestamp,
record.window);
Review Comment:
seems like this will reorder outputs. Can we preserve the order?
--
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]
