Github user sihuazhou commented on a diff in the pull request:
https://github.com/apache/flink/pull/5518#discussion_r169990041
--- Diff:
flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java
---
@@ -1991,43 +2034,87 @@ public int numStateEntries() {
return count;
}
- private static class RocksIteratorWrapper<K> implements Iterator<K> {
+ /**
+ * This class is not thread safety.
+ */
+ static class RocksIteratorWrapper<K> implements Iterator<K>,
AutoCloseable {
private final RocksIterator iterator;
private final String state;
private final TypeSerializer<K> keySerializer;
private final int keyGroupPrefixBytes;
+ private final byte[] namespaceBytes;
+ private final boolean ambiguousKeyPossible;
+ private K nextKey;
public RocksIteratorWrapper(
RocksIterator iterator,
String state,
TypeSerializer<K> keySerializer,
- int keyGroupPrefixBytes) {
+ int keyGroupPrefixBytes,
+ boolean ambiguousKeyPossible,
+ byte[] namespaceBytes) {
this.iterator = Preconditions.checkNotNull(iterator);
this.state = Preconditions.checkNotNull(state);
this.keySerializer =
Preconditions.checkNotNull(keySerializer);
this.keyGroupPrefixBytes =
Preconditions.checkNotNull(keyGroupPrefixBytes);
+ this.namespaceBytes =
Preconditions.checkNotNull(namespaceBytes);
+ this.nextKey = null;
+ this.ambiguousKeyPossible = ambiguousKeyPossible;
}
@Override
public boolean hasNext() {
- return iterator.isValid();
+ final int namespaceBytesLength = namespaceBytes.length;
+ final int basicLength = namespaceBytesLength +
keyGroupPrefixBytes;
+ while (nextKey == null && iterator.isValid()) {
+ try {
+ byte[] key = iterator.key();
+ if (key.length >= basicLength) {
+ if (isMatchingNameSpace(key)) {
+
ByteArrayInputStreamWithPos inputStream =
+ new
ByteArrayInputStreamWithPos(key, keyGroupPrefixBytes, key.length -
keyGroupPrefixBytes);
+
DataInputViewStreamWrapper dataInput = new
DataInputViewStreamWrapper(inputStream);
+ K value =
AbstractRocksDBState.AbstractRocksDBUtils.readKey(
+ keySerializer,
+ inputStream,
+ dataInput,
+
ambiguousKeyPossible);
+ nextKey = value;
+ }
+ }
+ iterator.next();
+ } catch (IOException e) {
+ throw new FlinkRuntimeException("Failed
to access state [" + state + "]", e);
+ }
+ }
+ return nextKey != null;
}
@Override
public K next() {
if (!hasNext()) {
throw new NoSuchElementException("Failed to
access state [" + state + "]");
}
- try {
- byte[] key = iterator.key();
- DataInputViewStreamWrapper dataInput =
new DataInputViewStreamWrapper(
- new ByteArrayInputStreamWithPos(key,
keyGroupPrefixBytes, key.length - keyGroupPrefixBytes));
- K value = keySerializer.deserialize(dataInput);
- iterator.next();
- return value;
- } catch (IOException e) {
- throw new FlinkRuntimeException("Failed to
access state [" + state + "]", e);
+
+ K tmpKey = nextKey;
+ nextKey = null;
+ return tmpKey;
+ }
+
+ private boolean isMatchingNameSpace(byte[] key) {
+
+ final int namespaceBytesLength = namespaceBytes.length;
+ for (int i = 1; i <= namespaceBytesLength; ++i) {
+ if (key[key.length - i] !=
namespaceBytes[namespaceBytesLength - i]) {
+ return false;
+ }
}
+ return true;
+ }
+
+ @Override
+ public void close() throws Exception {
--- End diff --
ð addressing
---
