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

    https://github.com/apache/flink/pull/5518#discussion_r169957582
  
    --- Diff: 
flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java
 ---
    @@ -253,22 +257,61 @@ public RocksDBKeyedStateBackend(
                this.restoredKvStateMetaInfos = new HashMap<>();
                this.materializedSstFiles = new TreeMap<>();
                this.backendUID = UUID.randomUUID();
    +           this.namespaceOutputStream = new 
ByteArrayOutputStreamWithPos(8);
                LOG.debug("Setting initial keyed backend uid for operator {} to 
{}.", this.operatorIdentifier, this.backendUID);
        }
     
        @Override
    -   public <N> Stream<K> getKeys(String state, N namespace) {
    +   public <N> Stream<K> getKeys(String state, N namespace, 
TypeSerializer<N> namespaceSerializer) {
                Tuple2<ColumnFamilyHandle, ?> columnInfo = 
kvStateInformation.get(state);
                if (columnInfo == null) {
                        return Stream.empty();
                }
     
    -           RocksIterator iterator = db.newIterator(columnInfo.f0);
    -           iterator.seekToFirst();
    +           RocksIterator iterator = null;
    +           try {
    +                   iterator = db.newIterator(columnInfo.f0);
    +                   iterator.seekToFirst();
    +
    +                   boolean ambiguousKeyPossible = 
AbstractRocksDBState.AbstractRocksDBUtils.isAmbiguousKeyPossible(keySerializer, 
namespaceSerializer);
    +                   final byte[] nameSpaceBytes;
     
    -           Iterable<K> iterable = () -> new 
RocksIteratorWrapper<>(iterator, state, keySerializer, keyGroupPrefixBytes);
    -           Stream<K> targetStream = 
StreamSupport.stream(iterable.spliterator(), false);
    -           return targetStream.onClose(iterator::close);
    +                   try {
    +                           namespaceOutputStream.reset();
    +                           
AbstractRocksDBState.AbstractRocksDBUtils.writeNameSpace(
    +                                   namespace,
    +                                   namespaceSerializer,
    +                                   namespaceOutputStream,
    +                                   new 
DataOutputViewStreamWrapper(namespaceOutputStream),
    +                                   ambiguousKeyPossible);
    +                           nameSpaceBytes = 
namespaceOutputStream.toByteArray();
    +                   } catch (IOException ex) {
    +                           throw new FlinkRuntimeException("Failed to get 
keys from RocksDB state backend.", ex);
    +                   }
    +
    +                   final RocksIteratorWrapper<K> iteratorWrapper = new 
RocksIteratorWrapper<>(iterator, state, keySerializer, keyGroupPrefixBytes,
    +                           ambiguousKeyPossible, nameSpaceBytes);
    +
    +                   Stream<K> targetStream = 
StreamSupport.stream(((Iterable<K>)()->iteratorWrapper).spliterator(), false);
    --- End diff --
    
    I think `Spliterators.spliteratorUnknownSize(iteratorWrapper, 
Spliterator.ORDERED)` might be the cleaner way to get a stream from the 
iterator.


---

Reply via email to