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

@Override
public PCollectionList<T> expand(PCollection<T> in) {
  final TupleTagList outputTags = partitionDoFn.getOutputTags();

  PCollectionTuple outputs =
      in.apply(
          ParDo.of(partitionDoFn)
              .withOutputTags(new TupleTag<Void>() {}, outputTags)
              .withSideInputs(partitionDoFn.getSideInputs()));

  PCollectionList<T> pcs = PCollectionList.empty(in.getPipeline());
  Coder<T> coder = in.getCoder();

  for (TupleTag<?> outputTag : outputTags.getAll()) {
    // All the tuple tags are actually TupleTag<T>
    // And all the collections are actually PCollection<T>
    @SuppressWarnings("unchecked")
    TupleTag<T> typedOutputTag = (TupleTag<T>) outputTag;
    pcs = pcs.and(outputs.get(typedOutputTag).setCoder(coder));
  }
  return pcs;
}
 

/** Transforms the input {@link PCollection} into a singleton {@link Map} per window. */
private <W extends BoundedWindow> PCollection<?> applyForSingletonFallback(
    PCollection<KV<K, V>> input) {
  @SuppressWarnings("unchecked")
  Coder<W> windowCoder = (Coder<W>) input.getWindowingStrategy().getWindowFn().windowCoder();

  @SuppressWarnings({"rawtypes", "unchecked"})
  KvCoder<K, V> inputCoder = (KvCoder) input.getCoder();

  @SuppressWarnings({"unchecked", "rawtypes"})
  Coder<Function<WindowedValue<V>, V>> transformCoder =
      (Coder) SerializableCoder.of(WindowedValueToValue.class);

  Coder<TransformedMap<K, WindowedValue<V>, V>> finalValueCoder =
      TransformedMapCoder.of(
          transformCoder,
          MapCoder.of(
              inputCoder.getKeyCoder(),
              FullWindowedValueCoder.of(inputCoder.getValueCoder(), windowCoder)));

  return BatchViewAsSingleton.applyForSingleton(
      runner, input, new ToMapDoFn<>(windowCoder), finalValueCoder, view);
}
 

/** Transforms the input {@link PCollection} into a singleton {@link Map} per window. */
private <W extends BoundedWindow> PCollection<?> applyForSingletonFallback(
    PCollection<KV<K, V>> input) {
  @SuppressWarnings("unchecked")
  Coder<W> windowCoder = (Coder<W>) input.getWindowingStrategy().getWindowFn().windowCoder();

  @SuppressWarnings({"rawtypes", "unchecked"})
  KvCoder<K, V> inputCoder = (KvCoder) input.getCoder();

  @SuppressWarnings({"unchecked", "rawtypes"})
  Coder<Function<Iterable<WindowedValue<V>>, Iterable<V>>> transformCoder =
      (Coder) SerializableCoder.of(IterableWithWindowedValuesToIterable.class);

  Coder<TransformedMap<K, Iterable<WindowedValue<V>>, Iterable<V>>> finalValueCoder =
      TransformedMapCoder.of(
          transformCoder,
          MapCoder.of(
              inputCoder.getKeyCoder(),
              IterableCoder.of(
                  FullWindowedValueCoder.of(inputCoder.getValueCoder(), windowCoder))));

  return BatchViewAsSingleton.applyForSingleton(
      runner, input, new ToMultimapDoFn<>(windowCoder), finalValueCoder, view);
}
 

@Override
public PCollection<KV<Integer, Iterable<KV<KV<K, W>, WindowedValue<V>>>>> expand(
    PCollection<KV<K, V>> input) {

  @SuppressWarnings("unchecked")
  Coder<W> windowCoder = (Coder<W>) input.getWindowingStrategy().getWindowFn().windowCoder();
  @SuppressWarnings("unchecked")
  KvCoder<K, V> inputCoder = (KvCoder<K, V>) input.getCoder();

  PCollection<KV<Integer, KV<KV<K, W>, WindowedValue<V>>>> keyedByHash;
  keyedByHash =
      input.apply(ParDo.of(new GroupByKeyHashAndSortByKeyAndWindowDoFn<K, V, W>(coder)));
  keyedByHash.setCoder(
      KvCoder.of(
          VarIntCoder.of(),
          KvCoder.of(
              KvCoder.of(inputCoder.getKeyCoder(), windowCoder),
              FullWindowedValueCoder.of(inputCoder.getValueCoder(), windowCoder))));

  return keyedByHash.apply(new GroupByKeyAndSortValuesOnly<>());
}
 

@Override
public PCollection<KV<K, TimestampedValue<V>>> expand(PCollection<KV<K, V>> input) {
  KvCoder<K, V> coder = (KvCoder<K, V>) input.getCoder();
  return input
      .apply(
          ParDo.of(
              new DoFn<KV<K, V>, KV<K, TimestampedValue<V>>>() {
                @ProcessElement
                public void processElement(
                    @Element KV<K, V> element,
                    @Timestamp Instant timestamp,
                    OutputReceiver<KV<K, TimestampedValue<V>>> r) {
                  r.output(
                      KV.of(
                          element.getKey(),
                          TimestampedValue.of(element.getValue(), timestamp)));
                }
              }))
      .setCoder(
          KvCoder.of(coder.getKeyCoder(), TimestampedValueCoder.of(coder.getValueCoder())));
}
 

@Override
public PCollection<T> expand(PCollection<? extends Iterable<T>> in) {
  Coder<? extends Iterable<T>> inCoder = in.getCoder();
  if (!(inCoder instanceof IterableLikeCoder)) {
    throw new IllegalArgumentException(
        "expecting the input Coder<Iterable> to be an IterableLikeCoder");
  }
  @SuppressWarnings("unchecked")
  Coder<T> elemCoder = ((IterableLikeCoder<T, ?>) inCoder).getElemCoder();

  return in.apply(
          "FlattenIterables",
          FlatMapElements.via(
              new SimpleFunction<Iterable<T>, Iterable<T>>() {
                @Override
                public Iterable<T> apply(Iterable<T> element) {
                  return element;
                }
              }))
      .setCoder(elemCoder);
}
 

@Test
public void testGloballyOutputCoder() {
  p.enableAbandonedNodeEnforcement(false);

  BigEndianLongCoder inputCoder = BigEndianLongCoder.of();

  PCollection<Long> output =
      p.apply(Create.of(1L, 2L).withCoder(inputCoder)).apply(Latest.globally());

  Coder<Long> outputCoder = output.getCoder();
  assertThat(outputCoder, instanceOf(NullableCoder.class));
  assertEquals(inputCoder, ((NullableCoder<?>) outputCoder).getValueCoder());
}
 

/**
 * Returns the value coder for the given PCollection. Assumes that the value coder is an instance
 * of {@code KvCoder<K, V>}.
 */
private <V> Coder<V> getValueCoder(PCollection<KV<K, V>> pCollection) {
  // Assumes that the PCollection uses a KvCoder.
  Coder<?> entryCoder = pCollection.getCoder();
  if (!(entryCoder instanceof KvCoder<?, ?>)) {
    throw new IllegalArgumentException("PCollection does not use a KvCoder");
  }
  @SuppressWarnings("unchecked")
  KvCoder<K, V> coder = (KvCoder<K, V>) entryCoder;
  return coder.getValueCoder();
}
 

/**
 * Constructs a {@link SingletonAssert} for the value of the provided {@link PCollection} with the
 * specified reason. The {@link PCollection} must have at most one value per key.
 *
 * <p>Note that the actual value must be coded by a {@link KvCoder}, not just any {@code Coder<K,
 * V>}.
 */
public static <K, V> SingletonAssert<Map<K, V>> thatMap(
    String reason, PCollection<KV<K, V>> actual) {
  @SuppressWarnings("unchecked")
  KvCoder<K, V> kvCoder = (KvCoder<K, V>) actual.getCoder();
  return new PCollectionViewAssert<>(
      actual,
      View.asMap(),
      MapCoder.of(kvCoder.getKeyCoder(), kvCoder.getValueCoder()),
      PAssertionSite.capture(reason));
}
 

PCollectionSingletonAssert(
    PCollection<T> actual,
    AssertionWindows rewindowingStrategy,
    SimpleFunction<Iterable<ValueInSingleWindow<T>>, Iterable<T>> paneExtractor,
    PAssertionSite site) {
  this.actual = actual;
  this.coder = actual.getCoder();
  this.rewindowingStrategy = rewindowingStrategy;
  this.paneExtractor = paneExtractor;
  this.site = site;
}
 

@Override
public PCollection<KV<K, ValueInSingleWindow<V>>> expand(PCollection<KV<K, V>> input) {
  KvCoder<K, V> coder = (KvCoder<K, V>) input.getCoder();
  return input
      .apply(
          ParDo.of(
              new DoFn<KV<K, V>, KV<K, ValueInSingleWindow<V>>>() {
                @ProcessElement
                public void processElement(
                    @Element KV<K, V> element,
                    @Timestamp Instant timestamp,
                    BoundedWindow window,
                    PaneInfo pane,
                    OutputReceiver<KV<K, ValueInSingleWindow<V>>> r) {
                  r.output(
                      KV.of(
                          element.getKey(),
                          ValueInSingleWindow.of(element.getValue(), timestamp, window, pane)));
                }
              }))
      .setCoder(
          KvCoder.of(
              coder.getKeyCoder(),
              ValueInSingleWindow.Coder.of(
                  coder.getValueCoder(),
                  input.getWindowingStrategy().getWindowFn().windowCoder())));
}
 

private static <K, InputT, OutputT> void doTranslate(
    PTransform<PCollection<KV<K, InputT>>, PCollection<KV<K, OutputT>>> transform,
    TransformHierarchy.Node node,
    TranslationContext ctx) {
  final PCollection<KV<K, InputT>> input = ctx.getInput(transform);

  final PCollection<KV<K, OutputT>> output = ctx.getOutput(transform);
  final TupleTag<KV<K, OutputT>> outputTag = ctx.getOutputTag(transform);

  @SuppressWarnings("unchecked")
  final WindowingStrategy<?, BoundedWindow> windowingStrategy =
      (WindowingStrategy<?, BoundedWindow>) input.getWindowingStrategy();

  final MessageStream<OpMessage<KV<K, InputT>>> inputStream = ctx.getMessageStream(input);

  final KvCoder<K, InputT> kvInputCoder = (KvCoder<K, InputT>) input.getCoder();
  final Coder<WindowedValue<KV<K, InputT>>> elementCoder = SamzaCoders.of(input);

  final SystemReduceFn<K, InputT, ?, OutputT, BoundedWindow> reduceFn =
      getSystemReduceFn(transform, input.getPipeline(), kvInputCoder);

  final MessageStream<OpMessage<KV<K, OutputT>>> outputStream =
      doTranslateGBK(
          inputStream,
          needRepartition(node, ctx),
          reduceFn,
          windowingStrategy,
          kvInputCoder,
          elementCoder,
          ctx.getTransformFullName(),
          ctx.getTransformId(),
          outputTag,
          input.isBounded());

  ctx.registerMessageStream(output, outputStream);
}
 

@Override
public PCollection<KV<K1, Iterable<KV<K2, V>>>> expand(PCollection<KV<K1, KV<K2, V>>> input) {
  @SuppressWarnings("unchecked")
  KvCoder<K1, KV<K2, V>> inputCoder = (KvCoder<K1, KV<K2, V>>) input.getCoder();
  return PCollection.createPrimitiveOutputInternal(
      input.getPipeline(),
      WindowingStrategy.globalDefault(),
      IsBounded.BOUNDED,
      KvCoder.of(inputCoder.getKeyCoder(), IterableCoder.of(inputCoder.getValueCoder())));
}
 

@Override
public void translateNode(
    View.CreatePCollectionView<ElemT, ViewT> transform, Twister2BatchTranslationContext context) {
  BatchTSet<WindowedValue<ElemT>> inputDataSet =
      context.getInputDataSet(context.getInput(transform));
  @SuppressWarnings("unchecked")
  AppliedPTransform<
          PCollection<ElemT>,
          PCollection<ElemT>,
          PTransform<PCollection<ElemT>, PCollection<ElemT>>>
      application =
          (AppliedPTransform<
                  PCollection<ElemT>,
                  PCollection<ElemT>,
                  PTransform<PCollection<ElemT>, PCollection<ElemT>>>)
              context.getCurrentTransform();
  org.apache.beam.sdk.values.PCollectionView<ViewT> input;
  PCollection<ElemT> inputPCol = context.getInput(transform);
  final KvCoder coder = (KvCoder) inputPCol.getCoder();
  Coder inputKeyCoder = coder.getKeyCoder();
  WindowingStrategy windowingStrategy = inputPCol.getWindowingStrategy();
  WindowFn windowFn = windowingStrategy.getWindowFn();
  final WindowedValue.WindowedValueCoder wvCoder =
      WindowedValue.FullWindowedValueCoder.of(coder.getValueCoder(), windowFn.windowCoder());
  BatchTSet<WindowedValue<ElemT>> inputGathered =
      inputDataSet
          .direct()
          .map(new MapToTupleFunction<>(inputKeyCoder, wvCoder))
          .allGather()
          .map(new ByteToWindowFunctionPrimitive(inputKeyCoder, wvCoder));
  try {
    input = CreatePCollectionViewTranslation.getView(application);
  } catch (IOException e) {
    throw new RuntimeException(e);
  }
  context.setSideInputDataSet(input.getTagInternal().getId(), inputGathered);
}
 

@Test
public void perKeyDefault() {
  PCollection<Long> stream = tp.apply(Create.of(smallStream));
  PCollection<Sketch<Long>> sketch =
      stream.apply(WithKeys.of(1)).apply(SketchFrequencies.perKey()).apply(Values.create());

  Coder<Long> coder = stream.getCoder();

  PAssert.thatSingleton("Verify number of hits", sketch)
      .satisfies(new VerifyStreamFrequencies<>(coder, distinctElems, frequencies));

  tp.run();
}
 

public <T> Coder<WindowedValue<T>> getWindowedInputCoder(PCollection<T> collection) {
  Coder<T> valueCoder = collection.getCoder();

  return WindowedValue.getFullCoder(
      valueCoder, collection.getWindowingStrategy().getWindowFn().windowCoder());
}
 

@VisibleForTesting
PCollection<KeyedWorkItem<K, KV<K, InputT>>> groupToKeyedWorkItem(
    PCollection<KV<K, InputT>> input) {

  WindowingStrategy<?, ?> inputWindowingStrategy = input.getWindowingStrategy();

  // A KvCoder is required since this goes through GBK. Further, WindowedValueCoder
  // is not registered by default, so we explicitly set the relevant coders.
  checkState(
      input.getCoder() instanceof KvCoder,
      "Input to a %s using state requires a %s, but the coder was %s",
      ParDo.class.getSimpleName(),
      KvCoder.class.getSimpleName(),
      input.getCoder());

  KvCoder<K, InputT> kvCoder = (KvCoder<K, InputT>) input.getCoder();
  Coder<K> keyCoder = kvCoder.getKeyCoder();
  Coder<? extends BoundedWindow> windowCoder =
      inputWindowingStrategy.getWindowFn().windowCoder();

  return input
      // Stash the original timestamps, etc, for when it is fed to the user's DoFn
      .apply("Reify timestamps", ParDo.of(new ReifyWindowedValueFn<>()))
      .setCoder(KvCoder.of(keyCoder, WindowedValue.getFullCoder(kvCoder, windowCoder)))

      // We are going to GBK to gather keys and windows but otherwise do not want
      // to alter the flow of data. This entails:
      //  - trigger as fast as possible
      //  - maintain the full timestamps of elements
      //  - ensure this GBK holds to the minimum of those timestamps (via TimestampCombiner)
      //  - discard past panes as it is "just a stream" of elements
      .apply(
          Window.<KV<K, WindowedValue<KV<K, InputT>>>>configure()
              .triggering(Repeatedly.forever(AfterPane.elementCountAtLeast(1)))
              .discardingFiredPanes()
              .withAllowedLateness(inputWindowingStrategy.getAllowedLateness())
              .withTimestampCombiner(TimestampCombiner.EARLIEST))

      // A full GBK to group by key _and_ window
      .apply("Group by key", GroupByKey.create())

      // Adapt to KeyedWorkItem; that is how this runner delivers timers
      .apply("To KeyedWorkItem", ParDo.of(new ToKeyedWorkItem<>()))
      .setCoder(KeyedWorkItemCoder.of(keyCoder, kvCoder, windowCoder))

      // Because of the intervening GBK, we may have abused the windowing strategy
      // of the input, which should be transferred to the output in a straightforward manner
      // according to what ParDo already does.
      .setWindowingStrategyInternal(inputWindowingStrategy);
}
 

@Override
public void translateNode(
    PTransform<PCollection<KV<K, InputT>>, PCollection<KV<K, Iterable<InputT>>>> transform,
    FlinkStreamingTranslationContext context) {

  PCollection<KV<K, InputT>> input = context.getInput(transform);

  @SuppressWarnings("unchecked")
  WindowingStrategy<?, BoundedWindow> windowingStrategy =
      (WindowingStrategy<?, BoundedWindow>) input.getWindowingStrategy();

  KvCoder<K, InputT> inputKvCoder = (KvCoder<K, InputT>) input.getCoder();

  SingletonKeyedWorkItemCoder<K, InputT> workItemCoder =
      SingletonKeyedWorkItemCoder.of(
          inputKvCoder.getKeyCoder(),
          inputKvCoder.getValueCoder(),
          input.getWindowingStrategy().getWindowFn().windowCoder());

  DataStream<WindowedValue<KV<K, InputT>>> inputDataStream = context.getInputDataStream(input);

  WindowedValue.FullWindowedValueCoder<SingletonKeyedWorkItem<K, InputT>>
      windowedWorkItemCoder =
          WindowedValue.getFullCoder(
              workItemCoder, input.getWindowingStrategy().getWindowFn().windowCoder());

  CoderTypeInformation<WindowedValue<SingletonKeyedWorkItem<K, InputT>>> workItemTypeInfo =
      new CoderTypeInformation<>(windowedWorkItemCoder);

  DataStream<WindowedValue<SingletonKeyedWorkItem<K, InputT>>> workItemStream =
      inputDataStream
          .flatMap(new ToKeyedWorkItem<>(context.getPipelineOptions()))
          .returns(workItemTypeInfo)
          .name("ToKeyedWorkItem");

  WorkItemKeySelector keySelector = new WorkItemKeySelector<>(inputKvCoder.getKeyCoder());

  KeyedStream<WindowedValue<SingletonKeyedWorkItem<K, InputT>>, ByteBuffer>
      keyedWorkItemStream =
          workItemStream.keyBy(new WorkItemKeySelector<>(inputKvCoder.getKeyCoder()));

  SystemReduceFn<K, InputT, Iterable<InputT>, Iterable<InputT>, BoundedWindow> reduceFn =
      SystemReduceFn.buffering(inputKvCoder.getValueCoder());

  Coder<WindowedValue<KV<K, Iterable<InputT>>>> outputCoder =
      context.getWindowedInputCoder(context.getOutput(transform));
  TypeInformation<WindowedValue<KV<K, Iterable<InputT>>>> outputTypeInfo =
      context.getTypeInfo(context.getOutput(transform));

  TupleTag<KV<K, Iterable<InputT>>> mainTag = new TupleTag<>("main output");

  String fullName = getCurrentTransformName(context);
  WindowDoFnOperator<K, InputT, Iterable<InputT>> doFnOperator =
      new WindowDoFnOperator<>(
          reduceFn,
          fullName,
          (Coder) windowedWorkItemCoder,
          mainTag,
          Collections.emptyList(),
          new DoFnOperator.MultiOutputOutputManagerFactory<>(mainTag, outputCoder),
          windowingStrategy,
          new HashMap<>(), /* side-input mapping */
          Collections.emptyList(), /* side inputs */
          context.getPipelineOptions(),
          inputKvCoder.getKeyCoder(),
          keySelector);

  // our operator expects WindowedValue<KeyedWorkItem> while our input stream
  // is WindowedValue<SingletonKeyedWorkItem>, which is fine but Java doesn't like it ...
  @SuppressWarnings("unchecked")
  SingleOutputStreamOperator<WindowedValue<KV<K, Iterable<InputT>>>> outDataStream =
      keyedWorkItemStream
          .transform(fullName, outputTypeInfo, (OneInputStreamOperator) doFnOperator)
          .uid(fullName);

  context.setOutputDataStream(context.getOutput(transform), outDataStream);
}
 

@Override
public void translateTransform(
    PTransform<PCollection<KV<K, InputT>>, PCollection<KV<K, OutputT>>> transform,
    TranslationContext context) {

  Combine.PerKey combineTransform = (Combine.PerKey) transform;
  @SuppressWarnings("unchecked")
  final PCollection<KV<K, InputT>> input = (PCollection<KV<K, InputT>>) context.getInput();
  @SuppressWarnings("unchecked")
  final PCollection<KV<K, OutputT>> output = (PCollection<KV<K, OutputT>>) context.getOutput();
  @SuppressWarnings("unchecked")
  final Combine.CombineFn<InputT, AccumT, OutputT> combineFn =
      (Combine.CombineFn<InputT, AccumT, OutputT>) combineTransform.getFn();
  WindowingStrategy<?, ?> windowingStrategy = input.getWindowingStrategy();

  Dataset<WindowedValue<KV<K, InputT>>> inputDataset = context.getDataset(input);

  KvCoder<K, InputT> inputCoder = (KvCoder<K, InputT>) input.getCoder();
  Coder<K> keyCoder = inputCoder.getKeyCoder();
  KvCoder<K, OutputT> outputKVCoder = (KvCoder<K, OutputT>) output.getCoder();
  Coder<OutputT> outputCoder = outputKVCoder.getValueCoder();

  KeyValueGroupedDataset<K, WindowedValue<KV<K, InputT>>> groupedDataset =
      inputDataset.groupByKey(KVHelpers.extractKey(), EncoderHelpers.fromBeamCoder(keyCoder));

  Coder<AccumT> accumulatorCoder = null;
  try {
    accumulatorCoder =
        combineFn.getAccumulatorCoder(
            input.getPipeline().getCoderRegistry(), inputCoder.getValueCoder());
  } catch (CannotProvideCoderException e) {
    throw new RuntimeException(e);
  }

  Dataset<Tuple2<K, Iterable<WindowedValue<OutputT>>>> combinedDataset =
      groupedDataset.agg(
          new AggregatorCombiner<K, InputT, AccumT, OutputT, BoundedWindow>(
                  combineFn, windowingStrategy, accumulatorCoder, outputCoder)
              .toColumn());

  // expand the list into separate elements and put the key back into the elements
  WindowedValue.WindowedValueCoder<KV<K, OutputT>> wvCoder =
      WindowedValue.FullWindowedValueCoder.of(
          outputKVCoder, input.getWindowingStrategy().getWindowFn().windowCoder());
  Dataset<WindowedValue<KV<K, OutputT>>> outputDataset =
      combinedDataset.flatMap(
          (FlatMapFunction<
                  Tuple2<K, Iterable<WindowedValue<OutputT>>>, WindowedValue<KV<K, OutputT>>>)
              tuple2 -> {
                K key = tuple2._1();
                Iterable<WindowedValue<OutputT>> windowedValues = tuple2._2();
                List<WindowedValue<KV<K, OutputT>>> result = new ArrayList<>();
                for (WindowedValue<OutputT> windowedValue : windowedValues) {
                  KV<K, OutputT> kv = KV.of(key, windowedValue.getValue());
                  result.add(
                      WindowedValue.of(
                          kv,
                          windowedValue.getTimestamp(),
                          windowedValue.getWindows(),
                          windowedValue.getPane()));
                }
                return result.iterator();
              },
          EncoderHelpers.fromBeamCoder(wvCoder));
  context.putDataset(output, outputDataset);
}
 

public <T> Coder<T> getInputCoder(PCollection<T> collection) {
  return collection.getCoder();
}