Examples with RocksStateKeysAndNamespaceIterator org.apache.flink.contrib.streaming.state.iterator.RocksStateKeysAndNamespaceIterator used on opensource projects

Example 1 with RocksStateKeysAndNamespaceIterator

use of org.apache.flink.contrib.streaming.state.iterator.RocksStateKeysAndNamespaceIterator in project flink by apache.

the class RocksDBRocksStateKeysAndNamespacesIteratorTest method testIteratorHelper.

@SuppressWarnings("unchecked")
<K> void testIteratorHelper(TypeSerializer<K> keySerializer, int maxKeyGroupNumber, Function<Integer, K> getKeyFunc) throws Exception {
    String testStateName = "aha";
    String namespace = "ns";
    try (RocksDBKeyedStateBackendTestFactory factory = new RocksDBKeyedStateBackendTestFactory()) {
        RocksDBKeyedStateBackend<K> keyedStateBackend = factory.create(tmp, keySerializer, maxKeyGroupNumber);
        ValueState<String> testState = keyedStateBackend.getPartitionedState(namespace, StringSerializer.INSTANCE, new ValueStateDescriptor<>(testStateName, String.class));
        // insert record
        for (int i = 0; i < 1000; ++i) {
            keyedStateBackend.setCurrentKey(getKeyFunc.apply(i));
            testState.update(String.valueOf(i));
        }
        DataOutputSerializer outputStream = new DataOutputSerializer(8);
        boolean ambiguousKeyPossible = CompositeKeySerializationUtils.isAmbiguousKeyPossible(keySerializer, StringSerializer.INSTANCE);
        CompositeKeySerializationUtils.writeNameSpace(namespace, StringSerializer.INSTANCE, outputStream, ambiguousKeyPossible);
        // already created with the state, should be closed with the backend
        ColumnFamilyHandle handle = keyedStateBackend.getColumnFamilyHandle(testStateName);
        try (RocksIteratorWrapper iterator = RocksDBOperationUtils.getRocksIterator(keyedStateBackend.db, handle, keyedStateBackend.getReadOptions());
            RocksStateKeysAndNamespaceIterator<K, String> iteratorWrapper = new RocksStateKeysAndNamespaceIterator<>(iterator, testStateName, keySerializer, StringSerializer.INSTANCE, keyedStateBackend.getKeyGroupPrefixBytes(), ambiguousKeyPossible)) {
            iterator.seekToFirst();
            // valid record
            List<Tuple2<Integer, String>> fetchedKeys = new ArrayList<>(1000);
            while (iteratorWrapper.hasNext()) {
                Tuple2 entry = iteratorWrapper.next();
                entry.f0 = Integer.parseInt(entry.f0.toString());
                fetchedKeys.add((Tuple2<Integer, String>) entry);
            }
            fetchedKeys.sort(Comparator.comparingInt(a -> a.f0));
            Assert.assertEquals(1000, fetchedKeys.size());
            for (int i = 0; i < 1000; ++i) {
                Assert.assertEquals(i, fetchedKeys.get(i).f0.intValue());
                Assert.assertEquals(namespace, fetchedKeys.get(i).f1);
            }
        }
    }
}