Github user StefanRRichter commented on a diff in the pull request:

    https://github.com/apache/flink/pull/5518#discussion_r169931821
  
    --- 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 --
    
    You can remove the `throws Exception` here. Then you can also still use the 
method reference `targetStream.onClose(iterable::close)` in line 296.


---

Reply via email to