scwhittle commented on code in PR #40543:
URL: https://github.com/apache/beam/pull/40543#discussion_r4478299851
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java:
##########
@@ -1114,6 +1181,342 @@ 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,
+ 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);
+ }
+ }
+
+ 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() {
Review Comment:
if we aren't leaking the references can we just clear objects instead of
reallocating?
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java:
##########
@@ -441,6 +455,26 @@ public WriteFiles<UserT, DestinationT, OutputT>
withBatchMaxBufferingDuration(
return
toBuilder().setBatchMaxBufferingDuration(batchMaxBufferingDuration).build();
}
+ /**
+ * Returns a new {@link WriteFiles} that writes large input bundles directly
to temporary files,
Review Comment:
that fuses with proceeding stages and writes large inputs from a single
bundle directly to temporary files
##########
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");
Review Comment:
how interesting is this? can this just be observed through the existing
pcollection metrics?
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/util/BundleLifter.java:
##########
@@ -57,89 +84,275 @@ public class BundleLifter<T> extends
PTransform<PCollection<T>, PCollectionTuple
final TupleTag<T> smallBatchTag;
final TupleTag<T> largeBatchTag;
- final int threshold;
+ final long threshold;
final SerializableFunction<T, Integer> elementSizer;
+ final @Nullable SerializableFunction<T, ? extends @Nullable Object> keyFn;
/**
- * A DoFn that buffers elements within a bundle and outputs them to
different tags in
- * finish_bundle based on the total bundle size.
+ * A DoFn that buffers elements per key within a bundle and outputs them to
the large tag once a
+ * key reaches the threshold, or to the small tag otherwise.
*
* @param <T> The type of elements being processed.
*/
static class BundleLiftDoFn<T> extends DoFn<T, Void> {
private static final Logger LOG =
LoggerFactory.getLogger(BundleLiftDoFn.class);
+ /** The key shared by all elements when no key function is set. */
+ private static final Object UNKEYED = new Object();
+
final TupleTag<T> smallBatchTag;
final TupleTag<T> largeBatchTag;
- final int threshold;
+ final long threshold;
+ final long maxBufferedBytes;
final SerializableFunction<T, Integer> elementSizer;
+ final @Nullable SerializableFunction<T, ? extends @Nullable Object> keyFn;
- private transient @MonotonicNonNull List<T> buffer;
- private transient long bundleSizeBytes;
- private transient @Nullable MultiOutputReceiver receiver;
+ // Per-bundle state, reset in startBundle.
+ private transient @MonotonicNonNull Map<Object, KeyBuffer<T>> buffers;
+ private transient @MonotonicNonNull Set<Object> liftedKeys;
+ // Reset in startBundle; not transient so SpotBugs doesn't flag it as
unrestored state.
+ private long bufferedBytes;
BundleLiftDoFn(
TupleTag<T> smallBatchTag,
TupleTag<T> largeBatchTag,
- int threshold,
+ long threshold,
SerializableFunction<T, Integer> elementSizer) {
+ this(smallBatchTag, largeBatchTag, threshold, elementSizer, null);
+ }
+
+ BundleLiftDoFn(
+ TupleTag<T> smallBatchTag,
+ TupleTag<T> largeBatchTag,
+ long threshold,
+ SerializableFunction<T, Integer> elementSizer,
+ @Nullable SerializableFunction<T, ? extends @Nullable Object> keyFn) {
this.smallBatchTag = smallBatchTag;
this.largeBatchTag = largeBatchTag;
this.threshold = threshold;
+ this.maxBufferedBytes = threshold > Long.MAX_VALUE / 2 ? Long.MAX_VALUE
: 2 * threshold;
Review Comment:
should we allow the caller to specify this?
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/util/BundleLifter.java:
##########
@@ -57,89 +84,275 @@ public class BundleLifter<T> extends
PTransform<PCollection<T>, PCollectionTuple
final TupleTag<T> smallBatchTag;
final TupleTag<T> largeBatchTag;
- final int threshold;
+ final long threshold;
final SerializableFunction<T, Integer> elementSizer;
+ final @Nullable SerializableFunction<T, ? extends @Nullable Object> keyFn;
/**
- * A DoFn that buffers elements within a bundle and outputs them to
different tags in
- * finish_bundle based on the total bundle size.
+ * A DoFn that buffers elements per key within a bundle and outputs them to
the large tag once a
+ * key reaches the threshold, or to the small tag otherwise.
*
* @param <T> The type of elements being processed.
*/
static class BundleLiftDoFn<T> extends DoFn<T, Void> {
private static final Logger LOG =
LoggerFactory.getLogger(BundleLiftDoFn.class);
+ /** The key shared by all elements when no key function is set. */
+ private static final Object UNKEYED = new Object();
+
final TupleTag<T> smallBatchTag;
final TupleTag<T> largeBatchTag;
- final int threshold;
+ final long threshold;
+ final long maxBufferedBytes;
final SerializableFunction<T, Integer> elementSizer;
+ final @Nullable SerializableFunction<T, ? extends @Nullable Object> keyFn;
- private transient @MonotonicNonNull List<T> buffer;
- private transient long bundleSizeBytes;
- private transient @Nullable MultiOutputReceiver receiver;
+ // Per-bundle state, reset in startBundle.
+ private transient @MonotonicNonNull Map<Object, KeyBuffer<T>> buffers;
+ private transient @MonotonicNonNull Set<Object> liftedKeys;
+ // Reset in startBundle; not transient so SpotBugs doesn't flag it as
unrestored state.
+ private long bufferedBytes;
BundleLiftDoFn(
TupleTag<T> smallBatchTag,
TupleTag<T> largeBatchTag,
- int threshold,
+ long threshold,
SerializableFunction<T, Integer> elementSizer) {
+ this(smallBatchTag, largeBatchTag, threshold, elementSizer, null);
+ }
+
+ BundleLiftDoFn(
+ TupleTag<T> smallBatchTag,
+ TupleTag<T> largeBatchTag,
+ long threshold,
+ SerializableFunction<T, Integer> elementSizer,
+ @Nullable SerializableFunction<T, ? extends @Nullable Object> keyFn) {
this.smallBatchTag = smallBatchTag;
this.largeBatchTag = largeBatchTag;
this.threshold = threshold;
+ this.maxBufferedBytes = threshold > Long.MAX_VALUE / 2 ? Long.MAX_VALUE
: 2 * threshold;
this.elementSizer = elementSizer;
+ this.keyFn = keyFn;
}
@StartBundle
public void startBundle() {
- buffer = new ArrayList<>();
- receiver = null;
- bundleSizeBytes = 0L;
+ buffers = new LinkedHashMap<>();
+ liftedKeys = new HashSet<>();
+ bufferedBytes = 0L;
}
@ProcessElement
- public void processElement(@Element T element, MultiOutputReceiver mor) {
- if (receiver == null) {
- receiver = mor;
+ public void processElement(
+ @Element T element,
+ @Timestamp Instant timestamp,
+ BoundedWindow window,
+ PaneInfo pane,
+ MultiOutputReceiver receiver) {
+ Map<Object, KeyBuffer<T>> buffers =
+ checkStateNotNull(this.buffers, "Buffers should be set by
startBundle.");
+ Set<Object> liftedKeys =
+ checkStateNotNull(this.liftedKeys, "Lifted keys should be set by
startBundle.");
+
+ Object key;
+ if (keyFn == null) {
+ key = UNKEYED;
+ } else {
+ @Nullable Object userKey = keyFn.apply(element);
+ if (userKey == null) {
+ receiver.get(smallBatchTag).output(element);
+ return;
+ }
+ key = new WindowedKey(window, pane, userKey);
+ }
+
+ if (liftedKeys.contains(key)) {
+ // This key already reached the threshold in this bundle; stream it
through.
+ receiver.get(largeBatchTag).output(element);
+ return;
+ }
+
+ KeyBuffer<T> buffer = buffers.get(key);
+ if (buffer == null) {
+ buffer = new KeyBuffer<>();
+ buffers.put(key, buffer);
+ }
+ long elementBytes = elementSizer.apply(element);
+ buffer.add(new BufferedElement<>(element, timestamp, window, pane),
elementBytes);
+ bufferedBytes += elementBytes;
+
+ if (buffer.sizeBytes >= threshold) {
+ buffers.remove(key);
+ bufferedBytes -= buffer.sizeBytes;
+ liftedKeys.add(key);
Review Comment:
liftedkeys is going to grow here.
What about instead transitioning to a finalized keys state tracked by a
boolean. Anything that isn't in lifted keys or buffers is sent to small keys
immediately.
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/util/BundleLifter.java:
##########
@@ -57,89 +84,275 @@ public class BundleLifter<T> extends
PTransform<PCollection<T>, PCollectionTuple
final TupleTag<T> smallBatchTag;
final TupleTag<T> largeBatchTag;
- final int threshold;
+ final long threshold;
final SerializableFunction<T, Integer> elementSizer;
+ final @Nullable SerializableFunction<T, ? extends @Nullable Object> keyFn;
/**
- * A DoFn that buffers elements within a bundle and outputs them to
different tags in
- * finish_bundle based on the total bundle size.
+ * A DoFn that buffers elements per key within a bundle and outputs them to
the large tag once a
+ * key reaches the threshold, or to the small tag otherwise.
*
* @param <T> The type of elements being processed.
*/
static class BundleLiftDoFn<T> extends DoFn<T, Void> {
private static final Logger LOG =
LoggerFactory.getLogger(BundleLiftDoFn.class);
+ /** The key shared by all elements when no key function is set. */
+ private static final Object UNKEYED = new Object();
+
final TupleTag<T> smallBatchTag;
final TupleTag<T> largeBatchTag;
- final int threshold;
+ final long threshold;
+ final long maxBufferedBytes;
final SerializableFunction<T, Integer> elementSizer;
+ final @Nullable SerializableFunction<T, ? extends @Nullable Object> keyFn;
- private transient @MonotonicNonNull List<T> buffer;
- private transient long bundleSizeBytes;
- private transient @Nullable MultiOutputReceiver receiver;
+ // Per-bundle state, reset in startBundle.
+ private transient @MonotonicNonNull Map<Object, KeyBuffer<T>> buffers;
+ private transient @MonotonicNonNull Set<Object> liftedKeys;
+ // Reset in startBundle; not transient so SpotBugs doesn't flag it as
unrestored state.
+ private long bufferedBytes;
BundleLiftDoFn(
TupleTag<T> smallBatchTag,
TupleTag<T> largeBatchTag,
- int threshold,
+ long threshold,
SerializableFunction<T, Integer> elementSizer) {
+ this(smallBatchTag, largeBatchTag, threshold, elementSizer, null);
+ }
+
+ BundleLiftDoFn(
+ TupleTag<T> smallBatchTag,
+ TupleTag<T> largeBatchTag,
+ long threshold,
+ SerializableFunction<T, Integer> elementSizer,
+ @Nullable SerializableFunction<T, ? extends @Nullable Object> keyFn) {
this.smallBatchTag = smallBatchTag;
this.largeBatchTag = largeBatchTag;
this.threshold = threshold;
+ this.maxBufferedBytes = threshold > Long.MAX_VALUE / 2 ? Long.MAX_VALUE
: 2 * threshold;
this.elementSizer = elementSizer;
+ this.keyFn = keyFn;
}
@StartBundle
public void startBundle() {
- buffer = new ArrayList<>();
- receiver = null;
- bundleSizeBytes = 0L;
+ buffers = new LinkedHashMap<>();
Review Comment:
should we prefer clearing these instead of resetting if we're not leaking
them out of this dofn?
##########
sdks/java/core/src/main/java/org/apache/beam/sdk/util/BundleLifter.java:
##########
@@ -57,89 +84,275 @@ public class BundleLifter<T> extends
PTransform<PCollection<T>, PCollectionTuple
final TupleTag<T> smallBatchTag;
final TupleTag<T> largeBatchTag;
- final int threshold;
+ final long threshold;
final SerializableFunction<T, Integer> elementSizer;
+ final @Nullable SerializableFunction<T, ? extends @Nullable Object> keyFn;
/**
- * A DoFn that buffers elements within a bundle and outputs them to
different tags in
- * finish_bundle based on the total bundle size.
+ * A DoFn that buffers elements per key within a bundle and outputs them to
the large tag once a
+ * key reaches the threshold, or to the small tag otherwise.
*
* @param <T> The type of elements being processed.
*/
static class BundleLiftDoFn<T> extends DoFn<T, Void> {
private static final Logger LOG =
LoggerFactory.getLogger(BundleLiftDoFn.class);
+ /** The key shared by all elements when no key function is set. */
+ private static final Object UNKEYED = new Object();
+
final TupleTag<T> smallBatchTag;
final TupleTag<T> largeBatchTag;
- final int threshold;
+ final long threshold;
+ final long maxBufferedBytes;
final SerializableFunction<T, Integer> elementSizer;
+ final @Nullable SerializableFunction<T, ? extends @Nullable Object> keyFn;
- private transient @MonotonicNonNull List<T> buffer;
- private transient long bundleSizeBytes;
- private transient @Nullable MultiOutputReceiver receiver;
+ // Per-bundle state, reset in startBundle.
+ private transient @MonotonicNonNull Map<Object, KeyBuffer<T>> buffers;
+ private transient @MonotonicNonNull Set<Object> liftedKeys;
+ // Reset in startBundle; not transient so SpotBugs doesn't flag it as
unrestored state.
+ private long bufferedBytes;
BundleLiftDoFn(
TupleTag<T> smallBatchTag,
TupleTag<T> largeBatchTag,
- int threshold,
+ long threshold,
SerializableFunction<T, Integer> elementSizer) {
+ this(smallBatchTag, largeBatchTag, threshold, elementSizer, null);
+ }
+
+ BundleLiftDoFn(
+ TupleTag<T> smallBatchTag,
+ TupleTag<T> largeBatchTag,
+ long threshold,
+ SerializableFunction<T, Integer> elementSizer,
+ @Nullable SerializableFunction<T, ? extends @Nullable Object> keyFn) {
this.smallBatchTag = smallBatchTag;
this.largeBatchTag = largeBatchTag;
this.threshold = threshold;
+ this.maxBufferedBytes = threshold > Long.MAX_VALUE / 2 ? Long.MAX_VALUE
: 2 * threshold;
this.elementSizer = elementSizer;
+ this.keyFn = keyFn;
}
@StartBundle
public void startBundle() {
- buffer = new ArrayList<>();
- receiver = null;
- bundleSizeBytes = 0L;
+ buffers = new LinkedHashMap<>();
+ liftedKeys = new HashSet<>();
+ bufferedBytes = 0L;
}
@ProcessElement
- public void processElement(@Element T element, MultiOutputReceiver mor) {
- if (receiver == null) {
- receiver = mor;
+ public void processElement(
+ @Element T element,
+ @Timestamp Instant timestamp,
+ BoundedWindow window,
+ PaneInfo pane,
+ MultiOutputReceiver receiver) {
+ Map<Object, KeyBuffer<T>> buffers =
+ checkStateNotNull(this.buffers, "Buffers should be set by
startBundle.");
+ Set<Object> liftedKeys =
+ checkStateNotNull(this.liftedKeys, "Lifted keys should be set by
startBundle.");
+
+ Object key;
+ if (keyFn == null) {
+ key = UNKEYED;
+ } else {
+ @Nullable Object userKey = keyFn.apply(element);
+ if (userKey == null) {
Review Comment:
should we use key = UNKEYED in this case instead?
##########
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:
ditto, do we need this beyond the existing pcollection metrics?
--
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]
