61 lines
1.7 KiB
C++
61 lines
1.7 KiB
C++
// (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
|
|
|
|
#pragma once
|
|
|
|
#include <faiss/invlists/InvertedLists.h>
|
|
|
|
#include <rocksdb/db.h>
|
|
|
|
namespace faiss_rocksdb {
|
|
|
|
struct RocksDBInvertedListsIterator : faiss::InvertedListsIterator {
|
|
RocksDBInvertedListsIterator(
|
|
rocksdb::DB* db,
|
|
size_t list_no,
|
|
size_t code_size);
|
|
virtual bool is_available() const override;
|
|
virtual void next() override;
|
|
virtual std::pair<faiss::idx_t, const uint8_t*> get_id_and_codes() override;
|
|
|
|
private:
|
|
std::unique_ptr<rocksdb::Iterator> it;
|
|
size_t list_no;
|
|
size_t code_size;
|
|
std::vector<uint8_t> codes; // buffer for returning codes in next()
|
|
};
|
|
|
|
struct RocksDBInvertedLists : faiss::InvertedLists {
|
|
RocksDBInvertedLists(
|
|
const char* db_directory,
|
|
size_t nlist,
|
|
size_t code_size);
|
|
|
|
size_t list_size(size_t list_no) const override;
|
|
const uint8_t* get_codes(size_t list_no) const override;
|
|
const faiss::idx_t* get_ids(size_t list_no) const override;
|
|
|
|
size_t add_entries(
|
|
size_t list_no,
|
|
size_t n_entry,
|
|
const faiss::idx_t* ids,
|
|
const uint8_t* code) override;
|
|
|
|
void update_entries(
|
|
size_t list_no,
|
|
size_t offset,
|
|
size_t n_entry,
|
|
const faiss::idx_t* ids,
|
|
const uint8_t* code) override;
|
|
|
|
void resize(size_t list_no, size_t new_size) override;
|
|
|
|
faiss::InvertedListsIterator* get_iterator(
|
|
size_t list_no,
|
|
void* inverted_list_context) const override;
|
|
|
|
private:
|
|
std::unique_ptr<rocksdb::DB> db_;
|
|
};
|
|
|
|
} // namespace faiss_rocksdb
|