diff --git a/examples/cpp/115_index_brute_force_hamming.cpp b/examples/cpp/115_index_brute_force_hamming.cpp new file mode 100644 index 0000000000..2df61c21ac --- /dev/null +++ b/examples/cpp/115_index_brute_force_hamming.cpp @@ -0,0 +1,126 @@ +// Copyright 2024-present the vsag project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#include +#include +#include +#include +#include + +// Example: Hamming distance over bit-packed binary vectors with the brute_force index. +// +// For binary vectors, "dim" is the number of logical bits. Each vector is stored as a +// contiguous byte array of ceil(dim / 8) bytes. Bits are packed least-significant-bit +// first inside each byte, and the unused high bits of the final byte are ignored by the +// Hamming distance computation: +// +// hamming(a, b) = popcount(a XOR b) + +int +main(int argc, char** argv) { + vsag::init(); + + /******************* Prepare Base Dataset *****************/ + int64_t num_vectors = 1000; + int64_t dim = 256; // number of logical bits + int64_t bytes_per_vector = (dim + 7) / 8; + + std::vector ids(num_vectors); + std::vector datas(num_vectors * bytes_per_vector); + std::mt19937 rng(47); + std::uniform_int_distribution distrib_byte(0, 255); + for (int64_t i = 0; i < num_vectors; ++i) { + ids[i] = i; + } + for (auto& byte : datas) { + byte = static_cast(distrib_byte(rng)); + } + + auto base = vsag::Dataset::Make(); + base->NumElements(num_vectors) + ->Dim(dim) + ->Ids(ids.data()) + ->BinaryVectors(datas.data()) + ->Owner(false); + + /******************* Create BruteForce Index *****************/ + std::string index_parameters = R"( + { + "dtype": "binary", + "metric_type": "hamming", + "dim": 256 + } + )"; + auto index = vsag::Factory::CreateIndex("brute_force", index_parameters).value(); + + /******************* Build BruteForce Index *****************/ + if (auto build_result = index->Build(base); build_result.has_value()) { + std::cout << "After Build(), Index BruteForce contains: " << index->GetNumElements() + << std::endl; + } else { + std::cerr << "Failed to build index: " << build_result.error().message << std::endl; + return -1; + } + + /******************* Prepare Query Dataset *****************/ + // Query with a copy of the last base vector flipped in two bits, so the exact + // Hamming distance to that vector must be 2. + std::vector query_vector(bytes_per_vector); + std::memcpy( + query_vector.data(), datas.data() + (num_vectors - 1) * bytes_per_vector, bytes_per_vector); + query_vector[0] ^= 0x05; // flip two bits + + auto query = vsag::Dataset::Make(); + query->NumElements(1)->Dim(dim)->BinaryVectors(query_vector.data())->Owner(false); + + /******************* KnnSearch For BruteForce Index *****************/ + int64_t topk = 10; + auto result = index->KnnSearch(query, topk, R"({})").value(); + + /******************* Verify and Print Search Result *****************/ + auto popcount = [](const uint8_t* x, const uint8_t* y, int64_t num_bytes) { + uint64_t distance = 0; + for (int64_t i = 0; i < num_bytes; ++i) { + uint8_t diff = static_cast(x[i] ^ y[i]); + while (diff != 0) { + diff = static_cast(diff & (diff - 1)); + ++distance; + } + } + return distance; + }; + + std::cout << "results (id: distance):" << std::endl; + for (int64_t i = 0; i < result->GetDim(); ++i) { + int64_t id = result->GetIds()[i]; + auto expected = + popcount(query_vector.data(), datas.data() + id * bytes_per_vector, bytes_per_vector); + if (static_cast(result->GetDistances()[i]) != expected) { + std::cerr << "distance mismatch for id " << id << std::endl; + return -1; + } + std::cout << id << ": " << result->GetDistances()[i] << std::endl; + } + + if (result->GetIds()[0] != num_vectors - 1 || result->GetDistances()[0] != 2.0F) { + std::cerr << "nearest neighbor verification failed" << std::endl; + return -1; + } + std::cout << "verified: nearest id " << result->GetIds()[0] << " has Hamming distance " + << result->GetDistances()[0] << std::endl; + + return 0; +} diff --git a/include/vsag/constants.h b/include/vsag/constants.h index cd76d9b557..fe724c8860 100644 --- a/include/vsag/constants.h +++ b/include/vsag/constants.h @@ -39,6 +39,7 @@ extern const char* const FLOAT32_VECTORS; extern const char* const FLOAT16_VECTORS; extern const char* const SPARSE_VECTORS; extern const char* const INT8_VECTORS; +extern const char* const BINARY_VECTORS; extern const char* const ATTRIBUTE_SETS; extern const char* const DATASET_PATHS; extern const char* const EXTRA_INFOS; @@ -53,10 +54,12 @@ extern const char* const SIMPLEFLAT_IDS; extern const char* const METRIC_L2; extern const char* const METRIC_COSINE; extern const char* const METRIC_IP; +extern const char* const METRIC_HAMMING; extern const char* const DATATYPE_FLOAT32; extern const char* const DATATYPE_FLOAT16; extern const char* const DATATYPE_BFLOAT16; extern const char* const DATATYPE_INT8; +extern const char* const DATATYPE_BINARY; extern const char* const DATATYPE_SPARSE; extern const char* const BLANK_INDEX; diff --git a/include/vsag/dataset.h b/include/vsag/dataset.h index a1f075cd49..2d4de1ac6e 100644 --- a/include/vsag/dataset.h +++ b/include/vsag/dataset.h @@ -252,6 +252,25 @@ class Dataset : public std::enable_shared_from_this { virtual const int8_t* GetInt8Vectors() const = 0; + /** + * @brief Sets the bit-packed binary vector array for the dataset. + * + * Each vector occupies ceil(Dim() / 8) bytes. Unused high bits in the final byte are ignored. + * + * @param vectors Pointer to the contiguous array of packed binary vectors. + * @return DatasetPtr A shared pointer to the dataset with updated binary vectors. + */ + virtual DatasetPtr + BinaryVectors(const uint8_t* vectors) = 0; + + /** + * @brief Retrieves the bit-packed binary vector array of the dataset. + * + * @return const uint8_t* Pointer to the array of packed binary vectors. + */ + virtual const uint8_t* + GetBinaryVectors() const = 0; + /** * @brief Sets the float32 vector array for the dataset. * diff --git a/python_bindings/index_binding.cpp b/python_bindings/index_binding.cpp index 8a58d70aa9..576dbc2432 100644 --- a/python_bindings/index_binding.cpp +++ b/python_bindings/index_binding.cpp @@ -21,6 +21,7 @@ #include "binding.h" #include "fmt/format.h" +#include "storage/serialization.h" #include "vsag/constants.h" #include "vsag/dataset.h" #include "vsag/vsag.h" @@ -32,7 +33,7 @@ to_int64(uint64_t value) { return static_cast(value); } -enum class DenseVectorKind { FLOAT32, FLOAT16, BFLOAT16, UNSUPPORTED }; +enum class DenseVectorKind { FLOAT32, FLOAT16, BFLOAT16, BINARY, UNSUPPORTED }; DenseVectorKind parse_dense_vector_kind(const std::string& parameters) { @@ -52,6 +53,9 @@ parse_dense_vector_kind(const std::string& parameters) { if (dtype_name == vsag::DATATYPE_BFLOAT16) { return DenseVectorKind::BFLOAT16; } + if (dtype_name == vsag::DATATYPE_BINARY) { + return DenseVectorKind::BINARY; + } return DenseVectorKind::UNSUPPORTED; } @@ -64,6 +68,8 @@ expected_numpy_dtype(DenseVectorKind kind) { return "numpy.float16"; case DenseVectorKind::BFLOAT16: return "numpy.uint16 (raw bfloat16 bits)"; + case DenseVectorKind::BINARY: + return "numpy.uint8 (bit-packed binary vectors)"; case DenseVectorKind::UNSUPPORTED: return "unsupported"; } @@ -79,6 +85,8 @@ expected_numpy_dtype_name(DenseVectorKind kind) { return "float16"; case DenseVectorKind::BFLOAT16: return "uint16"; + case DenseVectorKind::BINARY: + return "uint8"; case DenseVectorKind::UNSUPPORTED: return ""; } @@ -86,14 +94,22 @@ expected_numpy_dtype_name(DenseVectorKind kind) { } void -validate_dense_index_kind(DenseVectorKind kind, const char* operation_name) { +validate_dense_index_kind(DenseVectorKind kind, + const char* operation_name, + bool allow_unsupported = false) { if (kind == DenseVectorKind::UNSUPPORTED) { + if (allow_unsupported) { + // Unsupported (sparse) indexes reject dense queries inside C++, which the + // binding surfaces as RuntimeError; do not short-circuit with ValueError. + return; + } throw std::invalid_argument( - fmt::format("{} only supports indexes with dtype in [{}, {}, {}]", + fmt::format("{} only supports indexes with dtype in [{}, {}, {}, {}]", operation_name, vsag::DATATYPE_FLOAT32, vsag::DATATYPE_FLOAT16, - vsag::DATATYPE_BFLOAT16)); + vsag::DATATYPE_BFLOAT16, + vsag::DATATYPE_BINARY)); } } @@ -172,6 +188,10 @@ set_dense_vectors(const vsag::DatasetPtr& dataset, dataset->Float32Vectors(static_cast(buf.ptr)); return; } + if (kind == DenseVectorKind::BINARY) { + dataset->BinaryVectors(static_cast(buf.ptr)); + return; + } dataset->Float16Vectors(static_cast(buf.ptr)); } @@ -267,13 +287,23 @@ class Index { } } dense_vector_kind_ = parse_dense_vector_kind(parameters); + if (dense_vector_kind_ == DenseVectorKind::BINARY) { + const auto json = nlohmann::json::parse(parameters); + dim_ = json.at(vsag::PARAMETER_DIM).get(); + } } void Build(py::array vectors, py::array_t ids, uint64_t num_elements, uint64_t dim) { validate_dense_index_kind(dense_vector_kind_, "build"); - auto buf = - validate_dense_array(vectors, dense_vector_kind_, "vectors", num_elements * dim, dim); + if (dense_vector_kind_ == DenseVectorKind::BINARY && dim != dim_) { + throw std::invalid_argument( + fmt::format("dim ({}) must equal the index dimension ({})", dim, dim_)); + } + const uint64_t input_dim = + dense_vector_kind_ == DenseVectorKind::BINARY ? (dim + 7U) / 8U : dim; + auto buf = validate_dense_array( + vectors, dense_vector_kind_, "vectors", num_elements * input_dim, input_dim); validate_ids(ids, num_elements); auto dataset = vsag::Dataset::Make(); @@ -319,10 +349,15 @@ class Index { py::object KnnSearch(py::array vector, uint64_t k, std::string& parameters) { validate_dense_index_kind(dense_vector_kind_, "knn_search"); - auto buf = validate_dense_array(vector, dense_vector_kind_, "vector"); + const uint64_t query_size = + dense_vector_kind_ == DenseVectorKind::BINARY ? (dim_ + 7U) / 8U : 0; + auto buf = validate_dense_array(vector, dense_vector_kind_, "vector", query_size); auto query = vsag::Dataset::Make(); - query->NumElements(1)->Dim(to_int64(static_cast(buf.shape[0])))->Owner(false); + const uint64_t query_dim = dense_vector_kind_ == DenseVectorKind::BINARY + ? dim_ + : static_cast(buf.shape[0]); + query->NumElements(1)->Dim(to_int64(query_dim))->Owner(false); set_dense_vectors(query, buf, dense_vector_kind_); uint64_t ids_shape[1]{k}; @@ -355,9 +390,14 @@ class Index { py::tuple KnnSearchWithStatistics(const py::array& vector, uint64_t k, const std::string& parameters) { validate_dense_index_kind(dense_vector_kind_, "knn_search_with_statistics"); - auto buf = validate_dense_array(vector, dense_vector_kind_, "vector"); + const uint64_t query_size = + dense_vector_kind_ == DenseVectorKind::BINARY ? (dim_ + 7U) / 8U : 0; + auto buf = validate_dense_array(vector, dense_vector_kind_, "vector", query_size); auto query = vsag::Dataset::Make(); - query->NumElements(1)->Dim(to_int64(static_cast(buf.shape[0])))->Owner(false); + const uint64_t query_dim = dense_vector_kind_ == DenseVectorKind::BINARY + ? dim_ + : static_cast(buf.shape[0]); + query->NumElements(1)->Dim(to_int64(query_dim))->Owner(false); set_dense_vectors(query, buf, dense_vector_kind_); uint64_t ids_shape[1]{k}; uint64_t ids_strides[1]{sizeof(int64_t)}; @@ -457,10 +497,15 @@ class Index { py::object RangeSearch(py::array point, float threshold, std::string& parameters) { validate_dense_index_kind(dense_vector_kind_, "range_search"); - auto buf = validate_dense_array(point, dense_vector_kind_, "point"); + const uint64_t query_size = + dense_vector_kind_ == DenseVectorKind::BINARY ? (dim_ + 7U) / 8U : 0; + auto buf = validate_dense_array(point, dense_vector_kind_, "point", query_size); auto query = vsag::Dataset::Make(); - query->NumElements(1)->Dim(to_int64(static_cast(buf.shape[0])))->Owner(false); + const uint64_t query_dim = dense_vector_kind_ == DenseVectorKind::BINARY + ? dim_ + : static_cast(buf.shape[0]); + query->NumElements(1)->Dim(to_int64(query_dim))->Owner(false); set_dense_vectors(query, buf, dense_vector_kind_); py::array_t labels; @@ -487,9 +532,14 @@ class Index { float threshold, const std::string& parameters) { validate_dense_index_kind(dense_vector_kind_, "range_search_with_statistics"); - auto buf = validate_dense_array(point, dense_vector_kind_, "point"); + const uint64_t query_size = + dense_vector_kind_ == DenseVectorKind::BINARY ? (dim_ + 7U) / 8U : 0; + auto buf = validate_dense_array(point, dense_vector_kind_, "point", query_size); auto query = vsag::Dataset::Make(); - query->NumElements(1)->Dim(to_int64(static_cast(buf.shape[0])))->Owner(false); + const uint64_t query_dim = dense_vector_kind_ == DenseVectorKind::BINARY + ? dim_ + : static_cast(buf.shape[0]); + query->NumElements(1)->Dim(to_int64(query_dim))->Owner(false); set_dense_vectors(query, buf, dense_vector_kind_); auto result = index_->RangeSearch(query, threshold, parameters); @@ -537,8 +587,14 @@ class Index { void Add(py::array vectors, py::array_t ids, uint64_t num_elements, uint64_t dim) { validate_dense_index_kind(dense_vector_kind_, "add"); - auto buf = - validate_dense_array(vectors, dense_vector_kind_, "vectors", num_elements * dim, dim); + if (dense_vector_kind_ == DenseVectorKind::BINARY && dim != dim_) { + throw std::invalid_argument( + fmt::format("dim ({}) must equal the index dimension ({})", dim, dim_)); + } + const uint64_t input_dim = + dense_vector_kind_ == DenseVectorKind::BINARY ? (dim + 7U) / 8U : dim; + auto buf = validate_dense_array( + vectors, dense_vector_kind_, "vectors", num_elements * input_dim, input_dim); validate_ids(ids, num_elements); auto dataset = vsag::Dataset::Make(); @@ -570,9 +626,28 @@ class Index { } py::array_t - CalcDistancesById(const py::array_t& query, + CalcDistancesById(const py::array& query, const py::array_t& ids) { - auto buf_query = query.request(); + // Unsupported (sparse) indexes keep the legacy RuntimeError contract here. + validate_dense_index_kind(dense_vector_kind_, + "calc_distances_by_id", + /*allow_unsupported=*/true); + const bool is_binary = dense_vector_kind_ == DenseVectorKind::BINARY; + py::array validated_query = query; + if (not is_binary) { + using FloatQuery = py::array_t; + auto converted = FloatQuery::ensure(query); + if (not converted) { + throw std::invalid_argument("query must be convertible to a float32 array"); + } + validated_query = converted; + } + const uint64_t query_size = is_binary ? (dim_ + 7U) / 8U : 0; + auto buf_query = + validate_dense_array(validated_query, + is_binary ? DenseVectorKind::BINARY : DenseVectorKind::FLOAT32, + "query", + query_size); auto buf_ids = ids.request(); if (buf_query.ndim != 1 || buf_ids.ndim != 1) { @@ -589,10 +664,12 @@ class Index { auto dist_view = distances.mutable_unchecked<1>(); auto dataset = vsag::Dataset::Make(); - dataset->Owner(false) - ->NumElements(1) - ->Dim(buf_query.shape[0]) - ->Float32Vectors(query.data()); + dataset->Owner(false)->NumElements(1)->Dim(is_binary ? dim_ : buf_query.shape[0]); + if (is_binary) { + dataset->BinaryVectors(static_cast(buf_query.ptr)); + } else { + dataset->Float32Vectors(static_cast(buf_query.ptr)); + } auto result = index_->CalcDistancesById(dataset, ids.data(), count); if (!result.has_value()) { throw std::runtime_error( @@ -608,7 +685,7 @@ class Index { } py::array_t - CalDistanceById(const py::array_t& query, + CalDistanceById(const py::array& query, const py::array_t& ids) { return this->CalcDistancesById(query, ids); } @@ -623,12 +700,33 @@ class Index { void Load(const std::string& filename) { std::ifstream file(filename, std::ios::binary); - index_->Deserialize(file); + auto result = index_->Deserialize(file); + if (not result.has_value()) { + file.close(); + throw std::runtime_error(fmt::format("load failed: {}", result.error().message)); + } + if (dense_vector_kind_ == DenseVectorKind::BINARY) { + file.clear(); + file.seekg(0); + vsag::IOStreamReader reader(file); + auto footer = vsag::Footer::Parse(reader); + if (footer == nullptr) { + file.close(); + throw std::runtime_error("load failed: serialized binary index has no footer"); + } + const auto basic_info = footer->GetMetadata()->Get("basic_info"); + if (not basic_info.Contains(vsag::PARAMETER_DIM)) { + file.close(); + throw std::runtime_error("load failed: serialized binary index has no dimension"); + } + dim_ = static_cast(basic_info[vsag::PARAMETER_DIM].GetInt()); + } file.close(); } private: DenseVectorKind dense_vector_kind_{DenseVectorKind::FLOAT32}; + uint64_t dim_{0}; std::shared_ptr index_; }; @@ -647,6 +745,7 @@ bind_index(py::module_& module) { - float32 indexes accept numpy.float32 vectors - float16 indexes accept numpy.float16 vectors - bfloat16 indexes accept numpy.uint16 arrays containing raw bfloat16 bits + - binary indexes accept numpy.uint8 arrays containing ceil(dim / 8) packed bytes per vector )pbdoc") .def(py::init(), py::arg("name"), @@ -898,16 +997,17 @@ bind_index(py::module_& module) { Calculate distances between a query vector and vectors specified by IDs. Args: - query (numpy.ndarray): 1D array of float32 values representing the query vector + query (numpy.ndarray): 1D float32 query, or packed uint8 bytes for binary indexes ids (numpy.ndarray): 1D array of int64 IDs to calculate distances for Returns: numpy.ndarray: 1D float32 distances in ID order; missing IDs produce -1. Notes: - Query length must match the index dimension. Strided inputs are copied - to contiguous arrays. C++ operation failures raise RuntimeError rather - than returning missing-ID sentinels. Only one dense query is accepted. + Query length must match the index dimension (ceil(dim / 8) bytes for binary). + Strided non-binary inputs are copied to contiguous arrays. C++ operation failures + raise RuntimeError rather than returning missing-ID sentinels. Only one dense query + is accepted. )pbdoc") .def("cal_distance_by_id", &Index::CalDistanceById, diff --git a/src/algorithm/bruteforce/bruteforce.cpp b/src/algorithm/bruteforce/bruteforce.cpp index 35c9f693f7..db212d72f9 100644 --- a/src/algorithm/bruteforce/bruteforce.cpp +++ b/src/algorithm/bruteforce/bruteforce.cpp @@ -127,7 +127,7 @@ BruteForce::Train(const DatasetPtr& data) { if (is_multi_vector_) { this->train_multi_vector(data); } else { - this->inner_codes_->Train(data->GetFloat32Vectors(), data->GetNumElements()); + this->inner_codes_->Train(this->get_data(data, 0), data->GetNumElements()); } } @@ -225,7 +225,7 @@ BruteForce::Add(const DatasetPtr& data) { auto base_dim = data->GetDim(); CHECK_ARGUMENT(base_dim == dim_, fmt::format("base.dim({}) must be equal to index.dim({})", base_dim, dim_)); - CHECK_ARGUMENT(data->GetFloat32Vectors() != nullptr, "base.float_vector is nullptr"); + CHECK_ARGUMENT(this->get_data(data, 0) != nullptr, "base.vectors is nullptr"); { std::lock_guard lock(this->add_mutex_); @@ -253,7 +253,6 @@ BruteForce::Add(const DatasetPtr& data) { std::vector>> futures; const auto total = data->GetNumElements(); const auto* labels = data->GetIds(); - const auto* vectors = data->GetFloat32Vectors(); const auto* attrs = data->GetAttributeSets(); const auto* extra_info = data->GetExtraInfos(); const auto extra_info_size = data->GetExtraInfoSize(); @@ -274,14 +273,14 @@ BruteForce::Add(const DatasetPtr& data) { const auto* ei_ptr = extra_info == nullptr ? nullptr : extra_info + j * extra_info_size; if (this->thread_pool_ != nullptr) { auto future = this->thread_pool_->GeneralEnqueue(add_func, - vectors + j * dim_, + static_cast(this->get_data(data, j)), label, attrs == nullptr ? nullptr : attrs + j, ei_ptr); futures.emplace_back(std::move(future)); } else { if (auto add_res = add_func( - vectors + j * dim_, label, attrs == nullptr ? nullptr : attrs + j, ei_ptr); + static_cast(this->get_data(data, j)), label, attrs == nullptr ? nullptr : attrs + j, ei_ptr); add_res.has_value()) { failed_ids.emplace_back(add_res.value()); } @@ -697,13 +696,27 @@ BruteForce::make_search_computer(const DatasetPtr& query) const { CHECK_ARGUMENT(query_multi_vectors != nullptr, "query.multi_vectors is nullptr"); return this->inner_codes_->FactoryComputer(&query_multi_vectors[0]); } - return this->inner_codes_->FactoryComputer(query->GetFloat32Vectors()); + if (data_type_ == DataTypes::DATA_TYPE_BINARY) { + CHECK_ARGUMENT(query != nullptr, "query must not be null"); + this->validate_search_query(query); + CHECK_ARGUMENT(query->GetBinaryVectors() != nullptr, "query must contain binary vectors"); + } + return this->inner_codes_->FactoryComputer(this->get_data(query, 0)); } float BruteForce::CalcDistanceById(const DatasetPtr& query, int64_t id, bool calculate_precise_distance) const { + if (not is_multi_vector_ and data_type_ == DataTypes::DATA_TYPE_BINARY) { + CHECK_ARGUMENT(query != nullptr, "distance query must not be null"); + CHECK_ARGUMENT(query->GetNumElements() == 1, "single-ID distance requires one query"); + CHECK_ARGUMENT(query->GetDim() == dim_, "distance query dimension mismatch"); + const auto* vectors = query->GetBinaryVectors(); + CHECK_ARGUMENT(vectors != nullptr, "distance query must contain binary vectors"); + return this->CalcDistanceById( + reinterpret_cast(vectors), id, calculate_precise_distance); + } if (not is_multi_vector_) { return InnerIndexInterface::CalcDistanceById(query, id, calculate_precise_distance); } @@ -1151,7 +1164,7 @@ BruteForce::InitFeatures() { } } else { if (name != QUANTIZATION_TYPE_VALUE_FP32 and name != QUANTIZATION_TYPE_VALUE_BF16 and - name != QUANTIZATION_TYPE_VALUE_FP16) { + name != QUANTIZATION_TYPE_VALUE_FP16 and name != QUANTIZATION_TYPE_VALUE_BINARY) { this->index_feature_list_->SetFeature(IndexFeature::NEED_TRAIN); } else { this->index_feature_list_->SetFeatures( @@ -1161,8 +1174,9 @@ BruteForce::InitFeatures() { IndexFeature::SUPPORT_BATCH_CALC_DISTANCE_BY_ID, IndexFeature::SUPPORT_RANGE_SEARCH_WITH_ID_FILTER}); } - if (name == QUANTIZATION_TYPE_VALUE_FP32 and - (metric_ != MetricType::METRIC_TYPE_COSINE || this->inner_codes_->HoldMolds())) { + if (name == QUANTIZATION_TYPE_VALUE_BINARY or + (name == QUANTIZATION_TYPE_VALUE_FP32 and + (metric_ != MetricType::METRIC_TYPE_COSINE || this->inner_codes_->HoldMolds()))) { this->index_feature_list_->SetFeature(IndexFeature::SUPPORT_GET_RAW_VECTOR_BY_IDS); } } @@ -1254,6 +1268,9 @@ BruteForce::CheckAndMappingExternalParam(const JsonType& external_param, bool is_warp = external_param.Contains(WARP_MODE_MARKER); if (is_warp) { + CHECK_ARGUMENT(common_param.data_type_ != DataTypes::DATA_TYPE_BINARY, + "WARP does not support binary datatype"); + // Remove the marker key before mapping (it would cause "invalid config param" error) JsonType warp_external_param = external_param; warp_external_param.Erase(WARP_MODE_MARKER); @@ -1292,6 +1309,16 @@ BruteForce::CheckAndMappingExternalParam(const JsonType& external_param, } auto inner_json = build_default_brute_force_param(external_param, false); + if (common_param.data_type_ == DataTypes::DATA_TYPE_BINARY) { + if (not external_param.Contains(BRUTE_FORCE_BASE_QUANTIZATION_TYPE)) { + inner_json[BASE_CODES_KEY][QUANTIZATION_PARAMS_KEY][TYPE_KEY].SetString( + QUANTIZATION_TYPE_VALUE_BINARY); + } + if (not external_param.Contains(BRUTE_FORCE_PRECISE_QUANTIZATION_TYPE)) { + inner_json[PRECISE_CODES_KEY][QUANTIZATION_PARAMS_KEY][TYPE_KEY].SetString( + QUANTIZATION_TYPE_VALUE_BINARY); + } + } for (const auto& [key, value] : external_param.GetInnerJson()->items()) { (void)value; auto field = external_param[key]; @@ -1412,14 +1439,14 @@ BruteForce::UpdateVector(int64_t id, const DatasetPtr& new_base, bool force_upda auto base_dim = new_base->GetDim(); CHECK_ARGUMENT(base_dim == dim_, fmt::format("base.dim({}) must be equal to index.dim({})", base_dim, dim_)); - CHECK_ARGUMENT(new_base->GetFloat32Vectors() != nullptr, "base.float_vector is nullptr"); + CHECK_ARGUMENT(this->get_data(new_base, 0) != nullptr, "base.vectors is nullptr"); std::shared_lock add_lock(this->add_mutex_, std::defer_lock); std::shared_lock label_lock(this->label_lookup_mutex_, std::defer_lock); std::lock(add_lock, label_lock); std::unique_lock global_lock(this->global_mutex_); InnerIdType inner_id = this->label_table_->GetIdByLabel(id); - return this->inner_codes_->UpdateVector(new_base->GetFloat32Vectors(), inner_id); + return this->inner_codes_->UpdateVector(this->get_data(new_base, 0), inner_id); } void diff --git a/src/algorithm/hgraph/hgraph.cpp b/src/algorithm/hgraph/hgraph.cpp index 85470e408d..f6886299c1 100644 --- a/src/algorithm/hgraph/hgraph.cpp +++ b/src/algorithm/hgraph/hgraph.cpp @@ -196,6 +196,10 @@ HGraph::HGraph(const HGraphParameterPtr& hgraph_param, const vsag::IndexCommonPa optimizer_ = std::make_shared>(common_param); } check_and_init_raw_vector(hgraph_param->raw_vector_param, common_param); + if (data_type_ == DataTypes::DATA_TYPE_BINARY and + basic_flatten_codes_->GetQuantizerName() == QUANTIZATION_TYPE_VALUE_BINARY) { + this->index_feature_list_->SetFeature(IndexFeature::SUPPORT_GET_RAW_VECTOR_BY_IDS); + } if (this->using_dedup_storage()) { this->basic_flatten_codes_ = MakeCodeSlotFlattenAdapter( this->basic_flatten_codes_, this->code_slot_map_, allocator_, &this->total_count_); @@ -236,6 +240,9 @@ HGraph::Tune(const std::string& parameters, bool disable_future_tuning) { // map auto inner_json = map_hgraph_param(hgraph_json); + if (data_type_ == DataTypes::DATA_TYPE_BINARY) { + ApplyBinaryQuantizationDefaults(hgraph_json, inner_json); + } ValidateMRLEDim(hgraph_json, this->dim_); const bool requires_raw_vector = RequiresRawVectorForMRLERaBitQSplit(inner_json); if (requires_raw_vector) { @@ -245,11 +252,16 @@ HGraph::Tune(const std::string& parameters, bool disable_future_tuning) { // construct param obj auto hgraph_parameter = std::make_shared(); hgraph_parameter->FromJson(inner_json); + CHECK_ARGUMENT( + data_type_ != DataTypes::DATA_TYPE_BINARY or not hgraph_parameter->mci_parameters.enabled, + "binary HGraph does not support MCI"); auto inner_parameter = std::make_shared(); inner_parameter->FromJson(inner_json); // init new_basic_code obj auto common_param = this->basic_flatten_codes_->ExportCommonParam(); + // Deserialize restores dim_ but leaves the datacell's constructor-time common parameter stale. + common_param.dim_ = this->dim_; auto new_basic_code = FlattenInterface::MakeInstance(hgraph_parameter->base_codes_param, common_param); FlattenInterfacePtr new_precise_code; @@ -304,16 +316,28 @@ HGraph::Tune(const std::string& parameters, bool disable_future_tuning) { if (covers_active_ids(raw_vector_)) { tune_source = raw_vector_; } else if (covers_active_ids(high_precise_codes_) and - high_precise_codes_->GetQuantizerName() == QUANTIZATION_TYPE_VALUE_FP32) { + (high_precise_codes_->GetQuantizerName() == QUANTIZATION_TYPE_VALUE_FP32 or + high_precise_codes_->GetQuantizerName() == QUANTIZATION_TYPE_VALUE_BINARY)) { tune_source = high_precise_codes_; } else if (covers_active_ids(basic_flatten_codes_) and - basic_flatten_codes_->GetQuantizerName() == QUANTIZATION_TYPE_VALUE_FP32) { + (basic_flatten_codes_->GetQuantizerName() == QUANTIZATION_TYPE_VALUE_FP32 or + basic_flatten_codes_->GetQuantizerName() == QUANTIZATION_TYPE_VALUE_BINARY)) { tune_source = basic_flatten_codes_; } else { return false; } } + const auto is_binary_code = [](const FlattenInterfacePtr& code) { + return code != nullptr and code->GetQuantizerName() == QUANTIZATION_TYPE_VALUE_BINARY; + }; + const bool binary_tune_source = is_binary_code(tune_source); + const bool skip_float_train_data = + data_type_ == DataTypes::DATA_TYPE_BINARY and + (not is_tune_base_code or is_binary_code(new_basic_code)) and + (not is_tune_precise_code or is_binary_code(new_precise_code)) and + (not is_tune_raw_code or is_binary_code(new_raw_code)); + auto decode_tune_source = [&](InnerIdType inner_id, float* data) { auto buffer = tune_source->AcquireCodesById(inner_id); if (not buffer) { @@ -324,8 +348,9 @@ HGraph::Tune(const std::string& parameters, bool disable_future_tuning) { }; auto train_count = std::min(this->train_sample_count_, this->GetNumElements()); - Vector train_data(train_count * dim_, 0, allocator_); - if (is_tune_base_code or is_tune_precise_code or is_tune_raw_code) { + Vector train_data(skip_float_train_data ? 0 : train_count * dim_, 0, allocator_); + if (not skip_float_train_data and + (is_tune_base_code or is_tune_precise_code or is_tune_raw_code)) { for (InnerIdType i = 0; i < train_count; i++) { decode_tune_source(i, train_data.data() + i * dim_); } @@ -353,7 +378,21 @@ HGraph::Tune(const std::string& parameters, bool disable_future_tuning) { new_code = wrap_new_code(new_code); new_code->Train(train_data.data(), train_count); - Vector insert_buffer(dim_, 0, allocator_); + const bool packed_binary_rebuild = data_type_ == DataTypes::DATA_TYPE_BINARY and + binary_tune_source and is_binary_code(new_code); + if (packed_binary_rebuild) { + const auto expected_size = (static_cast(dim_) + 7U) / 8U; + CHECK_ARGUMENT(tune_source->GetQuantizerCodeSize() == expected_size and + new_code->GetQuantizerCodeSize() == expected_size, + "binary Tune source and destination packed widths must match dim"); + } + const uint64_t insert_buffer_bytes = packed_binary_rebuild + ? new_code->GetQuantizerCodeSize() + : static_cast(dim_) * sizeof(float); + // Binary decode/encode type-erases packed bytes through float*, so only the byte + // capacity changes; other dtypes retain their original float scratch layout. + Vector insert_buffer( + (insert_buffer_bytes + sizeof(float) - 1) / sizeof(float), 0, allocator_); for (int64_t i = 0; i < current_count; ++i) { decode_tune_source(i, insert_buffer.data()); new_code->InsertVector(static_cast(insert_buffer.data()), i); @@ -971,28 +1010,34 @@ HGraph::UpdateVector(int64_t id, const DatasetPtr& new_base, bool force_update) void* new_base_vec = nullptr; uint64_t data_size = 0; get_vectors(data_type_, dim_, new_base, &new_base_vec, &data_size); + // reject malformed replacements before constructing queries or re-encoding; + // a wrong declared dimension would make binary updates read past the buffer. + CHECK_ARGUMENT( + new_base->GetDim() == dim_, + fmt::format("new_base.dim({}) must be equal to index.dim({})", new_base->GetDim(), dim_)); + CHECK_ARGUMENT(new_base_vec != nullptr, "updated vector must not be null"); if (this->rabitq_fused_datacell_ != nullptr) { - CHECK_ARGUMENT(new_base->GetDim() == dim_, - "updated vector dimension must match the index dimension"); - CHECK_ARGUMENT(new_base_vec != nullptr, "updated vector must not be null"); this->validate_fused_vector_data(static_cast(new_base_vec), 1); } if (not force_update) { std::shared_lock label_lock(this->label_lookup_mutex_); + auto new_base_query = Dataset::Make(); + set_dataset(data_type_, dim_, new_base_query, new_base_vec, 1); float self_dist = 0.0F; if (this->rabitq_fused_datacell_ == nullptr) { // 1. check whether vectors are same Vector base_data(data_size, allocator_); GetVectorByInnerId(inner_id, reinterpret_cast(base_data.data())); - const float old_self_dist = - this->CalcDistanceById(reinterpret_cast(base_data.data()), id); - self_dist = this->CalcDistanceById(static_cast(new_base_vec), id); + auto old_base_query = Dataset::Make(); + set_dataset(data_type_, dim_, old_base_query, base_data.data(), 1); + const float old_self_dist = this->CalcDistanceById(old_base_query, id); + self_dist = this->CalcDistanceById(new_base_query, id); if (std::abs(old_self_dist - self_dist) < 1e-3) { return true; } } else { - self_dist = this->CalcDistanceById(static_cast(new_base_vec), id); + self_dist = this->CalcDistanceById(new_base_query, id); } // 2. check whether the neighborhood relationship is same @@ -1006,9 +1051,8 @@ HGraph::UpdateVector(int64_t id, const DatasetPtr& new_base, bool force_update) float neighbor_dist = 0; try { - neighbor_dist = - this->CalcDistanceById(static_cast(new_base_vec), - this->label_table_->GetLabelById(neighbor_inner_id)); + neighbor_dist = this->CalcDistanceById( + new_base_query, this->label_table_->GetLabelById(neighbor_inner_id)); } catch (const std::runtime_error& e) { // incase that neighbor has been deleted continue; @@ -1047,6 +1091,9 @@ HGraph::UpdateVector(int64_t id, const DatasetPtr& new_base, bool force_update) if (has_precise_reorder()) { update_status = update_status && high_precise_codes_->UpdateVector(new_base_vec, inner_id); } + if (update_status and create_new_raw_vector_ and raw_vector_ != nullptr) { + update_status = raw_vector_->UpdateVector(new_base_vec, inner_id); + } return update_status; } diff --git a/src/algorithm/hgraph/hgraph.h b/src/algorithm/hgraph/hgraph.h index a267452a16..c68011e617 100644 --- a/src/algorithm/hgraph/hgraph.h +++ b/src/algorithm/hgraph/hgraph.h @@ -62,6 +62,9 @@ namespace vsag { +void +ApplyBinaryQuantizationDefaults(const JsonType& user_json, JsonType& inner_json); + class ChunkedManifest; struct ComponentManifestEntry; class FlattenOptimizedBuildInterface; diff --git a/src/algorithm/hgraph/hgraph_build.cpp b/src/algorithm/hgraph/hgraph_build.cpp index ef09799e9a..63a090c100 100644 --- a/src/algorithm/hgraph/hgraph_build.cpp +++ b/src/algorithm/hgraph/hgraph_build.cpp @@ -131,9 +131,10 @@ HGraph::sample_train_dataset(const DatasetPtr& base) const { int64_t dim = base->GetDim(); if (rabitq_fused_datacell_ != nullptr) { return vsag::sample_train_data( - base, total_elements, dim, train_sample_count_, allocator_, 0x52425131U); + base, total_elements, dim, train_sample_count_, data_type_, allocator_, 0x52425131U); } - return vsag::sample_train_data(base, total_elements, dim, train_sample_count_, allocator_); + return vsag::sample_train_data( + base, total_elements, dim, train_sample_count_, data_type_, allocator_, std::nullopt); } void @@ -1419,8 +1420,8 @@ HGraph::InitFeatures() { auto name = this->basic_flatten_codes_->GetQuantizerName(); if (name != QUANTIZATION_TYPE_VALUE_FP32 and name != QUANTIZATION_TYPE_VALUE_BF16 and - name != QUANTIZATION_TYPE_VALUE_FP16 and name != QUANTIZATION_TYPE_VALUE_SQ8_LOCAL and - name != QUANTIZATION_TYPE_VALUE_SQ4_LOCAL) { + name != QUANTIZATION_TYPE_VALUE_FP16 and name != QUANTIZATION_TYPE_VALUE_BINARY and + name != QUANTIZATION_TYPE_VALUE_SQ8_LOCAL and name != QUANTIZATION_TYPE_VALUE_SQ4_LOCAL) { this->index_feature_list_->SetFeature(IndexFeature::NEED_TRAIN); } else { this->index_feature_list_->SetFeatures({ diff --git a/src/algorithm/hgraph/hgraph_param_mapping.cpp b/src/algorithm/hgraph/hgraph_param_mapping.cpp index 01f2fcf9f9..bb0fd1a5c0 100644 --- a/src/algorithm/hgraph/hgraph_param_mapping.cpp +++ b/src/algorithm/hgraph/hgraph_param_mapping.cpp @@ -158,6 +158,20 @@ build_default_hgraph_inner_param(const JsonType& external_json) { } // namespace +void +ApplyBinaryQuantizationDefaults(const JsonType& user_json, JsonType& inner_json) { + if (not user_json.Contains(HGRAPH_BASE_QUANTIZATION_TYPE)) { + inner_json[BASE_CODES_KEY][QUANTIZATION_PARAMS_KEY][TYPE_KEY].SetString( + QUANTIZATION_TYPE_VALUE_BINARY); + } + if (not user_json.Contains(HGRAPH_PRECISE_QUANTIZATION_TYPE)) { + inner_json[PRECISE_CODES_KEY][QUANTIZATION_PARAMS_KEY][TYPE_KEY].SetString( + QUANTIZATION_TYPE_VALUE_BINARY); + } + inner_json[RAW_VECTOR_KEY][QUANTIZATION_PARAMS_KEY][TYPE_KEY].SetString( + QUANTIZATION_TYPE_VALUE_BINARY); +} + JsonType HGraph::map_hgraph_param(const JsonType& hgraph_json) { CHECK_ARGUMENT(not hgraph_json.Contains(HGRAPH_MCI_KEY), @@ -376,6 +390,9 @@ ParamPtr HGraph::CheckAndMappingExternalParam(const JsonType& external_param, const IndexCommonParam& common_param) { auto inner_json = map_hgraph_param(external_param); + if (common_param.data_type_ == DataTypes::DATA_TYPE_BINARY) { + ApplyBinaryQuantizationDefaults(external_param, inner_json); + } if (common_param.data_type_ == DataTypes::DATA_TYPE_SPARSE) { inner_json[BASE_CODES_KEY][CODES_TYPE_KEY].SetString(SPARSE_CODES); inner_json[PRECISE_CODES_KEY][CODES_TYPE_KEY].SetString(SPARSE_CODES); @@ -390,6 +407,9 @@ HGraph::CheckAndMappingExternalParam(const JsonType& external_param, auto hgraph_parameter = std::make_shared(); hgraph_parameter->data_type = common_param.data_type_; hgraph_parameter->FromJson(inner_json); + CHECK_ARGUMENT(common_param.data_type_ != DataTypes::DATA_TYPE_BINARY or + not hgraph_parameter->mci_parameters.enabled, + "binary HGraph does not support MCI"); if (hgraph_parameter->rabitq_fused_datacell) { CHECK_ARGUMENT( // NOLINT(readability-simplify-boolean-expr) common_param.metric_ == MetricType::METRIC_TYPE_L2SQR or diff --git a/src/algorithm/hgraph/hgraph_search.cpp b/src/algorithm/hgraph/hgraph_search.cpp index 349b4351af..df42c2ecd4 100644 --- a/src/algorithm/hgraph/hgraph_search.cpp +++ b/src/algorithm/hgraph/hgraph_search.cpp @@ -719,7 +719,10 @@ HGraph::SearchWithRequest(const SearchRequest& request) const { if (is_range) { // Range search remains single-query only (validate_range_args enforces NumElements==1). if (not use_custom_distance) { + CHECK_ARGUMENT(query != nullptr, "query dataset cannot be null"); this->validate_range_args(query, request.radius_, request.limited_size_); + CHECK_ARGUMENT(get_data(query) != nullptr, + "query vector storage must match index data type"); } else { CHECK_ARGUMENT(query != nullptr, "query dataset cannot be null"); CHECK_ARGUMENT(query->GetNumElements() == 1, diff --git a/src/algorithm/inner_index_interface.cpp b/src/algorithm/inner_index_interface.cpp index 3782f48381..96dd4ce4ab 100644 --- a/src/algorithm/inner_index_interface.cpp +++ b/src/algorithm/inner_index_interface.cpp @@ -525,10 +525,17 @@ InnerIndexInterface::CalcDistancesById(const DatasetPtr& query, CHECK_ARGUMENT(not distance_count_overflows, "CalcDistancesById distance buffer size overflows"); const bool is_sparse = (query->GetSparseVectors() != nullptr); - const bool is_float_query = query->GetFloat32Vectors() != nullptr && !is_sparse && - query->GetMultiVectors() == nullptr && - query->GetInt8Vectors() == nullptr && - query->GetFloat16Vectors() == nullptr; + const bool is_float_query = + query->GetFloat32Vectors() != nullptr && !is_sparse && + query->GetMultiVectors() == nullptr && query->GetInt8Vectors() == nullptr && + query->GetBinaryVectors() == nullptr && query->GetFloat16Vectors() == nullptr; + // A float32-only Dataset selects the native float path below. For an index configured + // for packed binary vectors that would misread the floats as packed bits, so reject + // it the same way the single-ID Dataset path does. + const bool float_query_mismatches_binary = + is_float_query && data_type_ == DataTypes::DATA_TYPE_BINARY; + CHECK_ARGUMENT(not float_query_mismatches_binary, + "CalcDistancesById query vector storage must match index data type"); if (is_float_query) { CHECK_ARGUMENT(query->GetDim() == dim_, "CalcDistancesById query dimension mismatch"); } @@ -600,6 +607,10 @@ InnerIndexInterface::CalcDistancesById(const DatasetPtr& query, if (query->GetInt8Vectors() != nullptr) { sub->Int8Vectors(query->GetInt8Vectors() + q * query->GetDim()); } + if (query->GetBinaryVectors() != nullptr) { + const auto packed_dim = (query->GetDim() + 7) / 8; + sub->BinaryVectors(query->GetBinaryVectors() + q * packed_dim); + } if (query->GetFloat16Vectors() != nullptr) { sub->Float16Vectors(query->GetFloat16Vectors() + q * query->GetDim()); } @@ -788,6 +799,26 @@ InnerIndexInterface::GetVectorByIds(const int64_t* ids, return vectors; } + if (data_type_ == DataTypes::DATA_TYPE_BINARY) { + const auto packed_dim = static_cast((dim_ + 7) / 8); + auto* binary_vectors = + static_cast(allocator->Allocate(static_cast(count) * packed_dim)); + if (binary_vectors == nullptr) { + throw VsagException(ErrorType::NO_ENOUGH_MEMORY, + "failed to allocate memory for vectors"); + } + vectors->NumElements(count) + ->Dim(dim_) + ->BinaryVectors(binary_vectors) + ->Owner(/*auto release=*/not has_specified_allocator, allocator); + for (int64_t i = 0; i < count; ++i) { + InnerIdType inner_id = this->label_table_->GetIdByLabel(ids[i]); + this->GetVectorByInnerId(inner_id, + reinterpret_cast(binary_vectors + i * packed_dim)); + } + return vectors; + } + auto* float_vectors = static_cast(allocator->Allocate(sizeof(float) * count * dim_)); if (float_vectors == nullptr) { throw VsagException(ErrorType::NO_ENOUGH_MEMORY, "failed to allocate memory for vectors"); @@ -796,7 +827,7 @@ InnerIndexInterface::GetVectorByIds(const int64_t* ids, ->Dim(dim_) ->Float32Vectors(float_vectors) ->Owner(/*auto release=*/not has_specified_allocator, allocator); - for (int i = 0; i < count; ++i) { + for (int64_t i = 0; i < count; ++i) { InnerIdType inner_id = this->label_table_->GetIdByLabel(ids[i]); this->GetVectorByInnerId(inner_id, float_vectors + i * dim_); } @@ -879,12 +910,20 @@ InnerIndexInterface::get_data_by_ids_with_flag(const int64_t* ids, const auto item_per_thread = count / thread_count + static_cast(count % thread_count != 0); if ((selected_data_flag & DATA_FLAG_FLOAT32_VECTOR) != 0U) { - auto* fp32_data = reinterpret_cast( - this->allocator_->Allocate(count * this->dim_ * sizeof(float))); - dataset->Float32Vectors(fp32_data); + const bool is_binary = data_type_ == DataTypes::DATA_TYPE_BINARY; + const auto bytes_per_vector = + static_cast(is_binary ? (dim_ + 7) / 8 : dim_ * sizeof(float)); + auto* vector_data = static_cast( + this->allocator_->Allocate(static_cast(count) * bytes_per_vector)); + if (is_binary) { + dataset->BinaryVectors(vector_data); + } else { + dataset->Float32Vectors(reinterpret_cast(vector_data)); + } auto get_vec_func = [&](int64_t begin, int64_t end) { for (int64_t i = begin; i < end; ++i) { - this->GetVectorByInnerId(inner_ids[i], fp32_data + i * this->dim_); + this->GetVectorByInnerId( + inner_ids[i], reinterpret_cast(vector_data + i * bytes_per_vector)); } }; if (this->thread_pool_ != nullptr and thread_count > 1) { @@ -1350,6 +1389,10 @@ InnerIndexInterface::get_data(const DatasetPtr& dataset, int64_t index) const { const auto* ptr = dataset->GetInt8Vectors(); return ptr != nullptr ? ptr + index * dim_ : nullptr; } + if (data_type_ == DataTypes::DATA_TYPE_BINARY) { + const auto* ptr = dataset->GetBinaryVectors(); + return ptr != nullptr ? ptr + index * ((dim_ + 7) / 8) : nullptr; + } if (data_type_ == DataTypes::DATA_TYPE_FP16 || data_type_ == DataTypes::DATA_TYPE_BF16) { const auto* ptr = dataset->GetFloat16Vectors(); return ptr != nullptr ? ptr + index * dim_ : nullptr; diff --git a/src/algorithm/ivf/ivf_build.cpp b/src/algorithm/ivf/ivf_build.cpp index 21ba07b711..11d68ce4f4 100644 --- a/src/algorithm/ivf/ivf_build.cpp +++ b/src/algorithm/ivf/ivf_build.cpp @@ -156,8 +156,8 @@ IVF::Train(const DatasetPtr& data) { int64_t total_elements = data->GetNumElements(); int64_t dim = data->GetDim(); - DatasetPtr train_data = - vsag::sample_train_data(data, total_elements, dim, train_sample_count_, allocator_); + DatasetPtr train_data = vsag::sample_train_data( + data, total_elements, dim, train_sample_count_, DataTypes::DATA_TYPE_FLOAT, allocator_); int64_t sample_count = train_data->GetNumElements(); partition_strategy_->Train(train_data); diff --git a/src/algorithm/ivf/ivf_parameter_test.cpp b/src/algorithm/ivf/ivf_parameter_test.cpp index ea00d91f55..070d5e61f1 100644 --- a/src/algorithm/ivf/ivf_parameter_test.cpp +++ b/src/algorithm/ivf/ivf_parameter_test.cpp @@ -423,8 +423,12 @@ TEST_CASE("SampleTrainingData Function Test", "[ut][sample_train_data]") { ->Owner(false); // Test that small dataset is returned as is - auto result = - vsag::sample_train_data(small_dataset, small_count, small_dim, 10000, allocator.get()); + auto result = vsag::sample_train_data(small_dataset, + small_count, + small_dim, + 10000, + vsag::DataTypes::DATA_TYPE_FLOAT, + allocator.get()); REQUIRE(result == small_dataset); // Test with large dataset that should be sampled @@ -447,8 +451,12 @@ TEST_CASE("SampleTrainingData Function Test", "[ut][sample_train_data]") { ->Owner(false); // Test that large dataset is sampled - result = vsag::sample_train_data( - large_dataset, large_count, large_dim, sample_count, allocator.get()); + result = vsag::sample_train_data(large_dataset, + large_count, + large_dim, + sample_count, + vsag::DataTypes::DATA_TYPE_FLOAT, + allocator.get()); REQUIRE(result != large_dataset); REQUIRE(result->GetNumElements() == sample_count); REQUIRE(result->GetDim() == large_dim); @@ -472,8 +480,12 @@ TEST_CASE("SampleTrainingData Function Test", "[ut][sample_train_data]") { // When train_sample_count is less than min_train_size (512), // the function should use MIN_TRAIN_SIZE as the sample count const int64_t small_sample_count = 100; // Less than min_train_size (512) - result = vsag::sample_train_data( - normal_dataset, normal_count, large_dim, small_sample_count, allocator.get()); + result = vsag::sample_train_data(normal_dataset, + normal_count, + large_dim, + small_sample_count, + vsag::DataTypes::DATA_TYPE_FLOAT, + allocator.get()); REQUIRE(result != normal_dataset); REQUIRE(result->GetNumElements() == 512); // Should use min_train_size REQUIRE(result->GetDim() == large_dim); @@ -497,7 +509,8 @@ TEST_CASE("sample_train_data honors large explicit values", "[ut][sample_train_d ->Float32Vectors(data.data()) ->Owner(false); - auto result = vsag::sample_train_data(dataset, total, dim, requested, allocator.get()); + auto result = vsag::sample_train_data( + dataset, total, dim, requested, vsag::DataTypes::DATA_TYPE_FLOAT, allocator.get()); REQUIRE(result != dataset); REQUIRE(result->GetNumElements() == requested); REQUIRE(result->GetDim() == dim); diff --git a/src/analyzer/hgraph_analyzer.cpp b/src/analyzer/hgraph_analyzer.cpp index 9d8277e519..389cfe2a33 100644 --- a/src/analyzer/hgraph_analyzer.cpp +++ b/src/analyzer/hgraph_analyzer.cpp @@ -287,8 +287,20 @@ HGraphAnalyzer::calculate_quantization_result( } const auto* query = sample_datas.data() + static_cast(i) * dim_; float sample_error = 0.0F; - auto base_result = hgraph_->CalcDistancesById(query, result.data(), actual_topk, false); - auto precise_result = hgraph_->CalcDistancesById(query, result.data(), actual_topk, true); + DatasetPtr base_result = nullptr; + DatasetPtr precise_result = nullptr; + if (hgraph_->data_type_ == DataTypes::DATA_TYPE_BINARY) { + auto native_query = Dataset::Make(); + native_query->Dim(dim_)->NumElements(1)->Owner(false)->BinaryVectors( + reinterpret_cast(query)); + base_result = hgraph_->InnerIndexInterface::CalcDistancesById( + native_query, result.data(), actual_topk, false); + precise_result = hgraph_->InnerIndexInterface::CalcDistancesById( + native_query, result.data(), actual_topk, true); + } else { + base_result = hgraph_->CalcDistancesById(query, result.data(), actual_topk, false); + precise_result = hgraph_->CalcDistancesById(query, result.data(), actual_topk, true); + } if (base_result == nullptr or precise_result == nullptr) { continue; } @@ -340,6 +352,7 @@ HGraphAnalyzer::GetQuantizationInversionRatio(const std::string& search_param) { bool HGraphAnalyzer::SetQuery(const DatasetPtr& query) { + CHECK_ARGUMENT(query->GetDim() == dim_, "analysis query dimension mismatch"); if (query_sample_size_ != 0) { query_sample_ids_.clear(); query_sample_datas_.clear(); @@ -350,7 +363,19 @@ HGraphAnalyzer::SetQuery(const DatasetPtr& query) { query_sample_ids_.resize(query_sample_size_); query_sample_datas_.resize(static_cast::size_type>(query_sample_size_) * static_cast::size_type>(dim_)); - if (const float* query_vectors = query->GetFloat32Vectors()) { + if (hgraph_->data_type_ == DataTypes::DATA_TYPE_BINARY) { + const auto* query_vectors = query->GetBinaryVectors(); + if (query_vectors == nullptr) { + logger::error("Binary analysis query dataset has no packed binary vectors"); + return false; + } + const auto packed_dim = static_cast((dim_ + 7) / 8); + for (uint32_t i = 0; i < query_sample_size_; ++i) { + std::memcpy(query_sample_datas_.data() + static_cast(i) * dim_, + query_vectors + static_cast(i) * packed_dim, + packed_dim); + } + } else if (const float* query_vectors = query->GetFloat32Vectors()) { std::memcpy(query_sample_datas_.data(), query_vectors, static_cast(query_sample_size_ * dim_) * sizeof(float)); @@ -432,6 +457,9 @@ HGraphAnalyzer::calculate_search_result(const Vector& sample_datas, // INT8 Decode writes raw bytes into each sample's float-sized scratch row. std::memcpy(int8_query.data(), sample, dim_ * sizeof(int8_t)); query->Int8Vectors(int8_query.data()); + } else if (hgraph_->data_type_ == DataTypes::DATA_TYPE_BINARY) { + // Binary Decode likewise leaves packed bytes at the start of each scratch row. + query->BinaryVectors(reinterpret_cast(sample)); } else { query->Float32Vectors(sample); } @@ -569,7 +597,8 @@ HGraphAnalyzer::GetStats() { } JsonType HGraphAnalyzer::AnalyzeIndexBySearch(const SearchRequest& request) { - SetQuery(request.query_); + CHECK_ARGUMENT(SetQuery(request.query_), + "analysis query representation must match index dtype"); JsonType stats; stats["avg_distance_query"].SetFloat(GetQueryAvgDistance()); stats["recall_query"].SetFloat(GetQuerySearchRecall(request.params_str_)); diff --git a/src/constants.cpp b/src/constants.cpp index 182accc42e..b1c54e9473 100644 --- a/src/constants.cpp +++ b/src/constants.cpp @@ -37,6 +37,7 @@ const char* const FLOAT32_VECTORS = "f32_vectors"; const char* const FLOAT16_VECTORS = "f16_vectors"; const char* const SPARSE_VECTORS = "sparse_vectors"; const char* const INT8_VECTORS = "i8_vectors"; +const char* const BINARY_VECTORS = "binary_vectors"; const char* const ATTRIBUTE_SETS = "attribute_sets"; const char* const DATASET_PATHS = "paths"; const char* const EXTRA_INFOS = "extra_infos"; @@ -51,10 +52,12 @@ const char* const SIMPLEFLAT_IDS = "simpleflat_ids"; const char* const METRIC_L2 = "l2"; const char* const METRIC_COSINE = "cosine"; const char* const METRIC_IP = "ip"; +const char* const METRIC_HAMMING = "hamming"; const char* const DATATYPE_FLOAT32 = "float32"; const char* const DATATYPE_FLOAT16 = "float16"; const char* const DATATYPE_BFLOAT16 = "bfloat16"; const char* const DATATYPE_INT8 = "int8"; +const char* const DATATYPE_BINARY = "binary"; const char* const DATATYPE_SPARSE = "sparse"; const char* const BLANK_INDEX = "blank_index"; diff --git a/src/data_type.h b/src/data_type.h index 672e86fb00..696d60374d 100644 --- a/src/data_type.h +++ b/src/data_type.h @@ -22,6 +22,7 @@ enum class DataTypes { DATA_TYPE_FP16 = 2, DATA_TYPE_BF16 = 3, DATA_TYPE_SPARSE = 4, + DATA_TYPE_BINARY = 5, }; constexpr const char* @@ -37,6 +38,8 @@ ToString(DataTypes t) noexcept { return DATATYPE_SPARSE; case DataTypes::DATA_TYPE_BF16: return DATATYPE_BFLOAT16; + case DataTypes::DATA_TYPE_BINARY: + return DATATYPE_BINARY; } return "unknown"; } diff --git a/src/datacell/flatten_datacell.h b/src/datacell/flatten_datacell.h index bb1a6cdc5e..081b5d4dd1 100644 --- a/src/datacell/flatten_datacell.h +++ b/src/datacell/flatten_datacell.h @@ -396,6 +396,9 @@ FlattenDataCell::BatchInsertVector(const void* vectors, case DataTypes::DATA_TYPE_SPARSE: source_stride = sizeof(SparseVector); break; + case DataTypes::DATA_TYPE_BINARY: + source_stride = quantizer_->GetCodeSize(); + break; case DataTypes::DATA_TYPE_FLOAT: break; } diff --git a/src/datacell/flatten_datacell_factory.cpp b/src/datacell/flatten_datacell_factory.cpp index 0078463ef5..8279a448c8 100644 --- a/src/datacell/flatten_datacell_factory.cpp +++ b/src/datacell/flatten_datacell_factory.cpp @@ -29,6 +29,10 @@ FlattenInterfacePtr make_flatten_data_cell_cosine(const FlattenInterfaceParamPtr& param, const IndexCommonParam& common_param); +FlattenInterfacePtr +make_flatten_data_cell_hamming(const FlattenInterfaceParamPtr& param, + const IndexCommonParam& common_param); + FlattenInterfacePtr MakeFlattenDataCell(const FlattenInterfaceParamPtr& param, const IndexCommonParam& common_param) { if (common_param.metric_ == MetricType::METRIC_TYPE_L2SQR) { @@ -40,6 +44,9 @@ MakeFlattenDataCell(const FlattenInterfaceParamPtr& param, const IndexCommonPara if (common_param.metric_ == MetricType::METRIC_TYPE_COSINE) { return make_flatten_data_cell_cosine(param, common_param); } + if (common_param.metric_ == MetricType::METRIC_TYPE_HAMMING) { + return make_flatten_data_cell_hamming(param, common_param); + } return nullptr; } diff --git a/src/datacell/flatten_datacell_factory_hamming.cpp b/src/datacell/flatten_datacell_factory_hamming.cpp new file mode 100644 index 0000000000..0a1d1deb2f --- /dev/null +++ b/src/datacell/flatten_datacell_factory_hamming.cpp @@ -0,0 +1,25 @@ +// Copyright 2024-present the vsag project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "flatten_datacell_factory_impl.h" + +namespace vsag { + +FlattenInterfacePtr +make_flatten_data_cell_hamming(const FlattenInterfaceParamPtr& param, + const IndexCommonParam& common_param) { + return MakeFlattenDataCellForMetric(param, common_param); +} + +} // namespace vsag \ No newline at end of file diff --git a/src/datacell/flatten_datacell_factory_impl.h b/src/datacell/flatten_datacell_factory_impl.h index f7ef07041d..bcd7e7d305 100644 --- a/src/datacell/flatten_datacell_factory_impl.h +++ b/src/datacell/flatten_datacell_factory_impl.h @@ -17,6 +17,7 @@ #include "flatten_datacell.h" #include "inner_string_params.h" #include "io/io_headers.h" +#include "quantization/binary_quantizer.h" #include "quantization/int8_quantizer.h" #include "quantization/quantizer_adapter.h" #include "quantization/quantizer_headers.h" @@ -26,6 +27,32 @@ namespace vsag { +template +class BinaryFlattenDataCell final : public FlattenDataCell> { +public: + using Base = FlattenDataCell>; + using Base::Base; + + void + BatchInsertVector(const void* vectors, InnerIdType count, InnerIdType* idx_vec) override { + if (idx_vec != nullptr) { + bool ids_are_contiguous = count > 0; + for (InnerIdType i = 1; i < count and ids_are_contiguous; ++i) { + ids_are_contiguous = idx_vec[i] == idx_vec[0] + i; + } + if (not ids_are_contiguous) { + const auto packed_dim = this->quantizer_->GetCodeSize(); + const auto* data = static_cast(vectors); + for (InnerIdType i = 0; i < count; ++i) { + this->InsertVector(data + static_cast(i) * packed_dim, idx_vec[i]); + } + return; + } + } + Base::BatchInsertVector(vectors, count, idx_vec); + } +}; + template FlattenInterfacePtr MakeFlattenDataCellInstance(const FlattenInterfaceParamPtr& param, @@ -56,6 +83,19 @@ MakeFlattenDataCellInstance(const FlattenInterfaceParamPtr& param, const IndexCommonParam& common_param) { const std::string quantization_string = param->quantizer_parameter->GetTypeName(); + if constexpr (metric == MetricType::METRIC_TYPE_HAMMING) { + if (common_param.data_type_ != DataTypes::DATA_TYPE_BINARY || + quantization_string != QUANTIZATION_TYPE_VALUE_BINARY) { + throw VsagException(ErrorType::INVALID_ARGUMENT, + "Hamming distance requires binary dtype and binary quantization"); + } + if (param->name == FLATTEN_DATA_CELL) { + return std::make_shared>( + param->quantizer_parameter, param->io_parameter, common_param); + } + throw VsagException(ErrorType::INVALID_ARGUMENT, + fmt::format("Unknown flatten interface name: {}", param->name)); + } else { if (common_param.data_type_ == DataTypes::DATA_TYPE_INT8) { if (quantization_string == QUANTIZATION_TYPE_VALUE_INT8) { return MakeFlattenDataCellInstance, IOTmpl>(param, common_param); @@ -164,6 +204,7 @@ MakeFlattenDataCellInstance(const FlattenInterfaceParamPtr& param, throw VsagException(ErrorType::INVALID_ARGUMENT, fmt::format("Unsupported quantization type: {}", actual_quantization)); + } } template diff --git a/src/datacell/flatten_datacell_test.cpp b/src/datacell/flatten_datacell_test.cpp index cf31f70a27..8600df3462 100644 --- a/src/datacell/flatten_datacell_test.cpp +++ b/src/datacell/flatten_datacell_test.cpp @@ -1814,6 +1814,51 @@ TEST_CASE("RaBitQSplitDataCell hybrid IO (1bit in memory, supplement on disk)", } } +TEST_CASE("BinaryFlattenDataCell scattered insert uses restored packed stride", + "[ut][FlattenDataCell][binary][serialize]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + auto param = std::make_shared(); + param->FromJson(JsonType::Parse(R"({ + "io_params": { "type": "memory_io" }, + "quantization_params": { "type": "binary" } + })")); + + IndexCommonParam original_common; + original_common.allocator_ = allocator; + original_common.dim_ = 65; + original_common.metric_ = MetricType::METRIC_TYPE_HAMMING; + original_common.data_type_ = DataTypes::DATA_TYPE_BINARY; + + std::array packed{}; + for (uint64_t i = 0; i < packed.size(); ++i) { + packed[i] = static_cast(i * 11U + 3U); + } + packed[8] &= 0x01U; + packed[17] &= 0x01U; + + auto original = FlattenInterface::MakeInstance(param, original_common); + original->BatchInsertVector(packed.data(), 1); + std::stringstream stream; + IOStreamWriter writer(stream); + original->Serialize(writer); + + auto restored_common = original_common; + restored_common.dim_ = 128; + auto restored = FlattenInterface::MakeInstance(param, restored_common); + stream.seekg(0, std::ios::beg); + IOStreamReader reader(stream); + restored->Deserialize(reader); + REQUIRE(restored->code_size_ == 9); + + std::array ids = {0, 2}; + restored->BatchInsertVector(packed.data(), 2, ids.data()); + std::array actual{}; + REQUIRE(restored->GetCodesById(0, actual.data())); + REQUIRE(std::memcmp(actual.data(), packed.data(), actual.size()) == 0); + REQUIRE(restored->GetCodesById(2, actual.data())); + REQUIRE(std::memcmp(actual.data(), packed.data() + actual.size(), actual.size()) == 0); +} + TEST_CASE("RaBitQSplitDataCell optimized scalar-code build", "[ut][RaBitQSplitDataCell][optimized_build]") { auto allocator = SafeAllocator::FactoryDefaultAllocator(); diff --git a/src/dataset_impl.cpp b/src/dataset_impl.cpp index 72b985f067..29a8e7839d 100644 --- a/src/dataset_impl.cpp +++ b/src/dataset_impl.cpp @@ -507,6 +507,7 @@ DatasetImpl::~DatasetImpl() { // NOLINT ids, [this](const int64_t* values) { allocator_->Deallocate(void_ptr(values)); }); allocator_->Deallocate(void_ptr(DatasetImpl::GetDistances())); allocator_->Deallocate(void_ptr(DatasetImpl::GetInt8Vectors())); + allocator_->Deallocate(void_ptr(DatasetImpl::GetBinaryVectors())); allocator_->Deallocate(void_ptr(DatasetImpl::GetFloat16Vectors())); allocator_->Deallocate(void_ptr(DatasetImpl::GetFloat32Vectors())); allocator_->Deallocate(void_ptr(DatasetImpl::GetExtraInfos())); @@ -546,6 +547,7 @@ DatasetImpl::~DatasetImpl() { // NOLINT release_int64_metadata(ids, [](const int64_t* values) { delete[] values; }); delete[] DatasetImpl::GetDistances(); delete[] DatasetImpl::GetInt8Vectors(); + delete[] DatasetImpl::GetBinaryVectors(); delete[] DatasetImpl::GetFloat16Vectors(); delete[] DatasetImpl::GetFloat32Vectors(); delete[] DatasetImpl::GetExtraInfos(); @@ -621,6 +623,10 @@ DatasetImpl::DeepCopy(Allocator* allocator) const { copy_dataset->Int8Vectors( allocate_and_copy(this->GetInt8Vectors(), num_elements * dim, allocator_ref)); } + if (this->GetBinaryVectors() != nullptr) { + copy_dataset->BinaryVectors(allocate_and_copy( + this->GetBinaryVectors(), num_elements * ((dim + 7) / 8), allocator_ref)); + } if (this->GetFloat32Vectors() != nullptr) { copy_dataset->Float32Vectors( allocate_and_copy(this->GetFloat32Vectors(), num_elements * dim, allocator_ref)); @@ -944,6 +950,7 @@ DatasetImpl::Append(const DatasetPtr& other) { APPEND_DATA(IDS, int64_t*, Ids, 1); APPEND_DATA(DISTS, float*, Distances, dim); APPEND_DATA(INT8_VECTORS, int8_t*, Int8Vectors, dim); + APPEND_DATA(BINARY_VECTORS, uint8_t*, BinaryVectors, (dim + 7) / 8); APPEND_DATA(FLOAT16_VECTORS, uint16_t*, Float16Vectors, dim); APPEND_DATA(FLOAT32_VECTORS, float*, Float32Vectors, dim); APPEND_DATA(VECTOR_COUNTS, uint32_t*, VectorCounts, 1); diff --git a/src/dataset_impl.h b/src/dataset_impl.h index 93425fc419..15177e24e6 100644 --- a/src/dataset_impl.h +++ b/src/dataset_impl.h @@ -40,6 +40,7 @@ class DatasetImpl : public Dataset, public SearchMetricsDatasetAccessor { const float*, const char*, const int8_t*, + const uint8_t*, const uint16_t*, const int64_t*, const std::string*, @@ -164,6 +165,20 @@ class DatasetImpl : public Dataset, public SearchMetricsDatasetAccessor { return nullptr; } + DatasetPtr + BinaryVectors(const uint8_t* vectors) override { + this->data_[BINARY_VECTORS] = vectors; + return shared_from_this(); + } + + const uint8_t* + GetBinaryVectors() const override { + if (auto iter = this->data_.find(BINARY_VECTORS); iter != this->data_.end()) { + return std::get(iter->second); + } + return nullptr; + } + DatasetPtr Float32Vectors(const float* vectors) override { this->data_[FLOAT32_VECTORS] = vectors; diff --git a/src/factory/index_registry.cpp b/src/factory/index_registry.cpp index 4bb14472ce..7da088d90b 100644 --- a/src/factory/index_registry.cpp +++ b/src/factory/index_registry.cpp @@ -21,6 +21,7 @@ #include #include "common.h" +#include "vsag/constants.h" namespace vsag { namespace { @@ -76,6 +77,10 @@ create_registered_index(const std::string& index_name, "failed to create index(internal): null creator for ", normalized_name); } + const bool supports_binary = + normalized_name == INDEX_BRUTE_FORCE || normalized_name == INDEX_HGRAPH; + CHECK_ARGUMENT(index_common_params.data_type_ != DataTypes::DATA_TYPE_BINARY || supports_binary, + "binary dtype is only supported by brute_force and hgraph indexes"); return creator(parsed_params, index_common_params); } diff --git a/src/index_common_param.cpp b/src/index_common_param.cpp index 91128ddc8e..e8a33d039f 100644 --- a/src/index_common_param.cpp +++ b/src/index_common_param.cpp @@ -37,17 +37,20 @@ fill_datatype(IndexCommonParam& result, const JsonType& datatype_obj) { result.data_type_ = DataTypes::DATA_TYPE_BF16; } else if (datatype == DATATYPE_INT8) { result.data_type_ = DataTypes::DATA_TYPE_INT8; + } else if (datatype == DATATYPE_BINARY) { + result.data_type_ = DataTypes::DATA_TYPE_BINARY; } else if (datatype == DATATYPE_SPARSE) { result.data_type_ = DataTypes::DATA_TYPE_SPARSE; } else { throw VsagException( ErrorType::INVALID_ARGUMENT, - fmt::format("parameters[{}] must be one of [{}, {}, {}, {}, {}], now is {}", + fmt::format("parameters[{}] must be one of [{}, {}, {}, {}, {}, {}], now is {}", PARAMETER_DTYPE, DATATYPE_FLOAT32, DATATYPE_FLOAT16, DATATYPE_BFLOAT16, DATATYPE_INT8, + DATATYPE_BINARY, DATATYPE_SPARSE, datatype)); } @@ -86,13 +89,16 @@ fill_metrictype(IndexCommonParam& result, const JsonType& metric_obj) { result.metric_ = MetricType::METRIC_TYPE_IP; } else if (metric == METRIC_COSINE) { result.metric_ = MetricType::METRIC_TYPE_COSINE; + } else if (metric == METRIC_HAMMING) { + result.metric_ = MetricType::METRIC_TYPE_HAMMING; } else { throw VsagException(ErrorType::INVALID_ARGUMENT, - fmt::format("parameters[{}] must in [{}, {}, {}], now is {}", + fmt::format("parameters[{}] must in [{}, {}, {}, {}], now is {}", PARAMETER_METRIC_TYPE, METRIC_L2, METRIC_IP, METRIC_COSINE, + METRIC_HAMMING, metric)); } } @@ -149,6 +155,10 @@ IndexCommonParam::CheckAndCreate(JsonType& params, const std::shared_ptr + +#include "simd/bit_simd.h" + +namespace vsag { + +BinaryQuantizer::BinaryQuantizer(int dim, Allocator* allocator) + : Quantizer(dim, allocator) { + this->code_size_ = (static_cast(dim) + 7U) / 8U; + this->query_code_size_ = this->code_size_; + this->metric_ = MetricType::METRIC_TYPE_HAMMING; +} + +BinaryQuantizer::BinaryQuantizer(const BinaryQuantizerParamPtr& param, + const IndexCommonParam& common_param) + : BinaryQuantizer(common_param.dim_, common_param.allocator_.get()) { +} + +BinaryQuantizer::BinaryQuantizer(const QuantizerParamPtr& param, + const IndexCommonParam& common_param) + : BinaryQuantizer(common_param.dim_, common_param.allocator_.get()) { +} + +void +BinaryQuantizer::MaskUnusedBits(uint8_t* codes) const { + if (codes == nullptr or this->dim_ == 0) { + return; + } + const uint64_t used_bits = this->dim_ % 8U; + if (used_bits == 0) { + return; + } + codes[this->code_size_ - 1U] &= static_cast((1U << used_bits) - 1U); +} + +bool +BinaryQuantizer::TrainImpl(const float* data, uint64_t count) { + this->is_trained_ = true; + return true; +} + +bool +BinaryQuantizer::EncodeOneImpl(const float* data, uint8_t* codes) { + std::memcpy(codes, static_cast(data), this->code_size_); + this->MaskUnusedBits(codes); + return true; +} + +bool +BinaryQuantizer::EncodeBatchImpl(const float* data, uint8_t* codes, uint64_t count) { + const auto* source = reinterpret_cast(data); + for (uint64_t i = 0; i < count; ++i) { + std::memcpy(codes + i * this->code_size_, source + i * this->code_size_, this->code_size_); + this->MaskUnusedBits(codes + i * this->code_size_); + } + return true; +} + +bool +BinaryQuantizer::DecodeOneImpl(const uint8_t* codes, float* data) { + std::memcpy(static_cast(data), codes, this->code_size_); + return true; +} + +bool +BinaryQuantizer::DecodeBatchImpl(const uint8_t* codes, float* data, uint64_t count) { + auto* output = reinterpret_cast(data); + for (uint64_t i = 0; i < count; ++i) { + std::memcpy(output + i * this->code_size_, codes + i * this->code_size_, this->code_size_); + } + return true; +} + +float +BinaryQuantizer::ComputeImpl(const uint8_t* codes1, const uint8_t* codes2) { + return static_cast(HammingDistance(codes1, codes2, this->code_size_)); +} + +void +BinaryQuantizer::ProcessQueryImpl(const float* query, Computer& computer) const { + try { + if (computer.buf_ == nullptr and this->query_code_size_ > 0) { + computer.buf_ = + static_cast(this->allocator_->Allocate(this->query_code_size_)); + } + } catch (const std::bad_alloc& e) { + computer.buf_ = nullptr; + throw VsagException(ErrorType::NO_ENOUGH_MEMORY, "bad alloc when init computer buf"); + } + if (this->query_code_size_ == 0) { + return; + } + std::memcpy(computer.buf_, static_cast(query), this->query_code_size_); + this->MaskUnusedBits(computer.buf_); +} + +void +BinaryQuantizer::ComputeDistImpl(Computer& computer, + const uint8_t* codes, + float* dists) const { + *dists = static_cast(HammingDistance(codes, computer.buf_, this->code_size_)); +} + +void +BinaryQuantizer::ScanBatchDistImpl(Computer& computer, + uint64_t count, + const uint8_t* codes, + float* dists) const { + for (uint64_t i = 0; i < count; ++i) { + this->ComputeDistImpl(computer, codes + i * this->code_size_, dists + i); + } +} + +void +BinaryQuantizer::ComputeDistsBatch4Impl(Computer& computer, + const uint8_t* codes1, + const uint8_t* codes2, + const uint8_t* codes3, + const uint8_t* codes4, + float& dists1, + float& dists2, + float& dists3, + float& dists4) const { + dists1 = static_cast(HammingDistance(codes1, computer.buf_, this->code_size_)); + dists2 = static_cast(HammingDistance(codes2, computer.buf_, this->code_size_)); + dists3 = static_cast(HammingDistance(codes3, computer.buf_, this->code_size_)); + dists4 = static_cast(HammingDistance(codes4, computer.buf_, this->code_size_)); +} + +void +BinaryQuantizer::ReleaseComputerImpl(Computer& computer) const { + this->allocator_->Deallocate(computer.buf_); +} + +} // namespace vsag \ No newline at end of file diff --git a/src/quantization/binary_quantizer.h b/src/quantization/binary_quantizer.h new file mode 100644 index 0000000000..4900d08f4f --- /dev/null +++ b/src/quantization/binary_quantizer.h @@ -0,0 +1,105 @@ +// Copyright 2024-present the vsag project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include "binary_quantizer_parameter.h" +#include "index_common_param.h" +#include "inner_string_params.h" +#include "quantizer.h" + +namespace vsag { + +/*** + * @brief BinaryQuantizer stores packed binary vectors without converting them + * into floating-point vectors. + * + * A logical dimension of `dim` bits is stored in `ceil(dim / 8)` contiguous + * bytes. Bits are packed least-significant-bit first inside each byte, and the + * unused high bits of the final byte are cleared. Hamming distance is computed + * as `popcount(a XOR b)` over exactly those bytes. + */ +class BinaryQuantizer : public Quantizer { +public: + explicit BinaryQuantizer(int dim, Allocator* allocator); + + BinaryQuantizer(const BinaryQuantizerParamPtr& param, const IndexCommonParam& common_param); + + BinaryQuantizer(const QuantizerParamPtr& param, const IndexCommonParam& common_param); + + ~BinaryQuantizer() override = default; + + bool + TrainImpl(const float* data, uint64_t count); + + bool + EncodeOneImpl(const float* data, uint8_t* codes); + + bool + EncodeBatchImpl(const float* data, uint8_t* codes, uint64_t count); + + bool + DecodeOneImpl(const uint8_t* codes, float* data); + + bool + DecodeBatchImpl(const uint8_t* codes, float* data, uint64_t count); + + float + ComputeImpl(const uint8_t* codes1, const uint8_t* codes2); + + void + SerializeImpl(StreamWriter& writer) {}; + + void + DeserializeImpl(StreamReader& reader) { + this->query_code_size_ = this->code_size_; + }; + + void + ProcessQueryImpl(const float* query, Computer& computer) const; + + void + ComputeDistImpl(Computer& computer, const uint8_t* codes, float* dists) const; + + void + ScanBatchDistImpl(Computer& computer, + uint64_t count, + const uint8_t* codes, + float* dists) const; + + void + ComputeDistsBatch4Impl(Computer& computer, + const uint8_t* codes1, + const uint8_t* codes2, + const uint8_t* codes3, + const uint8_t* codes4, + float& dists1, + float& dists2, + float& dists3, + float& dists4) const; + + void + ReleaseComputerImpl(Computer& computer) const; + + [[nodiscard]] std::string + NameImpl() const { + return QUANTIZATION_TYPE_VALUE_BINARY; + } + +private: + void + MaskUnusedBits(uint8_t* codes) const; +}; + +} // namespace vsag diff --git a/src/quantization/binary_quantizer_parameter.cpp b/src/quantization/binary_quantizer_parameter.cpp new file mode 100644 index 0000000000..29f41be9a2 --- /dev/null +++ b/src/quantization/binary_quantizer_parameter.cpp @@ -0,0 +1,34 @@ +// Copyright 2024-present the vsag project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "binary_quantizer_parameter.h" + +#include "inner_string_params.h" + +namespace vsag { +BinaryQuantizerParameter::BinaryQuantizerParameter() + : QuantizerParameter(QUANTIZATION_TYPE_VALUE_BINARY) { +} + +void +BinaryQuantizerParameter::FromJson(const JsonType& json) { +} + +JsonType +BinaryQuantizerParameter::ToJson() const { + JsonType json; + json[TYPE_KEY].SetString(QUANTIZATION_TYPE_VALUE_BINARY); + return json; +} +} // namespace vsag \ No newline at end of file diff --git a/src/quantization/binary_quantizer_parameter.h b/src/quantization/binary_quantizer_parameter.h new file mode 100644 index 0000000000..3db0d151ee --- /dev/null +++ b/src/quantization/binary_quantizer_parameter.h @@ -0,0 +1,34 @@ +// Copyright 2024-present the vsag project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include "quantizer_parameter.h" + +namespace vsag { +DEFINE_POINTER2(BinaryQuantizerParam, BinaryQuantizerParameter); + +class BinaryQuantizerParameter : public QuantizerParameter { +public: + BinaryQuantizerParameter(); + + ~BinaryQuantizerParameter() override = default; + + void + FromJson(const JsonType& json) override; + + JsonType + ToJson() const override; +}; +} // namespace vsag \ No newline at end of file diff --git a/src/quantization/binary_quantizer_test.cpp b/src/quantization/binary_quantizer_test.cpp new file mode 100644 index 0000000000..259ec0c9c5 --- /dev/null +++ b/src/quantization/binary_quantizer_test.cpp @@ -0,0 +1,283 @@ +// Copyright 2024-present the vsag project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "binary_quantizer.h" + +#include +#include +#include +#include +#include + +#include "binary_quantizer_parameter.h" +#include "impl/allocator/safe_allocator.h" +#include "index_common_param.h" +#include "inner_string_params.h" +#include "quantizer_test.h" +#include "unittest.h" + +using namespace vsag; + +namespace { + +uint64_t +ReferenceHamming(const uint8_t* x, const uint8_t* y, uint64_t num_bytes) { + uint64_t result = 0; + for (uint64_t i = 0; i < num_bytes; ++i) { + result += static_cast(__builtin_popcount(static_cast(x[i] ^ y[i]))); + } + return result; +} + +// Generate packed binary vectors of `bits` logical bits each. Unless keep_dirty_tail is set, +// the unused high bits of the final byte are cleared so that ground-truth and stored codes agree. +std::vector +GeneratePackedBinary(uint64_t count, uint64_t bits, uint32_t seed, bool keep_dirty_tail = false) { + const uint64_t num_bytes = (bits + 7) / 8; + std::vector data(count * num_bytes); + std::mt19937 rng(seed); + for (auto& byte : data) { + byte = static_cast(rng()); + } + if (not keep_dirty_tail and bits % 8 != 0) { + const auto mask = static_cast((1U << (bits % 8U)) - 1U); + for (uint64_t i = 0; i < count; ++i) { + data[(i + 1) * num_bytes - 1] &= mask; + } + } + return data; +} + +uint64_t +TailMaskedByte(uint8_t byte, uint64_t bits, uint64_t byte_index, uint64_t num_bytes) { + if (bits % 8 != 0 and byte_index + 1 == num_bytes) { + return byte & static_cast((1U << (bits % 8U)) - 1U); + } + return byte; +} + +uint64_t +ReferenceHammingMasked(const uint8_t* x, const uint8_t* y, uint64_t num_bytes, uint64_t bits) { + uint64_t result = 0; + for (uint64_t i = 0; i < num_bytes; ++i) { + const auto diff = static_cast(TailMaskedByte(x[i], bits, i, num_bytes) ^ + TailMaskedByte(y[i], bits, i, num_bytes)); + result += static_cast(__builtin_popcount(diff)); + } + return result; +} + +} // namespace + +TEST_CASE("BinaryQuantizer Code Size and Parameter", "[ut][binary_quantizer]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + for (uint64_t bits : {1, 7, 8, 9, 63, 64, 65, 127, 128, 256}) { + BinaryQuantizer quantizer(static_cast(bits), allocator.get()); + REQUIRE(quantizer.GetCodeSize() == (bits + 7) / 8); + REQUIRE(quantizer.GetQueryCodeSize() == (bits + 7) / 8); + REQUIRE(quantizer.GetDim() == static_cast(bits)); + REQUIRE(quantizer.Metric() == MetricType::METRIC_TYPE_HAMMING); + REQUIRE(quantizer.Name() == QUANTIZATION_TYPE_VALUE_BINARY); + } + + // Parameter type round-trips through the quantizer parameter factory. + auto param = std::make_shared(); + auto json = param->ToJson(); + REQUIRE(json[TYPE_KEY].GetString() == QUANTIZATION_TYPE_VALUE_BINARY); + auto parsed = QuantizerParameter::GetQuantizerParameterByJson(json); + REQUIRE(parsed != nullptr); + REQUIRE(parsed->GetTypeName() == QUANTIZATION_TYPE_VALUE_BINARY); + REQUIRE(QuantizerParameter::IsValidQuantizationType(QUANTIZATION_TYPE_VALUE_BINARY)); + + // Constructing via IndexCommonParam matches the direct-dimension constructor. + IndexCommonParam common_param; + common_param.dim_ = 100; + common_param.allocator_ = allocator; + BinaryQuantizer from_common(param, common_param); + REQUIRE(from_common.GetCodeSize() == 13); +} + +TEST_CASE("BinaryQuantizer Encode Decode Compute", "[ut][binary_quantizer]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + const uint64_t count = 17; + for (uint64_t bits : {9, 64, 127, 256}) { + const uint64_t num_bytes = (bits + 7) / 8; + auto data = GeneratePackedBinary(count, bits, 47); + BinaryQuantizer quantizer(static_cast(bits), allocator.get()); + REQUIRE(quantizer.Train(nullptr, count)); + + std::vector codes(num_bytes * count, 0); + std::vector single(num_bytes, 0); + for (uint64_t i = 0; i < count; ++i) { + const auto* source = data.data() + i * num_bytes; + REQUIRE(quantizer.EncodeOne(reinterpret_cast(source), single.data())); + REQUIRE(std::memcmp(single.data(), source, num_bytes) == 0); + REQUIRE(quantizer.Compute(single.data(), codes.data() + i * num_bytes) == + static_cast(ReferenceHammingMasked( + source, codes.data() + i * num_bytes, num_bytes, bits))); + } + + REQUIRE(quantizer.EncodeBatch( + reinterpret_cast(data.data()), codes.data(), count)); + for (uint64_t i = 0; i < count * num_bytes; ++i) { + REQUIRE(codes[i] == data[i]); + } + + // Self distance is zero; pairwise distance is the exact differing-bit count. + for (uint64_t i = 0; i < count; ++i) { + REQUIRE(quantizer.Compute(codes.data() + i * num_bytes, codes.data() + i * num_bytes) == + 0.0F); + for (uint64_t j = 0; j < count; ++j) { + auto expected = static_cast(ReferenceHamming( + codes.data() + i * num_bytes, codes.data() + j * num_bytes, num_bytes)); + REQUIRE(quantizer.Compute(codes.data() + i * num_bytes, + codes.data() + j * num_bytes) == expected); + } + } + + // Decode preserves packed bytes so decoded vectors can be encoded again unchanged. + std::vector decoded(count * num_bytes, 0xFFU); + REQUIRE(quantizer.DecodeBatch( + codes.data(), reinterpret_cast(decoded.data()), count)); + REQUIRE(decoded == codes); + std::vector reencoded(count * num_bytes, 0); + REQUIRE(quantizer.EncodeBatch( + reinterpret_cast(decoded.data()), reencoded.data(), count)); + REQUIRE(reencoded == codes); + } +} + +TEST_CASE("BinaryQuantizer Ignores Unused Tail Bits", "[ut][binary_quantizer]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + constexpr uint64_t bits = 9; // two bytes, last byte uses 1 bit + const uint64_t num_bytes = 2; + BinaryQuantizer quantizer(static_cast(bits), allocator.get()); + REQUIRE(quantizer.Train(nullptr, 1)); + + const std::vector clean = {0x02U, 0x00U}; + const std::vector dirty = {0x02U, 0xFCU}; // differs only in unused bits + + std::vector codes1(num_bytes, 0); + std::vector codes2(num_bytes, 0); + quantizer.EncodeOne(reinterpret_cast(clean.data()), codes1.data()); + quantizer.EncodeOne(reinterpret_cast(dirty.data()), codes2.data()); + + // Stored codes have their unused tail bits cleared. + REQUIRE(codes2[1] == 0x00U); + REQUIRE(std::memcmp(codes1.data(), codes2.data(), num_bytes) == 0); + REQUIRE(quantizer.Compute(codes1.data(), codes2.data()) == 0.0F); + + // A dirty-tail query is masked during preprocessing as well. + auto computer = quantizer.FactoryComputer(); + computer->SetQuery(reinterpret_cast(dirty.data())); + float dist = -1.0F; + computer->ComputeDist(codes1.data(), &dist); + REQUIRE(dist == 0.0F); +} + +TEST_CASE("BinaryQuantizer Computer Distances", "[ut][binary_quantizer]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + const uint64_t count = 20; + for (uint64_t bits : {128, 257}) { + const uint64_t num_bytes = (bits + 7) / 8; + auto data = GeneratePackedBinary(count, bits, 1234); + auto query = GeneratePackedBinary(1, bits, 99); + BinaryQuantizer quantizer(static_cast(bits), allocator.get()); + REQUIRE(quantizer.Train(nullptr, count)); + std::vector codes(num_bytes * count, 0); + REQUIRE(quantizer.EncodeBatch( + reinterpret_cast(data.data()), codes.data(), count)); + + auto computer = quantizer.FactoryComputer(); + computer->SetQuery(reinterpret_cast(query.data())); + + std::vector dists(count, -1.0F); + computer->ScanBatchDists(count, codes.data(), dists.data()); + for (uint64_t i = 0; i < count; ++i) { + auto expected = static_cast( + ReferenceHamming(query.data(), codes.data() + i * num_bytes, num_bytes)); + REQUIRE(dists[i] == expected); + float single = -1.0F; + computer->ComputeDist(codes.data() + i * num_bytes, &single); + REQUIRE(single == expected); + } + + float d1 = -1.0F; + float d2 = -1.0F; + float d3 = -1.0F; + float d4 = -1.0F; + if (count >= 4) { + computer->ComputeDistsBatch4(codes.data(), + codes.data() + num_bytes, + codes.data() + 2 * num_bytes, + codes.data() + 3 * num_bytes, + d1, + d2, + d3, + d4); + for (uint64_t i = 0; i < 4; ++i) { + auto expected = static_cast( + ReferenceHamming(query.data(), codes.data() + i * num_bytes, num_bytes)); + const float actual = (i == 0) ? d1 : (i == 1) ? d2 : (i == 2) ? d3 : d4; + REQUIRE(actual == expected); + } + } + } +} + +TEST_CASE("BinaryQuantizer Serialize and Deserialize", "[ut][binary_quantizer]") { + auto allocator = SafeAllocator::FactoryDefaultAllocator(); + const uint64_t bits = 100; + const uint64_t num_bytes = (bits + 7) / 8; + const uint64_t count = 12; + auto data = GeneratePackedBinary(count, bits, 2026); + + BinaryQuantizer quantizer1(static_cast(bits), allocator.get()); + BinaryQuantizer quantizer2(64, allocator.get()); + REQUIRE(quantizer1.Train(nullptr, count)); + std::vector codes(num_bytes * count, 0); + REQUIRE( + quantizer1.EncodeBatch(reinterpret_cast(data.data()), codes.data(), count)); + test_serializion(quantizer1, quantizer2); + + REQUIRE(quantizer2.GetCodeSize() == quantizer1.GetCodeSize()); + REQUIRE(quantizer2.GetQueryCodeSize() == quantizer1.GetQueryCodeSize()); + REQUIRE(quantizer2.GetDim() == quantizer1.GetDim()); + REQUIRE(quantizer2.Metric() == MetricType::METRIC_TYPE_HAMMING); + + auto computer = quantizer2.FactoryComputer(); + computer->SetQuery(reinterpret_cast(data.data())); + float query_distance = -1.0F; + computer->ComputeDist(codes.data(), &query_distance); + REQUIRE(query_distance == 0.0F); + + std::vector restored_codes(num_bytes * count, 0); + REQUIRE(quantizer2.EncodeBatch( + reinterpret_cast(data.data()), restored_codes.data(), count)); + REQUIRE(std::memcmp(codes.data(), restored_codes.data(), codes.size()) == 0); + for (uint64_t i = 0; i < count; ++i) { + for (uint64_t j = 0; j < count; ++j) { + REQUIRE(quantizer2.Compute(codes.data() + i * num_bytes, + restored_codes.data() + j * num_bytes) == + static_cast(ReferenceHamming(codes.data() + i * num_bytes, + restored_codes.data() + j * num_bytes, + num_bytes))); + } + } + + BinaryQuantizer smaller_quantizer(64, allocator.get()); + test_serializion(smaller_quantizer, quantizer2); + REQUIRE(quantizer2.GetDim() == smaller_quantizer.GetDim()); + REQUIRE(quantizer2.GetQueryCodeSize() == smaller_quantizer.GetCodeSize()); +} diff --git a/src/quantization/fp32_quantizer.h b/src/quantization/fp32_quantizer.h index 131718c6c2..d1b135495a 100644 --- a/src/quantization/fp32_quantizer.h +++ b/src/quantization/fp32_quantizer.h @@ -67,7 +67,9 @@ class FP32Quantizer : public Quantizer> { SerializeImpl(StreamWriter& writer) {}; void - DeserializeImpl(StreamReader& reader) {}; + DeserializeImpl(StreamReader& reader) { + this->query_code_size_ = this->code_size_; + }; void ProcessQueryImpl(const float* query, Computer>& computer) const; diff --git a/src/quantization/fp32_quantizer_test.cpp b/src/quantization/fp32_quantizer_test.cpp index 2eb5d679a1..0d5af0d12f 100644 --- a/src/quantization/fp32_quantizer_test.cpp +++ b/src/quantization/fp32_quantizer_test.cpp @@ -104,9 +104,10 @@ void TestSerializeAndDeserializeMetricFP32(uint64_t dim, int count, float error = 1e-5) { auto allocator = SafeAllocator::FactoryDefaultAllocator(); FP32Quantizer quantizer1(dim, allocator.get()); - FP32Quantizer quantizer2(dim, allocator.get()); + FP32Quantizer quantizer2(dim / 2, allocator.get()); TestSerializeAndDeserialize, metric>( quantizer1, quantizer2, dim, count, error); + REQUIRE(quantizer2.GetQueryCodeSize() == quantizer1.GetQueryCodeSize()); } TEST_CASE("FP32 Serialize and Deserialize", "[ut][FP32Quantizer]") { diff --git a/src/quantization/int8_quantizer.h b/src/quantization/int8_quantizer.h index e40818d16a..5189987cfa 100644 --- a/src/quantization/int8_quantizer.h +++ b/src/quantization/int8_quantizer.h @@ -71,7 +71,9 @@ class INT8Quantizer : public Quantizer> { SerializeImpl(StreamWriter& writer) {}; void - DeserializeImpl(StreamReader& reader) {}; + DeserializeImpl(StreamReader& reader) { + this->query_code_size_ = this->code_size_; + }; void ProcessQueryImpl(const float* query, Computer>& computer) const; diff --git a/src/quantization/int8_quantizer_test.cpp b/src/quantization/int8_quantizer_test.cpp index 422f214312..1327d77e8c 100644 --- a/src/quantization/int8_quantizer_test.cpp +++ b/src/quantization/int8_quantizer_test.cpp @@ -264,6 +264,7 @@ TestSerializeAndDeserializeINT8(Quantizer>& quant1, test_serializion(quant1, quant2); REQUIRE(quant1.GetCodeSize() == quant2.GetCodeSize()); + REQUIRE(quant1.GetQueryCodeSize() == quant2.GetQueryCodeSize()); REQUIRE(quant1.GetDim() == quant2.GetDim()); TestQuantizerEncodeDecodeINT8(quant2, dim, count, error, false); @@ -283,7 +284,7 @@ void TestSerializeAndDeserializeMetricINT8(uint64_t dim, int count, float error = 1e-5) { auto allocator = SafeAllocator::FactoryDefaultAllocator(); INT8Quantizer quantizer1(dim, allocator.get()); - INT8Quantizer quantizer2(dim, allocator.get()); + INT8Quantizer quantizer2(dim / 2, allocator.get()); TestSerializeAndDeserializeINT8(quantizer1, quantizer2, dim, count, error); } diff --git a/src/quantization/quantizer_parameter.cpp b/src/quantization/quantizer_parameter.cpp index c0f14c40cf..062847f5dc 100644 --- a/src/quantization/quantizer_parameter.cpp +++ b/src/quantization/quantizer_parameter.cpp @@ -19,6 +19,7 @@ #include +#include "binary_quantizer_parameter.h" #include "fp32_quantizer_parameter.h" #include "inner_string_params.h" #include "product_quantization/pq_fastscan_quantizer_parameter.h" @@ -90,6 +91,9 @@ QuantizerParameter::GetQuantizerParameterByJson(const JsonType& json) { } else if (type_name == QUANTIZATION_TYPE_VALUE_INT8) { quantizer_param = std::make_shared(); quantizer_param->FromJson(json); + } else if (type_name == QUANTIZATION_TYPE_VALUE_BINARY) { + quantizer_param = std::make_shared(); + quantizer_param->FromJson(json); } else { throw VsagException(ErrorType::INVALID_ARGUMENT, fmt::format("invalid quantizer name {}", type_name)); @@ -114,7 +118,8 @@ QuantizerParameter::IsValidQuantizationType(const std::string& type_name) { QUANTIZATION_TYPE_VALUE_SPARSE, QUANTIZATION_TYPE_VALUE_PQFS, QUANTIZATION_TYPE_VALUE_TQ, - QUANTIZATION_TYPE_VALUE_INT8}; + QUANTIZATION_TYPE_VALUE_INT8, + QUANTIZATION_TYPE_VALUE_BINARY}; return valid_types.find(type_name) != valid_types.end(); } diff --git a/src/simd/bit_simd.h b/src/simd/bit_simd.h index 2362b48436..d7931e0e1c 100644 --- a/src/simd/bit_simd.h +++ b/src/simd/bit_simd.h @@ -50,4 +50,15 @@ extern BitOperatorType BitXor; using BitNotType = void (*)(const uint8_t* x, const uint64_t num_byte, uint8_t* result); extern BitNotType BitNot; + +/** Return the number of differing bits in two equally sized packed arrays. */ +inline uint64_t +HammingDistance(const uint8_t* x, const uint8_t* y, uint64_t num_byte) { + uint64_t distance = 0; + for (uint64_t i = 0; i < num_byte; ++i) { + const auto diff = static_cast(x[i] ^ y[i]); + distance += static_cast(__builtin_popcount(diff)); + } + return distance; +} } // namespace vsag diff --git a/src/utils/util_functions.cpp b/src/utils/util_functions.cpp index af5cdaaab3..95900c61fe 100644 --- a/src/utils/util_functions.cpp +++ b/src/utils/util_functions.cpp @@ -204,6 +204,9 @@ get_vectors(DataTypes type, } else if (type == DataTypes::DATA_TYPE_FP16 || type == DataTypes::DATA_TYPE_BF16) { *vectors_ptr = (void*)base->GetFloat16Vectors(); *data_size_ptr = dim * sizeof(uint16_t); + } else if (type == DataTypes::DATA_TYPE_BINARY) { + *vectors_ptr = (void*)base->GetBinaryVectors(); + *data_size_ptr = (dim + 7) / 8; } else { throw VsagException(ErrorType::INVALID_ARGUMENT, fmt::format("no support for this type: {}", (int)type)); @@ -225,6 +228,11 @@ set_dataset(DataTypes type, ->Dim(dim) ->Owner(false) ->NumElements(num_element); + } else if (type == DataTypes::DATA_TYPE_BINARY) { + base->BinaryVectors((uint8_t*)vectors_ptr) + ->Dim(dim) + ->Owner(false) + ->NumElements(num_element); } else { throw VsagException(ErrorType::INVALID_ARGUMENT, fmt::format("no support for this type: {}", (int)type)); @@ -236,6 +244,7 @@ sample_train_data(const vsag::DatasetPtr& data, int64_t total_elements, int64_t dim, int64_t train_sample_count, + DataTypes data_type, Allocator* allocator, std::optional random_seed) { const int64_t min_train_size = 512; @@ -257,7 +266,7 @@ sample_train_data(const vsag::DatasetPtr& data, safe_allocator = owned_allocator.get(); } - vsag::Vector sampled_data_buffer(safe_allocator); + vsag::Vector sampled_data_buffer(safe_allocator); vsag::Vector sampled_ids(safe_allocator); vsag::Vector sampled_indices(safe_allocator); sampled_indices.reserve(sample_count); @@ -274,22 +283,31 @@ sample_train_data(const vsag::DatasetPtr& data, sampled_indices[j] = i; } } - sampled_data_buffer.resize(sample_count * dim); - const auto* original_data = data->GetFloat32Vectors(); + const bool is_binary = data_type == DataTypes::DATA_TYPE_BINARY; + const auto bytes_per_vector = + static_cast(is_binary ? (dim + 7) / 8 : dim * sizeof(float)); + sampled_data_buffer.resize(static_cast(sample_count) * bytes_per_vector); + const auto* original_data = is_binary + ? data->GetBinaryVectors() + : reinterpret_cast(data->GetFloat32Vectors()); for (int64_t i = 0; i < sample_count; ++i) { - std::copy(original_data + sampled_indices[i] * dim, - original_data + (sampled_indices[i] + 1) * dim, - sampled_data_buffer.data() + i * dim); + const auto source_offset = static_cast(sampled_indices[i]) * bytes_per_vector; + const auto destination_offset = static_cast(i) * bytes_per_vector; + std::copy(original_data + source_offset, + original_data + source_offset + bytes_per_vector, + sampled_data_buffer.data() + destination_offset); } - auto* new_data_buffer = - static_cast(safe_allocator->Allocate(sample_count * dim * sizeof(float))); + auto* new_data_buffer = static_cast( + safe_allocator->Allocate(static_cast(sample_count) * bytes_per_vector)); std::copy(sampled_data_buffer.begin(), sampled_data_buffer.end(), new_data_buffer); auto sampled_dataset = std::make_shared(); - sampled_dataset->NumElements(sample_count) - ->Dim(dim) - ->Float32Vectors(new_data_buffer) - ->Owner(true, safe_allocator); + sampled_dataset->NumElements(sample_count)->Dim(dim)->Owner(true, safe_allocator); + if (is_binary) { + sampled_dataset->BinaryVectors(new_data_buffer); + } else { + sampled_dataset->Float32Vectors(reinterpret_cast(new_data_buffer)); + } if (data->GetIds() != nullptr) { sampled_ids.reserve(sample_count); const auto* original_ids = data->GetIds(); diff --git a/src/utils/util_functions.h b/src/utils/util_functions.h index 285d4fc1e0..944c3b850e 100644 --- a/src/utils/util_functions.h +++ b/src/utils/util_functions.h @@ -129,6 +129,7 @@ sample_train_data(const DatasetPtr& data, int64_t total_elements, int64_t dim, int64_t train_sample_count, + DataTypes data_type, Allocator* allocator = nullptr, std::optional random_seed = std::nullopt); diff --git a/src/utils/util_functions_test.cpp b/src/utils/util_functions_test.cpp index e5d9fc0d06..b233d5ad43 100644 --- a/src/utils/util_functions_test.cpp +++ b/src/utils/util_functions_test.cpp @@ -119,8 +119,10 @@ TEST_CASE("sample_train_data honors a fixed random seed", "[ut][UtilFunctions]") ->Owner(false); auto allocator = std::make_shared(); - auto first = sample_train_data(data, count, dim, sample_count, allocator.get(), 0x52425131U); - auto second = sample_train_data(data, count, dim, sample_count, allocator.get(), 0x52425131U); + auto first = sample_train_data( + data, count, dim, sample_count, DataTypes::DATA_TYPE_FLOAT, allocator.get(), 0x52425131U); + auto second = sample_train_data( + data, count, dim, sample_count, DataTypes::DATA_TYPE_FLOAT, allocator.get(), 0x52425131U); REQUIRE(first->GetNumElements() == sample_count); REQUIRE(second->GetNumElements() == sample_count); @@ -140,7 +142,8 @@ TEST_CASE("sample_train_data honors configured counts above the default", "[ut][ data->NumElements(count)->Dim(dim)->Float32Vectors(vectors.data())->Owner(false); auto allocator = std::make_shared(); - auto sampled = sample_train_data(data, count, dim, sample_count, allocator.get(), 0x52425131U); + auto sampled = sample_train_data( + data, count, dim, sample_count, DataTypes::DATA_TYPE_FLOAT, allocator.get(), 0x52425131U); REQUIRE(sampled->GetNumElements() == sample_count); } @@ -153,8 +156,12 @@ TEST_CASE("sample_train_data uses all vectors for an unbounded count", "[ut][Uti data->NumElements(count)->Dim(dim)->Float32Vectors(vectors.data())->Owner(false); auto allocator = std::make_shared(); - auto sampled = - sample_train_data(data, count, dim, std::numeric_limits::max(), allocator.get()); + auto sampled = sample_train_data(data, + count, + dim, + std::numeric_limits::max(), + DataTypes::DATA_TYPE_FLOAT, + allocator.get()); REQUIRE(sampled == data); REQUIRE(sampled->GetNumElements() == count); diff --git a/tests/python/test_binary_hamming.py b/tests/python/test_binary_hamming.py new file mode 100644 index 0000000000..8214b6998a --- /dev/null +++ b/tests/python/test_binary_hamming.py @@ -0,0 +1,278 @@ +# Copyright 2024-present the vsag project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Pytest tests for bit-packed binary vectors and Hamming distance.""" + +import json + +import numpy as np +import pyvsag +import pytest + +POPCOUNT = np.array([bin(i).count("1") for i in range(256)], dtype=np.uint16) + + +def hamming_matrix(queries: np.ndarray, base: np.ndarray, dim: int) -> np.ndarray: + """Exact Hamming distance with unused tail bits masked out.""" + bytes_per_vector = (dim + 7) // 8 + q = queries.reshape(-1, bytes_per_vector) + b = base.reshape(-1, bytes_per_vector) + xor = np.bitwise_xor(q[:, None, :], b[None, :, :]) + dists = POPCOUNT[xor].sum(axis=2, dtype=np.uint32) + used_tail_bits = dim % 8 + if used_tail_bits: + mask = (1 << used_tail_bits) - 1 + tail = POPCOUNT[(q[:, None, -1] ^ b[None, :, -1]) & (~mask & 0xFF)] + dists = dists - tail + return dists + + +def index_params(index_type: str, dim: int) -> str: + return json.dumps({"dtype": "binary", "metric_type": "hamming", "dim": dim}) + + +def make_binary_data(num_elements: int, dim: int, seed: int) -> np.ndarray: + rng = np.random.default_rng(seed) + bytes_per_vector = (dim + 7) // 8 + return rng.integers(0, 256, size=(num_elements, bytes_per_vector), dtype=np.uint8) + + +def test_bruteforce_binary_hamming_valid_input(): + dim = 137 # not a multiple of eight + num_elements = 400 + base = make_binary_data(num_elements, dim, 11) + ids = np.arange(num_elements, dtype=np.int64) + index = pyvsag.Index("brute_force", index_params("brute_force", dim)) + index.build(vectors=base, ids=ids, num_elements=num_elements, dim=dim) + + queries = make_binary_data(8, dim, 12) + ground = hamming_matrix(queries, base, dim) + for query, row in zip(queries, ground): + result_ids, result_dists = index.knn_search( + vector=query, k=10, parameters=json.dumps({}) + ) + expected = np.sort(row)[:10] + np.testing.assert_array_equal(result_dists, expected.astype(np.float32)) + for result_id, result_dist in zip(result_ids, result_dists): + assert row[result_id] == int(result_dist) + + threshold = float(np.sort(row)[49]) + range_ids, range_dists = index.range_search( + point=query, threshold=threshold, parameters=json.dumps({}) + ) + expected_range = np.count_nonzero(row <= threshold) + assert len(range_ids) == expected_range + assert np.all(range_dists <= threshold) + + by_id = index.calc_distances_by_id( + query=query, ids=np.array([0, 7, num_elements - 1], dtype=np.int64) + ) + np.testing.assert_array_equal( + by_id, row[[0, 7, num_elements - 1]].astype(np.float32) + ) + + +def test_bruteforce_binary_flat_buffer_and_add(): + dim = 64 + num_elements = 200 + bytes_per_vector = dim // 8 + base = make_binary_data(num_elements, dim, 21) + ids = np.arange(num_elements, dtype=np.int64) + index = pyvsag.Index("brute_force", index_params("brute_force", dim)) + index.build( + vectors=base.reshape(-1), ids=ids, num_elements=num_elements, dim=dim + ) + + extra = make_binary_data(20, dim, 22) + extra_ids = np.arange(num_elements, num_elements + 20, dtype=np.int64) + index.add( + vectors=extra.reshape(-1), ids=extra_ids, num_elements=20, dim=dim + ) + assert index.get_num_elements() == num_elements + 20 + + query = make_binary_data(1, dim, 23)[0] + all_base = np.concatenate([base, extra], axis=0) + row = hamming_matrix(query.reshape(1, -1), all_base, dim)[0] + result_ids, result_dists = index.knn_search( + vector=query, k=5, parameters=json.dumps({}) + ) + np.testing.assert_array_equal(result_dists, np.sort(row)[:5].astype(np.float32)) + assert bytes_per_vector == base.shape[1] + + +def test_binary_unused_tail_bits_are_ignored(): + dim = 9 # two bytes, only the low bit of the second byte is meaningful + rng = np.random.default_rng(31) + base = rng.integers(0, 256, size=(16, 2), dtype=np.uint8) + ids = np.arange(base.shape[0], dtype=np.int64) + index = pyvsag.Index("brute_force", index_params("brute_force", dim)) + index.build(vectors=base, ids=ids, num_elements=base.shape[0], dim=dim) + + query = base[3].copy() + query[1] ^= 0xFE # only unused tail bits differ + result_ids, result_dists = index.knn_search( + vector=query, k=1, parameters=json.dumps({}) + ) + assert result_ids[0] == 3 + assert result_dists[0] == 0.0 + + +@pytest.mark.parametrize("dim", [1, 7, 8, 9, 63, 64, 65, 127, 128]) +def test_binary_non_aligned_dimensions(dim): + num_elements = 32 + base = make_binary_data(num_elements, dim, dim + 100) + ids = np.arange(num_elements, dtype=np.int64) + index = pyvsag.Index("brute_force", index_params("brute_force", dim)) + index.build(vectors=base, ids=ids, num_elements=num_elements, dim=dim) + + query = make_binary_data(1, dim, dim + 200)[0] + row = hamming_matrix(query.reshape(1, -1), base, dim)[0] + result_ids, result_dists = index.knn_search( + vector=query, k=5, parameters=json.dumps({}) + ) + np.testing.assert_array_equal(result_dists, np.sort(row)[:5].astype(np.float32)) + + +def test_binary_rejects_invalid_dtype(): + dim = 64 + index = pyvsag.Index("brute_force", index_params("brute_force", dim)) + vectors = np.zeros((10, dim // 8), dtype=np.float32) + with pytest.raises(ValueError): + index.build( + vectors=vectors, + ids=np.arange(10, dtype=np.int64), + num_elements=10, + dim=dim, + ) + + +def test_binary_rejects_invalid_shape_and_query_length(): + dim = 65 # ceil(65 / 8) == 9 bytes + index = pyvsag.Index("brute_force", index_params("brute_force", dim)) + vectors = np.zeros((10, 8), dtype=np.uint8) # one byte short + with pytest.raises(ValueError): + index.build( + vectors=vectors, + ids=np.arange(10, dtype=np.int64), + num_elements=10, + dim=dim, + ) + + valid = np.zeros((10, 9), dtype=np.uint8) + index.build( + vectors=valid, + ids=np.arange(10, dtype=np.int64), + num_elements=10, + dim=dim, + ) + with pytest.raises(ValueError): + index.knn_search(vector=np.zeros(8, dtype=np.uint8), k=1, parameters="{}") + with pytest.raises(ValueError): + index.range_search( + point=np.zeros(10, dtype=np.uint8), threshold=1.0, parameters="{}" + ) + + +def test_binary_rejects_non_contiguous_arrays(): + dim = 64 + index = pyvsag.Index("brute_force", index_params("brute_force", dim)) + contiguous = np.zeros((10, dim // 8), dtype=np.uint8) + non_contiguous = np.asfortranarray(contiguous) + assert not non_contiguous.flags["C_CONTIGUOUS"] + with pytest.raises(ValueError): + index.build( + vectors=non_contiguous, + ids=np.arange(10, dtype=np.int64), + num_elements=10, + dim=dim, + ) + + +def test_binary_hgraph_load_refreshes_dimension(tmp_path): + source_dim = 10 + num_elements = 32 + base = make_binary_data(num_elements, source_dim, 51) + ids = np.arange(num_elements, dtype=np.int64) + source = pyvsag.Index("hgraph", index_params("hgraph", source_dim)) + source.build(vectors=base, ids=ids, num_elements=num_elements, dim=source_dim) + + filename = str(tmp_path / "binary_hgraph.bin") + source.save(filename) + restored = pyvsag.Index("hgraph", index_params("hgraph", source_dim - 1)) + restored.load(filename) + + query = base[0] + search_params = json.dumps({"hgraph": {"ef_search": 64}}) + result_ids, result_dists = restored.knn_search( + vector=query, k=1, parameters=search_params + ) + assert result_ids[0] == 0 + assert result_dists[0] == 0.0 + + range_ids, range_dists = restored.range_search( + point=query, threshold=0.0, parameters=search_params + ) + assert 0 in range_ids + assert range_dists[np.flatnonzero(range_ids == 0)[0]] == 0.0 + np.testing.assert_array_equal( + restored.calc_distances_by_id(query=query, ids=np.array([0], dtype=np.int64)), + np.array([0.0], dtype=np.float32), + ) + + extra = make_binary_data(1, source_dim, 52) + restored.add( + vectors=extra, + ids=np.array([num_elements], dtype=np.int64), + num_elements=1, + dim=source_dim, + ) + assert restored.get_num_elements() == num_elements + 1 + with pytest.raises(ValueError, match="dim"): + restored.add( + vectors=extra, + ids=np.array([num_elements + 1], dtype=np.int64), + num_elements=1, + dim=source_dim - 1, + ) + + +def test_binary_hgraph_recall_against_bruteforce(): + dim = 128 + num_elements = 600 + ids = np.arange(num_elements, dtype=np.int64) + base = make_binary_data(num_elements, dim, 41) + queries = make_binary_data(20, dim, 42) + + exact = pyvsag.Index("brute_force", index_params("brute_force", dim)) + exact.build(vectors=base, ids=ids, num_elements=num_elements, dim=dim) + + graph = pyvsag.Index("hgraph", index_params("hgraph", dim)) + graph.build(vectors=base, ids=ids, num_elements=num_elements, dim=dim) + + search_params = json.dumps({"hgraph": {"ef_search": 200}}) + recall_hits = 0 + for query in queries: + exact_ids, _ = exact.knn_search( + vector=query, k=10, parameters=json.dumps({}) + ) + graph_ids, graph_dists = graph.knn_search( + vector=query, k=10, parameters=search_params + ) + recall_hits += len(set(exact_ids.tolist()) & set(graph_ids.tolist())) + row = hamming_matrix(query.reshape(1, -1), base, dim)[0] + for result_id, result_dist in zip(graph_ids, graph_dists): + assert row[result_id] == int(result_dist) + + recall = recall_hits / (len(queries) * 10) + assert recall >= 0.8 diff --git a/tests/test_binary_hamming_hgraph.cpp b/tests/test_binary_hamming_hgraph.cpp new file mode 100644 index 0000000000..997ffaa923 --- /dev/null +++ b/tests/test_binary_hamming_hgraph.cpp @@ -0,0 +1,878 @@ +// Copyright 2024-present the vsag project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "vsag/vsag.h" + +namespace { + +uint64_t +ReferenceHamming(const uint8_t* x, const uint8_t* y, uint64_t num_bytes) { + uint64_t result = 0; + for (uint64_t i = 0; i < num_bytes; ++i) { + auto diff = static_cast(x[i] ^ y[i]); + while (diff != 0) { + diff = static_cast(diff & (diff - 1)); + ++result; + } + } + return result; +} + +std::vector +GeneratePacked(uint64_t count, uint64_t bits, uint64_t seed) { + const uint64_t bytes = (bits + 7) / 8; + std::vector data(count * bytes); + uint64_t state = seed; + for (auto& byte : data) { + state = state * 6364136223846793005ULL + 1442695040888963407ULL; + byte = static_cast(state >> 33); + } + if (bits % 8 != 0) { + const auto mask = static_cast((1U << (bits % 8U)) - 1U); + for (uint64_t i = 0; i < count; ++i) { + data[(i + 1) * bytes - 1] &= mask; + } + } + return data; +} + +class LargeAllocationRecorder : public vsag::Allocator { +public: + explicit LargeAllocationRecorder(uint64_t threshold) : threshold_(threshold) { + } + + std::string + Name() override { + return "LargeAllocationRecorder"; + } + + void* + Allocate(uint64_t size) override { + Record(size); + return std::malloc(size); + } + + void + Deallocate(void* ptr) override { + std::free(ptr); + } + + void* + Reallocate(void* ptr, uint64_t size) override { + Record(size); + return std::realloc(ptr, size); + } + + void + StartRecording() { + oversized_allocations_.store(0, std::memory_order_relaxed); + recording_.store(true, std::memory_order_release); + } + + void + StopRecording() { + recording_.store(false, std::memory_order_release); + } + + [[nodiscard]] uint64_t + OversizedAllocations() const { + return oversized_allocations_.load(std::memory_order_acquire); + } + +private: + void + Record(uint64_t size) { + if (recording_.load(std::memory_order_acquire) and size >= threshold_) { + oversized_allocations_.fetch_add(1, std::memory_order_relaxed); + } + } + + uint64_t threshold_; + std::atomic recording_{false}; + std::atomic oversized_allocations_{0}; +}; + +vsag::DatasetPtr +MakeDataset(const std::vector& data, const std::vector& ids, int64_t dim) { + return vsag::Dataset::Make() + ->NumElements(static_cast(ids.size())) + ->Dim(dim) + ->Ids(ids.data()) + ->BinaryVectors(data.data()) + ->Owner(false); +} + +} // namespace + +TEST_CASE("(PR) HGraph Binary Hamming Build Search Add Remove Serialize", + "[ft][build][hgraph][pr]") { + constexpr int64_t dim = 129; // non-byte-aligned logical bit dimension + constexpr int64_t count = 600; + constexpr uint64_t num_bytes = (dim + 7) / 8; + + const std::string params = R"( + { + "dtype": "binary", + "metric_type": "hamming", + "dim": 129, + "index_param": { + "base_quantization_type": "binary", + "support_remove": true + } + })"; + const std::string brute_params = R"( + { + "dtype": "binary", + "metric_type": "hamming", + "dim": 129 + })"; + const std::string search_params = R"({"hgraph":{"ef_search":200}})"; + + auto base_data = GeneratePacked(count, dim, 20260102); + std::vector ids(count); + std::iota(ids.begin(), ids.end(), 0); + + auto build_dataset = MakeDataset(base_data, ids, dim); + + auto brute_made = vsag::Factory::CreateIndex("brute_force", brute_params); + REQUIRE(brute_made.has_value()); + auto brute = brute_made.value(); + REQUIRE(brute->Build(build_dataset).has_value()); + + const std::vector retrieved_ids = {0, 7, count - 1}; + auto retrieved = + brute->GetDataByIds(retrieved_ids.data(), static_cast(retrieved_ids.size())); + REQUIRE(retrieved.has_value()); + REQUIRE(retrieved.value()->GetFloat32Vectors() == nullptr); + REQUIRE(retrieved.value()->GetBinaryVectors() != nullptr); + for (uint64_t i = 0; i < retrieved_ids.size(); ++i) { + const auto* expected = base_data.data() + retrieved_ids[i] * num_bytes; + const auto* actual = retrieved.value()->GetBinaryVectors() + i * num_bytes; + REQUIRE(std::memcmp(expected, actual, num_bytes) == 0); + } + auto raw_retrieved = + brute->GetRawVectorByIds(retrieved_ids.data(), static_cast(retrieved_ids.size())); + REQUIRE(raw_retrieved.has_value()); + REQUIRE(raw_retrieved.value()->GetFloat32Vectors() == nullptr); + REQUIRE(raw_retrieved.value()->GetBinaryVectors() != nullptr); + REQUIRE(std::memcmp(base_data.data() + 7 * num_bytes, + raw_retrieved.value()->GetBinaryVectors() + num_bytes, + num_bytes) == 0); + + // Auxiliary float storage must not divert a binary query to the float-only overload. + std::vector auxiliary_float(dim, 0.0F); + auto mixed_query = vsag::Dataset::Make() + ->NumElements(1) + ->Dim(dim) + ->BinaryVectors(base_data.data()) + ->Float32Vectors(auxiliary_float.data()) + ->Owner(false); + auto mixed_distance = brute->CalcDistanceById(mixed_query, 7); + REQUIRE(mixed_distance.has_value()); + const auto mixed_expected = + ReferenceHamming(base_data.data(), base_data.data() + 7 * num_bytes, num_bytes); + REQUIRE(mixed_distance.value() == static_cast(mixed_expected)); + + auto binary_query = vsag::Dataset::Make() + ->NumElements(1) + ->Dim(dim) + ->BinaryVectors(base_data.data()) + ->Owner(false); + auto batch_distances = brute->CalcDistancesById( + binary_query, retrieved_ids.data(), static_cast(retrieved_ids.size())); + REQUIRE(batch_distances.has_value()); + for (uint64_t i = 0; i < retrieved_ids.size(); ++i) { + const auto expected = ReferenceHamming( + base_data.data(), base_data.data() + retrieved_ids[i] * num_bytes, num_bytes); + REQUIRE(batch_distances.value()->GetDistances()[i] == static_cast(expected)); + } + + auto graph_made = vsag::Factory::CreateIndex("hgraph", params); + REQUIRE(graph_made.has_value()); + auto graph = graph_made.value(); + REQUIRE_FALSE(graph->CheckFeature(vsag::IndexFeature::NEED_TRAIN)); + REQUIRE(graph->CheckFeature(vsag::IndexFeature::SUPPORT_RANGE_SEARCH)); + REQUIRE(graph->CheckFeature(vsag::IndexFeature::SUPPORT_RANGE_SEARCH_WITH_ID_FILTER)); + REQUIRE(graph->Build(build_dataset).has_value()); + REQUIRE(graph->GetNumElements() == count); + auto graph_raw = + graph->GetRawVectorByIds(retrieved_ids.data(), static_cast(retrieved_ids.size())); + REQUIRE(graph_raw.has_value()); + REQUIRE(graph_raw.value()->GetFloat32Vectors() == nullptr); + REQUIRE(graph_raw.value()->GetBinaryVectors() != nullptr); + REQUIRE(std::memcmp(base_data.data() + 7 * num_bytes, + graph_raw.value()->GetBinaryVectors() + num_bytes, + num_bytes) == 0); + + auto queries = GeneratePacked(16, dim, 777); + uint64_t recall_hits = 0; + uint64_t expected_hits = 0; + for (int64_t q = 0; q < 16; ++q) { + const auto* query_vec = queries.data() + static_cast(q) * num_bytes; + auto query = + vsag::Dataset::Make()->NumElements(1)->Dim(dim)->BinaryVectors(query_vec)->Owner(false); + + auto exact = brute->KnnSearch(query, 10, "{}"); + auto approx = graph->KnnSearch(query, 10, search_params); + REQUIRE(exact.has_value()); + REQUIRE(approx.has_value()); + std::unordered_set exact_ids; + for (int64_t i = 0; i < exact.value()->GetDim(); ++i) { + exact_ids.insert(exact.value()->GetIds()[i]); + } + for (int64_t i = 0; i < approx.value()->GetDim(); ++i) { + const auto id = approx.value()->GetIds()[i]; + if (exact_ids.count(id) != 0) { + ++recall_hits; + } + ++expected_hits; + auto expected = ReferenceHamming( + query_vec, base_data.data() + static_cast(id) * num_bytes, num_bytes); + REQUIRE(approx.value()->GetDistances()[i] == static_cast(expected)); + } + + // Range search uses "distance <= threshold" for Hamming. + auto range = graph->RangeSearch(query, 40.0F, search_params); + REQUIRE(range.has_value()); + for (int64_t i = 0; i < range.value()->GetDim(); ++i) { + const auto id = range.value()->GetIds()[i]; + auto expected = ReferenceHamming( + query_vec, base_data.data() + static_cast(id) * num_bytes, num_bytes); + REQUIRE(expected <= 40); + REQUIRE(range.value()->GetDistances()[i] == static_cast(expected)); + } + + auto dist = graph->CalcDistanceById(query, q * 31); + REQUIRE(dist.has_value()); + auto expected = ReferenceHamming( + query_vec, base_data.data() + static_cast(q * 31) * num_bytes, num_bytes); + REQUIRE(dist.value() == static_cast(expected)); + } + const auto recall = static_cast(recall_hits) / static_cast(expected_hits); + REQUIRE(recall >= 0.75); + + // Forced updates must retain packed extraction and encoding. + auto replacement = GeneratePacked(1, dim, 999); + const std::vector replacement_id = {4}; + auto replacement_dataset = MakeDataset(replacement, replacement_id, dim); + auto updated = graph->UpdateVector(replacement_id[0], replacement_dataset, true); + REQUIRE(updated.has_value()); + REQUIRE(updated.value()); + auto replacement_query = vsag::Dataset::Make() + ->NumElements(1) + ->Dim(dim) + ->BinaryVectors(replacement.data()) + ->Owner(false); + auto updated_distance = graph->CalcDistanceById(replacement_query, replacement_id[0]); + REQUIRE(updated_distance.has_value()); + REQUIRE(updated_distance.value() == 0.0F); + + // Incremental add. + constexpr int64_t extra = 40; + auto extra_data = GeneratePacked(extra, dim, 888); + std::vector extra_ids(extra); + std::iota(extra_ids.begin(), extra_ids.end(), count); + auto add_dataset = MakeDataset(extra_data, extra_ids, dim); + REQUIRE(graph->Add(add_dataset).has_value()); + REQUIRE(graph->GetNumElements() == count + extra); + + // Remove some labels and make sure they no longer surface in search results. + std::vector removed_ids = {3, 77, count + 11}; + auto removed = graph->Remove(removed_ids); + REQUIRE(removed.has_value()); + REQUIRE(removed.value() == removed_ids.size()); + + auto remove_query = vsag::Dataset::Make() + ->NumElements(1) + ->Dim(dim) + ->BinaryVectors(queries.data()) + ->Owner(false); + auto after_remove = graph->KnnSearch(remove_query, 20, search_params); + REQUIRE(after_remove.has_value()); + for (int64_t i = 0; i < after_remove.value()->GetDim(); ++i) { + const auto id = after_remove.value()->GetIds()[i]; + REQUIRE(std::find(removed_ids.begin(), removed_ids.end(), id) == removed_ids.end()); + } + + // Serialize and deserialize restore the same exact binary distances. + std::stringstream stream; + REQUIRE(graph->SerializeStreaming(stream).has_value()); + auto restored_made = vsag::Factory::CreateIndex("hgraph", params); + REQUIRE(restored_made.has_value()); + auto restored = restored_made.value(); + std::stringstream restore_stream(stream.str()); + REQUIRE(restored->DeserializeStreaming(restore_stream).has_value()); + REQUIRE(restored->GetNumElements() == graph->GetNumElements()); + + auto original_result = graph->KnnSearch(remove_query, 10, search_params); + auto restored_result = restored->KnnSearch(remove_query, 10, search_params); + REQUIRE(original_result.has_value()); + REQUIRE(restored_result.has_value()); + REQUIRE(restored_result.value()->GetDim() == original_result.value()->GetDim()); + for (int64_t i = 0; i < restored_result.value()->GetDim(); ++i) { + REQUIRE(restored_result.value()->GetIds()[i] == original_result.value()->GetIds()[i]); + REQUIRE(restored_result.value()->GetDistances()[i] == + original_result.value()->GetDistances()[i]); + } +} + +TEST_CASE("(PR) HGraph Binary Hamming Ordinary Update And Tune", "[ft][hgraph][pr]") { + constexpr int64_t dim = 64; + constexpr int64_t count = 64; + constexpr uint64_t num_bytes = (dim + 7) / 8; + const std::string params = R"( + { + "dtype": "binary", + "metric_type": "hamming", + "dim": 64, + "index_param": {} + })"; + auto base_data = GeneratePacked(count, dim, 2468); + std::vector ids(count); + std::iota(ids.begin(), ids.end(), 0); + auto index_made = vsag::Factory::CreateIndex("hgraph", params); + REQUIRE(index_made.has_value()); + auto index = index_made.value(); + REQUIRE(index->Build(MakeDataset(base_data, ids, dim)).has_value()); + + const std::vector first_id = {0}; + std::vector unchanged(base_data.begin(), base_data.begin() + num_bytes); + auto unchanged_result = index->UpdateVector(first_id[0], MakeDataset(unchanged, first_id, dim)); + REQUIRE(unchanged_result.has_value()); + REQUIRE(unchanged_result.value()); + + const std::vector second_id = {1}; + std::vector changed(base_data.begin() + num_bytes, base_data.begin() + 2 * num_bytes); + changed[0] ^= 0x01U; + auto changed_result = index->UpdateVector(second_id[0], MakeDataset(changed, second_id, dim)); + REQUIRE(changed_result.has_value()); + + auto raw = index->GetRawVectorByIds(first_id.data(), 1); + REQUIRE(raw.has_value()); + REQUIRE(raw.value()->GetBinaryVectors() != nullptr); + REQUIRE(std::memcmp(base_data.data(), raw.value()->GetBinaryVectors(), num_bytes) == 0); + + auto tuned = index->Tune(R"({"index_param":{"use_reorder":true}})"); + REQUIRE(tuned.has_value()); + REQUIRE(tuned.value()); + auto query = vsag::Dataset::Make() + ->NumElements(1) + ->Dim(dim) + ->BinaryVectors(base_data.data()) + ->Owner(false); + auto distance = index->CalcDistanceById(query, first_id[0]); + REQUIRE(distance.has_value()); + REQUIRE(distance.value() == 0.0F); +} + +TEST_CASE("(PR) HGraph Binary Tune Uses Restored Dimension", "[ft][hgraph][pr]") { + constexpr int64_t dim = 65; + constexpr int64_t count = 32; + constexpr uint64_t num_bytes = (dim + 7) / 8; + const std::string params = R"( + { + "dtype": "binary", + "metric_type": "hamming", + "dim": 65, + "index_param": {} + })"; + const std::string different_dim_params = R"( + { + "dtype": "binary", + "metric_type": "hamming", + "dim": 128, + "index_param": {} + })"; + auto base_data = GeneratePacked(count, dim, 97531); + std::vector ids(count); + std::iota(ids.begin(), ids.end(), 0); + auto index_made = vsag::Factory::CreateIndex("hgraph", params); + REQUIRE(index_made.has_value()); + auto index = index_made.value(); + REQUIRE(index->Build(MakeDataset(base_data, ids, dim)).has_value()); + auto serialized = index->Serialize(); + REQUIRE(serialized.has_value()); + + auto restored_made = vsag::Factory::CreateIndex("hgraph", different_dim_params); + REQUIRE(restored_made.has_value()); + auto restored = restored_made.value(); + REQUIRE(restored->Deserialize(serialized.value()).has_value()); + auto tuned = restored->Tune(R"({"index_param":{"use_reorder":true}})"); + REQUIRE(tuned.has_value()); + REQUIRE(tuned.value()); + + constexpr int64_t query_id = 7; + auto query = vsag::Dataset::Make() + ->NumElements(1) + ->Dim(dim) + ->BinaryVectors(base_data.data() + query_id * num_bytes) + ->Owner(false); + auto precise_distance = restored->CalcDistanceById(query, ids[query_id], true); + REQUIRE(precise_distance.has_value()); + REQUIRE(precise_distance.value() == 0.0F); +} + +TEST_CASE("(PR) HGraph Binary Tune Skips Float Training Scratch", "[ft][hgraph][pr]") { + constexpr int64_t dim = 8192; + constexpr int64_t count = 64; + constexpr uint64_t float_train_scratch_bytes = + static_cast(count) * dim * sizeof(float); + LargeAllocationRecorder allocator(float_train_scratch_bytes); + const std::string params = R"( + { + "dtype": "binary", + "metric_type": "hamming", + "dim": 8192, + "index_param": {} + })"; + auto base_data = GeneratePacked(count, dim, 531); + std::vector ids(count); + std::iota(ids.begin(), ids.end(), 0); + auto index_made = vsag::Factory::CreateIndex("hgraph", params, &allocator); + REQUIRE(index_made.has_value()); + auto index = index_made.value(); + REQUIRE(index->Build(MakeDataset(base_data, ids, dim)).has_value()); + + allocator.StartRecording(); + auto tuned = index->Tune(R"({"index_param":{"use_reorder":true}})"); + allocator.StopRecording(); + REQUIRE(tuned.has_value()); + REQUIRE(tuned.value()); + REQUIRE(allocator.OversizedAllocations() == 0); + + auto query = vsag::Dataset::Make() + ->NumElements(1) + ->Dim(dim) + ->BinaryVectors(base_data.data()) + ->Owner(false); + auto distance = index->CalcDistanceById(query, ids.front()); + REQUIRE(distance.has_value()); + REQUIRE(distance.value() == 0.0F); +} + +TEST_CASE("(PR) HGraph Binary Raw Store Follows Forced Updates", "[ft][hgraph][pr]") { + constexpr int64_t dim = 8; + constexpr int64_t count = 32; + constexpr uint64_t num_bytes = (dim + 7) / 8; + const std::string params = R"( + { + "dtype": "binary", + "metric_type": "hamming", + "dim": 8, + "index_param": { + "store_raw_vector": true + } + })"; + auto base_data = GeneratePacked(count, dim, 1357); + std::vector ids(count); + std::iota(ids.begin(), ids.end(), 0); + auto index_made = vsag::Factory::CreateIndex("hgraph", params); + REQUIRE(index_made.has_value()); + auto index = index_made.value(); + REQUIRE(index->Build(MakeDataset(base_data, ids, dim)).has_value()); + + const std::vector update_id = {3}; + std::vector replacement(base_data.begin() + update_id[0] * num_bytes, + base_data.begin() + (update_id[0] + 1) * num_bytes); + replacement[0] ^= 0x01U; + auto updated = + index->UpdateVector(update_id[0], MakeDataset(replacement, update_id, dim), true); + REQUIRE(updated.has_value()); + REQUIRE(updated.value()); + + auto query = vsag::Dataset::Make() + ->NumElements(1) + ->Dim(dim) + ->BinaryVectors(replacement.data()) + ->Owner(false); + auto distance = index->CalcDistanceById(query, update_id[0]); + REQUIRE(distance.has_value()); + REQUIRE(distance.value() == 0.0F); + + auto stored = index->GetDataByIds(update_id.data(), 1); + REQUIRE(stored.has_value()); + REQUIRE(stored.value()->GetBinaryVectors() != nullptr); + REQUIRE(std::memcmp(replacement.data(), stored.value()->GetBinaryVectors(), num_bytes) == 0); + auto raw = index->GetRawVectorByIds(update_id.data(), 1); + REQUIRE(raw.has_value()); + REQUIRE(raw.value()->GetBinaryVectors() != nullptr); + REQUIRE(std::memcmp(replacement.data(), raw.value()->GetBinaryVectors(), num_bytes) == 0); + + auto tuned = index->Tune(R"({"index_param":{"use_reorder":true}})"); + REQUIRE(tuned.has_value()); + REQUIRE(tuned.value()); + replacement[0] ^= 0x02U; + updated = index->UpdateVector(update_id[0], MakeDataset(replacement, update_id, dim), true); + REQUIRE(updated.has_value()); + REQUIRE(updated.value()); + raw = index->GetRawVectorByIds(update_id.data(), 1); + REQUIRE(raw.has_value()); + REQUIRE(raw.value()->GetBinaryVectors() != nullptr); + REQUIRE(std::memcmp(replacement.data(), raw.value()->GetBinaryVectors(), num_bytes) == 0); + + // Reorder also exercises binary base/precise distances in the analyzer. + REQUIRE_NOTHROW(index->GetStats()); +} + +TEST_CASE("(PR) Float HGraph SQ8 Training Ignores Auxiliary Binary Storage", "[ft][hgraph][pr]") { + constexpr int64_t dim = 8; + constexpr int64_t count = 1024; + const std::string params = R"( + { + "dtype": "float32", + "metric_type": "l2", + "dim": 8, + "index_param": { + "base_quantization_type": "sq8", + "train_sample_count": 512 + } + })"; + std::vector vectors(count * dim); + for (int64_t i = 0; i < count; ++i) { + for (int64_t d = 0; d < dim; ++d) { + vectors[i * dim + d] = static_cast((i * 17 + d * 31) % 101) / 10.0F; + } + } + auto auxiliary_binary = GeneratePacked(count, dim, 9753); + std::vector ids(count); + std::iota(ids.begin(), ids.end(), 0); + auto mixed_dataset = vsag::Dataset::Make() + ->NumElements(count) + ->Dim(dim) + ->Ids(ids.data()) + ->Float32Vectors(vectors.data()) + ->BinaryVectors(auxiliary_binary.data()) + ->Owner(false); + auto index_made = vsag::Factory::CreateIndex("hgraph", params); + REQUIRE(index_made.has_value()); + auto index = index_made.value(); + REQUIRE(index->Build(mixed_dataset).has_value()); + + constexpr int64_t query_id = 17; + auto query = vsag::Dataset::Make() + ->NumElements(1) + ->Dim(dim) + ->Float32Vectors(vectors.data() + query_id * dim) + ->Owner(false); + auto base_distance = index->CalcDistanceById(query, query_id, false); + REQUIRE(base_distance.has_value()); + // Untrained SQ8 bounds produce a non-finite distance here; trained bounds must not. + REQUIRE(std::isfinite(base_distance.value())); + REQUIRE(base_distance.value() >= 0.0F); +} + +TEST_CASE("(PR) Binary HGraph Analysis Uses Packed Queries", "[ft][hgraph][pr]") { + constexpr int64_t dim = 65; + constexpr int64_t count = 32; + const std::string params = R"( + { + "dtype": "binary", + "metric_type": "hamming", + "dim": 65, + "index_param": {} + })"; + auto base_data = GeneratePacked(count, dim, 8642); + std::vector ids(count); + std::iota(ids.begin(), ids.end(), 0); + auto index_made = vsag::Factory::CreateIndex("hgraph", params); + REQUIRE(index_made.has_value()); + auto index = index_made.value(); + REQUIRE(index->Build(MakeDataset(base_data, ids, dim)).has_value()); + REQUIRE(index->CheckFeature(vsag::IndexFeature::SUPPORT_GET_STATS)); + + std::string stats; + REQUIRE_NOTHROW(stats = index->GetStats()); + REQUIRE_FALSE(stats.empty()); + + auto queries = vsag::Dataset::Make() + ->NumElements(4) + ->Dim(dim) + ->BinaryVectors(base_data.data()) + ->Owner(false); + vsag::SearchRequest request; + request.query_ = queries; + request.topk_ = 4; + request.params_str_ = R"({"hgraph":{"ef_search":32}})"; + std::string analysis; + REQUIRE_NOTHROW(analysis = index->AnalyzeIndexBySearch(request)); + REQUIRE_FALSE(analysis.empty()); + + auto short_data = GeneratePacked(4, dim - 1, 8643); + auto wrong_dim_queries = vsag::Dataset::Make() + ->NumElements(4) + ->Dim(dim - 1) + ->BinaryVectors(short_data.data()) + ->Owner(false); + request.query_ = wrong_dim_queries; + REQUIRE_THROWS(index->AnalyzeIndexBySearch(request)); + + auto missing_vectors = vsag::Dataset::Make()->NumElements(4)->Dim(dim)->Owner(false); + request.query_ = missing_vectors; + REQUIRE_THROWS(index->AnalyzeIndexBySearch(request)); +} + +TEST_CASE("(PR) Binary HGraph Rejects MCI During Validation", "[ft][hgraph][pr]") { + const std::string params = R"( + { + "dtype": "binary", + "metric_type": "hamming", + "dim": 64, + "index_param": { + "use_mci": true + } + })"; + auto index = vsag::Factory::CreateIndex("hgraph", params); + REQUIRE_FALSE(index.has_value()); + REQUIRE(index.error().type == vsag::ErrorType::INVALID_ARGUMENT); +} + +TEST_CASE("(PR) BruteForce Binary Hamming Update And Force Remove", "[ft][bruteforce][pr]") { + constexpr int64_t dim = 65; + constexpr uint64_t num_bytes = (dim + 7) / 8; + const std::string params = R"( + { + "dtype": "binary", + "metric_type": "hamming", + "dim": 65 + })"; + const std::vector ids = {10, 20, 30, 40}; + auto base_data = GeneratePacked(ids.size(), dim, 1234); + auto index_made = vsag::Factory::CreateIndex("brute_force", params); + REQUIRE(index_made.has_value()); + auto index = index_made.value(); + REQUIRE(index->Build(MakeDataset(base_data, ids, dim)).has_value()); + + auto replacement = GeneratePacked(1, dim, 5678); + replacement.back() ^= 0x80U; // Unused tail bits must be masked during re-encoding. + const std::vector update_id = {20}; + auto update = index->UpdateVector(update_id[0], MakeDataset(replacement, update_id, dim)); + REQUIRE(update.has_value()); + REQUIRE(update.value()); + auto replacement_query = vsag::Dataset::Make() + ->NumElements(1) + ->Dim(dim) + ->BinaryVectors(replacement.data()) + ->Owner(false); + auto updated_distance = index->CalcDistanceById(replacement_query, update_id[0]); + REQUIRE(updated_distance.has_value()); + REQUIRE(updated_distance.value() == 0.0F); + + // Removing a non-last slot decodes and reinserts the last packed row. + auto removed = index->Remove(ids.front(), vsag::RemoveMode::FORCE_REMOVE); + REQUIRE(removed.has_value()); + REQUIRE(removed.value() == 1); + auto moved_query = vsag::Dataset::Make() + ->NumElements(1) + ->Dim(dim) + ->BinaryVectors(base_data.data() + 3 * num_bytes) + ->Owner(false); + auto moved_distance = index->CalcDistanceById(moved_query, ids.back()); + REQUIRE(moved_distance.has_value()); + REQUIRE(moved_distance.value() == 0.0F); +} + +TEST_CASE("(PR) BruteForce Binary Rejects Malformed Queries Before Encoding", + "[ft][bruteforce][pr]") { + constexpr int64_t dim = 65; + const std::string params = R"( + { + "dtype": "binary", + "metric_type": "hamming", + "dim": 65 + })"; + const std::vector ids = {7}; + auto base_data = GeneratePacked(ids.size(), dim, 9012); + auto index_made = vsag::Factory::CreateIndex("brute_force", params); + REQUIRE(index_made.has_value()); + auto index = index_made.value(); + REQUIRE(index->Build(MakeDataset(base_data, ids, dim)).has_value()); + + std::vector short_data(1, 0); + auto wrong_dim_query = vsag::Dataset::Make() + ->NumElements(1) + ->Dim(8) + ->BinaryVectors(short_data.data()) + ->Owner(false); + auto wrong_dim_result = index->KnnSearch(wrong_dim_query, 1, "{}"); + REQUIRE_FALSE(wrong_dim_result.has_value()); + REQUIRE(wrong_dim_result.error().type == vsag::ErrorType::INVALID_ARGUMENT); + + std::vector float_data(dim, 0.0F); + auto missing_binary_query = vsag::Dataset::Make() + ->NumElements(1) + ->Dim(dim) + ->Float32Vectors(float_data.data()) + ->Owner(false); + auto missing_binary_result = index->KnnSearch(missing_binary_query, 1, "{}"); + REQUIRE_FALSE(missing_binary_result.has_value()); + REQUIRE(missing_binary_result.error().type == vsag::ErrorType::INVALID_ARGUMENT); +} + +TEST_CASE("(PR) Binary Unsupported Index Backends Are Rejected", "[ft][factory][pr]") { + const std::string params = R"( + { + "dtype": "binary", + "metric_type": "hamming", + "dim": 64 + })"; + for (const auto* name : + {"ivf", "lazy_hgraph", "pyramid", "sindi", "sindi_v2", "simq", "warp"}) { + INFO(name); + auto index = vsag::Factory::CreateIndex(name, params); + REQUIRE_FALSE(index.has_value()); + REQUIRE(index.error().type == vsag::ErrorType::INVALID_ARGUMENT); + } +} + +TEST_CASE("(PR) Binary HGraph Range And Batch Distances Reject Dense-Only Query", + "[ft][hgraph][bruteforce][pr]") { + constexpr int64_t dim = 65; + constexpr int64_t count = 16; + const std::string params = R"( + { + "dtype": "binary", + "metric_type": "hamming", + "dim": 65, + "index_param": { + "base_quantization_type": "binary" + } + })"; + const std::string brute_params = R"( + { + "dtype": "binary", + "metric_type": "hamming", + "dim": 65 + })"; + + auto base_data = GeneratePacked(count, dim, 3262); + std::vector ids(count); + std::iota(ids.begin(), ids.end(), 0); + + auto graph_made = vsag::Factory::CreateIndex("hgraph", params); + REQUIRE(graph_made.has_value()); + auto graph = graph_made.value(); + REQUIRE(graph->Build(MakeDataset(base_data, ids, dim)).has_value()); + + auto brute_made = vsag::Factory::CreateIndex("brute_force", brute_params); + REQUIRE(brute_made.has_value()); + auto brute = brute_made.value(); + REQUIRE(brute->Build(MakeDataset(base_data, ids, dim)).has_value()); + + // A float32-only Dataset must be rejected like KNN does, instead of decoding the + // floats as packed bits or crashing on a nullptr vector storage. + std::vector float_data(dim, 0.0F); + auto dense_query = vsag::Dataset::Make() + ->NumElements(1) + ->Dim(dim) + ->Float32Vectors(float_data.data()) + ->Owner(false); + + auto range_result = graph->RangeSearch(dense_query, 40.0F, "{}"); + REQUIRE_FALSE(range_result.has_value()); + REQUIRE(range_result.error().type == vsag::ErrorType::INVALID_ARGUMENT); + + vsag::DatasetPtr null_query; + auto null_range_result = graph->RangeSearch(null_query, 40.0F, "{}"); + REQUIRE_FALSE(null_range_result.has_value()); + REQUIRE(null_range_result.error().type == vsag::ErrorType::INVALID_ARGUMENT); + + auto batch_result = brute->CalcDistancesById(dense_query, ids.data(), count, true, -1); + REQUIRE_FALSE(batch_result.has_value()); + REQUIRE(batch_result.error().type == vsag::ErrorType::INVALID_ARGUMENT); + + auto topk_result = brute->CalcDistancesById(dense_query, ids.data(), count, true, 2); + REQUIRE_FALSE(topk_result.has_value()); + REQUIRE(topk_result.error().type == vsag::ErrorType::INVALID_ARGUMENT); + + auto graph_batch_result = graph->CalcDistancesById(dense_query, ids.data(), count, true, -1); + REQUIRE_FALSE(graph_batch_result.has_value()); + REQUIRE(graph_batch_result.error().type == vsag::ErrorType::INVALID_ARGUMENT); + + auto graph_topk_result = graph->CalcDistancesById(dense_query, ids.data(), count, true, 2); + REQUIRE_FALSE(graph_topk_result.has_value()); + REQUIRE(graph_topk_result.error().type == vsag::ErrorType::INVALID_ARGUMENT); +} + +TEST_CASE("(PR) HGraph Binary Update Rejects Malformed Replacement Before Encoding", + "[ft][hgraph][pr]") { + constexpr int64_t dim = 65; + constexpr int64_t count = 16; + const std::string params = R"( + { + "dtype": "binary", + "metric_type": "hamming", + "dim": 65, + "index_param": {} + })"; + auto base_data = GeneratePacked(count, dim, 4262); + std::vector ids(count); + std::iota(ids.begin(), ids.end(), 0); + auto index_made = vsag::Factory::CreateIndex("hgraph", params); + REQUIRE(index_made.has_value()); + auto index = index_made.value(); + REQUIRE(index->Build(MakeDataset(base_data, ids, dim)).has_value()); + + // A one-row, Dim(8) replacement is backed by a single packed byte, while the + // 65-bit index derives a nine-byte data size from its own dimension. Without + // validation, the ordinary path rebuilds the query at the index dimension and + // copies nine bytes out of the one-byte buffer, and the forced path re-encodes + // the same out-of-range data. Both must fail as INVALID_ARGUMENT before any + // of that can happen. + std::vector one_byte(1, 0x01U); + const std::vector update_id = {5}; + auto wrong_dim_dataset = MakeDataset(one_byte, update_id, 8); + + auto forced_result = index->UpdateVector(update_id[0], wrong_dim_dataset, true); + REQUIRE_FALSE(forced_result.has_value()); + REQUIRE(forced_result.error().type == vsag::ErrorType::INVALID_ARGUMENT); + + auto ordinary_result = index->UpdateVector(update_id[0], wrong_dim_dataset); + REQUIRE_FALSE(ordinary_result.has_value()); + REQUIRE(ordinary_result.error().type == vsag::ErrorType::INVALID_ARGUMENT); + + // A missing binary payload with the right dimension must be rejected too. + auto null_binary_dataset = vsag::Dataset::Make()->NumElements(1)->Dim(dim)->Owner(false); + auto null_forced_result = index->UpdateVector(update_id[0], null_binary_dataset, true); + REQUIRE_FALSE(null_forced_result.has_value()); + REQUIRE(null_forced_result.error().type == vsag::ErrorType::INVALID_ARGUMENT); + + auto null_ordinary_result = index->UpdateVector(update_id[0], null_binary_dataset); + REQUIRE_FALSE(null_ordinary_result.has_value()); + REQUIRE(null_ordinary_result.error().type == vsag::ErrorType::INVALID_ARGUMENT); + + // Well-formed replacements must keep working after the rejections. + auto replacement = GeneratePacked(1, dim, 555); + auto valid_result = index->UpdateVector(update_id[0], MakeDataset(replacement, update_id, dim)); + REQUIRE(valid_result.has_value()); + REQUIRE(valid_result.value()); + auto replacement_query = vsag::Dataset::Make() + ->NumElements(1) + ->Dim(dim) + ->BinaryVectors(replacement.data()) + ->Owner(false); + auto distance = index->CalcDistanceById(replacement_query, update_id[0]); + REQUIRE(distance.has_value()); + REQUIRE(distance.value() == 0.0F); +}