Examples with ColumnFamilyHandle org.rocksdb.ColumnFamilyHandle used on opensource projects

Example 46 with ColumnFamilyHandle

use of org.rocksdb.ColumnFamilyHandle in project flink by apache.

the class RocksDBPriorityQueueSetFactory method create.

@Nonnull
@Override
public <T extends HeapPriorityQueueElement & PriorityComparable<? super T> & Keyed<?>> KeyGroupedInternalPriorityQueue<T> create(@Nonnull String stateName, @Nonnull TypeSerializer<T> byteOrderedElementSerializer) {
    final RocksDBKeyedStateBackend.RocksDbKvStateInfo stateCFHandle = tryRegisterPriorityQueueMetaInfo(stateName, byteOrderedElementSerializer);
    final ColumnFamilyHandle columnFamilyHandle = stateCFHandle.columnFamilyHandle;
    return new KeyGroupPartitionedPriorityQueue<>(KeyExtractorFunction.forKeyedObjects(), PriorityComparator.forPriorityComparableObjects(), new KeyGroupPartitionedPriorityQueue.PartitionQueueSetFactory<T, RocksDBCachingPriorityQueueSet<T>>() {

        @Nonnull
        @Override
        public RocksDBCachingPriorityQueueSet<T> create(int keyGroupId, int numKeyGroups, @Nonnull KeyExtractorFunction<T> keyExtractor, @Nonnull PriorityComparator<T> elementPriorityComparator) {
            TreeOrderedSetCache orderedSetCache = new TreeOrderedSetCache(DEFAULT_CACHES_SIZE);
            return new RocksDBCachingPriorityQueueSet<>(keyGroupId, keyGroupPrefixBytes, db, readOptions, columnFamilyHandle, byteOrderedElementSerializer, sharedElementOutView, sharedElementInView, writeBatchWrapper, orderedSetCache);
        }
    }, keyGroupRange, numberOfKeyGroups);
}

Example 47 with ColumnFamilyHandle

use of org.rocksdb.ColumnFamilyHandle in project flink by apache.

the class RocksDBFullRestoreOperation method restoreKVStateData.

/**
 * Restore the KV-state / ColumnFamily data for all key-groups referenced by the current state
 * handle.
 */
private void restoreKVStateData(ThrowingIterator<KeyGroup> keyGroups, Map<Integer, ColumnFamilyHandle> columnFamilies) throws IOException, RocksDBException, StateMigrationException {
    // for all key-groups in the current state handle...
    try (RocksDBWriteBatchWrapper writeBatchWrapper = new RocksDBWriteBatchWrapper(this.rocksHandle.getDb(), writeBatchSize)) {
        ColumnFamilyHandle handle = null;
        while (keyGroups.hasNext()) {
            KeyGroup keyGroup = keyGroups.next();
            try (ThrowingIterator<KeyGroupEntry> groupEntries = keyGroup.getKeyGroupEntries()) {
                int oldKvStateId = -1;
                while (groupEntries.hasNext()) {
                    KeyGroupEntry groupEntry = groupEntries.next();
                    int kvStateId = groupEntry.getKvStateId();
                    if (kvStateId != oldKvStateId) {
                        oldKvStateId = kvStateId;
                        handle = columnFamilies.get(kvStateId);
                    }
                    writeBatchWrapper.put(handle, groupEntry.getKey(), groupEntry.getValue());
                }
            }
        }
    }
}

Example 48 with ColumnFamilyHandle

use of org.rocksdb.ColumnFamilyHandle in project flink by apache.

the class RocksDBWriteBatchWrapperTest method testWriteBatchWrapperFlushAfterCountExceed.

/**
 * Tests that {@link RocksDBWriteBatchWrapper} flushes after the kv count exceeds the
 * preconfigured value.
 */
@Test
public void testWriteBatchWrapperFlushAfterCountExceed() throws Exception {
    try (RocksDB db = RocksDB.open(folder.newFolder().getAbsolutePath());
        WriteOptions options = new WriteOptions().setDisableWAL(true);
        ColumnFamilyHandle handle = db.createColumnFamily(new ColumnFamilyDescriptor("test".getBytes()));
        RocksDBWriteBatchWrapper writeBatchWrapper = new RocksDBWriteBatchWrapper(db, options, 100, 50000)) {
        long initBatchSize = writeBatchWrapper.getDataSize();
        byte[] dummy = new byte[2];
        ThreadLocalRandom.current().nextBytes(dummy);
        for (int i = 1; i < 100; ++i) {
            writeBatchWrapper.put(handle, dummy, dummy);
            // each kv consumes 8 bytes
            assertEquals(initBatchSize + 8 * i, writeBatchWrapper.getDataSize());
        }
        writeBatchWrapper.put(handle, dummy, dummy);
        assertEquals(initBatchSize, writeBatchWrapper.getDataSize());
    }
}

Example 49 with ColumnFamilyHandle

use of org.rocksdb.ColumnFamilyHandle in project flink by apache.

the class EmbeddedRocksDBStateBackendTest method setupRocksKeyedStateBackend.

public void setupRocksKeyedStateBackend() throws Exception {
    blocker = new OneShotLatch();
    waiter = new OneShotLatch();
    testStreamFactory = new BlockerCheckpointStreamFactory(1024 * 1024);
    testStreamFactory.setBlockerLatch(blocker);
    testStreamFactory.setWaiterLatch(waiter);
    testStreamFactory.setAfterNumberInvocations(10);
    prepareRocksDB();
    RocksDBKeyedStateBackendBuilder keyedStateBackendBuilder = RocksDBTestUtils.builderForTestDB(TEMP_FOLDER.newFolder(), // injected
    IntSerializer.INSTANCE, spy(db), defaultCFHandle, optionsContainer.getColumnOptions()).setEnableIncrementalCheckpointing(enableIncrementalCheckpointing);
    if (enableIncrementalCheckpointing) {
        rocksDBStateUploader = spy(new RocksDBStateUploader(RocksDBOptions.CHECKPOINT_TRANSFER_THREAD_NUM.defaultValue()));
        keyedStateBackendBuilder.setRocksDBStateUploader(rocksDBStateUploader);
    }
    keyedStateBackend = keyedStateBackendBuilder.build();
    testState1 = keyedStateBackend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, new ValueStateDescriptor<>("TestState-1", Integer.class, 0));
    testState2 = keyedStateBackend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, new ValueStateDescriptor<>("TestState-2", String.class, ""));
    allCreatedCloseables = new ArrayList<>();
    doAnswer(new Answer<Object>() {

        @Override
        public Object answer(InvocationOnMock invocationOnMock) throws Throwable {
            RocksIterator rocksIterator = spy((RocksIterator) invocationOnMock.callRealMethod());
            allCreatedCloseables.add(rocksIterator);
            return rocksIterator;
        }
    }).when(keyedStateBackend.db).newIterator(any(ColumnFamilyHandle.class), any(ReadOptions.class));
    doAnswer(new Answer<Object>() {

        @Override
        public Object answer(InvocationOnMock invocationOnMock) throws Throwable {
            Snapshot snapshot = spy((Snapshot) invocationOnMock.callRealMethod());
            allCreatedCloseables.add(snapshot);
            return snapshot;
        }
    }).when(keyedStateBackend.db).getSnapshot();
    doAnswer(new Answer<Object>() {

        @Override
        public Object answer(InvocationOnMock invocationOnMock) throws Throwable {
            ColumnFamilyHandle snapshot = spy((ColumnFamilyHandle) invocationOnMock.callRealMethod());
            allCreatedCloseables.add(snapshot);
            return snapshot;
        }
    }).when(keyedStateBackend.db).createColumnFamily(any(ColumnFamilyDescriptor.class));
    for (int i = 0; i < 100; ++i) {
        keyedStateBackend.setCurrentKey(i);
        testState1.update(4200 + i);
        testState2.update("S-" + (4200 + i));
    }
}

Example 50 with ColumnFamilyHandle

use of org.rocksdb.ColumnFamilyHandle in project flink by apache.

the class EmbeddedRocksDBStateBackendTest method prepareRocksDB.

public void prepareRocksDB() throws Exception {
    String dbPath = new File(TEMP_FOLDER.newFolder(), DB_INSTANCE_DIR_STRING).getAbsolutePath();
    ColumnFamilyOptions columnOptions = optionsContainer.getColumnOptions();
    ArrayList<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(1);
    db = RocksDBOperationUtils.openDB(dbPath, Collections.emptyList(), columnFamilyHandles, columnOptions, optionsContainer.getDbOptions());
    defaultCFHandle = columnFamilyHandles.remove(0);
}