Java Code Examples for org.apache.beam.sdk.values.PCollection#createPrimitiveOutputInternal()

@Override
public PCollection<KV<K, Iterable<V>>> expand(
    PCollection<KV<K, Iterable<WindowedValue<V>>>> input) {
  @SuppressWarnings("unchecked")
  KvCoder<K, Iterable<WindowedValue<V>>> inputKvCoder =
      (KvCoder<K, Iterable<WindowedValue<V>>>) input.getCoder();

  Coder<K> keyCoder = inputKvCoder.getKeyCoder();
  Coder<Iterable<WindowedValue<V>>> inputValueCoder = inputKvCoder.getValueCoder();

  IterableCoder<WindowedValue<V>> inputIterableValueCoder =
      (IterableCoder<WindowedValue<V>>) inputValueCoder;
  Coder<WindowedValue<V>> inputIterableElementCoder = inputIterableValueCoder.getElemCoder();
  WindowedValueCoder<V> inputIterableWindowedValueCoder =
      (WindowedValueCoder<V>) inputIterableElementCoder;

  Coder<V> inputIterableElementValueCoder = inputIterableWindowedValueCoder.getValueCoder();
  Coder<Iterable<V>> outputValueCoder = IterableCoder.of(inputIterableElementValueCoder);
  Coder<KV<K, Iterable<V>>> outputKvCoder = KvCoder.of(keyCoder, outputValueCoder);

  return PCollection.createPrimitiveOutputInternal(
      input.getPipeline(), windowingStrategy, input.isBounded(), outputKvCoder);
}
 

private Object applyReadSourceTransform(
    PTransform<?, ?> transform, PCollection.IsBounded isBounded, StreamExecutionEnvironment env) {

  FlinkStreamingPipelineTranslator.StreamTransformTranslator<PTransform<?, ?>> translator =
      getReadSourceTranslator();
  FlinkStreamingTranslationContext ctx =
      new FlinkStreamingTranslationContext(env, PipelineOptionsFactory.create());

  Pipeline pipeline = Pipeline.create();
  PCollection<String> pc =
      PCollection.createPrimitiveOutputInternal(
          pipeline, WindowingStrategy.globalDefault(), isBounded, StringUtf8Coder.of());
  pc.setName("output");

  Map<TupleTag<?>, PValue> outputs = new HashMap<>();
  outputs.put(new TupleTag<>(), pc);
  AppliedPTransform<?, ?, ?> appliedTransform =
      AppliedPTransform.of(
          "test-transform", Collections.emptyMap(), outputs, transform, Pipeline.create());

  ctx.setCurrentTransform(appliedTransform);
  translator.translateNode(transform, ctx);

  return ctx.getInputDataStream(pc).getTransformation();
}
 

@Test
public void flattenWithDuplicateInputsWithDuplicates() {
  PCollection<Integer> duplicate =
      PCollection.createPrimitiveOutputInternal(
          p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED, VarIntCoder.of());
  AppliedPTransform application =
      AppliedPTransform.of(
          "Flatten",
          ImmutableMap.<TupleTag<?>, PValue>builder()
              .put(new TupleTag<Integer>(), duplicate)
              .put(new TupleTag<Integer>(), duplicate)
              .build(),
          Collections.singletonMap(
              new TupleTag<Integer>(),
              PCollection.createPrimitiveOutputInternal(
                  p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED, VarIntCoder.of())),
          Flatten.pCollections(),
          p);

  assertThat(PTransformMatchers.flattenWithDuplicateInputs().matches(application), is(true));
}
 

@Override
public final PCollection<T> expand(PBegin input) {
  source.validate();

  if (ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api")
      && !ExperimentalOptions.hasExperiment(
          input.getPipeline().getOptions(), "beam_fn_api_use_deprecated_read")) {
    // We don't use Create here since Create is defined as a BoundedSource and using it would
    // cause an infinite expansion loop. We can reconsider this if Create is implemented
    // directly as a SplittableDoFn.
    return input
        .getPipeline()
        .apply(Impulse.create())
        .apply(
            MapElements.into(new TypeDescriptor<BoundedSource<T>>() {}).via(element -> source))
        .setCoder(SerializableCoder.of(new TypeDescriptor<BoundedSource<T>>() {}))
        .apply(ParDo.of(new BoundedSourceAsSDFWrapperFn<>()))
        .setCoder(source.getOutputCoder());
  }

  return PCollection.createPrimitiveOutputInternal(
      input.getPipeline(),
      WindowingStrategy.globalDefault(),
      IsBounded.BOUNDED,
      source.getOutputCoder());
}
 

@Override
public final PCollection<ValueWithRecordId<T>> expand(PInput input) {
  return PCollection.createPrimitiveOutputInternal(
      input.getPipeline(),
      WindowingStrategy.globalDefault(),
      IsBounded.UNBOUNDED,
      ValueWithRecordId.ValueWithRecordIdCoder.of(source.getOutputCoder()));
}
 

@Override
public PCollection<T> expand(PBegin input) {
  runner.setClockSupplier(new TestClockSupplier());
  return PCollection.createPrimitiveOutputInternal(
      input.getPipeline(),
      WindowingStrategy.globalDefault(),
      IsBounded.UNBOUNDED,
      original.getValueCoder());
}
 

@SuppressWarnings({"rawtypes", "unchecked"})
@Override
public PCollection<KV<K, Iterable<WindowedValue<V>>>> expand(PCollection<KV<K, V>> input) {
  return PCollection.createPrimitiveOutputInternal(
      input.getPipeline(),
      input.getWindowingStrategy(),
      input.isBounded(),
      (Coder) GroupByKey.getOutputKvCoder(input.getCoder()));
}
 

@Override
public PCollection<T> expand(PBegin input) {
  return PCollection.createPrimitiveOutputInternal(
      input.getPipeline(),
      WindowingStrategy.globalDefault(),
      PCollection.IsBounded.UNBOUNDED,
      coder);
}
 

@Override
public PCollection<OutputT> expand(PCollection<? extends InputT> input) {
  return PCollection.createPrimitiveOutputInternal(
      input.getPipeline(),
      input.getWindowingStrategy(),
      input.isBounded(),
      outputCoder,
      onlyOutputTag);
}
 

public static PCollection<?> fromProto(
    RunnerApi.PCollection pCollection, Pipeline pipeline, RehydratedComponents components)
    throws IOException {

  Coder<?> coder = components.getCoder(pCollection.getCoderId());
  return PCollection.createPrimitiveOutputInternal(
      pipeline,
      components.getWindowingStrategy(pCollection.getWindowingStrategyId()),
      fromProto(pCollection.getIsBounded()),
      (Coder) coder);
}
 

@Override
public PCollection<T> expand(PCollectionList<T> inputs) {
  WindowingStrategy<?, ?> windowingStrategy;
  IsBounded isBounded = IsBounded.BOUNDED;
  if (!inputs.getAll().isEmpty()) {
    windowingStrategy = inputs.get(0).getWindowingStrategy();
    for (PCollection<?> input : inputs.getAll()) {
      WindowingStrategy<?, ?> other = input.getWindowingStrategy();
      if (!windowingStrategy.getWindowFn().isCompatible(other.getWindowFn())) {
        throw new IllegalStateException(
            "Inputs to Flatten had incompatible window windowFns: "
                + windowingStrategy.getWindowFn()
                + ", "
                + other.getWindowFn());
      }

      if (!windowingStrategy.getTrigger().isCompatible(other.getTrigger())) {
        throw new IllegalStateException(
            "Inputs to Flatten had incompatible triggers: "
                + windowingStrategy.getTrigger()
                + ", "
                + other.getTrigger());
      }
      isBounded = isBounded.and(input.isBounded());
    }
  } else {
    windowingStrategy = WindowingStrategy.globalDefault();
  }

  return PCollection.createPrimitiveOutputInternal(
      inputs.getPipeline(),
      windowingStrategy,
      isBounded,
      // Take coder from first collection. If there are none, will be left unspecified.
      inputs.getAll().isEmpty() ? null : inputs.get(0).getCoder());
}
 

@Test
public void producingOwnAndOthersOutputsFails() {
  PCollection<Long> created =
      PCollection.createPrimitiveOutputInternal(
          pipeline, WindowingStrategy.globalDefault(), IsBounded.BOUNDED, VarLongCoder.of());
  hierarchy.pushNode("Create", PBegin.in(pipeline), Create.of(1));
  hierarchy.setOutput(created);
  hierarchy.popNode();
  PCollectionList<Long> pcList = PCollectionList.of(created);

  final PCollectionList<Long> appended =
      pcList.and(
          PCollection.createPrimitiveOutputInternal(
                  pipeline,
                  WindowingStrategy.globalDefault(),
                  IsBounded.BOUNDED,
                  VarLongCoder.of())
              .setName("prim"));
  hierarchy.pushNode(
      "AddPc",
      pcList,
      new PTransform<PCollectionList<Long>, PCollectionList<Long>>() {
        @Override
        public PCollectionList<Long> expand(PCollectionList<Long> input) {
          return appended;
        }
      });
  thrown.expect(IllegalArgumentException.class);
  thrown.expectMessage("contains a primitive POutput produced by it");
  thrown.expectMessage("AddPc");
  thrown.expectMessage("Create");
  thrown.expectMessage(appended.expand().toString());
  hierarchy.setOutput(appended);
}
 

@Test
public void emptyCompositeSucceeds() {
  PCollection<Long> created =
      PCollection.createPrimitiveOutputInternal(
          pipeline, WindowingStrategy.globalDefault(), IsBounded.BOUNDED, VarLongCoder.of());
  TransformHierarchy.Node node = hierarchy.pushNode("Create", PBegin.in(pipeline), Create.of(1));
  hierarchy.setOutput(created);
  hierarchy.popNode();
  PCollectionList<Long> pcList = PCollectionList.of(created);

  TransformHierarchy.Node emptyTransform =
      hierarchy.pushNode(
          "Extract",
          pcList,
          new PTransform<PCollectionList<Long>, PCollection<Long>>() {
            @Override
            public PCollection<Long> expand(PCollectionList<Long> input) {
              return input.get(0);
            }
          });
  hierarchy.setOutput(created);
  hierarchy.popNode();
  assertThat(hierarchy.getProducer(created), equalTo(node));
  assertThat(
      "A Transform that produces non-primitive output should be composite",
      emptyTransform.isCompositeNode(),
      is(true));
}
 

@Override
public PCollection<Integer> expand(PCollection<Integer> input) {
  return PCollection.createPrimitiveOutputInternal(
      input.getPipeline(),
      WindowingStrategy.globalDefault(),
      input.isBounded(),
      input.getCoder());
}
 

@Override
public PCollection<KeyedWorkItem<K, V>> expand(PCollection<KV<K, V>> input) {
  return PCollection.createPrimitiveOutputInternal(
      input.getPipeline(),
      WindowingStrategy.globalDefault(),
      input.isBounded(),
      KeyedWorkItemCoder.of(
          GroupByKey.getKeyCoder(input.getCoder()),
          GroupByKey.getInputValueCoder(input.getCoder()),
          input.getWindowingStrategy().getWindowFn().windowCoder()));
}
 

@Override
public PCollection<PubsubMessage> expand(PBegin input) {
  Coder coder =
      transform.getNeedsMessageId()
          ? new PubsubMessageWithAttributesAndMessageIdCoder()
          : new PubsubMessageWithAttributesCoder();
  return PCollection.createPrimitiveOutputInternal(
      input.getPipeline(), WindowingStrategy.globalDefault(), IsBounded.UNBOUNDED, coder);
}
 

@Override
public PCollection<T> expand(PBegin input) {
  return PCollection.createPrimitiveOutputInternal(
      input.getPipeline(), WindowingStrategy.globalDefault(), IsBounded.BOUNDED, coder);
}
 

@Override
@SuppressWarnings("deprecation")
public PCollection<Iterable<ElemT>> expand(PCollection<Iterable<ElemT>> input) {
  return PCollection.createPrimitiveOutputInternal(
      input.getPipeline(), input.getWindowingStrategy(), input.isBounded(), input.getCoder());
}
 

@Override
public PCollection<KV<K, OutputT>> expand(PCollection<KV<K, Iterable<InputT>>> input) {
  return PCollection.createPrimitiveOutputInternal(
      input.getPipeline(), input.getWindowingStrategy(), input.isBounded(), outputCoder);
}
 

@Override
public PCollection<T> expand(PCollection<T> input) {
  return PCollection.createPrimitiveOutputInternal(
      input.getPipeline(), updatedStrategy, input.isBounded(), input.getCoder());
}