diff --git a/vector/docs/scalar_quantization_hnsw.md b/vector/docs/scalar_quantization_hnsw.md new file mode 100644 index 00000000..d604459e --- /dev/null +++ b/vector/docs/scalar_quantization_hnsw.md @@ -0,0 +1,257 @@ +# 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`. +- 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. + +## 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 transaction-local +dense cache 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 + transaction-local `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. Remove any transaction-local derived qemb cache entry. + +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 transaction-local 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`. + +Within one query execution, local HNSW scan states share a mutex while obtaining the quantized +embedding provider. This prevents duplicate dense cache construction and keeps storage-backed +provider construction serialized. + +This gives the cached quantized path the same hot-query shape as the regular cached baseline while +retaining a smaller quantized representation. + +## 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 transaction-local 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..5c491c57 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,169 @@ 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; + } + [[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 +258,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 +288,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..74ad3ab8 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,49 @@ 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(); + 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 +920,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 +936,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 +980,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 +1080,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 +1119,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 +1145,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 +1186,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 +1211,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 +1242,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 +1274,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 +1320,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 +1344,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 +1422,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 +1454,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 +1476,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 +1543,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));