Github user tillrohrmann commented on a diff in the pull request:
https://github.com/apache/flink/pull/6276#discussion_r201017209
--- Diff:
flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java
---
@@ -2579,4 +2604,85 @@ public static RocksIteratorWrapper getRocksIterator(
ReadOptions readOptions) {
return new
RocksIteratorWrapper(db.newIterator(columnFamilyHandle, readOptions));
}
+
+ /**
+ * Encapsulates the logic and resources in connection with creating
priority queue state structures.
+ */
+ class RocksDBPriorityQueueFactory implements PriorityQueueSetFactory,
AutoCloseable {
+
+ /** Default cache size per key-group. */
+ private static final int DEFAULT_CACHES_SIZE = 8 * 1024;
+
+ /** A shared buffer to serialize elements for the priority
queue. */
+ @Nonnull
+ private final ByteArrayOutputStreamWithPos
elementSerializationOutStream;
+
+ /** A shared adapter wrapper around
elementSerializationOutStream to become a {@link DataOutputView}. */
+ @Nonnull
+ private final DataOutputViewStreamWrapper
elementSerializationOutView;
+
+ /** A shared {@link RocksDBWriteBatchWrapper} to batch
modifications to priority queues. */
+ @Nonnull
+ private final RocksDBWriteBatchWrapper writeBatchWrapper;
+
+ /** Map to track all column families created to back priority
queues. */
+ @Nonnull
+ private final Map<String, ColumnFamilyHandle>
priorityQueueColumnFamilies;
+
+ RocksDBPriorityQueueFactory() {
+ this.elementSerializationOutStream = new
ByteArrayOutputStreamWithPos();
+ this.elementSerializationOutView = new
DataOutputViewStreamWrapper(elementSerializationOutStream);
+ this.writeBatchWrapper = new
RocksDBWriteBatchWrapper(db, writeOptions);
+ this.priorityQueueColumnFamilies = new HashMap<>();
+ }
+
+ @Override
+ public <T extends HeapPriorityQueueElement>
KeyGroupedInternalPriorityQueue<T> create(
+ @Nonnull String stateName,
+ @Nonnull TypeSerializer<T> byteOrderedElementSerializer,
+ @Nonnull Comparator<T> elementComparator,
+ @Nonnull KeyExtractorFunction<T> keyExtractor) {
+
+ final ColumnFamilyHandle columnFamilyHandle =
+
priorityQueueColumnFamilies.computeIfAbsent(stateName,
RocksDBKeyedStateBackend.this::createColumnFamily);
+
+ return new KeyGroupPartitionedPriorityQueue<>(
+ keyExtractor,
+ elementComparator,
+ new
KeyGroupPartitionedPriorityQueue.PartitionQueueSetFactory<T,
CachingInternalPriorityQueueSet<T>>() {
+ @Nonnull
+ @Override
+ public
CachingInternalPriorityQueueSet<T> create(
+ int keyGroupId,
+ int numKeyGroups,
+ @Nonnull Comparator<T>
elementComparator) {
--- End diff --
Indentation
---
