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

@Test
public void testCodingWithKvCoderClassToBeEncoded() throws IOException {
  final KryoRegistrar registrar =
      k -> {
        k.register(TestClass.class);
        k.register(ClassToBeEncoded.class);
      };

  final ListCoder<Void> listCoder = ListCoder.of(VoidCoder.of());
  final KvCoder<ClassToBeEncoded, List<Void>> kvCoder =
      KvCoder.of(KryoCoder.of(OPTIONS, registrar), listCoder);
  final List<Void> inputValue = new ArrayList<>();
  inputValue.add(null);
  inputValue.add(null);

  final ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();

  final ClassToBeEncoded inputKey = new ClassToBeEncoded("something", 1, 0.2);
  kvCoder.encode(KV.of(inputKey, inputValue), byteArrayOutputStream);

  final KV<ClassToBeEncoded, List<Void>> decoded =
      kvCoder.decode(new ByteArrayInputStream(byteArrayOutputStream.toByteArray()));

  assertNotNull(decoded);
  assertNotNull(decoded.getKey());
  assertEquals(inputKey, decoded.getKey());

  assertNotNull(decoded.getValue());
  assertEquals(inputValue, decoded.getValue());
}
 

private GrowthStateCoder(
    Coder<OutputT> outputCoder, Coder<TerminationStateT> terminationStateCoder) {
  this.outputCoder = outputCoder;
  this.terminationStateCoder = terminationStateCoder;
  this.timestampedOutputCoder =
      ListCoder.of(TimestampedValue.TimestampedValueCoder.of(outputCoder));
}
 

public SerializableArrayViaCoder(Coder<T> elementCoder, T[] value) {
  this.coder = ListCoder.of(elementCoder);
  this.value = value;
  try {
    this.encodedValue = CoderUtils.encodeToByteArray(coder, Arrays.asList(value));
  } catch (CoderException exc) {
    throw UserCodeException.wrap(exc);
  }
}
 

/** 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;
}
 

FlinkBroadcastBagState(
    OperatorStateBackend flinkStateBackend,
    StateTag<BagState<T>> address,
    StateNamespace namespace,
    Coder<T> coder) {
  super(flinkStateBackend, address.getId(), namespace, ListCoder.of(coder));

  this.namespace = namespace;
  this.address = address;
}
 

public Coder<List> getAccumulatorCoder() {
    AvroCoder valueCoder = null;
    if (avroSchemaStr != null) {
        valueCoder = AvroCoder.of(new Schema.Parser().parse(avroSchemaStr));
    }
    return (Coder<List>) (avroSchemaStr == null ? ListCoder.of(NullableCoder.of(StringUtf8Coder.of()))
            : ListCoder.of(NullableCoder.of(valueCoder)));
}
 

@Override
public Coder<List<AccumT>> getAccumulatorCoder(CoderRegistry registry, Coder<AccumT> inputCoder)
    throws CannotProvideCoderException {
  return ListCoder.of(accumCoder);
}
 

@Override
public Coder<List<AccumT>> getAccumulatorCoder(CoderRegistry registry, Coder<AccumT> inputCoder)
    throws CannotProvideCoderException {
  return ListCoder.of(accumCoder);
}
 

@Override
public Coder<List<T>> getAccumulatorCoder(CoderRegistry registry, Coder<T> inputCoder) {
  return ListCoder.of(inputCoder);
}
 

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

@Override
public Coder<List<T>> getAccumulatorCoder(CoderRegistry registry, Coder<T> inputCoder) {
  return ListCoder.of(inputCoder);
}
 

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

@Override
public Coder<List<T>> getAccumulatorCoder(CoderRegistry registry, Coder<T> inputCoder) {
  return ListCoder.of(inputCoder);
}
 

@Test
public void testFetchGlobalDataCacheOverflow() throws Exception {
  Coder<List<String>> coder = ListCoder.of(StringUtf8Coder.of());

  ByteString.Output stream = ByteString.newOutput();
  coder.encode(Arrays.asList("data1"), stream, Coder.Context.OUTER);
  ByteString encodedIterable1 = stream.toByteString();
  stream = ByteString.newOutput();
  coder.encode(Arrays.asList("data2"), stream, Coder.Context.OUTER);
  ByteString encodedIterable2 = stream.toByteString();

  Cache<StateFetcher.SideInputId, StateFetcher.SideInputCacheEntry> cache =
      CacheBuilder.newBuilder().build();

  StateFetcher fetcher = new StateFetcher(server, cache);

  PCollectionView<String> view1 =
      TestPipeline.create().apply(Create.empty(StringUtf8Coder.of())).apply(View.asSingleton());

  PCollectionView<String> view2 =
      TestPipeline.create().apply(Create.empty(StringUtf8Coder.of())).apply(View.asSingleton());

  String tag1 = view1.getTagInternal().getId();
  String tag2 = view2.getTagInternal().getId();

  // Test four calls in a row. First, fetch view1, then view2 (which evicts view1 from the cache),
  // then view 1 again twice.
  when(server.getSideInputData(any(Windmill.GlobalDataRequest.class)))
      .thenReturn(
          buildGlobalDataResponse(tag1, ByteString.EMPTY, true, encodedIterable1),
          buildGlobalDataResponse(tag2, ByteString.EMPTY, true, encodedIterable2),
          buildGlobalDataResponse(tag1, ByteString.EMPTY, true, encodedIterable1));

  assertEquals(
      "data1",
      fetcher
          .fetchSideInput(
              view1,
              GlobalWindow.INSTANCE,
              STATE_FAMILY,
              SideInputState.UNKNOWN,
              readStateSupplier)
          .orNull());
  assertEquals(
      "data2",
      fetcher
          .fetchSideInput(
              view2,
              GlobalWindow.INSTANCE,
              STATE_FAMILY,
              SideInputState.UNKNOWN,
              readStateSupplier)
          .orNull());
  cache.invalidateAll();
  assertEquals(
      "data1",
      fetcher
          .fetchSideInput(
              view1,
              GlobalWindow.INSTANCE,
              STATE_FAMILY,
              SideInputState.UNKNOWN,
              readStateSupplier)
          .orNull());
  assertEquals(
      "data1",
      fetcher
          .fetchSideInput(
              view1,
              GlobalWindow.INSTANCE,
              STATE_FAMILY,
              SideInputState.UNKNOWN,
              readStateSupplier)
          .orNull());

  ArgumentCaptor<Windmill.GlobalDataRequest> captor =
      ArgumentCaptor.forClass(Windmill.GlobalDataRequest.class);

  verify(server, times(3)).getSideInputData(captor.capture());
  verifyNoMoreInteractions(server);

  assertThat(
      captor.getAllValues(),
      contains(
          buildGlobalDataRequest(tag1, ByteString.EMPTY),
          buildGlobalDataRequest(tag2, ByteString.EMPTY),
          buildGlobalDataRequest(tag1, ByteString.EMPTY)));
}
 

private SparkBagState(StateNamespace namespace, StateTag<BagState<T>> address, Coder<T> coder) {
  super(namespace, address, ListCoder.of(coder));
}
 

@Override
public Coder<List<T>> getAccumulatorCoder(CoderRegistry registry, Coder<T> inputCoder) {
  return ListCoder.of(inputCoder);
}
 

public QuantileStateCoder(ComparatorT compareFn, Coder<T> elementCoder) {
  this.compareFn = compareFn;
  this.elementCoder = elementCoder;
  this.elementListCoder = ListCoder.of(elementCoder);
}
 

/**
 * Attempt to infer the type for some very common Apache Beam parameterized types.
 *
 * <p>TODO: Instead, build a TypeDescriptor so that the {@link CoderRegistry} is invoked for the
 * type instead of hard coding the coders for common types.
 */
private static Coder<?> inferCoderFromObject(
    CoderRegistry coderRegistry, SchemaRegistry schemaRegistry, Object o)
    throws CannotProvideCoderException {

  if (o == null) {
    return VoidCoder.of();
  }

  try {
    return SchemaCoder.of(
        schemaRegistry.getSchema(o.getClass()),
        TypeDescriptor.of(o.getClass()),
        (SerializableFunction) schemaRegistry.getToRowFunction(o.getClass()),
        (SerializableFunction) schemaRegistry.getFromRowFunction(o.getClass()));
  } catch (NoSuchSchemaException e) {
    // No schema.
  }

  if (o instanceof TimestampedValue) {
    return TimestampedValueCoder.of(
        inferCoderFromObject(coderRegistry, schemaRegistry, ((TimestampedValue) o).getValue()));
  } else if (o instanceof List) {
    return ListCoder.of(inferCoderFromObjects(coderRegistry, schemaRegistry, (Iterable) o));
  } else if (o instanceof Set) {
    return SetCoder.of(inferCoderFromObjects(coderRegistry, schemaRegistry, (Iterable) o));
  } else if (o instanceof Collection) {
    return CollectionCoder.of(inferCoderFromObjects(coderRegistry, schemaRegistry, (Iterable) o));
  } else if (o instanceof Iterable) {
    return IterableCoder.of(inferCoderFromObjects(coderRegistry, schemaRegistry, (Iterable) o));
  } else if (o instanceof Map) {
    return MapCoder.of(
        inferCoderFromObjects(coderRegistry, schemaRegistry, ((Map) o).keySet()),
        inferCoderFromObjects(coderRegistry, schemaRegistry, ((Map) o).entrySet()));
  } else if (o instanceof KV) {
    return KvCoder.of(
        inferCoderFromObject(coderRegistry, schemaRegistry, ((KV) o).getKey()),
        inferCoderFromObject(coderRegistry, schemaRegistry, ((KV) o).getValue()));
  } else {
    return coderRegistry.getCoder(o.getClass());
  }
}
 

public BoundedHeapCoder(int maximumSize, ComparatorT compareFn, Coder<T> elementCoder) {
  listCoder = ListCoder.of(elementCoder);
  this.compareFn = compareFn;
  this.maximumSize = maximumSize;
}
 

@Test
public void testNestedCoding() throws Throwable {
  Coder<List<ByteString>> listCoder = ListCoder.of(TEST_CODER);
  CoderProperties.coderDecodeEncodeContentsEqual(listCoder, TEST_VALUES);
  CoderProperties.coderDecodeEncodeContentsInSameOrder(listCoder, TEST_VALUES);
}