Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Fix for #22951 #22953

Closed
wants to merge 15 commits into from
Closed

Fix for #22951 #22953

Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
2eacbd4
Fix for #22951
nbali Aug 30, 2022
ff6b5a0
Compilation fix
nbali Aug 30, 2022
5f79ed6
Spotless fix
nbali Aug 30, 2022
bcd4ba9
Fix for #22951 - PR CR - using the same chunk size as the network layer
nbali Sep 6, 2022
040b744
Fix for #22951 - PR CR - added new test cases into GroupIntoBatchesTest
nbali Sep 6, 2022
b1b732c
Fix for #22951 - PR CR - added new test cases into GroupIntoBatchesTr…
nbali Sep 7, 2022
70a92b6
Fix for #22951 - PR CR - guaranteed element order for the newly intro…
nbali Sep 7, 2022
6abe4cd
#22591 Firing GroupIntoBatches earlier if byte size would go over the…
nbali Nov 15, 2022
3e53b0f
Spotless
nbali Nov 15, 2022
f00c85d
#22951 Adding byte size limit to WriteFiles transform as well
nbali Nov 15, 2022
5840e55
#22951 Enforcing the new expectations for the 'byteSize' limit in uni…
nbali Nov 16, 2022
6d016f1
Fixing State.readLater() usage + Unit test fix to reflect changed exp…
nbali Nov 16, 2022
9f874f7
Merge branch 'master' into fix-for-22951
nbali Nov 17, 2022
5a3ee7f
Merge branch 'master' into fix-for-22951
nbali Nov 18, 2022
fa1bd88
Apply suggestions from code review
lukecwik Nov 30, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
49 changes: 43 additions & 6 deletions .../test/java/org/apache/beam/runners/core/construction/GroupIntoBatchesTranslationTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.equalTo;

import java.util.Set;
import java.util.function.Function;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.model.pipeline.v1.RunnerApi.GroupIntoBatchesPayload;
import org.apache.beam.sdk.runners.AppliedPTransform;
Expand All @@ -32,6 +34,8 @@
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PValues;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Sets;
import org.joda.time.Duration;
import org.junit.Test;
import org.junit.runner.RunWith;
Expand All @@ -43,12 +47,45 @@
@RunWith(Parameterized.class)
public class GroupIntoBatchesTranslationTest {
@Parameters(name = "{index}: {0}")
public static Iterable<GroupIntoBatches<?, ?>> transform() {
return ImmutableList.of(
GroupIntoBatches.ofSize(5),
GroupIntoBatches.ofSize(5).withMaxBufferingDuration(Duration.ZERO),
GroupIntoBatches.ofSize(5).withMaxBufferingDuration(Duration.standardSeconds(10)),
GroupIntoBatches.ofByteSize(10).withMaxBufferingDuration(Duration.standardSeconds(10)));
public static Iterable<GroupIntoBatches<String, Integer>> transform() {
Set<Function<GroupIntoBatches<String, Integer>, GroupIntoBatches<String, Integer>>> creators =
ImmutableSet.of( //
$ -> GroupIntoBatches.ofSize(5), //
$ -> GroupIntoBatches.ofByteSize(10));

Set<Function<GroupIntoBatches<String, Integer>, GroupIntoBatches<String, Integer>>>
sizeModifiers =
ImmutableSet.of( //
gib -> gib, //
gib -> gib.withSize(5));

Set<Function<GroupIntoBatches<String, Integer>, GroupIntoBatches<String, Integer>>>
byteSizeModifiers =
ImmutableSet.of( //
gib -> gib, //
gib -> gib.withByteSize(10), //
gib -> gib.withByteSize(10, i -> (long) i));

Set<Function<GroupIntoBatches<String, Integer>, GroupIntoBatches<String, Integer>>>
maxBufferingDurationModifiers =
ImmutableSet.of( //
gib -> gib, //
gib -> gib.withMaxBufferingDuration(Duration.ZERO), //
gib -> gib.withMaxBufferingDuration(Duration.standardSeconds(10)));

return Sets.cartesianProduct(
creators, sizeModifiers, byteSizeModifiers, maxBufferingDurationModifiers)
.stream()
.map(
product -> {
GroupIntoBatches<String, Integer> groupIntoBatches = null;
for (Function<GroupIntoBatches<String, Integer>, GroupIntoBatches<String, Integer>>
fn : product) {
groupIntoBatches = fn.apply(groupIntoBatches);
}
return groupIntoBatches;
})
.collect(ImmutableList.toImmutableList());
}

@Parameter(0)
Expand Down
57 changes: 49 additions & 8 deletions sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/GroupIntoBatches.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,6 +117,11 @@ public class GroupIntoBatches<K, InputT>
*/
@AutoValue
public abstract static class BatchingParams<InputT> implements Serializable {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It looks like your adding support for GroupIntoBatches to limit on count and byte size at the same time.

Can you add tests that cover this new scenario to:

  • GroupIntoBatchesTest
  • GroupIntoBatchesTranslationTest

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

see 040b744 (#22953) and b1b732c (#22953)

public static <InputT> BatchingParams<InputT> createDefault() {
lukecwik marked this conversation as resolved.
Show resolved Hide resolved
return new AutoValue_GroupIntoBatches_BatchingParams(
Long.MAX_VALUE, Long.MAX_VALUE, null, Duration.ZERO);
}

public static <InputT> BatchingParams<InputT> create(
long batchSize,
long batchSizeBytes,
Expand Down Expand Up @@ -170,8 +175,7 @@ private GroupIntoBatches(BatchingParams<InputT> params) {
/** Aim to create batches each with the specified element count. */
public static <K, InputT> GroupIntoBatches<K, InputT> ofSize(long batchSize) {
Preconditions.checkState(batchSize < Long.MAX_VALUE);
return new GroupIntoBatches<>(
BatchingParams.create(batchSize, Long.MAX_VALUE, null, Duration.ZERO));
return new GroupIntoBatches<K, InputT>(BatchingParams.createDefault()).withSize(batchSize);
}

/**
Expand All @@ -185,9 +189,8 @@ public static <K, InputT> GroupIntoBatches<K, InputT> ofSize(long batchSize) {
* {@link #ofByteSize(long, SerializableFunction)} to specify code to calculate the byte size.
*/
public static <K, InputT> GroupIntoBatches<K, InputT> ofByteSize(long batchSizeBytes) {
Preconditions.checkState(batchSizeBytes < Long.MAX_VALUE);
return new GroupIntoBatches<>(
BatchingParams.create(Long.MAX_VALUE, batchSizeBytes, null, Duration.ZERO));
return new GroupIntoBatches<K, InputT>(BatchingParams.createDefault())
.withByteSize(batchSizeBytes);
}

/**
Expand All @@ -196,16 +199,54 @@ public static <K, InputT> GroupIntoBatches<K, InputT> ofByteSize(long batchSizeB
*/
public static <K, InputT> GroupIntoBatches<K, InputT> ofByteSize(
long batchSizeBytes, SerializableFunction<InputT, Long> getElementByteSize) {
Preconditions.checkState(batchSizeBytes < Long.MAX_VALUE);
return new GroupIntoBatches<>(
BatchingParams.create(Long.MAX_VALUE, batchSizeBytes, getElementByteSize, Duration.ZERO));
return new GroupIntoBatches<K, InputT>(BatchingParams.createDefault())
.withByteSize(batchSizeBytes, getElementByteSize);
}

/** Returns user supplied parameters for batching. */
public BatchingParams<InputT> getBatchingParams() {
return params;
}

@Override
public String toString() {
return super.toString() + ", params=" + params;
}

/** @see #ofSize(long) */
public GroupIntoBatches<K, InputT> withSize(long batchSize) {
Preconditions.checkState(batchSize < Long.MAX_VALUE);
return new GroupIntoBatches<>(
BatchingParams.create(
batchSize,
params.getBatchSizeBytes(),
params.getElementByteSize(),
params.getMaxBufferingDuration()));
}

/** @see #ofByteSize(long) */
public GroupIntoBatches<K, InputT> withByteSize(long batchSizeBytes) {
Preconditions.checkState(batchSizeBytes < Long.MAX_VALUE);
return new GroupIntoBatches<>(
BatchingParams.create(
params.getBatchSize(),
batchSizeBytes,
params.getElementByteSize(),
params.getMaxBufferingDuration()));
}

/** @see #ofByteSize(long, SerializableFunction) */
public GroupIntoBatches<K, InputT> withByteSize(
long batchSizeBytes, SerializableFunction<InputT, Long> getElementByteSize) {
Preconditions.checkState(batchSizeBytes < Long.MAX_VALUE);
return new GroupIntoBatches<>(
BatchingParams.create(
params.getBatchSize(),
batchSizeBytes,
getElementByteSize,
params.getMaxBufferingDuration()));
}

/**
* Sets a time limit (in processing time) on how long an incomplete batch of elements is allowed
* to be buffered. Once a batch is flushed to output, the timer is reset. The provided limit must
Expand Down