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

private TypeSerializerSchemaCompatibility<String> snapshotCompositeSerializerAndGetSchemaCompatibilityAfterRestore(
		TypeSerializer<?>[] initialNestedSerializers,
		TypeSerializer<?>[] newNestedSerializer,
		String initialOuterConfiguration,
		String newOuterConfiguration) throws IOException {
	TestCompositeTypeSerializer testSerializer =
		new TestCompositeTypeSerializer(initialOuterConfiguration, initialNestedSerializers);

	TypeSerializerSnapshot<String> testSerializerSnapshot = testSerializer.snapshotConfiguration();

	DataOutputSerializer out = new DataOutputSerializer(128);
	TypeSerializerSnapshot.writeVersionedSnapshot(out, testSerializerSnapshot);

	DataInputDeserializer in = new DataInputDeserializer(out.getCopyOfBuffer());
	testSerializerSnapshot = TypeSerializerSnapshot.readVersionedSnapshot(
		in, Thread.currentThread().getContextClassLoader());

	TestCompositeTypeSerializer newTestSerializer =
		new TestCompositeTypeSerializer(newOuterConfiguration, newNestedSerializer);
	return testSerializerSnapshot.resolveSchemaCompatibility(newTestSerializer);
}
 

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

/**
 * Serializes all values of the Iterable with the given serializer.
 *
 * @param entries         Key-value pairs to serialize
 * @param keySerializer   Serializer for UK
 * @param valueSerializer Serializer for UV
 * @param <UK>            Type of the keys
 * @param <UV>            Type of the values
 * @return Serialized values or <code>null</code> if values <code>null</code> or empty
 * @throws IOException On failure during serialization
 */
public static <UK, UV> byte[] serializeMap(Iterable<Map.Entry<UK, UV>> entries, TypeSerializer<UK> keySerializer, TypeSerializer<UV> valueSerializer) throws IOException {
	if (entries != null) {
		// Serialize
		DataOutputSerializer dos = new DataOutputSerializer(32);

		for (Map.Entry<UK, UV> entry : entries) {
			keySerializer.serialize(entry.getKey(), dos);

			if (entry.getValue() == null) {
				dos.writeBoolean(true);
			} else {
				dos.writeBoolean(false);
				valueSerializer.serialize(entry.getValue(), dos);
			}
		}

		return dos.getCopyOfBuffer();
	} else {
		return null;
	}
}
 

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

/**
 * Serializes all values of the Iterable with the given serializer.
 *
 * @param entries         Key-value pairs to serialize
 * @param keySerializer   Serializer for UK
 * @param valueSerializer Serializer for UV
 * @param <UK>            Type of the keys
 * @param <UV>            Type of the values
 * @return Serialized values or <code>null</code> if values <code>null</code> or empty
 * @throws IOException On failure during serialization
 */
public static <UK, UV> byte[] serializeMap(Iterable<Map.Entry<UK, UV>> entries, TypeSerializer<UK> keySerializer, TypeSerializer<UV> valueSerializer) throws IOException {
	if (entries != null) {
		// Serialize
		DataOutputSerializer dos = new DataOutputSerializer(32);

		for (Map.Entry<UK, UV> entry : entries) {
			keySerializer.serialize(entry.getKey(), dos);

			if (entry.getValue() == null) {
				dos.writeBoolean(true);
			} else {
				dos.writeBoolean(false);
				valueSerializer.serialize(entry.getValue(), dos);
			}
		}

		return dos.getCopyOfBuffer();
	} else {
		return null;
	}
}
 

/**
 * Serializes all values of the Iterable with the given serializer.
 *
 * @param entries         Key-value pairs to serialize
 * @param keySerializer   Serializer for UK
 * @param valueSerializer Serializer for UV
 * @param <UK>            Type of the keys
 * @param <UV>            Type of the values
 * @return Serialized values or <code>null</code> if values <code>null</code> or empty
 * @throws IOException On failure during serialization
 */
public static <UK, UV> byte[] serializeMap(Iterable<Map.Entry<UK, UV>> entries, TypeSerializer<UK> keySerializer, TypeSerializer<UV> valueSerializer) throws IOException {
	if (entries != null) {
		// Serialize
		DataOutputSerializer dos = new DataOutputSerializer(32);

		for (Map.Entry<UK, UV> entry : entries) {
			keySerializer.serialize(entry.getKey(), dos);

			if (entry.getValue() == null) {
				dos.writeBoolean(true);
			} else {
				dos.writeBoolean(false);
				valueSerializer.serialize(entry.getValue(), dos);
			}
		}

		return dos.getCopyOfBuffer();
	} else {
		return null;
	}
}
 

private TypeSerializerSchemaCompatibility<String> snapshotCompositeSerializerAndGetSchemaCompatibilityAfterRestore(
		TypeSerializer<?>[] initialNestedSerializers,
		TypeSerializer<?>[] newNestedSerializer,
		OuterSchemaCompatibility mockOuterSchemaCompatibilityResult) throws IOException {
	TestCompositeTypeSerializer testSerializer =
		new TestCompositeTypeSerializer(initialNestedSerializers);

	TypeSerializerSnapshot<String> testSerializerSnapshot = testSerializer.snapshotConfiguration();

	DataOutputSerializer out = new DataOutputSerializer(128);
	TypeSerializerSnapshot.writeVersionedSnapshot(out, testSerializerSnapshot);

	DataInputDeserializer in = new DataInputDeserializer(out.getCopyOfBuffer());
	testSerializerSnapshot = TypeSerializerSnapshot.readVersionedSnapshot(
		in, Thread.currentThread().getContextClassLoader());

	TestCompositeTypeSerializer newTestSerializer =
		new TestCompositeTypeSerializer(mockOuterSchemaCompatibilityResult, newNestedSerializer);
	return testSerializerSnapshot.resolveSchemaCompatibility(newTestSerializer);
}
 

@Override
public byte[] serialize(BucketState<BucketID> state) throws IOException {
	DataOutputSerializer out = new DataOutputSerializer(256);
	out.writeInt(MAGIC_NUMBER);
	serializeV1(state, out);
	return out.getCopyOfBuffer();
}
 

@Override
public byte[] serialize(PendingFileRecoverable pendingFileRecoverable) throws IOException {
	OutputStreamBasedPendingFileRecoverable outputStreamBasedPendingFileRecoverable = (OutputStreamBasedPendingFileRecoverable) pendingFileRecoverable;
	DataOutputSerializer dataOutputSerializer = new DataOutputSerializer(256);
	dataOutputSerializer.writeInt(MAGIC_NUMBER);
	serializeV1(outputStreamBasedPendingFileRecoverable, dataOutputSerializer);
	return dataOutputSerializer.getCopyOfBuffer();
}
 

@Test
public void testSerializeEmpty() throws IOException {
	final String testString = "beeeep!";

	SimpleVersionedSerializer<String> emptySerializer = new SimpleVersionedSerializer<String>() {

		@Override
		public int getVersion() {
			return 42;
		}

		@Override
		public byte[] serialize(String obj) throws IOException {
			return new byte[0];
		}

		@Override
		public String deserialize(int version, byte[] serialized) throws IOException {
			assertEquals(42, version);
			assertEquals(0, serialized.length);
			return testString;
		}
	};

	final DataOutputSerializer out = new DataOutputSerializer(32);
	SimpleVersionedSerialization.writeVersionAndSerialize(emptySerializer, "abc", out);
	final byte[] outBytes = out.getCopyOfBuffer();

	final byte[] bytes = SimpleVersionedSerialization.writeVersionAndSerialize(emptySerializer, "abc");
	assertArrayEquals(bytes, outBytes);

	final DataInputDeserializer in = new DataInputDeserializer(bytes);
	final String deserialized = SimpleVersionedSerialization.readVersionAndDeSerialize(emptySerializer, in);
	final String deserializedFromBytes = SimpleVersionedSerialization.readVersionAndDeSerialize(emptySerializer, outBytes);
	assertEquals(testString, deserialized);
	assertEquals(testString, deserializedFromBytes);
}
 

@Test
public void duplicatedSerializerCanDeserialize() throws IOException {
  ProtobufSerializer<SimpleMessage> serializer =
      ProtobufSerializer.forMessageGeneratedClass(SimpleMessage.class);

  DataOutputSerializer out = new DataOutputSerializer(512);
  serializer.serialize(originalMessage, out);

  DataInputDeserializer in = new DataInputDeserializer(out.getCopyOfBuffer());
  SimpleMessage foo = serializer.duplicate().deserialize(in);

  assertThat(foo, is(originalMessage));
}
 

@SuppressWarnings("unchecked")
private static <V extends Vector> void doVectorSerDeserTest(TypeSerializer ser, V vector) throws IOException {
	DataOutputSerializer out = new DataOutputSerializer(1024);
	ser.serialize(vector, out);
	DataInputDeserializer in = new DataInputDeserializer(out.getCopyOfBuffer());
	Vector deserialize = (Vector) ser.deserialize(in);
	Assert.assertEquals(vector.getClass(), deserialize.getClass());
	Assert.assertEquals(vector, deserialize);
}
 

/**
 * Serializes the value with the given serializer.
 *
 * @param value      Value of type T to serialize
 * @param serializer Serializer for T
 * @param <T>        Type of the value
 * @return Serialized value or <code>null</code> if value <code>null</code>
 * @throws IOException On failure during serialization
 */
public static <T> byte[] serializeValue(T value, TypeSerializer<T> serializer) throws IOException {
	if (value != null) {
		// Serialize
		DataOutputSerializer dos = new DataOutputSerializer(32);
		serializer.serialize(value, dos);
		return dos.getCopyOfBuffer();
	} else {
		return null;
	}
}
 

/**
 * Serializes the value with the given serializer.
 *
 * @param value      Value of type T to serialize
 * @param serializer Serializer for T
 * @param <T>        Type of the value
 * @return Serialized value or <code>null</code> if value <code>null</code>
 * @throws IOException On failure during serialization
 */
public static <T> byte[] serializeValue(T value, TypeSerializer<T> serializer) throws IOException {
	if (value != null) {
		// Serialize
		DataOutputSerializer dos = new DataOutputSerializer(32);
		serializer.serialize(value, dos);
		return dos.getCopyOfBuffer();
	} else {
		return null;
	}
}
 

@Test
public void duplicatedSerializerCanDeserialize() throws IOException {
  ProtobufSerializer<SimpleMessage> serializer =
      ProtobufSerializer.forMessageGeneratedClass(SimpleMessage.class);

  DataOutputSerializer out = new DataOutputSerializer(512);
  serializer.serialize(originalMessage, out);

  DataInputDeserializer in = new DataInputDeserializer(out.getCopyOfBuffer());
  SimpleMessage foo = serializer.duplicate().deserialize(in);

  assertThat(foo, is(originalMessage));
}
 

@Override
public byte[] serialize(BucketState<BucketID> state) throws IOException {
	DataOutputSerializer out = new DataOutputSerializer(256);
	out.writeInt(MAGIC_NUMBER);
	serializeV1(state, out);
	return out.getCopyOfBuffer();
}
 

@SuppressWarnings("unchecked")
@Override
public <N> Stream<K> getKeys(String state, N namespace) {
	RocksDbKvStateInfo columnInfo = kvStateInformation.get(state);
	if (columnInfo == null || !(columnInfo.metaInfo instanceof RegisteredKeyValueStateBackendMetaInfo)) {
		return Stream.empty();
	}

	RegisteredKeyValueStateBackendMetaInfo<N, ?> registeredKeyValueStateBackendMetaInfo =
		(RegisteredKeyValueStateBackendMetaInfo<N, ?>) columnInfo.metaInfo;

	final TypeSerializer<N> namespaceSerializer = registeredKeyValueStateBackendMetaInfo.getNamespaceSerializer();
	final DataOutputSerializer namespaceOutputView = new DataOutputSerializer(8);
	boolean ambiguousKeyPossible = RocksDBKeySerializationUtils.isAmbiguousKeyPossible(getKeySerializer(), namespaceSerializer);
	final byte[] nameSpaceBytes;
	try {
		RocksDBKeySerializationUtils.writeNameSpace(
			namespace,
			namespaceSerializer,
			namespaceOutputView,
			ambiguousKeyPossible);
		nameSpaceBytes = namespaceOutputView.getCopyOfBuffer();
	} catch (IOException ex) {
		throw new FlinkRuntimeException("Failed to get keys from RocksDB state backend.", ex);
	}

	RocksIteratorWrapper iterator = RocksDBOperationUtils.getRocksIterator(db, columnInfo.columnFamilyHandle);
	iterator.seekToFirst();

	final RocksStateKeysIterator<K> iteratorWrapper = new RocksStateKeysIterator<>(iterator, state, getKeySerializer(), keyGroupPrefixBytes,
		ambiguousKeyPossible, nameSpaceBytes);

	Stream<K> targetStream = StreamSupport.stream(Spliterators.spliteratorUnknownSize(iteratorWrapper, Spliterator.ORDERED), false);
	return targetStream.onClose(iteratorWrapper::close);
}
 

@Test
public void testSerializationRoundTrip() throws IOException {
	final SimpleVersionedSerializer<String> utfEncoder = new SimpleVersionedSerializer<String>() {

		private static final int VERSION = Integer.MAX_VALUE / 2; // version should occupy many bytes

		@Override
		public int getVersion() {
			return VERSION;
		}

		@Override
		public byte[] serialize(String str) throws IOException {
			return str.getBytes(StandardCharsets.UTF_8);
		}

		@Override
		public String deserialize(int version, byte[] serialized) throws IOException {
			assertEquals(VERSION, version);
			return new String(serialized, StandardCharsets.UTF_8);
		}
	};

	final String testString = "dugfakgs";
	final DataOutputSerializer out = new DataOutputSerializer(32);
	SimpleVersionedSerialization.writeVersionAndSerialize(utfEncoder, testString, out);
	final byte[] outBytes = out.getCopyOfBuffer();

	final byte[] bytes = SimpleVersionedSerialization.writeVersionAndSerialize(utfEncoder, testString);
	assertArrayEquals(bytes, outBytes);

	final DataInputDeserializer in = new DataInputDeserializer(bytes);
	final String deserialized = SimpleVersionedSerialization.readVersionAndDeSerialize(utfEncoder, in);
	final String deserializedFromBytes = SimpleVersionedSerialization.readVersionAndDeSerialize(utfEncoder, outBytes);
	assertEquals(testString, deserialized);
	assertEquals(testString, deserializedFromBytes);
}
 

@Test
public void testRestoreCompositeTypeSerializer() throws IOException {
	// the target compatibilities of the nested serializers doesn't matter,
	// because we're only testing the restore serializer
	TypeSerializer<?>[] testNestedSerializers = {
		new NestedSerializer(TargetCompatibility.COMPATIBLE_AS_IS),
		new NestedSerializer(TargetCompatibility.INCOMPATIBLE),
		new NestedSerializer(TargetCompatibility.COMPATIBLE_AFTER_MIGRATION)
	};

	TestCompositeTypeSerializer testSerializer = new TestCompositeTypeSerializer("outer-config", testNestedSerializers);

	TypeSerializerSnapshot<String> testSerializerSnapshot = testSerializer.snapshotConfiguration();

	DataOutputSerializer out = new DataOutputSerializer(128);
	TypeSerializerSnapshot.writeVersionedSnapshot(out, testSerializerSnapshot);

	DataInputDeserializer in = new DataInputDeserializer(out.getCopyOfBuffer());
	testSerializerSnapshot = TypeSerializerSnapshot.readVersionedSnapshot(
		in, Thread.currentThread().getContextClassLoader());

	// now, restore the composite type serializer;
	// the restored nested serializer should be a RestoredNestedSerializer
	testSerializer = (TestCompositeTypeSerializer) testSerializerSnapshot.restoreSerializer();
	Assert.assertTrue(testSerializer.getNestedSerializers()[0].getClass() == RestoredNestedSerializer.class);
	Assert.assertTrue(testSerializer.getNestedSerializers()[1].getClass() == RestoredNestedSerializer.class);
	Assert.assertTrue(testSerializer.getNestedSerializers()[2].getClass() == RestoredNestedSerializer.class);
}
 

<K> void testIteratorHelper(
	TypeSerializer<K> keySerializer,
	TypeSerializer namespaceSerializer,
	int maxKeyGroupNumber,
	Function<Integer, K> getKeyFunc) throws Exception {

	String testStateName = "aha";
	String namespace = "ns";

	String dbPath = tmp.newFolder().getAbsolutePath();
	String checkpointPath = tmp.newFolder().toURI().toString();
	RocksDBStateBackend backend = new RocksDBStateBackend(new FsStateBackend(checkpointPath), true);
	backend.setDbStoragePath(dbPath);

	Environment env = new DummyEnvironment("TestTask", 1, 0);
	RocksDBKeyedStateBackend<K> keyedStateBackend = (RocksDBKeyedStateBackend<K>) backend.createKeyedStateBackend(
		env,
		new JobID(),
		"Test",
		keySerializer,
		maxKeyGroupNumber,
		new KeyGroupRange(0, maxKeyGroupNumber - 1),
		mock(TaskKvStateRegistry.class),
		TtlTimeProvider.DEFAULT,
		new UnregisteredMetricsGroup(),
		Collections.emptyList(),
		new CloseableRegistry());

	try {
		ValueState<String> testState = keyedStateBackend.getPartitionedState(
			namespace,
			namespaceSerializer,
			new ValueStateDescriptor<String>(testStateName, String.class));

		// insert record
		for (int i = 0; i < 1000; ++i) {
			keyedStateBackend.setCurrentKey(getKeyFunc.apply(i));
			testState.update(String.valueOf(i));
		}

		DataOutputSerializer outputStream = new DataOutputSerializer(8);
		boolean ambiguousKeyPossible = RocksDBKeySerializationUtils.isAmbiguousKeyPossible(keySerializer, namespaceSerializer);
		RocksDBKeySerializationUtils.writeNameSpace(
			namespace,
			namespaceSerializer,
			outputStream,
			ambiguousKeyPossible);

		byte[] nameSpaceBytes = outputStream.getCopyOfBuffer();

		// already created with the state, should be closed with the backend
		ColumnFamilyHandle handle = keyedStateBackend.getColumnFamilyHandle(testStateName);

		try (
			RocksIteratorWrapper iterator = RocksDBOperationUtils.getRocksIterator(keyedStateBackend.db, handle);
			RocksStateKeysIterator<K> iteratorWrapper =
				new RocksStateKeysIterator<>(
					iterator,
					testStateName,
					keySerializer,
					keyedStateBackend.getKeyGroupPrefixBytes(),
					ambiguousKeyPossible,
					nameSpaceBytes)) {

			iterator.seekToFirst();

			// valid record
			List<Integer> fetchedKeys = new ArrayList<>(1000);
			while (iteratorWrapper.hasNext()) {
				fetchedKeys.add(Integer.parseInt(iteratorWrapper.next().toString()));
			}

			fetchedKeys.sort(Comparator.comparingInt(a -> a));
			Assert.assertEquals(1000, fetchedKeys.size());

			for (int i = 0; i < 1000; ++i) {
				Assert.assertEquals(i, fetchedKeys.get(i).intValue());
			}
		}
	} finally {
		if (keyedStateBackend != null) {
			keyedStateBackend.dispose();
		}
	}
}