You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by ke...@apache.org on 2017/06/06 20:28:10 UTC
[1/3] beam git commit: Fix RawPTransform translation
Repository: beam
Updated Branches:
refs/heads/master 7c608c32a -> 513c952fa
Fix RawPTransform translation
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/840492d9
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/840492d9
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/840492d9
Branch: refs/heads/master
Commit: 840492d9d8fb3b08cfe70a525655759fc1a31fdf
Parents: 7c608c3
Author: Kenneth Knowles <kl...@google.com>
Authored: Fri May 26 14:18:03 2017 -0700
Committer: Kenneth Knowles <kl...@google.com>
Committed: Tue Jun 6 13:10:33 2017 -0700
----------------------------------------------------------------------
.../construction/PTransformTranslation.java | 57 ++++++++++++++++----
runners/direct-java/pom.xml | 5 --
.../beam/runners/direct/DirectGroupByKey.java | 5 +-
.../direct/ParDoMultiOverrideFactory.java | 3 +-
.../direct/TestStreamEvaluatorFactory.java | 3 +-
.../direct/TransformEvaluatorRegistry.java | 8 +--
.../runners/direct/ViewOverrideFactory.java | 3 +-
7 files changed, 56 insertions(+), 28 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/840492d9/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PTransformTranslation.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PTransformTranslation.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PTransformTranslation.java
index 7c5c593..32ecf43 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PTransformTranslation.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PTransformTranslation.java
@@ -22,7 +22,6 @@ import static com.google.common.base.Preconditions.checkArgument;
import com.google.common.collect.ImmutableMap;
import com.google.protobuf.Any;
-import com.google.protobuf.Message;
import java.io.IOException;
import java.util.Collections;
import java.util.List;
@@ -115,7 +114,20 @@ public class PTransformTranslation {
// TODO: Display Data
PTransform<?, ?> transform = appliedPTransform.getTransform();
- if (KNOWN_PAYLOAD_TRANSLATORS.containsKey(transform.getClass())) {
+ // A RawPTransform directly vends its payload. Because it will generally be
+ // a subclass, we cannot do dictionary lookup in KNOWN_PAYLOAD_TRANSLATORS.
+ if (transform instanceof RawPTransform) {
+ RawPTransform<?, ?> rawPTransform = (RawPTransform<?, ?>) transform;
+
+ if (rawPTransform.getUrn() != null) {
+ FunctionSpec.Builder payload = FunctionSpec.newBuilder().setUrn(rawPTransform.getUrn());
+ @Nullable Any parameter = rawPTransform.getPayload();
+ if (parameter != null) {
+ payload.setParameter(parameter);
+ }
+ transformBuilder.setSpec(payload);
+ }
+ } else if (KNOWN_PAYLOAD_TRANSLATORS.containsKey(transform.getClass())) {
FunctionSpec payload =
KNOWN_PAYLOAD_TRANSLATORS
.get(transform.getClass())
@@ -145,6 +157,25 @@ public class PTransformTranslation {
}
/**
+ * Returns the URN for the transform if it is known, otherwise {@code null}.
+ */
+ @Nullable
+ public static String urnForTransformOrNull(PTransform<?, ?> transform) {
+
+ // A RawPTransform directly vends its URN. Because it will generally be
+ // a subclass, we cannot do dictionary lookup in KNOWN_PAYLOAD_TRANSLATORS.
+ if (transform instanceof RawPTransform) {
+ return ((RawPTransform) transform).getUrn();
+ }
+
+ TransformPayloadTranslator translator = KNOWN_PAYLOAD_TRANSLATORS.get(transform.getClass());
+ if (translator == null) {
+ return null;
+ }
+ return translator.getUrn(transform);
+ }
+
+ /**
* Returns the URN for the transform if it is known, otherwise throws.
*/
public static String urnForTransform(PTransform<?, ?> transform) {
@@ -176,13 +207,14 @@ public class PTransformTranslation {
* fully expanded in the pipeline proto.
*/
public abstract static class RawPTransform<
- InputT extends PInput, OutputT extends POutput, PayloadT extends Message>
+ InputT extends PInput, OutputT extends POutput>
extends PTransform<InputT, OutputT> {
+ @Nullable
public abstract String getUrn();
@Nullable
- PayloadT getPayload() {
+ public Any getPayload() {
return null;
}
}
@@ -190,24 +222,29 @@ public class PTransformTranslation {
/**
* A translator that uses the explicit URN and payload from a {@link RawPTransform}.
*/
- public static class RawPTransformTranslator<PayloadT extends Message>
- implements TransformPayloadTranslator<RawPTransform<?, ?, PayloadT>> {
+ public static class RawPTransformTranslator
+ implements TransformPayloadTranslator<RawPTransform<?, ?>> {
@Override
- public String getUrn(RawPTransform<?, ?, PayloadT> transform) {
+ public String getUrn(RawPTransform<?, ?> transform) {
return transform.getUrn();
}
@Override
public FunctionSpec translate(
- AppliedPTransform<?, ?, RawPTransform<?, ?, PayloadT>> transform,
+ AppliedPTransform<?, ?, RawPTransform<?, ?>> transform,
SdkComponents components) {
- PayloadT payload = transform.getTransform().getPayload();
+
+ // Anonymous composites have no spec
+ if (transform.getTransform().getUrn() == null) {
+ return null;
+ }
FunctionSpec.Builder transformSpec =
FunctionSpec.newBuilder().setUrn(getUrn(transform.getTransform()));
+ Any payload = transform.getTransform().getPayload();
if (payload != null) {
- transformSpec.setParameter(Any.pack(payload));
+ transformSpec.setParameter(payload);
}
return transformSpec.build();
http://git-wip-us.apache.org/repos/asf/beam/blob/840492d9/runners/direct-java/pom.xml
----------------------------------------------------------------------
diff --git a/runners/direct-java/pom.xml b/runners/direct-java/pom.xml
index cba4b09..bec2113 100644
--- a/runners/direct-java/pom.xml
+++ b/runners/direct-java/pom.xml
@@ -208,11 +208,6 @@
</dependency>
<dependency>
- <groupId>com.google.protobuf</groupId>
- <artifactId>protobuf-java</artifactId>
- </dependency>
-
- <dependency>
<groupId>joda-time</groupId>
<artifactId>joda-time</artifactId>
</dependency>
http://git-wip-us.apache.org/repos/asf/beam/blob/840492d9/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGroupByKey.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGroupByKey.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGroupByKey.java
index f239070..2fc0dd4 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGroupByKey.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGroupByKey.java
@@ -20,7 +20,6 @@ package org.apache.beam.runners.direct;
import static com.google.common.base.Preconditions.checkArgument;
-import com.google.protobuf.Message;
import org.apache.beam.runners.core.KeyedWorkItem;
import org.apache.beam.runners.core.KeyedWorkItemCoder;
import org.apache.beam.runners.core.construction.ForwardingPTransform;
@@ -74,7 +73,7 @@ class DirectGroupByKey<K, V>
static final class DirectGroupByKeyOnly<K, V>
extends PTransformTranslation.RawPTransform<
- PCollection<KV<K, V>>, PCollection<KeyedWorkItem<K, V>>, Message> {
+ PCollection<KV<K, V>>, PCollection<KeyedWorkItem<K, V>>> {
@Override
public PCollection<KeyedWorkItem<K, V>> expand(PCollection<KV<K, V>> input) {
return PCollection.createPrimitiveOutputInternal(
@@ -101,7 +100,7 @@ class DirectGroupByKey<K, V>
static final class DirectGroupAlsoByWindow<K, V>
extends PTransformTranslation.RawPTransform<
- PCollection<KeyedWorkItem<K, V>>, PCollection<KV<K, Iterable<V>>>, Message> {
+ PCollection<KeyedWorkItem<K, V>>, PCollection<KV<K, Iterable<V>>>> {
private final WindowingStrategy<?, ?> inputWindowingStrategy;
private final WindowingStrategy<?, ?> outputWindowingStrategy;
http://git-wip-us.apache.org/repos/asf/beam/blob/840492d9/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoMultiOverrideFactory.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoMultiOverrideFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoMultiOverrideFactory.java
index df2054b..858ea34 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoMultiOverrideFactory.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ParDoMultiOverrideFactory.java
@@ -19,7 +19,6 @@ package org.apache.beam.runners.direct;
import static com.google.common.base.Preconditions.checkState;
-import com.google.protobuf.Message;
import java.util.Map;
import org.apache.beam.runners.core.KeyedWorkItem;
import org.apache.beam.runners.core.KeyedWorkItemCoder;
@@ -172,7 +171,7 @@ class ParDoMultiOverrideFactory<InputT, OutputT>
static class StatefulParDo<K, InputT, OutputT>
extends PTransformTranslation.RawPTransform<
- PCollection<? extends KeyedWorkItem<K, KV<K, InputT>>>, PCollectionTuple, Message> {
+ PCollection<? extends KeyedWorkItem<K, KV<K, InputT>>>, PCollectionTuple> {
private final transient MultiOutput<KV<K, InputT>, OutputT> underlyingParDo;
private final transient PCollection<KV<K, InputT>> originalInput;
http://git-wip-us.apache.org/repos/asf/beam/blob/840492d9/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TestStreamEvaluatorFactory.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TestStreamEvaluatorFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TestStreamEvaluatorFactory.java
index b1db58f..2da7a71 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TestStreamEvaluatorFactory.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TestStreamEvaluatorFactory.java
@@ -22,7 +22,6 @@ import com.google.auto.value.AutoValue;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Supplier;
import com.google.common.collect.Iterables;
-import com.google.protobuf.Message;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
@@ -185,7 +184,7 @@ class TestStreamEvaluatorFactory implements TransformEvaluatorFactory {
static final String DIRECT_TEST_STREAM_URN = "urn:beam:directrunner:transforms:test_stream:v1";
static class DirectTestStream<T>
- extends PTransformTranslation.RawPTransform<PBegin, PCollection<T>, Message> {
+ extends PTransformTranslation.RawPTransform<PBegin, PCollection<T>> {
private final transient DirectRunner runner;
private final TestStream<T> original;
http://git-wip-us.apache.org/repos/asf/beam/blob/840492d9/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TransformEvaluatorRegistry.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TransformEvaluatorRegistry.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TransformEvaluatorRegistry.java
index d144b20..0c907df 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TransformEvaluatorRegistry.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/TransformEvaluatorRegistry.java
@@ -92,17 +92,17 @@ class TransformEvaluatorRegistry implements TransformEvaluatorFactory {
.<Class<? extends PTransform>, PTransformTranslation.TransformPayloadTranslator>builder()
.put(
DirectGroupByKey.DirectGroupByKeyOnly.class,
- new PTransformTranslation.RawPTransformTranslator<>())
+ new PTransformTranslation.RawPTransformTranslator())
.put(
DirectGroupByKey.DirectGroupAlsoByWindow.class,
new PTransformTranslation.RawPTransformTranslator())
.put(
ParDoMultiOverrideFactory.StatefulParDo.class,
- new PTransformTranslation.RawPTransformTranslator<>())
+ new PTransformTranslation.RawPTransformTranslator())
.put(
ViewOverrideFactory.WriteView.class,
- new PTransformTranslation.RawPTransformTranslator<>())
- .put(DirectTestStream.class, new PTransformTranslation.RawPTransformTranslator<>())
+ new PTransformTranslation.RawPTransformTranslator())
+ .put(DirectTestStream.class, new PTransformTranslation.RawPTransformTranslator())
.put(
SplittableParDoViaKeyedWorkItems.ProcessElements.class,
new SplittableParDoProcessElementsTranslator())
http://git-wip-us.apache.org/repos/asf/beam/blob/840492d9/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ViewOverrideFactory.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ViewOverrideFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ViewOverrideFactory.java
index 501b436..fdff63d 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ViewOverrideFactory.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/ViewOverrideFactory.java
@@ -18,7 +18,6 @@
package org.apache.beam.runners.direct;
-import com.google.protobuf.Message;
import java.util.Collections;
import java.util.Map;
import org.apache.beam.runners.core.construction.ForwardingPTransform;
@@ -95,7 +94,7 @@ class ViewOverrideFactory<ElemT, ViewT>
* to {@link ViewT}.
*/
static final class WriteView<ElemT, ViewT>
- extends RawPTransform<PCollection<Iterable<ElemT>>, PCollectionView<ViewT>, Message> {
+ extends RawPTransform<PCollection<Iterable<ElemT>>, PCollectionView<ViewT>> {
private final CreatePCollectionView<ElemT, ViewT> og;
WriteView(CreatePCollectionView<ElemT, ViewT> og) {
[2/3] beam git commit: Improve Splittable ParDo translation
Posted by ke...@apache.org.
Improve Splittable ParDo translation
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/1b00d95a
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/1b00d95a
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/1b00d95a
Branch: refs/heads/master
Commit: 1b00d95a1105d2611b985dc463da0884a6646354
Parents: 840492d
Author: Kenneth Knowles <kl...@google.com>
Authored: Thu May 25 06:29:16 2017 -0700
Committer: Kenneth Knowles <kl...@google.com>
Committed: Tue Jun 6 13:13:12 2017 -0700
----------------------------------------------------------------------
.../core/construction/ParDoTranslation.java | 20 +++++++++++
.../core/construction/SplittableParDo.java | 18 ++++++++--
.../core/construction/ParDoTranslationTest.java | 35 +++++++++++++++++++-
.../core/SplittableParDoViaKeyedWorkItems.java | 10 +++++-
.../direct/KeyedPValueTrackingVisitor.java | 2 +-
.../src/main/proto/beam_runner_api.proto | 3 ++
6 files changed, 82 insertions(+), 6 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/1b00d95a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ParDoTranslation.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ParDoTranslation.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ParDoTranslation.java
index fe66179..34e0d86 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ParDoTranslation.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ParDoTranslation.java
@@ -144,6 +144,7 @@ public class ParDoTranslation {
ParDoPayload.Builder builder = ParDoPayload.newBuilder();
builder.setDoFn(toProto(parDo.getFn(), parDo.getMainOutputTag()));
+ builder.setSplittable(signature.processElement().isSplittable());
for (PCollectionView<?> sideInput : parDo.getSideInputs()) {
builder.putSideInputs(sideInput.getTagInternal().getId(), toProto(sideInput));
}
@@ -496,6 +497,25 @@ public class ParDoTranslation {
.build();
}
+ private static <T> ParDoPayload getParDoPayload(AppliedPTransform<?, ?, ?> transform)
+ throws IOException {
+ return PTransformTranslation.toProto(
+ transform, Collections.<AppliedPTransform<?, ?, ?>>emptyList(), SdkComponents.create())
+ .getSpec()
+ .getParameter()
+ .unpack(ParDoPayload.class);
+ }
+
+ public static boolean usesStateOrTimers(AppliedPTransform<?, ?, ?> transform) throws IOException {
+ ParDoPayload payload = getParDoPayload(transform);
+ return payload.getStateSpecsCount() > 0 || payload.getTimerSpecsCount() > 0;
+ }
+
+ public static boolean isSplittable(AppliedPTransform<?, ?, ?> transform) throws IOException {
+ ParDoPayload payload = getParDoPayload(transform);
+ return payload.getSplittable();
+ }
+
private static ViewFn<?, ?> viewFnFromProto(SdkFunctionSpec viewFn)
throws InvalidProtocolBufferException {
FunctionSpec spec = viewFn.getSpec();
http://git-wip-us.apache.org/repos/asf/beam/blob/1b00d95a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SplittableParDo.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SplittableParDo.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SplittableParDo.java
index dfca7d2..665e39d 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SplittableParDo.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/SplittableParDo.java
@@ -22,6 +22,7 @@ import static com.google.common.base.Preconditions.checkNotNull;
import java.util.List;
import java.util.UUID;
+import org.apache.beam.runners.core.construction.PTransformTranslation.RawPTransform;
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.transforms.DoFn;
@@ -67,6 +68,12 @@ public class SplittableParDo<InputT, OutputT, RestrictionT>
public static final String SPLITTABLE_PROCESS_URN =
"urn:beam:runners_core:transforms:splittable_process:v1";
+ public static final String SPLITTABLE_PROCESS_KEYED_ELEMENTS_URN =
+ "urn:beam:runners_core:transforms:splittable_process_keyed_elements:v1";
+
+ public static final String SPLITTABLE_GBKIKWI_URN =
+ "urn:beam:runners_core:transforms:splittable_gbkikwi:v1";
+
/**
* Creates the transform for the given original multi-output {@link ParDo}.
*
@@ -133,11 +140,11 @@ public class SplittableParDo<InputT, OutputT, RestrictionT>
/**
* Runner-specific primitive {@link PTransform} that invokes the {@link DoFn.ProcessElement}
- * method for a splittable {@link DoFn} on each {@link ElementAndRestriction} of the input
- * {@link PCollection} of {@link KV KVs} keyed with arbitrary but globally unique keys.
+ * method for a splittable {@link DoFn} on each {@link ElementAndRestriction} of the input {@link
+ * PCollection} of {@link KV KVs} keyed with arbitrary but globally unique keys.
*/
public static class ProcessKeyedElements<InputT, OutputT, RestrictionT>
- extends PTransform<
+ extends RawPTransform<
PCollection<KV<String, ElementAndRestriction<InputT, RestrictionT>>>, PCollectionTuple> {
private final DoFn<InputT, OutputT> fn;
private final Coder<InputT> elementCoder;
@@ -227,6 +234,11 @@ public class SplittableParDo<InputT, OutputT, RestrictionT>
return outputs;
}
+
+ @Override
+ public String getUrn() {
+ return SPLITTABLE_PROCESS_KEYED_ELEMENTS_URN;
+ }
}
/**
http://git-wip-us.apache.org/repos/asf/beam/blob/1b00d95a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ParDoTranslationTest.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ParDoTranslationTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ParDoTranslationTest.java
index 46f6a80..a8490bf 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ParDoTranslationTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ParDoTranslationTest.java
@@ -51,6 +51,7 @@ import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.ParDo.MultiOutput;
import org.apache.beam.sdk.transforms.View;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
@@ -111,7 +112,11 @@ public class ParDoTranslationTest {
ParDo.of(new DropElementsFn())
.withOutputTags(
new TupleTag<Void>(),
- TupleTagList.of(new TupleTag<byte[]>() {}).and(new TupleTag<Integer>() {})));
+ TupleTagList.of(new TupleTag<byte[]>() {}).and(new TupleTag<Integer>() {})),
+ ParDo.of(new SplittableDropElementsFn())
+ .withOutputTags(
+ new TupleTag<Void>(),
+ TupleTagList.empty()));
}
@Parameter(0)
@@ -235,6 +240,34 @@ public class ParDoTranslationTest {
}
}
+ private static class SplittableDropElementsFn extends DoFn<KV<Long, String>, Void> {
+ @ProcessElement
+ public void proc(ProcessContext context, RestrictionTracker<Integer> restriction) {
+ context.output(null);
+ }
+
+ @GetInitialRestriction
+ public Integer restriction(KV<Long, String> elem) {
+ return 42;
+ }
+
+ @NewTracker
+ public RestrictionTracker<Integer> newTracker(Integer restriction) {
+ throw new UnsupportedOperationException("Should never be called; only to test translation");
+ }
+
+
+ @Override
+ public boolean equals(Object other) {
+ return other instanceof SplittableDropElementsFn;
+ }
+
+ @Override
+ public int hashCode() {
+ return SplittableDropElementsFn.class.hashCode();
+ }
+ }
+
@SuppressWarnings("unused")
private static class StateTimerDropElementsFn extends DoFn<KV<Long, String>, Void> {
private static final String BAG_STATE_ID = "bagState";
http://git-wip-us.apache.org/repos/asf/beam/blob/1b00d95a/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDoViaKeyedWorkItems.java
----------------------------------------------------------------------
diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDoViaKeyedWorkItems.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDoViaKeyedWorkItems.java
index b38e364..c4b086a 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDoViaKeyedWorkItems.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDoViaKeyedWorkItems.java
@@ -23,7 +23,9 @@ import java.util.List;
import java.util.Map;
import org.apache.beam.runners.core.construction.ElementAndRestriction;
import org.apache.beam.runners.core.construction.PTransformReplacements;
+import org.apache.beam.runners.core.construction.PTransformTranslation.RawPTransform;
import org.apache.beam.runners.core.construction.ReplacementOutputs;
+import org.apache.beam.runners.core.construction.SplittableParDo;
import org.apache.beam.runners.core.construction.SplittableParDo.ProcessKeyedElements;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.KvCoder;
@@ -67,12 +69,18 @@ public class SplittableParDoViaKeyedWorkItems {
* emits output immediately.
*/
public static class GBKIntoKeyedWorkItems<KeyT, InputT>
- extends PTransform<PCollection<KV<KeyT, InputT>>, PCollection<KeyedWorkItem<KeyT, InputT>>> {
+ extends RawPTransform<
+ PCollection<KV<KeyT, InputT>>, PCollection<KeyedWorkItem<KeyT, InputT>>> {
@Override
public PCollection<KeyedWorkItem<KeyT, InputT>> expand(PCollection<KV<KeyT, InputT>> input) {
return PCollection.createPrimitiveOutputInternal(
input.getPipeline(), WindowingStrategy.globalDefault(), input.isBounded());
}
+
+ @Override
+ public String getUrn() {
+ return SplittableParDo.SPLITTABLE_GBKIKWI_URN;
+ }
}
/** Overrides a {@link ProcessKeyedElements} into {@link SplittableProcessViaKeyedWorkItems}. */
http://git-wip-us.apache.org/repos/asf/beam/blob/1b00d95a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/KeyedPValueTrackingVisitor.java
----------------------------------------------------------------------
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/KeyedPValueTrackingVisitor.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/KeyedPValueTrackingVisitor.java
index f9b2dae..6eadaba 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/KeyedPValueTrackingVisitor.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/KeyedPValueTrackingVisitor.java
@@ -47,7 +47,7 @@ import org.apache.beam.sdk.values.TupleTag;
class KeyedPValueTrackingVisitor extends PipelineVisitor.Defaults {
private static final Set<Class<? extends PTransform>> PRODUCES_KEYED_OUTPUTS =
- ImmutableSet.of(
+ ImmutableSet.<Class<? extends PTransform>>of(
SplittableParDoViaKeyedWorkItems.GBKIntoKeyedWorkItems.class,
DirectGroupByKeyOnly.class,
DirectGroupAlsoByWindow.class);
http://git-wip-us.apache.org/repos/asf/beam/blob/1b00d95a/sdks/common/runner-api/src/main/proto/beam_runner_api.proto
----------------------------------------------------------------------
diff --git a/sdks/common/runner-api/src/main/proto/beam_runner_api.proto b/sdks/common/runner-api/src/main/proto/beam_runner_api.proto
index 87e33f0..039ecb0 100644
--- a/sdks/common/runner-api/src/main/proto/beam_runner_api.proto
+++ b/sdks/common/runner-api/src/main/proto/beam_runner_api.proto
@@ -220,6 +220,9 @@ message ParDoPayload {
// (Optional) A mapping of local timer names to timer specifications.
map<string, TimerSpec> timer_specs = 5;
+
+ // Whether the DoFn is splittable
+ bool splittable = 6;
}
// Parameters that a UDF might require.
[3/3] beam git commit: This closes #3282: [BEAM-3271] Improve
Splittable ParDo translation
Posted by ke...@apache.org.
This closes #3282: [BEAM-3271] Improve Splittable ParDo translation
Improve Splittable ParDo translation
Fix RawPTransform translation
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/513c952f
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/513c952f
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/513c952f
Branch: refs/heads/master
Commit: 513c952fa65dbddee167160f6efabb634133425d
Parents: 7c608c3 1b00d95
Author: Kenneth Knowles <kl...@google.com>
Authored: Tue Jun 6 13:13:40 2017 -0700
Committer: Kenneth Knowles <kl...@google.com>
Committed: Tue Jun 6 13:13:40 2017 -0700
----------------------------------------------------------------------
.../construction/PTransformTranslation.java | 57 ++++++++++++++++----
.../core/construction/ParDoTranslation.java | 20 +++++++
.../core/construction/SplittableParDo.java | 18 +++++--
.../core/construction/ParDoTranslationTest.java | 35 +++++++++++-
.../core/SplittableParDoViaKeyedWorkItems.java | 10 +++-
runners/direct-java/pom.xml | 5 --
.../beam/runners/direct/DirectGroupByKey.java | 5 +-
.../direct/KeyedPValueTrackingVisitor.java | 2 +-
.../direct/ParDoMultiOverrideFactory.java | 3 +-
.../direct/TestStreamEvaluatorFactory.java | 3 +-
.../direct/TransformEvaluatorRegistry.java | 8 +--
.../runners/direct/ViewOverrideFactory.java | 3 +-
.../src/main/proto/beam_runner_api.proto | 3 ++
13 files changed, 138 insertions(+), 34 deletions(-)
----------------------------------------------------------------------