Java Code Examples for org.rocksdb.RocksIterator#next()

@Override
public void iterate(byte[] from, EntryConsumer consumer)
    throws IOException {
  RocksIterator it = null;
  try {
    it = db.newIterator();
    if (from != null) {
      it.seek(from);
    } else {
      it.seekToFirst();
    }
    while (it.isValid()) {
      if (!consumer.consume(it.key(), it.value())) {
        break;
      }
      it.next();
    }
  } finally {
    if (it != null) {
      it.close();
    }
  }
}
 

@Test
public void testIterator() throws Exception {
    try (RocksDBMetronome db = new RocksDBMetronome.Builder()
            .setStoragePath(temporaryFolder.newFolder().toPath())
            .build()) {
        db.initialize();

        db.put(KEY, VALUE);
        db.put(KEY_2, VALUE_2);

        RocksIterator iterator = db.getIterator();
        iterator.seekToFirst();

        Map<String, byte[]> recovered = new HashMap<>();

        while (iterator.isValid()) {
            recovered.put(new String(iterator.key(), StandardCharsets.UTF_8), iterator.value());
            iterator.next();
        }

        assertEquals(2, recovered.size());
        assertArrayEquals(VALUE, recovered.get(new String(KEY, StandardCharsets.UTF_8)));
        assertArrayEquals(VALUE_2, recovered.get(new String(KEY_2, StandardCharsets.UTF_8)));
    }
}
 

@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 List<byte[]> pollByPrefix(final ColumnFamilyHandle cfh,final byte[] prefixKey,int nums){
    List<byte[]> values = new ArrayList<>();
    int count = 0;
    try{
        RocksIterator iterator = this.newIterator(cfh);
        WriteBatch writeBatch = new WriteBatch();
        for(iterator.seek(prefixKey);iterator.isValid();iterator.next()){
        	if(new String(iterator.key()).startsWith(new String(prefixKey))) {
        		values.add(iterator.value());
        		writeBatch.delete(cfh,iterator.key());
        		count++;
        	}
        	if(count>=nums) {
        		break;
        	}
        }
        if(count > 0){
            this.DB.write(WRITE_OPTIONS_SYNC,writeBatch);
        }
        log.debug("[RocksDB] -> succ while get by prefix,columnFamilyHandle:{}, pollByPrefix:{}",cfh.toString(),new String(prefixKey));
    }catch(Exception e){
        log.error("[RocksDB] ->  error while get by prefix, columnFamilyHandle:{}, pollByPrefix:{}, 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 ImmutablePair<byte[], byte[]> peekAround(int offset,
    byte[] from) throws IOException, IllegalArgumentException {
  RocksIterator it = null;
  try {
    it = db.newIterator();
    if (from == null) {
      it.seekToFirst();
    } else {
      it.seek(from);
    }
    if (!it.isValid()) {
      return null;
    }

    switch (offset) {
    case 0:
      break;
    case 1:
      it.next();
      break;
    case -1:
      it.prev();
      break;
    default:
      throw new IllegalArgumentException(
          "Position can only be -1, 0 " + "or 1, but found " + offset);
    }
    return it.isValid() ? new ImmutablePair<>(it.key(), it.value()) : null;
  } finally {
    if (it != null) {
      it.close();
    }
  }
}
 

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

/**
 * Perform a prefix delete and return stream of key-value pairs retrieved.
 *
 * @param columnFamilyName Column Family Name
 * @param prefix Prefix Key
 * @param <T> Type of value stored
 */
public <T extends Serializable> void prefixDelete(String columnFamilyName, String prefix) {
  ValidationUtils.checkArgument(!closed);
  LOG.info("Prefix DELETE (query=" + prefix + ") on " + columnFamilyName);
  final RocksIterator it = getRocksDB().newIterator(managedHandlesMap.get(columnFamilyName));
  it.seek(prefix.getBytes());
  // Find first and last keys to be deleted
  String firstEntry = null;
  String lastEntry = null;
  while (it.isValid() && new String(it.key()).startsWith(prefix)) {
    String result = new String(it.key());
    it.next();
    if (firstEntry == null) {
      firstEntry = result;
    }
    lastEntry = result;
  }
  it.close();

  if (null != firstEntry) {
    try {
      // This will not delete the last entry
      getRocksDB().deleteRange(managedHandlesMap.get(columnFamilyName), firstEntry.getBytes(), lastEntry.getBytes());
      // Delete the last entry
      getRocksDB().delete(lastEntry.getBytes());
    } catch (RocksDBException e) {
      LOG.error("Got exception performing range delete");
      throw new HoodieException(e);
    }
  }
}
 

@Benchmark
@SuppressWarnings("PMD.CloseResource")
public void readXxh64(final Reader r, final Blackhole bh) {
  long result = 0;
  final RocksIterator iterator = r.db.newIterator();
  iterator.seekToFirst();
  while (iterator.isValid()) {
    result += xx_r39().hashBytes(iterator.key());
    result += xx_r39().hashBytes(iterator.value());
    iterator.next();
  }
  bh.consume(result);
}
 

@Benchmark
@SuppressWarnings("PMD.CloseResource")
public void readSeq(final Reader r, final Blackhole bh) {
  final RocksIterator iterator = r.db.newIterator();
  iterator.seekToFirst();
  while (iterator.isValid()) {
    bh.consume(iterator.value());
    iterator.next();
  }
}
 

public void list() throws Exception {
    String crawldbPath = FilenameUtils.concat(crawlPath, "crawldb");
    RocksDB crawldbDatabase = RocksDBUtils.open(crawldbPath);
    RocksIterator crawldbIterator = crawldbDatabase.newIterator();

    for(crawldbIterator.seekToFirst(); crawldbIterator.isValid(); crawldbIterator.next()){
        CrawlDatum datum = RocksDBUtils.createCrawlDatum(crawldbIterator.key(), crawldbIterator.value());
        System.out.println(CrawlDatumFormater.datumToString(datum));
    }

    crawldbDatabase.close();

}
 

@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 StorageMapCursor<K, V> cursor(K from) {
    RocksIterator iterator = db.newIterator();
    if (from == null) {
        iterator.seekToFirst();
        return new RocksdbMapCursor<>(this, iterator);
    }
    for (iterator.seekToFirst(); iterator.isValid(); iterator.next()) {
        K key = k(iterator.key());
        if (keyType.compare(key, from) >= 0) {
            break;
        }
    }
    return new RocksdbMapCursor<>(this, iterator);
}
 

@Override
    public void merge() throws Exception {
        LOG.info("start merge");
        RocksDB crawldbDatabase = RocksDBUtils.openCrawldbDatabase(crawlPath);


        /*合并fetch库*/
        LOG.info("merge fetch database");
        RocksDB fetchDatabase = RocksDBUtils.openFetchDatabase(crawlPath);
        RocksIterator fetchIterator = fetchDatabase.newIterator();
        for(fetchIterator.seekToFirst(); fetchIterator.isValid(); fetchIterator.next()){
            crawldbDatabase.put(fetchIterator.key(), fetchIterator.value());
        }
        fetchDatabase.close();

        /*合并link库*/
        LOG.info("merge link database");
        RocksDB linkDatabase = RocksDBUtils.openLinkDatabase(crawlPath);
        RocksIterator linkIterator = linkDatabase.newIterator();
        for(linkIterator.seekToFirst(); linkIterator.isValid(); linkIterator.next()){
            if(crawldbDatabase.get(linkIterator.key()) == null){
                crawldbDatabase.put(linkIterator.key(), linkIterator.value());
            }
        }
        linkDatabase.close();

        LOG.info("end merge");
        crawldbDatabase.close();



//        env.removeDatabase(null, "fetch");
        RocksDBUtils.destroyFetchDatabase(crawlPath);
        LOG.debug("remove fetch database");
//        env.removeDatabase(null, "link");
        RocksDBUtils.destroyLinkDatabase(crawlPath);
        LOG.debug("remove link database");

    }
 

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

@Test
    public void testDeleteRange() {
        init();
        WriteBatch wb = new WriteBatch();
        ColumnFamilyHandle cfHandle = CFManager.CFH_DEFAULT;

        long st = System.currentTimeMillis();
        for(int i=100000; i<200000; i++) {
            wb.put(cfHandle, ("1324356527-" + i + "-5-5-345-356-234-232").getBytes(), "tasdfasdgasdfestfordb".getBytes());

            if(i % 30 == 0) {
                RDB.writeAsync(wb);
                wb.clear();
            }
        }
        for(int i=100000; i<200000; i++) {
            wb.put(cfHandle, ("1324356525-" + i + "-5-5-345-356-234-232").getBytes(), "tasdfasdgasdfestfordb".getBytes());

            if(i % 30 == 0) {
                RDB.writeAsync(wb);
                wb.clear();
            }
        }
        for(int i=100000; i<200000; i++) {
            wb.put(cfHandle, ("1324356529-" + i + "-5-5-345-356-234-232").getBytes(), "tasdfasdgasdfestfordb".getBytes());

            if(i % 30 == 0) {
                RDB.writeAsync(wb);
                wb.clear();
            }
        }
        RDB.writeAsync(wb);

        long ed = System.currentTimeMillis();
        System.out.println("write cost :" + (ed - st));

        try {
            TimeUnit.SECONDS.sleep(2);
        } catch (InterruptedException e) {
            e.printStackTrace();
        }

        long start = System.currentTimeMillis();
        RocksIterator it = RDB.newIterator(cfHandle);
        byte[] now = "1324356527".getBytes();
        long count = 0;
        for(it.seek(now); it.isValid(); it.next()) {
//            System.out.println(new String(it.key()) + " " + new String(it.value()));
            count++;
            if(count == 100000)
                break;
        }
        it.close();
        long end = System.currentTimeMillis();
        System.out.println("cost : " + (end - start) + " count:" +count);
        RDB.deleteFilesInRange(CFManager.CFH_DEFAULT, "132435653".getBytes(), "1324356529".getBytes());

        count = 0;
        it = RDB.newIterator(cfHandle);
        now = "1324356525".getBytes();
        for(it.seek(now); it.isValid(); it.next()) {
//            System.out.println(new String(it.key()) + " " + new String(it.value()));
            count++;
            if(count == 100000)
                break;
        }
        it.close();
        end = System.currentTimeMillis();
        System.out.println("cost : " + (end - start) + " count:" +count);

        destructor();
    }
 

@Override
public Map<ByteString, ByteString> getByPrefix(ByteString key, int max_reply, boolean allow_partial)
{
  ByteString key_str = prefix.concat(key);

  LinkedList<ByteString> set = new LinkedList<>();
Map<ByteString, ByteString> map = new HashMap<>(16,0.5f);

  int count = 0;
  RocksIterator it = db.newIterator();

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

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

      ByteString k = curr_key.substring(prefix.size());
      
     	map.put(k, ByteString.copyFrom(it.value()));
		count++;

      if (count > max_reply)
      {
        if (allow_partial) return map;
        else throw new DBTooManyResultsException();
      }

      it.next();
    }
  }
  finally
  {
    it.dispose();
  }

  return map;

}
 

@Override
public Set<String> senses() {
  return new AbstractSet<String>() {
    @Override
    public Iterator<String> iterator() {
      RocksIterator rocksIterator = rocksDB.newIterator();
      rocksIterator.seekToFirst();
      return new Iterator<String>() {
        boolean closed = false;

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

        @Override
        public String next() {
          if (!rocksIterator.isValid()) {
            throw new NoSuchElementException();
          }
          String sense = new String(rocksIterator.key(), StandardCharsets.UTF_8);
          rocksIterator.next();
          if (!rocksIterator.isValid()) {
            closed = true;
            rocksIterator.close();
          }
          return sense;
        }

        @Override
        protected void finalize() {
          if (!closed) {
            rocksIterator.close();
          }
        }
      };
    }

    @Override
    public boolean contains(Object o) {
      return o != null && o instanceof String && containsSense((String) o);
    }

    @Override
    public int size() {
      return RocksDBSenseVectors.this.size();
    }
  };
}
 

@Override
public Collection<SparseVector> vectors() {
  return new AbstractCollection<SparseVector>() {
    @Override
    public Iterator<SparseVector> iterator() {
      RocksIterator rocksIterator = rocksDB.newIterator();
      rocksIterator.seekToFirst();
      return new Iterator<SparseVector>() {
        boolean closed = false;

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

        @Override
        public SparseVector next() {
          if (!rocksIterator.isValid()) {
            throw new NoSuchElementException();
          }
          SparseVector sparseVector = new SparseVector(rocksIterator.value());
          rocksIterator.next();
          if (!rocksIterator.isValid()) {
            closed = true;
            rocksIterator.close();
          }
          return sparseVector;
        }

        @Override
        protected void finalize() {
          if (!closed) {
            rocksIterator.close();
          }
        }
      };
    }

    @Override
    public int size() {
      return RocksDBSenseVectors.this.size();
    }
  };
}