Search in sources :

Example 1 with StreamingCreatePCollectionViewFactory

use of org.apache.beam.runners.dataflow.StreamingViewOverrides.StreamingCreatePCollectionViewFactory in project beam by apache.

the class DataflowRunner method getOverrides.

private List<PTransformOverride> getOverrides(boolean streaming) {
    ImmutableList.Builder<PTransformOverride> overridesBuilder = ImmutableList.builder();
    // Create is implemented in terms of a Read, so it must precede the override to Read in
    // streaming
    overridesBuilder.add(PTransformOverride.of(PTransformMatchers.flattenWithDuplicateInputs(), DeduplicatedFlattenFactory.create())).add(PTransformOverride.of(PTransformMatchers.emptyFlatten(), EmptyFlattenAsCreateFactory.instance()));
    if (streaming) {
        if (!hasExperiment(options, "enable_custom_pubsub_source")) {
            overridesBuilder.add(PTransformOverride.of(PTransformMatchers.classEqualTo(PubsubUnboundedSource.class), new ReflectiveRootOverrideFactory(StreamingPubsubIORead.class, this)));
        }
        if (!hasExperiment(options, "enable_custom_pubsub_sink")) {
            overridesBuilder.add(PTransformOverride.of(PTransformMatchers.classEqualTo(PubsubUnboundedSink.class), new StreamingPubsubIOWriteOverrideFactory(this)));
        }
        overridesBuilder.add(// must precede it
        PTransformOverride.of(PTransformMatchers.classEqualTo(Read.Bounded.class), new ReflectiveRootOverrideFactory(StreamingBoundedRead.class, this))).add(PTransformOverride.of(PTransformMatchers.classEqualTo(Read.Unbounded.class), new ReflectiveRootOverrideFactory(StreamingUnboundedRead.class, this))).add(PTransformOverride.of(PTransformMatchers.classEqualTo(View.CreatePCollectionView.class), new StreamingCreatePCollectionViewFactory()));
    } else {
        overridesBuilder.add(PTransformOverride.of(PTransformMatchers.stateOrTimerParDoMulti(), BatchStatefulParDoOverrides.multiOutputOverrideFactory())).add(PTransformOverride.of(PTransformMatchers.stateOrTimerParDoSingle(), BatchStatefulParDoOverrides.singleOutputOverrideFactory())).add(PTransformOverride.of(PTransformMatchers.classEqualTo(Combine.GloballyAsSingletonView.class), new BatchCombineGloballyAsSingletonViewFactory(this))).add(PTransformOverride.of(PTransformMatchers.classEqualTo(View.AsMap.class), new ReflectiveOneToOneOverrideFactory(BatchViewOverrides.BatchViewAsMap.class, this))).add(PTransformOverride.of(PTransformMatchers.classEqualTo(View.AsMultimap.class), new ReflectiveOneToOneOverrideFactory(BatchViewOverrides.BatchViewAsMultimap.class, this))).add(PTransformOverride.of(PTransformMatchers.classEqualTo(View.AsSingleton.class), new ReflectiveOneToOneOverrideFactory(BatchViewOverrides.BatchViewAsSingleton.class, this))).add(PTransformOverride.of(PTransformMatchers.classEqualTo(View.AsList.class), new ReflectiveOneToOneOverrideFactory(BatchViewOverrides.BatchViewAsList.class, this))).add(PTransformOverride.of(PTransformMatchers.classEqualTo(View.AsIterable.class), new ReflectiveOneToOneOverrideFactory(BatchViewOverrides.BatchViewAsIterable.class, this)));
    }
    overridesBuilder.add(PTransformOverride.of(PTransformMatchers.classEqualTo(Reshuffle.class), new ReshuffleOverrideFactory())).add(PTransformOverride.of(PTransformMatchers.classEqualTo(Combine.GroupedValues.class), new PrimitiveCombineGroupedValuesOverrideFactory())).add(PTransformOverride.of(PTransformMatchers.classEqualTo(ParDo.SingleOutput.class), new PrimitiveParDoSingleFactory()));
    return overridesBuilder.build();
}
Also used : BatchCombineGloballyAsSingletonViewFactory(org.apache.beam.runners.dataflow.BatchViewOverrides.BatchCombineGloballyAsSingletonViewFactory) ImmutableList(com.google.common.collect.ImmutableList) PTransformOverride(org.apache.beam.sdk.runners.PTransformOverride) View(org.apache.beam.sdk.transforms.View) PCollectionView(org.apache.beam.sdk.values.PCollectionView) GroupedValues(org.apache.beam.sdk.transforms.Combine.GroupedValues) StreamingCreatePCollectionViewFactory(org.apache.beam.runners.dataflow.StreamingViewOverrides.StreamingCreatePCollectionViewFactory) Read(org.apache.beam.sdk.io.Read) ParDo(org.apache.beam.sdk.transforms.ParDo)

Aggregations

ImmutableList (com.google.common.collect.ImmutableList)1 BatchCombineGloballyAsSingletonViewFactory (org.apache.beam.runners.dataflow.BatchViewOverrides.BatchCombineGloballyAsSingletonViewFactory)1 StreamingCreatePCollectionViewFactory (org.apache.beam.runners.dataflow.StreamingViewOverrides.StreamingCreatePCollectionViewFactory)1 Read (org.apache.beam.sdk.io.Read)1 PTransformOverride (org.apache.beam.sdk.runners.PTransformOverride)1 GroupedValues (org.apache.beam.sdk.transforms.Combine.GroupedValues)1 ParDo (org.apache.beam.sdk.transforms.ParDo)1 View (org.apache.beam.sdk.transforms.View)1 PCollectionView (org.apache.beam.sdk.values.PCollectionView)1