You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@beam.apache.org by GitBox <gi...@apache.org> on 2021/05/24 16:18:48 UTC

[GitHub] [beam] reuvenlax commented on a change in pull request #14852: [BEAM-12378] GroupIntoBatches improvements

reuvenlax commented on a change in pull request #14852:
URL: https://github.com/apache/beam/pull/14852#discussion_r638092467



##########
File path: sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java
##########
@@ -101,29 +111,69 @@
    * GroupIntoBatches} and {@link GroupIntoBatches.WithShardedKey}.
    */
   @AutoValue
-  public abstract static class BatchingParams implements Serializable {
-    public static BatchingParams create(long batchSize, Duration maxBufferingDuration) {
-      return new AutoValue_GroupIntoBatches_BatchingParams(batchSize, maxBufferingDuration);
+  public abstract static class BatchingParams<InputT> implements Serializable {
+    public static <InputT> BatchingParams<InputT> create(
+        long batchSize,
+        long batchSizeByes,

Review comment:
       done

##########
File path: sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java
##########
@@ -101,29 +111,69 @@
    * GroupIntoBatches} and {@link GroupIntoBatches.WithShardedKey}.
    */
   @AutoValue
-  public abstract static class BatchingParams implements Serializable {
-    public static BatchingParams create(long batchSize, Duration maxBufferingDuration) {
-      return new AutoValue_GroupIntoBatches_BatchingParams(batchSize, maxBufferingDuration);
+  public abstract static class BatchingParams<InputT> implements Serializable {

Review comment:
       added

##########
File path: sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/GroupIntoBatchesTest.java
##########
@@ -126,6 +127,99 @@ public Void apply(Iterable<KV<String, Iterable<String>>> input) {
     pipeline.run();
   }
 
+  @Test
+  @Category({NeedsRunner.class, UsesTimersInParDo.class, UsesStatefulParDo.class})
+  public void testInGlobalWindowBatchSizeByteSize() {
+    PCollection<KV<String, Iterable<String>>> collection =
+        pipeline
+            .apply("Input data", Create.of(data))
+            .apply(GroupIntoBatches.ofByteSize(BATCH_SIZE_BYTES))
+            // set output coder
+            .setCoder(KvCoder.of(StringUtf8Coder.of(), IterableCoder.of(StringUtf8Coder.of())));
+    PAssert.that("Incorrect batch size in one or more elements", collection)
+        .satisfies(
+            new SerializableFunction<Iterable<KV<String, Iterable<String>>>, Void>() {
+
+              private boolean checkBatchSizes(Iterable<KV<String, Iterable<String>>> listToCheck) {
+                for (KV<String, Iterable<String>> element : listToCheck) {
+                  long byteSize = 0;
+                  for (String str : element.getValue()) {
+                    if (byteSize >= BATCH_SIZE_BYTES) {
+                      // We already reached the batch size, so extra elements are not expected.
+                      return false;
+                    }
+                    try {
+                      byteSize += StringUtf8Coder.of().getEncodedElementByteSize(str);
+                    } catch (Exception e) {
+                      throw new RuntimeException(e);
+                    }
+                  }
+                }
+                return true;
+              }
+
+              @Override
+              public Void apply(Iterable<KV<String, Iterable<String>>> input) {
+                assertTrue(checkBatchSizes(input));
+                return null;
+              }
+            });
+    PAssert.thatSingleton("Incorrect collection size", collection.apply("Count", Count.globally()))
+        .isEqualTo(3L);
+    pipeline.run();
+  }
+
+  @Test
+  @Category({NeedsRunner.class, UsesTimersInParDo.class, UsesStatefulParDo.class})
+  public void testInGlobalWindowBatchSizeByteSizeFn() {
+    PCollection<KV<String, Iterable<String>>> collection =
+        pipeline
+            .apply("Input data", Create.of(data))
+            .apply(
+                GroupIntoBatches.ofByteSize(
+                    BATCH_SIZE_BYTES,
+                    s -> {
+                      try {
+                        return StringUtf8Coder.of().getEncodedElementByteSize(s);
+                      } catch (Exception e) {
+                        throw new RuntimeException(e);
+                      }
+                    }))

Review comment:
       ddone




-- 
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.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org