Skip to content

Commit

Permalink
TODO: PR this separately
Browse files Browse the repository at this point in the history
  • Loading branch information
@laurynas-biveinis
laurynas-biveinis committed Apr 26, 2024
1 parent 1af3e9f commit 1927fc6
Show file tree
Hide file tree
Showing 3 changed files with 41 additions and 29 deletions.
4 changes: 2 additions & 2 deletions storage/rocksdb/rdb_datadic.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3557,8 +3557,8 @@ uint Rdb_key_def::setup_vector_index(const TABLE &tbl,
}

return create_vector_index(cmd_srv_helper, tbl_def.base_dbname(),
m_vector_index_config, m_cf_handle, m_index_number,
m_vector_index);
m_vector_index_config, m_cf_handle, *this,
m_index_number, m_vector_index);
}

// See Rdb_charset_space_info::spaces_xfrm
Expand Down
62 changes: 36 additions & 26 deletions storage/rocksdb/rdb_vector_db.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -175,8 +175,8 @@ class Rdb_faiss_inverted_list_write_context {
class Rdb_vector_iterator : public faiss::InvertedListsIterator {
public:
Rdb_vector_iterator(Rdb_faiss_inverted_list_context *context,
Index_id index_id, rocksdb::ColumnFamilyHandle *const cf,
const uint code_size, const size_t list_id)
Index_id index_id, rocksdb::ColumnFamilyHandle &cf,
const Rdb_key_def &kd, uint code_size, size_t list_id)
: m_context(context),
m_index_id(index_id),
m_list_id(list_id),
Expand All @@ -188,10 +188,10 @@ class Rdb_vector_iterator : public faiss::InvertedListsIterator {
Rdb_string_writer upper_key_writer;
write_inverted_list_key(upper_key_writer, index_id, list_id + 1);
m_iterator_upper_bound_key.PinSelf(upper_key_writer.to_slice());
m_iterator.reset(rdb_tx_get_iterator(
context->m_thd, cf, /* skip_bloom_filter */ true,
m_iterator = rdb_tx_get_iterator(
context->m_thd, cf, kd, /* skip_bloom_filter */ true,
m_iterator_lower_bound_key, m_iterator_upper_bound_key,
/* snapshot */ nullptr, TABLE_TYPE::USER_TABLE));
/* snapshot */ nullptr, TABLE_TYPE::USER_TABLE);
m_iterator->SeekToFirst();
}

Expand Down Expand Up @@ -268,14 +268,21 @@ class Rdb_vector_iterator : public faiss::InvertedListsIterator {
*/
class Rdb_faiss_inverted_list : public faiss::InvertedLists {
public:
Rdb_faiss_inverted_list(Index_id index_id,
rocksdb::ColumnFamilyHandle *const cf, uint nlist,
uint code_size)
: InvertedLists(nlist, code_size), m_index_id(index_id), m_cf(cf) {
Rdb_faiss_inverted_list(Index_id index_id, rocksdb::ColumnFamilyHandle &cf,
const Rdb_key_def &kd, uint nlist, uint code_size)
: InvertedLists(nlist, code_size),
m_index_id(index_id),
m_cf(cf),
m_kd(kd) {
use_iterator = true;
}
~Rdb_faiss_inverted_list() override = default;

Rdb_faiss_inverted_list(const Rdb_faiss_inverted_list &) = delete;
Rdb_faiss_inverted_list(Rdb_faiss_inverted_list &&) = delete;
Rdb_faiss_inverted_list &operator=(const Rdb_faiss_inverted_list &) = delete;
Rdb_faiss_inverted_list &operator=(Rdb_faiss_inverted_list &&) = delete;

size_t list_size(size_t list_no) const override {
throw std::runtime_error(std::string("unexpected function call ") +
__PRETTY_FUNCTION__);
Expand All @@ -288,7 +295,7 @@ class Rdb_faiss_inverted_list : public faiss::InvertedLists {
return new Rdb_vector_iterator(
reinterpret_cast<Rdb_faiss_inverted_list_context *>(
inverted_list_context),
m_index_id, m_cf, code_size, list_no);
m_index_id, m_cf, m_kd, code_size, list_no);
}

const uint8_t *get_codes(size_t list_no) const override {
Expand All @@ -314,7 +321,7 @@ class Rdb_faiss_inverted_list : public faiss::InvertedLists {

rocksdb::Slice value_slice(reinterpret_cast<const char *>(code), code_size);
rocksdb::Status status =
context->m_write_batch->Put(m_cf, key_writer.to_slice(), value_slice);
context->m_write_batch->Put(&m_cf, key_writer.to_slice(), value_slice);
if (!status.ok()) {
LogPluginErrMsg(
INFORMATION_LEVEL, ER_LOG_PRINTF_MSG,
Expand Down Expand Up @@ -347,15 +354,19 @@ class Rdb_faiss_inverted_list : public faiss::InvertedLists {

private:
Index_id m_index_id;
rocksdb::ColumnFamilyHandle *const m_cf;
rocksdb::ColumnFamilyHandle &m_cf;
const Rdb_key_def &m_kd;
};

class Rdb_vector_index_ivf : public Rdb_vector_index {
public:
Rdb_vector_index_ivf(const FB_vector_index_config index_def,
std::shared_ptr<rocksdb::ColumnFamilyHandle> cf_handle,
const Index_id index_id)
: m_index_id{index_id}, m_index_def{index_def}, m_cf_handle{cf_handle} {}
const Rdb_key_def &kd, const Index_id index_id)
: m_index_id{index_id},
m_index_def{index_def},
m_cf_handle{cf_handle},
m_kd{kd} {}

virtual ~Rdb_vector_index_ivf() override = default;

Expand Down Expand Up @@ -441,8 +452,8 @@ class Rdb_vector_index_ivf : public Rdb_vector_index {
for (std::size_t i = 0; i < m_list_size_stats.size(); i++) {
std::size_t list_size = 0;
Rdb_faiss_inverted_list_context context(thd);
Rdb_vector_iterator vector_iter(&context, m_index_id, m_cf_handle.get(),
m_index_l2->code_size, i);
Rdb_vector_iterator vector_iter(&context, m_index_id, *m_cf_handle.get(),
m_kd, m_index_l2->code_size, i);
while (vector_iter.is_available()) {
uint rtn = vector_iter.get_pk_and_codes(pk, codes);
if (rtn) {
Expand Down Expand Up @@ -519,7 +530,7 @@ class Rdb_vector_index_ivf : public Rdb_vector_index {

// create inverted list
m_inverted_list = std::make_unique<Rdb_faiss_inverted_list>(
m_index_id, m_cf_handle.get(), m_index_l2->nlist,
m_index_id, *m_cf_handle.get(), m_kd, m_index_l2->nlist,
m_index_l2->code_size);
m_index_l2->replace_invlists(m_inverted_list.get());
m_index_ip->replace_invlists(m_inverted_list.get());
Expand Down Expand Up @@ -585,6 +596,7 @@ class Rdb_vector_index_ivf : public Rdb_vector_index {
Index_id m_index_id;
FB_vector_index_config m_index_def;
std::shared_ptr<rocksdb::ColumnFamilyHandle> m_cf_handle;
const Rdb_key_def &m_kd;
std::atomic<uint> m_hit{0};
std::unique_ptr<faiss::IndexFlatL2> m_quantizer;
std::unique_ptr<faiss::IndexIVF> m_index_l2;
Expand Down Expand Up @@ -679,13 +691,14 @@ uint create_vector_index(Rdb_cmd_srv_helper &cmd_srv_helper,
const std::string &db_name,
const FB_vector_index_config index_def,
std::shared_ptr<rocksdb::ColumnFamilyHandle> cf_handle,
const Rdb_key_def &kd,
const Index_id index_id,
std::unique_ptr<Rdb_vector_index> &index) {
if (index_def.type() == FB_VECTOR_INDEX_TYPE::FLAT ||
index_def.type() == FB_VECTOR_INDEX_TYPE::IVFFLAT ||
index_def.type() == FB_VECTOR_INDEX_TYPE::IVFPQ) {
index =
std::make_unique<Rdb_vector_index_ivf>(index_def, cf_handle, index_id);
std::make_unique<Rdb_vector_index_ivf>(index_def, cf_handle, kd, index_id);
} else {
assert(false);
return HA_ERR_UNSUPPORTED;
Expand All @@ -696,14 +709,11 @@ uint create_vector_index(Rdb_cmd_srv_helper &cmd_srv_helper,
#else

// dummy implementation for non-fbvectordb builds
uint create_vector_index(Rdb_cmd_srv_helper &cmd_srv_helper [[maybe_unused]],
const std::string &db_name [[maybe_unused]],
const FB_vector_index_config index_def
[[maybe_unused]],
std::shared_ptr<rocksdb::ColumnFamilyHandle> cf_handle
[[maybe_unused]],
const Index_id index_id [[maybe_unused]],
std::unique_ptr<Rdb_vector_index> &index) {
uint create_vector_index(Rdb_cmd_srv_helper &, const std::string &,
const FB_vector_index_config,
std::shared_ptr<rocksdb::ColumnFamilyHandle>,
const Rdb_key_def, const Index_id,
std::unique_ptr<Rdb_vector_index> &) {
index = nullptr;
return HA_ERR_UNSUPPORTED;
}
Expand Down
4 changes: 3 additions & 1 deletion storage/rocksdb/rdb_vector_db.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,8 @@

namespace myrocks {

class Rdb_key_def;

/** for infomation schema */
class Rdb_vector_index_info {
public:
Expand Down Expand Up @@ -112,7 +114,7 @@ uint create_vector_index(Rdb_cmd_srv_helper &cmd_srv_helper,
const std::string &db_name,
const FB_vector_index_config index_def,
std::shared_ptr<rocksdb::ColumnFamilyHandle> cf_handle,
const Index_id index_id,
const Rdb_key_def &kd, const Index_id index_id,
std::unique_ptr<Rdb_vector_index> &index);

/**
Expand Down

0 comments on commit 1927fc6

Please sign in to comment.