Skip to content

Commit

Permalink
This closes #1495
Browse files Browse the repository at this point in the history
  • Loading branch information
@tgroh
tgroh committed Dec 6, 2016
2 parents 077d911 + ec1eff3 commit ca6ab6c
Show file tree
Hide file tree
Showing 17 changed files with 292 additions and 237 deletions.
18 changes: 10 additions & 8 deletions ...ct-java/src/test/java/org/apache/beam/runners/direct/BoundedReadEvaluatorFactoryTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,7 @@ public class BoundedReadEvaluatorFactoryTest {
private BoundedReadEvaluatorFactory factory;
@Mock private EvaluationContext context;
private BundleFactory bundleFactory;
private AppliedPTransform<?, ?, ?> longsProducer;

@Before
public void setup() {
Expand All @@ -92,6 +93,7 @@ public void setup() {
new BoundedReadEvaluatorFactory(
context, Long.MAX_VALUE /* minimum size for dynamic splits */);
bundleFactory = ImmutableListBundleFactory.create();
longsProducer = DirectGraphs.getProducer(longs);
}

@Test
Expand All @@ -102,11 +104,11 @@ public void boundedSourceInMemoryTransformEvaluatorProducesElements() throws Exc

Collection<CommittedBundle<?>> initialInputs =
new BoundedReadEvaluatorFactory.InputProvider(context)
.getInitialInputs(longs.getProducingTransformInternal(), 1);
.getInitialInputs(longsProducer, 1);
List<WindowedValue<?>> outputs = new ArrayList<>();
for (CommittedBundle<?> shardBundle : initialInputs) {
TransformEvaluator<?> evaluator =
factory.forApplication(longs.getProducingTransformInternal(), null);
factory.forApplication(longsProducer, null);
for (WindowedValue<?> shard : shardBundle.getElements()) {
evaluator.processElement((WindowedValue) shard);
}
Expand Down Expand Up @@ -141,7 +143,7 @@ public void boundedSourceEvaluatorProducesDynamicSplits() throws Exception {
}
PCollection<Long> read =
TestPipeline.create().apply(Read.from(new TestSource<>(VarLongCoder.of(), 5, elems)));
AppliedPTransform<?, ?, ?> transform = read.getProducingTransformInternal();
AppliedPTransform<?, ?, ?> transform = DirectGraphs.getProducer(read);
Collection<CommittedBundle<?>> unreadInputs =
new BoundedReadEvaluatorFactory.InputProvider(context).getInitialInputs(transform, 1);

Expand Down Expand Up @@ -191,7 +193,7 @@ public void boundedSourceEvaluatorDynamicSplitsUnsplittable() throws Exception {
PCollection<Long> read =
TestPipeline.create()
.apply(Read.from(SourceTestUtils.toUnsplittableSource(CountingSource.upTo(10L))));
AppliedPTransform<?, ?, ?> transform = read.getProducingTransformInternal();
AppliedPTransform<?, ?, ?> transform = DirectGraphs.getProducer(read);

when(context.createRootBundle()).thenReturn(bundleFactory.createRootBundle());
when(context.createRootBundle()).thenReturn(bundleFactory.createRootBundle());
Expand Down Expand Up @@ -238,7 +240,7 @@ public UncommittedBundle<?> answer(InvocationOnMock invocation) throws Throwable
});
Collection<CommittedBundle<?>> initialInputs =
new BoundedReadEvaluatorFactory.InputProvider(context)
.getInitialInputs(longs.getProducingTransformInternal(), 3);
.getInitialInputs(longsProducer, 3);

assertThat(initialInputs, hasSize(allOf(greaterThanOrEqualTo(3), lessThanOrEqualTo(4))));

Expand Down Expand Up @@ -271,7 +273,7 @@ public void boundedSourceInMemoryTransformEvaluatorShardsOfSource() throws Excep
CommittedBundle<BoundedSourceShard<Long>> shards = rootBundle.commit(Instant.now());

TransformEvaluator<BoundedSourceShard<Long>> evaluator =
factory.forApplication(longs.getProducingTransformInternal(), shards);
factory.forApplication(longsProducer, shards);
for (WindowedValue<BoundedSourceShard<Long>> shard : shards.getElements()) {
UncommittedBundle<Long> outputBundle = bundleFactory.createBundle(longs);
when(context.createBundle(longs)).thenReturn(outputBundle);
Expand Down Expand Up @@ -299,7 +301,7 @@ public void boundedSourceEvaluatorClosesReader() throws Exception {

TestPipeline p = TestPipeline.create();
PCollection<Long> pcollection = p.apply(Read.from(source));
AppliedPTransform<?, ?, ?> sourceTransform = pcollection.getProducingTransformInternal();
AppliedPTransform<?, ?, ?> sourceTransform = DirectGraphs.getProducer(pcollection);

UncommittedBundle<Long> output = bundleFactory.createBundle(pcollection);
when(context.createBundle(pcollection)).thenReturn(output);
Expand All @@ -320,7 +322,7 @@ public void boundedSourceEvaluatorNoElementsClosesReader() throws Exception {

TestPipeline p = TestPipeline.create();
PCollection<Long> pcollection = p.apply(Read.from(source));
AppliedPTransform<?, ?, ?> sourceTransform = pcollection.getProducingTransformInternal();
AppliedPTransform<?, ?, ?> sourceTransform = DirectGraphs.getProducer(pcollection);

UncommittedBundle<Long> output = bundleFactory.createBundle(pcollection);
when(context.createBundle(pcollection)).thenReturn(output);
Expand Down
69 changes: 47 additions & 22 deletions runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectGraphVisitorTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,25 +19,34 @@

import static org.hamcrest.Matchers.emptyIterable;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.hasSize;
import static org.hamcrest.Matchers.is;
import static org.junit.Assert.assertThat;
import static org.junit.Assert.assertTrue;

import com.google.common.collect.Iterables;
import java.io.Serializable;
import java.util.ArrayList;
import java.util.List;
import org.apache.beam.sdk.io.CountingInput;
import org.apache.beam.sdk.io.CountingSource;
import org.apache.beam.sdk.io.Read;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.AppliedPTransform;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.Flatten;
import org.apache.beam.sdk.transforms.Flatten.FlattenPCollectionList;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.View;
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionList;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.PDone;
import org.apache.beam.sdk.values.PInput;
import org.apache.beam.sdk.values.POutput;
import org.hamcrest.Matchers;
import org.junit.Rule;
import org.junit.Test;
Expand Down Expand Up @@ -80,26 +89,36 @@ public void processElement(DoFn<String, String>.ProcessContext c)
@Test
public void getRootTransformsContainsPBegins() {
PCollection<String> created = p.apply(Create.of("foo", "bar"));
PCollection<Long> counted = p.apply(CountingInput.upTo(1234L));
PCollection<Long> counted = p.apply(Read.from(CountingSource.upTo(1234L)));
PCollection<Long> unCounted = p.apply(CountingInput.unbounded());
p.traverseTopologically(visitor);
DirectGraph graph = visitor.getGraph();
assertThat(graph.getRootTransforms(), hasSize(3));
List<PTransform<?, ?>> unapplied = new ArrayList<>();
assertThat(
visitor.getGraph().getRootTransforms(),
graph.getRootTransforms(),
Matchers.<AppliedPTransform<?, ?, ?>>containsInAnyOrder(
created.getProducingTransformInternal(),
counted.getProducingTransformInternal(),
unCounted.getProducingTransformInternal()));
graph.getProducer(created), graph.getProducer(counted), graph.getProducer(unCounted)));
for (AppliedPTransform<?, ?, ?> root : graph.getRootTransforms()) {
assertTrue(root.getInput() instanceof PBegin);
assertThat(root.getOutput(), Matchers.<POutput>isOneOf(created, counted, unCounted));
}
}

@Test
public void getRootTransformsContainsEmptyFlatten() {
PCollection<String> empty =
PCollectionList.<String>empty(p).apply(Flatten.<String>pCollections());
FlattenPCollectionList<String> flatten = Flatten.pCollections();
PCollectionList<String> emptyList = PCollectionList.empty(p);
PCollection<String> empty = emptyList.apply(flatten);
p.traverseTopologically(visitor);
DirectGraph graph = visitor.getGraph();
assertThat(
visitor.getGraph().getRootTransforms(),
Matchers.<AppliedPTransform<?, ?, ?>>containsInAnyOrder(
empty.getProducingTransformInternal()));
graph.getRootTransforms(),
Matchers.<AppliedPTransform<?, ?, ?>>containsInAnyOrder(graph.getProducer(empty)));
AppliedPTransform<?, ?, ?> onlyRoot = Iterables.getOnlyElement(graph.getRootTransforms());
assertThat(onlyRoot.getTransform(), Matchers.<PTransform<?, ?>>equalTo(flatten));
assertThat(onlyRoot.getInput(), Matchers.<PInput>equalTo(emptyList));
assertThat(onlyRoot.getOutput(), Matchers.<POutput>equalTo(empty));
}

@Test
Expand All @@ -121,16 +140,20 @@ public void processElement(DoFn<String, String>.ProcessContext c)

p.traverseTopologically(visitor);

DirectGraph graph = visitor.getGraph();
AppliedPTransform<?, ?, ?> transformedProducer =
graph.getProducer(transformed);
AppliedPTransform<?, ?, ?> flattenedProducer =
graph.getProducer(flattened);

assertThat(
visitor.getGraph().getPrimitiveConsumers(created),
graph.getPrimitiveConsumers(created),
Matchers.<AppliedPTransform<?, ?, ?>>containsInAnyOrder(
transformed.getProducingTransformInternal(),
flattened.getProducingTransformInternal()));
transformedProducer, flattenedProducer));
assertThat(
visitor.getGraph().getPrimitiveConsumers(transformed),
Matchers.<AppliedPTransform<?, ?, ?>>containsInAnyOrder(
flattened.getProducingTransformInternal()));
assertThat(visitor.getGraph().getPrimitiveConsumers(flattened), emptyIterable());
graph.getPrimitiveConsumers(transformed),
Matchers.<AppliedPTransform<?, ?, ?>>containsInAnyOrder(flattenedProducer));
assertThat(graph.getPrimitiveConsumers(flattened), emptyIterable());
}

@Test
Expand All @@ -142,12 +165,14 @@ public void getValueToConsumersWithDuplicateInputSucceeds() {

p.traverseTopologically(visitor);

DirectGraph graph = visitor.getGraph();
AppliedPTransform<?, ?, ?> flattenedProducer = graph.getProducer(flattened);

assertThat(
visitor.getGraph().getPrimitiveConsumers(created),
Matchers.<AppliedPTransform<?, ?, ?>>containsInAnyOrder(
flattened.getProducingTransformInternal(),
flattened.getProducingTransformInternal()));
assertThat(visitor.getGraph().getPrimitiveConsumers(flattened), emptyIterable());
graph.getPrimitiveConsumers(created),
Matchers.<AppliedPTransform<?, ?, ?>>containsInAnyOrder(flattenedProducer,
flattenedProducer));
assertThat(graph.getPrimitiveConsumers(flattened), emptyIterable());
}

@Test
Expand Down
35 changes: 35 additions & 0 deletions runners/direct-java/src/test/java/org/apache/beam/runners/direct/DirectGraphs.java
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.beam.runners.direct;

import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.transforms.AppliedPTransform;
import org.apache.beam.sdk.values.PValue;

/** Test utilities for the {@link DirectRunner}. */
final class DirectGraphs {
public static DirectGraph getGraph(Pipeline p) {
DirectGraphVisitor visitor = new DirectGraphVisitor();
p.traverseTopologically(visitor);
return visitor.getGraph();
}

public static AppliedPTransform<?, ?, ?> getProducer(PValue value) {
return getGraph(value.getPipeline()).getProducer(value);
}
}

0 comments on commit ca6ab6c

Please sign in to comment.