org.rocksdb.RocksIterator Java Examples

@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());
    }
}
 

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 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;
}
 

@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 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 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();
    }
}
 

public List<byte[]> range(String key, long start, long end) throws KitDBException {
    try (CloseLock ignored = checkClose()) {
        byte[] key_b = getKey(key);

        List<byte[]> list = new ArrayList<>();

        MetaV metaV = getMeta(key_b);
        if (metaV == null) {
            return list;
        }
        ValueK valueK_seek = new ValueK(key_b.length, key_b, metaV.getVersion(), start);
        try (final RocksIterator iterator = newIterator(SstColumnFamily.DEFAULT)) {
            ValueKD valueKD = valueK_seek.convertValueBytes();
            byte[] heads = valueKD.toHeadBytes();
            iterator.seek(valueKD.toBytes());
            long index = 0;
            while (iterator.isValid() && index < end) {
                byte[] key_bs = iterator.key();
                if (!BytesUtil.checkHead(heads, key_bs)) break;
                ValueK valueK = ValueKD.build(key_bs).convertValue();
                index = valueK.getIndex();
                list.add(iterator.value());
                iterator.next();
            }
        } catch (Exception e) {
            throw e;
        }
        return list;
    }
}
 

@Override
public MappingIterator mappingIterator() {
  RocksIterator rocksIterator = terms.newIterator();
  rocksIterator.seekToFirst();
  return new MappingIterator() {
    @Override
    public boolean isValid() {
      return rocksIterator.isValid();
    }

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

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

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

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

RocksDBRangeIterator(String storeName,
                     RocksIterator iter,
                     Set<KeyValueIterator<byte[], byte[]>> openIterators,
                     byte[] from,
                     boolean fromInclusive,
                     byte[] to,
                     boolean toInclusive,
                     boolean isDescending,
                     Comparator<byte[]> comparator) {
    super(storeName, iter, openIterators, isDescending);
    this.rawFromKey = from;
    if (rawFromKey == null) {
        if (isDescending) {
            iter.seekToLast();
        } else {
            iter.seekToFirst();
        }
    } else {
        if (isDescending) {
            iter.seekForPrev(rawFromKey);
        } else {
            iter.seek(rawFromKey);
        }
    }
    this.fromInclusive = fromInclusive;
    if (rawFromKey != null && !fromInclusive) {
        checkAndSkipFrom = true;
    }

    this.rawToKey = to;
    this.toInclusive = toInclusive;
    this.comparator = isDescending ? Collections.reverseOrder(comparator) : comparator;
}
 

@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");
        }
    };
}
 

public Set<ByteString> getSet(ByteString key, int max_reply)
{
ByteString dbKey = getDBKey(key);

  HashSet<ByteString> set = new HashSet<>();
  int count = 0;
  RocksIterator it = db.newIterator();

  try
  {
    it.seek(dbKey.toByteArray());

    while(it.isValid())
    {
		ByteString curr_key = ByteString.copyFrom(it.key());
      if (!curr_key.startsWith(dbKey)) break;

      ByteString v = curr_key.substring(dbKey.size());
      set.add(v);
      count++;

      if (count > max_reply) throw new DBTooManyResultsException();

      it.next();

    }

  }
  finally
  {
    it.dispose();
  }


  return set;

}
 

@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 List<ByteString> getSet(ByteString key, int max_reply)
{
ByteString dbKey = getDBKey(key);

  LinkedList<ByteString> set = new LinkedList<>();
  int count = 0;
  try(RocksIterator it = db.newIterator())
  {
    it.seek(dbKey.toByteArray());

    while(it.isValid())
    {
		ByteString curr_key = ByteString.copyFrom(it.key());
      if (!curr_key.startsWith(dbKey)) break;

      ByteString v = curr_key.substring(dbKey.size());
      set.add(v);
      count++;

      if (count > max_reply) throw new DBTooManyResultsException();

      it.next();
    }
  }

  return set;

}
 

@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();
  }
}
 

private static Map<Long,AccessPolicy> loadPolicies(RocksDB db, ColumnFamilyHandle policyCf) {
    Map<Long,AccessPolicy> map = new TreeMap<>();
    try (RocksIterator it = db.newIterator(policyCf)) {
        it.seekToFirst();
        while (it.isValid()) {
            AccessPolicy policy = GSON.fromJson(new String(it.value(), UTF_8), AccessPolicy.class);
            map.put(policy.id, policy);
            it.next();
        }
    }
    return map;
}
 

@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();
}
 

@MustBeClosed
public static <K, V> RocksDbIterator<K, V> create(
    final RocksDbColumn<K, V> column,
    final RocksIterator rocksIt,
    final Predicate<K> continueTest,
    final Supplier<Boolean> isDatabaseClosed) {
  return new RocksDbIterator<>(column, rocksIt, continueTest, isDatabaseClosed);
}
 

/**
 * Clear the array storage
 */
@Override
public void clear() {
    RocksIterator iterator = db.newIterator();
    while (iterator.isValid())
        try {
            db.remove(iterator.key());
        } catch (RocksDBException e) {
            throw new RuntimeException(e);
        }
    iterator.close();
    size = 0;
}