From 314e3da822ba59f8ac7133cbf47c61fe429a7361 Mon Sep 17 00:00:00 2001 From: Thomas Groh Date: Fri, 9 Dec 2016 15:52:15 -0800 Subject: [PATCH] Stop expanding PValues in DirectRunner visitors A PValue always expands to itself, and these calls are unneccessary. --- .../beam/runners/direct/DirectGraphVisitor.java | 14 ++++++-------- .../runners/direct/KeyedPValueTrackingVisitor.java | 2 +- 2 files changed, 7 insertions(+), 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 index 4f38bce1e9d98..0283d0352240d 100644 --- 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/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 4161f9e25def8..7f85169478a1c 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 @@ -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); } }