diff --git a/storage/rocksdb/rdb_datadic.cc b/storage/rocksdb/rdb_datadic.cc index be9a785469ce..7c136228de6a 100644 --- a/storage/rocksdb/rdb_datadic.cc +++ b/storage/rocksdb/rdb_datadic.cc @@ -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 diff --git a/storage/rocksdb/rdb_vector_db.cc b/storage/rocksdb/rdb_vector_db.cc index a04dafd4f29b..e0080be53dad 100644 --- a/storage/rocksdb/rdb_vector_db.cc +++ b/storage/rocksdb/rdb_vector_db.cc @@ -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), @@ -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(); } @@ -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__); @@ -288,7 +295,7 @@ class Rdb_faiss_inverted_list : public faiss::InvertedLists { return new Rdb_vector_iterator( reinterpret_cast( 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 { @@ -314,7 +321,7 @@ class Rdb_faiss_inverted_list : public faiss::InvertedLists { rocksdb::Slice value_slice(reinterpret_cast(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, @@ -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 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; @@ -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) { @@ -519,7 +530,7 @@ class Rdb_vector_index_ivf : public Rdb_vector_index { // create inverted list m_inverted_list = std::make_unique( - 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()); @@ -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 m_cf_handle; + const Rdb_key_def &m_kd; std::atomic m_hit{0}; std::unique_ptr m_quantizer; std::unique_ptr m_index_l2; @@ -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 cf_handle, + const Rdb_key_def &kd, const Index_id index_id, std::unique_ptr &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(index_def, cf_handle, index_id); + std::make_unique(index_def, cf_handle, kd, index_id); } else { assert(false); return HA_ERR_UNSUPPORTED; @@ -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 cf_handle - [[maybe_unused]], - const Index_id index_id [[maybe_unused]], - std::unique_ptr &index) { +uint create_vector_index(Rdb_cmd_srv_helper &, const std::string &, + const FB_vector_index_config, + std::shared_ptr, + const Rdb_key_def, const Index_id, + std::unique_ptr &) { index = nullptr; return HA_ERR_UNSUPPORTED; } diff --git a/storage/rocksdb/rdb_vector_db.h b/storage/rocksdb/rdb_vector_db.h index d710f857b328..0a0b60470116 100644 --- a/storage/rocksdb/rdb_vector_db.h +++ b/storage/rocksdb/rdb_vector_db.h @@ -30,6 +30,8 @@ namespace myrocks { +class Rdb_key_def; + /** for infomation schema */ class Rdb_vector_index_info { public: @@ -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 cf_handle, - const Index_id index_id, + const Rdb_key_def &kd, const Index_id index_id, std::unique_ptr &index); /**