From 907df8fda357d6a2f92892c4a66bf0a110b34b74 Mon Sep 17 00:00:00 2001 From: rahul-iyer Date: Mon, 6 Jul 2026 20:21:59 -0700 Subject: [PATCH] Implement scalar quantization storage for HNSW Persist SQ8 and SQ16 embeddings in an internal _QEMB node table so quantized index state participates in Ladybug's normal MVCC, WAL, and checkpoint lifecycle. Use the durable quantized rows for HNSW construction and on-disk search, with batched table-backed reads when caching is disabled and an aligned dense cache when caching is enabled. Preserve optional full-precision reranking and normalize cosine vectors for the quantized distance path. Keep the shared dense cache transactionally safe by tracking dirty writer transactions and invalidating it from commit and rollback callbacks. Add coverage for quantized index behavior and document storage layout, query flow, cache semantics, and benchmark results. --- vector/docs/scalar_quantization_hnsw.md | 354 +++++++++++ .../src/catalog/hnsw_index_catalog_entry.cpp | 13 +- vector/src/function/create_hnsw_index.cpp | 54 +- vector/src/function/drop_hnsw_index.cpp | 23 +- vector/src/function/query_hnsw_index.cpp | 106 +++- vector/src/include/index/hnsw_config.h | 24 + vector/src/include/index/hnsw_graph.h | 127 +++- vector/src/include/index/hnsw_index.h | 43 +- vector/src/include/index/hnsw_index_utils.h | 24 +- vector/src/index/hnsw_config.cpp | 79 +++ vector/src/index/hnsw_graph.cpp | 587 +++++++++++++++++- vector/src/index/hnsw_index.cpp | 327 ++++++++-- vector/src/index/hnsw_index_utils.cpp | 297 ++++++++- vector/test/test_files/error.test | 6 + vector/test/test_files/small.test | 66 ++ 15 files changed, 2031 insertions(+), 99 deletions(-) create mode 100644 vector/docs/scalar_quantization_hnsw.md diff --git a/vector/docs/scalar_quantization_hnsw.md b/vector/docs/scalar_quantization_hnsw.md new file mode 100644 index 00000000..1f93ee41 --- /dev/null +++ b/vector/docs/scalar_quantization_hnsw.md @@ -0,0 +1,354 @@ +# Scalar Quantization For HNSW Indexes + +This document describes the scalar quantization design for the vector extension HNSW index. The +implementation supports SQ8 and SQ16 quantized distance evaluation while keeping quantized +embeddings in Ladybug storage instead of a raw sidecar file. + +## Goals + +- Reduce memory and distance-computation cost for HNSW search by storing quantized vectors. +- Keep quantized embeddings transactionally consistent with the base node table and HNSW graph. +- Support both `cache_embeddings := true` and `cache_embeddings := false`, with the optimized + query path focused on `cache_embeddings := true`. +- Preserve existing HNSW graph storage: upper and lower graph layers remain internal rel tables. +- Keep full-precision rerank optional through `use_full_precision_rerank := true`. + +## Non-Goals + +- Quantization is not supported for dot-product metrics. It is available for L2, L2SQ, and cosine. +- The HNSW graph itself is not quantized. Only the embedding representation used for distance + evaluation is quantized. +- The in-memory quantized cache is not the durable source of truth. It is a query-time accelerator. + +## User-Facing Configuration + +Scalar quantization is enabled through `CREATE_VECTOR_INDEX`: + +```cypher +CALL CREATE_VECTOR_INDEX( + 'embeddings', + 'idx', + 'vec', + metric := 'l2', + quantization := 'sq8' +); +``` + +Supported values are: + +- `quantization := 'sq8'` +- `quantization := 'sq16'` + +When quantization is enabled, query distance evaluation uses the quantized representation unless +full-precision rerank is enabled. With `use_full_precision_rerank := true`, HNSW candidate +generation uses quantized distances, then final candidate ordering is recomputed against the +original full-precision embedding column. + +## Physical Storage + +Each quantized HNSW index owns an internal node table in addition to its existing graph tables. + +For base table id `T` and index name `idx`, the table name is: + +```text +___QEMB +``` + +The table schema is: + +```cypher +CREATE NODE TABLE ___QEMB ( + id INT64, + valid UINT8, + scale FLOAT, + norm_sq FLOAT, + payload INT8[dimension] | INT16[dimension], + PRIMARY KEY (id) +); +``` + +Column meaning: + +- `id`: source node offset. The quantized table row identity mirrors the base table offset. +- `valid`: `1` when the source embedding is non-null and visible to the writer, `0` for null rows. +- `scale`: per-vector scalar. For L2 and L2SQ this is the reconstruction scale + `max_abs / max_quantized_value`. For cosine this is the normalized-dot scale + `1 / sqrt(sum(q_i^2))`. +- `norm_sq`: squared norm of the quantized vector representation. L2 and L2SQ use this in the + reconstructed distance formula. Cosine keeps it for metadata compatibility but uses the + normalized-dot scale in the hot path. +- `payload`: SQ8 or SQ16 vector payload. + +This table is an internal catalog entry. It is created and dropped with the vector index and is +not part of the public user schema. + +## Page Layout + +`_QEMB` does not introduce a custom packed disk-page format. It is a normal internal `NodeTable`, +so its columns use the existing node-group, column-chunk, buffer-manager, MVCC, WAL, and checkpoint +machinery. The logical row at offset `N` corresponds to source node offset `N`, but the values are +stored column-wise rather than as one contiguous quantized record: + +```text +logical _QEMB row N + +id[N] | valid[N] | scale[N] | norm_sq[N] | payload[N][dimension] + | | | | | + +---------+----------+-----------+-----------------+ + normal NodeTable column chunks and storage pages +``` + +`TableBackedQuantizedEmbeddings` resolves an offset to its committed or uncommitted node group. +`lookupMultiple` reuses the current scan state while offsets remain in the same source and node +group, and reinitializes it when either changes. The selected `valid`, `scale`, `norm_sq`, and +`payload` values are then read through the normal storage and buffer-manager path. The fixed-size +SQ8/SQ16 `payload` remains an `ARRAY` column; it is not a pointer into a separate `.qemb` file. + +When `cache_embeddings := true`, `CachedQuantizedColumn` materializes a separate derived layout for +distance evaluation: + +```text +cache slot N + +nullMask[N] scales[N] normSqs[N] views[N] + | + v +aligned payload base + N * stride +``` + +The payload buffer base is 32-byte aligned. For dimension `D`, SQ8 stores `D` payload bytes with +`stride = alignUp(D, 32)`, while SQ16 stores `2 * D` payload bytes with `stride = 2 * D`. Scale and +norm metadata live in separate dense arrays, and each `QuantizedEmbeddingView` combines those +values with the payload address for the same offset. This cache layout is not durable: it is rebuilt +from `_QEMB` after invalidation. + +## Catalog And Index Metadata + +`HNSWStorageInfo` stores the table ids for all durable index-owned structures: + +- upper HNSW graph rel table +- lower HNSW graph rel table +- quantized embeddings node table, when quantization is enabled + +The quantized table id is appended to the serialized storage info. Deserialization treats it as +optional so older non-quantized or pre-quantized index metadata remains readable. + +## In-Memory Types + +The implementation uses three embedding providers: + +- `QuantizedInMemEmbeddings` + Used while building a new in-memory HNSW graph during `CREATE_VECTOR_INDEX` when embeddings are + cached. It quantizes from the input vectors and avoids storage lookups during the hot build loop. + +- `TableBackedQuantizedEmbeddings` + Reads and writes quantized embeddings through the internal `_QEMB` node table. This is the + durable, transactional representation used when populating or maintaining quantized rows. It is + also the on-disk query provider when `cache_embeddings := false`. Query-time reads use storage + `lookup`/`lookupMultiple`, so hot pages are managed by the normal buffer manager. + +- `CachedQuantizedEmbeddings` + A reader over `CachedQuantizedColumn`, which materializes `_QEMB` rows into an aligned dense + quantized cache for on-disk queries when `cache_embeddings := true`. This mirrors the baseline + cached embedding path: storage remains the source of truth, but HNSW neighbor distance + evaluation reads contiguous in-memory views. + +The durable source of truth is `_QEMB`. With `cache_embeddings := false`, query-time reads stay in +the storage layer. With `cache_embeddings := true`, `_QEMB` is used to hydrate a shared dense cache +owned by the on-disk index before serving hot HNSW distance calls. + +## Create Index Flow + +`CREATE_VECTOR_INDEX` creates all internal storage in one rewrite: + +1. Create the upper graph rel table. +2. Create the lower graph rel table. +3. If quantization is enabled, create the `_QEMB` node table. +4. Build the HNSW graph. +5. Register the `OnDiskHNSWIndex` with `HNSWStorageInfo`, including the quantized table id. +6. Rebuild quantized rows by scanning the base embedding column and writing one `_QEMB` row per + source offset. +7. Force checkpoint so the index-owned tables are persisted together. + +Null embeddings are written as `valid = 0` rows. This preserves the invariant that source offset +`N` maps to quantized row `N`, even when the source embedding is null. + +## Query Flow + +When a query uses a quantized index: + +1. The query vector is quantized into the same SQ8 or SQ16 format. +2. `OnDiskHNSWIndex` chooses an embedding provider: + - `cache_embeddings := true` uses `CachedQuantizedEmbeddings`, backed by a dense + `CachedQuantizedColumn` populated from `_QEMB`. + - `cache_embeddings := false` uses `TableBackedQuantizedEmbeddings`. + - `TableBackedQuantizedEmbeddings::getEmbeddings` batches neighbor vector reads with + `_QEMB` `lookupMultiple`, matching the baseline storage-backed access pattern more closely + than one point lookup per neighbor. +3. HNSW graph traversal reads neighbors from the existing graph rel tables. +4. Distance evaluation reads quantized embedding views from the selected provider. +5. If full-precision rerank is enabled, the final candidates are rescored using `OnDiskEmbeddings` + over the original embedding column. + +For quantized on-disk query, `cache_embeddings := true` does not cache full-precision embeddings. +It caches the quantized representation in memory after hydrating it from durable `_QEMB` rows. If +full-precision rerank is enabled, the original embedding column is read separately for reranking. + +## Distance Evaluation + +L2 and L2SQ use scalar reconstruction from the quantized dot product: + +```text +l2sq = scale_l^2 * norm_l + scale_r^2 * norm_r - 2 * scale_l * scale_r * dot(q_l, q_r) +``` + +`metric := 'l2'` returns `sqrt(max(0, l2sq))`; `metric := 'l2sq'` returns `max(0, l2sq)`. + +Cosine uses a MariaDB-style normalized quantized dot product. During quantization for cosine +indexes, `scale` is stored as: + +```text +scale = 1 / sqrt(sum(q_i^2)) +``` + +The hot path then evaluates cosine distance as: + +```text +distance = 1 - clamp(scale_l * scale_r * dot(q_l, q_r), -1, 1) +``` + +This avoids computing a denominator from both vector norms for every candidate comparison. Query +vectors are quantized with the same cosine-specific scale semantics as `_QEMB` rows, so cached, +table-backed, and in-memory build paths use the same formula. + +## Insert, Update, And Finalize + +Committed inserts write both graph state and quantized state in the same transaction: + +1. Read the inserted embedding from the insert vector. +2. Insert the vector into the HNSW graph if it is non-null. +3. Write the matching `_QEMB` row through `TableBackedQuantizedEmbeddings`. +4. Write `valid = 0` for null embeddings. +5. Mark the shared derived qemb cache dirty for the writer transaction. + +Updates use the existing HNSW update model: the old table row version becomes invisible and the new +vector is inserted as a new version. The quantized table follows the inserted version. + +`finalize()` catches up rows that were not part of the checkpointed graph. For quantized indexes it +writes quantized rows only for the offsets it is finalizing. It does not rebuild the full `_QEMB` +table, which avoids duplicate primary-key writes for rows already maintained by insert. + +## Checkpoint And Recovery + +Checkpoint includes: + +- upper graph rel table +- lower graph rel table +- quantized embeddings node table, when present + +Because `_QEMB` is a normal internal node table, it uses the same storage, WAL, MVCC, checkpoint, +and recovery paths as other node tables. This is the main reason the design uses an internal table +instead of a raw `.qemb` sidecar file. + +After reload, `OnDiskHNSWIndex::load` reconstructs the index from serialized storage info and +loads the quantized table pointer from `quantizedEmbeddingsTableID`. + +## Drop Index Flow + +`DROP_VECTOR_INDEX` drops all index-owned internal structures: + +1. upper graph rel table +2. lower graph rel table +3. `_QEMB` node table, when the index config has quantization enabled +4. index catalog entry + +There is no filesystem sidecar cleanup because quantized embeddings are stored through Ladybug +storage. + +## Cache And Buffer Semantics + +`cache_embeddings := true` does not mean cache full-precision embeddings for quantized indexes. +For on-disk quantized query, it creates a dense qemb cache from `_QEMB` and uses that cache for +HNSW distance evaluation. This is analogous to the baseline cached-embedding path, but the cached +payload is SQ8/SQ16 rather than full-precision floats. + +The durable source of truth remains the `_QEMB` table. `TableBackedQuantizedEmbeddings` uses +transactional `NodeTable` lookup APIs, including `lookupMultiple` for batched neighbor reads, when +`cache_embeddings := false`. `CachedQuantizedEmbeddings` still checks source-table visibility while +serving dense cached views when `cache_embeddings := true`. + +The persistent scan state retains a `const Transaction*` because quantized lookup only reads +transaction visibility and local-storage metadata. `NodeTable::lookupMultiple` currently accepts a +non-const `Transaction*` as part of its scan-state initialization API, so this call site uses a +contained `const_cast`. This remains safe only while `lookupMultiple` and its callees treat the +transaction as read-only. If that call chain gains transaction mutations, the scan state must +instead retain and pass a non-const transaction rather than casting away constness. + +The dense cache is stored in `OnDiskHNSWIndex` as shared cache state. It is protected by a mutex, +has a commit timestamp version, and tracks writer transactions that dirtied `_QEMB`. A query can +reuse the cache only when the cache is not dirty and the query transaction start timestamp is at or +after the cache version. If the cache is absent or too small, the first eligible query rebuilds it +from `_QEMB`. + +When a transaction writes `_QEMB`, the index registers commit and rollback callbacks on the +transaction: + +- On commit, the cache is reset and the cache version advances to the maximum of its current value + and the commit timestamp. +- On rollback, the dirty writer is removed without advancing the cache version. +- The cache remains dirty while any tracked writer transaction is still unresolved. + +Commit callbacks take the cache-state mutex, so their updates are serialized even when multiple +transactions finish concurrently. The callbacks can still acquire the mutex out of commit timestamp +order; taking the maximum keeps the version watermark monotonic and prevents an older callback from +moving it backward. Tracking every dirty transaction separately keeps the cache unavailable until +the last concurrent writer has committed or rolled back. + +This avoids reusing a dense cache across committed qemb changes while allowing rollback-only writes +to clear the dirty state without forcing a new committed version. + +This gives the cached quantized path the same hot-query shape as the regular cached baseline while +retaining a smaller quantized representation. + +![Quantized HNSW cache read and concurrent writer flow](images/quantized_hnsw_cache_transactions.svg) + +[Editable Excalidraw source](images/quantized_hnsw_cache_transactions.excalidraw) + +## Benchmark Results + +The following run uses `cache_embeddings := true`, `efs = 64`, `k = 10`, 200 speed queries repeated +3 times, and 100 accuracy queries on `openai_small_50k`: + + +| Variant | Load | Index Build | Ingest | Query Latency | DB Size | Bytes/Vec | Query RSS | Recall@10 | Precision@10 | +|---|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| base | 5.862s | 51.000s | 56.862s | 17.812 ms/query | 606.5 MiB | 12718.7 | 746.6 MiB | 0.9960 | 0.9960 | +| sq8 | 5.765s | 13.111s | 18.875s | 16.366 ms/query | 723.4 MiB | 15170.2 | 327.1 MiB | 0.9460 | 0.9460 | +| sq16 | 5.421s | 23.337s | 28.758s | 16.444 ms/query | 908.0 MiB | 19041.5 | 489.4 MiB | 0.9950 | 0.9950 | + +Speedup vs baseline: + +| Variant | Load | Index Build | Ingest | Query Latency | Query RSS Reduction | +|---|---:|---:|---:|---:|---:| +| sq8 | 1.02x | 3.89x | 3.01x | 1.09x | 2.28x lower | +| sq16 | 1.08x | 2.19x | 1.98x | 1.08x | 1.53x lower | + +Storage overhead vs baseline: + +| Variant | DB Size | Bytes/Vec | +|---|---:|---:| +| sq8 | 1.19x | 1.19x | +| sq16 | 1.50x | 1.50x | + +## Tradeoffs And Follow-Ups + +- `TableBackedQuantizedEmbeddings` uses storage lookups during search. This is correct, + transactional, page-aware, and buffer-pool-backed, but it still pays lookup/list-decoding overhead + compared with the dense cache used by `cache_embeddings := true`. +- A dense `CachedQuantizedColumn` avoids storage lookup in the hot path, but it adds cache + population cost and shared transaction-aware invalidation requirements. +- `norm_sq` is currently stored as `FLOAT`, matching the table schema. SQ16 with very high + dimensionality may benefit from wider norm storage if precision becomes measurable in recall or + ranking tests. +- Dead HNSW graph edges can still accumulate under the existing delete/update model. Quantization + follows existing visibility rules but does not compact the graph. diff --git a/vector/src/catalog/hnsw_index_catalog_entry.cpp b/vector/src/catalog/hnsw_index_catalog_entry.cpp index 74babcad..24c4f8af 100644 --- a/vector/src/catalog/hnsw_index_catalog_entry.cpp +++ b/vector/src/catalog/hnsw_index_catalog_entry.cpp @@ -40,9 +40,20 @@ std::string HNSWIndexAuxInfo::toCypher(const IndexCatalogEntry& indexEntry, auto propertyName = tableEntry->getProperty(indexEntry.getPropertyIDs()[0]).getName(); auto metricName = HNSWIndexConfig::metricToString(config.metric); cypher += std::format("CALL CREATE_VECTOR_INDEX('{}', '{}', '{}', mu := {}, ml := {}, " - "pu := {}, metric := '{}', alpha := {}, efc := {});", + "pu := {}, metric := '{}', alpha := {}, efc := {}", tableName, indexEntry.getIndexName(), propertyName, config.mu, config.ml, config.pu, metricName, config.alpha, config.efc); + if (!config.cacheEmbeddingsColumn) { + cypher += ", cache_embeddings := false"; + } + if (config.quantization != QuantizationType::NONE) { + cypher += std::format(", quantization := '{}'", + HNSWIndexConfig::quantizationToString(config.quantization)); + } + if (config.storeFullPrecisionEmbeddings) { + cypher += ", use_full_precision_rerank := true"; + } + cypher += ");"; return cypher; } diff --git a/vector/src/function/create_hnsw_index.cpp b/vector/src/function/create_hnsw_index.cpp index 16d5f096..61f71ef5 100644 --- a/vector/src/function/create_hnsw_index.cpp +++ b/vector/src/function/create_hnsw_index.cpp @@ -34,7 +34,7 @@ CreateInMemHNSWSharedState::CreateInMemHNSWSharedState(const CreateHNSWIndexBind ->getTable(bindData.tableEntry->getTableID()) ->cast()}, numNodes{bindData.numRows}, bindData{&bindData} { - storage::IndexInfo dummyIndexInfo{"", "", bindData.tableEntry->getTableID(), + storage::IndexInfo dummyIndexInfo{bindData.indexName, "", bindData.tableEntry->getTableID(), {bindData.tableEntry->getColumnID(bindData.propertyID)}, {PhysicalTypeID::ARRAY}, false, false}; hnswIndex = std::make_shared(bindData.context, dummyIndexInfo, @@ -285,15 +285,39 @@ static void finalizeHNSWTableFinalizeFunc(const ExecutionContext* context, {columnID}, {PhysicalTypeID::ARRAY}, hnswIndexType.constraintType == storage::IndexConstraintType::PRIMARY, hnswIndexType.definitionType == storage::IndexDefinitionType::BUILTIN}; + auto storageManager = storage::StorageManager::Get(*clientContext); auto upperTableID = upperTableEntry.getSingleRelEntryInfo().oid; auto lowerTableID = lowerTableEntry.getSingleRelEntryInfo().oid; + auto quantizedEmbeddingsTableID = common::INVALID_TABLE_ID; + storage::NodeTable* quantizedEmbeddingsTable = nullptr; + if (bindData->config.quantization != QuantizationType::NONE) { + auto quantizedEmbeddingsTableName = + HNSWIndexUtils::getQuantizedEmbeddingsTableName(nodeTableID, bindData->indexName); + auto quantizedEmbeddingsTableEntry = + catalog->getTableCatalogEntry(transaction, quantizedEmbeddingsTableName, true) + ->ptrCast(); + quantizedEmbeddingsTableID = quantizedEmbeddingsTableEntry->getTableID(); + quantizedEmbeddingsTable = + storageManager->getTable(quantizedEmbeddingsTableID)->ptrCast(); + } auto storageInfo = std::make_unique(upperTableID, lowerTableID, - index->getUpperEntryPoint(), index->getLowerEntryPoint(), bindData->numRows); + quantizedEmbeddingsTableID, index->getUpperEntryPoint(), index->getLowerEntryPoint(), + bindData->numRows); auto onDiskIndex = std::make_unique(context->clientContext, indexInfo, std::move(storageInfo), bindData->config.copy()); - auto storageManager = storage::StorageManager::Get(*clientContext); auto nodeTable = storageManager->getTable(nodeTableID)->ptrCast(); nodeTable->addIndex(std::move(onDiskIndex)); + if (bindData->config.quantization != QuantizationType::NONE) { + DASSERT(quantizedEmbeddingsTable != nullptr); + const auto& columnType = nodeTable->getColumn(columnID).getDataType(); + const auto typeInfo = + columnType.getExtraTypeInfo()->constPtrCast(); + auto quantizedEmbeddings = TableBackedQuantizedEmbeddings(context->clientContext, + common::ArrayTypeInfo{typeInfo->getChildType().copy(), typeInfo->getNumElements()}, + *nodeTable, columnID, *quantizedEmbeddingsTable, bindData->config.quantization, + bindData->config.metric, bindData->numRows); + quantizedEmbeddings.rebuild(context->clientContext); + } index->moveToPartitionState(*hnswSharedState->partitionerSharedState); transaction->setForceCheckpoint(); } @@ -346,8 +370,21 @@ static std::string rewriteCreateHNSWQuery(main::ClientContext& context, HNSWIndexUtils::getUpperGraphTableName(tableID, indexName), tableName, tableName); query += std::format("CREATE REL TABLE {} (FROM {} TO {}) WITH (storage_direction='fwd');", HNSWIndexUtils::getLowerGraphTableName(tableID, indexName), tableName, tableName); - std::string params; auto& config = hnswBindData->config; + if (config.quantization != QuantizationType::NONE) { + const auto& columnType = + hnswBindData->tableEntry->getProperty(hnswBindData->propertyID).getType(); + const auto typeInfo = + columnType.getExtraTypeInfo()->constPtrCast(); + const auto payloadType = + config.quantization == QuantizationType::SQ8 ? "INT8" : "INT16"; + query += std::format( + "CREATE NODE TABLE {} (id INT64, valid UINT8, scale FLOAT, norm_sq FLOAT, payload " + "{}[{}], PRIMARY KEY (id));", + HNSWIndexUtils::getQuantizedEmbeddingsTableName(tableID, indexName), payloadType, + typeInfo->getNumElements()); + } + std::string params; params += std::format("mu := {}, ", config.mu); params += std::format("ml := {}, ", config.ml); params += std::format("efc := {}, ", config.efc); @@ -356,8 +393,15 @@ static std::string rewriteCreateHNSWQuery(main::ClientContext& context, params += std::format("pu := {}, ", config.pu); params += std::format("cache_embeddings := {}", config.cacheEmbeddingsColumn ? "true" : "false"); + if (config.quantization != QuantizationType::NONE) { + params += std::format(", quantization := '{}'", + HNSWIndexConfig::quantizationToString(config.quantization)); + } + if (config.storeFullPrecisionEmbeddings) { + params += ", use_full_precision_rerank := true"; + } auto columnName = hnswBindData->tableEntry->getProperty(hnswBindData->propertyID).getName(); - if (config.cacheEmbeddingsColumn) { + if (config.cacheEmbeddingsColumn && config.quantization == QuantizationType::NONE) { query += std::format("CALL _CACHE_ARRAY_COLUMN_LOCALLY('{}', '{}');", tableName, columnName); } diff --git a/vector/src/function/drop_hnsw_index.cpp b/vector/src/function/drop_hnsw_index.cpp index 9c715a2b..a663038c 100644 --- a/vector/src/function/drop_hnsw_index.cpp +++ b/vector/src/function/drop_hnsw_index.cpp @@ -1,5 +1,6 @@ #include "catalog/catalog.h" #include "catalog/catalog_entry/node_table_catalog_entry.h" +#include "catalog/hnsw_index_catalog_entry.h" #include "common/exception/binder.h" #include "function/hnsw_index_functions.h" #include "function/table/bind_data.h" @@ -18,15 +19,18 @@ namespace vector_extension { struct DropHNSWIndexBindData final : TableFuncBindData { catalog::NodeTableCatalogEntry* tableEntry; std::string indexName; + bool hasQuantizedEmbeddingsTable; bool skipAfterBind; DropHNSWIndexBindData(catalog::NodeTableCatalogEntry* tableEntry, std::string indexName, - bool skipAfterBind = false) + bool hasQuantizedEmbeddingsTable = false, bool skipAfterBind = false) : TableFuncBindData{0}, tableEntry{tableEntry}, indexName{std::move(indexName)}, + hasQuantizedEmbeddingsTable{hasQuantizedEmbeddingsTable}, skipAfterBind{skipAfterBind} {} std::unique_ptr copy() const override { - return std::make_unique(tableEntry, indexName, skipAfterBind); + return std::make_unique(tableEntry, indexName, + hasQuantizedEmbeddingsTable, skipAfterBind); } }; @@ -38,9 +42,15 @@ static std::unique_ptr bindFunc(main::ClientContext* context, const auto nodeTableEntry = HNSWIndexUtils::bindNodeTable(*context, tableName); if (!HNSWIndexUtils::validateIndexExistence(*context, nodeTableEntry, indexName, HNSWIndexUtils::IndexOperation::DROP, config.conflictAction)) { - return std::make_unique(nullptr, indexName, true); + return std::make_unique(nullptr, indexName, false, true); } - return std::make_unique(nodeTableEntry, indexName); + auto transaction = transaction::Transaction::Get(*context); + auto indexEntry = + catalog::Catalog::Get(*context)->getIndex(transaction, nodeTableEntry->getTableID(), + indexName); + const auto auxInfo = indexEntry->getAuxInfo().cast(); + return std::make_unique(nodeTableEntry, indexName, + auxInfo.config.quantization != QuantizationType::NONE); } static common::offset_t internalTableFunc(const TableFuncInput& input, TableFuncOutput&) { @@ -74,6 +84,11 @@ static std::string dropHNSWIndexTables(main::ClientContext& context, HNSWIndexUtils::getUpperGraphTableName(nodeTableID, dropHNSWIndexBindData->indexName)); query += std::format("DROP TABLE {};", HNSWIndexUtils::getLowerGraphTableName(nodeTableID, dropHNSWIndexBindData->indexName)); + if (dropHNSWIndexBindData->hasQuantizedEmbeddingsTable) { + query += std::format("DROP TABLE {};", + HNSWIndexUtils::getQuantizedEmbeddingsTableName(nodeTableID, + dropHNSWIndexBindData->indexName)); + } if (requireNewTransaction) { query += "COMMIT;"; } diff --git a/vector/src/function/query_hnsw_index.cpp b/vector/src/function/query_hnsw_index.cpp index 7066f4ba..23e3c1c3 100644 --- a/vector/src/function/query_hnsw_index.cpp +++ b/vector/src/function/query_hnsw_index.cpp @@ -1,3 +1,6 @@ +#include +#include + #include "binder/binder.h" #include "binder/expression/expression_util.h" #include "binder/expression/literal_expression.h" @@ -207,8 +210,8 @@ static const LogicalType& getIndexColumnType(const NodeTableCatalogEntry& nodeEn // This struct wraps a vector of embedding data // It exists so that we can match the interface for on-disk HNSW search template -struct HNSWQueryVector : GetEmbeddingsScanState { - HNSWQueryVector(main::ClientContext* context, std::shared_ptr queryExpression, +struct HNSWRawQueryVector : GetEmbeddingsScanState { + HNSWRawQueryVector(main::ClientContext* context, std::shared_ptr queryExpression, const LogicalType& indexType, uint64_t dimension) : data(getQueryVector(context, std::move(queryExpression), indexType, dimension)) {} @@ -224,6 +227,81 @@ struct HNSWQueryVector : GetEmbeddingsScanState { std::vector data; }; +template +static void quantizeQueryVector(const std::vector& src, QuantizationType quantization, + MetricType metric, uint8_t* payloadDst, float& scaleDst, float& normSqDst) { + constexpr float zeroScale = 0.0f; + constexpr float zeroNormSq = 0.0f; + double maxAbs = 0.0; + for (const auto value : src) { + maxAbs = std::max(maxAbs, std::abs(static_cast(value))); + } + if (maxAbs == 0.0) { + scaleDst = zeroScale; + normSqDst = zeroNormSq; + std::memset(payloadDst, 0, + HNSWIndexUtils::getQuantizedCachedEmbeddingPayloadBytes(src.size(), quantization)); + return; + } + const auto maxQuantizedValue = quantization == QuantizationType::SQ8 ? 127.0 : 32767.0; + const auto scale = static_cast(maxAbs / maxQuantizedValue); + float normSq = 0.0f; + if (quantization == QuantizationType::SQ8) { + for (auto i = 0u; i < src.size(); ++i) { + const auto quantized = + static_cast(std::llround(static_cast(src[i]) / scale)); + const auto clamped = static_cast(std::clamp(quantized, -127, 127)); + payloadDst[i] = static_cast(clamped); + normSq += static_cast(clamped) * static_cast(clamped); + } + normSqDst = normSq; + scaleDst = metric == MetricType::Cosine && normSq > 0.0f ? 1.0f / std::sqrt(normSq) : + scale; + return; + } + auto* payloadValues = reinterpret_cast(payloadDst); + for (auto i = 0u; i < src.size(); ++i) { + const auto quantized = + static_cast(std::llround(static_cast(src[i]) / scale)); + const auto clamped = static_cast(std::clamp(quantized, -32767, 32767)); + normSq += static_cast(clamped) * static_cast(clamped); + payloadValues[i] = clamped; + } + normSqDst = normSq; + scaleDst = metric == MetricType::Cosine && normSq > 0.0f ? 1.0f / std::sqrt(normSq) : scale; +} + +template +struct HNSWQuantizedQueryVector : GetEmbeddingsScanState { + HNSWQuantizedQueryVector(main::ClientContext* context, + std::shared_ptr queryExpression, const LogicalType& indexType, + uint64_t dimension, QuantizationType quantization, MetricType metric) + : data(HNSWIndexUtils::getQuantizedCachedEmbeddingStride(dimension, quantization) + 31), + view{nullptr, 0.0f, 0.0f} { + const auto rawVector = + getQueryVector(context, std::move(queryExpression), indexType, dimension); + alignedData = reinterpret_cast( + common::ceilDiv(reinterpret_cast(data.data()), + static_cast(32)) * + 32); + quantizeQueryVector(rawVector, quantization, metric, alignedData, view.scale, view.normSq); + view.payload = alignedData; + } + + void* getEmbeddingPtr([[maybe_unused]] const EmbeddingHandle& handle) override { + DASSERT(!handle.isNull()); + DASSERT(handle.offsetInData == 0); + return &view; + } + + void addEmbedding(const EmbeddingHandle&) override {}; + void reclaimEmbedding(const EmbeddingHandle&) override {}; + + std::vector data; + uint8_t* alignedData; + QuantizedEmbeddingView view; +}; + static offset_t tableFunc(const TableFuncInput& input, TableFuncOutput& output) { const auto localState = input.localState->ptrCast(); const auto bindData = input.bindData->constPtrCast(); @@ -242,12 +320,22 @@ static offset_t tableFunc(const TableFuncInput& input, TableFuncOutput& output) TypeUtils::visit( indexType, [&](T) { - auto queryVector = HNSWQueryVector(input.context->clientContext, + auto exactQueryVector = HNSWRawQueryVector(input.context->clientContext, bindData->queryExpression, index.getElementType(), dimension); + auto exactQueryVectorHandle = EmbeddingHandle{0, &exactQueryVector}; + if (index.getQuantization() == QuantizationType::NONE) { + localState->result = index.search( + transaction::Transaction::Get(*input.context->clientContext), + exactQueryVectorHandle, exactQueryVectorHandle, localState->searchState); + return; + } + auto queryVector = HNSWQuantizedQueryVector(input.context->clientContext, + bindData->queryExpression, index.getElementType(), dimension, + index.getQuantization(), index.getMetric()); auto queryVectorHandle = EmbeddingHandle{0, &queryVector}; localState->result = index.search(transaction::Transaction::Get(*input.context->clientContext), - queryVectorHandle, localState->searchState); + queryVectorHandle, exactQueryVectorHandle, localState->searchState); }, [&](auto) { UNREACHABLE_CODE; }); } @@ -301,9 +389,17 @@ std::unique_ptr initQueryHNSWLocalState( catalog ->getTableCatalogEntry(transaction::Transaction::Get(*context), lowerRelTableName, true) ->ptrCast(); + const auto& indexConfig = + hnswBindData->indexEntry->getAuxInfo().cast().config; + auto indexOpt = hnswSharedState->nodeTable->getIndex(hnswBindData->indexEntry->getIndexName()); + DASSERT(indexOpt.has_value()); + auto& index = indexOpt.value()->cast(); + auto quantizedEmbeddings = + index.getOrCreateQuantizedEmbeddings(context, hnswSharedState->numNodes); HNSWSearchState searchState{context, hnswBindData->nodeTableEntry, upperRelTableEntry, lowerRelTableEntry, *hnswSharedState->nodeTable, hnswBindData->indexColumnID, - hnswSharedState->numNodes, static_cast(k), hnswBindData->config}; + hnswSharedState->numNodes, static_cast(k), hnswBindData->config, indexConfig, + true, std::move(quantizedEmbeddings)}; const auto tableID = hnswBindData->nodeTableEntry->getTableID(); auto& semiMasks = hnswSharedState->semiMasks; if (semiMasks.containsTableID(tableID)) { diff --git a/vector/src/include/index/hnsw_config.h b/vector/src/include/index/hnsw_config.h index b0ef6bb9..f0431972 100644 --- a/vector/src/include/index/hnsw_config.h +++ b/vector/src/include/index/hnsw_config.h @@ -11,6 +11,8 @@ namespace vector_extension { enum class MetricType : uint8_t { Cosine = 0, L2 = 1, L2_SQUARE = 2, DotProduct = 3 }; +enum class QuantizationType : uint8_t { NONE = 0, SQ8 = 1, SQ16 = 2 }; + // We use this ratio to calculate the max degree of the upper/lower graph based on the user provided // max degree value for the upper/lower graph, respectively. static constexpr double DEFAULT_DEGREE_THRESHOLD_RATIO = 1.25; @@ -80,6 +82,22 @@ struct CacheEmbeddings { static constexpr bool DEFAULT_VALUE = true; }; +struct Quantization { + static constexpr const char* NAME = "quantization"; + static constexpr common::LogicalTypeID TYPE = common::LogicalTypeID::STRING; + static constexpr QuantizationType DEFAULT_VALUE = QuantizationType::NONE; + + static void validate(const std::string& value); +}; + +struct UseFullPrecisionRerank { + static constexpr const char* NAME = "use_full_precision_rerank"; + static constexpr common::LogicalTypeID TYPE = common::LogicalTypeID::BOOL; + static constexpr bool DEFAULT_VALUE = false; + + static void validate(bool value); +}; + struct SkipIfExists { static constexpr const char* NAME = "skip_if_exists"; static constexpr common::LogicalTypeID TYPE = common::LogicalTypeID::BOOL; @@ -136,6 +154,8 @@ struct HNSWIndexConfig { double alpha = Alpha::DEFAULT_VALUE; int64_t efc = Efc::DEFAULT_VALUE; bool cacheEmbeddingsColumn = CacheEmbeddings::DEFAULT_VALUE; + QuantizationType quantization = Quantization::DEFAULT_VALUE; + bool storeFullPrecisionEmbeddings = UseFullPrecisionRerank::DEFAULT_VALUE; common::ConflictAction conflictAction = SkipIfExists::DEFAULT_VALUE; HNSWIndexConfig() = default; @@ -149,14 +169,18 @@ struct HNSWIndexConfig { static HNSWIndexConfig deserialize(common::Deserializer& deSer); static std::string metricToString(MetricType metric); + static std::string quantizationToString(QuantizationType quantization); private: HNSWIndexConfig(const HNSWIndexConfig& other) : mu{other.mu}, ml{other.ml}, pu{other.pu}, metric{other.metric}, alpha{other.alpha}, efc{other.efc}, cacheEmbeddingsColumn(other.cacheEmbeddingsColumn), + quantization{other.quantization}, + storeFullPrecisionEmbeddings{other.storeFullPrecisionEmbeddings}, conflictAction(other.conflictAction) {} static MetricType getMetricType(const std::string& metricName); + static QuantizationType getQuantizationType(const std::string& quantizationName); }; struct DropHNSWConfig { diff --git a/vector/src/include/index/hnsw_graph.h b/vector/src/include/index/hnsw_graph.h index 55d44081..37a71e1e 100644 --- a/vector/src/include/index/hnsw_graph.h +++ b/vector/src/include/index/hnsw_graph.h @@ -2,16 +2,24 @@ #include #include +#include +#include +#include #include "common/data_chunk/data_chunk.h" +#include "index/hnsw_index_utils.h" #include "processor/operator/base_partitioner_shared_state.h" #include "storage/buffer_manager/memory_manager.h" #include "storage/local_cached_column.h" +#include "transaction/transaction.h" namespace lbug { namespace storage { struct NodeTableScanState; } +namespace main { +class ClientContext; +} namespace vector_extension { struct EmbeddingColumnInfo { @@ -73,7 +81,9 @@ class HNSWIndexEmbeddings { GetEmbeddingsScanState& scanState) const = 0; common::length_t getDimension() const { return info.getDimension(); } - virtual std::unique_ptr constructScanState() const = 0; + virtual metric_func_t getMetricFunction(MetricType metric) const; + virtual std::unique_ptr constructScanState( + transaction::Transaction* transaction = nullptr) const = 0; protected: EmbeddingColumnInfo info; @@ -89,12 +99,122 @@ class InMemEmbeddings final : public HNSWIndexEmbeddings { std::vector getEmbeddings(std::span offset, GetEmbeddingsScanState& scanState) const override; - std::unique_ptr constructScanState() const override; + std::unique_ptr constructScanState( + transaction::Transaction* transaction = nullptr) const override; private: storage::CachedColumn* data; }; +class QuantizedInMemEmbeddings final : public HNSWIndexEmbeddings { +public: + QuantizedInMemEmbeddings(const main::ClientContext* context, common::ArrayTypeInfo typeInfo, + storage::NodeTable& table, common::column_id_t columnID, QuantizationType quantization, + MetricType metric); + + EmbeddingHandle getEmbedding(common::offset_t offset, + GetEmbeddingsScanState& scanState) const override; + std::vector getEmbeddings(std::span offset, + GetEmbeddingsScanState& scanState) const override; + std::unique_ptr constructScanState( + transaction::Transaction* transaction = nullptr) const override; + + metric_func_t getMetricFunction(MetricType metric) const override; + +private: + QuantizationType quantization; + uint64_t strideBytes; + std::vector nullMask; + std::unique_ptr payloadBuffer; + uint8_t* payloadDataPtr; + std::vector scales; + std::vector normSqs; + std::vector views; +}; + +class CachedQuantizedColumn final : public transaction::LocalCacheObject { +public: + static std::string getKey(common::table_id_t quantizedTableID) { + return std::format("quantized-{}", quantizedTableID); + } + + CachedQuantizedColumn(const main::ClientContext* context, common::ArrayTypeInfo typeInfo, + storage::NodeTable& sourceTable, storage::NodeTable& quantizedTable, + common::column_id_t sourceColumnID, QuantizationType quantization, MetricType metric, + common::offset_t numNodes); + DELETE_BOTH_COPY(CachedQuantizedColumn); + + bool isNull(common::offset_t offset) const { return nullMask[offset] != 0; } + const QuantizedEmbeddingView* getViews() const { return views.data(); } + common::offset_t getNumNodes() const { return nullMask.size(); } + +private: + void setQuantizedRecord(common::offset_t offset, uint8_t valid, float scale, float normSq, + const uint8_t* payload); + + common::ArrayTypeInfo typeInfo; + uint64_t strideBytes; + uint64_t payloadBytes; + std::vector nullMask; + std::unique_ptr payloadBuffer; + uint8_t* payloadDataPtr; + std::vector scales; + std::vector normSqs; + std::vector views; +}; + +class CachedQuantizedEmbeddings final : public HNSWIndexEmbeddings { +public: + CachedQuantizedEmbeddings(common::ArrayTypeInfo typeInfo, storage::NodeTable& sourceTable, + std::shared_ptr cachedColumn, QuantizationType quantization); + + EmbeddingHandle getEmbedding(common::offset_t offset, + GetEmbeddingsScanState& scanState) const override; + std::vector getEmbeddings(std::span offset, + GetEmbeddingsScanState& scanState) const override; + std::unique_ptr constructScanState( + transaction::Transaction* transaction = nullptr) const override; + metric_func_t getMetricFunction(MetricType metric) const override; + +private: + QuantizationType quantization; + storage::NodeTable& nodeTable; + std::shared_ptr cachedColumn; +}; + +class TableBackedQuantizedEmbeddings final : public HNSWIndexEmbeddings { +public: + TableBackedQuantizedEmbeddings(const main::ClientContext* context, + common::ArrayTypeInfo typeInfo, storage::NodeTable& sourceTable, + common::column_id_t sourceColumnID, storage::NodeTable& quantizedTable, + QuantizationType quantization, MetricType metric, common::offset_t numNodes); + + EmbeddingHandle getEmbedding(common::offset_t offset, + GetEmbeddingsScanState& scanState) const override; + std::vector getEmbeddings(std::span offset, + GetEmbeddingsScanState& scanState) const override; + std::unique_ptr constructScanState( + transaction::Transaction* transaction = nullptr) const override; + metric_func_t getMetricFunction(MetricType metric) const override; + + void rebuild(const main::ClientContext* context) const; + void writeEmbedding(transaction::Transaction* transaction, common::offset_t offset, + const EmbeddingHandle& embedding) const; + +private: + void writeQuantizedRecord(transaction::Transaction* transaction, common::offset_t offset, + uint8_t valid, float scale, float normSq, const uint8_t* payload) const; + + QuantizationType quantization; + MetricType metric; + storage::NodeTable& sourceTable; + common::column_id_t sourceColumnID; + storage::NodeTable& quantizedTable; + transaction::Transaction* defaultTransaction; + common::offset_t numNodes; + uint64_t strideBytes; +}; + class OnDiskEmbeddingScanState final : public GetEmbeddingsScanState { public: OnDiskEmbeddingScanState(const transaction::Transaction* transaction, @@ -129,7 +249,8 @@ class OnDiskEmbeddings final : public HNSWIndexEmbeddings { GetEmbeddingsScanState& scanState) const override; std::vector getEmbeddings(std::span offset, GetEmbeddingsScanState& scanState) const override; - std::unique_ptr constructScanState() const override; + std::unique_ptr constructScanState( + transaction::Transaction* transaction = nullptr) const override; private: transaction::Transaction* transaction; diff --git a/vector/src/include/index/hnsw_index.h b/vector/src/include/index/hnsw_index.h index 002bb1f9..a3ac95bd 100644 --- a/vector/src/include/index/hnsw_index.h +++ b/vector/src/include/index/hnsw_index.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include "common/random_engine.h" @@ -25,6 +26,7 @@ namespace vector_extension { struct HNSWIndexPartitionerSharedState; struct CreateInMemHNSWLocalState; +struct QuantizedEmbeddingsCacheState; struct MinNodePriorityQueueComparator { bool operator()(const NodeWithDistance& l, const NodeWithDistance& r) const { @@ -61,18 +63,21 @@ struct VisitedState { struct HNSWStorageInfo final : storage::IndexStorageInfo { common::table_id_t upperRelTableID; common::table_id_t lowerRelTableID; + common::table_id_t quantizedEmbeddingsTableID; common::offset_t upperEntryPoint; common::offset_t lowerEntryPoint; common::offset_t numCheckpointedNodes; HNSWStorageInfo() : upperRelTableID{common::INVALID_TABLE_ID}, lowerRelTableID{common::INVALID_TABLE_ID}, + quantizedEmbeddingsTableID{common::INVALID_TABLE_ID}, upperEntryPoint{common::INVALID_OFFSET}, lowerEntryPoint{common::INVALID_OFFSET}, numCheckpointedNodes{0} {} HNSWStorageInfo(common::table_id_t upperRelTableID, common::table_id_t lowerRelTableID, - common::offset_t upperEntryPoint, common::offset_t lowerEntryPoint, - common::offset_t numCheckpointedNodes) + common::table_id_t quantizedEmbeddingsTableID, common::offset_t upperEntryPoint, + common::offset_t lowerEntryPoint, common::offset_t numCheckpointedNodes) : upperRelTableID{upperRelTableID}, lowerRelTableID{lowerRelTableID}, + quantizedEmbeddingsTableID{quantizedEmbeddingsTableID}, upperEntryPoint{upperEntryPoint}, lowerEntryPoint{lowerEntryPoint}, numCheckpointedNodes{numCheckpointedNodes} {} @@ -90,8 +95,9 @@ class HNSWIndex : public storage::Index { HNSWIndexConfig config, common::ArrayTypeInfo typeInfo) : Index{std::move(indexInfo), std::move(storageInfo)}, config{std::move(config)}, typeInfo{std::move(typeInfo)} { - metricFunc = + exactMetricFunc = HNSWIndexUtils::getMetricsFunction(this->config.metric, this->typeInfo.getChildType()); + metricFunc = exactMetricFunc; } ~HNSWIndex() override = default; @@ -103,6 +109,8 @@ class HNSWIndex : public storage::Index { } common::LogicalType getElementType() const { return typeInfo.getChildType().copy(); } + QuantizationType getQuantization() const { return config.quantization; } + MetricType getMetric() const { return config.metric; } static std::vector popTopK(max_node_priority_queue_t& result, common::length_t k); @@ -129,6 +137,7 @@ class HNSWIndex : public storage::Index { HNSWIndexConfig config; common::ArrayTypeInfo typeInfo; metric_func_t metricFunc; + metric_func_t exactMetricFunc; common::RandomEngine randomEngine; }; @@ -256,8 +265,11 @@ enum class SearchType : uint8_t { struct HNSWSearchState { VisitedState visited; - std::unique_ptr embeddings; - OnDiskEmbeddingScanState embeddingScanState; + std::shared_ptr embeddings; + std::unique_ptr embeddingScanState; + metric_func_t metricFunc; + std::shared_ptr exactEmbeddings; + std::unique_ptr exactEmbeddingScanState; uint64_t k; QueryHNSWConfig config; uint64_t ef; @@ -274,8 +286,9 @@ struct HNSWSearchState { HNSWSearchState(main::ClientContext* context, catalog::TableCatalogEntry* nodeTableEntry, catalog::TableCatalogEntry* upperRelTableEntry, catalog::TableCatalogEntry* lowerRelTableEntry, storage::NodeTable& nodeTable, - common::column_id_t columnID, common::offset_t numNodes, uint64_t k, - QueryHNSWConfig config); + common::column_id_t columnID, common::offset_t numNodes, uint64_t k, QueryHNSWConfig config, + const HNSWIndexConfig& indexConfig, bool useQuantizedEmbeddings, + std::shared_ptr cachedEmbeddings = nullptr); bool isMasked(common::offset_t offset) const { return !hasMask() || semiMask->isMasked(offset); @@ -315,7 +328,8 @@ class OnDiskHNSWIndex final : public HNSWIndex { std::unique_ptr storageInfo, HNSWIndexConfig config); std::vector search(transaction::Transaction* transaction, - const EmbeddingHandle& queryVector, HNSWSearchState& searchState) const; + const EmbeddingHandle& queryVector, const EmbeddingHandle& exactQueryVector, + HNSWSearchState& searchState) const; static std::unique_ptr load(main::ClientContext* context, storage::StorageManager* storageManager, storage::IndexInfo indexInfo, @@ -338,7 +352,10 @@ class OnDiskHNSWIndex final : public HNSWIndex { } void finalize(main::ClientContext*) override; - void checkpoint(main::ClientContext* context, storage::PageAllocator& pageAllocator, storage::ShadowFile& shadowFile) override; + void checkpoint(main::ClientContext* context, storage::PageAllocator& pageAllocator, + storage::ShadowFile& shadowFile) override; + std::shared_ptr getOrCreateQuantizedEmbeddings( + const main::ClientContext* context, common::offset_t numNodes) const; private: common::offset_t searchNNInUpperLayer(const EmbeddingHandle& queryVector, @@ -347,7 +364,8 @@ class OnDiskHNSWIndex final : public HNSWIndex { const EmbeddingHandle& queryVector, common::offset_t entryNode, HNSWSearchState& searchState, bool isUpperLayer) const; std::vector searchBruteForce(transaction::Transaction* transaction, - const EmbeddingHandle& queryVector, HNSWSearchState& searchState) const; + const EmbeddingHandle& approxQueryVector, const EmbeddingHandle& exactQueryVector, + HNSWSearchState& searchState) const; std::vector searchFromCheckpointed(transaction::Transaction* transaction, const EmbeddingHandle& queryVector, HNSWSearchState& searchState) const; void searchFromUnCheckpointed(transaction::Transaction* transaction, @@ -416,6 +434,9 @@ class OnDiskHNSWIndex final : public HNSWIndex { void initSearchCandidates(const EmbeddingHandle& queryVector, HNSWSearchState& searchState, min_node_priority_queue_t& candidates, max_node_priority_queue_t& results) const; + void exactRerankResults(const EmbeddingHandle& exactQueryVector, HNSWSearchState& searchState, + std::vector& result) const; + void markQuantizedEmbeddingsCacheDirty(transaction::Transaction* transaction) const; static SearchType getFilteredSearchType(transaction::Transaction* transaction, const HNSWSearchState& searchState); @@ -426,6 +447,8 @@ class OnDiskHNSWIndex final : public HNSWIndex { storage::MemoryManager* mm; storage::NodeTable& nodeTable; + storage::NodeTable* quantizedEmbeddingsTable; + std::shared_ptr quantizedEmbeddingsCacheState; storage::RelTable* upperRelTable; storage::RelTable* lowerRelTable; }; diff --git a/vector/src/include/index/hnsw_index_utils.h b/vector/src/include/index/hnsw_index_utils.h index fdd9f2a0..079f4817 100644 --- a/vector/src/include/index/hnsw_index_utils.h +++ b/vector/src/include/index/hnsw_index_utils.h @@ -18,6 +18,12 @@ template concept VectorElementType = std::is_floating_point_v; using metric_func_t = std::function; +struct QuantizedEmbeddingView { + const uint8_t* payload; + float scale; + float normSq; +}; + struct HNSWIndexUtils { enum class IndexOperation { CREATE, QUERY, DROP }; @@ -36,7 +42,19 @@ struct HNSWIndexUtils { static void validateAutoTransaction(const main::ClientContext& context, const std::string& funcName); - static metric_func_t getMetricsFunction(MetricType metric, const common::LogicalType& type); + static metric_func_t getMetricsFunction(MetricType metric, const common::LogicalType& type, + QuantizationType quantization = QuantizationType::NONE); + static metric_func_t getQuantizedCachedMetricsFunction(MetricType metric, + QuantizationType quantization); + static uint64_t getQuantizedEmbeddingPayloadOffset(QuantizationType quantization); + static uint64_t getQuantizedEmbeddingStride(uint64_t dimension, + QuantizationType quantization); + static uint64_t getQuantizedCachedEmbeddingStride(uint64_t dimension, + QuantizationType quantization); + static uint64_t getQuantizedEmbeddingPayloadBytes(uint64_t dimension, + QuantizationType quantization); + static uint64_t getQuantizedCachedEmbeddingPayloadBytes(uint64_t dimension, + QuantizationType quantization); static void validateColumnType(const catalog::TableCatalogEntry& tableEntry, const std::string& columnName); @@ -49,6 +67,10 @@ struct HNSWIndexUtils { const std::string& indexName) { return "_" + std::to_string(tableID) + "_" + indexName + "_LOWER"; } + static std::string getQuantizedEmbeddingsTableName(common::table_id_t tableID, + const std::string& indexName) { + return "_" + std::to_string(tableID) + "_" + indexName + "_QEMB"; + } template static T* getEmbedding(const storage::ColumnChunkData& embeddings, common::offset_t offset, diff --git a/vector/src/index/hnsw_config.cpp b/vector/src/index/hnsw_config.cpp index 47315b50..a3bfc20a 100644 --- a/vector/src/index/hnsw_config.cpp +++ b/vector/src/index/hnsw_config.cpp @@ -100,6 +100,17 @@ void UseKnn::validate(bool value) { (void)value; } +void Quantization::validate(const std::string& value) { + const auto lowerCaseValue = common::StringUtils::getLower(value); + if (lowerCaseValue != "none" && lowerCaseValue != "sq8" && lowerCaseValue != "sq16") { + throw common::BinderException{"Quantization must be one of NONE, SQ8 or SQ16."}; + } +} + +void UseFullPrecisionRerank::validate(bool value) { + (void)value; +} + HNSWIndexConfig::HNSWIndexConfig(const function::optional_params_t& optionalParams) { for (auto& [name, value] : optionalParams) { auto lowerCaseName = common::StringUtils::getLower(name); @@ -131,6 +142,15 @@ HNSWIndexConfig::HNSWIndexConfig(const function::optional_params_t& optionalPara } else if (CacheEmbeddings::NAME == lowerCaseName) { value.validateType(CacheEmbeddings::TYPE); cacheEmbeddingsColumn = value.getValue(); + } else if (Quantization::NAME == lowerCaseName) { + value.validateType(Quantization::TYPE); + auto quantizationName = value.getValue(); + Quantization::validate(quantizationName); + quantization = getQuantizationType(quantizationName); + } else if (UseFullPrecisionRerank::NAME == lowerCaseName) { + value.validateType(UseFullPrecisionRerank::TYPE); + storeFullPrecisionEmbeddings = value.getValue(); + UseFullPrecisionRerank::validate(storeFullPrecisionEmbeddings); } else if (SkipIfExists::NAME == lowerCaseName) { value.validateType(SkipIfExists::TYPE); conflictAction = value.getValue() ? @@ -141,6 +161,28 @@ HNSWIndexConfig::HNSWIndexConfig(const function::optional_params_t& optionalPara name, CreateVectorIndexFunction::name)}; } } + if (quantization != QuantizationType::NONE && metric == MetricType::DotProduct) { + throw common::BinderException{ + "Quantization currently supports COSINE, L2 and L2SQ metrics only."}; + } +} + +std::string HNSWIndexConfig::quantizationToString(QuantizationType quantization) { + switch (quantization) { + case QuantizationType::NONE: { + return "none"; + } + case QuantizationType::SQ8: { + return "sq8"; + } + case QuantizationType::SQ16: { + return "sq16"; + } + default: { + throw common::RuntimeException(std::format("Unknown quantization type {}.", + static_cast(quantization))); + } + } } std::string HNSWIndexConfig::metricToString(MetricType metric) { @@ -175,6 +217,12 @@ void HNSWIndexConfig::serialize(common::Serializer& ser) const { ser.serializeValue(alpha); ser.writeDebuggingInfo("efc"); ser.serializeValue(efc); + ser.writeDebuggingInfo("cacheEmbeddingsColumn"); + ser.serializeValue(cacheEmbeddingsColumn); + ser.writeDebuggingInfo("quantization"); + ser.serializeValue(static_cast(quantization)); + ser.writeDebuggingInfo("useFullPrecisionRerank"); + ser.serializeValue(storeFullPrecisionEmbeddings); } HNSWIndexConfig HNSWIndexConfig::deserialize(common::Deserializer& deSer) { @@ -192,6 +240,23 @@ HNSWIndexConfig HNSWIndexConfig::deserialize(common::Deserializer& deSer) { deSer.deserializeValue(config.alpha); deSer.validateDebuggingInfo(debuggingInfo, "efc"); deSer.deserializeValue(config.efc); + if (deSer.finished()) { + return config; + } + deSer.validateDebuggingInfo(debuggingInfo, "cacheEmbeddingsColumn"); + deSer.deserializeValue(config.cacheEmbeddingsColumn); + if (deSer.finished()) { + return config; + } + deSer.validateDebuggingInfo(debuggingInfo, "quantization"); + uint8_t quantization = 0; + deSer.deserializeValue(quantization); + config.quantization = static_cast(quantization); + if (deSer.finished()) { + return config; + } + deSer.validateDebuggingInfo(debuggingInfo, "useFullPrecisionRerank"); + deSer.deserializeValue(config.storeFullPrecisionEmbeddings); return config; } @@ -213,6 +278,20 @@ MetricType HNSWIndexConfig::getMetricType(const std::string& metricName) { UNREACHABLE_CODE; } +QuantizationType HNSWIndexConfig::getQuantizationType(const std::string& quantizationName) { + const auto lowerQuantizationName = common::StringUtils::getLower(quantizationName); + if (lowerQuantizationName == "none") { + return QuantizationType::NONE; + } + if (lowerQuantizationName == "sq8") { + return QuantizationType::SQ8; + } + if (lowerQuantizationName == "sq16") { + return QuantizationType::SQ16; + } + UNREACHABLE_CODE; +} + DropHNSWConfig::DropHNSWConfig(const function::optional_params_t& optionalParams) { for (auto& [name, value] : optionalParams) { auto lowerCaseName = common::StringUtils::getLower(name); diff --git a/vector/src/index/hnsw_graph.cpp b/vector/src/index/hnsw_graph.cpp index d5a7cec2..d630ff4a 100644 --- a/vector/src/index/hnsw_graph.cpp +++ b/vector/src/index/hnsw_graph.cpp @@ -1,6 +1,10 @@ #include "index/hnsw_graph.h" +#include +#include + #include "index/hnsw_index_utils.h" +#include "main/client_context.h" #include "storage/local_cached_column.h" #include "storage/table/list_chunk_data.h" #include "storage/table/node_table.h" @@ -12,6 +16,12 @@ using namespace lbug::storage; namespace lbug { namespace vector_extension { +static uint8_t* alignPointer(uint8_t* ptr, uint64_t alignment) { + const auto addr = reinterpret_cast(ptr); + const auto alignedAddr = common::ceilDiv(addr, alignment) * alignment; + return reinterpret_cast(alignedAddr); +} + EmbeddingHandle::EmbeddingHandle(common::offset_t offset, GetEmbeddingsScanState* lifetimeManager) : offsetInData(offset), lifetimeManager(lifetimeManager) { if (lifetimeManager) { @@ -40,6 +50,10 @@ EmbeddingHandle& EmbeddingHandle::operator=(EmbeddingHandle&& o) noexcept { return *this; } +metric_func_t HNSWIndexEmbeddings::getMetricFunction(MetricType metric) const { + return HNSWIndexUtils::getMetricsFunction(metric, info.typeInfo.getChildType()); +} + struct InMemEmbeddingLocalState : GetEmbeddingsScanState { explicit InMemEmbeddingLocalState(const CachedColumn* data, const common::LogicalType& type) : data(data), type(type) {} @@ -53,6 +67,172 @@ struct InMemEmbeddingLocalState : GetEmbeddingsScanState { const common::LogicalType& type; }; +struct QuantizedInMemEmbeddingLocalState : GetEmbeddingsScanState { + explicit QuantizedInMemEmbeddingLocalState(const QuantizedEmbeddingView* views) + : views(views) {} + + void* getEmbeddingPtr(const EmbeddingHandle& handle) override { + return const_cast(&views[handle.offsetInData]); + } + + void addEmbedding(const EmbeddingHandle&) override {} + void reclaimEmbedding(const EmbeddingHandle&) override {} + + const QuantizedEmbeddingView* views; +}; + +struct QuantizedOnDiskEmbeddingScanState final : GetEmbeddingsScanState { + QuantizedOnDiskEmbeddingScanState(const transaction::Transaction* transaction, + const storage::NodeTable* nodeTable, const QuantizedEmbeddingView* views) + : transaction(transaction), nodeTable(nodeTable), views(views) {} + + void* getEmbeddingPtr(const EmbeddingHandle& handle) override { + return const_cast(&views[handle.offsetInData]); + } + + void addEmbedding(const EmbeddingHandle&) override {} + void reclaimEmbedding(const EmbeddingHandle&) override {} + + const transaction::Transaction* transaction; + const storage::NodeTable* nodeTable; + const QuantizedEmbeddingView* views; +}; + +struct PersistentQuantizedEmbeddingScanState final : GetEmbeddingsScanState { + struct OwnedEmbedding { + std::vector payload; + QuantizedEmbeddingView view; + uint64_t refCount = 0; + }; + + PersistentQuantizedEmbeddingScanState(const transaction::Transaction* transaction, + const storage::NodeTable* sourceTable, storage::NodeTable* quantizedTable, + uint64_t strideBytes, uint64_t payloadBytes) + : transaction(transaction), sourceTable(sourceTable), quantizedTable(quantizedTable), + nodeIDVector(common::LogicalType::INTERNAL_ID()), validVector(common::LogicalType::UINT8()), + scaleVector(common::LogicalType::FLOAT()), normSqVector(common::LogicalType::FLOAT()), + payloadVector(quantizedTable->getColumn(4).getDataType().copy()), + strideBytes(strideBytes), payloadBytes(payloadBytes) { + auto state = std::make_shared(); + nodeIDVector.state = state; + validVector.state = state; + scaleVector.state = state; + normSqVector.state = state; + payloadVector.state = state; + scanState = std::make_unique(&nodeIDVector, + std::vector{&validVector, &scaleVector, &normSqVector, &payloadVector}, state); + scanState->setToTable(transaction, quantizedTable, {1, 2, 3, 4}, {}); + } + + void* getEmbeddingPtr(const EmbeddingHandle& handle) override { + return &ownedEmbeddings.at(handle.offsetInData).view; + } + + void addEmbedding(const EmbeddingHandle& handle) override { + if (loadEmbedding(handle.offsetInData)) { + ownedEmbeddings.at(handle.offsetInData).refCount++; + } + } + + bool loadEmbedding(common::offset_t offset) { + if (auto it = ownedEmbeddings.find(offset); it != ownedEmbeddings.end()) { + return true; + } + if (missingEmbeddings.contains(offset)) { + return false; + } + nodeIDVector.setValue(0, common::internalID_t{offset, quantizedTable->getTableID()}); + nodeIDVector.state->getSelVectorUnsafe().setToUnfiltered(1); + quantizedTable->initScanState(const_cast(transaction), + *scanState, quantizedTable->getTableID(), offset); + if (!quantizedTable->lookup(transaction, *scanState)) { + missingEmbeddings.insert(offset); + return false; + } + if (validVector.isNull(0) || validVector.getValue(0) == 0) { + missingEmbeddings.insert(offset); + return false; + } + loadEmbeddingFromOutput(offset, 0); + return true; + } + + void loadEmbeddings(std::span offsets) { + common::sel_t numSelectedValues = 0; + auto& selVector = nodeIDVector.state->getSelVectorUnsafe(); + selVector.setToFiltered(); + for (auto i = 0u; i < offsets.size(); ++i) { + const auto offset = offsets[i]; + if (ownedEmbeddings.contains(offset) || missingEmbeddings.contains(offset)) { + continue; + } + if (!sourceTable->isVisible(transaction, offset)) { + missingEmbeddings.insert(offset); + continue; + } + nodeIDVector.setValue(i, common::internalID_t{offset, quantizedTable->getTableID()}); + selVector[numSelectedValues++] = i; + } + selVector.setSelSize(numSelectedValues); + if (numSelectedValues == 0) { + return; + } + // lookupMultiple takes Transaction* for scan-state initialization, although this path only + // reads transaction visibility metadata. Keep this cast safe by ensuring lookupMultiple and + // its callees do not mutate Transaction; pass a non-const transaction here if that changes. + [[maybe_unused]] const auto lookupSuccess = + quantizedTable->lookupMultiple(const_cast(transaction), + *scanState); + DASSERT(lookupSuccess); + for (auto i = 0u; i < selVector.getSelSize(); ++i) { + const auto pos = selVector[i]; + const auto offset = offsets[pos]; + if (validVector.isNull(pos) || validVector.getValue(pos) == 0 || + scaleVector.isNull(pos) || normSqVector.isNull(pos) || payloadVector.isNull(pos)) { + missingEmbeddings.insert(offset); + continue; + } + loadEmbeddingFromOutput(offset, pos); + } + } + + void loadEmbeddingFromOutput(common::offset_t offset, common::sel_t pos) { + auto owned = OwnedEmbedding{}; + owned.payload.resize(strideBytes); + const auto entry = payloadVector.getValue(pos); + std::memcpy(owned.payload.data(), common::ListVector::getListValues(&payloadVector, entry), + payloadBytes); + const auto scale = scaleVector.getValue(pos); + const auto normSq = normSqVector.getValue(pos); + owned.view = QuantizedEmbeddingView{owned.payload.data(), scale, normSq}; + ownedEmbeddings.emplace(offset, std::move(owned)); + } + + void reclaimEmbedding(const EmbeddingHandle& handle) override { + auto it = ownedEmbeddings.find(handle.offsetInData); + if (it == ownedEmbeddings.end()) { + return; + } + if (--it->second.refCount == 0) { + ownedEmbeddings.erase(it); + } + } + + const transaction::Transaction* transaction; + const storage::NodeTable* sourceTable; + storage::NodeTable* quantizedTable; + common::ValueVector nodeIDVector; + common::ValueVector validVector; + common::ValueVector scaleVector; + common::ValueVector normSqVector; + common::ValueVector payloadVector; + std::unique_ptr scanState; + uint64_t strideBytes; + uint64_t payloadBytes; + std::unordered_map ownedEmbeddings; + std::unordered_set missingEmbeddings; +}; + void* InMemEmbeddingLocalState::getEmbeddingPtr(const EmbeddingHandle& handle) { void* val = nullptr; if (!handle.isNull()) { @@ -81,7 +261,8 @@ InMemEmbeddings::InMemEmbeddings(transaction::Transaction* transaction, } } -std::unique_ptr InMemEmbeddings::constructScanState() const { +std::unique_ptr InMemEmbeddings::constructScanState( + transaction::Transaction*) const { return std::make_unique(data, info.typeInfo.getChildType()); } @@ -110,15 +291,413 @@ std::vector InMemEmbeddings::getEmbeddings( return ret; } +template +static void quantizeEmbedding(const T* src, uint64_t dimension, QuantizationType quantization, + MetricType metric, uint8_t* payloadDst, float& scaleDst, float& normSqDst) { + constexpr float zeroScale = 0.0f; + constexpr float zeroNormSq = 0.0f; + double maxAbs = 0.0; + for (auto i = 0u; i < dimension; ++i) { + maxAbs = std::max(maxAbs, std::abs(static_cast(src[i]))); + } + if (maxAbs == 0.0) { + scaleDst = zeroScale; + normSqDst = zeroNormSq; + std::memset(payloadDst, 0, + HNSWIndexUtils::getQuantizedCachedEmbeddingPayloadBytes(dimension, quantization)); + return; + } + const auto maxQuantizedValue = quantization == QuantizationType::SQ8 ? 127.0 : 32767.0; + const auto scale = static_cast(maxAbs / maxQuantizedValue); + float normSq = 0.0f; + if (quantization == QuantizationType::SQ8) { + for (auto i = 0u; i < dimension; ++i) { + const auto quantized = + static_cast(std::llround(static_cast(src[i]) / scale)); + const auto clamped = static_cast(std::clamp(quantized, -127, 127)); + payloadDst[i] = static_cast(clamped); + normSq += static_cast(clamped) * static_cast(clamped); + } + normSqDst = normSq; + scaleDst = metric == MetricType::Cosine && normSq > 0.0f ? 1.0f / std::sqrt(normSq) : + scale; + return; + } + auto* payloadValues = reinterpret_cast(payloadDst); + for (auto i = 0u; i < dimension; ++i) { + const auto quantized = + static_cast(std::llround(static_cast(src[i]) / scale)); + const auto clamped = static_cast(std::clamp(quantized, -32767, 32767)); + normSq += static_cast(clamped) * static_cast(clamped); + payloadValues[i] = clamped; + } + normSqDst = normSq; + scaleDst = metric == MetricType::Cosine && normSq > 0.0f ? 1.0f / std::sqrt(normSq) : scale; +} + +static std::vector initializeNullMask(common::offset_t numRows) { + return std::vector(numRows, 1); +} + +QuantizedInMemEmbeddings::QuantizedInMemEmbeddings(const main::ClientContext* context, + common::ArrayTypeInfo typeInfo, storage::NodeTable& table, common::column_id_t columnID, + QuantizationType quantization, MetricType metric) + : HNSWIndexEmbeddings{std::move(typeInfo)}, quantization(quantization), + strideBytes( + HNSWIndexUtils::getQuantizedCachedEmbeddingStride(info.getDimension(), quantization)), + nullMask(initializeNullMask(table.getNumTotalRows(transaction::Transaction::Get(*context)))), + payloadDataPtr(nullptr) { + const auto numRows = table.getNumTotalRows(transaction::Transaction::Get(*context)); + scales.resize(numRows, 0.0f); + normSqs.resize(numRows, 0.0f); + views.resize(numRows); + const auto alignmentPadding = 31; + payloadBuffer = storage::MemoryManager::Get(*context)->allocateBuffer(false, + numRows * strideBytes + alignmentPadding); + payloadDataPtr = alignPointer(payloadBuffer->getData(), 32); + auto onDiskEmbeddings = OnDiskEmbeddings(transaction::Transaction::Get(*context), + storage::MemoryManager::Get(*context), + common::ArrayTypeInfo{info.typeInfo.getChildType().copy(), info.getDimension()}, table, + columnID); + auto scanState = onDiskEmbeddings.constructScanState(); + for (common::offset_t offset = 0; offset < numRows; ++offset) { + const auto embedding = onDiskEmbeddings.getEmbedding(offset, *scanState); + if (embedding.isNull()) { + nullMask[offset] = 1; + continue; + } + nullMask[offset] = 0; + const auto dst = payloadDataPtr + offset * strideBytes; + switch (info.typeInfo.getChildType().getLogicalTypeID()) { + case common::LogicalTypeID::FLOAT: { + quantizeEmbedding(static_cast(embedding.getPtr()), info.getDimension(), + quantization, metric, dst, scales[offset], normSqs[offset]); + } break; + case common::LogicalTypeID::DOUBLE: { + quantizeEmbedding(static_cast(embedding.getPtr()), info.getDimension(), + quantization, metric, dst, scales[offset], normSqs[offset]); + } break; + default: { + UNREACHABLE_CODE; + } + } + views[offset] = QuantizedEmbeddingView{dst, scales[offset], normSqs[offset]}; + } +} + +std::unique_ptr QuantizedInMemEmbeddings::constructScanState( + transaction::Transaction*) const { + return std::make_unique(views.data()); +} + +EmbeddingHandle QuantizedInMemEmbeddings::getEmbedding(common::offset_t offset, + GetEmbeddingsScanState& scanState) const { + if (nullMask[offset] != 0) { + return EmbeddingHandle::createNullHandle(); + } + return EmbeddingHandle{offset, &scanState}; +} + +std::vector QuantizedInMemEmbeddings::getEmbeddings( + std::span offsets, GetEmbeddingsScanState& scanState) const { + std::vector ret; + ret.reserve(offsets.size()); + for (const auto offset : offsets) { + ret.push_back(getEmbedding(offset, scanState)); + } + return ret; +} + +metric_func_t QuantizedInMemEmbeddings::getMetricFunction(MetricType metric) const { + return HNSWIndexUtils::getQuantizedCachedMetricsFunction(metric, quantization); +} + +CachedQuantizedColumn::CachedQuantizedColumn(const main::ClientContext* context, + common::ArrayTypeInfo typeInfo, storage::NodeTable& sourceTable, + storage::NodeTable& quantizedTable, common::column_id_t sourceColumnID, + QuantizationType quantization, MetricType metric, common::offset_t numNodes) + : LocalCacheObject{getKey(quantizedTable.getTableID())}, typeInfo{std::move(typeInfo)}, + strideBytes( + HNSWIndexUtils::getQuantizedCachedEmbeddingStride(this->typeInfo.getNumElements(), + quantization)), + payloadBytes(HNSWIndexUtils::getQuantizedCachedEmbeddingPayloadBytes( + this->typeInfo.getNumElements(), quantization)), + nullMask(initializeNullMask(numNodes)), payloadDataPtr(nullptr), scales(numNodes, 0.0f), + normSqs(numNodes, 0.0f), views(numNodes) { + const auto alignmentPadding = 31; + payloadBuffer = storage::MemoryManager::Get(*context)->allocateBuffer(false, + numNodes * strideBytes + alignmentPadding); + payloadDataPtr = alignPointer(payloadBuffer->getData(), 32); + auto transaction = transaction::Transaction::Get(*context); + + auto state = std::make_shared(); + auto nodeIDVector = common::ValueVector(common::LogicalType::INTERNAL_ID()); + auto validVector = common::ValueVector(common::LogicalType::UINT8()); + auto scaleVector = common::ValueVector(common::LogicalType::FLOAT()); + auto normSqVector = common::ValueVector(common::LogicalType::FLOAT()); + auto payloadVector = common::ValueVector(quantizedTable.getColumn(4).getDataType().copy(), + storage::MemoryManager::Get(*context)); + nodeIDVector.state = state; + validVector.state = state; + scaleVector.state = state; + normSqVector.state = state; + payloadVector.state = state; + auto scanState = storage::NodeTableScanState(&nodeIDVector, + std::vector{&validVector, &scaleVector, &normSqVector, &payloadVector}, state); + scanState.setToTable(transaction, &quantizedTable, {1, 2, 3, 4}, {}); + scanState.source = storage::TableScanSource::COMMITTED; + for (auto nodeGroupIdx = 0u; nodeGroupIdx < quantizedTable.getNumCommittedNodeGroups(); + ++nodeGroupIdx) { + scanState.nodeGroupIdx = nodeGroupIdx; + quantizedTable.initScanState(transaction, scanState); + while (quantizedTable.scan(transaction, scanState)) { + const auto& selVector = scanState.outState->getSelVector(); + selVector.forEach([&](auto pos) { + if (validVector.isNull(pos) || scaleVector.isNull(pos) || + normSqVector.isNull(pos) || payloadVector.isNull(pos)) { + return; + } + const auto entry = payloadVector.getValue(pos); + setQuantizedRecord(nodeIDVector.readNodeOffset(pos), + validVector.getValue(pos), + scaleVector.getValue(pos), normSqVector.getValue(pos), + common::ListVector::getListValues(&payloadVector, entry)); + }); + } + } + + const auto numCommittedRows = quantizedTable.getNumTotalRows(nullptr); + if (numCommittedRows < numNodes) { + auto onDiskEmbeddings = OnDiskEmbeddings(transaction, storage::MemoryManager::Get(*context), + common::ArrayTypeInfo{this->typeInfo.getChildType().copy(), + this->typeInfo.getNumElements()}, + sourceTable, sourceColumnID); + auto pointLookupState = onDiskEmbeddings.constructScanState(transaction); + std::vector payload(payloadBytes, 0); + for (auto offset = numCommittedRows; offset < numNodes; ++offset) { + const auto embedding = onDiskEmbeddings.getEmbedding(offset, *pointLookupState); + if (embedding.isNull()) { + continue; + } + auto scale = float{0.0f}; + auto normSq = float{0.0f}; + switch (this->typeInfo.getChildType().getLogicalTypeID()) { + case common::LogicalTypeID::FLOAT: { + quantizeEmbedding(static_cast(embedding.getPtr()), + this->typeInfo.getNumElements(), quantization, metric, payload.data(), scale, + normSq); + } break; + case common::LogicalTypeID::DOUBLE: { + quantizeEmbedding(static_cast(embedding.getPtr()), + this->typeInfo.getNumElements(), quantization, metric, payload.data(), scale, + normSq); + } break; + default: { + UNREACHABLE_CODE; + } + } + setQuantizedRecord(offset, 1, scale, normSq, payload.data()); + } + } +} + +void CachedQuantizedColumn::setQuantizedRecord(common::offset_t offset, uint8_t valid, float scale, + float normSq, const uint8_t* payload) { + DASSERT(offset < getNumNodes()); + const auto dst = payloadDataPtr + offset * strideBytes; + if (valid == 0) { + nullMask[offset] = 1; + scales[offset] = 0.0f; + normSqs[offset] = 0.0f; + views[offset] = QuantizedEmbeddingView{dst, scales[offset], normSqs[offset]}; + return; + } + std::memcpy(dst, payload, payloadBytes); + scales[offset] = scale; + normSqs[offset] = normSq; + nullMask[offset] = 0; + views[offset] = QuantizedEmbeddingView{dst, scales[offset], normSqs[offset]}; +} + +CachedQuantizedEmbeddings::CachedQuantizedEmbeddings(common::ArrayTypeInfo typeInfo, + storage::NodeTable& sourceTable, std::shared_ptr cachedColumn, + QuantizationType quantization) + : HNSWIndexEmbeddings{std::move(typeInfo)}, quantization(quantization), nodeTable(sourceTable), + cachedColumn(std::move(cachedColumn)) {} + +std::unique_ptr CachedQuantizedEmbeddings::constructScanState( + transaction::Transaction* transaction) const { + DASSERT(transaction != nullptr); + return std::make_unique(transaction, &nodeTable, + cachedColumn->getViews()); +} + +EmbeddingHandle CachedQuantizedEmbeddings::getEmbedding(common::offset_t offset, + GetEmbeddingsScanState& scanState) const { + if (offset >= cachedColumn->getNumNodes()) { + return EmbeddingHandle::createNullHandle(); + } + const auto& quantizedScanState = scanState.cast(); + if (!quantizedScanState.nodeTable->isVisibleNoLock(quantizedScanState.transaction, offset)) { + return EmbeddingHandle::createNullHandle(); + } + if (cachedColumn->isNull(offset)) { + return EmbeddingHandle::createNullHandle(); + } + return EmbeddingHandle{offset, &scanState}; +} + +metric_func_t CachedQuantizedEmbeddings::getMetricFunction(MetricType metric) const { + return HNSWIndexUtils::getQuantizedCachedMetricsFunction(metric, quantization); +} + +std::vector CachedQuantizedEmbeddings::getEmbeddings( + std::span offsets, GetEmbeddingsScanState& scanState) const { + std::vector ret; + ret.reserve(offsets.size()); + for (const auto offset : offsets) { + ret.push_back(getEmbedding(offset, scanState)); + } + return ret; +} + +TableBackedQuantizedEmbeddings::TableBackedQuantizedEmbeddings(const main::ClientContext* context, + common::ArrayTypeInfo typeInfo, storage::NodeTable& sourceTable, + common::column_id_t sourceColumnID, storage::NodeTable& quantizedTable, + QuantizationType quantization, MetricType metric, common::offset_t numNodes) + : HNSWIndexEmbeddings{std::move(typeInfo)}, quantization(quantization), + metric(metric), sourceTable(sourceTable), sourceColumnID(sourceColumnID), + quantizedTable(quantizedTable), + defaultTransaction(context == nullptr ? nullptr : transaction::Transaction::Get(*context)), + numNodes(numNodes), + strideBytes( + HNSWIndexUtils::getQuantizedCachedEmbeddingStride(info.getDimension(), quantization)) {} + +void TableBackedQuantizedEmbeddings::writeEmbedding(transaction::Transaction* transaction, + common::offset_t offset, const EmbeddingHandle& embedding) const { + std::vector payload(strideBytes, 0); + auto valid = uint8_t{0}; + auto scale = float{0.0f}; + auto normSq = float{0.0f}; + if (!embedding.isNull()) { + valid = 1; + switch (info.typeInfo.getChildType().getLogicalTypeID()) { + case common::LogicalTypeID::FLOAT: { + quantizeEmbedding(static_cast(embedding.getPtr()), info.getDimension(), + quantization, metric, payload.data(), scale, normSq); + } break; + case common::LogicalTypeID::DOUBLE: { + quantizeEmbedding(static_cast(embedding.getPtr()), info.getDimension(), + quantization, metric, payload.data(), scale, normSq); + } break; + default: { + UNREACHABLE_CODE; + } + } + } + writeQuantizedRecord(transaction, offset, valid, scale, normSq, payload.data()); +} + +void TableBackedQuantizedEmbeddings::rebuild(const main::ClientContext* context) const { + auto transaction = transaction::Transaction::Get(*context); + auto onDiskEmbeddings = OnDiskEmbeddings(transaction, storage::MemoryManager::Get(*context), + common::ArrayTypeInfo{info.typeInfo.getChildType().copy(), info.getDimension()}, + sourceTable, sourceColumnID); + auto scanState = onDiskEmbeddings.constructScanState(transaction); + for (common::offset_t offset = 0; offset < numNodes; ++offset) { + const auto embedding = onDiskEmbeddings.getEmbedding(offset, *scanState); + writeEmbedding(transaction, offset, embedding); + } +} + +void TableBackedQuantizedEmbeddings::writeQuantizedRecord(transaction::Transaction* transaction, + common::offset_t offset, uint8_t valid, float scale, float normSq, + const uint8_t* payload) const { + auto state = common::DataChunkState::getSingleValueDataChunkState(); + auto nodeIDVector = common::ValueVector(common::LogicalType::INTERNAL_ID()); + auto idVector = common::ValueVector(common::LogicalType::INT64()); + auto validVector = common::ValueVector(common::LogicalType::UINT8()); + auto scaleVector = common::ValueVector(common::LogicalType::FLOAT()); + auto normSqVector = common::ValueVector(common::LogicalType::FLOAT()); + auto payloadVector = common::ValueVector(quantizedTable.getColumn(4).getDataType().copy(), + storage::MemoryManager::Get(*transaction->getClientContext())); + nodeIDVector.state = state; + idVector.state = state; + validVector.state = state; + scaleVector.state = state; + normSqVector.state = state; + payloadVector.state = state; + nodeIDVector.setNull(0, false); + idVector.setValue(0, static_cast(offset)); + validVector.setValue(0, valid); + scaleVector.setValue(0, scale); + normSqVector.setValue(0, normSq); + const auto entry = common::ListVector::addList(&payloadVector, info.getDimension()); + payloadVector.setValue(0, entry); + const auto payloadBytes = + HNSWIndexUtils::getQuantizedCachedEmbeddingPayloadBytes(info.getDimension(), quantization); + std::memcpy(common::ListVector::getListValues(&payloadVector, entry), payload, payloadBytes); + std::vector propertyVectors{ + &idVector, &validVector, &scaleVector, &normSqVector, &payloadVector}; + auto insertState = + storage::NodeTableInsertState(nodeIDVector, idVector, std::move(propertyVectors)); + quantizedTable.initInsertState(transaction->getClientContext(), insertState); + quantizedTable.insert(transaction, insertState); +} + +std::unique_ptr TableBackedQuantizedEmbeddings::constructScanState( + transaction::Transaction* transaction) const { + auto scanTransaction = transaction == nullptr ? defaultTransaction : transaction; + DASSERT(scanTransaction != nullptr); + const auto payloadBytes = + HNSWIndexUtils::getQuantizedCachedEmbeddingPayloadBytes(info.getDimension(), quantization); + return std::make_unique(scanTransaction, &sourceTable, + &quantizedTable, strideBytes, payloadBytes); +} + +EmbeddingHandle TableBackedQuantizedEmbeddings::getEmbedding(common::offset_t offset, + GetEmbeddingsScanState& scanState) const { + if (offset >= numNodes) { + return EmbeddingHandle::createNullHandle(); + } + const auto& persistentScanState = scanState.cast(); + if (!sourceTable.isVisibleNoLock(persistentScanState.transaction, offset)) { + return EmbeddingHandle::createNullHandle(); + } + if (!const_cast(persistentScanState).loadEmbedding( + offset)) { + return EmbeddingHandle::createNullHandle(); + } + return EmbeddingHandle{offset, &scanState}; +} + +std::vector TableBackedQuantizedEmbeddings::getEmbeddings( + std::span offsets, GetEmbeddingsScanState& scanState) const { + auto& persistentScanState = scanState.cast(); + persistentScanState.loadEmbeddings(offsets); + std::vector ret; + ret.reserve(offsets.size()); + for (const auto offset : offsets) { + ret.push_back(getEmbedding(offset, scanState)); + } + return ret; +} + +metric_func_t TableBackedQuantizedEmbeddings::getMetricFunction(MetricType metric) const { + return HNSWIndexUtils::getQuantizedCachedMetricsFunction(metric, quantization); +} + OnDiskEmbeddings::OnDiskEmbeddings(transaction::Transaction* transaction, storage::MemoryManager* mm, common::ArrayTypeInfo typeInfo, storage::NodeTable& nodeTable, common::column_id_t columnID) : HNSWIndexEmbeddings(std::move(typeInfo)), transaction(transaction), mm(mm), nodeTable(nodeTable), columnID(columnID) {} -std::unique_ptr OnDiskEmbeddings::constructScanState() const { - return std::make_unique(transaction, mm, nodeTable, columnID, - info.getDimension()); +std::unique_ptr OnDiskEmbeddings::constructScanState( + transaction::Transaction* transaction_) const { + return std::make_unique(transaction_ == nullptr ? transaction : + transaction_, + mm, nodeTable, columnID, info.getDimension()); } EmbeddingHandle OnDiskEmbeddings::getEmbedding(common::offset_t offset, diff --git a/vector/src/index/hnsw_index.cpp b/vector/src/index/hnsw_index.cpp index 79fe88a0..cd5d2300 100644 --- a/vector/src/index/hnsw_index.cpp +++ b/vector/src/index/hnsw_index.cpp @@ -1,5 +1,6 @@ #include "index/hnsw_index.h" +#include #include #include "catalog/catalog.h" @@ -22,6 +23,14 @@ static constexpr double NAVIX_ONE_HOP_LOCAL_SELECTIVITY_THRESHOLD = 0.4; static constexpr double NAVIX_TWO_HOP_SELECTED_NEIGHBOR_ESTIMATE_RATIO = 0.4; static constexpr int64_t NAVIX_DIRECTED_FIRST_HOP_DISTANCE_COST_FACTOR = 2; +struct QuantizedEmbeddingsCacheState { + std::mutex mtx; + std::shared_ptr cache; + common::transaction_t version = transaction::Transaction::DUMMY_START_TIMESTAMP; + bool dirty = false; + std::unordered_set dirtyTransactions; +}; + InMemHNSWLayer::InMemHNSWLayer(MemoryManager* mm, InMemHNSWLayerInfo info) : entryPoint{common::INVALID_OFFSET}, info{info} { graph = std::make_unique(mm, info.numNodes, info.degreeThresholdToShrink); @@ -314,6 +323,11 @@ static std::unique_ptr constructEmbeddingsColumn( common::column_id_t columnID, const HNSWIndexConfig& config) { auto transaction = Transaction::Get(*context); if (config.cacheEmbeddingsColumn) { + if (config.quantization != QuantizationType::NONE) { + return std::make_unique(context, + common::ArrayTypeInfo{typeInfo.getChildType().copy(), typeInfo.getNumElements()}, + table, columnID, config.quantization, config.metric); + } return std::make_unique(transaction, common::ArrayTypeInfo{typeInfo.getChildType().copy(), typeInfo.getNumElements()}, table.getTableID(), columnID); @@ -355,14 +369,15 @@ InMemHNSWIndex::InMemHNSWIndex(const main::ClientContext* context, IndexInfo ind lowerGraphSelectionMap = std::make_unique(numNodes); upperGraphSelectionMap = std::make_unique(numNodes, upperLayerSelectionMask.get()); + auto inMemMetricFunc = embeddings->getMetricFunction(this->config.metric); lowerLayer = std::make_unique(MemoryManager::Get(*context), - InMemHNSWLayerInfo{numNodes, embeddings.get(), this->metricFunc, + InMemHNSWLayerInfo{numNodes, embeddings.get(), inMemMetricFunc, getDegreeThresholdToShrink(this->config.ml), this->config.ml, this->config.alpha, this->config.efc, *lowerGraphSelectionMap}); upperLayer = std::make_unique(MemoryManager::Get(*context), - InMemHNSWLayerInfo{upperLayerSelectionMask->countNulls(), embeddings.get(), - this->metricFunc, getDegreeThresholdToShrink(this->config.mu), this->config.mu, - this->config.alpha, this->config.efc, *upperGraphSelectionMap}); + InMemHNSWLayerInfo{upperLayerSelectionMask->countNulls(), embeddings.get(), inMemMetricFunc, + getDegreeThresholdToShrink(this->config.mu), this->config.mu, this->config.alpha, + this->config.efc, *upperGraphSelectionMap}); } // NOLINTNEXTLINE(readability-make-member-function-const): Semantically non-const function. @@ -401,6 +416,7 @@ std::shared_ptr HNSWStorageInfo::serialize() const { serializer.write(upperEntryPoint); serializer.write(lowerEntryPoint); serializer.write(numCheckpointedNodes); + serializer.write(quantizedEmbeddingsTableID); return bufferWriter; } @@ -408,6 +424,7 @@ std::unique_ptr HNSWStorageInfo::deserialize( std::unique_ptr reader) { common::table_id_t upperRelTableID = common::INVALID_TABLE_ID; common::table_id_t lowerRelTableID = common::INVALID_TABLE_ID; + common::table_id_t quantizedEmbeddingsTableID = common::INVALID_TABLE_ID; common::offset_t upperEntryPoint = common::INVALID_OFFSET; common::offset_t lowerEntryPoint = common::INVALID_OFFSET; common::offset_t checkpointedNodeOffset = common::INVALID_OFFSET; @@ -417,22 +434,47 @@ std::unique_ptr HNSWStorageInfo::deserialize( deSer.deserializeValue(upperEntryPoint); deSer.deserializeValue(lowerEntryPoint); deSer.deserializeValue(checkpointedNodeOffset); - return std::make_unique(upperRelTableID, lowerRelTableID, upperEntryPoint, - lowerEntryPoint, checkpointedNodeOffset); + if (!deSer.finished()) { + deSer.deserializeValue(quantizedEmbeddingsTableID); + } + return std::make_unique(upperRelTableID, lowerRelTableID, + quantizedEmbeddingsTableID, upperEntryPoint, lowerEntryPoint, checkpointedNodeOffset); } HNSWSearchState::HNSWSearchState(main::ClientContext* context, catalog::TableCatalogEntry* nodeTableEntry, catalog::TableCatalogEntry* upperRelTableEntry, catalog::TableCatalogEntry* lowerRelTableEntry, NodeTable& nodeTable, - common::column_id_t columnID, common::offset_t numNodes, uint64_t k, QueryHNSWConfig config) - : visited{numNodes}, embeddings{std::make_unique(Transaction::Get(*context), - MemoryManager::Get(*context), getArrayTypeInfo(nodeTable, columnID), - nodeTable, columnID)}, - embeddingScanState{Transaction::Get(*context), MemoryManager::Get(*context), nodeTable, - columnID, embeddings->getDimension()}, - k{k}, config{config}, semiMask{nullptr}, upperRelTableEntry{upperRelTableEntry}, - lowerRelTableEntry{lowerRelTableEntry}, searchType{SearchType::UNFILTERED}, - nbrScanState{nullptr}, secondHopNbrScanState{nullptr} { + common::column_id_t columnID, common::offset_t numNodes, uint64_t k, QueryHNSWConfig config, + const HNSWIndexConfig& indexConfig, bool useQuantizedEmbeddings, + std::shared_ptr cachedEmbeddings) + : visited{numNodes}, k{k}, config{config}, semiMask{nullptr}, + upperRelTableEntry{upperRelTableEntry}, lowerRelTableEntry{lowerRelTableEntry}, + searchType{SearchType::UNFILTERED}, nbrScanState{nullptr}, secondHopNbrScanState{nullptr} { + auto transaction = Transaction::Get(*context); + auto typeInfo = getArrayTypeInfo(nodeTable, columnID); + if (useQuantizedEmbeddings && indexConfig.quantization != QuantizationType::NONE) { + if (cachedEmbeddings != nullptr) { + embeddings = std::move(cachedEmbeddings); + } else { + embeddings = std::make_shared(context, + common::ArrayTypeInfo{typeInfo.getChildType().copy(), typeInfo.getNumElements()}, + nodeTable, columnID, indexConfig.quantization, indexConfig.metric); + } + embeddingScanState = embeddings->constructScanState(transaction); + if (indexConfig.storeFullPrecisionEmbeddings) { + exactEmbeddings = std::make_shared(transaction, + MemoryManager::Get(*context), + common::ArrayTypeInfo{typeInfo.getChildType().copy(), typeInfo.getNumElements()}, + nodeTable, columnID); + exactEmbeddingScanState = exactEmbeddings->constructScanState(transaction); + } + } else { + embeddings = std::make_shared(transaction, MemoryManager::Get(*context), + common::ArrayTypeInfo{typeInfo.getChildType().copy(), typeInfo.getNumElements()}, + nodeTable, columnID); + embeddingScanState = embeddings->constructScanState(transaction); + } + metricFunc = embeddings->getMetricFunction(indexConfig.metric); ef = std::max(k, static_cast(config.efs)); graph::NativeGraphEntry lowerGraphEntry{{nodeTableEntry}, {lowerRelTableEntry}}; lowerGraph = std::make_unique(context, std::move(lowerGraphEntry)); @@ -446,7 +488,7 @@ OnDiskHNSWIndex::HNSWInsertState::HNSWInsertState(main::ClientContext* context, common::column_id_t columnID, uint64_t degree) : searchState{context, nodeTableEntry, upperRelTableEntry, lowerRelTableEntry, nodeTable, columnID, nodeTable.getNumTotalRows(Transaction::Get(*context)), degree, - QueryHNSWConfig{}} { + QueryHNSWConfig{}, HNSWIndexConfig{}, false} { std::vector insertTypes; insertTypes.push_back(common::LogicalType::INTERNAL_ID()); insertTypes.push_back(common::LogicalType::INTERNAL_ID()); @@ -473,7 +515,9 @@ OnDiskHNSWIndex::OnDiskHNSWIndex(const main::ClientContext* context, IndexInfo i StorageManager::Get(*context)->getTable(indexInfo.tableID)->cast(), indexInfo.columnIDs[0])}, mm{MemoryManager::Get(*context)}, - nodeTable{StorageManager::Get(*context)->getTable(indexInfo.tableID)->cast()} { + nodeTable{StorageManager::Get(*context)->getTable(indexInfo.tableID)->cast()}, + quantizedEmbeddingsTable{nullptr}, + quantizedEmbeddingsCacheState{std::make_shared()} { DASSERT(this->indexInfo.columnIDs.size() == 1); DASSERT(nodeTable.getColumn(this->indexInfo.columnIDs[0]).getDataType().getLogicalTypeID() == common::LogicalTypeID::ARRAY); @@ -481,6 +525,50 @@ OnDiskHNSWIndex::OnDiskHNSWIndex(const main::ClientContext* context, IndexInfo i const auto& hnswStorageInfo = this->storageInfo->cast(); lowerRelTable = storageManager->getTable(hnswStorageInfo.lowerRelTableID)->ptrCast(); upperRelTable = storageManager->getTable(hnswStorageInfo.upperRelTableID)->ptrCast(); + if (hnswStorageInfo.quantizedEmbeddingsTableID != common::INVALID_TABLE_ID) { + quantizedEmbeddingsTable = + storageManager->getTable(hnswStorageInfo.quantizedEmbeddingsTableID) + ->ptrCast(); + } +} + +std::shared_ptr OnDiskHNSWIndex::getOrCreateQuantizedEmbeddings( + const main::ClientContext* context, common::offset_t numNodes) const { + if (config.quantization == QuantizationType::NONE) { + return nullptr; + } + DASSERT(quantizedEmbeddingsTable != nullptr); + if (config.cacheEmbeddingsColumn) { + auto transaction = Transaction::Get(*context); + auto shouldUseCache = false; + std::shared_ptr cache; + { + std::lock_guard lck{quantizedEmbeddingsCacheState->mtx}; + shouldUseCache = !quantizedEmbeddingsCacheState->dirty && + transaction->getStartTS() >= quantizedEmbeddingsCacheState->version; + if (shouldUseCache) { + if (quantizedEmbeddingsCacheState->cache == nullptr || + quantizedEmbeddingsCacheState->cache->getNumNodes() < numNodes) { + quantizedEmbeddingsCacheState->cache = + std::make_shared(context, + common::ArrayTypeInfo{typeInfo.getChildType().copy(), + typeInfo.getNumElements()}, + nodeTable, *quantizedEmbeddingsTable, indexInfo.columnIDs[0], + config.quantization, config.metric, numNodes); + } + cache = quantizedEmbeddingsCacheState->cache; + } + } + if (shouldUseCache) { + return std::make_shared( + common::ArrayTypeInfo{typeInfo.getChildType().copy(), typeInfo.getNumElements()}, + nodeTable, std::move(cache), config.quantization); + } + } + return std::make_shared(context, + common::ArrayTypeInfo{typeInfo.getChildType().copy(), typeInfo.getNumElements()}, + nodeTable, indexInfo.columnIDs[0], *quantizedEmbeddingsTable, config.quantization, + config.metric, numNodes); } std::unique_ptr OnDiskHNSWIndex::load(main::ClientContext* context, StorageManager*, @@ -497,23 +585,28 @@ std::unique_ptr OnDiskHNSWIndex::load(main::ClientContext* context, Stora } std::vector OnDiskHNSWIndex::search(Transaction* transaction, - const EmbeddingHandle& queryVector, HNSWSearchState& searchState) const { + const EmbeddingHandle& queryVector, const EmbeddingHandle& exactQueryVector, + HNSWSearchState& searchState) const { if (isBruteForceSearch(searchState.config)) { - auto result = searchBruteForce(transaction, queryVector, searchState); + auto result = searchBruteForce(transaction, queryVector, exactQueryVector, searchState); result.resize(std::min(result.size(), searchState.k)); return result; } auto result = searchFromCheckpointed(transaction, queryVector, searchState); searchFromUnCheckpointed(transaction, queryVector, searchState, result); if (result.size() < searchState.k) { - result = searchBruteForce(transaction, queryVector, searchState); + result = searchBruteForce(transaction, queryVector, exactQueryVector, searchState); + } + if (searchState.exactEmbeddings != nullptr) { + exactRerankResults(exactQueryVector, searchState, result); } result.resize(std::min(result.size(), searchState.k)); return result; } std::vector OnDiskHNSWIndex::searchBruteForce(Transaction* transaction, - const EmbeddingHandle& queryVector, HNSWSearchState& searchState) const { + const EmbeddingHandle& approxQueryVector, const EmbeddingHandle& exactQueryVector, + HNSWSearchState& searchState) const { const auto numTotalRows = nodeTable.getNumTotalRows(transaction); std::vector result; result.reserve( @@ -523,12 +616,21 @@ std::vector OnDiskHNSWIndex::searchBruteForce(Transaction* tra continue; } const auto vector = - searchState.embeddings->getEmbedding(offset, searchState.embeddingScanState); + searchState.exactEmbeddings ? + searchState.exactEmbeddings->getEmbedding(offset, + *searchState.exactEmbeddingScanState) : + searchState.embeddings->getEmbedding(offset, *searchState.embeddingScanState); if (vector.isNull()) { continue; } - const auto dist = metricFunc(queryVector.getPtr(), vector.getPtr(), - searchState.embeddings->getDimension()); + double dist; + if (searchState.exactEmbeddings != nullptr) { + dist = exactMetricFunc(exactQueryVector.getPtr(), vector.getPtr(), + searchState.embeddings->getDimension()); + } else { + dist = searchState.metricFunc(approxQueryVector.getPtr(), vector.getPtr(), + searchState.embeddings->getDimension()); + } result.emplace_back(offset, dist); } std::ranges::sort(result, [](const NodeWithDistance& l, const NodeWithDistance& r) { @@ -565,11 +667,11 @@ void OnDiskHNSWIndex::searchFromUnCheckpointed(Transaction* transaction, // TODO(Guodong): Perhaps should switch to scan instead of lookup here. for (auto offset = hnswStorageInfo.numCheckpointedNodes; offset < numTotalRows; offset++) { const auto vector = - searchState.embeddings->getEmbedding(offset, searchState.embeddingScanState); + searchState.embeddings->getEmbedding(offset, *searchState.embeddingScanState); if (vector.isNull()) { continue; // Skip null or deleted values. } - auto dist = metricFunc(queryVector.getPtr(), vector.getPtr(), + auto dist = searchState.metricFunc(queryVector.getPtr(), vector.getPtr(), searchState.embeddings->getDimension()); result.emplace_back(offset, dist); } @@ -578,6 +680,26 @@ void OnDiskHNSWIndex::searchFromUnCheckpointed(Transaction* transaction, }); } +void OnDiskHNSWIndex::exactRerankResults(const EmbeddingHandle& exactQueryVector, + HNSWSearchState& searchState, std::vector& result) const { + if (config.quantization == QuantizationType::NONE || searchState.exactEmbeddings == nullptr) { + return; + } + for (auto& node : result) { + const auto vector = searchState.exactEmbeddings->getEmbedding(node.nodeOffset, + *searchState.exactEmbeddingScanState); + if (vector.isNull()) { + node.distance = std::numeric_limits::max(); + continue; + } + node.distance = exactMetricFunc(exactQueryVector.getPtr(), vector.getPtr(), + searchState.embeddings->getDimension()); + } + std::ranges::sort(result, [](const NodeWithDistance& l, const NodeWithDistance& r) { + return l.distance < r.distance; + }); +} + static std::tuple getIndexTableCatalogEntries(const catalog::Catalog* catalog, const Transaction* transaction, @@ -638,16 +760,34 @@ void OnDiskHNSWIndex::commitInsert(Transaction* transaction, DASSERT(nodeIDVector.state->getSelSize() == dataVectors[0]->state->getSelSize()); auto& hnswInsertState = insertState.cast(); auto commitInsertScanState = std::make_unique(dataVectors[0]); + std::unique_ptr quantizedEmbeddings; + if (config.quantization != QuantizationType::NONE) { + DASSERT(quantizedEmbeddingsTable != nullptr); + quantizedEmbeddings = std::make_unique( + nullptr, + common::ArrayTypeInfo{typeInfo.getChildType().copy(), typeInfo.getNumElements()}, + nodeTable, indexInfo.columnIDs[0], *quantizedEmbeddingsTable, config.quantization, + config.metric, nodeTable.getNumTotalRows(transaction)); + } for (size_t i = 0; i < nodeIDVector.state->getSelSize(); ++i) { auto pos = nodeIDVector.state->getSelVector()[i]; const auto offset = nodeIDVector.readNodeOffset(pos); auto valuePos = dataVectors[0]->state->getSelVector()[i]; if (dataVectors[0]->isNull(valuePos)) { + if (quantizedEmbeddings != nullptr) { + auto nullHandle = EmbeddingHandle::createNullHandle(); + quantizedEmbeddings->writeEmbedding(transaction, offset, nullHandle); + markQuantizedEmbeddingsCacheDirty(transaction); + } continue; // Skip null or deleted values. } EmbeddingHandle handle{dataVectors[0]->getValue(valuePos).offset, commitInsertScanState.get()}; insertInternal(transaction, offset, handle, hnswInsertState); + if (quantizedEmbeddings != nullptr) { + quantizedEmbeddings->writeEmbedding(transaction, offset, handle); + markQuantizedEmbeddingsCacheDirty(transaction); + } auto& hnswStorageInfo = storageInfo->cast(); hnswStorageInfo.numCheckpointedNodes = std::max(hnswStorageInfo.numCheckpointedNodes, offset + 1); @@ -671,7 +811,8 @@ void OnDiskHNSWIndex::finalize(main::ClientContext* context) { auto& hnswStorageInfo = storageInfo->cast(); auto transaction = Transaction::Get(*context); const auto numTotalRows = nodeTable.getNumTotalRows(transaction); - if (numTotalRows == hnswStorageInfo.numCheckpointedNodes) { + const auto numCheckpointedNodes = hnswStorageInfo.numCheckpointedNodes; + if (numTotalRows == numCheckpointedNodes) { return; } auto [nodeTableEntry, upperRelTableEntry, lowerRelTableEntry] = @@ -681,9 +822,21 @@ void OnDiskHNSWIndex::finalize(main::ClientContext* context) { indexInfo.columnIDs[0], embeddingDim); const auto insertState = std::make_unique(context, nodeTableEntry, upperRelTableEntry, lowerRelTableEntry, nodeTable, indexInfo.columnIDs[0], config.ml); + std::unique_ptr quantizedEmbeddings; + if (config.quantization != QuantizationType::NONE) { + DASSERT(quantizedEmbeddingsTable != nullptr); + quantizedEmbeddings = std::make_unique(context, + common::ArrayTypeInfo{typeInfo.getChildType().copy(), typeInfo.getNumElements()}, + nodeTable, indexInfo.columnIDs[0], *quantizedEmbeddingsTable, config.quantization, + config.metric, numTotalRows); + } // TODO(Guodong): Perhaps should switch to scan instead of lookup here. - for (auto offset = hnswStorageInfo.numCheckpointedNodes; offset < numTotalRows; offset++) { + for (auto offset = numCheckpointedNodes; offset < numTotalRows; offset++) { const auto vector = insertState->searchState.embeddings->getEmbedding(offset, *scanState); + if (quantizedEmbeddings != nullptr) { + quantizedEmbeddings->writeEmbedding(transaction, offset, vector); + markQuantizedEmbeddingsCacheDirty(transaction); + } if (vector.isNull()) { continue; } @@ -698,12 +851,51 @@ void OnDiskHNSWIndex::finalize(main::ClientContext* context) { hnswStorageInfo.numCheckpointedNodes = numTotalRows; } +void OnDiskHNSWIndex::markQuantizedEmbeddingsCacheDirty(Transaction* transaction) const { + if (quantizedEmbeddingsTable == nullptr || !config.cacheEmbeddingsColumn) { + return; + } + const auto transactionID = transaction->getID(); + auto shouldRegisterCallbacks = false; + { + std::lock_guard lck{quantizedEmbeddingsCacheState->mtx}; + quantizedEmbeddingsCacheState->dirty = true; + shouldRegisterCallbacks = + quantizedEmbeddingsCacheState->dirtyTransactions.insert(transactionID).second; + } + if (!shouldRegisterCallbacks) { + return; + } + const auto cacheState = quantizedEmbeddingsCacheState; + transaction->pushCommitCallback([cacheState, transactionID](Transaction& transaction) { + std::lock_guard lck{cacheState->mtx}; + cacheState->cache.reset(); + // Commit callbacks are serialized by mtx but may run out of commit timestamp order. + // Keep the watermark monotonic; dirty remains set until every writer has finished. + cacheState->version = std::max(cacheState->version, transaction.getCommitTS()); + cacheState->dirtyTransactions.erase(transactionID); + cacheState->dirty = !cacheState->dirtyTransactions.empty(); + }); + transaction->pushRollbackCallback([cacheState, transactionID](Transaction&) { + std::lock_guard lck{cacheState->mtx}; + cacheState->dirtyTransactions.erase(transactionID); + cacheState->dirty = !cacheState->dirtyTransactions.empty(); + }); +} + void OnDiskHNSWIndex::checkpoint(main::ClientContext* context, storage::PageAllocator& pageAllocator, [[maybe_unused]] storage::ShadowFile& shadowFile) { + const auto& hnswStorageInfo = storageInfo->cast(); auto [nodeTableEntry, upperRelTableEntry, lowerRelTableEntry] = getIndexTableCatalogEntries( catalog::Catalog::Get(*context), &DUMMY_CHECKPOINT_TRANSACTION, indexInfo); upperRelTable->checkpoint(context, upperRelTableEntry, pageAllocator); lowerRelTable->checkpoint(context, lowerRelTableEntry, pageAllocator); + if (hnswStorageInfo.quantizedEmbeddingsTableID != common::INVALID_TABLE_ID) { + DASSERT(quantizedEmbeddingsTable != nullptr); + auto quantizedEmbeddingsTableEntry = catalog::Catalog::Get(*context)->getTableCatalogEntry( + &DUMMY_CHECKPOINT_TRANSACTION, hnswStorageInfo.quantizedEmbeddingsTableID); + quantizedEmbeddingsTable->checkpoint(context, quantizedEmbeddingsTableEntry, pageAllocator); + } } void OnDiskHNSWIndex::insertInternal(Transaction* transaction, common::offset_t offset, @@ -730,11 +922,11 @@ common::offset_t OnDiskHNSWIndex::searchNNInUpperLayer(const EmbeddingHandle& qu double lastMinDist = std::numeric_limits::max(); const auto& embeddings = searchState.embeddings; const auto currNodeVector = - embeddings->getEmbedding(currentNodeOffset, searchState.embeddingScanState); + embeddings->getEmbedding(currentNodeOffset, *searchState.embeddingScanState); double minDist = 0.0; if (!currNodeVector.isNull()) { - minDist = - metricFunc(queryVector.getPtr(), currNodeVector.getPtr(), embeddings->getDimension()); + minDist = searchState.metricFunc(queryVector.getPtr(), currNodeVector.getPtr(), + embeddings->getDimension()); } const auto scanState = searchState.upperGraph->prepareRelScan(*searchState.upperRelTableEntry, hnswStorageInfo.upperRelTableID, indexInfo.tableID, {} /* relProperties */); @@ -746,10 +938,10 @@ common::offset_t OnDiskHNSWIndex::searchNNInUpperLayer(const EmbeddingHandle& qu neighborChunk.forEach([&](auto neighbors, auto, auto i) { auto neighbor = neighbors[i]; const auto nbrVector = - embeddings->getEmbedding(neighbor.offset, searchState.embeddingScanState); + embeddings->getEmbedding(neighbor.offset, *searchState.embeddingScanState); if (!nbrVector.isNull()) { - const auto dist = metricFunc(queryVector.getPtr(), nbrVector.getPtr(), - embeddings->getDimension()); + const auto dist = searchState.metricFunc(queryVector.getPtr(), + nbrVector.getPtr(), embeddings->getDimension()); if (dist < minDist) { minDist = dist; currentNodeOffset = neighbor.offset; @@ -790,9 +982,9 @@ std::vector OnDiskHNSWIndex::searchKNNInLayer(Transaction* tra initLayerSearchState(transaction, searchState, isUpperLayer); const auto entryVector = - searchState.embeddings->getEmbedding(entryNode, searchState.embeddingScanState); + searchState.embeddings->getEmbedding(entryNode, *searchState.embeddingScanState); if (!entryVector.isNull()) { - auto dist = metricFunc(queryVector.getPtr(), entryVector.getPtr(), + auto dist = searchState.metricFunc(queryVector.getPtr(), entryVector.getPtr(), searchState.embeddings->getDimension()); candidates.push({entryNode, dist}); if (searchState.isMasked(entryNode)) { @@ -890,12 +1082,12 @@ void OnDiskHNSWIndex::initSearchCandidates(const EmbeddingHandle& queryVector, } searchState.visited.add(candidate); const auto candidateVector = - searchState.embeddings->getEmbedding(candidate, searchState.embeddingScanState); + searchState.embeddings->getEmbedding(candidate, *searchState.embeddingScanState); if (candidateVector.isNull()) { continue; } - auto candidateDist = metricFunc(queryVector.getPtr(), candidateVector.getPtr(), - searchState.embeddings->getDimension()); + auto candidateDist = searchState.metricFunc(queryVector.getPtr(), + candidateVector.getPtr(), searchState.embeddings->getDimension()); candidates.push({candidate, candidateDist}); results.push({candidate, candidateDist}); } @@ -929,9 +1121,10 @@ void OnDiskHNSWIndex::oneHopSearch(const EmbeddingHandle& queryVector, const auto nbr = neighbors[i]; if (!searchState.visited.contains(nbr.offset) && searchState.isMasked(nbr.offset)) { const auto nbrVector = searchState.embeddings->getEmbedding(nbr.offset, - searchState.embeddingScanState); + *searchState.embeddingScanState); processNbrNodeInKNNSearch(queryVector, nbrVector, nbr.offset, searchState.ef, - searchState.visited, metricFunc, searchState.embeddings->getDimension(), + searchState.visited, searchState.metricFunc, + searchState.embeddings->getDimension(), candidates, results); } }); @@ -954,9 +1147,10 @@ void OnDiskHNSWIndex::navixFilteredSearch(const EmbeddingHandle& queryVector, for (const auto nbr : firstHopNbrs) { if (!searchState.visited.contains(nbr) && searchState.isMasked(nbr)) { const auto nbrVector = - searchState.embeddings->getEmbedding(nbr, searchState.embeddingScanState); + searchState.embeddings->getEmbedding(nbr, *searchState.embeddingScanState); processNbrNodeInKNNSearch(queryVector, nbrVector, nbr, searchState.ef, - searchState.visited, metricFunc, searchState.embeddings->getDimension(), + searchState.visited, searchState.metricFunc, + searchState.embeddings->getDimension(), candidates, results); } } @@ -994,10 +1188,10 @@ void OnDiskHNSWIndex::randomFilteredSearch(const EmbeddingHandle& queryVector, } if (!searchState.visited.contains(nbr) && searchState.isMasked(nbr)) { const auto nbrVector = - searchState.embeddings->getEmbedding(nbr, searchState.embeddingScanState); + searchState.embeddings->getEmbedding(nbr, *searchState.embeddingScanState); processNbrNodeInKNNSearch(queryVector, nbrVector, nbr, searchState.ef, - searchState.visited, metricFunc, searchState.embeddings->getDimension(), candidates, - results); + searchState.visited, searchState.metricFunc, + searchState.embeddings->getDimension(), candidates, results); } nbrsToExplore.push_back(nbr); } @@ -1019,9 +1213,10 @@ void OnDiskHNSWIndex::randomFilteredSearch(const EmbeddingHandle& queryVector, } if (!searchState.visited.contains(nbr) && searchState.isMasked(nbr)) { const auto nbrVector = - searchState.embeddings->getEmbedding(nbr, searchState.embeddingScanState); + searchState.embeddings->getEmbedding(nbr, *searchState.embeddingScanState); processNbrNodeInKNNSearch(queryVector, nbrVector, nbr, searchState.ef, - searchState.visited, metricFunc, searchState.embeddings->getDimension(), + searchState.visited, searchState.metricFunc, + searchState.embeddings->getDimension(), candidates, results); } }); @@ -1049,10 +1244,10 @@ min_node_priority_queue_t OnDiskHNSWIndex::collectFirstHopNbrsDirected( const auto neighbor = neighbors[i]; auto nbrOffset = neighbor.offset; if (!searchState.visited.contains(nbrOffset)) { - const auto nbrVector = - searchState.embeddings->getEmbedding(nbrOffset, searchState.embeddingScanState); + const auto nbrVector = searchState.embeddings->getEmbedding(nbrOffset, + *searchState.embeddingScanState); if (!nbrVector.isNull()) { - auto dist = metricFunc(queryVector.getPtr(), nbrVector.getPtr(), + auto dist = searchState.metricFunc(queryVector.getPtr(), nbrVector.getPtr(), searchState.embeddings->getDimension()); candidatesForSecHop.push({nbrOffset, dist}); if (searchState.isMasked(nbrOffset)) { @@ -1081,9 +1276,9 @@ min_node_priority_queue_t OnDiskHNSWIndex::collectFirstHopNbrsDirected( for (const auto nbrOffset : firstHopNbrs) { if (!searchState.visited.contains(nbrOffset)) { const auto nbrVector = - searchState.embeddings->getEmbedding(nbrOffset, searchState.embeddingScanState); + searchState.embeddings->getEmbedding(nbrOffset, *searchState.embeddingScanState); if (!nbrVector.isNull()) { - auto dist = metricFunc(queryVector.getPtr(), nbrVector.getPtr(), + auto dist = searchState.metricFunc(queryVector.getPtr(), nbrVector.getPtr(), searchState.embeddings->getDimension()); candidatesForSecHop.push({nbrOffset, dist}); if (searchState.isMasked(nbrOffset)) { @@ -1127,9 +1322,10 @@ common::offset_vec_t OnDiskHNSWIndex::collectFirstHopNbrsBlind(const EmbeddingHa if (searchState.isMasked(nbr.offset)) { numVisitedNbrs++; auto nbrVector = searchState.embeddings->getEmbedding(nbr.offset, - searchState.embeddingScanState); + *searchState.embeddingScanState); processNbrNodeInKNNSearch(queryVector, nbrVector, nbr.offset, searchState.ef, - searchState.visited, metricFunc, searchState.embeddings->getDimension(), + searchState.visited, searchState.metricFunc, + searchState.embeddings->getDimension(), candidates, results); } } @@ -1150,9 +1346,10 @@ common::offset_vec_t OnDiskHNSWIndex::collectFirstHopNbrsBlind(const EmbeddingHa if (searchState.isMasked(nbr)) { numVisitedNbrs++; auto nbrVector = - searchState.embeddings->getEmbedding(nbr, searchState.embeddingScanState); + searchState.embeddings->getEmbedding(nbr, *searchState.embeddingScanState); processNbrNodeInKNNSearch(queryVector, nbrVector, nbr, searchState.ef, - searchState.visited, metricFunc, searchState.embeddings->getDimension(), + searchState.visited, searchState.metricFunc, + searchState.embeddings->getDimension(), candidates, results); } } @@ -1227,7 +1424,8 @@ void OnDiskHNSWIndex::createRels(Transaction* transaction, common::offset_t offs void OnDiskHNSWIndex::shrinkForNode(Transaction* transaction, common::offset_t offset, bool isUpperLayer, common::length_t maxDegree, HNSWInsertState& insertState) { const auto& embeddings = *insertState.searchState.embeddings; - auto& embeddingScanState = insertState.searchState.embeddingScanState; + const auto approxMetricFunc = embeddings.getMetricFunction(config.metric); + auto& embeddingScanState = *insertState.searchState.embeddingScanState; const auto vector = embeddings.getEmbedding(offset, embeddingScanState); DASSERT(!vector.isNull()); const auto& searchState = insertState.searchState; @@ -1258,8 +1456,8 @@ void OnDiskHNSWIndex::shrinkForNode(Transaction* transaction, common::offset_t o if (nbrVectors[i].isNull()) { continue; } - auto dist = - metricFunc(vector.getPtr(), nbrVectors[i].getPtr(), embeddings.getDimension()); + auto dist = approxMetricFunc(vector.getPtr(), nbrVectors[i].getPtr(), + embeddings.getDimension()); nbrs.emplace_back(nbrOffsets[i], dist, std::move(nbrVectors[i])); } } @@ -1280,8 +1478,8 @@ void OnDiskHNSWIndex::shrinkForNode(Transaction* transaction, common::offset_t o for (const auto j : keptNbrs) { DASSERT(checkEmbeddingValidity(nbrs, i, embeddings, embeddingScanState)); DASSERT(checkEmbeddingValidity(nbrs, j, embeddings, embeddingScanState)); - const auto dist = metricFunc(nbrs[i].embedding.getPtr(), nbrs[j].embedding.getPtr(), - embeddings.getDimension()); + const auto dist = approxMetricFunc(nbrs[i].embedding.getPtr(), + nbrs[j].embedding.getPtr(), embeddings.getDimension()); if (config.alpha * dist < nbrs[i].getDist()) { keepNbr = false; break; @@ -1347,9 +1545,10 @@ bool OnDiskHNSWIndex::searchOverSecondHopNbrs(const EmbeddingHandle& queryVector auto nbr = neighbors[i]; if (!searchState.visited.contains(nbr.offset) && searchState.isMasked(nbr.offset)) { auto nbrVector = searchState.embeddings->getEmbedding(nbr.offset, - searchState.embeddingScanState); + *searchState.embeddingScanState); processNbrNodeInKNNSearch(queryVector, nbrVector, nbr.offset, ef, - searchState.visited, metricFunc, searchState.embeddings->getDimension(), + searchState.visited, searchState.metricFunc, + searchState.embeddings->getDimension(), candidates, results); numVisitedNbrs++; if (numVisitedNbrs >= config.ml) { diff --git a/vector/src/index/hnsw_index_utils.cpp b/vector/src/index/hnsw_index_utils.cpp index b9c723fd..924b5461 100644 --- a/vector/src/index/hnsw_index_utils.cpp +++ b/vector/src/index/hnsw_index_utils.cpp @@ -1,5 +1,9 @@ #include "index/hnsw_index_utils.h" +#include +#include +#include + #include "binder/binder.h" #include "catalog/catalog.h" #include "catalog/catalog_entry/node_table_catalog_entry.h" @@ -9,6 +13,13 @@ #include "transaction/transaction_context.h" #include +#if defined(__x86_64__) || defined(_M_X64) +#include +#endif +#if defined(__ARM_NEON) || defined(__ARM_NEON__) +#include +#endif + namespace lbug { namespace vector_extension { @@ -113,8 +124,248 @@ static metric_func_t computeDistanceFuncDispatch(const common::LogicalType& type } } -metric_func_t HNSWIndexUtils::getMetricsFunction(MetricType metric, - const common::LogicalType& type) { +template +struct ScalarQuantizedVector { + explicit ScalarQuantizedVector(const T* vector, uint32_t dimension, int8_t maxQuantizedValue) + : values(dimension), scale{0.0} { + double maxAbs = 0.0; + for (auto i = 0u; i < dimension; ++i) { + maxAbs = std::max(maxAbs, std::abs(static_cast(vector[i]))); + } + if (maxAbs == 0.0) { + return; + } + scale = maxAbs / maxQuantizedValue; + for (auto i = 0u; i < dimension; ++i) { + const auto quantized = + static_cast(std::llround(static_cast(vector[i]) / scale)); + values[i] = static_cast( + std::clamp(quantized, -maxQuantizedValue, maxQuantizedValue)); + } + } + + std::vector values; + double scale; +}; + +template +static double computeDistanceI8(const int8_t* left, const int8_t* right, uint32_t dimension) { + simsimd_distance_t distance = 0.0; + Func(reinterpret_cast(left), reinterpret_cast(right), + dimension, &distance); + return distance; +} + +static double computeDotI16(const int16_t* left, const int16_t* right, uint32_t dimension) { +#if defined(__AVX2__) + __m256i sum64Lo = _mm256_setzero_si256(); + __m256i sum64Hi = _mm256_setzero_si256(); + uint32_t i = 0; + for (; i + 16 <= dimension; i += 16) { + const auto leftVec = _mm256_loadu_si256(reinterpret_cast(left + i)); + const auto rightVec = _mm256_loadu_si256(reinterpret_cast(right + i)); + const auto pairwise32 = _mm256_madd_epi16(leftVec, rightVec); + const auto pairwise32Lo = _mm256_castsi256_si128(pairwise32); + const auto pairwise32Hi = _mm256_extracti128_si256(pairwise32, 1); + sum64Lo = _mm256_add_epi64(sum64Lo, _mm256_cvtepi32_epi64(pairwise32Lo)); + sum64Hi = _mm256_add_epi64(sum64Hi, _mm256_cvtepi32_epi64(pairwise32Hi)); + } + alignas(32) int64_t partialSums[8]; + _mm256_store_si256(reinterpret_cast<__m256i*>(partialSums), sum64Lo); + _mm256_store_si256(reinterpret_cast<__m256i*>(partialSums + 4), sum64Hi); + int64_t dot = 0; + for (const auto partial : partialSums) { + dot += partial; + } + for (; i < dimension; ++i) { + dot += static_cast(left[i]) * static_cast(right[i]); + } + return static_cast(dot); +#elif defined(__ARM_NEON) || defined(__ARM_NEON__) + int64x2_t sum64Lo = vdupq_n_s64(0); + int64x2_t sum64Hi = vdupq_n_s64(0); + uint32_t i = 0; + for (; i + 8 <= dimension; i += 8) { + const auto leftVec = vld1q_s16(left + i); + const auto rightVec = vld1q_s16(right + i); + const auto mulLo = vmull_s16(vget_low_s16(leftVec), vget_low_s16(rightVec)); + const auto mulHi = vmull_s16(vget_high_s16(leftVec), vget_high_s16(rightVec)); + sum64Lo = vaddq_s64(sum64Lo, vpaddlq_s32(mulLo)); + sum64Hi = vaddq_s64(sum64Hi, vpaddlq_s32(mulHi)); + } + int64_t dot = vgetq_lane_s64(sum64Lo, 0) + vgetq_lane_s64(sum64Lo, 1) + + vgetq_lane_s64(sum64Hi, 0) + vgetq_lane_s64(sum64Hi, 1); + for (; i < dimension; ++i) { + dot += static_cast(left[i]) * static_cast(right[i]); + } + return static_cast(dot); +#else + int64_t dot = 0; + for (auto i = 0u; i < dimension; ++i) { + dot += static_cast(left[i]) * static_cast(right[i]); + } + return static_cast(dot); +#endif +} + +static constexpr uint64_t QUANTIZED_HEADER_BYTES = sizeof(float) * 2; +static constexpr uint64_t QUANTIZED_SIMD_ALIGNMENT = 32; + +static uint64_t alignUp(uint64_t value, uint64_t alignment) { + return common::ceilDiv(value, alignment) * alignment; +} + +static uint64_t getQuantizedPayloadOffsetInternal(QuantizationType quantization) { + return quantization == QuantizationType::SQ8 ? + alignUp(QUANTIZED_HEADER_BYTES, QUANTIZED_SIMD_ALIGNMENT) : + QUANTIZED_HEADER_BYTES; +} + +static uint64_t getQuantizedEmbeddingPayloadBytesInternal(uint64_t dimension, + QuantizationType quantization) { + switch (quantization) { + case QuantizationType::SQ8: + return dimension; + case QuantizationType::SQ16: + return dimension * sizeof(int16_t); + default: + return 0; + } +} + +static uint64_t getQuantizedCachedEmbeddingPayloadBytesInternal(uint64_t dimension, + QuantizationType quantization) { + return getQuantizedEmbeddingPayloadBytesInternal(dimension, quantization); +} + +template +static double computeScalarQuantizedDistance(const void* left, const void* right, + uint32_t dimension, MetricType metric, int8_t maxQuantizedValue) { + const auto* leftVector = static_cast(left); + const auto* rightVector = static_cast(right); + const ScalarQuantizedVector leftQuantized{leftVector, dimension, maxQuantizedValue}; + const ScalarQuantizedVector rightQuantized{rightVector, dimension, maxQuantizedValue}; + switch (metric) { + case MetricType::Cosine: { + if (leftQuantized.scale == 0.0 || rightQuantized.scale == 0.0) { + return 1.0; + } + const auto leftNormSq = computeDistanceI8(leftQuantized.values.data(), + leftQuantized.values.data(), dimension); + const auto rightNormSq = computeDistanceI8(rightQuantized.values.data(), + rightQuantized.values.data(), dimension); + if (leftNormSq == 0.0 || rightNormSq == 0.0) { + return 1.0; + } + const auto dot = computeDistanceI8(leftQuantized.values.data(), + rightQuantized.values.data(), dimension); + const auto similarity = dot / std::sqrt(leftNormSq * rightNormSq); + return 1.0 - std::clamp(similarity, -1.0, 1.0); + } + case MetricType::L2: { + const auto leftNormSq = computeDistanceI8(leftQuantized.values.data(), + leftQuantized.values.data(), dimension); + const auto rightNormSq = computeDistanceI8(rightQuantized.values.data(), + rightQuantized.values.data(), dimension); + const auto dot = computeDistanceI8(leftQuantized.values.data(), + rightQuantized.values.data(), dimension); + const auto l2sq = leftQuantized.scale * leftQuantized.scale * leftNormSq + + rightQuantized.scale * rightQuantized.scale * rightNormSq - + 2.0 * leftQuantized.scale * rightQuantized.scale * dot; + return std::sqrt(std::max(0.0, l2sq)); + } + case MetricType::L2_SQUARE: { + const auto leftNormSq = computeDistanceI8(leftQuantized.values.data(), + leftQuantized.values.data(), dimension); + const auto rightNormSq = computeDistanceI8(rightQuantized.values.data(), + rightQuantized.values.data(), dimension); + const auto dot = computeDistanceI8(leftQuantized.values.data(), + rightQuantized.values.data(), dimension); + return std::max(0.0, leftQuantized.scale * leftQuantized.scale * leftNormSq + + rightQuantized.scale * rightQuantized.scale * rightNormSq - + 2.0 * leftQuantized.scale * rightQuantized.scale * dot); + } + default: { + UNREACHABLE_CODE; + } + } +} + +static double computeCachedSQ8Distance(const void* left, const void* right, uint32_t dimension, + MetricType metric) { + const auto& leftView = *static_cast(left); + const auto& rightView = *static_cast(right); + const auto* leftValues = reinterpret_cast(leftView.payload); + const auto* rightValues = reinterpret_cast(rightView.payload); + switch (metric) { + case MetricType::Cosine: { + if (leftView.scale == 0.0f || rightView.scale == 0.0f) { + return 1.0; + } + const auto dot = computeDistanceI8(leftValues, rightValues, dimension); + const auto similarity = static_cast(leftView.scale) * + static_cast(rightView.scale) * dot; + return 1.0 - std::clamp(similarity, -1.0, 1.0); + } + case MetricType::L2: { + const auto dot = computeDistanceI8(leftValues, rightValues, dimension); + const auto l2sq = leftView.scale * leftView.scale * leftView.normSq + + rightView.scale * rightView.scale * rightView.normSq - + 2.0 * leftView.scale * rightView.scale * dot; + return std::sqrt(std::max(0.0, l2sq)); + } + case MetricType::L2_SQUARE: { + const auto dot = computeDistanceI8(leftValues, rightValues, dimension); + return std::max(0.0, leftView.scale * leftView.scale * leftView.normSq + + rightView.scale * rightView.scale * rightView.normSq - + 2.0 * leftView.scale * rightView.scale * dot); + } + default: { + UNREACHABLE_CODE; + } + } +} + +static double computeCachedSQ16Distance(const void* left, const void* right, uint32_t dimension, + MetricType metric) { + const auto& leftView = *static_cast(left); + const auto& rightView = *static_cast(right); + const auto* leftValues = reinterpret_cast(leftView.payload); + const auto* rightValues = reinterpret_cast(rightView.payload); + switch (metric) { + case MetricType::Cosine: { + if (leftView.scale == 0.0f || rightView.scale == 0.0f) { + return 1.0; + } + const auto dot = computeDotI16(leftValues, rightValues, dimension); + const auto similarity = static_cast(leftView.scale) * + static_cast(rightView.scale) * dot; + return 1.0 - std::clamp(similarity, -1.0, 1.0); + } + case MetricType::L2: { + const auto dot = computeDotI16(leftValues, rightValues, dimension); + const auto l2sq = leftView.scale * leftView.scale * leftView.normSq + + rightView.scale * rightView.scale * rightView.normSq - + 2.0 * leftView.scale * rightView.scale * dot; + return std::sqrt(std::max(0.0, l2sq)); + } + case MetricType::L2_SQUARE: { + const auto dot = computeDotI16(leftValues, rightValues, dimension); + return std::max(0.0, leftView.scale * leftView.scale * leftView.normSq + + rightView.scale * rightView.scale * rightView.normSq - + 2.0 * leftView.scale * rightView.scale * dot); + } + default: { + UNREACHABLE_CODE; + } + } +} + +metric_func_t HNSWIndexUtils::getMetricsFunction(MetricType metric, const common::LogicalType& type, + QuantizationType quantization) { + if (quantization != QuantizationType::NONE && metric != MetricType::DotProduct) { + return HNSWIndexUtils::getQuantizedCachedMetricsFunction(metric, quantization); + } switch (metric) { case MetricType::Cosine: { return computeDistanceFuncDispatch(type); @@ -134,6 +385,48 @@ metric_func_t HNSWIndexUtils::getMetricsFunction(MetricType metric, } } +metric_func_t HNSWIndexUtils::getQuantizedCachedMetricsFunction(MetricType metric, + QuantizationType quantization) { + if (quantization == QuantizationType::SQ8) { + return [metric](const void* left, const void* right, uint32_t dimension) { + return computeCachedSQ8Distance(left, right, dimension, metric); + }; + } + if (quantization == QuantizationType::SQ16) { + return [metric](const void* left, const void* right, uint32_t dimension) { + return computeCachedSQ16Distance(left, right, dimension, metric); + }; + } + throw common::RuntimeException("Quantized cached metric function requires SQ8 or SQ16."); +} + +uint64_t HNSWIndexUtils::getQuantizedEmbeddingPayloadOffset(QuantizationType quantization) { + return getQuantizedPayloadOffsetInternal(quantization); +} + +uint64_t HNSWIndexUtils::getQuantizedEmbeddingStride(uint64_t dimension, + QuantizationType quantization) { + return alignUp(getQuantizedPayloadOffsetInternal(quantization) + + getQuantizedEmbeddingPayloadBytesInternal(dimension, quantization), + quantization == QuantizationType::SQ8 ? QUANTIZED_SIMD_ALIGNMENT : 1); +} + +uint64_t HNSWIndexUtils::getQuantizedCachedEmbeddingStride(uint64_t dimension, + QuantizationType quantization) { + return alignUp(getQuantizedCachedEmbeddingPayloadBytesInternal(dimension, quantization), + quantization == QuantizationType::SQ8 ? QUANTIZED_SIMD_ALIGNMENT : 1); +} + +uint64_t HNSWIndexUtils::getQuantizedEmbeddingPayloadBytes(uint64_t dimension, + QuantizationType quantization) { + return getQuantizedEmbeddingPayloadBytesInternal(dimension, quantization); +} + +uint64_t HNSWIndexUtils::getQuantizedCachedEmbeddingPayloadBytes(uint64_t dimension, + QuantizationType quantization) { + return getQuantizedCachedEmbeddingPayloadBytesInternal(dimension, quantization); +} + void HNSWIndexUtils::validateColumnType(const common::LogicalType& type) { if (type.getLogicalTypeID() == common::LogicalTypeID::ARRAY) { auto& childType = diff --git a/vector/test/test_files/error.test b/vector/test/test_files/error.test index da749b5c..51ad9dee 100644 --- a/vector/test/test_files/error.test +++ b/vector/test/test_files/error.test @@ -67,6 +67,12 @@ Binder exception: Alpha must be a double greater than or equal to 1. -STATEMENT CALL CREATE_VECTOR_INDEX('embeddings', 'e_hnsw_index', 'vec', metric := 'invalid'); ---- error Binder exception: Metric must be one of COSINE, L2, L2SQ, DOTPRODUCT or IP. +-STATEMENT CALL CREATE_VECTOR_INDEX('embeddings', 'e_hnsw_index', 'vec', quantization := 'invalid'); +---- error +Binder exception: Quantization must be one of NONE, SQ8 or SQ16. +-STATEMENT CALL CREATE_VECTOR_INDEX('embeddings', 'e_hnsw_index', 'vec', metric := 'dotproduct', quantization := 'sq8'); +---- error +Binder exception: Quantization currently supports COSINE, L2 and L2SQ metrics only. -STATEMENT CALL CREATE_VECTOR_INDEX('embeddings', 'e_hnsw_index', 'vec'); ---- ok -STATEMENT CALL QUERY_VECTOR_INDEX('embeddings', 'e_hnsw_index', CAST([0.0459,0.0439,0.0251,0.1,0.2,0.3,0.4,0.4], 'FLOAT[8]'), 10, unknown_param := 1) RETURN *; diff --git a/vector/test/test_files/small.test b/vector/test/test_files/small.test index 0bc42b48..88b4a3b0 100644 --- a/vector/test/test_files/small.test +++ b/vector/test/test_files/small.test @@ -114,6 +114,72 @@ 444 133 +-CASE ScalarQuantization +-LOAD_DYNAMIC_EXTENSION vector +-STATEMENT CREATE NODE TABLE embeddings (id int64, vec FLOAT[4], PRIMARY KEY (id)); +---- ok +-STATEMENT CREATE (:embeddings {id: 0, vec: [0.0, 1.0, 0.0, 0.0]}); +---- ok +-STATEMENT CREATE (:embeddings {id: 1, vec: [1.0, 0.0, 0.0, 0.0]}); +---- ok +-STATEMENT CREATE (:embeddings {id: 2, vec: [0.9, 0.1, 0.0, 0.0]}); +---- ok +-STATEMENT CREATE (:embeddings {id: 3, vec: [-1.0, 0.0, 0.0, 0.0]}); +---- ok +-STATEMENT CALL CREATE_VECTOR_INDEX('embeddings', 'sq8_index', 'vec', metric := 'l2', quantization := 'sq8'); +---- ok +-STATEMENT CALL QUERY_VECTOR_INDEX('embeddings', 'sq8_index', [1.0, 0.0, 0.0, 0.0], 2, search_type := 'naive') RETURN node.id ORDER BY distance; +-CHECK_ORDER +---- 2 +1 +2 +-RELOADDB +-LOAD_DYNAMIC_EXTENSION vector +-STATEMENT CALL QUERY_VECTOR_INDEX('embeddings', 'sq8_index', [1.0, 0.0, 0.0, 0.0], 2, search_type := 'naive') RETURN node.id ORDER BY distance; +-CHECK_ORDER +---- 2 +1 +2 +-STATEMENT CALL CREATE_VECTOR_INDEX('embeddings', 'sq8_full_index', 'vec', metric := 'l2', quantization := 'sq8', use_full_precision_rerank := true); +---- ok +-STATEMENT CALL QUERY_VECTOR_INDEX('embeddings', 'sq8_full_index', [1.0, 0.0, 0.0, 0.0], 2, search_type := 'naive') RETURN node.id ORDER BY distance; +-CHECK_ORDER +---- 2 +1 +2 +-STATEMENT CALL DROP_VECTOR_INDEX('embeddings', 'sq8_full_index'); +---- ok +-STATEMENT CALL DROP_VECTOR_INDEX('embeddings', 'sq8_index'); +---- ok +-STATEMENT CALL CREATE_VECTOR_INDEX('embeddings', 'sq16_index', 'vec', metric := 'l2sq', quantization := 'sq16', cache_embeddings := false); +---- ok +-STATEMENT CALL QUERY_VECTOR_INDEX('embeddings', 'sq16_index', [1.0, 0.0, 0.0, 0.0], 2, search_type := 'naive') RETURN node.id ORDER BY distance; +-CHECK_ORDER +---- 2 +1 +2 +-RELOADDB +-LOAD_DYNAMIC_EXTENSION vector +-STATEMENT CALL QUERY_VECTOR_INDEX('embeddings', 'sq16_index', [1.0, 0.0, 0.0, 0.0], 2, search_type := 'naive') RETURN node.id ORDER BY distance; +-CHECK_ORDER +---- 2 +1 +2 +-STATEMENT CALL CREATE_VECTOR_INDEX('embeddings', 'sq8_cos_index', 'vec', metric := 'cosine', quantization := 'sq8'); +---- ok +-STATEMENT CALL QUERY_VECTOR_INDEX('embeddings', 'sq8_cos_index', [1.0, 0.0, 0.0, 0.0], 2, search_type := 'naive') RETURN node.id ORDER BY distance; +-CHECK_ORDER +---- 2 +1 +2 +-RELOADDB +-LOAD_DYNAMIC_EXTENSION vector +-STATEMENT CALL QUERY_VECTOR_INDEX('embeddings', 'sq8_cos_index', [1.0, 0.0, 0.0, 0.0], 2, search_type := 'naive') RETURN node.id ORDER BY distance; +-CHECK_ORDER +---- 2 +1 +2 + -CASE 8DimL2Sq -LOAD_DYNAMIC_EXTENSION vector -STATEMENT CREATE NODE TABLE embeddings (id int64, vec FLOAT[8], PRIMARY KEY (id));