Java Code Examples for org.apache.flink.core.memory.DataOutputSerializer#clear()

@Test
public void testKeySerializationAndDeserialization() throws Exception {
	final DataOutputSerializer outputView = new DataOutputSerializer(8);
	final DataInputDeserializer inputView = new DataInputDeserializer();

	// test for key
	for (int orgKey = 0; orgKey < 100; ++orgKey) {
		outputView.clear();
		RocksDBKeySerializationUtils.writeKey(orgKey, IntSerializer.INSTANCE, outputView, false);
		inputView.setBuffer(outputView.getCopyOfBuffer());
		int deserializedKey = RocksDBKeySerializationUtils.readKey(IntSerializer.INSTANCE, inputView, false);
		Assert.assertEquals(orgKey, deserializedKey);

		RocksDBKeySerializationUtils.writeKey(orgKey, IntSerializer.INSTANCE, outputView, true);
		inputView.setBuffer(outputView.getCopyOfBuffer());
		deserializedKey = RocksDBKeySerializationUtils.readKey(IntSerializer.INSTANCE, inputView, true);
		Assert.assertEquals(orgKey, deserializedKey);
	}
}
 

@Test
public void testNamespaceSerializationAndDeserialization() throws Exception {
	final DataOutputSerializer outputView = new DataOutputSerializer(8);
	final DataInputDeserializer inputView = new DataInputDeserializer();

	for (int orgNamespace = 0; orgNamespace < 100; ++orgNamespace) {
		outputView.clear();
		RocksDBKeySerializationUtils.writeNameSpace(orgNamespace, IntSerializer.INSTANCE, outputView, false);
		inputView.setBuffer(outputView.getCopyOfBuffer());
		int deserializedNamepsace = RocksDBKeySerializationUtils.readNamespace(IntSerializer.INSTANCE, inputView, false);
		Assert.assertEquals(orgNamespace, deserializedNamepsace);

		RocksDBKeySerializationUtils.writeNameSpace(orgNamespace, IntSerializer.INSTANCE, outputView, true);
		inputView.setBuffer(outputView.getCopyOfBuffer());
		deserializedNamepsace = RocksDBKeySerializationUtils.readNamespace(IntSerializer.INSTANCE, inputView, true);
		Assert.assertEquals(orgNamespace, deserializedNamepsace);
	}
}
 

/**
 * Converts a list of checkpoints into an array of SerializedCheckpointData.
 *
 * @param checkpoints The checkpoints to be converted into IdsCheckpointData.
 * @param serializer The serializer to serialize the IDs.
 * @param outputBuffer The reusable serialization buffer.
 * @param <T> The type of the ID.
 * @return An array of serializable SerializedCheckpointData, one per entry in the queue.
 *
 * @throws IOException Thrown, if the serialization fails.
 */
public static <T> SerializedCheckpointData[] fromDeque(
		ArrayDeque<Tuple2<Long, Set<T>>> checkpoints,
		TypeSerializer<T> serializer,
		DataOutputSerializer outputBuffer) throws IOException {

	SerializedCheckpointData[] serializedCheckpoints = new SerializedCheckpointData[checkpoints.size()];

	int pos = 0;
	for (Tuple2<Long, Set<T>> checkpoint : checkpoints) {
		outputBuffer.clear();
		Set<T> checkpointIds = checkpoint.f1;

		for (T id : checkpointIds) {
			serializer.serialize(id, outputBuffer);
		}

		serializedCheckpoints[pos++] = new SerializedCheckpointData(
				checkpoint.f0, outputBuffer.getCopyOfBuffer(), checkpointIds.size());
	}

	return serializedCheckpoints;
}
 

@Test
public void testKeySerializationAndDeserialization() throws Exception {
	final DataOutputSerializer outputView = new DataOutputSerializer(8);
	final DataInputDeserializer inputView = new DataInputDeserializer();

	// test for key
	for (int orgKey = 0; orgKey < 100; ++orgKey) {
		outputView.clear();
		RocksDBKeySerializationUtils.writeKey(orgKey, IntSerializer.INSTANCE, outputView, false);
		inputView.setBuffer(outputView.getCopyOfBuffer());
		int deserializedKey = RocksDBKeySerializationUtils.readKey(IntSerializer.INSTANCE, inputView, false);
		Assert.assertEquals(orgKey, deserializedKey);

		RocksDBKeySerializationUtils.writeKey(orgKey, IntSerializer.INSTANCE, outputView, true);
		inputView.setBuffer(outputView.getCopyOfBuffer());
		deserializedKey = RocksDBKeySerializationUtils.readKey(IntSerializer.INSTANCE, inputView, true);
		Assert.assertEquals(orgKey, deserializedKey);
	}
}
 

@Test
public void testNamespaceSerializationAndDeserialization() throws Exception {
	final DataOutputSerializer outputView = new DataOutputSerializer(8);
	final DataInputDeserializer inputView = new DataInputDeserializer();

	for (int orgNamespace = 0; orgNamespace < 100; ++orgNamespace) {
		outputView.clear();
		RocksDBKeySerializationUtils.writeNameSpace(orgNamespace, IntSerializer.INSTANCE, outputView, false);
		inputView.setBuffer(outputView.getCopyOfBuffer());
		int deserializedNamepsace = RocksDBKeySerializationUtils.readNamespace(IntSerializer.INSTANCE, inputView, false);
		Assert.assertEquals(orgNamespace, deserializedNamepsace);

		RocksDBKeySerializationUtils.writeNameSpace(orgNamespace, IntSerializer.INSTANCE, outputView, true);
		inputView.setBuffer(outputView.getCopyOfBuffer());
		deserializedNamepsace = RocksDBKeySerializationUtils.readNamespace(IntSerializer.INSTANCE, inputView, true);
		Assert.assertEquals(orgNamespace, deserializedNamepsace);
	}
}
 

/**
 * Converts a list of checkpoints into an array of SerializedCheckpointData.
 *
 * @param checkpoints The checkpoints to be converted into IdsCheckpointData.
 * @param serializer The serializer to serialize the IDs.
 * @param outputBuffer The reusable serialization buffer.
 * @param <T> The type of the ID.
 * @return An array of serializable SerializedCheckpointData, one per entry in the queue.
 *
 * @throws IOException Thrown, if the serialization fails.
 */
public static <T> SerializedCheckpointData[] fromDeque(
		ArrayDeque<Tuple2<Long, Set<T>>> checkpoints,
		TypeSerializer<T> serializer,
		DataOutputSerializer outputBuffer) throws IOException {

	SerializedCheckpointData[] serializedCheckpoints = new SerializedCheckpointData[checkpoints.size()];

	int pos = 0;
	for (Tuple2<Long, Set<T>> checkpoint : checkpoints) {
		outputBuffer.clear();
		Set<T> checkpointIds = checkpoint.f1;

		for (T id : checkpointIds) {
			serializer.serialize(id, outputBuffer);
		}

		serializedCheckpoints[pos++] = new SerializedCheckpointData(
				checkpoint.f0, outputBuffer.getCopyOfBuffer(), checkpointIds.size());
	}

	return serializedCheckpoints;
}
 

@Test
public void testKeySerializationAndDeserialization() throws Exception {
	final DataOutputSerializer outputView = new DataOutputSerializer(8);
	final DataInputDeserializer inputView = new DataInputDeserializer();

	// test for key
	for (int orgKey = 0; orgKey < 100; ++orgKey) {
		outputView.clear();
		RocksDBKeySerializationUtils.writeKey(orgKey, IntSerializer.INSTANCE, outputView, false);
		inputView.setBuffer(outputView.getCopyOfBuffer());
		int deserializedKey = RocksDBKeySerializationUtils.readKey(IntSerializer.INSTANCE, inputView, false);
		Assert.assertEquals(orgKey, deserializedKey);

		RocksDBKeySerializationUtils.writeKey(orgKey, IntSerializer.INSTANCE, outputView, true);
		inputView.setBuffer(outputView.getCopyOfBuffer());
		deserializedKey = RocksDBKeySerializationUtils.readKey(IntSerializer.INSTANCE, inputView, true);
		Assert.assertEquals(orgKey, deserializedKey);
	}
}
 

@Test
public void testNamespaceSerializationAndDeserialization() throws Exception {
	final DataOutputSerializer outputView = new DataOutputSerializer(8);
	final DataInputDeserializer inputView = new DataInputDeserializer();

	for (int orgNamespace = 0; orgNamespace < 100; ++orgNamespace) {
		outputView.clear();
		RocksDBKeySerializationUtils.writeNameSpace(orgNamespace, IntSerializer.INSTANCE, outputView, false);
		inputView.setBuffer(outputView.getCopyOfBuffer());
		int deserializedNamepsace = RocksDBKeySerializationUtils.readNamespace(IntSerializer.INSTANCE, inputView, false);
		Assert.assertEquals(orgNamespace, deserializedNamepsace);

		RocksDBKeySerializationUtils.writeNameSpace(orgNamespace, IntSerializer.INSTANCE, outputView, true);
		inputView.setBuffer(outputView.getCopyOfBuffer());
		deserializedNamepsace = RocksDBKeySerializationUtils.readNamespace(IntSerializer.INSTANCE, inputView, true);
		Assert.assertEquals(orgNamespace, deserializedNamepsace);
	}
}
 

/**
 * Converts a list of checkpoints into an array of SerializedCheckpointData.
 *
 * @param checkpoints The checkpoints to be converted into IdsCheckpointData.
 * @param serializer The serializer to serialize the IDs.
 * @param outputBuffer The reusable serialization buffer.
 * @param <T> The type of the ID.
 * @return An array of serializable SerializedCheckpointData, one per entry in the queue.
 *
 * @throws IOException Thrown, if the serialization fails.
 */
public static <T> SerializedCheckpointData[] fromDeque(
		ArrayDeque<Tuple2<Long, Set<T>>> checkpoints,
		TypeSerializer<T> serializer,
		DataOutputSerializer outputBuffer) throws IOException {

	SerializedCheckpointData[] serializedCheckpoints = new SerializedCheckpointData[checkpoints.size()];

	int pos = 0;
	for (Tuple2<Long, Set<T>> checkpoint : checkpoints) {
		outputBuffer.clear();
		Set<T> checkpointIds = checkpoint.f1;

		for (T id : checkpointIds) {
			serializer.serialize(id, outputBuffer);
		}

		serializedCheckpoints[pos++] = new SerializedCheckpointData(
				checkpoint.f0, outputBuffer.getCopyOfBuffer(), checkpointIds.size());
	}

	return serializedCheckpoints;
}
 

/**
 * Migrate only the state value, that is the "value" that is stored in RocksDB. We don't migrate
 * the key here, which is made up of key group, key, namespace and map key
 * (in case of MapState).
 */
private <N, S extends State, SV> void migrateStateValues(
	StateDescriptor<S, SV> stateDesc,
	Tuple2<ColumnFamilyHandle, RegisteredKeyValueStateBackendMetaInfo<N, SV>> stateMetaInfo) throws Exception {

	if (stateDesc.getType() == StateDescriptor.Type.MAP) {
		throw new StateMigrationException("The new serializer for a MapState requires state migration in order for the job to proceed." +
			" However, migration for MapState currently isn't supported.");
	}

	LOG.info(
		"Performing state migration for state {} because the state serializer's schema, i.e. serialization format, has changed.",
		stateDesc);

	// we need to get an actual state instance because migration is different
	// for different state types. For example, ListState needs to deal with
	// individual elements
	StateFactory stateFactory = STATE_FACTORIES.get(stateDesc.getClass());
	if (stateFactory == null) {
		String message = String.format("State %s is not supported by %s",
			stateDesc.getClass(), this.getClass());
		throw new FlinkRuntimeException(message);
	}
	State state = stateFactory.createState(
		stateDesc,
		stateMetaInfo,
		RocksDBKeyedStateBackend.this);
	if (!(state instanceof AbstractRocksDBState)) {
		throw new FlinkRuntimeException(
			"State should be an AbstractRocksDBState but is " + state);
	}

	@SuppressWarnings("unchecked")
	AbstractRocksDBState<?, ?, SV> rocksDBState = (AbstractRocksDBState<?, ?, SV>) state;

	Snapshot rocksDBSnapshot = db.getSnapshot();
	try (
		RocksIteratorWrapper iterator = RocksDBOperationUtils.getRocksIterator(db, stateMetaInfo.f0);
		RocksDBWriteBatchWrapper batchWriter = new RocksDBWriteBatchWrapper(db, getWriteOptions())
	) {
		iterator.seekToFirst();

		DataInputDeserializer serializedValueInput = new DataInputDeserializer();
		DataOutputSerializer migratedSerializedValueOutput = new DataOutputSerializer(512);
		while (iterator.isValid()) {
			serializedValueInput.setBuffer(iterator.value());

			rocksDBState.migrateSerializedValue(
				serializedValueInput,
				migratedSerializedValueOutput,
				stateMetaInfo.f1.getPreviousStateSerializer(),
				stateMetaInfo.f1.getStateSerializer());

			batchWriter.put(stateMetaInfo.f0, iterator.key(), migratedSerializedValueOutput.getCopyOfBuffer());

			migratedSerializedValueOutput.clear();
			iterator.next();
		}
	} finally {
		db.releaseSnapshot(rocksDBSnapshot);
		rocksDBSnapshot.close();
	}
}
 

/**
 * Migrate only the state value, that is the "value" that is stored in RocksDB. We don't migrate
 * the key here, which is made up of key group, key, namespace and map key
 * (in case of MapState).
 */
@SuppressWarnings("unchecked")
private <N, S extends State, SV> void migrateStateValues(
	StateDescriptor<S, SV> stateDesc,
	Tuple2<ColumnFamilyHandle, RegisteredKeyValueStateBackendMetaInfo<N, SV>> stateMetaInfo) throws Exception {

	if (stateDesc.getType() == StateDescriptor.Type.MAP) {
		TypeSerializerSnapshot<SV> previousSerializerSnapshot = stateMetaInfo.f1.getPreviousStateSerializerSnapshot();
		checkState(previousSerializerSnapshot != null, "the previous serializer snapshot should exist.");
		checkState(previousSerializerSnapshot instanceof MapSerializerSnapshot, "previous serializer snapshot should be a MapSerializerSnapshot.");

		TypeSerializer<SV> newSerializer = stateMetaInfo.f1.getStateSerializer();
		checkState(newSerializer instanceof MapSerializer, "new serializer should be a MapSerializer.");

		MapSerializer<?, ?> mapSerializer = (MapSerializer<?, ?>) newSerializer;
		MapSerializerSnapshot<?, ?> mapSerializerSnapshot = (MapSerializerSnapshot<?, ?>) previousSerializerSnapshot;
		if (!checkMapStateKeySchemaCompatibility(mapSerializerSnapshot, mapSerializer)) {
			throw new StateMigrationException(
				"The new serializer for a MapState requires state migration in order for the job to proceed, since the key schema has changed. However, migration for MapState currently only allows value schema evolutions.");
		}
	}

	LOG.info(
		"Performing state migration for state {} because the state serializer's schema, i.e. serialization format, has changed.",
		stateDesc);

	// we need to get an actual state instance because migration is different
	// for different state types. For example, ListState needs to deal with
	// individual elements
	StateFactory stateFactory = STATE_FACTORIES.get(stateDesc.getClass());
	if (stateFactory == null) {
		String message = String.format("State %s is not supported by %s",
			stateDesc.getClass(), this.getClass());
		throw new FlinkRuntimeException(message);
	}
	State state = stateFactory.createState(
		stateDesc,
		stateMetaInfo,
		RocksDBKeyedStateBackend.this);
	if (!(state instanceof AbstractRocksDBState)) {
		throw new FlinkRuntimeException(
			"State should be an AbstractRocksDBState but is " + state);
	}

	@SuppressWarnings("unchecked")
	AbstractRocksDBState<?, ?, SV> rocksDBState = (AbstractRocksDBState<?, ?, SV>) state;

	Snapshot rocksDBSnapshot = db.getSnapshot();
	try (
		RocksIteratorWrapper iterator = RocksDBOperationUtils.getRocksIterator(db, stateMetaInfo.f0);
		RocksDBWriteBatchWrapper batchWriter = new RocksDBWriteBatchWrapper(db, getWriteOptions())
	) {
		iterator.seekToFirst();

		DataInputDeserializer serializedValueInput = new DataInputDeserializer();
		DataOutputSerializer migratedSerializedValueOutput = new DataOutputSerializer(512);
		while (iterator.isValid()) {
			serializedValueInput.setBuffer(iterator.value());

			rocksDBState.migrateSerializedValue(
				serializedValueInput,
				migratedSerializedValueOutput,
				stateMetaInfo.f1.getPreviousStateSerializer(),
				stateMetaInfo.f1.getStateSerializer());

			batchWriter.put(stateMetaInfo.f0, iterator.key(), migratedSerializedValueOutput.getCopyOfBuffer());

			migratedSerializedValueOutput.clear();
			iterator.next();
		}
	} finally {
		db.releaseSnapshot(rocksDBSnapshot);
		rocksDBSnapshot.close();
	}
}
 

/**
 * Migrate only the state value, that is the "value" that is stored in RocksDB. We don't migrate
 * the key here, which is made up of key group, key, namespace and map key
 * (in case of MapState).
 */
@SuppressWarnings("unchecked")
private <N, S extends State, SV> void migrateStateValues(
	StateDescriptor<S, SV> stateDesc,
	Tuple2<ColumnFamilyHandle, RegisteredKeyValueStateBackendMetaInfo<N, SV>> stateMetaInfo) throws Exception {

	if (stateDesc.getType() == StateDescriptor.Type.MAP) {
		TypeSerializerSnapshot<SV> previousSerializerSnapshot = stateMetaInfo.f1.getPreviousStateSerializerSnapshot();
		checkState(previousSerializerSnapshot != null, "the previous serializer snapshot should exist.");
		checkState(previousSerializerSnapshot instanceof MapSerializerSnapshot, "previous serializer snapshot should be a MapSerializerSnapshot.");

		TypeSerializer<SV> newSerializer = stateMetaInfo.f1.getStateSerializer();
		checkState(newSerializer instanceof MapSerializer, "new serializer should be a MapSerializer.");

		MapSerializer<?, ?> mapSerializer = (MapSerializer<?, ?>) newSerializer;
		MapSerializerSnapshot<?, ?> mapSerializerSnapshot = (MapSerializerSnapshot<?, ?>) previousSerializerSnapshot;
		if (!checkMapStateKeySchemaCompatibility(mapSerializerSnapshot, mapSerializer)) {
			throw new StateMigrationException(
				"The new serializer for a MapState requires state migration in order for the job to proceed, since the key schema has changed. However, migration for MapState currently only allows value schema evolutions.");
		}
	}

	LOG.info(
		"Performing state migration for state {} because the state serializer's schema, i.e. serialization format, has changed.",
		stateDesc);

	// we need to get an actual state instance because migration is different
	// for different state types. For example, ListState needs to deal with
	// individual elements
	StateFactory stateFactory = STATE_FACTORIES.get(stateDesc.getClass());
	if (stateFactory == null) {
		String message = String.format("State %s is not supported by %s",
			stateDesc.getClass(), this.getClass());
		throw new FlinkRuntimeException(message);
	}
	State state = stateFactory.createState(
		stateDesc,
		stateMetaInfo,
		RocksDBKeyedStateBackend.this);
	if (!(state instanceof AbstractRocksDBState)) {
		throw new FlinkRuntimeException(
			"State should be an AbstractRocksDBState but is " + state);
	}

	@SuppressWarnings("unchecked")
	AbstractRocksDBState<?, ?, SV> rocksDBState = (AbstractRocksDBState<?, ?, SV>) state;

	Snapshot rocksDBSnapshot = db.getSnapshot();
	try (
		RocksIteratorWrapper iterator = RocksDBOperationUtils.getRocksIterator(db, stateMetaInfo.f0, readOptions);
		RocksDBWriteBatchWrapper batchWriter = new RocksDBWriteBatchWrapper(db, getWriteOptions(), getWriteBatchSize())
	) {
		iterator.seekToFirst();

		DataInputDeserializer serializedValueInput = new DataInputDeserializer();
		DataOutputSerializer migratedSerializedValueOutput = new DataOutputSerializer(512);
		while (iterator.isValid()) {
			serializedValueInput.setBuffer(iterator.value());

			rocksDBState.migrateSerializedValue(
				serializedValueInput,
				migratedSerializedValueOutput,
				stateMetaInfo.f1.getPreviousStateSerializer(),
				stateMetaInfo.f1.getStateSerializer());

			batchWriter.put(stateMetaInfo.f0, iterator.key(), migratedSerializedValueOutput.getCopyOfBuffer());

			migratedSerializedValueOutput.clear();
			iterator.next();
		}
	} finally {
		db.releaseSnapshot(rocksDBSnapshot);
		rocksDBSnapshot.close();
	}
}