use of org.apache.beam.sdk.runners.AppliedPTransform in project beam by apache.
the class WindowIntoTranslationTest method testToFromProto.
@Test
public void testToFromProto() throws InvalidProtocolBufferException {
pipeline.apply(GenerateSequence.from(0)).apply(Window.<Long>into((WindowFn) windowFn));
final AtomicReference<AppliedPTransform<?, ?, Assign<?>>> assign = new AtomicReference<>(null);
pipeline.traverseTopologically(new PipelineVisitor.Defaults() {
@Override
public void visitPrimitiveTransform(Node node) {
if (node.getTransform() instanceof Window.Assign) {
checkState(assign.get() == null);
assign.set((AppliedPTransform<?, ?, Assign<?>>) node.toAppliedPTransform(getPipeline()));
}
}
});
checkState(assign.get() != null);
SdkComponents components = SdkComponents.create();
WindowIntoPayload payload = WindowIntoTranslation.toProto(assign.get().getTransform(), components);
assertEquals(windowFn, WindowIntoTranslation.getWindowFn(payload));
}
use of org.apache.beam.sdk.runners.AppliedPTransform in project beam by apache.
the class CombineTranslationTest method testToFromProto.
@Test
public void testToFromProto() throws Exception {
PCollection<Integer> input = pipeline.apply(Create.of(1, 2, 3));
input.apply(Combine.globally(combineFn));
final AtomicReference<AppliedPTransform<?, ?, Combine.PerKey<?, ?, ?>>> combine = new AtomicReference<>();
pipeline.traverseTopologically(new PipelineVisitor.Defaults() {
@Override
public void leaveCompositeTransform(Node node) {
if (node.getTransform() instanceof Combine.PerKey) {
checkState(combine.get() == null);
combine.set((AppliedPTransform) node.toAppliedPTransform(getPipeline()));
}
}
});
checkState(combine.get() != null);
SdkComponents sdkComponents = SdkComponents.create();
CombinePayload combineProto = CombineTranslation.toProto(combine.get(), sdkComponents);
RunnerApi.Components componentsProto = sdkComponents.toComponents();
assertEquals(combineFn.getAccumulatorCoder(pipeline.getCoderRegistry(), input.getCoder()), CombineTranslation.getAccumulatorCoder(combineProto, componentsProto));
assertEquals(combineFn, CombineTranslation.getCombineFn(combineProto));
}
use of org.apache.beam.sdk.runners.AppliedPTransform in project beam by apache.
the class PTransformMatchersTest method flattenWithDuplicateInputsWithoutDuplicates.
@Test
public void flattenWithDuplicateInputsWithoutDuplicates() {
AppliedPTransform application = AppliedPTransform.<PCollectionList<Object>, PCollection<Object>, Flatten.PCollections<Object>>of("Flatten", Collections.<TupleTag<?>, PValue>singletonMap(new TupleTag<Object>(), PCollection.createPrimitiveOutputInternal(p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED)), Collections.<TupleTag<?>, PValue>singletonMap(new TupleTag<Object>(), PCollection.createPrimitiveOutputInternal(p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED)), Flatten.pCollections(), p);
assertThat(PTransformMatchers.flattenWithDuplicateInputs().matches(application), is(false));
}
use of org.apache.beam.sdk.runners.AppliedPTransform in project beam by apache.
the class PTransformMatchersTest method emptyFlattenWithEmptyFlatten.
@Test
public void emptyFlattenWithEmptyFlatten() {
AppliedPTransform application = AppliedPTransform.<PCollectionList<Object>, PCollection<Object>, Flatten.PCollections<Object>>of("EmptyFlatten", Collections.<TupleTag<?>, PValue>emptyMap(), Collections.<TupleTag<?>, PValue>singletonMap(new TupleTag<Object>(), PCollection.createPrimitiveOutputInternal(p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED)), Flatten.pCollections(), p);
assertThat(PTransformMatchers.emptyFlatten().matches(application), is(true));
}
use of org.apache.beam.sdk.runners.AppliedPTransform in project beam by apache.
the class PTransformMatchersTest method flattenWithDuplicateInputsWithDuplicates.
@Test
public void flattenWithDuplicateInputsWithDuplicates() {
PCollection<Object> duplicate = PCollection.createPrimitiveOutputInternal(p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED);
AppliedPTransform application = AppliedPTransform.<PCollectionList<Object>, PCollection<Object>, Flatten.PCollections<Object>>of("Flatten", ImmutableMap.<TupleTag<?>, PValue>builder().put(new TupleTag<Object>(), duplicate).put(new TupleTag<Object>(), duplicate).build(), Collections.<TupleTag<?>, PValue>singletonMap(new TupleTag<Object>(), PCollection.createPrimitiveOutputInternal(p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED)), Flatten.pCollections(), p);
assertThat(PTransformMatchers.flattenWithDuplicateInputs().matches(application), is(true));
}
Aggregations