This closes #1570

apache · Dec 12, 2016 · 59f1fb2 · 59f1fb2
2 parents 5d619e8 + 8ef74a7
commit 59f1fb2
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 9 deletions.
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGraphVisitor.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectGraphVisitor.java
@@ -99,14 +99,12 @@ public void visitValue(PValue value, TransformHierarchy.Node producer) {
     if (!producers.containsKey(value)) {
       producers.put(value, appliedTransform);
     }
-    for (PValue expandedValue : value.expand()) {
-      if (expandedValue instanceof PCollectionView) {
-        views.add((PCollectionView<?>) expandedValue);
-      }
-      if (!producers.containsKey(expandedValue)) {
-        producers.put(value, appliedTransform);
-      }
-    }
+   if (value instanceof PCollectionView) {
+     views.add((PCollectionView<?>) value);
+   }
+   if (!producers.containsKey(value)) {
+     producers.put(value, appliedTransform);
+   }
   }
 
   private AppliedPTransform<?, ?, ?> getAppliedTransform(TransformHierarchy.Node node) {

diff --git a/.../direct-java/src/main/java/org/apache/beam/runners/direct/KeyedPValueTrackingVisitor.java b/.../direct-java/src/main/java/org/apache/beam/runners/direct/KeyedPValueTrackingVisitor.java
@@ -84,7 +84,7 @@ public void visitPrimitiveTransform(TransformHierarchy.Node node) {}
   @Override
   public void visitValue(PValue value, TransformHierarchy.Node producer) {
     if (producesKeyedOutputs.contains(producer.getTransform().getClass())) {
-      keyedValues.addAll(value.expand());
+      keyedValues.add(value);
     }
   }