Search in sources :

Example 1 with BiMap

use of org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.BiMap in project beam by apache.

the class PCollectionOutputTagVisitor method visitValue.

@Override
public void visitValue(PValue value, Node producer) {
    for (Entry<ProjectionProducer<PTransform<?, ?>>, Map<PCollection<?>, FieldAccessDescriptor>> entry : pCollFieldAccess.entrySet()) {
        FieldAccessDescriptor fieldAccess = entry.getValue().get(value);
        if (fieldAccess == null) {
            continue;
        }
        BiMap<PCollection<?>, TupleTag<?>> outputs = ImmutableBiMap.copyOf(producer.getOutputs()).inverse();
        TupleTag<?> tag = outputs.get(value);
        Preconditions.checkArgumentNotNull(tag, "PCollection %s not found in outputs of producer %s", value, producer);
        ImmutableMap.Builder<TupleTag<?>, FieldAccessDescriptor> tagEntryBuilder = tagFieldAccess.build().get(entry.getKey());
        if (tagEntryBuilder == null) {
            tagEntryBuilder = ImmutableMap.builder();
            tagFieldAccess.put(entry.getKey(), tagEntryBuilder);
        }
        tagEntryBuilder.put(tag, fieldAccess);
    }
}
Also used : FieldAccessDescriptor(org.apache.beam.sdk.schemas.FieldAccessDescriptor) PCollection(org.apache.beam.sdk.values.PCollection) ProjectionProducer(org.apache.beam.sdk.schemas.ProjectionProducer) TupleTag(org.apache.beam.sdk.values.TupleTag) ImmutableMap(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap) ImmutableBiMap(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableBiMap) Map(java.util.Map) BiMap(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.BiMap) ImmutableMap(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap)

Aggregations

Map (java.util.Map)1 FieldAccessDescriptor (org.apache.beam.sdk.schemas.FieldAccessDescriptor)1 ProjectionProducer (org.apache.beam.sdk.schemas.ProjectionProducer)1 PCollection (org.apache.beam.sdk.values.PCollection)1 TupleTag (org.apache.beam.sdk.values.TupleTag)1 BiMap (org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.BiMap)1 ImmutableBiMap (org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableBiMap)1 ImmutableMap (org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap)1