Skip to content

Commit 1927fc6

Browse files
laurynas-biveinislaurynas-biveinis
committed
TODO: PR this separately
1 parent 1af3e9f commit 1927fc6

File tree

3 files changed

+41
-29
lines changed

3 files changed

+41
-29
lines changed

‎storage/rocksdb/rdb_datadic.cc‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3557,8 +3557,8 @@ uint Rdb_key_def::setup_vector_index(const TABLE &tbl,
35573557
}
35583558

35593559
return create_vector_index(cmd_srv_helper, tbl_def.base_dbname(),
3560-
m_vector_index_config, m_cf_handle, m_index_number,
3561-
m_vector_index);
3560+
m_vector_index_config, m_cf_handle, *this,
3561+
m_index_number, m_vector_index);
35623562
}
35633563

35643564
// See Rdb_charset_space_info::spaces_xfrm

‎storage/rocksdb/rdb_vector_db.cc‎

Lines changed: 36 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -175,8 +175,8 @@ class Rdb_faiss_inverted_list_write_context {
175175
class Rdb_vector_iterator : public faiss::InvertedListsIterator {
176176
public:
177177
Rdb_vector_iterator(Rdb_faiss_inverted_list_context *context,
178-
Index_id index_id, rocksdb::ColumnFamilyHandle *const cf,
179-
const uint code_size, const size_t list_id)
178+
Index_id index_id, rocksdb::ColumnFamilyHandle &cf,
179+
const Rdb_key_def &kd, uint code_size, size_t list_id)
180180
: m_context(context),
181181
m_index_id(index_id),
182182
m_list_id(list_id),
@@ -188,10 +188,10 @@ class Rdb_vector_iterator : public faiss::InvertedListsIterator {
188188
Rdb_string_writer upper_key_writer;
189189
write_inverted_list_key(upper_key_writer, index_id, list_id + 1);
190190
m_iterator_upper_bound_key.PinSelf(upper_key_writer.to_slice());
191-
m_iterator.reset(rdb_tx_get_iterator(
192-
context->m_thd, cf, /* skip_bloom_filter */ true,
191+
m_iterator = rdb_tx_get_iterator(
192+
context->m_thd, cf, kd, /* skip_bloom_filter */ true,
193193
m_iterator_lower_bound_key, m_iterator_upper_bound_key,
194-
/* snapshot */ nullptr, TABLE_TYPE::USER_TABLE));
194+
/* snapshot */ nullptr, TABLE_TYPE::USER_TABLE);
195195
m_iterator->SeekToFirst();
196196
}
197197

@@ -268,14 +268,21 @@ class Rdb_vector_iterator : public faiss::InvertedListsIterator {
268268
*/
269269
class Rdb_faiss_inverted_list : public faiss::InvertedLists {
270270
public:
271-
Rdb_faiss_inverted_list(Index_id index_id,
272-
rocksdb::ColumnFamilyHandle *const cf, uint nlist,
273-
uint code_size)
274-
: InvertedLists(nlist, code_size), m_index_id(index_id), m_cf(cf) {
271+
Rdb_faiss_inverted_list(Index_id index_id, rocksdb::ColumnFamilyHandle &cf,
272+
const Rdb_key_def &kd, uint nlist, uint code_size)
273+
: InvertedLists(nlist, code_size),
274+
m_index_id(index_id),
275+
m_cf(cf),
276+
m_kd(kd) {
275277
use_iterator = true;
276278
}
277279
~Rdb_faiss_inverted_list() override = default;
278280

281+
Rdb_faiss_inverted_list(const Rdb_faiss_inverted_list &) = delete;
282+
Rdb_faiss_inverted_list(Rdb_faiss_inverted_list &&) = delete;
283+
Rdb_faiss_inverted_list &operator=(const Rdb_faiss_inverted_list &) = delete;
284+
Rdb_faiss_inverted_list &operator=(Rdb_faiss_inverted_list &&) = delete;
285+
279286
size_t list_size(size_t list_no) const override {
280287
throw std::runtime_error(std::string("unexpected function call ") +
281288
__PRETTY_FUNCTION__);
@@ -288,7 +295,7 @@ class Rdb_faiss_inverted_list : public faiss::InvertedLists {
288295
return new Rdb_vector_iterator(
289296
reinterpret_cast<Rdb_faiss_inverted_list_context *>(
290297
inverted_list_context),
291-
m_index_id, m_cf, code_size, list_no);
298+
m_index_id, m_cf, m_kd, code_size, list_no);
292299
}
293300

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

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

348355
private:
349356
Index_id m_index_id;
350-
rocksdb::ColumnFamilyHandle *const m_cf;
357+
rocksdb::ColumnFamilyHandle &m_cf;
358+
const Rdb_key_def &m_kd;
351359
};
352360

353361
class Rdb_vector_index_ivf : public Rdb_vector_index {
354362
public:
355363
Rdb_vector_index_ivf(const FB_vector_index_config index_def,
356364
std::shared_ptr<rocksdb::ColumnFamilyHandle> cf_handle,
357-
const Index_id index_id)
358-
: m_index_id{index_id}, m_index_def{index_def}, m_cf_handle{cf_handle} {}
365+
const Rdb_key_def &kd, const Index_id index_id)
366+
: m_index_id{index_id},
367+
m_index_def{index_def},
368+
m_cf_handle{cf_handle},
369+
m_kd{kd} {}
359370

360371
virtual ~Rdb_vector_index_ivf() override = default;
361372

@@ -441,8 +452,8 @@ class Rdb_vector_index_ivf : public Rdb_vector_index {
441452
for (std::size_t i = 0; i < m_list_size_stats.size(); i++) {
442453
std::size_t list_size = 0;
443454
Rdb_faiss_inverted_list_context context(thd);
444-
Rdb_vector_iterator vector_iter(&context, m_index_id, m_cf_handle.get(),
445-
m_index_l2->code_size, i);
455+
Rdb_vector_iterator vector_iter(&context, m_index_id, *m_cf_handle.get(),
456+
m_kd, m_index_l2->code_size, i);
446457
while (vector_iter.is_available()) {
447458
uint rtn = vector_iter.get_pk_and_codes(pk, codes);
448459
if (rtn) {
@@ -519,7 +530,7 @@ class Rdb_vector_index_ivf : public Rdb_vector_index {
519530

520531
// create inverted list
521532
m_inverted_list = std::make_unique<Rdb_faiss_inverted_list>(
522-
m_index_id, m_cf_handle.get(), m_index_l2->nlist,
533+
m_index_id, *m_cf_handle.get(), m_kd, m_index_l2->nlist,
523534
m_index_l2->code_size);
524535
m_index_l2->replace_invlists(m_inverted_list.get());
525536
m_index_ip->replace_invlists(m_inverted_list.get());
@@ -585,6 +596,7 @@ class Rdb_vector_index_ivf : public Rdb_vector_index {
585596
Index_id m_index_id;
586597
FB_vector_index_config m_index_def;
587598
std::shared_ptr<rocksdb::ColumnFamilyHandle> m_cf_handle;
599+
const Rdb_key_def &m_kd;
588600
std::atomic<uint> m_hit{0};
589601
std::unique_ptr<faiss::IndexFlatL2> m_quantizer;
590602
std::unique_ptr<faiss::IndexIVF> m_index_l2;
@@ -679,13 +691,14 @@ uint create_vector_index(Rdb_cmd_srv_helper &cmd_srv_helper,
679691
const std::string &db_name,
680692
const FB_vector_index_config index_def,
681693
std::shared_ptr<rocksdb::ColumnFamilyHandle> cf_handle,
694+
const Rdb_key_def &kd,
682695
const Index_id index_id,
683696
std::unique_ptr<Rdb_vector_index> &index) {
684697
if (index_def.type() == FB_VECTOR_INDEX_TYPE::FLAT ||
685698
index_def.type() == FB_VECTOR_INDEX_TYPE::IVFFLAT ||
686699
index_def.type() == FB_VECTOR_INDEX_TYPE::IVFPQ) {
687700
index =
688-
std::make_unique<Rdb_vector_index_ivf>(index_def, cf_handle, index_id);
701+
std::make_unique<Rdb_vector_index_ivf>(index_def, cf_handle, kd, index_id);
689702
} else {
690703
assert(false);
691704
return HA_ERR_UNSUPPORTED;
@@ -696,14 +709,11 @@ uint create_vector_index(Rdb_cmd_srv_helper &cmd_srv_helper,
696709
#else
697710

698711
// dummy implementation for non-fbvectordb builds
699-
uint create_vector_index(Rdb_cmd_srv_helper &cmd_srv_helper [[maybe_unused]],
700-
const std::string &db_name [[maybe_unused]],
701-
const FB_vector_index_config index_def
702-
[[maybe_unused]],
703-
std::shared_ptr<rocksdb::ColumnFamilyHandle> cf_handle
704-
[[maybe_unused]],
705-
const Index_id index_id [[maybe_unused]],
706-
std::unique_ptr<Rdb_vector_index> &index) {
712+
uint create_vector_index(Rdb_cmd_srv_helper &, const std::string &,
713+
const FB_vector_index_config,
714+
std::shared_ptr<rocksdb::ColumnFamilyHandle>,
715+
const Rdb_key_def, const Index_id,
716+
std::unique_ptr<Rdb_vector_index> &) {
707717
index = nullptr;
708718
return HA_ERR_UNSUPPORTED;
709719
}

‎storage/rocksdb/rdb_vector_db.h‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,8 @@
3030

3131
namespace myrocks {
3232

33+
class Rdb_key_def;
34+
3335
/** for infomation schema */
3436
class Rdb_vector_index_info {
3537
public:
@@ -112,7 +114,7 @@ uint create_vector_index(Rdb_cmd_srv_helper &cmd_srv_helper,
112114
const std::string &db_name,
113115
const FB_vector_index_config index_def,
114116
std::shared_ptr<rocksdb::ColumnFamilyHandle> cf_handle,
115-
const Index_id index_id,
117+
const Rdb_key_def &kd, const Index_id index_id,
116118
std::unique_ptr<Rdb_vector_index> &index);
117119

118120
/**

0 commit comments

Comments
 (0)