org.apache.flink.core.memory.DataOutputViewStreamWrapper Java Examples

@Override
public void setOutputType(TypeInformation<R> outTypeInfo, ExecutionConfig executionConfig) {
	// out type is not used, just use this for the execution config
	accSerializer = accTypeInformation.createSerializer(executionConfig);

	ByteArrayOutputStream baos = new ByteArrayOutputStream();
	DataOutputViewStreamWrapper out = new DataOutputViewStreamWrapper(baos);

	try {
		accSerializer.serialize(initialValue, out);
	} catch (IOException ioe) {
		throw new RuntimeException("Unable to serialize initial value of type " +
			initialValue.getClass().getSimpleName() + " of fold window function.", ioe);
	}

	serializedInitialValue = baos.toByteArray();
}
 

@Test
public void testKeyedStateMetaInfoSerialization() throws Exception {

	String name = "test";
	TypeSerializer<?> namespaceSerializer = LongSerializer.INSTANCE;
	TypeSerializer<?> stateSerializer = DoubleSerializer.INSTANCE;

	StateMetaInfoSnapshot metaInfo = new RegisteredKeyValueStateBackendMetaInfo<>(
		StateDescriptor.Type.VALUE, name, namespaceSerializer, stateSerializer).snapshot();

	byte[] serialized;
	try (ByteArrayOutputStreamWithPos out = new ByteArrayOutputStreamWithPos()) {
		StateMetaInfoSnapshotReadersWriters.getWriter().
			writeStateMetaInfoSnapshot(metaInfo, new DataOutputViewStreamWrapper(out));
		serialized = out.toByteArray();
	}

	try (ByteArrayInputStreamWithPos in = new ByteArrayInputStreamWithPos(serialized)) {
		final StateMetaInfoReader reader = StateMetaInfoSnapshotReadersWriters.getReader(
			CURRENT_STATE_META_INFO_SNAPSHOT_VERSION, StateMetaInfoSnapshotReadersWriters.StateTypeHint.KEYED_STATE);
		metaInfo = reader.readStateMetaInfoSnapshot(
			new DataInputViewStreamWrapper(in), Thread.currentThread().getContextClassLoader());
	}

	Assert.assertEquals(name, metaInfo.getName());
}
 

/**
 * Write a list of serializers and their corresponding config snapshots to the provided
 * data output view. This method writes in a fault tolerant way, so that when read again
 * using {@link #readSerializersAndConfigsWithResilience(DataInputView, ClassLoader)}, if
 * deserialization of the serializer fails, its configuration snapshot will remain intact.
 *
 * <p>Specifically, all written serializers and their config snapshots are indexed by their
 * offset positions within the serialized bytes. The serialization format is as follows:
 * <ul>
 *     <li>1. number of serializer and configuration snapshot pairs.</li>
 *     <li>2. offsets of each serializer and configuration snapshot, in order.</li>
 *     <li>3. total number of bytes for the serialized serializers and the config snapshots.</li>
 *     <li>4. serialized serializers and the config snapshots.</li>
 * </ul>
 *
 * @param out the data output view.
 * @param serializersAndConfigs serializer and configuration snapshot pairs
 *
 * @throws IOException
 */
public static void writeSerializersAndConfigsWithResilience(
		DataOutputView out,
		List<Tuple2<TypeSerializer<?>, TypeSerializerSnapshot<?>>> serializersAndConfigs) throws IOException {

	try (
		ByteArrayOutputStreamWithPos bufferWithPos = new ByteArrayOutputStreamWithPos();
		DataOutputViewStreamWrapper bufferWrapper = new DataOutputViewStreamWrapper(bufferWithPos)) {

		out.writeInt(serializersAndConfigs.size());
		for (Tuple2<TypeSerializer<?>, TypeSerializerSnapshot<?>> serAndConfSnapshot : serializersAndConfigs) {
			out.writeInt(bufferWithPos.getPosition());
			writeSerializer(bufferWrapper, serAndConfSnapshot.f0);

			out.writeInt(bufferWithPos.getPosition());
			TypeSerializerSnapshotSerializationUtil.writeSerializerSnapshot(
				bufferWrapper, (TypeSerializerSnapshot) serAndConfSnapshot.f1, serAndConfSnapshot.f0);
		}

		out.writeInt(bufferWithPos.getPosition());
		out.write(bufferWithPos.getBuf(), 0, bufferWithPos.getPosition());
	}
}
 

public FailingCollectionSource(
	TypeSerializer<T> serializer,
	Iterable<T> elements,
	int failureAfterNumElements) throws IOException {
	ByteArrayOutputStream baos = new ByteArrayOutputStream();
	DataOutputViewStreamWrapper wrapper = new DataOutputViewStreamWrapper(baos);

	int count = 0;
	try {
		for (T element : elements) {
			serializer.serialize(element, wrapper);
			count++;
		}
	}
	catch (Exception e) {
		throw new IOException("Serializing the source elements failed: " + e.getMessage(), e);
	}

	this.serializer = serializer;
	this.elementsSerialized = baos.toByteArray();
	this.numElements = count;
	checkArgument(failureAfterNumElements > 0);
	this.failureAfterNumElements = failureAfterNumElements;
	this.checkpointedEmittedNums = new HashMap<>();
}
 

@Test
public void testReadMismatchVersion() throws Exception {

	String payload = "test";

	TestWriteable testWriteable = new TestWriteable(1, payload);
	byte[] serialized;
	try (ByteArrayOutputStreamWithPos out = new ByteArrayOutputStreamWithPos()) {
		testWriteable.write(new DataOutputViewStreamWrapper(out));
		serialized = out.toByteArray();
	}

	testWriteable = new TestWriteable(2);
	try (ByteArrayInputStreamWithPos in = new ByteArrayInputStreamWithPos(serialized)) {
		testWriteable.read(new DataInputViewStreamWrapper(in));
		Assert.fail("Version mismatch expected.");
	} catch (VersionMismatchException ignored) {

	}

	Assert.assertEquals(null, testWriteable.getData());
}
 

public FailingCollectionSource(
	TypeSerializer<T> serializer,
	Iterable<T> elements,
	int failureAfterNumElements) throws IOException {
	ByteArrayOutputStream baos = new ByteArrayOutputStream();
	DataOutputViewStreamWrapper wrapper = new DataOutputViewStreamWrapper(baos);

	int count = 0;
	try {
		for (T element : elements) {
			serializer.serialize(element, wrapper);
			count++;
		}
	}
	catch (Exception e) {
		throw new IOException("Serializing the source elements failed: " + e.getMessage(), e);
	}

	this.serializer = serializer;
	this.elementsSerialized = baos.toByteArray();
	this.numElements = count;
	checkArgument(failureAfterNumElements > 0);
	this.failureAfterNumElements = failureAfterNumElements;
	this.checkpointedEmittedNums = new HashMap<>();
}
 

/**
 * Clones the given writable using the {@link IOReadableWritable serialization}.
 *
 * @param original Object to clone
 * @param <T> Type of the object to clone
 * @return Cloned object
 * @throws IOException Thrown is the serialization fails.
 */
public static <T extends IOReadableWritable> T createCopyWritable(T original) throws IOException {
	if (original == null) {
		return null;
	}

	final ByteArrayOutputStream baos = new ByteArrayOutputStream();
	try (DataOutputViewStreamWrapper out = new DataOutputViewStreamWrapper(baos)) {
		original.write(out);
	}

	final ByteArrayInputStream bais = new ByteArrayInputStream(baos.toByteArray());
	try (DataInputViewStreamWrapper in = new DataInputViewStreamWrapper(bais)) {

		@SuppressWarnings("unchecked")
		T copy = (T) instantiate(original.getClass());
		copy.read(in);
		return copy;
	}
}
 

@Override
public void setOutputType(TypeInformation<R> outTypeInfo, ExecutionConfig executionConfig) {
	// out type is not used, just use this for the execution config
	accSerializer = accTypeInformation.createSerializer(executionConfig);

	ByteArrayOutputStream baos = new ByteArrayOutputStream();
	DataOutputViewStreamWrapper out = new DataOutputViewStreamWrapper(baos);

	try {
		accSerializer.serialize(initialValue, out);
	} catch (IOException ioe) {
		throw new RuntimeException("Unable to serialize initial value of type " +
			initialValue.getClass().getSimpleName() + " of fold window function.", ioe);
	}

	serializedInitialValue = baos.toByteArray();
}
 

/**
 * Verifies that reading and writing serializers work correctly.
 */
@Test
public void testSerializerSerialization() throws Exception {

	TypeSerializer<?> serializer = IntSerializer.INSTANCE;

	byte[] serialized;
	try (ByteArrayOutputStreamWithPos out = new ByteArrayOutputStreamWithPos()) {
		TypeSerializerSerializationUtil.writeSerializer(new DataOutputViewStreamWrapper(out), serializer);
		serialized = out.toByteArray();
	}

	TypeSerializer<?> deserializedSerializer;
	try (ByteArrayInputStreamWithPos in = new ByteArrayInputStreamWithPos(serialized)) {
		deserializedSerializer = TypeSerializerSerializationUtil.tryReadSerializer(
			new DataInputViewStreamWrapper(in), Thread.currentThread().getContextClassLoader());
	}

	Assert.assertEquals(serializer, deserializedSerializer);
}
 

/**
 * Copied from HeapListState.getSerializedValue(Object, Object).
 */
private byte[] serializeInitValue(List<Long> toSerialize) throws IOException {
	TypeSerializer<Long> serializer = listStateDesc.getElementSerializer();

	ByteArrayOutputStream baos = new ByteArrayOutputStream();
	DataOutputViewStreamWrapper view = new DataOutputViewStreamWrapper(baos);

	// write the same as RocksDB writes lists, with one ',' separator
	for (int i = 0; i < toSerialize.size(); i++) {
		serializer.serialize(toSerialize.get(i), view);
		if (i < toSerialize.size() - 1) {
			view.writeByte(',');
		}
	}
	view.flush();

	return baos.toByteArray();
}
 

public FromElementsFunction(TypeSerializer<T> serializer, Iterable<T> elements) throws IOException {
	ByteArrayOutputStream baos = new ByteArrayOutputStream();
	DataOutputViewStreamWrapper wrapper = new DataOutputViewStreamWrapper(baos);

	int count = 0;
	try {
		for (T element : elements) {
			serializer.serialize(element, wrapper);
			count++;
		}
	}
	catch (Exception e) {
		throw new IOException("Serializing the source elements failed: " + e.getMessage(), e);
	}

	this.serializer = serializer;
	this.elementsSerialized = baos.toByteArray();
	this.numElements = count;
}
 

/**
 * Verifies that deserializing config snapshots fail if the config class could not be found.
 */
@Test(expected = IOException.class)
public void testFailsWhenConfigurationSnapshotClassNotFound() throws Exception {
	byte[] serializedConfig;
	try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
		TypeSerializerSnapshotSerializationUtil.writeSerializerSnapshot(
			new DataOutputViewStreamWrapper(out),
			new TypeSerializerSerializationUtilTest.TestConfigSnapshot<>(123, "foobar"),
			StringSerializer.INSTANCE);
		serializedConfig = out.toByteArray();
	}

	try (ByteArrayInputStream in = new ByteArrayInputStream(serializedConfig)) {
		// read using a dummy classloader
		TypeSerializerSnapshotSerializationUtil.readSerializerSnapshot(
			new DataInputViewStreamWrapper(in), new URLClassLoader(new URL[0], null), null);
	}

	fail("Expected a ClassNotFoundException wrapped in IOException");
}
 

@Override
public void setOutputType(TypeInformation<R> outTypeInfo, ExecutionConfig executionConfig) {
	accSerializer = accTypeInformation.createSerializer(executionConfig);

	ByteArrayOutputStream baos = new ByteArrayOutputStream();
	DataOutputViewStreamWrapper out = new DataOutputViewStreamWrapper(baos);

	try {
		accSerializer.serialize(initialValue, out);
	} catch (IOException ioe) {
		throw new RuntimeException("Unable to serialize initial value of type " +
			initialValue.getClass().getSimpleName() + " of fold window function.", ioe);
	}

	serializedInitialValue = baos.toByteArray();
}
 

@Test
public void testKeyGroupSerializationAndDeserialization() throws Exception {
	ByteArrayOutputStreamWithPos outputStream = new ByteArrayOutputStreamWithPos(8);
	DataOutputView outputView = new DataOutputViewStreamWrapper(outputStream);

	for (int keyGroupPrefixBytes = 1; keyGroupPrefixBytes <= 2; ++keyGroupPrefixBytes) {
		for (int orgKeyGroup = 0; orgKeyGroup < 128; ++orgKeyGroup) {
			outputStream.reset();
			RocksDBKeySerializationUtils.writeKeyGroup(orgKeyGroup, keyGroupPrefixBytes, outputView);
			int deserializedKeyGroup = RocksDBKeySerializationUtils.readKeyGroup(
				keyGroupPrefixBytes,
				new DataInputViewStreamWrapper(new ByteArrayInputStreamWithPos(outputStream.toByteArray())));
			Assert.assertEquals(orgKeyGroup, deserializedKeyGroup);
		}
	}
}
 

@Override
public int compare(TestElement o1, TestElement o2) {

	ByteArrayOutputStreamWithPos os = new ByteArrayOutputStreamWithPos();
	DataOutputViewStreamWrapper ow = new DataOutputViewStreamWrapper(os);
	try {
		TestElementSerializer.INSTANCE.serialize(o1, ow);
		byte[] a1 = os.toByteArray();
		os.reset();
		TestElementSerializer.INSTANCE.serialize(o2, ow);
		byte[] a2 = os.toByteArray();
		return UnsignedBytes.lexicographicalComparator().compare(a1, a2);
	} catch (Exception e) {
		throw new RuntimeException(e);
	}
}
 

/**
 * Verifies deserialization failure cases when reading a serializer from bytes, in the
 * case of a {@link InvalidClassException}.
 */
@Test
public void testSerializerSerializationWithInvalidClass() throws Exception {

	TypeSerializer<?> serializer = IntSerializer.INSTANCE;

	byte[] serialized;
	try (ByteArrayOutputStreamWithPos out = new ByteArrayOutputStreamWithPos()) {
		TypeSerializerSerializationUtil.writeSerializer(new DataOutputViewStreamWrapper(out), serializer);
		serialized = out.toByteArray();
	}

	TypeSerializer<?> deserializedSerializer;

	try (ByteArrayInputStreamWithPos in = new ByteArrayInputStreamWithPos(serialized)) {
		deserializedSerializer = TypeSerializerSerializationUtil.tryReadSerializer(
			new DataInputViewStreamWrapper(in),
			new ArtificialCNFExceptionThrowingClassLoader(
				Thread.currentThread().getContextClassLoader(),
				Collections.singleton(IntSerializer.class.getName())),
			true);
	}
	Assert.assertTrue(deserializedSerializer instanceof UnloadableDummyTypeSerializer);
}
 

@Test
public void testReadSameVersion() throws Exception {

	String payload = "test";

	TestWriteable testWriteable = new TestWriteable(1, payload);
	byte[] serialized;
	try (ByteArrayOutputStreamWithPos out = new ByteArrayOutputStreamWithPos()) {
		testWriteable.write(new DataOutputViewStreamWrapper(out));
		serialized = out.toByteArray();
	}

	testWriteable = new TestWriteable(1);
	try (ByteArrayInputStreamWithPos in = new ByteArrayInputStreamWithPos(serialized)) {
		testWriteable.read(new DataInputViewStreamWrapper(in));
	}

	Assert.assertEquals(payload, testWriteable.getData());
}
 

@Override
public void setOutputType(TypeInformation<R> outTypeInfo, ExecutionConfig executionConfig) {
	accSerializer = accTypeInformation.createSerializer(executionConfig);

	ByteArrayOutputStream baos = new ByteArrayOutputStream();
	DataOutputViewStreamWrapper out = new DataOutputViewStreamWrapper(baos);

	try {
		accSerializer.serialize(initialValue, out);
	} catch (IOException ioe) {
		throw new RuntimeException("Unable to serialize initial value of type " +
			initialValue.getClass().getSimpleName() + " of fold window function.", ioe);
	}

	serializedInitialValue = baos.toByteArray();
}
 

@Test
public void testReadMismatchVersion() throws Exception {

	String payload = "test";

	TestWriteable testWriteable = new TestWriteable(1, payload);
	byte[] serialized;
	try (ByteArrayOutputStreamWithPos out = new ByteArrayOutputStreamWithPos()) {
		testWriteable.write(new DataOutputViewStreamWrapper(out));
		serialized = out.toByteArray();
	}

	testWriteable = new TestWriteable(2);
	try (ByteArrayInputStreamWithPos in = new ByteArrayInputStreamWithPos(serialized)) {
		testWriteable.read(new DataInputViewStreamWrapper(in));
		Assert.fail("Version mismatch expected.");
	} catch (VersionMismatchException ignored) {

	}

	Assert.assertEquals(null, testWriteable.getData());
}
 

/**
 * Verifies that deserializing config snapshots fail if the config class could not be found.
 */
@Test(expected = IOException.class)
public void testFailsWhenConfigurationSnapshotClassNotFound() throws Exception {
	byte[] serializedConfig;
	try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
		TypeSerializerSnapshotSerializationUtil.writeSerializerSnapshot(
			new DataOutputViewStreamWrapper(out),
			new TypeSerializerSerializationUtilTest.TestConfigSnapshot<>(123, "foobar"),
			StringSerializer.INSTANCE);
		serializedConfig = out.toByteArray();
	}

	try (ByteArrayInputStream in = new ByteArrayInputStream(serializedConfig)) {
		// read using a dummy classloader
		TypeSerializerSnapshotSerializationUtil.readSerializerSnapshot(
			new DataInputViewStreamWrapper(in), new URLClassLoader(new URL[0], null), null);
	}

	fail("Expected a ClassNotFoundException wrapped in IOException");
}
 

/**
 * Clones the given writable using the {@link IOReadableWritable serialization}.
 *
 * @param original Object to clone
 * @param <T> Type of the object to clone
 * @return Cloned object
 * @throws IOException Thrown is the serialization fails.
 */
public static <T extends IOReadableWritable> T createCopyWritable(T original) throws IOException {
	if (original == null) {
		return null;
	}

	final ByteArrayOutputStream baos = new ByteArrayOutputStream();
	try (DataOutputViewStreamWrapper out = new DataOutputViewStreamWrapper(baos)) {
		original.write(out);
	}

	final ByteArrayInputStream bais = new ByteArrayInputStream(baos.toByteArray());
	try (DataInputViewStreamWrapper in = new DataInputViewStreamWrapper(bais)) {

		@SuppressWarnings("unchecked")
		T copy = (T) instantiate(original.getClass());
		copy.read(in);
		return copy;
	}
}
 

@Override
public void setOutputType(TypeInformation<OUT> outTypeInfo, ExecutionConfig executionConfig) {
	outTypeSerializer = outTypeInfo.createSerializer(executionConfig);

	ByteArrayOutputStream baos = new ByteArrayOutputStream();
	DataOutputViewStreamWrapper out = new DataOutputViewStreamWrapper(baos);

	try {
		outTypeSerializer.serialize(initialValue, out);
	} catch (IOException ioe) {
		throw new RuntimeException("Unable to serialize initial value of type " +
				initialValue.getClass().getSimpleName() + " of fold operator.", ioe);
	}

	serializedInitialValue = baos.toByteArray();
}
 

/**
 * Copied from HeapListState.getSerializedValue(Object, Object).
 */
private byte[] serializeInitValue(List<Long> toSerialize) throws IOException {
	TypeSerializer<Long> serializer = listStateDesc.getElementSerializer();

	ByteArrayOutputStream baos = new ByteArrayOutputStream();
	DataOutputViewStreamWrapper view = new DataOutputViewStreamWrapper(baos);

	// write the same as RocksDB writes lists, with one ',' separator
	for (int i = 0; i < toSerialize.size(); i++) {
		serializer.serialize(toSerialize.get(i), view);
		if (i < toSerialize.size() - 1) {
			view.writeByte(',');
		}
	}
	view.flush();

	return baos.toByteArray();
}
 

@Test
public void testKeyGroupSerializationAndDeserialization() throws Exception {
	ByteArrayOutputStreamWithPos outputStream = new ByteArrayOutputStreamWithPos(8);
	DataOutputView outputView = new DataOutputViewStreamWrapper(outputStream);

	for (int keyGroupPrefixBytes = 1; keyGroupPrefixBytes <= 2; ++keyGroupPrefixBytes) {
		for (int orgKeyGroup = 0; orgKeyGroup < 128; ++orgKeyGroup) {
			outputStream.reset();
			RocksDBKeySerializationUtils.writeKeyGroup(orgKeyGroup, keyGroupPrefixBytes, outputView);
			int deserializedKeyGroup = RocksDBKeySerializationUtils.readKeyGroup(
				keyGroupPrefixBytes,
				new DataInputViewStreamWrapper(new ByteArrayInputStreamWithPos(outputStream.toByteArray())));
			Assert.assertEquals(orgKeyGroup, deserializedKeyGroup);
		}
	}
}
 

/**
 * This tests that resource can be released for a successful snapshot.
 */
@Test
public void testReleaseForSuccessfulSnapshot() throws IOException {
	int numberOfKeyGroups = 10;
	CopyOnWriteStateTable<Integer, Integer, Float> table = createStateTableForSnapshotRelease(numberOfKeyGroups);

	ByteArrayOutputStreamWithPos byteArrayOutputStreamWithPos = new ByteArrayOutputStreamWithPos();
	DataOutputView dataOutputView = new DataOutputViewStreamWrapper(byteArrayOutputStreamWithPos);

	CopyOnWriteStateTableSnapshot<Integer, Integer, Float> snapshot = table.stateSnapshot();
	for (int group = 0; group < numberOfKeyGroups; group++) {
		snapshot.writeStateInKeyGroup(dataOutputView, group);
		// resource used by one key group should be released after the snapshot is successful
		Assert.assertTrue(isResourceReleasedForKeyGroup(table, group));
	}
	snapshot.release();
	verifyResourceIsReleasedForAllKeyGroup(table, 1);
}
 

/**
 * Verifies that deserializing config snapshots fail if the config class could not be found.
 */
@Test(expected = IOException.class)
public void testFailsWhenConfigurationSnapshotClassNotFound() throws Exception {
	byte[] serializedConfig;
	try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
		TypeSerializerSnapshotSerializationUtil.writeSerializerSnapshot(
			new DataOutputViewStreamWrapper(out),
			new TypeSerializerSerializationUtilTest.TestConfigSnapshot<>(123, "foobar"),
			StringSerializer.INSTANCE);
		serializedConfig = out.toByteArray();
	}

	try (ByteArrayInputStream in = new ByteArrayInputStream(serializedConfig)) {
		// read using a dummy classloader
		TypeSerializerSnapshotSerializationUtil.readSerializerSnapshot(
			new DataInputViewStreamWrapper(in), new URLClassLoader(new URL[0], null), null);
	}

	fail("Expected a ClassNotFoundException wrapped in IOException");
}
 

private void writeSnapshotToOutputStream(
	@Nonnull CheckpointStreamWithResultProvider checkpointStreamWithResultProvider,
	@Nonnull KeyGroupRangeOffsets keyGroupRangeOffsets) throws IOException, InterruptedException {

	final List<Tuple2<RocksIteratorWrapper, Integer>> kvStateIterators =
		new ArrayList<>(metaData.size());
	final DataOutputView outputView =
		new DataOutputViewStreamWrapper(checkpointStreamWithResultProvider.getCheckpointOutputStream());
	final ReadOptions readOptions = new ReadOptions();
	try {
		readOptions.setSnapshot(snapshot);
		writeKVStateMetaData(kvStateIterators, readOptions, outputView);
		writeKVStateData(kvStateIterators, checkpointStreamWithResultProvider, keyGroupRangeOffsets);
	} finally {

		for (Tuple2<RocksIteratorWrapper, Integer> kvStateIterator : kvStateIterators) {
			IOUtils.closeQuietly(kvStateIterator.f0);
		}

		IOUtils.closeQuietly(readOptions);
	}
}
 

private IterationEventWithAggregators pipeThroughSerialization(IterationEventWithAggregators event) {
	try {
		ByteArrayOutputStream baos = new ByteArrayOutputStream();
		event.write(new DataOutputViewStreamWrapper(baos));

		byte[] data = baos.toByteArray();
		baos.close();

		DataInputViewStreamWrapper in = new DataInputViewStreamWrapper(new ByteArrayInputStream(data));
		IterationEventWithAggregators newEvent = event.getClass().newInstance();
		newEvent.read(in);
		in.close();

		return newEvent;
	}
	catch (Exception e) {
		System.err.println(e.getMessage());
		e.printStackTrace();
		Assert.fail("Test threw an exception: " + e.getMessage());
		return null;
	}
}
 

private IterationEventWithAggregators pipeThroughSerialization(IterationEventWithAggregators event) {
	try {
		ByteArrayOutputStream baos = new ByteArrayOutputStream();
		event.write(new DataOutputViewStreamWrapper(baos));

		byte[] data = baos.toByteArray();
		baos.close();

		DataInputViewStreamWrapper in = new DataInputViewStreamWrapper(new ByteArrayInputStream(data));
		IterationEventWithAggregators newEvent = event.getClass().newInstance();
		newEvent.read(in);
		in.close();

		return newEvent;
	}
	catch (Exception e) {
		System.err.println(e.getMessage());
		e.printStackTrace();
		Assert.fail("Test threw an exception: " + e.getMessage());
		return null;
	}
}
 

@Test
public void testReadWriteMissingKeyGroups() throws Exception {
	final KeyGroupRange keyRange = new KeyGroupRange(0, 2);
	KeyedStateCheckpointOutputStream stream = createStream(keyRange);

	DataOutputView dov = new DataOutputViewStreamWrapper(stream);
	stream.startNewKeyGroup(1);
	dov.writeInt(1);

	KeyGroupsStateHandle fullHandle = stream.closeAndGetHandle();

	int count = 0;
	try (FSDataInputStream in = fullHandle.openInputStream()) {
		DataInputView div = new DataInputViewStreamWrapper(in);
		for (int kg : fullHandle.getKeyGroupRange()) {
			long off = fullHandle.getOffsetForKeyGroup(kg);
			if (off >= 0) {
				in.seek(off);
				Assert.assertEquals(1, div.readInt());
				++count;
			}
		}
	}

	Assert.assertEquals(1, count);
}