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

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

public CloseableIterator<GeoWaveRow> dataIndexIterator(
    final byte[] startDataId,
    final byte[] endDataId) {
  final RocksDB readDb = getReadDb();
  if (readDb == null) {
    return new CloseableIterator.Empty<>();
  }
  final ReadOptions options;
  final RocksIterator it;
  if (endDataId == null) {
    options = null;
    it = readDb.newIterator();
  } else {
    options =
        new ReadOptions().setIterateUpperBound(
            new Slice(ByteArrayUtils.getNextPrefix(endDataId)));
    it = readDb.newIterator(options);
  }
  if (startDataId == null) {
    it.seekToFirst();
  } else {
    it.seek(startDataId);
  }
  return new DataIndexRowIterator(options, it, adapterId, visibilityEnabled);
}
 

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

public boolean deleteByPrefix(final ColumnFamilyHandle cfh,final byte[] prefixKey,boolean sync){
    try{
        RocksIterator iterator = this.newIterator(cfh);
        int item = 0;
        WriteBatch writeBatch = new WriteBatch();
        for(iterator.seek(prefixKey);iterator.isValid();iterator.next()){
        	if(new String(iterator.key()).startsWith(new String(prefixKey))) {
        		writeBatch.delete(cfh,iterator.key());
                item++;
        	}
        }
        if(item > 0){
            this.DB.write(sync?WRITE_OPTIONS_SYNC:WRITE_OPTIONS_ASYNC,writeBatch);
        }
        log.debug("[RocksDB] -> succ while delete by prefix,columnFamilyHandle:{}, prefixKey:{}, nums:{}",cfh.toString(),new String(prefixKey),item);
    }catch(RocksDBException e){
        log.error("[RocksDB] ->  error while delete by prefix, columnFamilyHandle:{}, prefixKey:{}, err:{}",
                cfh.toString(), new String(prefixKey), e);
        return false;
    }
    return true;
}
 

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

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;

}
 

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

/**
 * 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);
    }
  }
}
 

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

protected KeyIterator getKeyIterator(byte[] head) {
    RocksIterator iterator = newIterator(SstColumnFamily.META);
    iterator.seek(head);
    KeyIterator keyIterator = new KeyIterator(iterator, head);
    return keyIterator;
}
 

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

}
 

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

private CloseableIterator<GeoWaveMetadata> prefixIterator(final byte[] prefix) {
  final ReadOptions options = new ReadOptions().setPrefixSameAsStart(true);
  final RocksIterator it = db.newIterator(options);
  it.seek(prefix);
  return new RocksDBMetadataIterator(options, it, requiresTimestamp, visibilityEnabled);
}
 

private List<Map.Entry<byte[], byte[]>> getRangeKVs(byte[] startKey,
    int count, boolean sequential,
    MetadataKeyFilters.MetadataKeyFilter... filters)
    throws IOException, IllegalArgumentException {
  List<Map.Entry<byte[], byte[]>> result = new ArrayList<>();
  long start = System.currentTimeMillis();
  if (count < 0) {
    throw new IllegalArgumentException(
        "Invalid count given " + count + ", count must be greater than 0");
  }
  RocksIterator it = null;
  try {
    it = db.newIterator();
    if (startKey == null) {
      it.seekToFirst();
    } else {
      if(get(startKey) == null) {
        // Key not found, return empty list
        return result;
      }
      it.seek(startKey);
    }
    while(it.isValid() && result.size() < count) {
      byte[] currentKey = it.key();
      byte[] currentValue = it.value();

      it.prev();
      final byte[] prevKey = it.isValid() ? it.key() : null;

      it.seek(currentKey);
      it.next();
      final byte[] nextKey = it.isValid() ? it.key() : null;

      if (filters == null) {
        result.add(new AbstractMap.SimpleImmutableEntry<>(currentKey,
            currentValue));
      } else {
        if (Arrays.asList(filters).stream()
            .allMatch(entry -> entry.filterKey(prevKey,
                currentKey, nextKey))) {
          result.add(new AbstractMap.SimpleImmutableEntry<>(currentKey,
              currentValue));
        } else {
          if (result.size() > 0 && sequential) {
            // if the caller asks for a sequential range of results,
            // and we met a dis-match, abort iteration from here.
            // if result is empty, we continue to look for the first match.
            break;
          }
        }
      }
    }
  } finally {
    if (it != null) {
      it.close();
    }
    long end = System.currentTimeMillis();
    long timeConsumed = end - start;
    if (LOG.isDebugEnabled()) {
      if (filters != null) {
        for (MetadataKeyFilters.MetadataKeyFilter filter : filters) {
          int scanned = filter.getKeysScannedNum();
          int hinted = filter.getKeysHintedNum();
          if (scanned > 0 || hinted > 0) {
            LOG.debug(
                "getRangeKVs ({}) numOfKeysScanned={}, numOfKeysHinted={}",
                filter.getClass().getSimpleName(), filter.getKeysScannedNum(),
                filter.getKeysHintedNum());
          }
        }
      }
      LOG.debug("Time consumed for getRangeKVs() is {}ms,"
          + " result length is {}.", timeConsumed, result.size());
    }
  }
  return result;
}
 

public Set<Sha256Hash> getSet(String key, int max_reply)
{
  String s = name + "/" + key + "/";

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

  try
  {
    it.seek(s.getBytes());

    while(it.isValid())
    { 
      String curr_key = new String(it.key());
      if (!curr_key.startsWith(s)) break;

      String hash_string = curr_key.substring(s.length());
      set.add(Sha256Hash.wrap(hash_string));
      count++;

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

      it.next();
            
    }

  }
  finally
  {
    it.dispose();
  } 
  /*catch(RocksDBException e)
  { 
    throw new RuntimeException(e);
  }*/


  return set;

}