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

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

@Override
public long getFirstLogIndex() {
    this.readLock.lock();
    RocksIterator it = null;
    try {
        if (this.hasLoadFirstLogIndex) {
            return this.firstLogIndex;
        }
        checkState();
        it = this.db.newIterator(this.defaultHandle, this.totalOrderReadOptions);
        it.seekToFirst();
        if (it.isValid()) {
            final long ret = Bits.getLong(it.key(), 0);
            saveFirstLogIndex(ret);
            setFirstLogIndex(ret);
            return ret;
        }
        return 1L;
    } finally {
        if (it != null) {
            it.close();
        }
        this.readLock.unlock();
    }
}
 

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

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

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

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

    }
 

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

RocksDBIterator(final String storeName,
                final RocksIterator iter,
                final Set<KeyValueIterator<byte[], byte[]>> openIterators,
                final boolean isDescending) {
    this.storeName = storeName;
    this.iter = iter;
    this.openIterators = openIterators;
    this.isDescending = isDescending;
    if (isDescending) {
        iter.seekToLast();
    } else {
        iter.seekToFirst();
    }
}
 

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

}
 

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

public CloseableIterator<GeoWaveMetadata> iterator() {
  final RocksIterator it = db.newIterator();
  it.seekToFirst();
  return new RocksDBMetadataIterator(it, requiresTimestamp, visibilityEnabled);
}
 

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

@Override
public Stream<byte[]> streamKeys() {
  final RocksIterator rocksIterator = db.newIterator();
  rocksIterator.seekToFirst();
  return RocksDbKeyIterator.create(rocksIterator).toStream();
}
 

@Override
public Stream<byte[]> streamKeys(final ColumnFamilyHandle segmentHandle) {
  final RocksIterator rocksIterator = db.newIterator(segmentHandle);
  rocksIterator.seekToFirst();
  return RocksDbKeyIterator.create(rocksIterator).toStream();
}