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

@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 Entry getEntry(RocksIterator iterator) throws KitDBException {
    try (CloseLock ignored = checkClose()) {
        byte[] key_bs = iterator.key();
        if (key_bs == null) {
            return null;
        }

        KeyD keyD = KeyD.build(key_bs);
        Entry entry = new Entry(new String(keyD.key, charset), iterator.value());
        return entry;
    }
}
 

@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(sData.value);
        return entry;
    }
}
 

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

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

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

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

protected void merge(Pair<RocksDB, Map<COLUMN_FAMILIES, ColumnFamilyHandle>> dbAndHandles)
    throws RocksDBException, IOException, ClassNotFoundException {
  LOGGER.info("Beginning merge on Pubchem CID");
  RocksDB db = dbAndHandles.getLeft();
  ColumnFamilyHandle pubchemIdCFH = dbAndHandles.getRight().get(COLUMN_FAMILIES.CID_TO_HASHES);
  ColumnFamilyHandle meshCFH = dbAndHandles.getRight().get(COLUMN_FAMILIES.HASH_TO_MESH);
  ColumnFamilyHandle synonymCFH = dbAndHandles.getRight().get(COLUMN_FAMILIES.HASH_TO_SYNONYMS);
  ColumnFamilyHandle synonymTypeCFH = dbAndHandles.getRight().get(COLUMN_FAMILIES.HASH_TO_SYNONYM_TYPE);
  ColumnFamilyHandle mergeResultsCFH = dbAndHandles.getRight().get(COLUMN_FAMILIES.CID_TO_SYNONYMS);

  RocksIterator cidIterator = db.newIterator(pubchemIdCFH);
  // With help from https://github.com/facebook/rocksdb/wiki/Basic-Operations
  int processed = 0;
  for (cidIterator.seekToFirst(); cidIterator.isValid(); cidIterator.next()) {
    byte[] key = cidIterator.key();
    byte[] val = cidIterator.value();
    String pubchemId = new String(key, UTF8);
    List<String> hashes;
    try (ObjectInputStream ois = new ObjectInputStream(new ByteArrayInputStream(val))) {
      // We know all our values so far have been lists of strings, so this should be completely safe.
      hashes = (List<String>) ois.readObject();
    }

    PubchemSynonyms pubchemSynonyms = new PubchemSynonyms(pubchemId);

    /* The hash keys are based on synonym value, which we can manually compute with:
     *   $ echo -n  'dimethyltin(iv)' | md5
     * This means that MeSH ids are linked to synonyms rather than pubchem ids.  We need to look up each cid-linked
     * hash in both the MeSH and synonym collections, as the key may legitimately exist in both (and serve to link
     * cid to synonym and cid to MeSH). */
    for (String hash : hashes) {
      /* Note: these ids are not proper MeSH topic ids, but are internal MeSH ids found in the RDF and TTL
       * representations of the MeSH corpus.  You can find them in the MeSH .nt or .xml files, but they won't turn up
       * anything on the MeSH website. */
      List<String> meshIds = getValueAsObject(db, meshCFH, hash);
      if (meshIds != null) {
        pubchemSynonyms.addMeSHIds(meshIds);
      }

      List<String> synonyms = getValueAsObject(db, synonymCFH, hash);
      // There are, surprisingly, some dangling hashes in the DB!  Handle them gracefully.
      if (synonyms == null) {
        LOGGER.warn("Dangling synonym hash reference, adding empty list in place of value: cid = %s, hash = %s",
            pubchemId, hash);
        synonyms = Collections.emptyList();
      }

      List<String> synonymTypeStrings = getValueAsObject(db, synonymTypeCFH, hash);
      Set<PC_SYNONYM_TYPES> synonymTypes = DEFAULT_SYNONYM_DATA_TYPES;
      if (synonymTypeStrings != null) {
        synonymTypes = synonymTypeStrings.stream().map(PC_SYNONYM_TYPES::valueOf).collect(Collectors.toSet());
      }

      if (synonymTypes.size() == 0) {
        LOGGER.warn("Found zero synonym types for synonym, defaulting to %s: %s %s, synonyms = %s",
            PC_SYNONYM_TYPES.UNKNOWN.name(), pubchemId, hash, StringUtils.join(synonyms, ", "));
      }
      /* It turns out that *lots* of synonyms are duplicated as depositor supplied names, so don't complain about it
       * here.  For performance sake we might want to consider changing the data model of PubchemSynonyms to reduce
       * synonym string duplication, as the current model is pretty inefficient. */

      for (PC_SYNONYM_TYPES synonymType : synonymTypes) {
        for (String synonym : synonyms) {
          // Let the PubchemSynonyms object do the de-duplication for us rather than reducing `synonyms` to a Set.
          pubchemSynonyms.addSynonym(synonymType, synonym);
        }
      }
    }

    try (ByteArrayOutputStream bos = new ByteArrayOutputStream();
         ObjectOutputStream oo = new ObjectOutputStream(bos)) {
      oo.writeObject(pubchemSynonyms);
      oo.flush();

      db.put(mergeResultsCFH, key, bos.toByteArray());
    }

    processed++;
    if (processed % 100000 == 0) {
      LOGGER.info("Merged %d entries on Pubchem compound id", processed);
    }
  }
  LOGGER.info("Merge complete, %d entries processed", processed);
}
 

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

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;

}