Repository: beam Updated Branches: refs/heads/master 3b04d1355 -> 0bc375634
Add fast path to ParDoTranslation for known ParDo class Project: http://git-wip-us.apache.org/repos/asf/beam/repo Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/9b30bcaa Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/9b30bcaa Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/9b30bcaa Branch: refs/heads/master Commit: 9b30bcaa707738da46bd238d3823717e266a8056 Parents: 8e391d9 Author: Kenneth Knowles <[email protected]> Authored: Tue Sep 12 16:22:18 2017 -0700 Committer: Kenneth Knowles <[email protected]> Committed: Tue Sep 12 16:22:18 2017 -0700 ---------------------------------------------------------------------- .../core/construction/ParDoTranslation.java | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/beam/blob/9b30bcaa/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 6ae95e4..5372563 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 @@ -219,6 +219,11 @@ public class ParDoTranslation { } public static DoFn<?, ?> getDoFn(AppliedPTransform<?, ?, ?> application) throws IOException { + PTransform<?, ?> transform = application.getTransform(); + if (transform instanceof ParDo.MultiOutput) { + return ((ParDo.MultiOutput<?, ?>) transform).getFn(); + } + return getDoFn(getParDoPayload(application)); } @@ -229,11 +234,20 @@ public class ParDoTranslation { public static TupleTag<?> getMainOutputTag(AppliedPTransform<?, ?, ?> application) throws IOException { + PTransform<?, ?> transform = application.getTransform(); + if (transform instanceof ParDo.MultiOutput) { + return ((ParDo.MultiOutput<?, ?>) transform).getMainOutputTag(); + } + return getMainOutputTag(getParDoPayload(application)); } public static TupleTagList getAdditionalOutputTags(AppliedPTransform<?, ?, ?> application) throws IOException { + PTransform<?, ?> transform = application.getTransform(); + if (transform instanceof ParDo.MultiOutput) { + return ((ParDo.MultiOutput<?, ?>) transform).getAdditionalOutputTags(); + } RunnerApi.PTransform protoTransform = PTransformTranslation.toProto(application, SdkComponents.create()); @@ -253,6 +267,10 @@ public class ParDoTranslation { public static List<PCollectionView<?>> getSideInputs(AppliedPTransform<?, ?, ?> application) throws IOException { + PTransform<?, ?> transform = application.getTransform(); + if (transform instanceof ParDo.MultiOutput) { + return ((ParDo.MultiOutput<?, ?>) transform).getSideInputs(); + } SdkComponents sdkComponents = SdkComponents.create(); RunnerApi.PTransform parDoProto =
