add splittable URNs as primitives

cosmoskitten · Sep 10, 2018 · 80a1102 · 80a1102
1 parent d3772a8
commit 80a1102
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/...n-java/src/main/java/org/apache/beam/runners/core/construction/PTransformTranslation.java b/...n-java/src/main/java/org/apache/beam/runners/core/construction/PTransformTranslation.java
@@ -92,6 +92,12 @@ public class PTransformTranslation {
   public static final String RESHUFFLE_URN = getUrn(StandardPTransforms.Composites.RESHUFFLE);
   public static final String WRITE_FILES_TRANSFORM_URN =
       getUrn(StandardPTransforms.Composites.WRITE_FILES);
+
+  // SplittableParDoComponents
+  public static final String SPLITTABLE_PAIR_WITH_RESTRICTION_URN =
+      getUrn(SplittableParDoComponents.PAIR_WITH_RESTRICTION);
+  public static final String SPLITTABLE_SPLIT_RESTRICTION_URN =
+      getUrn(SplittableParDoComponents.SPLIT_RESTRICTION);
   public static final String SPLITTABLE_PROCESS_KEYED_URN =
       getUrn(SplittableParDoComponents.PROCESS_KEYED_ELEMENTS);
   public static final String SPLITTABLE_PROCESS_ELEMENTS_URN =

diff --git a/...java/src/main/java/org/apache/beam/runners/core/construction/graph/QueryablePipeline.java b/...java/src/main/java/org/apache/beam/runners/core/construction/graph/QueryablePipeline.java
@@ -27,6 +27,8 @@
 import static org.apache.beam.runners.core.construction.PTransformTranslation.MAP_WINDOWS_TRANSFORM_URN;
 import static org.apache.beam.runners.core.construction.PTransformTranslation.PAR_DO_TRANSFORM_URN;
 import static org.apache.beam.runners.core.construction.PTransformTranslation.READ_TRANSFORM_URN;
+import static org.apache.beam.runners.core.construction.PTransformTranslation.SPLITTABLE_PROCESS_ELEMENTS_URN;
+import static org.apache.beam.runners.core.construction.PTransformTranslation.SPLITTABLE_PROCESS_KEYED_URN;
 import static org.apache.beam.runners.core.construction.PTransformTranslation.TEST_STREAM_TRANSFORM_URN;
 
 import com.google.common.annotations.VisibleForTesting;
@@ -145,7 +147,9 @@ static Collection<String> getPrimitiveTransformIds(RunnerApi.Components componen
           TEST_STREAM_TRANSFORM_URN,
           MAP_WINDOWS_TRANSFORM_URN,
           READ_TRANSFORM_URN,
-          CREATE_VIEW_TRANSFORM_URN);
+          CREATE_VIEW_TRANSFORM_URN,
+          SPLITTABLE_PROCESS_KEYED_URN,
+          SPLITTABLE_PROCESS_ELEMENTS_URN);
 
   /** Returns true if the provided transform is a primitive. */
   private static boolean isPrimitiveTransform(PTransform transform) {