How to use

registerKryoType

method

in

org.apache.flink.api.common.ExecutionConfig

/**
 * Registers the given type with the serialization stack. If the type is eventually
 * serialized as a POJO, then the type is registered with the POJO serializer. If the
 * type ends up being serialized with Kryo, then it will be registered at Kryo to make
 * sure that only tags are written.
 *
 * @param type The class of the type to register.
 */
public void registerType(Class<?> type) {
  if (type == null) {
    throw new NullPointerException("Cannot register null type class.");
  }
  TypeInformation<?> typeInfo = TypeExtractor.createTypeInfo(type);
  if (typeInfo instanceof PojoTypeInfo) {
    config.registerPojoType(type);
  } else {
    config.registerKryoType(type);
  }
}

public static void recursivelyRegisterType(Class<?> type, ExecutionConfig config, Set<Class<?>> alreadySeen) {
  // don't register or remember primitives
  if (type == null || type.isPrimitive() || type == Object.class) {
    return;
  }
  
  // prevent infinite recursion for recursive types
  if (!alreadySeen.add(type)) {
    return;
  }
  
  if (type.isArray()) {
    recursivelyRegisterType(type.getComponentType(), config, alreadySeen);
  }
  else {
    config.registerKryoType(type);
    // add serializers for Avro type if necessary
    AvroUtils.getAvroUtils().addAvroSerializersIfRequired(config, type);
    Field[] fields = type.getDeclaredFields();
    for (Field field : fields) {
      if (Modifier.isStatic(field.getModifiers()) || Modifier.isTransient(field.getModifiers())) {
        continue;
      }
      Type fieldType = field.getGenericType();
      recursivelyRegisterGenericType(fieldType, config, alreadySeen);
    }
  }
}

/**
 * Registers the given type with the serialization stack. If the type is
 * eventually serialized as a POJO, then the type is registered with the
 * POJO serializer. If the type ends up being serialized with Kryo, then it
 * will be registered at Kryo to make sure that only tags are written.
 *
 * @param type
 *         The class of the type to register.
 */
public void registerType(Class<?> type) {
  if (type == null) {
    throw new NullPointerException("Cannot register null type class.");
  }
  TypeInformation<?> typeInfo = TypeExtractor.createTypeInfo(type);
  if (typeInfo instanceof PojoTypeInfo) {
    config.registerPojoType(type);
  } else {
    config.registerKryoType(type);
  }
}

@Test
public void testDoubleTypeRegistration() {
  ExecutionConfig config = new ExecutionConfig();
  List<Class<?>> types = Arrays.<Class<?>>asList(Double.class, Integer.class, Double.class);
  List<Class<?>> expectedTypes = Arrays.<Class<?>>asList(Double.class, Integer.class);
  for (Class<?> tpe: types) {
    config.registerKryoType(tpe);
  }
  int counter = 0;
  for (Class<?> tpe: config.getRegisteredKryoTypes()){
    assertEquals(tpe, expectedTypes.get(counter++));
  }
  assertEquals(expectedTypes.size(), counter);
}

@Test
public void testFoldingStateInstantiation() throws Exception {
  final ExecutionConfig config = new ExecutionConfig();
  config.registerKryoType(Path.class);
  final AtomicReference<Object> descriptorCapture = new AtomicReference<>();
  StreamingRuntimeContext context = new StreamingRuntimeContext(
      createDescriptorCapturingMockOp(descriptorCapture, config),
      createMockEnvironment(),
      Collections.<String, Accumulator<?, ?>>emptyMap());
  @SuppressWarnings("unchecked")
  FoldFunction<String, TaskInfo> folder = (FoldFunction<String, TaskInfo>) mock(FoldFunction.class);
  FoldingStateDescriptor<String, TaskInfo> descr =
      new FoldingStateDescriptor<>("name", null, folder, TaskInfo.class);
  context.getFoldingState(descr);
  FoldingStateDescriptor<?, ?> descrIntercepted = (FoldingStateDescriptor<?, ?>) descriptorCapture.get();
  TypeSerializer<?> serializer = descrIntercepted.getSerializer();
  // check that the Path class is really registered, i.e., the execution config was applied
  assertTrue(serializer instanceof KryoSerializer);
  assertTrue(((KryoSerializer<?>) serializer).getKryo().getRegistration(Path.class).getId() > 0);
}

@Test
public void testReducingStateInstantiation() throws Exception {
  final ExecutionConfig config = new ExecutionConfig();
  config.registerKryoType(Path.class);
  final AtomicReference<Object> descriptorCapture = new AtomicReference<>();
  StreamingRuntimeContext context = new StreamingRuntimeContext(
      createDescriptorCapturingMockOp(descriptorCapture, config),
      createMockEnvironment(),
      Collections.<String, Accumulator<?, ?>>emptyMap());
  @SuppressWarnings("unchecked")
  ReduceFunction<TaskInfo> reducer = (ReduceFunction<TaskInfo>) mock(ReduceFunction.class);
  ReducingStateDescriptor<TaskInfo> descr =
      new ReducingStateDescriptor<>("name", reducer, TaskInfo.class);
  context.getReducingState(descr);
  StateDescriptor<?, ?> descrIntercepted = (StateDescriptor<?, ?>) descriptorCapture.get();
  TypeSerializer<?> serializer = descrIntercepted.getSerializer();
  // check that the Path class is really registered, i.e., the execution config was applied
  assertTrue(serializer instanceof KryoSerializer);
  assertTrue(((KryoSerializer<?>) serializer).getKryo().getRegistration(Path.class).getId() > 0);
}

@Test
public void testAggregatingStateInstantiation() throws Exception {
  final ExecutionConfig config = new ExecutionConfig();
  config.registerKryoType(Path.class);
  final AtomicReference<Object> descriptorCapture = new AtomicReference<>();
  StreamingRuntimeContext context = new StreamingRuntimeContext(
      createDescriptorCapturingMockOp(descriptorCapture, config),
      createMockEnvironment(),
      Collections.<String, Accumulator<?, ?>>emptyMap());
  @SuppressWarnings("unchecked")
  AggregateFunction<String, TaskInfo, String> aggregate = (AggregateFunction<String, TaskInfo, String>) mock(AggregateFunction.class);
  AggregatingStateDescriptor<String, TaskInfo, String> descr =
      new AggregatingStateDescriptor<>("name", aggregate, TaskInfo.class);
  context.getAggregatingState(descr);
  AggregatingStateDescriptor<?, ?, ?> descrIntercepted = (AggregatingStateDescriptor<?, ?, ?>) descriptorCapture.get();
  TypeSerializer<?> serializer = descrIntercepted.getSerializer();
  // check that the Path class is really registered, i.e., the execution config was applied
  assertTrue(serializer instanceof KryoSerializer);
  assertTrue(((KryoSerializer<?>) serializer).getKryo().getRegistration(Path.class).getId() > 0);
}

@Test
public void testMapStateInstantiation() throws Exception {
  final ExecutionConfig config = new ExecutionConfig();
  config.registerKryoType(Path.class);
  final AtomicReference<Object> descriptorCapture = new AtomicReference<>();
  StreamingRuntimeContext context = new StreamingRuntimeContext(
      createDescriptorCapturingMockOp(descriptorCapture, config),
      createMockEnvironment(),
      Collections.<String, Accumulator<?, ?>>emptyMap());
  MapStateDescriptor<String, TaskInfo> descr =
      new MapStateDescriptor<>("name", String.class, TaskInfo.class);
  context.getMapState(descr);
  MapStateDescriptor<?, ?> descrIntercepted = (MapStateDescriptor<?, ?>) descriptorCapture.get();
  TypeSerializer<?> valueSerializer = descrIntercepted.getValueSerializer();
  // check that the Path class is really registered, i.e., the execution config was applied
  assertTrue(valueSerializer instanceof KryoSerializer);
  assertTrue(((KryoSerializer<?>) valueSerializer).getKryo().getRegistration(Path.class).getId() > 0);
}

@Test
public void testValueStateInstantiation() throws Exception {
  final ExecutionConfig config = new ExecutionConfig();
  config.registerKryoType(Path.class);
  final AtomicReference<Object> descriptorCapture = new AtomicReference<>();
  StreamingRuntimeContext context = new StreamingRuntimeContext(
      createDescriptorCapturingMockOp(descriptorCapture, config),
      createMockEnvironment(),
      Collections.<String, Accumulator<?, ?>>emptyMap());
  ValueStateDescriptor<TaskInfo> descr = new ValueStateDescriptor<>("name", TaskInfo.class);
  context.getState(descr);
  StateDescriptor<?, ?> descrIntercepted = (StateDescriptor<?, ?>) descriptorCapture.get();
  TypeSerializer<?> serializer = descrIntercepted.getSerializer();
  // check that the Path class is really registered, i.e., the execution config was applied
  assertTrue(serializer instanceof KryoSerializer);
  assertTrue(((KryoSerializer<?>) serializer).getKryo().getRegistration(Path.class).getId() > 0);
}

@Test
public void testListStateInstantiation() throws Exception {
  final ExecutionConfig config = new ExecutionConfig();
  config.registerKryoType(Path.class);
  final AtomicReference<Object> descriptorCapture = new AtomicReference<>();
  StreamingRuntimeContext context = new StreamingRuntimeContext(
      createDescriptorCapturingMockOp(descriptorCapture, config),
      createMockEnvironment(),
      Collections.<String, Accumulator<?, ?>>emptyMap());
  ListStateDescriptor<TaskInfo> descr = new ListStateDescriptor<>("name", TaskInfo.class);
  context.getListState(descr);
  ListStateDescriptor<?> descrIntercepted = (ListStateDescriptor<?>) descriptorCapture.get();
  TypeSerializer<?> serializer = descrIntercepted.getSerializer();
  // check that the Path class is really registered, i.e., the execution config was applied
  assertTrue(serializer instanceof ListSerializer);
  TypeSerializer<?> elementSerializer = descrIntercepted.getElementSerializer();
  assertTrue(elementSerializer instanceof KryoSerializer);
  assertTrue(((KryoSerializer<?>) elementSerializer).getKryo().getRegistration(Path.class).getId() > 0);
}

executionConfig.registerKryoType(TestClassA.class);
executionConfig.registerKryoType(TestClassB.class);
executionConfig.registerKryoType(TestClassB.class); // test with B registered before A
executionConfig.registerKryoType(TestClassA.class);

ExecutionConfig executionConfig = new ExecutionConfig();
executionConfig.registerTypeWithKryoSerializer(TestRecord.class, new TestRecordSerializer());
executionConfig.registerKryoType(TestRecord.class);

@Test
public void testInitializeSerializerAfterSerializationWithCustomConfig() throws Exception {
  // guard our test assumptions.
  assertEquals("broken test assumption", -1,
      new KryoSerializer<>(String.class, new ExecutionConfig()).getKryo()
          .getRegistration(File.class).getId());
  final ExecutionConfig config = new ExecutionConfig();
  config.registerKryoType(File.class);
  final TestStateDescriptor<Path> original = new TestStateDescriptor<>("test", Path.class);
  TestStateDescriptor<Path> clone = CommonTestUtils.createCopySerializable(original);
  clone.initializeSerializerUnlessSet(config);
  // serialized one (later initialized) carries the registration
  assertTrue(((KryoSerializer<?>) clone.getSerializer()).getKryo()
      .getRegistration(File.class).getId() > 0);
}

private void registerKryoTypes(Fields fields) {
  if(fields.hasTypes()) {
    Class[] fieldTypeClasses = fields.getTypesClasses();
    for(Class fieldTypeClass : fieldTypeClasses) {
      if(!fieldTypeClass.isPrimitive() &&
          !fieldTypeClass.equals(String.class) &&
          !Writable.class.isAssignableFrom(fieldTypeClass)) {
        // register type if it is neither a primitive, String, or Writable
        env.getConfig().registerKryoType(fieldTypeClass);
      }
    }
  }
}

/**
 * Registers the given type with the serialization stack. If the type is eventually
 * serialized as a POJO, then the type is registered with the POJO serializer. If the
 * type ends up being serialized with Kryo, then it will be registered at Kryo to make
 * sure that only tags are written.
 *
 * @param type The class of the type to register.
 */
public void registerType(Class<?> type) {
  if (type == null) {
    throw new NullPointerException("Cannot register null type class.");
  }
  TypeInformation<?> typeInfo = TypeExtractor.createTypeInfo(type);
  if (typeInfo instanceof PojoTypeInfo) {
    config.registerPojoType(type);
  } else {
    config.registerKryoType(type);
  }
}

/**
 * Registers the given type with the serialization stack. If the type is eventually
 * serialized as a POJO, then the type is registered with the POJO serializer. If the
 * type ends up being serialized with Kryo, then it will be registered at Kryo to make
 * sure that only tags are written.
 *
 * @param type The class of the type to register.
 */
public void registerType(Class<?> type) {
  if (type == null) {
    throw new NullPointerException("Cannot register null type class.");
  }
  TypeInformation<?> typeInfo = TypeExtractor.createTypeInfo(type);
  if (typeInfo instanceof PojoTypeInfo) {
    config.registerPojoType(type);
  } else {
    config.registerKryoType(type);
  }
}

/**
 * Registers the given type with the serialization stack. If the type is
 * eventually serialized as a POJO, then the type is registered with the
 * POJO serializer. If the type ends up being serialized with Kryo, then it
 * will be registered at Kryo to make sure that only tags are written.
 *
 * @param type
 *         The class of the type to register.
 */
public void registerType(Class<?> type) {
  if (type == null) {
    throw new NullPointerException("Cannot register null type class.");
  }
  TypeInformation<?> typeInfo = TypeExtractor.createTypeInfo(type);
  if (typeInfo instanceof PojoTypeInfo) {
    config.registerPojoType(type);
  } else {
    config.registerKryoType(type);
  }
}

/**
 * Registers the given type with the serialization stack. If the type is
 * eventually serialized as a POJO, then the type is registered with the
 * POJO serializer. If the type ends up being serialized with Kryo, then it
 * will be registered at Kryo to make sure that only tags are written.
 *
 * @param type
 *         The class of the type to register.
 */
public void registerType(Class<?> type) {
  if (type == null) {
    throw new NullPointerException("Cannot register null type class.");
  }
  TypeInformation<?> typeInfo = TypeExtractor.createTypeInfo(type);
  if (typeInfo instanceof PojoTypeInfo) {
    config.registerPojoType(type);
  } else {
    config.registerKryoType(type);
  }
}

/**
 * Registers the given type with the serialization stack. If the type is eventually
 * serialized as a POJO, then the type is registered with the POJO serializer. If the
 * type ends up being serialized with Kryo, then it will be registered at Kryo to make
 * sure that only tags are written.
 *
 * @param type The class of the type to register.
 */
public void registerType(Class<?> type) {
  if (type == null) {
    throw new NullPointerException("Cannot register null type class.");
  }
  TypeInformation<?> typeInfo = TypeExtractor.createTypeInfo(type);
  if (typeInfo instanceof PojoTypeInfo) {
    config.registerPojoType(type);
  } else {
    config.registerKryoType(type);
  }
}

/**
 * Configures the Flink program for this step
 */
public Configuration createInitializedConfig( FlowProcess<Configuration> flowProcess, Configuration parentConfig ) {
  this.env.getConfig().registerKryoType(Tuple.class);
  Configuration config = parentConfig == null ? new JobConf() : HadoopUtil.copyJobConf( parentConfig );
  config.set( "cascading.flow.step.num", Integer.toString( getOrdinal() ) );
  HadoopUtil.setIsInflow(config);
  this.setConfig(config);
  return config;
}