Examples with PCollection org.apache.beam.sdk.values.PCollection used on opensource projects

Example 86 with PCollection

use of org.apache.beam.sdk.values.PCollection in project beam by apache.

the class PTransformMatchersTest method classEqualToDoesNotMatchSubclass.

@Test
public void classEqualToDoesNotMatchSubclass() {
    class MyPTransform extends PTransform<PCollection<KV<String, Integer>>, PCollection<Integer>> {

        @Override
        public PCollection<Integer> expand(PCollection<KV<String, Integer>> input) {
            return PCollection.createPrimitiveOutputInternal(input.getPipeline(), input.getWindowingStrategy(), input.isBounded(), VarIntCoder.of());
        }
    }
    PTransformMatcher matcher = PTransformMatchers.classEqualTo(MyPTransform.class);
    MyPTransform subclass = new MyPTransform() {
    };
    assertThat(subclass.getClass(), not(Matchers.<Class<?>>equalTo(MyPTransform.class)));
    assertThat(subclass, instanceOf(MyPTransform.class));
    AppliedPTransform<?, ?, ?> application = getAppliedTransform(subclass);
    assertThat(matcher.matches(application), is(false));
}

Example 87 with PCollection

use of org.apache.beam.sdk.values.PCollection in project beam by apache.

the class ReplacementOutputsTest method singletonSucceeds.

@Test
public void singletonSucceeds() {
    Map<PCollection<?>, ReplacementOutput> replacements = ReplacementOutputs.singleton(PValues.expandValue(ints), replacementInts);
    assertThat(replacements, Matchers.hasKey(replacementInts));
    ReplacementOutput replacement = replacements.get(replacementInts);
    Map.Entry<TupleTag<?>, PValue> taggedInts = Iterables.getOnlyElement(ints.expand().entrySet());
    assertThat(replacement.getOriginal().getTag(), equalTo(taggedInts.getKey()));
    assertThat(replacement.getOriginal().getValue(), equalTo(taggedInts.getValue()));
    assertThat(replacement.getReplacement().getValue(), equalTo(replacementInts));
}

Example 88 with PCollection

use of org.apache.beam.sdk.values.PCollection in project beam by apache.

the class FieldAccessVisitor method visitPrimitiveTransform.

@Override
public void visitPrimitiveTransform(Node node) {
    Map<PCollection<?>, FieldAccessDescriptor> currentFieldAccess = getFieldAccess(node);
    for (Entry<PCollection<?>, FieldAccessDescriptor> entry : currentFieldAccess.entrySet()) {
        FieldAccessDescriptor previousFieldAccess = pCollectionFieldAccess.get(entry.getKey());
        FieldAccessDescriptor newFieldAccess = previousFieldAccess == null ? entry.getValue() : FieldAccessDescriptor.union(ImmutableList.of(previousFieldAccess, entry.getValue()));
        pCollectionFieldAccess.put(entry.getKey(), newFieldAccess);
    }
}

Example 89 with PCollection

use of org.apache.beam.sdk.values.PCollection in project beam by apache.

the class FieldAccessVisitor method getFieldAccess.

private static Map<PCollection<?>, FieldAccessDescriptor> getFieldAccess(Node node) {
    PTransform<?, ?> transform = node.getTransform();
    HashMap<PCollection<?>, FieldAccessDescriptor> access = new HashMap<>();
    if (transform instanceof MultiOutput) {
        // Get main input pcoll.
        Set<PCollection<?>> mainInputs = node.getInputs().entrySet().stream().filter((entry) -> !transform.getAdditionalInputs().containsKey(entry.getKey())).map(Entry::getValue).collect(Collectors.toSet());
        PCollection<?> mainInput = Iterables.getOnlyElement(mainInputs);
        // Get field access.
        DoFn<?, ?> fn = ((MultiOutput<?, ?>) transform).getFn();
        FieldAccessDescriptor fields = ParDo.getDoFnSchemaInformation(fn, mainInput).getFieldAccessDescriptor();
        // Record field access.
        access.put(mainInput, fields);
    }
    // For every input without field access info, we must assume all fields need to be accessed.
    for (PCollection<?> input : node.getInputs().values()) {
        if (!access.containsKey(input)) {
            access.put(input, FieldAccessDescriptor.withAllFields());
        }
    }
    return ImmutableMap.copyOf(access);
}

Example 90 with PCollection

use of org.apache.beam.sdk.values.PCollection in project beam by apache.

the class ProjectionProducerVisitor method enterCompositeTransform.

@Override
public CompositeBehavior enterCompositeTransform(Node node) {
    PTransform<?, ?> transform = node.getTransform();
    // TODO(BEAM-13658) Support inputs other than PBegin.
    if (!node.getInputs().isEmpty()) {
        return CompositeBehavior.DO_NOT_ENTER_TRANSFORM;
    }
    if (!(transform instanceof ProjectionProducer)) {
        return CompositeBehavior.ENTER_TRANSFORM;
    }
    ProjectionProducer<PTransform<?, ?>> pushdownProjector = (ProjectionProducer<PTransform<?, ?>>) transform;
    if (!pushdownProjector.supportsProjectionPushdown()) {
        return CompositeBehavior.ENTER_TRANSFORM;
    }
    ImmutableMap.Builder<PCollection<?>, FieldAccessDescriptor> builder = ImmutableMap.builder();
    for (PCollection<?> output : node.getOutputs().values()) {
        FieldAccessDescriptor fieldAccess = pCollectionFieldAccess.get(output);
        if (fieldAccess != null && !fieldAccess.getAllFields()) {
            builder.put(output, fieldAccess);
        }
    }
    Map<PCollection<?>, FieldAccessDescriptor> localOpportunities = builder.build();
    if (localOpportunities.isEmpty()) {
        return CompositeBehavior.ENTER_TRANSFORM;
    }
    pushdownOpportunities.put(pushdownProjector, localOpportunities);
    // If there are nested PushdownProjector implementations, apply only the outermost one.
    return CompositeBehavior.DO_NOT_ENTER_TRANSFORM;
}