[Feature] Add scalar quantization support for vector extension

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;
 }
 

vector/src/function/create_hnsw_index.cpp

@@ -34,7 +34,7 @@ CreateInMemHNSWSharedState::CreateInMemHNSWSharedState(const CreateHNSWIndexBind
                     ->getTable(bindData.tableEntry->getTableID())
                     ->cast<storage::NodeTable>()},
       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<InMemHNSWIndex>(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<catalog::NodeTableCatalogEntry>();
+        quantizedEmbeddingsTableID = quantizedEmbeddingsTableEntry->getTableID();
+        quantizedEmbeddingsTable =
+            storageManager->getTable(quantizedEmbeddingsTableID)->ptrCast<storage::NodeTable>();
+    }
     auto storageInfo = std::make_unique<HNSWStorageInfo>(upperTableID, lowerTableID,
-        index->getUpperEntryPoint(), index->getLowerEntryPoint(), bindData->numRows);
+        quantizedEmbeddingsTableID, index->getUpperEntryPoint(), index->getLowerEntryPoint(),
+        bindData->numRows);
     auto onDiskIndex = std::make_unique<OnDiskHNSWIndex>(context->clientContext, indexInfo,
         std::move(storageInfo), bindData->config.copy());
-    auto storageManager = storage::StorageManager::Get(*clientContext);
     auto nodeTable = storageManager->getTable(nodeTableID)->ptrCast<storage::NodeTable>();
     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<common::ArrayTypeInfo>();
+        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<common::ArrayTypeInfo>();
+        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);
     }

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<TableFuncBindData> copy() const override {
-        return std::make_unique<DropHNSWIndexBindData>(tableEntry, indexName, skipAfterBind);
+        return std::make_unique<DropHNSWIndexBindData>(tableEntry, indexName,
+            hasQuantizedEmbeddingsTable, skipAfterBind);
     }
 };
 
@@ -38,9 +42,15 @@ static std::unique_ptr<TableFuncBindData> 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<DropHNSWIndexBindData>(nullptr, indexName, true);
+        return std::make_unique<DropHNSWIndexBindData>(nullptr, indexName, false, true);
     }
-    return std::make_unique<DropHNSWIndexBindData>(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<HNSWIndexAuxInfo>();
+    return std::make_unique<DropHNSWIndexBindData>(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;";
     }

vector/src/function/query_hnsw_index.cpp

@@ -1,3 +1,6 @@
+#include <cmath>
+#include <cstring>
+
 #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<typename T>
-struct HNSWQueryVector : GetEmbeddingsScanState {
-    HNSWQueryVector(main::ClientContext* context, std::shared_ptr<Expression> queryExpression,
+struct HNSWRawQueryVector : GetEmbeddingsScanState {
+    HNSWRawQueryVector(main::ClientContext* context, std::shared_ptr<Expression> queryExpression,
         const LogicalType& indexType, uint64_t dimension)
         : data(getQueryVector<T>(context, std::move(queryExpression), indexType, dimension)) {}
 
@@ -224,6 +227,81 @@ struct HNSWQueryVector : GetEmbeddingsScanState {
     std::vector<T> data;
 };
 
+template<VectorElementType T>
+static void quantizeQueryVector(const std::vector<T>& 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<double>(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<float>(maxAbs / maxQuantizedValue);
+    float normSq = 0.0f;
+    if (quantization == QuantizationType::SQ8) {
+        for (auto i = 0u; i < src.size(); ++i) {
+            const auto quantized =
+                static_cast<int64_t>(std::llround(static_cast<double>(src[i]) / scale));
+            const auto clamped = static_cast<int8_t>(std::clamp<int64_t>(quantized, -127, 127));
+            payloadDst[i] = static_cast<uint8_t>(clamped);
+            normSq += static_cast<float>(clamped) * static_cast<float>(clamped);
+        }
+        normSqDst = normSq;
+        scaleDst = metric == MetricType::Cosine && normSq > 0.0f ? 1.0f / std::sqrt(normSq) :
+                                                                    scale;
+        return;
+    }
+    auto* payloadValues = reinterpret_cast<int16_t*>(payloadDst);
+    for (auto i = 0u; i < src.size(); ++i) {
+        const auto quantized =
+            static_cast<int64_t>(std::llround(static_cast<double>(src[i]) / scale));
+        const auto clamped = static_cast<int16_t>(std::clamp<int64_t>(quantized, -32767, 32767));
+        normSq += static_cast<float>(clamped) * static_cast<float>(clamped);
+        payloadValues[i] = clamped;
+    }
+    normSqDst = normSq;
+    scaleDst = metric == MetricType::Cosine && normSq > 0.0f ? 1.0f / std::sqrt(normSq) : scale;
+}
+
+template<typename T>
+struct HNSWQuantizedQueryVector : GetEmbeddingsScanState {
+    HNSWQuantizedQueryVector(main::ClientContext* context,
+        std::shared_ptr<Expression> 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<T>(context, std::move(queryExpression), indexType, dimension);
+        alignedData = reinterpret_cast<uint8_t*>(
+            common::ceilDiv<uintptr_t>(reinterpret_cast<uintptr_t>(data.data()),
+                static_cast<uintptr_t>(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<uint8_t> data;
+    uint8_t* alignedData;
+    QuantizedEmbeddingView view;
+};
+
 static offset_t tableFunc(const TableFuncInput& input, TableFuncOutput& output) {
     const auto localState = input.localState->ptrCast<QueryHNSWLocalState>();
     const auto bindData = input.bindData->constPtrCast<QueryHNSWIndexBindData>();
@@ -242,12 +320,22 @@ static offset_t tableFunc(const TableFuncInput& input, TableFuncOutput& output)
         TypeUtils::visit(
             indexType,
             [&]<VectorElementType T>(T) {
-                auto queryVector = HNSWQueryVector<T>(input.context->clientContext,
+                auto exactQueryVector = HNSWRawQueryVector<T>(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<T>(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<TableFuncLocalState> initQueryHNSWLocalState(
         catalog
             ->getTableCatalogEntry(transaction::Transaction::Get(*context), lowerRelTableName, true)
             ->ptrCast<RelGroupCatalogEntry>();
+    const auto& indexConfig =
+        hnswBindData->indexEntry->getAuxInfo().cast<HNSWIndexAuxInfo>().config;
+    auto indexOpt = hnswSharedState->nodeTable->getIndex(hnswBindData->indexEntry->getIndexName());
+    DASSERT(indexOpt.has_value());
+    auto& index = indexOpt.value()->cast<OnDiskHNSWIndex>();
+    auto quantizedEmbeddings =
+        index.getOrCreateQuantizedEmbeddings(context, hnswSharedState->numNodes);
     HNSWSearchState searchState{context, hnswBindData->nodeTableEntry, upperRelTableEntry,
         lowerRelTableEntry, *hnswSharedState->nodeTable, hnswBindData->indexColumnID,
-        hnswSharedState->numNodes, static_cast<uint64_t>(k), hnswBindData->config};
+        hnswSharedState->numNodes, static_cast<uint64_t>(k), hnswBindData->config, indexConfig,
+        true, std::move(quantizedEmbeddings)};
     const auto tableID = hnswBindData->nodeTableEntry->getTableID();
     auto& semiMasks = hnswSharedState->semiMasks;
     if (semiMasks.containsTableID(tableID)) {

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 {