tejasiyer-dev commented on code in PR #40543:
URL: https://github.com/apache/beam/pull/40543#discussion_r4569443688
##########
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:
Done. Allocated once in @Setup and cleared in @StartBundle.
--
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]
