You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by jk...@apache.org on 2017/09/27 22:29:39 UTC
[1/4] beam git commit: Adds ReadableFile.toString() for debugging
Repository: beam
Updated Branches:
refs/heads/master 7e3f591b2 -> 393e56310
Adds ReadableFile.toString() for debugging
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/a6e20017
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/a6e20017
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/a6e20017
Branch: refs/heads/master
Commit: a6e20017e0aadf6d0015696c8d8a22cec6d48077
Parents: b74644a
Author: Eugene Kirpichov <ki...@google.com>
Authored: Mon Sep 25 11:58:53 2017 -0700
Committer: Eugene Kirpichov <ki...@google.com>
Committed: Wed Sep 27 15:08:38 2017 -0700
----------------------------------------------------------------------
sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileIO.java | 5 +++++
1 file changed, 5 insertions(+)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/a6e20017/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileIO.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileIO.java
index 6b75370..7df4bde 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileIO.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileIO.java
@@ -152,6 +152,11 @@ public class FileIO {
public String readFullyAsUTF8String() throws IOException {
return new String(readFullyAsBytes(), StandardCharsets.UTF_8);
}
+
+ @Override
+ public String toString() {
+ return "ReadableFile{metadata=" + metadata + ", compression=" + compression + '}';
+ }
}
/**
[3/4] beam git commit: PAssert improvements.
Posted by jk...@apache.org.
PAssert improvements.
- Captures stack trace by introducing a SerializableThrowable.
Fixes an incorrect test of this.
- PAssert.thatSingletonIterable, thatMap/Multimap/Singleton no longer
require that the collection is produced by a trigger that promises
a single firing. thatSingletonIterable checks that the iterable is a
singleton by other means. thatMap/Multimap/Singleton don't need this
requirement at all.
PaneExtractors.onlyPane() is now used only when the user explicitly
specifies inOnlyPane().
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/30e7d15b
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/30e7d15b
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/30e7d15b
Branch: refs/heads/master
Commit: 30e7d15b3ab68d0c129dbd2be76e77346f2e1f38
Parents: 7e3f591
Author: Eugene Kirpichov <ki...@google.com>
Authored: Mon Sep 25 11:59:04 2017 -0700
Committer: Eugene Kirpichov <ki...@google.com>
Committed: Wed Sep 27 15:08:38 2017 -0700
----------------------------------------------------------------------
.../org/apache/beam/sdk/testing/PAssert.java | 63 +++++++-------------
.../apache/beam/sdk/testing/PaneExtractors.java | 25 +++++---
.../beam/sdk/testing/SuccessOrFailure.java | 41 ++++++++-----
.../apache/beam/sdk/testing/PAssertTest.java | 41 +++++++------
.../beam/sdk/testing/PaneExtractorsTest.java | 7 +--
5 files changed, 91 insertions(+), 86 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/30e7d15b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PAssert.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PAssert.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PAssert.java
index ed80f2f..d2ad67d 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PAssert.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PAssert.java
@@ -31,7 +31,6 @@ import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.Map;
-import java.util.NoSuchElementException;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.Pipeline.PipelineVisitor;
import org.apache.beam.sdk.PipelineRunner;
@@ -74,8 +73,6 @@ import org.apache.beam.sdk.values.PDone;
import org.apache.beam.sdk.values.ValueInSingleWindow;
import org.apache.beam.sdk.values.WindowingStrategy;
import org.joda.time.Duration;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
/**
* An assertion on the contents of a {@link PCollection} incorporated into the pipeline. Such an
@@ -105,8 +102,6 @@ import org.slf4j.LoggerFactory;
* <p>JUnit and Hamcrest must be linked in by any code that uses PAssert.
*/
public class PAssert {
-
- private static final Logger LOG = LoggerFactory.getLogger(PAssert.class);
public static final String SUCCESS_COUNTER = "PAssertSuccess";
public static final String FAILURE_COUNTER = "PAssertFailure";
private static final Counter successCounter = Metrics.counter(
@@ -170,10 +165,6 @@ public class PAssert {
return new PAssertionSite(message, new Throwable().getStackTrace());
}
- PAssertionSite() {
- this(null, new StackTraceElement[0]);
- }
-
PAssertionSite(String message, StackTraceElement[] creationStackTrace) {
this.message = message;
this.creationStackTrace = creationStackTrace;
@@ -381,15 +372,6 @@ public class PAssert {
*/
public static <T> IterableAssert<T> thatSingletonIterable(
String reason, PCollection<? extends Iterable<T>> actual) {
-
- try {
- } catch (NoSuchElementException | IllegalArgumentException exc) {
- throw new IllegalArgumentException(
- "PAssert.<T>thatSingletonIterable requires a PCollection<Iterable<T>>"
- + " with a Coder<Iterable<T>> where getCoderArguments() yields a"
- + " single Coder<T> to apply to the elements.");
- }
-
@SuppressWarnings("unchecked") // Safe covariant cast
PCollection<Iterable<T>> actualIterables = (PCollection<Iterable<T>>) actual;
@@ -581,7 +563,7 @@ public class PAssert {
@SafeVarargs
final PCollectionContentsAssert<T> containsInAnyOrder(
SerializableMatcher<? super T>... elementMatchers) {
- return satisfies(SerializableMatchers.<T>containsInAnyOrder(elementMatchers));
+ return satisfies(SerializableMatchers.containsInAnyOrder(elementMatchers));
}
/**
@@ -592,7 +574,7 @@ public class PAssert {
private PCollectionContentsAssert<T> satisfies(
AssertRelation<Iterable<T>, Iterable<T>> relation, Iterable<T> expectedElements) {
return satisfies(
- new CheckRelationAgainstExpected<Iterable<T>>(
+ new CheckRelationAgainstExpected<>(
relation, expectedElements, IterableCoder.of(actual.getCoder())));
}
@@ -668,7 +650,10 @@ public class PAssert {
public PCollectionSingletonIterableAssert(
PCollection<Iterable<T>> actual, PAssertionSite site) {
this(
- actual, IntoGlobalWindow.<Iterable<T>>of(), PaneExtractors.<Iterable<T>>onlyPane(), site);
+ actual,
+ IntoGlobalWindow.<Iterable<T>>of(),
+ PaneExtractors.<Iterable<T>>allPanes(),
+ site);
}
public PCollectionSingletonIterableAssert(
@@ -753,7 +738,7 @@ public class PAssert {
private PCollectionSingletonIterableAssert<T> satisfies(
AssertRelation<Iterable<T>, Iterable<T>> relation, Iterable<T> expectedElements) {
return satisfies(
- new CheckRelationAgainstExpected<Iterable<T>>(
+ new CheckRelationAgainstExpected<>(
relation, expectedElements, IterableCoder.of(elementCoder)));
}
}
@@ -777,8 +762,12 @@ public class PAssert {
Coder<ViewT> coder,
PAssertionSite site) {
this(
- actual, view, IntoGlobalWindow.<ElemT>of(), PaneExtractors.<ElemT>onlyPane(), coder, site
- );
+ actual,
+ view,
+ IntoGlobalWindow.<ElemT>of(),
+ PaneExtractors.<ElemT>allPanes(),
+ coder,
+ site);
}
private PCollectionViewAssert(
@@ -798,7 +787,7 @@ public class PAssert {
@Override
public PCollectionViewAssert<ElemT, ViewT> inOnlyPane(BoundedWindow window) {
- return inPane(window, PaneExtractors.<ElemT>onlyPane());
+ return inPane(window, PaneExtractors.<ElemT>onlyPane(site));
}
@Override
@@ -841,7 +830,7 @@ public class PAssert {
.getPipeline()
.apply(
"PAssert$" + (assertCount++),
- new OneSideInputAssert<ViewT>(
+ new OneSideInputAssert<>(
CreateActual.from(actual, rewindowActuals, paneExtractor, view),
rewindowActuals.<Integer>windowDummy(),
checkerFn,
@@ -857,7 +846,7 @@ public class PAssert {
*/
private PCollectionViewAssert<ElemT, ViewT> satisfies(
AssertRelation<ViewT, ViewT> relation, final ViewT expectedValue) {
- return satisfies(new CheckRelationAgainstExpected<ViewT>(relation, expectedValue, coder));
+ return satisfies(new CheckRelationAgainstExpected<>(relation, expectedValue, coder));
}
/**
@@ -1224,11 +1213,7 @@ public class PAssert {
@ProcessElement
public void processElement(ProcessContext c) {
- try {
- c.output(doChecks(site, c.element(), checkerFn));
- } catch (Throwable t) {
- throw t;
- }
+ c.output(doChecks(site, c.element(), checkerFn));
}
}
@@ -1262,13 +1247,11 @@ public class PAssert {
PAssertionSite site,
ActualT actualContents,
SerializableFunction<ActualT, Void> checkerFn) {
- SuccessOrFailure result = SuccessOrFailure.success();
try {
checkerFn.apply(actualContents);
+ return SuccessOrFailure.success();
} catch (Throwable t) {
- result = SuccessOrFailure.failure(site, t.getMessage());
- } finally {
- return result;
+ return SuccessOrFailure.failure(site, t);
}
}
@@ -1329,7 +1312,7 @@ public class PAssert {
}
public AssertContainsInAnyOrder(Iterable<T> expected) {
- this(Lists.<T>newArrayList(expected));
+ this(Lists.newArrayList(expected));
}
@Override
@@ -1356,7 +1339,7 @@ public class PAssert {
private static class AssertIsEqualToRelation<T> implements AssertRelation<T, T> {
@Override
public SerializableFunction<T, Void> assertFor(T expected) {
- return new AssertIsEqualTo<T>(expected);
+ return new AssertIsEqualTo<>(expected);
}
}
@@ -1366,7 +1349,7 @@ public class PAssert {
private static class AssertNotEqualToRelation<T> implements AssertRelation<T, T> {
@Override
public SerializableFunction<T, Void> assertFor(T expected) {
- return new AssertNotEqualTo<T>(expected);
+ return new AssertNotEqualTo<>(expected);
}
}
@@ -1378,7 +1361,7 @@ public class PAssert {
implements AssertRelation<Iterable<T>, Iterable<T>> {
@Override
public SerializableFunction<Iterable<T>, Void> assertFor(Iterable<T> expectedElements) {
- return new AssertContainsInAnyOrder<T>(expectedElements);
+ return new AssertContainsInAnyOrder<>(expectedElements);
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/30e7d15b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PaneExtractors.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PaneExtractors.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PaneExtractors.java
index f88efcb..8ff35f3 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PaneExtractors.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/PaneExtractors.java
@@ -17,8 +17,6 @@
*/
package org.apache.beam.sdk.testing;
-import static com.google.common.base.Preconditions.checkState;
-
import java.util.ArrayList;
import java.util.List;
import org.apache.beam.sdk.transforms.PTransform;
@@ -42,8 +40,9 @@ final class PaneExtractors {
private PaneExtractors() {
}
- static <T> SimpleFunction<Iterable<ValueInSingleWindow<T>>, Iterable<T>> onlyPane() {
- return new ExtractOnlyPane<>();
+ static <T> SimpleFunction<Iterable<ValueInSingleWindow<T>>, Iterable<T>> onlyPane(
+ PAssert.PAssertionSite site) {
+ return new ExtractOnlyPane<>(site);
}
static <T> SimpleFunction<Iterable<ValueInSingleWindow<T>>, Iterable<T>> onTimePane() {
@@ -68,15 +67,23 @@ final class PaneExtractors {
private static class ExtractOnlyPane<T>
extends SimpleFunction<Iterable<ValueInSingleWindow<T>>, Iterable<T>> {
+ private final PAssert.PAssertionSite site;
+
+ private ExtractOnlyPane(PAssert.PAssertionSite site) {
+ this.site = site;
+ }
+
@Override
public Iterable<T> apply(Iterable<ValueInSingleWindow<T>> input) {
List<T> outputs = new ArrayList<>();
for (ValueInSingleWindow<T> value : input) {
- checkState(value.getPane().isFirst() && value.getPane().isLast(),
- "Expected elements to be produced by a trigger that fires at most once, but got"
- + "a value in a pane that is %s. Actual Pane Info: %s",
- value.getPane().isFirst() ? "not the last pane" : "not the first pane",
- value.getPane());
+ if (!value.getPane().isFirst() || !value.getPane().isLast()) {
+ throw site.wrap(
+ String.format(
+ "Expected elements to be produced by a trigger that fires at most once, but got "
+ + "a value %s in a pane that is %s.",
+ value, value.getPane().isFirst() ? "not the last pane" : "not the first pane"));
+ }
outputs.add(value.getValue());
}
return outputs;
http://git-wip-us.apache.org/repos/asf/beam/blob/30e7d15b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/SuccessOrFailure.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/SuccessOrFailure.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/SuccessOrFailure.java
index 04e3c35..79e83d6 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/SuccessOrFailure.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/testing/SuccessOrFailure.java
@@ -18,6 +18,8 @@
package org.apache.beam.sdk.testing;
import com.google.common.base.MoreObjects;
+import java.io.IOException;
+import java.io.ObjectInputStream;
import java.io.Serializable;
import javax.annotation.Nullable;
import org.apache.beam.sdk.coders.DefaultCoder;
@@ -28,25 +30,36 @@ import org.apache.beam.sdk.coders.SerializableCoder;
*/
@DefaultCoder(SerializableCoder.class)
public final class SuccessOrFailure implements Serializable {
- // TODO Add a SerializableThrowable. instead of relying on PAssertionSite.(BEAM-1898)
+ private static final class SerializableThrowable implements Serializable {
+ private final Throwable throwable;
+ private final StackTraceElement[] stackTrace;
+
+ private SerializableThrowable(Throwable t) {
+ this.throwable = t;
+ this.stackTrace = (t == null) ? null : t.getStackTrace();
+ }
+
+ private void readObject(ObjectInputStream is) throws IOException, ClassNotFoundException {
+ is.defaultReadObject();
+ if (throwable != null) {
+ throwable.setStackTrace(stackTrace);
+ }
+ }
+ }
private final boolean isSuccess;
@Nullable
private final PAssert.PAssertionSite site;
@Nullable
- private final String message;
-
- private SuccessOrFailure() {
- this(true, null, null);
- }
+ private final SerializableThrowable throwable;
private SuccessOrFailure(
boolean isSuccess,
@Nullable PAssert.PAssertionSite site,
- @Nullable String message) {
+ @Nullable Throwable throwable) {
this.isSuccess = isSuccess;
this.site = site;
- this.message = message;
+ this.throwable = new SerializableThrowable(throwable);
}
public boolean isSuccess() {
@@ -55,7 +68,7 @@ public final class SuccessOrFailure implements Serializable {
@Nullable
public AssertionError assertionError() {
- return site == null ? null : site.wrap(message);
+ return site == null ? null : site.wrap(throwable.throwable);
}
public static SuccessOrFailure success() {
@@ -63,19 +76,15 @@ public final class SuccessOrFailure implements Serializable {
}
public static SuccessOrFailure failure(@Nullable PAssert.PAssertionSite site,
- @Nullable String message) {
- return new SuccessOrFailure(false, site, message);
- }
-
- public static SuccessOrFailure failure(@Nullable PAssert.PAssertionSite site) {
- return new SuccessOrFailure(false, site, null);
+ @Nullable Throwable t) {
+ return new SuccessOrFailure(false, site, t);
}
@Override
public String toString() {
return MoreObjects.toStringHelper(this)
.add("isSuccess", isSuccess())
- .addValue(message)
+ .addValue(throwable)
.omitNullValues()
.toString();
}
http://git-wip-us.apache.org/repos/asf/beam/blob/30e7d15b/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PAssertTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PAssertTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PAssertTest.java
index 491f001..2a79060 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PAssertTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PAssertTest.java
@@ -121,27 +121,34 @@ public class PAssertTest implements Serializable {
}
}
+ private void throwNestedError() {
+ throw new RuntimeException("Nested error");
+ }
+
+ private void throwWrappedError() {
+ try {
+ throwNestedError();
+ } catch (Exception e) {
+ throw new RuntimeException("Wrapped error", e);
+ }
+ }
+
@Test
- public void testFailureEncodedDecoded() throws IOException {
- AssertionError error = null;
+ public void testFailureWithExceptionEncodedDecoded() throws IOException {
+ Throwable error;
try {
- assertEquals(0, 1);
- } catch (AssertionError e) {
+ throwWrappedError();
+ throw new IllegalStateException("Should have failed");
+ } catch (Throwable e) {
error = e;
}
- SuccessOrFailure failure = SuccessOrFailure.failure(
- new PAssert.PAssertionSite(error.getMessage(), error.getStackTrace()));
- SerializableCoder<SuccessOrFailure> coder = SerializableCoder.of(SuccessOrFailure.class);
-
- byte[] encoded = CoderUtils.encodeToByteArray(coder, failure);
- SuccessOrFailure res = CoderUtils.decodeFromByteArray(coder, encoded);
-
- // Should compare strings, because throwables are not directly comparable.
- assertEquals("Encode-decode failed SuccessOrFailure",
- failure.assertionError().toString(), res.assertionError().toString());
- String resultStacktrace = Throwables.getStackTraceAsString(res.assertionError());
- String failureStacktrace = Throwables.getStackTraceAsString(failure.assertionError());
- assertThat(resultStacktrace, is(failureStacktrace));
+ SuccessOrFailure failure =
+ SuccessOrFailure.failure(PAssert.PAssertionSite.capture("here"), error);
+ SuccessOrFailure res = CoderUtils.clone(SerializableCoder.of(SuccessOrFailure.class), failure);
+ assertEquals(
+ "Encode-decode failed SuccessOrFailure",
+ Throwables.getStackTraceAsString(failure.assertionError()),
+ Throwables.getStackTraceAsString(res.assertionError()));
}
@Test
http://git-wip-us.apache.org/repos/asf/beam/blob/30e7d15b/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PaneExtractorsTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PaneExtractorsTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PaneExtractorsTest.java
index 8801bde..1d8390e 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PaneExtractorsTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/testing/PaneExtractorsTest.java
@@ -43,7 +43,7 @@ public class PaneExtractorsTest {
@Test
public void onlyPaneNoFiring() {
SerializableFunction<Iterable<ValueInSingleWindow<Integer>>, Iterable<Integer>> extractor =
- PaneExtractors.onlyPane();
+ PaneExtractors.onlyPane(PAssert.PAssertionSite.capture(""));
Iterable<ValueInSingleWindow<Integer>> noFiring =
ImmutableList.of(
ValueInSingleWindow.of(
@@ -56,7 +56,7 @@ public class PaneExtractorsTest {
@Test
public void onlyPaneOnlyOneFiring() {
SerializableFunction<Iterable<ValueInSingleWindow<Integer>>, Iterable<Integer>> extractor =
- PaneExtractors.onlyPane();
+ PaneExtractors.onlyPane(PAssert.PAssertionSite.capture(""));
Iterable<ValueInSingleWindow<Integer>> onlyFiring =
ImmutableList.of(
ValueInSingleWindow.of(
@@ -70,7 +70,7 @@ public class PaneExtractorsTest {
@Test
public void onlyPaneMultiplePanesFails() {
SerializableFunction<Iterable<ValueInSingleWindow<Integer>>, Iterable<Integer>> extractor =
- PaneExtractors.onlyPane();
+ PaneExtractors.onlyPane(PAssert.PAssertionSite.capture(""));
Iterable<ValueInSingleWindow<Integer>> multipleFiring =
ImmutableList.of(
ValueInSingleWindow.of(
@@ -89,7 +89,6 @@ public class PaneExtractorsTest {
GlobalWindow.INSTANCE,
PaneInfo.createPane(false, false, Timing.LATE, 2L, 1L)));
- thrown.expect(IllegalStateException.class);
thrown.expectMessage("trigger that fires at most once");
extractor.apply(multipleFiring);
}
[4/4] beam git commit: This closes #3890: Introduces
Reshuffle.viaRandomKey()
Posted by jk...@apache.org.
This closes #3890: Introduces Reshuffle.viaRandomKey()
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/393e5631
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/393e5631
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/393e5631
Branch: refs/heads/master
Commit: 393e5631054a81ae1fdcd304f81cc68cf53d3422
Parents: 7e3f591 a6e2001
Author: Eugene Kirpichov <ki...@google.com>
Authored: Wed Sep 27 15:08:59 2017 -0700
Committer: Eugene Kirpichov <ki...@google.com>
Committed: Wed Sep 27 15:08:59 2017 -0700
----------------------------------------------------------------------
.../java/org/apache/beam/sdk/io/FileIO.java | 12 +++-
.../beam/sdk/io/ReadAllViaFileBasedSource.java | 29 +--------
.../org/apache/beam/sdk/testing/PAssert.java | 63 +++++++-------------
.../apache/beam/sdk/testing/PaneExtractors.java | 25 +++++---
.../beam/sdk/testing/SuccessOrFailure.java | 41 ++++++++-----
.../apache/beam/sdk/transforms/Reshuffle.java | 47 +++++++++++++++
.../apache/beam/sdk/testing/PAssertTest.java | 41 +++++++------
.../beam/sdk/testing/PaneExtractorsTest.java | 7 +--
.../beam/sdk/io/gcp/bigquery/BigQueryIO.java | 12 +---
.../io/gcp/bigquery/StreamingWriteTables.java | 10 ++--
.../beam/sdk/io/gcp/datastore/DatastoreV1.java | 27 +++------
.../beam/sdk/io/gcp/spanner/SpannerIO.java | 27 ++-------
.../sdk/io/gcp/datastore/DatastoreV1Test.java | 15 ++---
.../sdk/io/gcp/datastore/SplitQueryFnIT.java | 5 +-
.../org/apache/beam/sdk/io/jdbc/JdbcIO.java | 24 +-------
15 files changed, 178 insertions(+), 207 deletions(-)
----------------------------------------------------------------------
[2/4] beam git commit: Introduces Reshuffle.viaRandomKey()
Posted by jk...@apache.org.
Introduces Reshuffle.viaRandomKey()
It's a commonly used pattern for breaking fusion
https://cloud.google.com/dataflow/service/dataflow-service-desc#fusion-optimization
viaRandomKey() only abstracts away the current commonly used pattern.
It has the same caveats as using Reshuffle.of() directly - the semantics
are technically not guaranteed by the Beam model, but it works in
practice, and this is the pattern we keep recommending to users.
The naming is deliberately operational rather than semantic, to
emphasize that we don't have the semantics figured out, and the
transform promises only that it expands into exactly the sequence
"pair with random key, reshuffle, drop key".
The goal of this change is just to reduce copy-paste.
See prior discussion at
https://lists.apache.org/thread.html/ac34c9ac665a8d9f67b0254015e44c59ea65ecc1360d4014b95d3b2e@%3Cdev.beam.apache.org%3E
This change also converts several existing usages to use it, and adds another
one in Match.
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/b74644a2
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/b74644a2
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/b74644a2
Branch: refs/heads/master
Commit: b74644a2bccd7fca90da4912781196c804e41325
Parents: 30e7d15
Author: Eugene Kirpichov <ki...@google.com>
Authored: Fri Sep 22 15:24:36 2017 -0700
Committer: Eugene Kirpichov <ki...@google.com>
Committed: Wed Sep 27 15:08:38 2017 -0700
----------------------------------------------------------------------
.../java/org/apache/beam/sdk/io/FileIO.java | 7 ++-
.../beam/sdk/io/ReadAllViaFileBasedSource.java | 29 +-----------
.../apache/beam/sdk/transforms/Reshuffle.java | 47 ++++++++++++++++++++
.../beam/sdk/io/gcp/bigquery/BigQueryIO.java | 12 +----
.../io/gcp/bigquery/StreamingWriteTables.java | 10 ++---
.../beam/sdk/io/gcp/datastore/DatastoreV1.java | 27 ++++-------
.../beam/sdk/io/gcp/spanner/SpannerIO.java | 27 +++--------
.../sdk/io/gcp/datastore/DatastoreV1Test.java | 15 +++----
.../sdk/io/gcp/datastore/SplitQueryFnIT.java | 5 +--
.../org/apache/beam/sdk/io/jdbc/JdbcIO.java | 24 +---------
10 files changed, 82 insertions(+), 121 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/b74644a2/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileIO.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileIO.java
index c909c3c..6b75370 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileIO.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileIO.java
@@ -37,6 +37,7 @@ import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.Reshuffle;
import org.apache.beam.sdk.transforms.Values;
import org.apache.beam.sdk.transforms.Watch;
import org.apache.beam.sdk.transforms.Watch.Growth.TerminationCondition;
@@ -305,12 +306,13 @@ public class FileIO {
@Override
public PCollection<MatchResult.Metadata> expand(PCollection<String> input) {
+ PCollection<MatchResult.Metadata> res;
if (getConfiguration().getWatchInterval() == null) {
- return input.apply(
+ res = input.apply(
"Match filepatterns",
ParDo.of(new MatchFn(getConfiguration().getEmptyMatchTreatment())));
} else {
- return input
+ res = input
.apply(
"Continuously match filepatterns",
Watch.growthOf(new MatchPollFn())
@@ -318,6 +320,7 @@ public class FileIO {
.withTerminationPerInput(getConfiguration().getWatchTerminationCondition()))
.apply(Values.<MatchResult.Metadata>create());
}
+ return res.apply(Reshuffle.<MatchResult.Metadata>viaRandomKey());
}
private static class MatchFn extends DoFn<String, MatchResult.Metadata> {
http://git-wip-us.apache.org/repos/asf/beam/blob/b74644a2/sdks/java/core/src/main/java/org/apache/beam/sdk/io/ReadAllViaFileBasedSource.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/ReadAllViaFileBasedSource.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/ReadAllViaFileBasedSource.java
index 03cdbb1..c53f405 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/ReadAllViaFileBasedSource.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/ReadAllViaFileBasedSource.java
@@ -18,7 +18,6 @@
package org.apache.beam.sdk.io;
import java.io.IOException;
-import java.util.concurrent.ThreadLocalRandom;
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.io.FileIO.ReadableFile;
@@ -30,7 +29,6 @@ import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.Reshuffle;
import org.apache.beam.sdk.transforms.SerializableFunction;
-import org.apache.beam.sdk.transforms.Values;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
@@ -62,36 +60,11 @@ public class ReadAllViaFileBasedSource<T>
public PCollection<T> expand(PCollection<ReadableFile> input) {
return input
.apply("Split into ranges", ParDo.of(new SplitIntoRangesFn(desiredBundleSizeBytes)))
- .apply("Reshuffle", new ReshuffleWithUniqueKey<KV<ReadableFile, OffsetRange>>())
+ .apply("Reshuffle", Reshuffle.<KV<ReadableFile, OffsetRange>>viaRandomKey())
.apply("Read ranges", ParDo.of(new ReadFileRangesFn<T>(createSource)))
.setCoder(coder);
}
- private static class ReshuffleWithUniqueKey<T>
- extends PTransform<PCollection<T>, PCollection<T>> {
- @Override
- public PCollection<T> expand(PCollection<T> input) {
- return input
- .apply("Unique key", ParDo.of(new AssignUniqueKeyFn<T>()))
- .apply("Reshuffle", Reshuffle.<Integer, T>of())
- .apply("Values", Values.<T>create());
- }
- }
-
- private static class AssignUniqueKeyFn<T> extends DoFn<T, KV<Integer, T>> {
- private int index;
-
- @Setup
- public void setup() {
- this.index = ThreadLocalRandom.current().nextInt();
- }
-
- @ProcessElement
- public void process(ProcessContext c) {
- c.output(KV.of(++index, c.element()));
- }
- }
-
private static class SplitIntoRangesFn extends DoFn<ReadableFile, KV<ReadableFile, OffsetRange>> {
private final long desiredBundleSizeBytes;
http://git-wip-us.apache.org/repos/asf/beam/blob/b74644a2/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reshuffle.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reshuffle.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reshuffle.java
index 3b7122c..68e4560 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reshuffle.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reshuffle.java
@@ -17,6 +17,8 @@
*/
package org.apache.beam.sdk.transforms;
+import java.util.concurrent.ThreadLocalRandom;
+import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.annotations.Internal;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.ReshuffleTrigger;
@@ -55,6 +57,15 @@ public class Reshuffle<K, V> extends PTransform<PCollection<KV<K, V>>, PCollecti
return new Reshuffle<K, V>();
}
+ /**
+ * Encapsulates the sequence "pair input with unique key, apply {@link
+ * Reshuffle#of}, drop the key" commonly used to break fusion.
+ */
+ @Experimental
+ public static <T> ViaRandomKey<T> viaRandomKey() {
+ return new ViaRandomKey<T>();
+ }
+
@Override
public PCollection<KV<K, V>> expand(PCollection<KV<K, V>> input) {
WindowingStrategy<?, ?> originalStrategy = input.getWindowingStrategy();
@@ -94,4 +105,40 @@ public class Reshuffle<K, V> extends PTransform<PCollection<KV<K, V>>, PCollecti
"RestoreOriginalTimestamps",
ReifyTimestamps.<K, V>extractFromValues());
}
+
+ /** Implementation of {@link #viaRandomKey()}. */
+ public static class ViaRandomKey<T> extends PTransform<PCollection<T>, PCollection<T>> {
+ private ViaRandomKey() {}
+
+ @Override
+ public PCollection<T> expand(PCollection<T> input) {
+ return input
+ .apply("Pair with random key", ParDo.of(new AssignShardFn<T>()))
+ .apply(Reshuffle.<Integer, T>of())
+ .apply(Values.<T>create());
+ }
+
+ private static class AssignShardFn<T> extends DoFn<T, KV<Integer, T>> {
+ private int shard;
+
+ @Setup
+ public void setup() {
+ shard = ThreadLocalRandom.current().nextInt();
+ }
+
+ @ProcessElement
+ public void processElement(ProcessContext context) {
+ ++shard;
+ // Smear the shard into something more random-looking, to avoid issues
+ // with runners that don't properly hash the key being shuffled, but rely
+ // on it being random-looking. E.g. Spark takes the Java hashCode() of keys,
+ // which for Integer is a no-op and it is an issue:
+ // http://hydronitrogen.com/poor-hash-partitioning-of-timestamps-integers-and-longs-in-
+ // spark.html
+ // This hashing strategy is copied from com.google.common.collect.Hashing.smear().
+ int hashOfShard = 0x1b873593 * Integer.rotateLeft(shard * 0xcc9e2d51, 15);
+ context.output(KV.of(hashOfShard, context.element()));
+ }
+ }
+ }
}
http://git-wip-us.apache.org/repos/asf/beam/blob/b74644a2/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java
index e0b86b6..9e1dbfe 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java
@@ -81,9 +81,7 @@ import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.Reshuffle;
import org.apache.beam.sdk.transforms.SerializableFunction;
import org.apache.beam.sdk.transforms.SimpleFunction;
-import org.apache.beam.sdk.transforms.Values;
import org.apache.beam.sdk.transforms.View;
-import org.apache.beam.sdk.transforms.WithKeys;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.util.Transport;
import org.apache.beam.sdk.util.gcsfs.GcsPath;
@@ -598,15 +596,7 @@ public class BigQueryIO {
rows =
tuple
.get(filesTag)
- .apply(
- WithKeys.of(
- new SerializableFunction<String, String>() {
- public String apply(String s) {
- return s;
- }
- }))
- .apply(Reshuffle.<String, String>of())
- .apply(Values.<String>create())
+ .apply(Reshuffle.<String>viaRandomKey())
.apply(
"ReadFiles",
ParDo.of(
http://git-wip-us.apache.org/repos/asf/beam/blob/b74644a2/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteTables.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteTables.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteTables.java
index fa5b3ce..44563c0 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteTables.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StreamingWriteTables.java
@@ -79,10 +79,11 @@ public class StreamingWriteTables extends PTransform<
// get good batching into BigQuery's insert calls, and enough that we can max out the
// streaming insert quota.
PCollection<KV<ShardedKey<String>, TableRowInfo>> tagged =
- input.apply("ShardTableWrites", ParDo.of
- (new GenerateShardedTable(50)))
- .setCoder(KvCoder.of(ShardedKeyCoder.of(StringUtf8Coder.of()), TableRowJsonCoder.of()))
- .apply("TagWithUniqueIds", ParDo.of(new TagWithUniqueIds()));
+ input
+ .apply("ShardTableWrites", ParDo.of(new GenerateShardedTable(50)))
+ .setCoder(KvCoder.of(ShardedKeyCoder.of(StringUtf8Coder.of()), TableRowJsonCoder.of()))
+ .apply("TagWithUniqueIds", ParDo.of(new TagWithUniqueIds()))
+ .setCoder(KvCoder.of(ShardedKeyCoder.of(StringUtf8Coder.of()), TableRowInfoCoder.of()));
// To prevent having the same TableRow processed more than once with regenerated
// different unique ids, this implementation relies on "checkpointing", which is
@@ -91,7 +92,6 @@ public class StreamingWriteTables extends PTransform<
TupleTag<Void> mainOutputTag = new TupleTag<>("mainOutput");
TupleTag<TableRow> failedInsertsTag = new TupleTag<>("failedInserts");
PCollectionTuple tuple = tagged
- .setCoder(KvCoder.of(ShardedKeyCoder.of(StringUtf8Coder.of()), TableRowInfoCoder.of()))
.apply(Reshuffle.<ShardedKey<String>, TableRowInfo>of())
// Put in the global window to ensure that DynamicDestinations side inputs are accessed
// correctly.
http://git-wip-us.apache.org/repos/asf/beam/blob/b74644a2/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java
index 04f02fc..eb0c26f 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java
@@ -80,13 +80,11 @@ import org.apache.beam.sdk.options.ValueProvider;
import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
-import org.apache.beam.sdk.transforms.Flatten;
-import org.apache.beam.sdk.transforms.GroupByKey;
import org.apache.beam.sdk.transforms.MapElements;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.Reshuffle;
import org.apache.beam.sdk.transforms.SimpleFunction;
-import org.apache.beam.sdk.transforms.Values;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.transforms.display.DisplayData.Builder;
import org.apache.beam.sdk.transforms.display.HasDisplayData;
@@ -631,18 +629,10 @@ public class DatastoreV1 {
.apply(ParDo.of(new GqlQueryTranslateFn(v1Options)));
}
- PCollection<KV<Integer, Query>> splitQueries = inputQuery
- .apply(ParDo.of(new SplitQueryFn(v1Options, getNumQuerySplits())));
-
- PCollection<Query> shardedQueries = splitQueries
- .apply(GroupByKey.<Integer, Query>create())
- .apply(Values.<Iterable<Query>>create())
- .apply(Flatten.<Query>iterables());
-
- PCollection<Entity> entities = shardedQueries
- .apply(ParDo.of(new ReadFn(v1Options)));
-
- return entities;
+ return inputQuery
+ .apply("Split", ParDo.of(new SplitQueryFn(v1Options, getNumQuerySplits())))
+ .apply("Reshuffle", Reshuffle.<Query>viaRandomKey())
+ .apply("Read", ParDo.of(new ReadFn(v1Options)));
}
@Override
@@ -757,7 +747,7 @@ public class DatastoreV1 {
* keys and outputs them as {@link KV}.
*/
@VisibleForTesting
- static class SplitQueryFn extends DoFn<Query, KV<Integer, Query>> {
+ static class SplitQueryFn extends DoFn<Query, Query> {
private final V1Options options;
// number of splits to make for a given query
private final int numSplits;
@@ -789,12 +779,11 @@ public class DatastoreV1 {
@ProcessElement
public void processElement(ProcessContext c) throws Exception {
- int key = 1;
Query query = c.element();
// If query has a user set limit, then do not split.
if (query.hasLimit()) {
- c.output(KV.of(key, query));
+ c.output(query);
return;
}
@@ -818,7 +807,7 @@ public class DatastoreV1 {
// assign unique keys to query splits.
for (Query subquery : querySplits) {
- c.output(KV.of(key++, subquery));
+ c.output(query);
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/b74644a2/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java
index 54759fb..26ac884 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java
@@ -34,7 +34,6 @@ import com.google.common.annotations.VisibleForTesting;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
-import java.util.UUID;
import javax.annotation.Nullable;
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.options.ValueProvider;
@@ -43,10 +42,7 @@ import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.Reshuffle;
-import org.apache.beam.sdk.transforms.SerializableFunction;
-import org.apache.beam.sdk.transforms.Values;
import org.apache.beam.sdk.transforms.View;
-import org.apache.beam.sdk.transforms.WithKeys;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
@@ -287,27 +283,16 @@ public class SpannerIO {
@Override
public PCollection<Struct> expand(PCollection<ReadOperation> input) {
- PCollection<ReadOperation> reshuffled =
- input
- .apply(
- "Pair with random key",
- WithKeys.of(
- new SerializableFunction<ReadOperation, String>() {
- @Override
- public String apply(ReadOperation ignored) {
- return UUID.randomUUID().toString();
- }
- }))
- .apply("Reshuffle", Reshuffle.<String, ReadOperation>of())
- .apply("Strip keys", Values.<ReadOperation>create());
List<PCollectionView<Transaction>> sideInputs =
getTransaction() == null
? Collections.<PCollectionView<Transaction>>emptyList()
: Collections.singletonList(getTransaction());
- return reshuffled.apply(
- "Execute queries",
- ParDo.of(new NaiveSpannerReadFn(getSpannerConfig(), getTransaction()))
- .withSideInputs(sideInputs));
+ return input
+ .apply(Reshuffle.<ReadOperation>viaRandomKey())
+ .apply(
+ "Execute queries",
+ ParDo.of(new NaiveSpannerReadFn(getSpannerConfig(), getTransaction()))
+ .withSideInputs(sideInputs));
}
}
http://git-wip-us.apache.org/repos/asf/beam/blob/b74644a2/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1Test.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1Test.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1Test.java
index 729ac6b..58bab21 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1Test.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1Test.java
@@ -611,7 +611,7 @@ public class DatastoreV1Test {
.thenReturn(splitQuery(QUERY, numSplits));
SplitQueryFn splitQueryFn = new SplitQueryFn(V_1_OPTIONS, numSplits, mockDatastoreFactory);
- DoFnTester<Query, KV<Integer, Query>> doFnTester = DoFnTester.of(splitQueryFn);
+ DoFnTester<Query, Query> doFnTester = DoFnTester.of(splitQueryFn);
/**
* Although Datastore client is marked transient in {@link SplitQueryFn}, when injected through
* mock factory using a when clause for unit testing purposes, it is not serializable
@@ -619,10 +619,9 @@ public class DatastoreV1Test {
* doFn from being serialized.
*/
doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
- List<KV<Integer, Query>> queries = doFnTester.processBundle(QUERY);
+ List<Query> queries = doFnTester.processBundle(QUERY);
assertEquals(queries.size(), numSplits);
- verifyUniqueKeys(queries);
verify(mockQuerySplitter, times(1)).getSplits(
eq(QUERY), any(PartitionId.class), eq(numSplits), any(Datastore.class));
verifyZeroInteractions(mockDatastore);
@@ -657,12 +656,11 @@ public class DatastoreV1Test {
.thenReturn(splitQuery(QUERY, expectedNumSplits));
SplitQueryFn splitQueryFn = new SplitQueryFn(V_1_OPTIONS, numSplits, mockDatastoreFactory);
- DoFnTester<Query, KV<Integer, Query>> doFnTester = DoFnTester.of(splitQueryFn);
+ DoFnTester<Query, Query> doFnTester = DoFnTester.of(splitQueryFn);
doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
- List<KV<Integer, Query>> queries = doFnTester.processBundle(QUERY);
+ List<Query> queries = doFnTester.processBundle(QUERY);
assertEquals(queries.size(), expectedNumSplits);
- verifyUniqueKeys(queries);
verify(mockQuerySplitter, times(1)).getSplits(
eq(QUERY), any(PartitionId.class), eq(expectedNumSplits), any(Datastore.class));
verify(mockDatastore, times(1)).runQuery(latestTimestampRequest);
@@ -679,12 +677,11 @@ public class DatastoreV1Test {
.build();
SplitQueryFn splitQueryFn = new SplitQueryFn(V_1_OPTIONS, 10, mockDatastoreFactory);
- DoFnTester<Query, KV<Integer, Query>> doFnTester = DoFnTester.of(splitQueryFn);
+ DoFnTester<Query, Query> doFnTester = DoFnTester.of(splitQueryFn);
doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
- List<KV<Integer, Query>> queries = doFnTester.processBundle(queryWithLimit);
+ List<Query> queries = doFnTester.processBundle(queryWithLimit);
assertEquals(queries.size(), 1);
- verifyUniqueKeys(queries);
verifyNoMoreInteractions(mockDatastore);
verifyNoMoreInteractions(mockQuerySplitter);
}
http://git-wip-us.apache.org/repos/asf/beam/blob/b74644a2/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/SplitQueryFnIT.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/SplitQueryFnIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/SplitQueryFnIT.java
index 5b1066a..fa391cc 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/SplitQueryFnIT.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/SplitQueryFnIT.java
@@ -27,7 +27,6 @@ import javax.annotation.Nullable;
import org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.Read.SplitQueryFn;
import org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.Read.V1Options;
import org.apache.beam.sdk.transforms.DoFnTester;
-import org.apache.beam.sdk.values.KV;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
@@ -87,9 +86,9 @@ public class SplitQueryFnIT {
SplitQueryFn splitQueryFn = new SplitQueryFn(
V1Options.from(projectId, namespace, null), 0);
- DoFnTester<Query, KV<Integer, Query>> doFnTester = DoFnTester.of(splitQueryFn);
+ DoFnTester<Query, Query> doFnTester = DoFnTester.of(splitQueryFn);
- List<KV<Integer, Query>> queries = doFnTester.processBundle(query.build());
+ List<Query> queries = doFnTester.processBundle(query.build());
assertEquals(queries.size(), expectedNumSplits);
}
http://git-wip-us.apache.org/repos/asf/beam/blob/b74644a2/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcIO.java b/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcIO.java
index 14b766e..6ac6341 100644
--- a/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcIO.java
+++ b/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcIO.java
@@ -25,7 +25,6 @@ import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
-import java.util.concurrent.ThreadLocalRandom;
import javax.annotation.Nullable;
import javax.sql.DataSource;
import org.apache.beam.sdk.annotations.Experimental;
@@ -38,10 +37,8 @@ import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.Reshuffle;
import org.apache.beam.sdk.transforms.SerializableFunctions;
-import org.apache.beam.sdk.transforms.Values;
import org.apache.beam.sdk.transforms.View;
import org.apache.beam.sdk.transforms.display.DisplayData;
-import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionView;
@@ -647,26 +644,7 @@ public class JdbcIO {
}
})
.withSideInputs(empty));
- return materialized
- .apply(
- "Pair with random key",
- ParDo.of(
- new DoFn<T, KV<Integer, T>>() {
- private int shard;
-
- @Setup
- public void setup() {
- shard = ThreadLocalRandom.current().nextInt();
- }
-
- @ProcessElement
- public void processElement(ProcessContext context) {
- context.output(KV.of(++shard, context.element()));
- }
- }))
- .apply(Reshuffle.<Integer, T>of())
- .apply(Values.<T>create());
-
+ return materialized.apply(Reshuffle.<T>viaRandomKey());
}
}
}