tejasiyer-dev commented on code in PR #40543:
URL: https://github.com/apache/beam/pull/40543#discussion_r4569492786
##########
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:
I added a keysFinalized boolean. Once the memory cap is first exceeded, any
key that isn't already lifted or buffered goes straight to the small output for
the rest of the bundle
--
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]
