Java Code Examples for org.rocksdb.RocksIterator

@Override
public void clear(final ColumnFamilyHandle segmentHandle) {
  try (final RocksIterator rocksIterator = db.newIterator(segmentHandle)) {
    rocksIterator.seekToFirst();
    if (rocksIterator.isValid()) {
      final byte[] firstKey = rocksIterator.key();
      rocksIterator.seekToLast();
      if (rocksIterator.isValid()) {
        final byte[] lastKey = rocksIterator.key();
        db.deleteRange(segmentHandle, firstKey, lastKey);
        db.delete(segmentHandle, lastKey);
      }
    }
  } catch (final RocksDBException e) {
    throw new StorageException(e);
  }
}
 

@Override
public void removeWords(Collection<Integer> indexes) {
  try (WriteBatch writeBatch = new WriteBatch()) {
    try (RocksIterator rocksIterator = rocksDB.newIterator()) {
      rocksIterator.seekToFirst();
      while (rocksIterator.isValid()) {
        SparseVector sparseVector = new SparseVector(rocksIterator.value());
        sparseVector.removeAll(indexes);
        writeBatch.put(rocksIterator.key(), sparseVector.toBytes());
      }
    }
    rocksDB.write(new WriteOptions(), writeBatch);
  } catch (RocksDBException e) {
    throw new RuntimeException(e);
  }
}
 

public Records(RocksDB db, ColumnFamilyHandle columnFamilyHandle, byte[] startKey, RecordConstructor<T> constructor, Predicate<T> scope, boolean reverse, long cap) {
    final RocksIterator it = db.newIterator(columnFamilyHandle);
    it.seek(startKey);
    if (reverse) {
        if (it.isValid()) {
            it.prev();
        } else {
            it.seekToLast();
        }
    }
    this.constructor = constructor;
    this.scope = scope;
    this.it = it;
    this.reverse = reverse;
    this.cap = cap;
}
 

public List<byte[]> getByPrefix(final ColumnFamilyHandle cfh,final byte[] prefixKey){
    List<byte[]> values = new ArrayList<>();
    try{
        RocksIterator iterator = this.newIterator(cfh);
        for(iterator.seek(prefixKey);iterator.isValid();iterator.next()){
        	if(new String(iterator.key()).startsWith(new String(prefixKey))) {
        		values.add(iterator.value());
        	}
        }
        log.debug("[RocksDB] -> succ while get by prefix,columnFamilyHandle:{}, prefixKey:{}",cfh.toString(),new String(prefixKey));
    }catch(Exception e){
        log.error("[RocksDB] ->  error while get by prefix, columnFamilyHandle:{}, prefixKey:{}, err:{}",
                cfh.toString(), new String(prefixKey), e);
    }
    return values;
}
 

public int getCountByPrefix(final ColumnFamilyHandle cfh,final byte[] prefixKey){
    int count = 0;
    try{
        RocksIterator iterator = this.newIterator(cfh);
        for(iterator.seek(prefixKey);iterator.isValid();iterator.next()){
        	if(new String(iterator.key()).startsWith(new String(prefixKey))) {
        		count++;
        	}
        }
        log.debug("[RocksDB] -> succ while get count by prefix,columnFamilyHandle:{}, prefixKey:{}",cfh.toString(),new String(prefixKey));
    }catch(Exception e){
        log.error("[RocksDB] ->  error while get count by prefix, columnFamilyHandle:{}, prefixKey:{}, err:{}",
                cfh.toString(), new String(prefixKey), e);
    }
    return count;
}
 

@Override
public Collection<K> getAllKeys(TimeWindow window) {
    try {
        ColumnFamilyHandle handler = getColumnFamilyHandle(window);
        Collection<K> keys = new ArrayList<K>();
        RocksIterator itr = rocksDb.newIterator(handler);
        itr.seekToFirst();
        while (itr.isValid()) {
            keys.add((K) serializer.deserialize(itr.key()));
            itr.next();
        }
        return keys;
    } catch (RocksDBException e) {
        LOG.error("Failed to get all keys for timeWindow={}", window);
        throw new RuntimeException(e.getMessage());
    }
}
 

@Override
public Set<byte[]> allKeys() throws RuntimeException {
    if (quitIfNotAlive()) {
        return null;
    }
    resetDbLock.readLock().lock();
    Set<byte[]> result = Sets.newHashSet();
    try (final RocksIterator iter = database.newIterator()) {
        for (iter.seekToFirst(); iter.isValid(); iter.next()) {
            result.add(iter.key());
        }
        return result;
    } finally {
        resetDbLock.readLock().unlock();
    }
}
 

@Override
public long getTotal() throws RuntimeException {
    if (quitIfNotAlive()) {
        return 0;
    }
    resetDbLock.readLock().lock();
    try (RocksIterator iterator = database.newIterator()) {
        long total = 0;
        for (iterator.seekToFirst(); iterator.isValid(); iterator.next()) {
            total++;
        }
        return total;
    } finally {
        resetDbLock.readLock().unlock();
    }
}
 

public Map<byte[], byte[]> getNext(byte[] key, long limit) {
    if (quitIfNotAlive()) {
        return null;
    }
    if (limit <= 0) {
        return Collections.emptyMap();
    }
    resetDbLock.readLock().lock();
    try (RocksIterator iter = database.newIterator()) {
        Map<byte[], byte[]> result = new HashMap<>();
        long i = 0;
        for (iter.seek(key); iter.isValid() && i < limit; iter.next(), i++) {
            result.put(iter.key(), iter.value());
        }
        return result;
    } finally {
        resetDbLock.readLock().unlock();
    }
}
 

public Set<byte[]> getlatestValues(long limit) {
    if (quitIfNotAlive()) {
        return null;
    }
    if (limit <= 0) {
        return Sets.newHashSet();
    }
    resetDbLock.readLock().lock();
    try (RocksIterator iter = database.newIterator()) {
        Set<byte[]> result = Sets.newHashSet();
        long i = 0;
        for (iter.seekToLast(); iter.isValid() && i < limit; iter.prev(), i++) {
            result.add(iter.value());
        }
        return result;
    } finally {
        resetDbLock.readLock().unlock();
    }
}
 

public Set<byte[]> getValuesPrev(byte[] key, long limit) {
    if (quitIfNotAlive()) {
        return null;
    }
    if (limit <= 0) {
        return Sets.newHashSet();
    }
    resetDbLock.readLock().lock();
    try (RocksIterator iter = database.newIterator()) {
        Set<byte[]> result = Sets.newHashSet();
        long i = 0;
        byte[] data = getData(key);
        if (Objects.nonNull(data)) {
            result.add(data);
            i++;
        }
        for (iter.seekForPrev(key); iter.isValid() && i < limit; iter.prev(), i++) {
            result.add(iter.value());
        }
        return result;
    } finally {
        resetDbLock.readLock().unlock();
    }
}
 

public Map<byte[], byte[]> getPrevious(byte[] key, long limit, int precision) {
    if (quitIfNotAlive()) {
        return null;
    }
    if (limit <= 0 || key.length < precision) {
        return Collections.emptyMap();
    }
    resetDbLock.readLock().lock();
    try (RocksIterator iterator = database.newIterator()) {
        Map<byte[], byte[]> result = new HashMap<>();
        long i = 0;
        for (iterator.seekToFirst(); iterator.isValid() && i++ < limit; iterator.next()) {

            if (iterator.key().length >= precision) {
                if (ByteUtil.less(ByteUtil.parseBytes(key, 0, precision),
                        ByteUtil.parseBytes(iterator.key(), 0, precision))) {
                    break;
                }
                result.put(iterator.key(), iterator.value());
            }
        }
        return result;
    } finally {
        resetDbLock.readLock().unlock();
    }
}
 

@Override
public boolean isEmpty() throws IOException {
  RocksIterator it = null;
  try {
    it = db.newIterator();
    it.seekToFirst();
    return !it.isValid();
  } finally {
    if (it != null) {
      it.close();
    }
  }
}
 

@Override
public MappingIterator mappingIterator() {
  RocksIterator rocksIterator = indices.newIterator();
  return new MappingIterator() {
    @Override
    public void close() {
      rocksIterator.close();
    }

    @Override
    public boolean isValid() {
      return rocksIterator.isValid();
    }

    @Override
    public int identifier() {
      byte[] value = rocksIterator.value();
      return ByteBuffer.wrap(value).getInt();
    }

    @Override
    public String string() {
      byte[] key = rocksIterator.key();
      return new String(key, StandardCharsets.UTF_8);
    }

    @Override
    public void next() {
      rocksIterator.next();
    }
  };
}
 

@Override
public Collection<K> getAllKeys() {
    Collection<K> keys = new ArrayList<K>();
    RocksIterator itr = rocksDb.newIterator();
    itr.seekToFirst();
    while (itr.isValid()) {
        keys.add((K) serializer.deserialize(itr.key()));
        itr.next();
    }
    return keys;
}
 

@Override
public Entry getEntry(RocksIterator iterator) throws KitDBException {
    try (CloseLock ignored = checkClose()) {
        byte[] key_bs = iterator.key();
        if (key_bs == null) {
            return null;
        }
        SData sData = SDataD.build(key_bs).convertValue();
        Entry entry = new Entry(ArrayKits.bytesToLong(iterator.value()), sData.value);
        return entry;
    }
}
 

@Override
public void scan(final byte[] startKey, final byte[] endKey, final int limit,
                 @SuppressWarnings("unused") final boolean readOnlySafe, final boolean returnValue,
                 final KVStoreClosure closure) {
    final Timer.Context timeCtx = getTimeContext("SCAN");
    final List<KVEntry> entries = Lists.newArrayList();
    final int maxCount = normalizeLimit(limit);
    final Lock readLock = this.readWriteLock.readLock();
    readLock.lock();
    try (final RocksIterator it = this.db.newIterator()) {
        if (startKey == null) {
            it.seekToFirst();
        } else {
            it.seek(startKey);
        }
        int count = 0;
        while (it.isValid() && count++ < maxCount) {
            final byte[] key = it.key();
            if (endKey != null && BytesUtil.compare(key, endKey) >= 0) {
                break;
            }
            entries.add(new KVEntry(key, returnValue ? it.value() : null));
            it.next();
        }
        setSuccess(closure, entries);
    } catch (final Exception e) {
        LOG.error("Fail to [SCAN], range: ['[{}, {})'], {}.", BytesUtil.toHex(startKey), BytesUtil.toHex(endKey),
            StackTraceUtil.stackTrace(e));
        setFailure(closure, "Fail to [SCAN]");
    } finally {
        readLock.unlock();
        timeCtx.stop();
    }
}
 

@Override
public long getApproximateKeysInRange(final byte[] startKey, final byte[] endKey) {
    // TODO This is a sad code, the performance is too damn bad
    final Timer.Context timeCtx = getTimeContext("APPROXIMATE_KEYS");
    final Lock readLock = this.readWriteLock.readLock();
    readLock.lock();
    final Snapshot snapshot = this.db.getSnapshot();
    try (final ReadOptions readOptions = new ReadOptions()) {
        readOptions.setSnapshot(snapshot);
        try (final RocksIterator it = this.db.newIterator(readOptions)) {
            if (startKey == null) {
                it.seekToFirst();
            } else {
                it.seek(startKey);
            }
            long approximateKeys = 0;
            for (;;) {
                // The accuracy is 100, don't ask more
                for (int i = 0; i < 100; i++) {
                    if (!it.isValid()) {
                        return approximateKeys;
                    }
                    it.next();
                    ++approximateKeys;
                }
                if (endKey != null && BytesUtil.compare(it.key(), endKey) >= 0) {
                    return approximateKeys;
                }
            }
        }
    } finally {
        // Nothing to release, rocksDB never own the pointer for a snapshot.
        snapshot.close();
        // The pointer to the snapshot is released by the database instance.
        this.db.releaseSnapshot(snapshot);
        readLock.unlock();
        timeCtx.stop();
    }
}
 

public DataIndexRowIterator(
    final ReadOptions options,
    final RocksIterator it,
    final short adapterId,
    final boolean visiblityEnabled) {
  super(options, it);
  this.adapterId = adapterId;
  visibilityEnabled = visiblityEnabled;
}
 

public CloseableIterator<GeoWaveRow> iterator(final ByteArrayRange range) {
  final RocksDB readDb = getReadDb();
  if (readDb == null) {
    return new CloseableIterator.Empty<>();
  }
  final ReadOptions options;
  final RocksIterator it;
  if (range.getEnd() == null) {
    options = null;
    it = readDb.newIterator();
  } else {
    options = new ReadOptions().setIterateUpperBound(new Slice(range.getEndAsNextPrefix()));
    it = readDb.newIterator(options);
  }
  if (range.getStart() == null) {
    it.seekToFirst();
  } else {
    it.seek(range.getStart());
  }

  return new RocksDBRowIterator(
      options,
      it,
      adapterId,
      partition,
      requiresTimestamp,
      visibilityEnabled);
}
 

@Override
public long getLastLogIndex() {
    this.readLock.lock();
    checkState();
    try (final RocksIterator it = this.db.newIterator(this.defaultHandle, this.totalOrderReadOptions)) {
        it.seekToLast();
        if (it.isValid()) {
            return Bits.getLong(it.key(), 0);
        }
        return 0L;
    } finally {
        this.readLock.unlock();
    }
}
 

@Test
    public void backupDb() throws RocksDBException {
//        String originPath = dbFolder.getRoot().getAbsolutePath();
//        String backupPath = backupFolder.getRoot().getAbsolutePath();
        String originPath = "/tmp/rocksdb";
        String backupPath = "/tmp/rocksdb_backup";
        System.out.println("originPath=" + originPath);
        System.out.println("backupPath=" + backupPath);
        // Open empty database.
        try (final Options opt = new Options().setCreateIfMissing(true);
             final RocksDB db = RocksDB.open(opt, originPath)) {

            // Fill database with some test values
            prepareDatabase(db);

            try (RocksIterator it = db.newIterator()) {
                for (it.seekToFirst(); it.isValid(); it.next()) {
                    System.out.println(originPath + ":" + new String(it.key()) + ":" + new String(it.value()));
                }
            }

            // Create two backups
            try (final BackupableDBOptions bopt = new BackupableDBOptions(backupPath);
                 final BackupEngine be = BackupEngine.open(opt.getEnv(), bopt)) {
                be.createNewBackup(db, false);
                be.createNewBackup(db, true);
                verifyNumberOfValidBackups(be, 2);
            }
        }
    }
 

@Test
    public void backupDb2() throws RocksDBException {
//        String originPath = dbFolder.getRoot().getAbsolutePath();
//        String backupPath = backupFolder.getRoot().getAbsolutePath();
        String originPath = "/tmp/rocksdb";
        String originPath2 = "/tmp/rocksdb2";
        String backupPath = "/tmp/rocksdb_backup";
        System.out.println("originPath=" + originPath);
        System.out.println("backupPath=" + backupPath);
        // Open empty database.
        try (final Options opt = new Options().setCreateIfMissing(true);
             final RocksDB db = RocksDB.open(opt, originPath)) {

            // Fill database with some test values
            prepareDatabase(db);

            try (RocksIterator it = db.newIterator()) {
                for (it.seekToFirst(); it.isValid(); it.next()) {
                    System.out.println(originPath + ":" + new String(it.key()) + ":" + new String(it.value()));
                }
            }

            // Create two backups
            try (final BackupableDBOptions bopt = new BackupableDBOptions(backupPath);
                 final BackupEngine be = BackupEngine.open(opt.getEnv(), bopt)) {
                be.createNewBackup(db, true);
                be.createNewBackup(db, true);

                //restore the backup
                final List<BackupInfo> backupInfo = verifyNumberOfValidBackups(be, 2);
                // restore db from first backup
                be.restoreDbFromBackup(backupInfo.get(0).backupId(), originPath2, originPath2, new RestoreOptions(true));
                // Open database again.
                RocksDB db2 = RocksDB.open(opt, originPath2);

                try (RocksIterator it = db2.newIterator()) {
                    for (it.seekToFirst(); it.isValid(); it.next()) {
                        System.out.println(originPath2 + ":" + new String(it.key()) + ":" + new String(it.value()));
                    }
                }

                db2.close();
            }
        }


    }
 

@Override
public Iterator<String[]> iterator() {
    final RocksIterator rocksIterator = getRocksIterator();
    rocksIterator.seekToFirst();

    return new Iterator<String[]>() {
        int counter;

        @Override
        public boolean hasNext() {
            boolean valid = rocksIterator.isValid();
            if (!valid) {
                rocksIterator.close();
            }
            return valid;
        }

        @Override
        public String[] next() {
            counter++;
            if (counter % 100000 == 0) {
                logger.info("scanned {} rows from rocksDB", counter);
            }
            String[] result = rowEncoder.decode(new KV(rocksIterator.key(), rocksIterator.value()));
            rocksIterator.next();
            return result;
        }

        @Override
        public void remove() {
            throw new UnsupportedOperationException("not support operation");
        }
    };
}
 

@Override
public boolean isEmpty() {
    try (RocksIterator iterator = db.newIterator()) {
        iterator.seekToFirst();
        if (iterator.isValid()) {
            return true;
        }
    }
    return false;
}
 

@Override
public Iterator<String[]> iterator() {
    final RocksIterator rocksIterator = getRocksIterator();
    rocksIterator.seekToFirst();

    return new Iterator<String[]>() {
        int counter;

        @Override
        public boolean hasNext() {
            boolean valid = rocksIterator.isValid();
            if (!valid) {
                rocksIterator.close();
            }
            return valid;
        }

        @Override
        public String[] next() {
            counter++;
            if (counter % 100000 == 0) {
                logger.info("scanned {} rows from rocksDB", counter);
            }
            String[] result = rowEncoder.decode(new KV(rocksIterator.key(), rocksIterator.value()));
            rocksIterator.next();
            return result;
        }

        @Override
        public void remove() {
            throw new UnsupportedOperationException("not support operation");
        }
    };
}
 

@Override
public boolean isEmpty() {
    check();

    try (RocksIterator itr = db.newIterator()) {
        itr.seekToFirst();

        // check if there is at least one valid item
        return !itr.isValid();
    } catch (Exception e) {
        LOG.error("Unable to extract information from database " + this.toString() + ".", e);
    }

    return true;
}
 

/**
 * @implNote Building two wrappers for the same {@link RocksIterator} will lead to
 *     inconsistent behavior.
 */
RocksDBIteratorWrapper(final ReadOptions readOptions, final RocksIterator iterator) {
    this.readOptions = readOptions;
    this.iterator = iterator;
    iterator.seekToFirst();
    closed = false;
}
 

@Override
public <K, V> Optional<ColumnEntry<K, V>> getFloorEntry(RocksDbColumn<K, V> column, final K key) {
  assertOpen();
  final byte[] keyBytes = column.getKeySerializer().serialize(key);
  final Consumer<RocksIterator> setupIterator = it -> it.seekForPrev(keyBytes);
  try (final Stream<ColumnEntry<K, V>> stream = createStream(column, setupIterator)) {
    return stream.findFirst();
  }
}
 

@SuppressWarnings("MustBeClosedChecker")
@MustBeClosed
private <K, V> Stream<ColumnEntry<K, V>> createStream(
    RocksDbColumn<K, V> column,
    Consumer<RocksIterator> setupIterator,
    Predicate<K> continueTest) {
  final ColumnFamilyHandle handle = columnHandles.get(column);
  final RocksIterator rocksDbIterator = db.newIterator(handle);
  setupIterator.accept(rocksDbIterator);
  return RocksDbIterator.create(column, rocksDbIterator, continueTest, closed::get).toStream();
}