Java Code Examples for org.apache.beam.sdk.coders.IterableCoder#of()

public AggregatorCombiner(
    Combine.CombineFn<InputT, AccumT, OutputT> combineFn,
    WindowingStrategy<?, ?> windowingStrategy,
    Coder<AccumT> accumulatorCoder,
    Coder<OutputT> outputCoder) {
  this.combineFn = combineFn;
  this.windowingStrategy = (WindowingStrategy<InputT, W>) windowingStrategy;
  this.timestampCombiner = windowingStrategy.getTimestampCombiner();
  this.accumulatorCoder =
      IterableCoder.of(
          WindowedValue.FullWindowedValueCoder.of(
              accumulatorCoder, windowingStrategy.getWindowFn().windowCoder()));
  this.outputCoder =
      IterableCoder.of(
          WindowedValue.FullWindowedValueCoder.of(
              outputCoder, windowingStrategy.getWindowFn().windowCoder()));
}
 

@Test
public void serializableAccumulatorSerializationTest()
    throws IOException, ClassNotFoundException {
  Iterable<WindowedValue<Integer>> accumulatedValue =
      Arrays.asList(winVal(0), winVal(1), winVal(3), winVal(4));

  final WindowedValue.FullWindowedValueCoder<Integer> wvaCoder =
      WindowedValue.FullWindowedValueCoder.of(
          BigEndianIntegerCoder.of(), GlobalWindow.Coder.INSTANCE);

  final IterableCoder<WindowedValue<Integer>> iterAccumCoder = IterableCoder.of(wvaCoder);

  ValueAndCoderLazySerializable<Iterable<WindowedValue<Integer>>> accUnderTest =
      ValueAndCoderLazySerializable.of(accumulatedValue, iterAccumCoder);

  ByteArrayOutputStream inMemOut = new ByteArrayOutputStream();
  ObjectOutputStream oos = new ObjectOutputStream(inMemOut);
  oos.writeObject(accUnderTest);

  ObjectInputStream ois = new ObjectInputStream(new ByteArrayInputStream(inMemOut.toByteArray()));

  @SuppressWarnings("unchecked")
  ValueAndCoderLazySerializable<Iterable<WindowedValue<Integer>>> materialized =
      (ValueAndCoderLazySerializable<Iterable<WindowedValue<Integer>>>) ois.readObject();
  assertEquals(accumulatedValue, materialized.getOrDecode(iterAccumCoder));
}
 

UpdateStateByKeyFunction(
    final List<Integer> sourceIds,
    final WindowingStrategy<?, W> windowingStrategy,
    final FullWindowedValueCoder<InputT> wvCoder,
    final Coder<K> keyCoder,
    final SerializablePipelineOptions options,
    final String logPrefix) {
  this.wvCoder = wvCoder;
  this.keyCoder = keyCoder;
  this.sourceIds = sourceIds;
  this.timerDataCoder = timerDataCoderOf(windowingStrategy);
  this.windowingStrategy = windowingStrategy;
  this.options = options;
  this.itrWvCoder = IterableCoder.of(wvCoder);
  this.logPrefix = logPrefix;
  this.wvKvIterCoder =
      windowedValueKeyValueCoderOf(
          keyCoder,
          wvCoder.getValueCoder(),
          ((FullWindowedValueCoder<InputT>) wvCoder).getWindowCoder());
}
 

@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);
}
 

static CoderTranslator<IterableCoder<?>> iterable() {
  return new SimpleStructuredCoderTranslator<IterableCoder<?>>() {
    @Override
    public List<? extends Coder<?>> getComponents(IterableCoder<?> from) {
      return Collections.singletonList(from.getElemCoder());
    }

    @Override
    public IterableCoder<?> fromComponents(List<Coder<?>> components) {
      return IterableCoder.of(components.get(0));
    }
  };
}
 

@SuppressWarnings("unchecked")
WindowedAccumulatorCoder(
    Function<InputT, ValueT> toValue,
    Coder<BoundedWindow> windowCoder,
    Comparator<BoundedWindow> windowComparator,
    Coder<AccumT> accumCoder,
    WindowedAccumulator.Type type) {

  this.toValue = toValue;
  this.accumCoder = WindowedValue.FullWindowedValueCoder.of(accumCoder, windowCoder);
  this.windowComparator = windowComparator;
  this.wrap = IterableCoder.of(this.accumCoder);
  this.type = type;
}
 

private static <ReadT, WriteT>
    TransformEvaluator<View.CreatePCollectionView<ReadT, WriteT>> createPCollView() {
  return new TransformEvaluator<View.CreatePCollectionView<ReadT, WriteT>>() {
    @Override
    public void evaluate(
        View.CreatePCollectionView<ReadT, WriteT> transform, EvaluationContext context) {
      Iterable<? extends WindowedValue<?>> iter =
          context.getWindowedValues(context.getInput(transform));
      PCollectionView<WriteT> output = transform.getView();
      Coder<Iterable<WindowedValue<?>>> coderInternal =
          (Coder)
              IterableCoder.of(
                  WindowedValue.getFullCoder(
                      output.getCoderInternal(),
                      output.getWindowingStrategyInternal().getWindowFn().windowCoder()));

      @SuppressWarnings("unchecked")
      Iterable<WindowedValue<?>> iterCast = (Iterable<WindowedValue<?>>) iter;

      context.putPView(output, iterCast, coderInternal);
    }

    @Override
    public String toNativeString() {
      return "<createPCollectionView>";
    }
  };
}
 

/**
 * Returns a {@link CloudObjectTranslator} that produces a {@link CloudObject} that is of kind
 * "stream".
 */
public static CloudObjectTranslator<IterableCoder> stream() {
  return new CloudObjectTranslator<IterableCoder>() {
    @Override
    public CloudObject toCloudObject(IterableCoder target, SdkComponents sdkComponents) {
      CloudObject result = CloudObject.forClassName(CloudObjectKinds.KIND_STREAM);
      Structs.addBoolean(result, PropertyNames.IS_STREAM_LIKE, true);
      return addComponents(
          result, Collections.<Coder<?>>singletonList(target.getElemCoder()), sdkComponents);
    }

    @Override
    public IterableCoder fromCloudObject(CloudObject object) {
      List<Coder<?>> components = getComponents(object);
      checkArgument(components.size() == 1, "Expecting 1 component, got %s", components.size());
      return IterableCoder.of(components.get(0));
    }

    @Override
    public Class<? extends IterableCoder> getSupportedClass() {
      return IterableCoder.class;
    }

    @Override
    public String cloudObjectClassName() {
      return CloudObjectKinds.KIND_STREAM;
    }
  };
}
 

/**
 * Creates a {@link ReduceFnTester} for the given {@link WindowingStrategy}, creating a {@link
 * TriggerStateMachine} from its {@link Trigger}.
 */
public static <W extends BoundedWindow>
    ReduceFnTester<Integer, Iterable<Integer>, W> nonCombining(
        WindowingStrategy<?, W> windowingStrategy) throws Exception {
  return new ReduceFnTester<>(
      windowingStrategy,
      TriggerStateMachines.stateMachineForTrigger(
          TriggerTranslation.toProto(windowingStrategy.getTrigger())),
      SystemReduceFn.buffering(VarIntCoder.of()),
      IterableCoder.of(VarIntCoder.of()),
      PipelineOptionsFactory.create(),
      NullSideInputReader.empty());
}
 

private KeyedWorkItemCoder(
    Coder<K> keyCoder, Coder<ElemT> elemCoder, Coder<? extends BoundedWindow> windowCoder) {
  this.keyCoder = keyCoder;
  this.elemCoder = elemCoder;
  this.windowCoder = windowCoder;
  this.timersCoder = IterableCoder.of(TimerDataCoderV2.of(windowCoder));
  this.elemsCoder = IterableCoder.of(FullWindowedValueCoder.of(elemCoder, windowCoder));
}
 

@Test
public void serializableAccumulatorKryoTest() {
  Iterable<WindowedValue<Integer>> accumulatedValue =
      Arrays.asList(winVal(0), winVal(1), winVal(3), winVal(4));

  final WindowedValue.FullWindowedValueCoder<Integer> wvaCoder =
      WindowedValue.FullWindowedValueCoder.of(
          BigEndianIntegerCoder.of(), GlobalWindow.Coder.INSTANCE);

  final IterableCoder<WindowedValue<Integer>> iterAccumCoder = IterableCoder.of(wvaCoder);

  ValueAndCoderLazySerializable<Iterable<WindowedValue<Integer>>> accUnderTest =
      ValueAndCoderLazySerializable.of(accumulatedValue, iterAccumCoder);

  ValueAndCoderKryoSerializer kryoSerializer = new ValueAndCoderKryoSerializer();
  Kryo kryo = new Kryo();
  kryo.register(ValueAndCoderLazySerializable.class, kryoSerializer);

  ByteArrayOutputStream inMemOut = new ByteArrayOutputStream();
  Output out = new Output(inMemOut);
  kryo.writeObject(out, accUnderTest);
  out.close();

  Input input = new Input(new ByteArrayInputStream(inMemOut.toByteArray()));

  @SuppressWarnings("unchecked")
  ValueAndCoderLazySerializable<Iterable<WindowedValue<Integer>>> materialized =
      (ValueAndCoderLazySerializable<Iterable<WindowedValue<Integer>>>)
          kryo.readObject(input, ValueAndCoderLazySerializable.class);
  input.close();

  assertEquals(accumulatedValue, materialized.getOrDecode(iterAccumCoder));
}
 

private static byte[] mapAsBytes(Map<String, String> stringMap) throws IOException {
  IterableCoder<KV<String, String>> coder =
      IterableCoder.of(KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()));
  List<KV<String, String>> stringList =
      stringMap.entrySet().stream()
          .map(kv -> KV.of(kv.getKey(), kv.getValue()))
          .collect(Collectors.toList());
  ByteArrayOutputStream baos = new ByteArrayOutputStream();
  coder.encode(stringList, baos);
  return baos.toByteArray();
}
 

/** Returns the coder used for a given primitive type. */
public static <T> Coder<T> coderForFieldType(FieldType fieldType) {
  Coder<T> coder;
  switch (fieldType.getTypeName()) {
    case ROW:
      coder = (Coder<T>) SchemaCoder.of(fieldType.getRowSchema());
      break;
    case ARRAY:
      coder = (Coder<T>) ListCoder.of(coderForFieldType(fieldType.getCollectionElementType()));
      break;
    case ITERABLE:
      coder =
          (Coder<T>) IterableCoder.of(coderForFieldType(fieldType.getCollectionElementType()));
      break;
    case MAP:
      coder =
          (Coder<T>)
              MapCoder.of(
                  coderForFieldType(fieldType.getMapKeyType()),
                  coderForFieldType(fieldType.getMapValueType()));
      break;
    case LOGICAL_TYPE:
      coder =
          new LogicalTypeCoder(
              fieldType.getLogicalType(),
              coderForFieldType(fieldType.getLogicalType().getBaseType()));
      break;
    default:
      coder = (Coder<T>) CODER_MAP.get(fieldType.getTypeName());
  }
  Preconditions.checkNotNull(coder, "Unexpected field type " + fieldType.getTypeName());
  if (fieldType.getNullable()) {
    coder = NullableCoder.of(coder);
  }
  return coder;
}
 

@Override
public void addCollectionToSingletonOutput(
    PCollection<?> inputValue, String outputName, PCollectionView<?> outputValue) {
  translator.producers.put(outputValue, translator.currentTransform);
  Coder<?> inputValueCoder = checkNotNull(translator.outputCoders.get(inputValue));
  // The inputValueCoder for the input PCollection should be some
  // WindowedValueCoder of the input PCollection's element
  // coder.
  checkState(inputValueCoder instanceof WindowedValue.WindowedValueCoder);
  // The outputValueCoder for the output should be an
  // IterableCoder of the inputValueCoder. This is a property
  // of the backend "CollectionToSingleton" step.
  Coder<?> outputValueCoder = IterableCoder.of(inputValueCoder);
  addOutput(outputName, outputValue, outputValueCoder);
}
 

private static byte[] listAsBytes(List<String> stringList) throws IOException {
  IterableCoder<String> coder = IterableCoder.of(StringUtf8Coder.of());
  ByteArrayOutputStream baos = new ByteArrayOutputStream();
  coder.encode(stringList, baos);
  return baos.toByteArray();
}
 

@Override
public Coder<Iterable<T>> getDefaultOutputCoder(CoderRegistry registry, Coder<T> inputCoder) {
  return IterableCoder.of(inputCoder);
}
 

@SuppressWarnings("rawtypes")
private IterableCoder tagListCoder(int unionTag) {
  return IterableCoder.of(unionCoder.getElementCoders().get(unionTag));
}
 

private <K, V> SingleOutputStreamOperator<WindowedValue<KV<K, Iterable<V>>>> addGBK(
    DataStream<WindowedValue<KV<K, V>>> inputDataStream,
    WindowingStrategy<?, ?> windowingStrategy,
    WindowedValueCoder<KV<K, V>> windowedInputCoder,
    String operatorName,
    StreamingTranslationContext context) {
  KvCoder<K, V> inputElementCoder = (KvCoder<K, V>) windowedInputCoder.getValueCoder();

  SingletonKeyedWorkItemCoder<K, V> workItemCoder =
      SingletonKeyedWorkItemCoder.of(
          inputElementCoder.getKeyCoder(),
          inputElementCoder.getValueCoder(),
          windowingStrategy.getWindowFn().windowCoder());

  WindowedValue.FullWindowedValueCoder<SingletonKeyedWorkItem<K, V>> windowedWorkItemCoder =
      WindowedValue.getFullCoder(workItemCoder, windowingStrategy.getWindowFn().windowCoder());

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

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

  WorkItemKeySelector<K, V> keySelector =
      new WorkItemKeySelector<>(inputElementCoder.getKeyCoder());

  KeyedStream<WindowedValue<SingletonKeyedWorkItem<K, V>>, ByteBuffer> keyedWorkItemStream =
      workItemStream.keyBy(keySelector);

  SystemReduceFn<K, V, Iterable<V>, Iterable<V>, BoundedWindow> reduceFn =
      SystemReduceFn.buffering(inputElementCoder.getValueCoder());

  Coder<Iterable<V>> accumulatorCoder = IterableCoder.of(inputElementCoder.getValueCoder());

  Coder<WindowedValue<KV<K, Iterable<V>>>> outputCoder =
      WindowedValue.getFullCoder(
          KvCoder.of(inputElementCoder.getKeyCoder(), accumulatorCoder),
          windowingStrategy.getWindowFn().windowCoder());

  TypeInformation<WindowedValue<KV<K, Iterable<V>>>> outputTypeInfo =
      new CoderTypeInformation<>(outputCoder);

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

  WindowDoFnOperator<K, V, Iterable<V>> doFnOperator =
      new WindowDoFnOperator<>(
          reduceFn,
          operatorName,
          (Coder) windowedWorkItemCoder,
          mainTag,
          Collections.emptyList(),
          new DoFnOperator.MultiOutputOutputManagerFactory(mainTag, outputCoder),
          windowingStrategy,
          new HashMap<>(), /* side-input mapping */
          Collections.emptyList(), /* side inputs */
          context.getPipelineOptions(),
          inputElementCoder.getKeyCoder(),
          (KeySelector) keySelector /* key selector */);

  SingleOutputStreamOperator<WindowedValue<KV<K, Iterable<V>>>> outputDataStream =
      keyedWorkItemStream.transform(
          operatorName, outputTypeInfo, (OneInputStreamOperator) doFnOperator);

  return outputDataStream;
}
 

@Override
public Coder<Iterable<T>> getDefaultOutputCoder(CoderRegistry registry, Coder<T> inputCoder) {
  return IterableCoder.of(inputCoder);
}
 

@Test
public void testCompoundCodersForExternalConfiguration() throws Exception {
  ExternalTransforms.ExternalConfigurationPayload.Builder builder =
      ExternalTransforms.ExternalConfigurationPayload.newBuilder();

  builder.putConfiguration(
      "config_key1",
      ExternalTransforms.ConfigValue.newBuilder()
          .addCoderUrn(BeamUrns.getUrn(RunnerApi.StandardCoders.Enum.VARINT))
          .setPayload(ByteString.copyFrom(new byte[] {1}))
          .build());

  List<byte[]> byteList =
      ImmutableList.of("testing", "compound", "coders").stream()
          .map(str -> str.getBytes(Charsets.UTF_8))
          .collect(Collectors.toList());
  IterableCoder<byte[]> compoundCoder = IterableCoder.of(ByteArrayCoder.of());
  ByteArrayOutputStream baos = new ByteArrayOutputStream();
  compoundCoder.encode(byteList, baos);

  builder.putConfiguration(
      "config_key2",
      ExternalTransforms.ConfigValue.newBuilder()
          .addCoderUrn(BeamUrns.getUrn(RunnerApi.StandardCoders.Enum.ITERABLE))
          .addCoderUrn(BeamUrns.getUrn(RunnerApi.StandardCoders.Enum.BYTES))
          .setPayload(ByteString.copyFrom(baos.toByteArray()))
          .build());

  List<KV<byte[], Long>> byteKvList =
      ImmutableList.of("testing", "compound", "coders").stream()
          .map(str -> KV.of(str.getBytes(Charsets.UTF_8), (long) str.length()))
          .collect(Collectors.toList());
  IterableCoder<KV<byte[], Long>> compoundCoder2 =
      IterableCoder.of(KvCoder.of(ByteArrayCoder.of(), VarLongCoder.of()));
  ByteArrayOutputStream baos2 = new ByteArrayOutputStream();
  compoundCoder2.encode(byteKvList, baos2);

  builder.putConfiguration(
      "config_key3",
      ExternalTransforms.ConfigValue.newBuilder()
          .addCoderUrn(BeamUrns.getUrn(RunnerApi.StandardCoders.Enum.ITERABLE))
          .addCoderUrn(BeamUrns.getUrn(RunnerApi.StandardCoders.Enum.KV))
          .addCoderUrn(BeamUrns.getUrn(RunnerApi.StandardCoders.Enum.BYTES))
          .addCoderUrn(BeamUrns.getUrn(RunnerApi.StandardCoders.Enum.VARINT))
          .setPayload(ByteString.copyFrom(baos2.toByteArray()))
          .build());

  List<KV<List<Long>, byte[]>> byteKvListWithListKey =
      ImmutableList.of("testing", "compound", "coders").stream()
          .map(
              str ->
                  KV.of(
                      Collections.singletonList((long) str.length()),
                      str.getBytes(Charsets.UTF_8)))
          .collect(Collectors.toList());
  Coder compoundCoder3 =
      IterableCoder.of(KvCoder.of(IterableCoder.of(VarLongCoder.of()), ByteArrayCoder.of()));
  ByteArrayOutputStream baos3 = new ByteArrayOutputStream();
  compoundCoder3.encode(byteKvListWithListKey, baos3);

  builder.putConfiguration(
      "config_key4",
      ExternalTransforms.ConfigValue.newBuilder()
          .addCoderUrn(BeamUrns.getUrn(RunnerApi.StandardCoders.Enum.ITERABLE))
          .addCoderUrn(BeamUrns.getUrn(RunnerApi.StandardCoders.Enum.KV))
          .addCoderUrn(BeamUrns.getUrn(RunnerApi.StandardCoders.Enum.ITERABLE))
          .addCoderUrn(BeamUrns.getUrn(RunnerApi.StandardCoders.Enum.VARINT))
          .addCoderUrn(BeamUrns.getUrn(RunnerApi.StandardCoders.Enum.BYTES))
          .setPayload(ByteString.copyFrom(baos3.toByteArray()))
          .build());

  ExternalTransforms.ExternalConfigurationPayload externalConfig = builder.build();
  TestConfig config = new TestConfig();
  ExpansionService.ExternalTransformRegistrarLoader.populateConfiguration(config, externalConfig);

  assertThat(config.configKey1, Matchers.is(1L));
  assertArrayEquals(Iterables.toArray(config.configKey2, byte[].class), byteList.toArray());
  assertArrayEquals(Iterables.toArray(config.configKey3, KV.class), byteKvList.toArray());
  assertArrayEquals(
      Iterables.toArray(config.configKey4, KV.class), byteKvListWithListKey.toArray());
}