vector_mhnsw.cc 42.1 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18
/*
   Copyright (c) 2024, MariaDB plc

   This program is free software; you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation; version 2 of the License.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program; if not, write to the Free Software
   Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1335  USA
*/

#include <my_global.h>
19
#include "key.h"                                // key_copy()
20
#include "create_options.h"
21
#include "vector_mhnsw.h"
22
#include "item_vectorfunc.h"
Sergei Golubchik's avatar
Sergei Golubchik committed
23
#include <scope.h>
24 25 26
#include <my_atomic_wrapper.h>
#include "bloom_filters.h"

27 28
// Algorithm parameters
static constexpr float alpha = 1.1f;
29
static constexpr float generosity = 1.1f;
30 31
static constexpr uint ef_construction= 10;

32 33 34 35 36 37 38 39 40 41 42 43 44 45 46
static ulonglong mhnsw_cache_size;
static MYSQL_SYSVAR_ULONGLONG(cache_size, mhnsw_cache_size,
       PLUGIN_VAR_RQCMDARG, "Size of the cache for the MHNSW vector index",
       nullptr, nullptr, 16*1024*1024, 1024*1024, SIZE_T_MAX, 1);
static MYSQL_THDVAR_UINT(min_limit, PLUGIN_VAR_RQCMDARG,
       "Defines the minimal number of result candidates to look for in the "
       "vector index for ORDER BY ... LIMIT N queries. The search will never "
       "search for less rows than that, even if LIMIT is smaller. "
       "This notably improves the search quality at low LIMIT values, "
       "at the expense of search time", nullptr, nullptr, 20, 1, 65535, 1);
static MYSQL_THDVAR_UINT(max_edges_per_node, PLUGIN_VAR_RQCMDARG,
       "Larger values means slower INSERT, larger index size and higher "
       "memory consumption, but better search results",
       nullptr, nullptr, 6, 3, 200, 1);

Sergei Golubchik's avatar
Sergei Golubchik committed
47 48 49 50 51 52 53
enum metric_type : uint { EUCLIDEAN, COSINE };
static const char *distance_function_names[]= { "euclidean", "cosine", nullptr };
static TYPELIB distance_functions= CREATE_TYPELIB_FOR(distance_function_names);
static MYSQL_THDVAR_ENUM(distance_function, PLUGIN_VAR_RQCMDARG,
       "Distance function to build the vector index for",
       nullptr, nullptr, EUCLIDEAN, &distance_functions);

54 55 56
struct ha_index_option_struct
{
  ulonglong M; // option struct does not support uint
Sergei Golubchik's avatar
Sergei Golubchik committed
57
  metric_type metric;
58 59
};

60 61 62 63
enum Graph_table_fields {
  FIELD_LAYER, FIELD_TREF, FIELD_VEC, FIELD_NEIGHBORS
};
enum Graph_table_indices {
Sergei Golubchik's avatar
Sergei Golubchik committed
64
  IDX_TREF, IDX_LAYER
65 66
};

67
class MHNSW_Context;
68 69 70
class FVectorNode;

/*
71
  One vector, an array of coordinates in ctx->vec_len dimensions
72
*/
73 74
#pragma pack(push, 1)
struct FVector
75
{
76 77 78 79 80 81 82 83 84 85 86 87 88 89
  static constexpr size_t data_header= sizeof(float);
  static constexpr size_t alloc_header= data_header + sizeof(float);

  float abs2, scale;
  int16_t dims[4];

  uchar *data() const { return (uchar*)(&scale); }

  static size_t data_size(size_t n)
  { return data_header + n*2; }

  static size_t data_to_value_size(size_t data_size)
  { return (data_size - data_header)*2; }

Sergei Golubchik's avatar
Sergei Golubchik committed
90
  static const FVector *create(metric_type metric, void *mem, const void *src, size_t src_len)
91 92 93 94 95 96 97 98 99 100 101 102
  {
    float scale=0, *v= (float *)src;
    size_t vec_len= src_len / sizeof(float);
    for (size_t i= 0; i < vec_len; i++)
      if (std::abs(scale) < std::abs(get_float(v + i)))
        scale= get_float(v + i);

    FVector *vec= align_ptr(mem);
    vec->scale= scale ? scale/32767 : 1;
    for (size_t i= 0; i < vec_len; i++)
      vec->dims[i] = static_cast<int16_t>(std::round(get_float(v + i) / vec->scale));
    vec->postprocess(vec_len);
Sergei Golubchik's avatar
Sergei Golubchik committed
103 104 105 106 107 108
    if (metric == COSINE)
    {
      if (vec->abs2 > 0.0f)
        vec->scale/= std::sqrt(vec->abs2);
      vec->abs2= 1.0f;
    }
109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154
    return vec;
  }

  void postprocess(size_t vec_len)
  {
    fix_tail(vec_len);
    abs2= scale * scale * dot_product(dims, dims, vec_len) / 2;
  }

#ifdef AVX2_IMPLEMENTATION
  /************* AVX2 *****************************************************/
  static constexpr size_t AVX2_bytes= 256/8;
  static constexpr size_t AVX2_dims= AVX2_bytes/sizeof(int16_t);

  AVX2_IMPLEMENTATION
  static float dot_product(const int16_t *v1, const int16_t *v2, size_t len)
  {
    typedef float v8f __attribute__((vector_size(AVX2_bytes)));
    union { v8f v; __m256 i; } tmp;
    __m256i *p1= (__m256i*)v1;
    __m256i *p2= (__m256i*)v2;
    v8f d= {0};
    for (size_t i= 0; i < (len + AVX2_dims-1)/AVX2_dims; p1++, p2++, i++)
    {
      tmp.i= _mm256_cvtepi32_ps(_mm256_madd_epi16(*p1, *p2));
      d+= tmp.v;
    }
    return d[0] + d[1] + d[2] + d[3] + d[4] + d[5] + d[6] + d[7];
  }

  AVX2_IMPLEMENTATION
  static size_t alloc_size(size_t n)
  { return alloc_header + MY_ALIGN(n*2, AVX2_bytes) + AVX2_bytes - 1; }

  AVX2_IMPLEMENTATION
  static FVector *align_ptr(void *ptr)
  { return (FVector*)(MY_ALIGN(((intptr)ptr) + alloc_header, AVX2_bytes)
                      - alloc_header); }

  AVX2_IMPLEMENTATION
  void fix_tail(size_t vec_len)
  {
    bzero(dims + vec_len, (MY_ALIGN(vec_len, AVX2_dims) - vec_len)*2);
  }
#endif

Sergei Golubchik's avatar
Sergei Golubchik committed
155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186
#ifdef AVX512_IMPLEMENTATION
  /************* AVX512 ****************************************************/
  static constexpr size_t AVX512_bytes= 512/8;
  static constexpr size_t AVX512_dims= AVX512_bytes/sizeof(int16_t);

  AVX512_IMPLEMENTATION
  static float dot_product(const int16_t *v1, const int16_t *v2, size_t len)
  {
    __m512i *p1= (__m512i*)v1;
    __m512i *p2= (__m512i*)v2;
    __m512 d= _mm512_setzero_ps();
    for (size_t i= 0; i < (len + AVX512_dims-1)/AVX512_dims; p1++, p2++, i++)
      d= _mm512_add_ps(d, _mm512_cvtepi32_ps(_mm512_madd_epi16(*p1, *p2)));
    return _mm512_reduce_add_ps(d);
  }

  AVX512_IMPLEMENTATION
  static size_t alloc_size(size_t n)
  { return alloc_header + MY_ALIGN(n*2, AVX512_bytes) + AVX512_bytes - 1; }

  AVX512_IMPLEMENTATION
  static FVector *align_ptr(void *ptr)
  { return (FVector*)(MY_ALIGN(((intptr)ptr) + alloc_header, AVX512_bytes)
                      - alloc_header); }

  AVX512_IMPLEMENTATION
  void fix_tail(size_t vec_len)
  {
    bzero(dims + vec_len, (MY_ALIGN(vec_len, AVX512_dims) - vec_len)*2);
  }
#endif

187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210
  /************* no-SIMD default ******************************************/
  DEFAULT_IMPLEMENTATION
  static float dot_product(const int16_t *v1, const int16_t *v2, size_t len)
  {
    int64_t d= 0;
    for (size_t i= 0; i < len; i++)
      d+= int32_t(v1[i]) * int32_t(v2[i]);
    return static_cast<float>(d);
  }

  DEFAULT_IMPLEMENTATION
  static size_t alloc_size(size_t n) { return alloc_header + n*2; }

  DEFAULT_IMPLEMENTATION
  static FVector *align_ptr(void *ptr) { return (FVector*)ptr; }

  DEFAULT_IMPLEMENTATION
  void fix_tail(size_t) {  }

  float distance_to(const FVector *other, size_t vec_len) const
  {
    return abs2 + other->abs2 - scale * other->scale *
           dot_product(dims, other->dims, vec_len);
  }
211
};
212
#pragma pack(pop)
213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238

/*
  An array of pointers to graph nodes

  It's mainly used to store all neighbors of a given node on a given layer.

  An array is fixed size, 2*M for the zero layer, M for other layers
  see MHNSW_Context::max_neighbors().

  Number of neighbors is zero-padded to multiples of 8 (for SIMD Bloom filter).

  Also used as a simply array of nodes in search_layer, the array size
  then is defined by ef or efConstruction.
*/
struct Neighborhood: public Sql_alloc
{
  FVectorNode **links;
  size_t num;
  FVectorNode **init(FVectorNode **ptr, size_t n)
  {
    num= 0;
    links= ptr;
    n= MY_ALIGN(n, 8);
    bzero(ptr, n*sizeof(*ptr));
    return ptr + n;
  }
239 240
};

241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263

/*
  One node in a graph = one row in the graph table

  stores a vector itself, ref (= position) in the graph (= hlindex)
  table, a ref in the main table, and an array of Neighborhood's, one
  per layer.

  It's lazily initialized, may know only gref, everything else is
  loaded on demand.

  On the other hand, on INSERT the new node knows everything except
  gref - which only becomes known after ha_write_row.

  Allocated on memroot in two chunks. One is the same size for all nodes
  and stores FVectorNode object, gref, tref, and vector. The second
  stores neighbors, all Neighborhood's together, its size depends
  on the number of layers this node is on.

  There can be millions of nodes in the cache and the cache size
  is constrained by mhnsw_cache_size, so every byte matters here
*/
#pragma pack(push, 1)
264
class FVectorNode
265 266
{
private:
267
  MHNSW_Context *ctx;
268

269
  const FVector *make_vec(const void *v);
270
  int alloc_neighborhood(uint8_t layer);
271
public:
272
  const FVector *vec= nullptr;
273 274
  Neighborhood *neighbors= nullptr;
  uint8_t max_layer;
275
  bool stored:1, deleted:1;
276 277 278 279

  FVectorNode(MHNSW_Context *ctx_, const void *gref_);
  FVectorNode(MHNSW_Context *ctx_, const void *tref_, uint8_t layer,
              const void *vec_);
280
  float distance_to(const FVector *other) const;
281 282 283 284 285 286 287 288
  int load(TABLE *graph);
  int load_from_record(TABLE *graph);
  int save(TABLE *graph);
  size_t tref_len() const;
  size_t gref_len() const;
  uchar *gref() const;
  uchar *tref() const;
  void push_neighbor(size_t layer, FVectorNode *v);
289 290

  static uchar *get_key(const FVectorNode *elem, size_t *key_len, my_bool);
291
};
292
#pragma pack(pop)
293

294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311
/*
  Shared algorithm context. The graph.

  Stored in TABLE_SHARE and on TABLE_SHARE::mem_root.
  Stores the complete graph in MHNSW_Context::root,
  The mapping gref->FVectorNode is in the node_cache.
  Both root and node_cache are protected by a cache_lock, but it's
  needed when loading nodes and is not used when the whole graph is in memory.
  Graph can be traversed concurrently by different threads, as traversal
  changes neither nodes nor the ctx.
  Nodes can be loaded concurrently by different threads, this is protected
  by a partitioned node_lock.
  reference counter allows flushing the graph without interrupting
  concurrent searches.
  MyISAM automatically gets exclusive write access because of the TL_WRITE,
  but InnoDB has to use a dedicated ctx->commit_lock for that
*/
class MHNSW_Context : public Sql_alloc
312
{
313 314 315 316 317 318 319 320 321 322 323 324
  std::atomic<uint> refcnt{0};
  mysql_mutex_t cache_lock;
  mysql_mutex_t node_lock[8];

  void cache_internal(FVectorNode *node)
  {
    DBUG_ASSERT(node->stored);
    node_cache.insert(node);
  }
  void *alloc_node_internal()
  {
    return alloc_root(&root, sizeof(FVectorNode) + gref_len + tref_len
325
                      + FVector::alloc_size(vec_len));
326 327 328
  }

protected:
329
  MEM_ROOT root;
330 331 332 333
  Hash_set<FVectorNode> node_cache{PSI_INSTRUMENT_MEM, FVectorNode::get_key};

public:
  mysql_rwlock_t commit_lock;
334
  size_t vec_len= 0;
335
  size_t byte_len= 0;
336 337 338 339 340
  Atomic_relaxed<double> ef_power{0.6}; // for the bloom filter size heuristic
  FVectorNode *start= 0;
  const uint tref_len;
  const uint gref_len;
  const uint M;
Sergei Golubchik's avatar
Sergei Golubchik committed
341
  metric_type metric;
342 343 344 345

  MHNSW_Context(TABLE *t)
    : tref_len(t->file->ref_length),
      gref_len(t->hlindex->file->ref_length),
Sergei Golubchik's avatar
Sergei Golubchik committed
346 347
      M(static_cast<uint>(t->s->key_info[t->s->keys].option_struct->M)),
      metric(t->s->key_info[t->s->keys].option_struct->metric)
348 349 350 351 352 353 354
  {
    mysql_rwlock_init(PSI_INSTRUMENT_ME, &commit_lock);
    mysql_mutex_init(PSI_INSTRUMENT_ME, &cache_lock, MY_MUTEX_INIT_FAST);
    for (uint i=0; i < array_elements(node_lock); i++)
      mysql_mutex_init(PSI_INSTRUMENT_ME, node_lock + i, MY_MUTEX_INIT_SLOW);
    init_alloc_root(PSI_INSTRUMENT_MEM, &root, 1024*1024, 0, MYF(0));
  }
355

356 357 358 359 360 361 362 363
  virtual ~MHNSW_Context()
  {
    free_root(&root, MYF(0));
    mysql_rwlock_destroy(&commit_lock);
    mysql_mutex_destroy(&cache_lock);
    for (size_t i=0; i < array_elements(node_lock); i++)
      mysql_mutex_destroy(node_lock + i);
  }
364

365
  uint lock_node(FVectorNode *ptr)
366
  {
367 368 369 370 371
    ulong nr1= 1, nr2= 4;
    my_hash_sort_bin(0, (uchar*)&ptr, sizeof(ptr), &nr1, &nr2);
    uint ticket= nr1 % array_elements(node_lock);
    mysql_mutex_lock(node_lock + ticket);
    return ticket;
372
  }
373

374
  void unlock_node(uint ticket)
375
  {
376 377 378 379 380 381
    mysql_mutex_unlock(node_lock + ticket);
  }

  uint max_neighbors(size_t layer) const
  {
    return (layer ? 1 : 2) * M; // heuristic from the paper
382 383
  }

384 385 386
  void set_lengths(size_t len)
  {
    byte_len= len;
387
    vec_len= len / sizeof(float);
388
  }
389 390 391 392

  static int acquire(MHNSW_Context **ctx, TABLE *table, bool for_update);
  static MHNSW_Context *get_from_share(TABLE_SHARE *share, TABLE *table);

Sergei Golubchik's avatar
Sergei Golubchik committed
393
  virtual void reset(TABLE_SHARE *share)
394
  {
395
    share->lock_share();
396 397 398 399 400
    if (static_cast<MHNSW_Context*>(share->hlindex->hlindex_data) == this)
    {
      share->hlindex->hlindex_data= nullptr;
      --refcnt;
    }
401
    share->unlock_share();
402 403 404 405 406 407 408 409 410 411 412 413
  }

  void release(TABLE *table)
  {
    return release(table->file->has_transactions(), table->s);
  }

  virtual void release(bool can_commit, TABLE_SHARE *share)
  {
    if (can_commit)
      mysql_rwlock_unlock(&commit_lock);
    if (root_size(&root) > mhnsw_cache_size)
Sergei Golubchik's avatar
Sergei Golubchik committed
414
      reset(share);
415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465
    if (--refcnt == 0)
      this->~MHNSW_Context(); // XXX reuse
  }

  FVectorNode *get_node(const void *gref)
  {
    mysql_mutex_lock(&cache_lock);
    FVectorNode *node= node_cache.find(gref, gref_len);
    if (!node)
    {
      node= new (alloc_node_internal()) FVectorNode(this, gref);
      cache_internal(node);
    }
    mysql_mutex_unlock(&cache_lock);
    return node;
  }

  /* used on INSERT, gref isn't known, so cannot cache the node yet */
  void *alloc_node()
  {
    mysql_mutex_lock(&cache_lock);
    auto p= alloc_node_internal();
    mysql_mutex_unlock(&cache_lock);
    return p;
  }

  /* explicitly cache the node after alloc_node() */
  void cache_node(FVectorNode *node)
  {
    mysql_mutex_lock(&cache_lock);
    cache_internal(node);
    mysql_mutex_unlock(&cache_lock);
  }

  /* find the node without creating, only used on merging trx->ctx */
  FVectorNode *find_node(const void *gref)
  {
    mysql_mutex_lock(&cache_lock);
    FVectorNode *node= node_cache.find(gref, gref_len);
    mysql_mutex_unlock(&cache_lock);
    return node;
  }

  void *alloc_neighborhood(size_t max_layer)
  {
    mysql_mutex_lock(&cache_lock);
    auto p= alloc_root(&root, sizeof(Neighborhood)*(max_layer+1) +
             sizeof(FVectorNode*)*(MY_ALIGN(M, 4)*2 + MY_ALIGN(M,8)*max_layer));
    mysql_mutex_unlock(&cache_lock);
    return p;
  }
466
};
467

468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485
/*
  This is a non-shared context that exists within one transaction.

  At the end of the transaction it's either discarded (on rollback)
  or merged into the shared ctx (on commit).

  trx's are stored in thd->ha_data[] in a single-linked list,
  one instance of trx per TABLE_SHARE and allocated on the
  thd->transaction->mem_root
*/
class MHNSW_Trx : public MHNSW_Context
{
public:
  TABLE_SHARE *table_share;
  bool list_of_nodes_is_lost= false;
  MHNSW_Trx *next= nullptr;

  MHNSW_Trx(TABLE *table) : MHNSW_Context(table), table_share(table->s) {}
Sergei Golubchik's avatar
Sergei Golubchik committed
486
  void reset(TABLE_SHARE *) override
487 488 489 490 491 492 493 494 495
  {
    node_cache.clear();
    free_root(&root, MYF(0));
    start= 0;
    list_of_nodes_is_lost= true;
  }
  void release(bool, TABLE_SHARE *) override
  {
    if (root_size(&root) > mhnsw_cache_size)
Sergei Golubchik's avatar
Sergei Golubchik committed
496
      reset(nullptr);
497 498
  }

499
  static MHNSW_Trx *get_from_thd(TABLE *table, bool for_update);
500 501 502 503

  // it's okay in a transaction-local cache, there's no concurrent access
  Hash_set<FVectorNode> &get_cache() { return node_cache; }

504 505 506 507
  static transaction_participant tp;
  static int do_commit(THD *thd, bool);
  static int do_savepoint_rollback(THD *thd, void *);
  static int do_rollback(THD *thd, bool);
508 509
};

510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525
struct transaction_participant MHNSW_Trx::tp=
{
  0, 0, 0,
  nullptr,                        /* close_connection */
  [](THD *, void *){ return 0; }, /* savepoint_set */
  MHNSW_Trx::do_savepoint_rollback,
  [](THD *thd){ return true; },   /*savepoint_rollback_can_release_mdl*/
  nullptr,                        /*savepoint_release*/
  MHNSW_Trx::do_commit, MHNSW_Trx::do_rollback,
  nullptr,                        /* prepare */
  nullptr,                        /* recover */
  nullptr, nullptr,               /* commit/rollback_by_xid */
  nullptr, nullptr,               /* recover_rollback_by_xid/recovery_done */
  nullptr, nullptr, nullptr,      /* snapshot, commit/prepare_ordered */
  nullptr, nullptr                /* checkpoint, versioned */
};
526

527
int MHNSW_Trx::do_savepoint_rollback(THD *thd, void *)
528
{
529
  for (auto trx= static_cast<MHNSW_Trx*>(thd_get_ha_data(thd, &tp));
530
       trx; trx= trx->next)
Sergei Golubchik's avatar
Sergei Golubchik committed
531
    trx->reset(nullptr);
532 533 534
  return 0;
}

535
int MHNSW_Trx::do_rollback(THD *thd, bool)
536 537
{
  MHNSW_Trx *trx_next;
538
  for (auto trx= static_cast<MHNSW_Trx*>(thd_get_ha_data(thd, &tp));
539 540 541 542 543
       trx; trx= trx_next)
  {
    trx_next= trx->next;
    trx->~MHNSW_Trx();
  }
544
  thd_set_ha_data(current_thd, &tp, nullptr);
545 546 547
  return 0;
}

548
int MHNSW_Trx::do_commit(THD *thd, bool)
549 550
{
  MHNSW_Trx *trx_next;
551
  for (auto trx= static_cast<MHNSW_Trx*>(thd_get_ha_data(thd, &tp));
552 553 554 555 556 557 558 559
       trx; trx= trx_next)
  {
    trx_next= trx->next;
    auto ctx= MHNSW_Context::get_from_share(trx->table_share, nullptr);
    if (ctx)
    {
      mysql_rwlock_wrlock(&ctx->commit_lock);
      if (trx->list_of_nodes_is_lost)
Sergei Golubchik's avatar
Sergei Golubchik committed
560
        ctx->reset(trx->table_share);
561 562 563 564 565 566 567 568 569 570 571 572 573 574
      else
      {
        // consider copying nodes from trx to shared cache when it makes sense
        // for ann_benchmarks it does not
        // also, consider flushing only changed nodes (a flag in the node)
        for (FVectorNode &from : trx->get_cache())
          if (FVectorNode *node= ctx->find_node(from.gref()))
            node->vec= nullptr;
        ctx->start= nullptr;
      }
      ctx->release(true, trx->table_share);
    }
    trx->~MHNSW_Trx();
  }
575
  thd_set_ha_data(current_thd, &tp, nullptr);
576 577 578
  return 0;
}

579
MHNSW_Trx *MHNSW_Trx::get_from_thd(TABLE *table, bool for_update)
580
{
581 582 583 584 585 586 587 588
  if (!table->file->has_transactions())
      return NULL;

  THD *thd= table->in_use;
  auto trx= static_cast<MHNSW_Trx*>(thd_get_ha_data(thd, &tp));
  if (!for_update && !trx)
    return NULL;

589 590 591 592
  while (trx && trx->table_share != table->s) trx= trx->next;
  if (!trx)
  {
    trx= new (&thd->transaction->mem_root) MHNSW_Trx(table);
593 594
    trx->next= static_cast<MHNSW_Trx*>(thd_get_ha_data(thd, &tp));
    thd_set_ha_data(thd, &tp, trx);
595 596 597
    if (!trx->next)
    {
      bool all= thd_test_options(thd, OPTION_NOT_AUTOCOMMIT | OPTION_BEGIN);
598
      trans_register_ha(thd, all, &tp, 0);
599 600 601 602 603 604
    }
  }
  return trx;
}

MHNSW_Context *MHNSW_Context::get_from_share(TABLE_SHARE *share, TABLE *table)
605
{
606
  share->lock_share();
607 608 609 610 611 612 613 614 615 616
  auto ctx= static_cast<MHNSW_Context*>(share->hlindex->hlindex_data);
  if (!ctx && table)
  {
    ctx= new (&share->hlindex->mem_root) MHNSW_Context(table);
    if (!ctx) return nullptr;
    share->hlindex->hlindex_data= ctx;
    ctx->refcnt++;
  }
  if (ctx)
    ctx->refcnt++;
617
  share->unlock_share();
618
  return ctx;
619 620
}

621
int MHNSW_Context::acquire(MHNSW_Context **ctx, TABLE *table, bool for_update)
622
{
623 624
  TABLE *graph= table->hlindex;

625
  if (!(*ctx= MHNSW_Trx::get_from_thd(table, for_update)))
626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643
  {
    *ctx= MHNSW_Context::get_from_share(table->s, table);
    if (table->file->has_transactions())
      mysql_rwlock_rdlock(&(*ctx)->commit_lock);
  }

  if ((*ctx)->start)
    return 0;

  if (int err= graph->file->ha_index_init(IDX_LAYER, 1))
    return err;

  int err= graph->file->ha_index_last(graph->record[0]);
  graph->file->ha_index_end();
  if (err)
    return err;

  graph->file->position(graph->record[0]);
644
  (*ctx)->set_lengths(FVector::data_to_value_size(graph->field[FIELD_VEC]->value_length()));
645 646 647 648
  (*ctx)->start= (*ctx)->get_node(graph->file->ref);
  return (*ctx)->start->load_from_record(graph);
}

649 650
/* copy the vector, preprocessed as needed */
const FVector *FVectorNode::make_vec(const void *v)
651
{
Sergei Golubchik's avatar
Sergei Golubchik committed
652
  return FVector::create(ctx->metric, tref() + tref_len(), v, ctx->byte_len);
653
}
654

655
FVectorNode::FVectorNode(MHNSW_Context *ctx_, const void *gref_)
656
  : ctx(ctx_), stored(true), deleted(false)
657
{
658
  memcpy(gref(), gref_, gref_len());
659
}
660

661 662
FVectorNode::FVectorNode(MHNSW_Context *ctx_, const void *tref_, uint8_t layer,
                         const void *vec_)
663
  : ctx(ctx_), stored(false), deleted(false)
664
{
665
  DBUG_ASSERT(tref_);
666 667 668 669
  memset(gref(), 0xff, gref_len()); // important: larger than any real gref
  memcpy(tref(), tref_, tref_len());
  vec= make_vec(vec_);

670
  alloc_neighborhood(layer);
671
}
672

673
float FVectorNode::distance_to(const FVector *other) const
674
{
675
  return vec->distance_to(other, ctx->vec_len);
676
}
677

678
int FVectorNode::alloc_neighborhood(uint8_t layer)
679
{
680 681
  if (neighbors)
    return 0;
682
  max_layer= layer;
683 684 685 686
  neighbors= (Neighborhood*)ctx->alloc_neighborhood(layer);
  auto ptr= (FVectorNode**)(neighbors + (layer+1));
  for (size_t i= 0; i <= layer; i++)
    ptr= neighbors[i].init(ptr, ctx->max_neighbors(i));
687 688 689
  return 0;
}

690
int FVectorNode::load(TABLE *graph)
691
{
692
  if (likely(vec))
693
    return 0;
694

695 696 697 698 699 700
  DBUG_ASSERT(stored);
  // trx: consider loading nodes from shared, when it makes sense
  // for ann_benchmarks it does not
  if (int err= graph->file->ha_rnd_pos(graph->record[0], gref()))
    return err;
  return load_from_record(graph);
701 702
}

703
int FVectorNode::load_from_record(TABLE *graph)
704
{
705 706 707 708 709 710 711 712
  DBUG_ASSERT(ctx->byte_len);

  uint ticket= ctx->lock_node(this);
  SCOPE_EXIT([this, ticket](){ ctx->unlock_node(ticket); });

  if (vec)
    return 0;

713
  String buf, *v= graph->field[FIELD_TREF]->val_str(&buf);
714 715 716 717 718 719 720
  deleted= graph->field[FIELD_TREF]->is_null();
  if (!deleted)
  {
    if (unlikely(v->length() != tref_len()))
      return my_errno= HA_ERR_CRASHED;
    memcpy(tref(), v->ptr(), v->length());
  }
721 722 723

  v= graph->field[FIELD_VEC]->val_str(&buf);
  if (unlikely(!v))
724
    return my_errno= HA_ERR_CRASHED;
725

726
  if (v->length() != FVector::data_size(ctx->vec_len))
727
    return my_errno= HA_ERR_CRASHED;
728 729 730
  FVector *vec_ptr= FVector::align_ptr(tref() + tref_len());
  memcpy(vec_ptr->data(), v->ptr(), v->length());
  vec_ptr->postprocess(ctx->vec_len);
731 732 733

  longlong layer= graph->field[FIELD_LAYER]->val_int();
  if (layer > 100) // 10e30 nodes at M=2, more at larger M's
734
    return my_errno= HA_ERR_CRASHED;
735

736 737
  if (int err= alloc_neighborhood(static_cast<uint8_t>(layer)))
    return err;
738 739 740

  v= graph->field[FIELD_NEIGHBORS]->val_str(&buf);
  if (unlikely(!v))
741
    return my_errno= HA_ERR_CRASHED;
742 743 744 745 746 747

  // <N> <gref> <gref> ... <N> ...etc...
  uchar *ptr= (uchar*)v->ptr(), *end= ptr + v->length();
  for (size_t i=0; i <= max_layer; i++)
  {
    if (unlikely(ptr >= end))
748
      return my_errno= HA_ERR_CRASHED;
749
    size_t grefs= *ptr++;
750 751 752 753 754
    if (unlikely(ptr + grefs * gref_len() > end))
      return my_errno= HA_ERR_CRASHED;
    neighbors[i].num= grefs;
    for (size_t j=0; j < grefs; j++, ptr+= gref_len())
      neighbors[i].links[j]= ctx->get_node(ptr);
755
  }
756
  vec= vec_ptr; // must be done at the very end
757
  return 0;
758 759
}

760
void FVectorNode::push_neighbor(size_t layer, FVectorNode *other)
761
{
762 763
  DBUG_ASSERT(neighbors[layer].num < ctx->max_neighbors(layer));
  neighbors[layer].links[neighbors[layer].num++]= other;
764
}
765

766 767 768 769
size_t FVectorNode::tref_len() const { return ctx->tref_len; }
size_t FVectorNode::gref_len() const { return ctx->gref_len; }
uchar *FVectorNode::gref() const { return (uchar*)(this+1); }
uchar *FVectorNode::tref() const { return gref() + gref_len(); }
770

771 772
uchar *FVectorNode::get_key(const FVectorNode *elem, size_t *key_len, my_bool)
{
773 774
  *key_len= elem->gref_len();
  return elem->gref();
775 776
}

777 778
/* one visited node during the search. caches the distance to target */
struct Visited : public Sql_alloc
779
{
780 781 782 783
  FVectorNode *node;
  const float distance_to_target;
  Visited(FVectorNode *n, float d) : node(n), distance_to_target(d) {}
  static int cmp(void *, const Visited* a, const Visited *b)
784
  {
785 786
    return a->distance_to_target < b->distance_to_target ? -1 :
           a->distance_to_target > b->distance_to_target ?  1 : 0;
787
  }
788 789 790 791
};

/*
  a factory to create Visited and keep track of already seen nodes
792

793 794 795 796 797 798
  note that PatternedSimdBloomFilter works in blocks of 8 elements,
  so on insert they're accumulated in nodes[], on search the caller
  provides 8 addresses at once. we record 0x0 as "seen" so that
  the caller could pad the input with nullptr's
*/
class VisitedSet
799
{
800
  MEM_ROOT *root;
801
  const FVector *target;
802 803 804 805 806
  PatternedSimdBloomFilter<FVectorNode> map;
  const FVectorNode *nodes[8]= {0,0,0,0,0,0,0,0};
  size_t idx= 1; // to record 0 in the filter
  public:
  uint count= 0;
807
  VisitedSet(MEM_ROOT *root, const FVector *target, uint size) :
808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826
    root(root), target(target), map(size, 0.01f) {}
  Visited *create(FVectorNode *node)
  {
    auto *v= new (root) Visited(node, node->distance_to(target));
    insert(node);
    count++;
    return v;
  }
  void insert(const FVectorNode *n)
  {
    nodes[idx++]= n;
    if (idx == 8) flush();
  }
  void flush() {
    if (idx) map.Insert(nodes);
    idx=0;
  }
  uint8_t seen(FVectorNode **nodes) { return map.Query(nodes); }
};
827 828


829 830 831 832 833 834 835 836 837
/*
  selects best neighbors from the list of candidates plus one extra candidate

  one extra candidate is specified separately to avoid appending it to
  the Neighborhood candidates, which might be already at its max size.
*/
static int select_neighbors(MHNSW_Context *ctx, TABLE *graph, size_t layer,
                            FVectorNode &target, const Neighborhood &candidates,
                            FVectorNode *extra_candidate,
838
                            size_t max_neighbor_connections)
839
{
840
  Queue<Visited> pq; // working queue
841

842 843
  if (pq.init(10000, false, Visited::cmp))
    return my_errno= HA_ERR_OUT_OF_MEM;
844

845 846 847 848
  MEM_ROOT * const root= graph->in_use->mem_root;
  auto discarded= (Visited**)my_safe_alloca(sizeof(Visited**)*max_neighbor_connections);
  size_t discarded_num= 0;
  Neighborhood &neighbors= target.neighbors[layer];
849

850
  for (size_t i=0; i < candidates.num; i++)
851
  {
852 853 854
    FVectorNode *node= candidates.links[i];
    if (int err= node->load(graph))
      return err;
855
    pq.push(new (root) Visited(node, node->distance_to(target.vec)));
856
  }
857
  if (extra_candidate)
858
    pq.push(new (root) Visited(extra_candidate, extra_candidate->distance_to(target.vec)));
859 860

  DBUG_ASSERT(pq.elements());
861
  neighbors.num= 0;
862

863
  while (pq.elements() && neighbors.num < max_neighbor_connections)
864
  {
865 866 867
    Visited *vec= pq.pop();
    FVectorNode * const node= vec->node;
    const float target_dista= vec->distance_to_target / alpha;
868
    bool discard= false;
869
    for (size_t i=0; i < neighbors.num; i++)
870
      if ((discard= node->distance_to(neighbors.links[i]->vec) < target_dista))
871
        break;
872
    if (!discard)
873 874 875
      target.push_neighbor(layer, node);
    else if (discarded_num + neighbors.num < max_neighbor_connections)
      discarded[discarded_num++]= vec;
876 877
  }

878 879
  for (size_t i=0; i < discarded_num && neighbors.num < max_neighbor_connections; i++)
    target.push_neighbor(layer, discarded[i]->node);
880

881
  my_safe_afree(discarded, sizeof(Visited**)*max_neighbor_connections);
Sergei Golubchik's avatar
Sergei Golubchik committed
882
  return 0;
883 884
}

Sergei Golubchik's avatar
Sergei Golubchik committed
885

886
int FVectorNode::save(TABLE *graph)
887
{
888 889
  DBUG_ASSERT(vec);
  DBUG_ASSERT(neighbors);
890

891 892
  restore_record(graph, s->default_values);
  graph->field[FIELD_LAYER]->store(max_layer, false);
893 894 895 896 897 898 899
  if (deleted)
    graph->field[FIELD_TREF]->set_null();
  else
  {
    graph->field[FIELD_TREF]->set_notnull();
    graph->field[FIELD_TREF]->store_binary(tref(), tref_len());
  }
900
  graph->field[FIELD_VEC]->store_binary(vec->data(), FVector::data_size(ctx->vec_len));
901

902 903
  size_t total_size= 0;
  for (size_t i=0; i <= max_layer; i++)
904
    total_size+= 1 + gref_len() * neighbors[i].num;
905

906 907 908
  uchar *neighbor_blob= static_cast<uchar *>(my_safe_alloca(total_size));
  uchar *ptr= neighbor_blob;
  for (size_t i= 0; i <= max_layer; i++)
909
  {
910 911 912
    *ptr++= (uchar)(neighbors[i].num);
    for (size_t j= 0; j < neighbors[i].num; j++, ptr+= gref_len())
      memcpy(ptr, neighbors[i].links[j]->gref(), gref_len());
913
  }
914
  graph->field[FIELD_NEIGHBORS]->store_binary(neighbor_blob, total_size);
915

916 917
  int err;
  if (stored)
918
  {
919
    if (!(err= graph->file->ha_rnd_pos(graph->record[1], gref())))
920
    {
921 922 923
      err= graph->file->ha_update_row(graph->record[1], graph->record[0]);
      if (err == HA_ERR_RECORD_IS_THE_SAME)
        err= 0;
924
    }
925
  }
926 927
  else
  {
928
    err= graph->file->ha_write_row(graph->record[0]);
929
    graph->file->position(graph->record[0]);
930 931 932
    memcpy(gref(), graph->file->ref, gref_len());
    stored= true;
    ctx->cache_node(this);
933 934
  }
  my_safe_afree(neighbor_blob, total_size);
935
  return err;
936 937
}

938 939
static int update_second_degree_neighbors(MHNSW_Context *ctx, TABLE *graph,
                                          size_t layer, FVectorNode *node)
940
{
941 942 943 944
  const uint max_neighbors= ctx->max_neighbors(layer);
  // it seems that one could update nodes in the gref order
  // to avoid InnoDB deadlocks, but it produces no noticeable effect
  for (size_t i=0; i < node->neighbors[layer].num; i++)
945
  {
946 947 948 949 950 951 952 953 954 955
    FVectorNode *neigh= node->neighbors[layer].links[i];
    Neighborhood &neighneighbors= neigh->neighbors[layer];
    if (neighneighbors.num < max_neighbors)
      neigh->push_neighbor(layer, node);
    else
      if (int err= select_neighbors(ctx, graph, layer, *neigh, neighneighbors,
                                    node, max_neighbors))
        return err;
    if (int err= neigh->save(graph))
      return err;
956
  }
957
  return 0;
958 959
}

960
static int search_layer(MHNSW_Context *ctx, TABLE *graph, const FVector *target,
Sergei Golubchik's avatar
Sergei Golubchik committed
961 962
                        Neighborhood *start_nodes, uint result_size,
                        size_t layer, Neighborhood *result, bool construction)
963
{
964 965
  DBUG_ASSERT(start_nodes->num > 0);
  result->num= 0;
966

967
  MEM_ROOT * const root= graph->in_use->mem_root;
Sergei Golubchik's avatar
Sergei Golubchik committed
968 969 970
  Queue<Visited> candidates, best;
  bool skip_deleted;
  uint ef= result_size;
971

Sergei Golubchik's avatar
Sergei Golubchik committed
972 973 974 975 976 977 978 979 980 981
  if (construction)
  {
    skip_deleted= false;
    if (ef > 1)
      ef= std::max(ef_construction, ef);
  }
  else
  {
    skip_deleted= layer == 0;
    if (ef > 1 || layer == 0)
982
      ef= std::max(THDVAR(graph->in_use, min_limit), ef);
Sergei Golubchik's avatar
Sergei Golubchik committed
983
  }
984

985 986 987 988 989 990 991 992
  // WARNING! heuristic here
  const double est_heuristic= 8 * std::sqrt(ctx->max_neighbors(layer));
  const uint est_size= static_cast<uint>(est_heuristic * std::pow(ef, ctx->ef_power));
  VisitedSet visited(root, target, est_size);

  candidates.init(10000, false, Visited::cmp);
  best.init(ef, true, Visited::cmp);

Sergei Golubchik's avatar
Sergei Golubchik committed
993
  DBUG_ASSERT(start_nodes->num <= result_size);
994
  for (size_t i=0; i < start_nodes->num; i++)
995
  {
996 997
    Visited *v= visited.create(start_nodes->links[i]);
    candidates.push(v);
998 999
    if (skip_deleted && v->node->deleted)
      continue;
Sergei Golubchik's avatar
Sergei Golubchik committed
1000
    best.push(v);
1001 1002
  }

1003 1004
  float furthest_best= best.is_empty() ? FLT_MAX
                       : best.top()->distance_to_target * generosity;
1005 1006
  while (candidates.elements())
  {
1007
    const Visited &cur= *candidates.pop();
Sergei Golubchik's avatar
Sergei Golubchik committed
1008
    if (cur.distance_to_target > furthest_best && best.is_full())
1009
      break; // All possible candidates are worse than what we have
1010

1011 1012 1013 1014 1015
    visited.flush();

    Neighborhood &neighbors= cur.node->neighbors[layer];
    FVectorNode **links= neighbors.links, **end= links + neighbors.num;
    for (; links < end; links+= 8)
1016
    {
1017 1018
      uint8_t res= visited.seen(links);
      if (res == 0xff)
1019 1020
        continue;

1021
      for (size_t i= 0; i < 8; i++)
1022
      {
1023 1024 1025 1026 1027
        if (res & (1 << i))
          continue;
        if (int err= links[i]->load(graph))
          return err;
        Visited *v= visited.create(links[i]);
Sergei Golubchik's avatar
Sergei Golubchik committed
1028
        if (!best.is_full())
1029 1030
        {
          candidates.push(v);
1031 1032
          if (skip_deleted && v->node->deleted)
            continue;
1033
          best.push(v);
1034
          furthest_best= best.top()->distance_to_target * generosity;
1035 1036 1037
        }
        else if (v->distance_to_target < furthest_best)
        {
1038
          candidates.safe_push(v);
1039 1040
          if (skip_deleted && v->node->deleted)
            continue;
1041 1042 1043 1044 1045
          if (v->distance_to_target < best.top()->distance_to_target)
          {
            best.replace_top(v);
            furthest_best= best.top()->distance_to_target * generosity;
          }
1046
        }
1047 1048 1049
      }
    }
  }
1050 1051 1052 1053 1054
  if (ef > 1 && visited.count*2 > est_size)
  {
    double ef_power= std::log(visited.count*2/est_heuristic) / std::log(ef);
    set_if_bigger(ctx->ef_power, ef_power); // not atomic, but it's ok
  }
1055

Sergei Golubchik's avatar
Sergei Golubchik committed
1056 1057 1058
  while (best.elements() > result_size)
    best.pop();

1059 1060 1061
  result->num= best.elements();
  for (FVectorNode **links= result->links + result->num; best.elements();)
    *--links= best.pop()->node;
1062

1063
  return 0;
1064 1065
}

1066

Sergei Golubchik's avatar
Sergei Golubchik committed
1067 1068 1069 1070 1071
static int bad_value_on_insert(Field *f)
{
  my_error(ER_TRUNCATED_WRONG_VALUE_FOR_FIELD, MYF(0), "vector", "...",
           f->table->s->db.str, f->table->s->table_name.str, f->field_name.str,
           f->table->in_use->get_stmt_da()->current_row_for_warning());
1072
  return my_errno= HA_ERR_GENERIC;
Sergei Golubchik's avatar
Sergei Golubchik committed
1073 1074
}

1075

1076 1077
int mhnsw_insert(TABLE *table, KEY *keyinfo)
{
Sergei Golubchik's avatar
Sergei Golubchik committed
1078
  THD *thd= table->in_use;
1079 1080
  TABLE *graph= table->hlindex;
  MY_BITMAP *old_map= dbug_tmp_use_all_columns(table, &table->read_set);
1081 1082
  Field *vec_field= keyinfo->key_part->field;
  String buf, *res= vec_field->val_str(&buf);
1083
  MHNSW_Context *ctx;
1084 1085 1086 1087 1088

  /* metadata are checked on open */
  DBUG_ASSERT(graph);
  DBUG_ASSERT(keyinfo->algorithm == HA_KEY_ALG_VECTOR);
  DBUG_ASSERT(keyinfo->usable_key_parts == 1);
1089 1090
  DBUG_ASSERT(vec_field->binary());
  DBUG_ASSERT(vec_field->cmp_type() == STRING_RESULT);
1091
  DBUG_ASSERT(res); // ER_INDEX_CANNOT_HAVE_NULL
1092
  DBUG_ASSERT(table->file->ref_length <= graph->field[FIELD_TREF]->field_length);
1093

Sergei Golubchik's avatar
Sergei Golubchik committed
1094 1095 1096 1097
  // XXX returning an error here will rollback the insert in InnoDB
  // but in MyISAM the row will stay inserted, making the index out of sync:
  // invalid vector values are present in the table but cannot be found
  // via an index. The easiest way to fix it is with a VECTOR(N) type
1098
  if (res->length() == 0 || res->length() % 4)
Sergei Golubchik's avatar
Sergei Golubchik committed
1099
    return bad_value_on_insert(vec_field);
1100

1101 1102
  table->file->position(table->record[0]);

1103 1104 1105
  int err= MHNSW_Context::acquire(&ctx, table, true);
  SCOPE_EXIT([ctx, table](){ ctx->release(table); });
  if (err)
1106
  {
1107 1108
    if (err != HA_ERR_END_OF_FILE)
      return err;
Sergei Golubchik's avatar
Sergei Golubchik committed
1109

1110
    // First insert!
1111 1112 1113 1114 1115 1116
    ctx->set_lengths(res->length());
    FVectorNode *target= new (ctx->alloc_node())
                   FVectorNode(ctx, table->file->ref, 0, res->ptr());
    if (!((err= target->save(graph))))
      ctx->start= target;
    return err;
1117 1118
  }

1119 1120
  if (ctx->byte_len != res->length())
    return bad_value_on_insert(vec_field);
Sergei Golubchik's avatar
Sergei Golubchik committed
1121

1122 1123 1124
  MEM_ROOT_SAVEPOINT memroot_sv;
  root_make_savepoint(thd->mem_root, &memroot_sv);
  SCOPE_EXIT([memroot_sv](){ root_free_to_savepoint(&memroot_sv); });
Sergei Golubchik's avatar
Sergei Golubchik committed
1125

Sergei Golubchik's avatar
Sergei Golubchik committed
1126
  const size_t max_found= ctx->max_neighbors(0);
1127
  Neighborhood candidates, start_nodes;
Sergei Golubchik's avatar
Sergei Golubchik committed
1128 1129
  candidates.init(thd->alloc<FVectorNode*>(max_found + 7), max_found);
  start_nodes.init(thd->alloc<FVectorNode*>(max_found + 7), max_found);
1130
  start_nodes.links[start_nodes.num++]= ctx->start;
1131

1132 1133 1134 1135 1136
  const double NORMALIZATION_FACTOR= 1 / std::log(ctx->M);
  double log= -std::log(my_rnd(&thd->rand)) * NORMALIZATION_FACTOR;
  const uint8_t max_layer= start_nodes.links[0]->max_layer;
  uint8_t target_layer= std::min<uint8_t>(static_cast<uint8_t>(std::floor(log)), max_layer + 1);
  int cur_layer;
Sergei Golubchik's avatar
Sergei Golubchik committed
1137

1138 1139
  FVectorNode *target= new (ctx->alloc_node())
                 FVectorNode(ctx, table->file->ref, target_layer, res->ptr());
Sergei Golubchik's avatar
Sergei Golubchik committed
1140

1141 1142 1143
  if (int err= graph->file->ha_rnd_init(0))
    return err;
  SCOPE_EXIT([graph](){ graph->file->ha_rnd_end(); });
Sergei Golubchik's avatar
Sergei Golubchik committed
1144

1145
  for (cur_layer= max_layer; cur_layer > target_layer; cur_layer--)
1146
  {
1147 1148
    if (int err= search_layer(ctx, graph, target->vec, &start_nodes, 1,
                              cur_layer, &candidates, false))
1149 1150
      return err;
    std::swap(start_nodes, candidates);
1151
  }
1152

1153
  for (; cur_layer >= 0; cur_layer--)
1154
  {
1155
    uint max_neighbors= ctx->max_neighbors(cur_layer);
1156
    if (int err= search_layer(ctx, graph, target->vec, &start_nodes,
Sergei Golubchik's avatar
Sergei Golubchik committed
1157
                              max_neighbors, cur_layer, &candidates, true))
1158 1159 1160 1161 1162 1163
      return err;

    if (int err= select_neighbors(ctx, graph, cur_layer, *target, candidates,
                                  0, max_neighbors))
      return err;
    std::swap(start_nodes, candidates);
1164
  }
1165

1166 1167
  if (int err= target->save(graph))
    return err;
1168

1169 1170 1171 1172
  if (target_layer > max_layer)
    ctx->start= target;

  for (cur_layer= target_layer; cur_layer >= 0; cur_layer--)
1173
  {
1174 1175
    if (int err= update_second_degree_neighbors(ctx, graph, cur_layer, target))
      return err;
1176 1177
  }

1178 1179
  dbug_tmp_restore_column_map(&table->read_set, old_map);

Sergei Golubchik's avatar
Sergei Golubchik committed
1180
  return 0;
1181 1182
}

1183

1184 1185
int mhnsw_first(TABLE *table, KEY *keyinfo, Item *dist, ulonglong limit)
{
Sergei Golubchik's avatar
Sergei Golubchik committed
1186
  THD *thd= table->in_use;
1187
  TABLE *graph= table->hlindex;
Sergei Golubchik's avatar
Sergei Golubchik committed
1188 1189 1190
  auto *fun= static_cast<Item_func_vec_distance_common*>(dist->real_item());
  DBUG_ASSERT(fun);

Sergei Golubchik's avatar
Sergei Golubchik committed
1191
  String buf, *res= fun->get_const_arg()->val_str(&buf);
1192
  MHNSW_Context *ctx;
1193

1194
  if (int err= table->file->ha_rnd_init(0))
1195
    return err;
1196

1197 1198 1199
  int err= MHNSW_Context::acquire(&ctx, table, false);
  SCOPE_EXIT([ctx, table](){ ctx->release(table); });
  if (err)
1200 1201
    return err;

1202
  Neighborhood candidates, start_nodes;
Sergei Golubchik's avatar
Sergei Golubchik committed
1203 1204
  candidates.init(thd->alloc<FVectorNode*>(limit + 7), limit);
  start_nodes.init(thd->alloc<FVectorNode*>(limit + 7), limit);
Sergei Golubchik's avatar
Sergei Golubchik committed
1205

Sergei Golubchik's avatar
Sergei Golubchik committed
1206 1207
  // one could put all max_layer nodes in start_nodes
  // but it has no effect on the recall or speed
1208
  start_nodes.links[start_nodes.num++]= ctx->start;
Sergei Golubchik's avatar
Sergei Golubchik committed
1209 1210 1211 1212

  /*
    if the query vector is NULL or invalid, VEC_DISTANCE will return
    NULL, so the result is basically unsorted, we can return rows
1213
    in any order. Let's use some hardcoded value here
Sergei Golubchik's avatar
Sergei Golubchik committed
1214
  */
1215
  if (!res || ctx->byte_len != res->length())
1216 1217 1218 1219 1220 1221
  {
    res= &buf;
    buf.alloc(ctx->byte_len);
    for (size_t i=0; i < ctx->vec_len; i++)
      ((float*)buf.ptr())[i]= i == 0;
  }
1222 1223

  const longlong max_layer= start_nodes.links[0]->max_layer;
Sergei Golubchik's avatar
Sergei Golubchik committed
1224
  auto target= FVector::create(ctx->metric, thd->alloc(FVector::alloc_size(ctx->vec_len)),
1225
                               res->ptr(), res->length());
1226 1227 1228 1229

  if (int err= graph->file->ha_rnd_init(0))
    return err;
  SCOPE_EXIT([graph](){ graph->file->ha_rnd_end(); });
Sergei Golubchik's avatar
Sergei Golubchik committed
1230

1231 1232
  for (size_t cur_layer= max_layer; cur_layer > 0; cur_layer--)
  {
1233
    if (int err= search_layer(ctx, graph, target, &start_nodes, 1, cur_layer,
1234
                              &candidates, false))
1235 1236
      return err;
    std::swap(start_nodes, candidates);
1237 1238
  }

Sergei Golubchik's avatar
Sergei Golubchik committed
1239 1240
  if (int err= search_layer(ctx, graph, target, &start_nodes,
                            static_cast<uint>(limit), 0, &candidates, false))
1241
    return err;
1242

1243 1244 1245
  if (limit > candidates.num)
    limit= candidates.num;
  size_t context_size= limit * ctx->tref_len + sizeof(ulonglong);
1246
  char *context= thd->alloc(context_size);
1247 1248
  graph->context= context;

1249 1250 1251
  *(ulonglong*)context= limit;
  context+= context_size;

1252
  for (size_t i=0; limit--; i++)
1253
  {
1254 1255
    context-= ctx->tref_len;
    memcpy(context, candidates.links[i]->tref(), ctx->tref_len);
1256 1257
  }
  DBUG_ASSERT(context - sizeof(ulonglong) == graph->context);
1258

Sergei Golubchik's avatar
Sergei Golubchik committed
1259
  return mhnsw_next(table);
1260 1261 1262 1263
}

int mhnsw_next(TABLE *table)
{
1264 1265
  uchar *ref= (uchar*)(table->hlindex->context);
  if (ulonglong *limit= (ulonglong*)ref)
1266
  {
1267 1268
    ref+= sizeof(ulonglong) + (--*limit) * table->file->ref_length;
    return table->file->ha_rnd_pos(table->record[0], ref);
1269
  }
1270 1271 1272 1273 1274 1275 1276 1277 1278 1279 1280
  return my_errno= HA_ERR_END_OF_FILE;
}

void mhnsw_free(TABLE_SHARE *share)
{
  TABLE_SHARE *graph_share= share->hlindex;
  if (!graph_share->hlindex_data)
    return;

  static_cast<MHNSW_Context*>(graph_share->hlindex_data)->~MHNSW_Context();
  graph_share->hlindex_data= 0;
1281
}
1282

Sergei Golubchik's avatar
Sergei Golubchik committed
1283
int mhnsw_invalidate(TABLE *table, const uchar *rec, KEY *keyinfo)
1284 1285 1286
{
  TABLE *graph= table->hlindex;
  handler *h= table->file;
Sergei Golubchik's avatar
Sergei Golubchik committed
1287 1288
  MHNSW_Context *ctx;
  bool use_ctx= !MHNSW_Context::acquire(&ctx, table, true);
1289 1290 1291 1292 1293

  /* metadata are checked on open */
  DBUG_ASSERT(graph);
  DBUG_ASSERT(keyinfo->algorithm == HA_KEY_ALG_VECTOR);
  DBUG_ASSERT(keyinfo->usable_key_parts == 1);
Sergei Golubchik's avatar
Sergei Golubchik committed
1294
  DBUG_ASSERT(h->ref_length <= graph->field[FIELD_TREF]->field_length);
1295 1296 1297 1298

  // target record:
  h->position(rec);
  graph->field[FIELD_TREF]->set_notnull();
Sergei Golubchik's avatar
Sergei Golubchik committed
1299
  graph->field[FIELD_TREF]->store_binary(h->ref, h->ref_length);
1300

Sergei Golubchik's avatar
Sergei Golubchik committed
1301 1302 1303
  uchar *key= (uchar*)alloca(graph->key_info[IDX_TREF].key_length);
  key_copy(key, graph->record[0], &graph->key_info[IDX_TREF],
           graph->key_info[IDX_TREF].key_length);
1304

Sergei Golubchik's avatar
Sergei Golubchik committed
1305 1306 1307 1308 1309 1310 1311 1312
  if (int err= graph->file->ha_index_read_idx_map(graph->record[1], IDX_TREF,
                                        key, HA_WHOLE_KEY, HA_READ_KEY_EXACT))
   return err;

  restore_record(graph, record[1]);
  graph->field[FIELD_TREF]->set_null();
  if (int err= graph->file->ha_update_row(graph->record[1], graph->record[0]))
    return err;
1313

Sergei Golubchik's avatar
Sergei Golubchik committed
1314
  if (use_ctx)
1315
  {
Sergei Golubchik's avatar
Sergei Golubchik committed
1316 1317 1318 1319
    graph->file->position(graph->record[0]);
    FVectorNode *node= ctx->get_node(graph->file->ref);
    node->deleted= true;
    ctx->release(table);
1320 1321
  }

Sergei Golubchik's avatar
Sergei Golubchik committed
1322 1323 1324
  return 0;
}

Sergei Golubchik's avatar
Sergei Golubchik committed
1325
int mhnsw_delete_all(TABLE *table, KEY *keyinfo, bool truncate)
Sergei Golubchik's avatar
Sergei Golubchik committed
1326 1327 1328 1329 1330 1331 1332
{
  TABLE *graph= table->hlindex;

  /* metadata are checked on open */
  DBUG_ASSERT(graph);
  DBUG_ASSERT(keyinfo->algorithm == HA_KEY_ALG_VECTOR);
  DBUG_ASSERT(keyinfo->usable_key_parts == 1);
1333

Sergei Golubchik's avatar
Sergei Golubchik committed
1334 1335
  if (int err= truncate ? graph->file->truncate()
                        : graph->file->delete_all_rows())
Sergei Golubchik's avatar
Sergei Golubchik committed
1336 1337 1338 1339 1340 1341 1342 1343
   return err;

  MHNSW_Context *ctx;
  if (!MHNSW_Context::acquire(&ctx, table, true))
  {
    ctx->reset(table->s);
    ctx->release(table);
  }
1344 1345 1346 1347

  return 0;
}

1348 1349 1350
const LEX_CSTRING mhnsw_hlindex_table_def(THD *thd, uint ref_length)
{
  const char templ[]="CREATE TABLE i (                   "
1351 1352 1353 1354
                     "  layer tinyint not null,          "
                     "  tref varbinary(%u),              "
                     "  vec blob not null,               "
                     "  neighbors blob not null,         "
Sergei Golubchik's avatar
Sergei Golubchik committed
1355 1356
                     "  unique (tref),                   "
                     "  key (layer))                     ";
1357 1358
  size_t len= sizeof(templ) + 32;
  char *s= thd->alloc(len);
1359
  len= my_snprintf(s, len, templ, ref_length);
1360 1361
  return {s, len};
}
1362

Sergei Golubchik's avatar
Sergei Golubchik committed
1363 1364 1365 1366 1367 1368 1369
bool mhnsw_uses_distance(const TABLE *table, KEY *keyinfo, const Item *dist)
{
  if (keyinfo->option_struct->metric == EUCLIDEAN)
    return dynamic_cast<const Item_func_vec_distance_euclidean*>(dist) != NULL;
  return dynamic_cast<const Item_func_vec_distance_cosine*>(dist) != NULL;
}

1370 1371 1372 1373 1374 1375 1376
/*
  Declare the plugin and index options
*/

ha_create_table_option mhnsw_index_options[]=
{
  HA_IOPTION_SYSVAR("max_edges_per_node", M, max_edges_per_node),
Sergei Golubchik's avatar
Sergei Golubchik committed
1377
  HA_IOPTION_SYSVAR("distance_function", metric, distance_function),
1378 1379 1380 1381 1382 1383 1384 1385 1386 1387 1388 1389 1390 1391 1392 1393 1394 1395 1396 1397 1398 1399 1400 1401 1402 1403 1404 1405
  HA_IOPTION_END
};

st_plugin_int *mhnsw_plugin;

static int mhnsw_init(void *p)
{
  mhnsw_plugin= (st_plugin_int *)p;
  mhnsw_plugin->data= &MHNSW_Trx::tp;
  if (setup_transaction_participant(mhnsw_plugin))
    return 1;

  return resolve_sysvar_table_options(mhnsw_index_options);
}

static int mhnsw_deinit(void *)
{
  free_sysvar_table_options(mhnsw_index_options);
  return 0;
}

static struct st_mysql_storage_engine mhnsw_daemon=
{ MYSQL_DAEMON_INTERFACE_VERSION };

static struct st_mysql_sys_var *mhnsw_sys_vars[]=
{
  MYSQL_SYSVAR(cache_size),
  MYSQL_SYSVAR(max_edges_per_node),
Sergei Golubchik's avatar
Sergei Golubchik committed
1406
  MYSQL_SYSVAR(distance_function),
1407 1408 1409 1410 1411 1412 1413 1414 1415 1416 1417 1418 1419
  MYSQL_SYSVAR(min_limit),
  NULL
};

maria_declare_plugin(mhnsw)
{
  MYSQL_DAEMON_PLUGIN,
  &mhnsw_daemon, "mhnsw", "MariaDB plc",
  "A plugin for mhnsw vector index algorithm",
  PLUGIN_LICENSE_GPL, mhnsw_init, mhnsw_deinit, 0x0100, NULL,
  mhnsw_sys_vars, "1.0", MariaDB_PLUGIN_MATURITY_STABLE
}
maria_declare_plugin_end;