use of org.apache.flink.core.memory.DataInputViewStreamWrapper in project flink by apache.
the class StateInitializationContextImplTest method getOperatorStateStreams.
@Test
public void getOperatorStateStreams() throws Exception {
int i = 0;
int s = 0;
for (StatePartitionStreamProvider streamProvider : initializationContext.getRawOperatorStateInputs()) {
if (0 == i % 4) {
++i;
}
Assert.assertNotNull(streamProvider);
try (InputStream is = streamProvider.getStream()) {
DataInputView div = new DataInputViewStreamWrapper(is);
int val = div.readInt();
Assert.assertEquals(i * NUM_HANDLES + s, val);
}
++s;
if (s == i % 4) {
s = 0;
++i;
}
}
}
use of org.apache.flink.core.memory.DataInputViewStreamWrapper in project flink by apache.
the class KeyedStateCheckpointOutputStreamTest method verifyRead.
private static void verifyRead(KeyGroupsStateHandle fullHandle, KeyGroupRange keyRange) throws IOException {
int count = 0;
try (FSDataInputStream in = fullHandle.openInputStream()) {
DataInputView div = new DataInputViewStreamWrapper(in);
for (int kg : fullHandle.keyGroups()) {
long off = fullHandle.getOffsetForKeyGroup(kg);
in.seek(off);
Assert.assertEquals(kg, div.readInt());
++count;
}
}
Assert.assertEquals(keyRange.getNumberOfKeyGroups(), count);
}
use of org.apache.flink.core.memory.DataInputViewStreamWrapper in project flink by apache.
the class HeapKeyedStateBackend method restorePartitionedState.
@SuppressWarnings({ "unchecked" })
private void restorePartitionedState(Collection<KeyGroupsStateHandle> state) throws Exception {
final Map<Integer, String> kvStatesById = new HashMap<>();
int numRegisteredKvStates = 0;
stateTables.clear();
for (KeyGroupsStateHandle keyGroupsHandle : state) {
if (keyGroupsHandle == null) {
continue;
}
FSDataInputStream fsDataInputStream = keyGroupsHandle.openInputStream();
cancelStreamRegistry.registerClosable(fsDataInputStream);
try {
DataInputViewStreamWrapper inView = new DataInputViewStreamWrapper(fsDataInputStream);
KeyedBackendSerializationProxy serializationProxy = new KeyedBackendSerializationProxy(userCodeClassLoader);
serializationProxy.read(inView);
List<KeyedBackendSerializationProxy.StateMetaInfo<?, ?>> metaInfoList = serializationProxy.getNamedStateSerializationProxies();
for (KeyedBackendSerializationProxy.StateMetaInfo<?, ?> metaInfoSerializationProxy : metaInfoList) {
StateTable<K, ?, ?> stateTable = stateTables.get(metaInfoSerializationProxy.getStateName());
//important: only create a new table we did not already create it previously
if (null == stateTable) {
RegisteredBackendStateMetaInfo<?, ?> registeredBackendStateMetaInfo = new RegisteredBackendStateMetaInfo<>(metaInfoSerializationProxy);
stateTable = newStateTable(registeredBackendStateMetaInfo);
stateTables.put(metaInfoSerializationProxy.getStateName(), stateTable);
kvStatesById.put(numRegisteredKvStates, metaInfoSerializationProxy.getStateName());
++numRegisteredKvStates;
}
}
for (Tuple2<Integer, Long> groupOffset : keyGroupsHandle.getGroupRangeOffsets()) {
int keyGroupIndex = groupOffset.f0;
long offset = groupOffset.f1;
fsDataInputStream.seek(offset);
int writtenKeyGroupIndex = inView.readInt();
Preconditions.checkState(writtenKeyGroupIndex == keyGroupIndex, "Unexpected key-group in restore.");
for (int i = 0; i < metaInfoList.size(); i++) {
int kvStateId = inView.readShort();
StateTable<K, ?, ?> stateTable = stateTables.get(kvStatesById.get(kvStateId));
StateTableByKeyGroupReader keyGroupReader = StateTableByKeyGroupReaders.readerForVersion(stateTable, serializationProxy.getRestoredVersion());
keyGroupReader.readMappingsInKeyGroup(inView, keyGroupIndex);
}
}
} finally {
cancelStreamRegistry.unregisterClosable(fsDataInputStream);
IOUtils.closeQuietly(fsDataInputStream);
}
}
}
use of org.apache.flink.core.memory.DataInputViewStreamWrapper in project flink by apache.
the class EventWithAggregatorsTest method pipeThroughSerialization.
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;
}
}
Aggregations