From efab064676de092cd40e4332cf0cb206ab8e9893 Mon Sep 17 00:00:00 2001 From: luoxiaojian Date: Thu, 30 Apr 2026 10:22:25 +0800 Subject: [PATCH] feat: refac entity and impl Vamana. (#371) --- python/tests/test_hnsw_contiguous_memory.py | 415 ++++++++++ python/tests/test_vamana.py | 481 +++++++++++ python/zvec/__init__.py | 4 + python/zvec/__init__.pyi | 11 + python/zvec/model/param/__init__.py | 4 + python/zvec/model/param/__init__.pyi | 18 +- src/binding/python/CMakeLists.txt | 3 + .../python/model/param/python_param.cc | 234 +++++- src/binding/python/model/python_collection.cc | 14 +- src/binding/python/typing/python_type.cc | 1 + src/core/algorithm/CMakeLists.txt | 1 + src/core/algorithm/hnsw/hnsw_algorithm.cc | 121 +-- src/core/algorithm/hnsw/hnsw_algorithm.h | 45 +- src/core/algorithm/hnsw/hnsw_chunk.h | 1 + src/core/algorithm/hnsw/hnsw_entity.h | 136 +++ src/core/algorithm/hnsw/hnsw_params.h | 3 + src/core/algorithm/hnsw/hnsw_streamer.cc | 193 +++-- src/core/algorithm/hnsw/hnsw_streamer.h | 28 +- .../algorithm/hnsw/hnsw_streamer_entity.cc | 212 +++++ .../algorithm/hnsw/hnsw_streamer_entity.h | 477 ++++++++++- src/core/algorithm/vamana/CMakeLists.txt | 11 + src/core/algorithm/vamana/vamana_algorithm.cc | 465 +++++++++++ src/core/algorithm/vamana/vamana_algorithm.h | 115 +++ src/core/algorithm/vamana/vamana_context.cc | 159 ++++ src/core/algorithm/vamana/vamana_context.h | 315 +++++++ .../algorithm/vamana/vamana_dist_calculator.h | 197 +++++ src/core/algorithm/vamana/vamana_entity.cc | 203 +++++ src/core/algorithm/vamana/vamana_entity.h | 352 ++++++++ .../algorithm/vamana/vamana_index_provider.h | 117 +++ src/core/algorithm/vamana/vamana_params.h | 79 ++ src/core/algorithm/vamana/vamana_streamer.cc | 714 ++++++++++++++++ src/core/algorithm/vamana/vamana_streamer.h | 192 +++++ .../vamana/vamana_streamer_entity.cc | 773 ++++++++++++++++++ .../algorithm/vamana/vamana_streamer_entity.h | 635 ++++++++++++++ src/core/interface/index_factory.cc | 36 + src/core/interface/index_param.cc | 46 ++ src/core/interface/indexes/hnsw_index.cc | 24 + src/core/interface/indexes/vamana_index.cc | 105 +++ src/core/utility/buffer_storage.cc | 5 + src/db/collection.cc | 49 ++ .../column/vector_column/engine_helper.hpp | 53 ++ .../vector_column/vector_column_indexer.h | 6 + src/db/index/common/proto_converter.cc | 41 +- src/db/index/common/proto_converter.h | 5 + src/db/index/common/schema.cc | 3 +- src/db/index/common/type_helper.h | 6 + src/db/proto/zvec.proto | 20 + .../zvec/core/framework/index_storage.h | 5 + src/include/zvec/core/interface/constants.h | 7 + src/include/zvec/core/interface/index.h | 25 +- src/include/zvec/core/interface/index_param.h | 42 + .../core/interface/index_param_builders.h | 55 +- src/include/zvec/db/collection.h | 9 + src/include/zvec/db/index_params.h | 142 +++- src/include/zvec/db/query_params.h | 25 + src/include/zvec/db/type.h | 1 + .../squared_euclidean.cc | 83 +- tests/core/algorithm/CMakeLists.txt | 1 + .../core/algorithm/hnsw/hnsw_streamer_test.cc | 207 +++++ tests/core/algorithm/vamana/CMakeLists.txt | 14 + .../algorithm/vamana/vamana_streamer_test.cc | 721 ++++++++++++++++ tests/core/interface/CMakeLists.txt | 2 +- tests/core/interface/index_interface_test.cc | 391 +++++++++ tools/core/CMakeLists.txt | 12 +- 64 files changed, 8669 insertions(+), 201 deletions(-) create mode 100644 python/tests/test_hnsw_contiguous_memory.py create mode 100644 python/tests/test_vamana.py create mode 100644 src/core/algorithm/vamana/CMakeLists.txt create mode 100644 src/core/algorithm/vamana/vamana_algorithm.cc create mode 100644 src/core/algorithm/vamana/vamana_algorithm.h create mode 100644 src/core/algorithm/vamana/vamana_context.cc create mode 100644 src/core/algorithm/vamana/vamana_context.h create mode 100644 src/core/algorithm/vamana/vamana_dist_calculator.h create mode 100644 src/core/algorithm/vamana/vamana_entity.cc create mode 100644 src/core/algorithm/vamana/vamana_entity.h create mode 100644 src/core/algorithm/vamana/vamana_index_provider.h create mode 100644 src/core/algorithm/vamana/vamana_params.h create mode 100644 src/core/algorithm/vamana/vamana_streamer.cc create mode 100644 src/core/algorithm/vamana/vamana_streamer.h create mode 100644 src/core/algorithm/vamana/vamana_streamer_entity.cc create mode 100644 src/core/algorithm/vamana/vamana_streamer_entity.h create mode 100644 src/core/interface/indexes/vamana_index.cc create mode 100644 tests/core/algorithm/vamana/CMakeLists.txt create mode 100644 tests/core/algorithm/vamana/vamana_streamer_test.cc diff --git a/python/tests/test_hnsw_contiguous_memory.py b/python/tests/test_hnsw_contiguous_memory.py new file mode 100644 index 0000000..cb23a01 --- /dev/null +++ b/python/tests/test_hnsw_contiguous_memory.py @@ -0,0 +1,415 @@ +# Copyright 2025-present the zvec 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. +""" +Tests for the ``use_contiguous_memory`` optimization on ``HnswIndexParam``. + +The HNSW streamer supports two allocation strategies for graph nodes: + +* ``use_contiguous_memory=False`` (default): each node allocates its own + linked buffer. Lower peak memory usage, worse cache locality. +* ``use_contiguous_memory=True``: a single contiguous arena holds every + node. Higher peak memory usage, better cache locality and search + throughput. + +These tests exercise the Python surface end-to-end and make sure that +when a collection is created / reopened with ``use_contiguous_memory=True`` +the underlying HNSW streamer entity is constructed correctly and serves +search traffic. +""" + +from __future__ import annotations + +import pickle +import sys + +import numpy as np +import pytest + +import zvec +from zvec import ( + Collection, + CollectionOption, + CollectionSchema, + Doc, + FieldSchema, + HnswIndexParam, + HnswQueryParam, + InvertIndexParam, + VectorQuery, + VectorSchema, +) +from zvec.typing import DataType, IndexType, MetricType, QuantizeType + + +DIMENSION = 32 +NUM_DOCS = 128 +TOPK = 5 + +# --------------------------------------------------------------------------- + + +def _debug_hnsw_storage_mode(coll: Collection, column: str = "dense") -> str: + """Return the internal HNSW entity storage mode for ``column``. + + Exposes the debug-only introspection hook on the pybind11 ``_Collection``. + Only meaningful after ``optimize()`` has built a persisted HNSW index; on + a pure writing segment it will raise ``KeyError``. + """ + underlying = coll._obj # type: ignore[attr-defined] + return underlying._debug_hnsw_storage_mode(column) + + +def _build_schema(name: str, *, use_contiguous_memory: bool) -> CollectionSchema: + """Create a simple schema with a single FP32 HNSW vector column.""" + return CollectionSchema( + name=name, + fields=[ + FieldSchema( + "id", + DataType.INT64, + nullable=False, + index_param=InvertIndexParam(enable_range_optimization=True), + ), + ], + vectors=[ + VectorSchema( + "dense", + DataType.VECTOR_FP32, + dimension=DIMENSION, + index_param=HnswIndexParam( + metric_type=MetricType.IP, + m=16, + ef_construction=100, + use_contiguous_memory=use_contiguous_memory, + ), + ), + ], + ) + + +def _generate_docs(rng: np.random.Generator, num: int = NUM_DOCS) -> list[Doc]: + """Produce deterministic documents for insertion.""" + docs: list[Doc] = [] + for i in range(num): + vec = rng.standard_normal(DIMENSION).astype(np.float32) + docs.append( + Doc( + id=str(i), + fields={"id": i}, + vectors={"dense": vec.tolist()}, + ) + ) + return docs + + +def _assert_query_matches(coll: Collection, query_vec: list[float]) -> list[str]: + """Run a top-k vector query and return the returned ids in order.""" + vector_query = VectorQuery( + field_name="dense", + vector=query_vec, + param=HnswQueryParam(ef=128), + ) + hits = coll.query(vectors=vector_query, topk=TOPK) + # Expect a single result group for the single vector query. + assert hits is not None, "query returned None" + assert len(hits) >= 1, f"expected at least one hit, got {hits!r}" + return [doc.id for doc in hits] + + +# --------------------------------------------------------------------------- +# 1) Pure Python surface: construction / property / to_dict / repr / pickle +# --------------------------------------------------------------------------- + + +class TestHnswIndexParamContiguousMemorySurface: + """Verify the Python binding exposes ``use_contiguous_memory`` correctly.""" + + def test_default_is_false(self): + param = HnswIndexParam() + assert param.use_contiguous_memory is False + + def test_custom_true(self): + param = HnswIndexParam(use_contiguous_memory=True) + assert param.use_contiguous_memory is True + assert param.type == IndexType.HNSW + # other fields keep their default values + assert param.m == 50 + assert param.ef_construction == 500 + + def test_to_dict_includes_use_contiguous_memory(self): + param = HnswIndexParam( + metric_type=MetricType.L2, + m=16, + ef_construction=100, + quantize_type=QuantizeType.FP16, + use_contiguous_memory=True, + ) + data = param.to_dict() + assert data["use_contiguous_memory"] is True + # Make sure existing fields are still present. + assert data["metric_type"] == "L2" + assert data["m"] == 16 + assert data["ef_construction"] == 100 + assert data["quantize_type"] == "FP16" + + def test_repr_contains_flag(self): + on = repr(HnswIndexParam(use_contiguous_memory=True)) + off = repr(HnswIndexParam(use_contiguous_memory=False)) + assert "use_contiguous_memory" in on + assert "use_contiguous_memory" in off + assert "true" in on + assert "false" in off + + def test_readonly_property(self): + param = HnswIndexParam(use_contiguous_memory=True) + if sys.version_info >= (3, 11): + match_pattern = r"(can't set attribute|has no setter|readonly attribute)" + else: + match_pattern = r"can't set attribute" + with pytest.raises(AttributeError, match=match_pattern): + param.use_contiguous_memory = False # type: ignore[misc] + + def test_pickle_roundtrip(self): + original = HnswIndexParam( + metric_type=MetricType.COSINE, + m=24, + ef_construction=150, + quantize_type=QuantizeType.INT8, + use_contiguous_memory=True, + ) + restored = pickle.loads(pickle.dumps(original)) + assert restored.use_contiguous_memory is True + assert restored.metric_type == MetricType.COSINE + assert restored.m == 24 + assert restored.ef_construction == 150 + assert restored.quantize_type == QuantizeType.INT8 + + +# --------------------------------------------------------------------------- +# 2) End-to-end: create collection, insert, query with contiguous memory on +# --------------------------------------------------------------------------- + + +@pytest.fixture +def rng() -> np.random.Generator: + return np.random.default_rng(seed=42) + + +# NOTE: the ``enable_mmap=False`` (BufferPool) variant is intentionally +# omitted from this fixture. Building a persisted HNSW index via +# ``optimize()`` / ``create_vector_index`` / ``drop_vector_index`` +# currently requires mmap-backed storage, because the BufferPool backend +# has not implemented the ``create_new`` semantics yet and the guard in +# ``SegmentImpl::merge_vector_indexer`` rejects that combination. Once +# BufferPool gains write support, re-add ``False`` to ``params`` (and +# drop the guard in segment.cc) so these end-to-end tests cover both +# storage modes again. +@pytest.fixture(params=[True], ids=["mmap_on"]) +def collection_option(request) -> CollectionOption: + return CollectionOption(read_only=False, enable_mmap=request.param) + + +# Building a new persisted HNSW index currently requires mmap-backed storage +# because the BufferPool backend has not implemented `create_new` semantics +# yet. Collections opened with ``enable_mmap=False`` therefore cannot run +# optimize()/create_vector_index/drop_vector_index. Tests use this fixture +# to know which behaviour to assert, and once BufferPool gains write support +# the guard in segment.cc (and these branches) can be removed together. +@pytest.fixture +def build_index_supported(collection_option: CollectionOption) -> bool: + return bool(collection_option.enable_mmap) + + +# Error message fragments emitted by the NotSupported guard in +# SegmentImpl::merge_vector_indexer / drop_vector_index. If the C++ message +# changes, update these together. +_BUILD_NOT_SUPPORTED_FRAGMENTS = ("not yet supported", "enable_mmap=false") + + +class TestHnswContiguousMemoryEndToEnd: + """End-to-end: schema -> create_and_open -> insert -> query works.""" + + def test_create_with_contiguous_memory_and_query( + self, + tmp_path_factory, + collection_option, + rng, + ): + """With the flag on, the schema round-trips and search works end-to-end. + + After ``optimize()`` the writing segment is compacted into a persisted + segment backed by the configured HNSW entity. We assert both the + user-observable behaviour (schema + search) and, via the debug hook, + that the entity type actually honours ``use_contiguous_memory``. + """ + schema = _build_schema("hnsw_contig_create", use_contiguous_memory=True) + + path = tmp_path_factory.mktemp("zvec") / "hnsw_contig_create" + coll = zvec.create_and_open( + path=str(path), schema=schema, option=collection_option + ) + try: + # Schema round-trips with the flag set. + vec_schema = coll.schema.vectors[0] + assert vec_schema.index_param.use_contiguous_memory is True + + docs = _generate_docs(rng) + insert_result = coll.insert(docs=docs) + for r in insert_result: + assert r.ok(), f"insert failed: code={r.code()}" + assert coll.stats.doc_count == NUM_DOCS + + # Build persisted HNSW index; this is where the contiguous entity + # is actually instantiated. + coll.optimize() + assert _debug_hnsw_storage_mode(coll) == "contiguous", ( + "use_contiguous_memory=True should produce a contiguous entity" + ) + + # Pick an existing vector as the query; top-1 must be itself. + query_vec = docs[0].vector("dense") + ids = _assert_query_matches(coll, query_vec) + assert ids[0] == "0", f"expected self-recall, got top-1 id={ids[0]}" + finally: + coll.destroy() + + def test_create_without_contiguous_memory_uses_mmap_entity( + self, + tmp_path_factory, + collection_option, + rng, + ): + """Baseline: when the flag is omitted the default (mmap) entity is used.""" + schema = _build_schema("hnsw_contig_default", use_contiguous_memory=False) + path = tmp_path_factory.mktemp("zvec") / "hnsw_contig_default" + coll = zvec.create_and_open( + path=str(path), schema=schema, option=collection_option + ) + try: + vec_schema = coll.schema.vectors[0] + assert vec_schema.index_param.use_contiguous_memory is False + + docs = _generate_docs(rng) + for r in coll.insert(docs=docs): + assert r.ok() + assert coll.stats.doc_count == NUM_DOCS + + coll.optimize() + # With the flag off and mmap on, the persisted entity must be the + # default mmap layout — specifically, not the contiguous arena. + assert _debug_hnsw_storage_mode(coll) == "mmap", ( + "use_contiguous_memory=False + enable_mmap=True should " + "produce the mmap entity" + ) + + # Search still functions with the default entity backing. + query_vec = docs[0].vector("dense") + ids = _assert_query_matches(coll, query_vec) + assert ids[0] == "0" + finally: + coll.destroy() + + def test_close_and_reopen_with_contiguous_memory( + self, + tmp_path_factory, + collection_option, + rng, + ): + """Reopening a collection must preserve the ``use_contiguous_memory`` flag. + + The core property: the flag survives the schema persist/reload + round-trip so the HNSW streamer entity — constructed lazily on first + persisted-segment build — honours the user's choice. We run + ``optimize()`` after reopen and confirm the contiguous entity was + materialized. + """ + schema = _build_schema("hnsw_contig_reopen", use_contiguous_memory=True) + path = tmp_path_factory.mktemp("zvec") / "hnsw_contig_reopen" + path_str = str(path) + + created = zvec.create_and_open( + path=path_str, schema=schema, option=collection_option + ) + docs = _generate_docs(rng) + for r in created.insert(docs=docs): + assert r.ok() + assert created.stats.doc_count == NUM_DOCS + # Persist pending writes so that reopen reconstructs state from disk. + created.flush() + del created # close the handle + + reopened = zvec.open(path=path_str, option=collection_option) + try: + assert reopened is not None + assert reopened.stats.doc_count == NUM_DOCS + + # Schema persisted the flag across the reopen boundary. + vec_schema = reopened.schema.vectors[0] + assert vec_schema.index_param.use_contiguous_memory is True + + reopened.optimize() + assert _debug_hnsw_storage_mode(reopened) == "contiguous" + + # Entity actually works: exact self-recall + fetch parity. + query_vec = docs[7].vector("dense") + ids = _assert_query_matches(reopened, query_vec) + assert ids[0] == "7" + + fetched = reopened.fetch([d.id for d in docs[:10]]) + assert len(fetched) == 10 + finally: + reopened.destroy() + + def test_result_parity_with_and_without_contiguous_memory( + self, + tmp_path_factory, + rng, + ): + """ + Two collections built from the same documents must return the same + top-k neighbors regardless of whether contiguous memory is enabled: + the flag is a memory-layout optimization and must not alter recall + for identical graph construction parameters on the same data. + """ + docs = _generate_docs(rng) + query_vec = docs[3].vector("dense") + + def _build_and_query(tag: str, flag: bool) -> list[str]: + schema = _build_schema(f"hnsw_parity_{tag}", use_contiguous_memory=flag) + option = CollectionOption(read_only=False, enable_mmap=True) + path = tmp_path_factory.mktemp("zvec") / f"hnsw_parity_{tag}" + coll = zvec.create_and_open(path=str(path), schema=schema, option=option) + try: + for r in coll.insert(docs=docs): + assert r.ok() + coll.optimize() + expected_mode = "contiguous" if flag else "mmap" + assert _debug_hnsw_storage_mode(coll) == expected_mode, ( + f"{tag}: unexpected entity type" + ) + return _assert_query_matches(coll, query_vec) + finally: + coll.destroy() + + ids_off = _build_and_query("off", flag=False) + ids_on = _build_and_query("on", flag=True) + + # The graph is built with the same (m, ef_construction, data, order), + # so top-k results must match exactly. + assert ids_on == ids_off, ( + f"top-{TOPK} results diverged between use_contiguous_memory modes: " + f"on={ids_on}, off={ids_off}" + ) + # Sanity: self-recall is still perfect. + assert ids_on[0] == "3" diff --git a/python/tests/test_vamana.py b/python/tests/test_vamana.py new file mode 100644 index 0000000..359a9d8 --- /dev/null +++ b/python/tests/test_vamana.py @@ -0,0 +1,481 @@ +# Copyright 2025-present the zvec 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. +""" +Tests for the Python entry point of the Vamana (DiskANN) dense vector index. + +Mirrors the structure of ``test_hnsw_contiguous_memory.py`` (the closest +hnsw dense reference), and is split into two parts: + +1. **Surface tests** — verify that ``VamanaIndexParam`` / ``VamanaQueryParam`` + are correctly bound: construction defaults, readonly properties, + ``to_dict``, ``__repr__``, pickle round-trip, and that they appear in the + public ``zvec`` namespace with the expected ``IndexType.VAMANA`` value. + +2. **End-to-end tests** — build a collection that uses Vamana on a dense + FP32 column, insert deterministic documents, then run a top-k query + through ``VamanaQueryParam`` on both the writer segment and the + persisted (post-``optimize()``) segment. +""" + +from __future__ import annotations + +import pickle +import sys + +import numpy as np +import pytest + +import zvec +from zvec import ( + Collection, + CollectionOption, + CollectionSchema, + Doc, + FieldSchema, + InvertIndexParam, + VamanaIndexParam, + VamanaQueryParam, + VectorQuery, + VectorSchema, +) +from zvec.typing import DataType, IndexType, MetricType, QuantizeType + +DIMENSION = 32 +NUM_DOCS = 128 +TOPK = 5 + +# Defaults pulled from src/include/zvec/core/interface/constants.h. Keep +# in sync with kDefaultVamana* if the engine defaults ever change. +DEFAULT_MAX_DEGREE = 64 +DEFAULT_SEARCH_LIST_SIZE = 100 +DEFAULT_ALPHA = 1.2 +DEFAULT_EF_SEARCH = 200 +DEFAULT_SATURATE_GRAPH = False + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _build_schema( + name: str, + *, + metric_type: MetricType = MetricType.IP, + max_degree: int = 32, + search_list_size: int = 64, + alpha: float = 1.2, + use_contiguous_memory: bool = False, +) -> CollectionSchema: + """Create a simple schema with a single FP32 Vamana vector column.""" + return CollectionSchema( + name=name, + fields=[ + FieldSchema( + "id", + DataType.INT64, + nullable=False, + index_param=InvertIndexParam(enable_range_optimization=True), + ), + ], + vectors=[ + VectorSchema( + "dense", + DataType.VECTOR_FP32, + dimension=DIMENSION, + index_param=VamanaIndexParam( + metric_type=metric_type, + max_degree=max_degree, + search_list_size=search_list_size, + alpha=alpha, + use_contiguous_memory=use_contiguous_memory, + ), + ), + ], + ) + + +def _generate_docs(rng: np.random.Generator, num: int = NUM_DOCS) -> list[Doc]: + """Produce deterministic documents for insertion.""" + docs: list[Doc] = [] + for i in range(num): + vec = rng.standard_normal(DIMENSION).astype(np.float32) + docs.append( + Doc( + id=str(i), + fields={"id": i}, + vectors={"dense": vec.tolist()}, + ) + ) + return docs + + +def _query_topk( + coll: Collection, query_vec: list[float], *, ef_search: int = 64 +) -> list[str]: + """Run a top-k vector query and return the returned ids in order.""" + vector_query = VectorQuery( + field_name="dense", + vector=query_vec, + param=VamanaQueryParam(ef_search=ef_search), + ) + hits = coll.query(vectors=vector_query, topk=TOPK) + assert hits is not None, "query returned None" + assert len(hits) >= 1, f"expected at least one hit, got {hits!r}" + return [doc.id for doc in hits] + + +# --------------------------------------------------------------------------- +# 1) Surface: construction / property / to_dict / repr / pickle / namespace +# --------------------------------------------------------------------------- + + +class TestVamanaIndexParamSurface: + """Verify the Python binding for ``VamanaIndexParam``.""" + + def test_defaults(self): + param = VamanaIndexParam() + assert param.type == IndexType.VAMANA + assert param.metric_type == MetricType.IP + assert param.max_degree == DEFAULT_MAX_DEGREE + assert param.search_list_size == DEFAULT_SEARCH_LIST_SIZE + assert param.alpha == pytest.approx(DEFAULT_ALPHA) + assert param.saturate_graph is DEFAULT_SATURATE_GRAPH + assert param.use_contiguous_memory is False + assert param.use_id_map is False + assert param.quantize_type == QuantizeType.UNDEFINED + + def test_custom_construction(self): + param = VamanaIndexParam( + metric_type=MetricType.COSINE, + max_degree=48, + search_list_size=128, + alpha=1.5, + saturate_graph=True, + use_contiguous_memory=True, + use_id_map=False, + quantize_type=QuantizeType.INT8, + ) + assert param.type == IndexType.VAMANA + assert param.metric_type == MetricType.COSINE + assert param.max_degree == 48 + assert param.search_list_size == 128 + assert param.alpha == pytest.approx(1.5) + assert param.saturate_graph is True + assert param.use_contiguous_memory is True + assert param.use_id_map is False + assert param.quantize_type == QuantizeType.INT8 + + def test_to_dict_includes_all_fields(self): + param = VamanaIndexParam( + metric_type=MetricType.L2, + max_degree=32, + search_list_size=80, + alpha=1.3, + saturate_graph=True, + use_contiguous_memory=True, + use_id_map=False, + quantize_type=QuantizeType.FP16, + ) + data = param.to_dict() + assert data["type"] == "VAMANA" + assert data["metric_type"] == "L2" + assert data["max_degree"] == 32 + assert data["search_list_size"] == 80 + assert data["alpha"] == pytest.approx(1.3) + assert data["saturate_graph"] is True + assert data["use_contiguous_memory"] is True + assert data["use_id_map"] is False + assert data["quantize_type"] == "FP16" + + def test_repr_contains_key_fields(self): + text = repr( + VamanaIndexParam( + metric_type=MetricType.COSINE, + max_degree=24, + search_list_size=72, + alpha=1.4, + saturate_graph=True, + use_contiguous_memory=True, + ) + ) + # Spot-check the most diagnostic fields are rendered. + assert "VAMANA" in text + assert "COSINE" in text + assert "max_degree" in text and "24" in text + assert "search_list_size" in text and "72" in text + assert "alpha" in text + assert "saturate_graph" in text and "true" in text + assert "use_contiguous_memory" in text and "true" in text + + @pytest.mark.parametrize( + "field, kwargs", + [ + ("max_degree", dict(max_degree=99)), + ("search_list_size", dict(search_list_size=99)), + ("alpha", dict(alpha=1.7)), + ("saturate_graph", dict(saturate_graph=True)), + ("use_contiguous_memory", dict(use_contiguous_memory=True)), + ("use_id_map", dict(use_id_map=True)), + ], + ) + def test_readonly_properties(self, field, kwargs): + param = VamanaIndexParam(**kwargs) + if sys.version_info >= (3, 11): + match_pattern = r"(can't set attribute|has no setter|readonly attribute)" + else: + match_pattern = r"can't set attribute" + with pytest.raises(AttributeError, match=match_pattern): + setattr(param, field, getattr(param, field)) + + def test_pickle_roundtrip(self): + original = VamanaIndexParam( + metric_type=MetricType.COSINE, + max_degree=48, + search_list_size=120, + alpha=1.4, + saturate_graph=True, + use_contiguous_memory=True, + use_id_map=False, + quantize_type=QuantizeType.INT8, + ) + restored = pickle.loads(pickle.dumps(original)) + assert restored.type == IndexType.VAMANA + assert restored.metric_type == MetricType.COSINE + assert restored.max_degree == 48 + assert restored.search_list_size == 120 + assert restored.alpha == pytest.approx(1.4) + assert restored.saturate_graph is True + assert restored.use_contiguous_memory is True + assert restored.use_id_map is False + assert restored.quantize_type == QuantizeType.INT8 + # to_dict equality is the strongest end-to-end equivalence we have. + assert restored.to_dict() == original.to_dict() + + +class TestVamanaQueryParamSurface: + """Verify the Python binding for ``VamanaQueryParam``.""" + + def test_defaults(self): + q = VamanaQueryParam() + assert q.type == IndexType.VAMANA + assert q.ef_search == DEFAULT_EF_SEARCH + assert q.radius == pytest.approx(0.0) + assert q.is_linear is False + assert q.is_using_refiner is False + + def test_custom_construction(self): + q = VamanaQueryParam( + ef_search=300, radius=0.5, is_linear=True, is_using_refiner=True + ) + assert q.type == IndexType.VAMANA + assert q.ef_search == 300 + assert q.radius == pytest.approx(0.5) + assert q.is_linear is True + assert q.is_using_refiner is True + + def test_repr_contains_key_fields(self): + text = repr(VamanaQueryParam(ef_search=128, radius=0.25)) + assert "VAMANA" in text + assert "ef_search" in text and "128" in text + assert "radius" in text + + def test_readonly_ef_search(self): + q = VamanaQueryParam(ef_search=100) + if sys.version_info >= (3, 11): + match_pattern = r"(can't set attribute|has no setter|readonly attribute)" + else: + match_pattern = r"can't set attribute" + with pytest.raises(AttributeError, match=match_pattern): + q.ef_search = 200 # type: ignore[misc] + + def test_pickle_roundtrip(self): + original = VamanaQueryParam( + ef_search=256, radius=0.3, is_linear=False, is_using_refiner=True + ) + restored = pickle.loads(pickle.dumps(original)) + assert restored.type == IndexType.VAMANA + assert restored.ef_search == 256 + assert restored.radius == pytest.approx(0.3) + assert restored.is_linear is False + assert restored.is_using_refiner is True + + +class TestVamanaPublicNamespace: + """The Vamana entry points must be importable from the top-level ``zvec``.""" + + def test_top_level_exports(self): + assert zvec.VamanaIndexParam is VamanaIndexParam + assert zvec.VamanaQueryParam is VamanaQueryParam + assert "VamanaIndexParam" in zvec.__all__ + assert "VamanaQueryParam" in zvec.__all__ + + def test_index_type_enum_member(self): + # Sanity: the IndexType enum exposes VAMANA and it is what the + # bound params advertise. + assert IndexType.VAMANA is not None + assert VamanaIndexParam().type == IndexType.VAMANA + assert VamanaQueryParam().type == IndexType.VAMANA + + +# --------------------------------------------------------------------------- +# 2) End-to-end: create collection, insert, query through the writer segment +# --------------------------------------------------------------------------- + + +@pytest.fixture +def rng() -> np.random.Generator: + return np.random.default_rng(seed=42) + + +# Mirror the hnsw dense test fixture: only the mmap-backed variant is +# currently usable for vector index construction. BufferPool (enable_mmap= +# False) is intentionally omitted because the same write-path guard in +# ``SegmentImpl::merge_vector_indexer`` rejects that combination. +@pytest.fixture(params=[True], ids=["mmap_on"]) +def collection_option(request) -> CollectionOption: + return CollectionOption(read_only=False, enable_mmap=request.param) + + +class TestVamanaEndToEnd: + """End-to-end: schema -> create_and_open -> insert -> query works.""" + + def test_schema_round_trip(self, tmp_path_factory, collection_option): + """The Vamana index params survive the schema persist path.""" + schema = _build_schema( + "vamana_schema_rt", + metric_type=MetricType.COSINE, + max_degree=32, + search_list_size=80, + alpha=1.3, + use_contiguous_memory=True, + ) + path = tmp_path_factory.mktemp("zvec") / "vamana_schema_rt" + coll = zvec.create_and_open( + path=str(path), schema=schema, option=collection_option + ) + try: + vec_schema = coll.schema.vectors[0] + ip = vec_schema.index_param + assert ip.type == IndexType.VAMANA + assert ip.metric_type == MetricType.COSINE + assert ip.max_degree == 32 + assert ip.search_list_size == 80 + assert ip.alpha == pytest.approx(1.3) + assert ip.use_contiguous_memory is True + finally: + coll.destroy() + + def test_insert_and_query_self_recall( + self, tmp_path_factory, collection_option, rng + ): + """Top-1 of a query equal to an inserted vector must be that vector. + + Exercises the writer-segment Vamana streamer end-to-end through the + Python entry point: ``VamanaIndexParam`` for build and + ``VamanaQueryParam`` for search. + """ + schema = _build_schema("vamana_e2e_recall") + path = tmp_path_factory.mktemp("zvec") / "vamana_e2e_recall" + coll = zvec.create_and_open( + path=str(path), schema=schema, option=collection_option + ) + try: + docs = _generate_docs(rng) + for r in coll.insert(docs=docs): + assert r.ok(), f"insert failed: code={r.code()}" + assert coll.stats.doc_count == NUM_DOCS + + # Self-recall: query with the i-th inserted vector, expect id i + # to be the top result. + for probe in (0, 7, 42, NUM_DOCS - 1): + query_vec = docs[probe].vector("dense") + ids = _query_topk(coll, query_vec) + assert ids[0] == str(probe), ( + f"expected self-recall at probe={probe}, got top-1 id={ids[0]} " + f"(top-{TOPK}={ids})" + ) + finally: + coll.destroy() + + def test_query_param_ef_search_affects_only_quality( + self, tmp_path_factory, collection_option, rng + ): + """``ef_search`` is a search-time knob and must not crash for any + sensible value. Larger ``ef_search`` should be at least as good as + smaller for self-recall.""" + schema = _build_schema("vamana_e2e_ef") + path = tmp_path_factory.mktemp("zvec") / "vamana_e2e_ef" + coll = zvec.create_and_open( + path=str(path), schema=schema, option=collection_option + ) + try: + docs = _generate_docs(rng) + for r in coll.insert(docs=docs): + assert r.ok() + + query_vec = docs[3].vector("dense") + ids_small = _query_topk(coll, query_vec, ef_search=16) + ids_large = _query_topk(coll, query_vec, ef_search=256) + + # Both should self-recall the probe vector at top-1. + assert ids_small[0] == "3" + assert ids_large[0] == "3" + assert len(ids_small) == TOPK + assert len(ids_large) == TOPK + finally: + coll.destroy() + + def test_optimize_then_query(self, tmp_path_factory, collection_option, rng): + """The persisted Vamana segment built by ``optimize()`` must serve + queries correctly. + + Until the cmake fix to force-load ``core_knn_vamana_static`` into the + ``_zvec`` pybind module, this path failed at ``VamanaStreamer`` + creation because the global factory registration in + ``vamana_streamer.cc`` was never linked in. This test pins down the + regression. + """ + schema = _build_schema("vamana_e2e_optimize") + path = tmp_path_factory.mktemp("zvec") / "vamana_e2e_optimize" + coll = zvec.create_and_open( + path=str(path), schema=schema, option=collection_option + ) + try: + docs = _generate_docs(rng) + for r in coll.insert(docs=docs): + assert r.ok() + assert coll.stats.doc_count == NUM_DOCS + + # Snapshot the writer-segment top-k for a probe vector. + query_vec = docs[5].vector("dense") + ids_pre = _query_topk(coll, query_vec) + assert ids_pre[0] == "5" + + # Trigger persisted segment build. Pre-fix this raised + # RuntimeError("Failed to create index"). + coll.optimize() + + # Persisted segment must still serve queries with the same + # top-1 self-recall guarantee. We do not assert full top-k + # equality with the writer segment because the persisted + # streamer may visit nodes in a different order; top-1 self- + # recall is the strong invariant. + ids_post = _query_topk(coll, query_vec) + assert ids_post[0] == "5", ( + f"post-optimize top-1 should still be probe id, got {ids_post}" + ) + assert len(ids_post) == TOPK + finally: + coll.destroy() diff --git a/python/zvec/__init__.py b/python/zvec/__init__.py index 40db92b..298bf30 100644 --- a/python/zvec/__init__.py +++ b/python/zvec/__init__.py @@ -70,6 +70,8 @@ from .model.param import ( IVFIndexParam, IVFQueryParam, OptimizeOption, + VamanaIndexParam, + VamanaQueryParam, ) from .model.param.vector_query import VectorQuery @@ -122,6 +124,8 @@ __all__ = [ "HnswQueryParam", "HnswRabitqQueryParam", "IVFQueryParam", + "VamanaIndexParam", + "VamanaQueryParam", # Extensions "DenseEmbeddingFunction", "SparseEmbeddingFunction", diff --git a/python/zvec/__init__.pyi b/python/zvec/__init__.pyi index 97a57a1..2c47b6d 100644 --- a/python/zvec/__init__.pyi +++ b/python/zvec/__init__.pyi @@ -26,6 +26,8 @@ from .model.param import ( IVFIndexParam, IVFQueryParam, OptimizeOption, + VamanaIndexParam, + VamanaQueryParam, ) from .model.param.vector_query import VectorQuery from .model.schema import CollectionSchema, CollectionStats, FieldSchema, VectorSchema @@ -73,6 +75,8 @@ __all__: list = [ "RrfReRanker", "Status", "StatusCode", + "VamanaIndexParam", + "VamanaQueryParam", "VectorQuery", "VectorSchema", "WeightedReRanker", @@ -122,6 +126,13 @@ class _Collection: def Stats(self) -> schema.CollectionStats: ... def Update(self, arg0: collections.abc.Sequence[_Doc]) -> list[typing.Status]: ... def Upsert(self, arg0: collections.abc.Sequence[_Doc]) -> list[typing.Status]: ... + def _debug_hnsw_storage_mode(self, column_name: str) -> str: + """Debug-only: returns the storage mode of the HNSW entity on the + given vector column. One of 'mmap', 'buffer_pool', 'contiguous'. + Raises KeyError if no HNSW index exists on the column, or + ValueError if the column's index is not an HNSW index. Intended + for introspection and testing only; not part of the stable API.""" + def __getstate__(self) -> tuple: ... def __setstate__(self, arg0: tuple) -> None: ... diff --git a/python/zvec/model/param/__init__.py b/python/zvec/model/param/__init__.py index c613edf..5758218 100644 --- a/python/zvec/model/param/__init__.py +++ b/python/zvec/model/param/__init__.py @@ -27,6 +27,8 @@ from _zvec.param import ( IVFIndexParam, IVFQueryParam, OptimizeOption, + VamanaIndexParam, + VamanaQueryParam, ) __all__ = [ @@ -43,4 +45,6 @@ __all__ = [ "IndexOption", "InvertIndexParam", "OptimizeOption", + "VamanaIndexParam", + "VamanaQueryParam", ] diff --git a/python/zvec/model/param/__init__.pyi b/python/zvec/model/param/__init__.pyi index cd1491e..31b43f6 100644 --- a/python/zvec/model/param/__init__.pyi +++ b/python/zvec/model/param/__init__.pyi @@ -198,6 +198,10 @@ class HnswIndexParam(VectorIndexParam): quantize_type (QuantizeType): Optional quantization type for vector compression (e.g., FP16, INT8). Default is `QuantizeType.UNDEFINED` to disable quantization. + use_contiguous_memory (bool): If True, the HNSW streamer allocates a + single contiguous memory arena for all graph nodes, improving cache + locality and search throughput at the cost of peak memory usage. + Default is False. Examples: >>> from zvec.typing import MetricType, QuantizeType @@ -205,10 +209,11 @@ class HnswIndexParam(VectorIndexParam): ... metric_type=MetricType.COSINE, ... m=16, ... ef_construction=200, - ... quantize_type=QuantizeType.INT8 + ... quantize_type=QuantizeType.INT8, + ... use_contiguous_memory=True, ... ) >>> print(params) - {'metric_type': 'IP', 'm': 16, 'ef_construction': 200, 'quantize_type': 'INT8'} + {'metric_type': 'IP', 'm': 16, 'ef_construction': 200, 'quantize_type': 'INT8', 'use_contiguous_memory': True} """ def __getstate__(self) -> tuple: ... @@ -218,6 +223,7 @@ class HnswIndexParam(VectorIndexParam): m: typing.SupportsInt = 50, ef_construction: typing.SupportsInt = 500, quantize_type: _zvec.typing.QuantizeType = ..., + use_contiguous_memory: bool = False, ) -> None: ... def __repr__(self) -> str: ... def __setstate__(self, arg0: tuple) -> None: ... @@ -238,6 +244,14 @@ class HnswIndexParam(VectorIndexParam): int: Maximum number of neighbors per node in upper layers. """ + @property + def use_contiguous_memory(self) -> bool: + """ + bool: Whether to allocate a single contiguous memory arena for all + HNSW graph nodes. Improves cache locality and search throughput at + the cost of peak memory usage. Defaults to False. + """ + class HnswQueryParam(QueryParam): """ diff --git a/src/binding/python/CMakeLists.txt b/src/binding/python/CMakeLists.txt index 7e51691..d17f562 100644 --- a/src/binding/python/CMakeLists.txt +++ b/src/binding/python/CMakeLists.txt @@ -27,6 +27,7 @@ if (CMAKE_SYSTEM_NAME STREQUAL "Linux") $ $ $ + $ $ $ $ @@ -46,6 +47,7 @@ elseif (APPLE) -Wl,-force_load,$ -Wl,-force_load,$ -Wl,-force_load,$ + -Wl,-force_load,$ -Wl,-force_load,$ -Wl,-force_load,$ -Wl,-force_load,$ @@ -63,6 +65,7 @@ elseif (MSVC) core_knn_hnsw_static core_knn_hnsw_sparse_static core_knn_ivf_static + core_knn_vamana_static core_knn_cluster_static core_mix_reducer_static core_metric_static diff --git a/src/binding/python/model/param/python_param.cc b/src/binding/python/model/param/python_param.cc index 6ad7e1b..268246c 100644 --- a/src/binding/python/model/param/python_param.cc +++ b/src/binding/python/model/param/python_param.cc @@ -33,6 +33,8 @@ static std::string index_type_to_string(const IndexType type) { return "HNSW"; case IndexType::HNSW_RABITQ: return "HNSW_RABITQ"; + case IndexType::VAMANA: + return "VAMANA"; default: return "UNDEFINED"; } @@ -325,24 +327,31 @@ Examples: ... metric_type=MetricType.COSINE, ... m=16, ... ef_construction=200, - ... quantize_type=QuantizeType.INT8 + ... quantize_type=QuantizeType.INT8, + ... use_contiguous_memory=True, ... ) >>> print(params) - {'metric_type': 'IP', 'm': 16, 'ef_construction': 200, 'quantize_type': 'INT8'} + {'metric_type': 'IP', 'm': 16, 'ef_construction': 200, 'quantize_type': 'INT8', 'use_contiguous_memory': True} )pbdoc"); hnsw_params - .def(py::init(), + .def(py::init(), py::arg("metric_type") = MetricType::IP, py::arg("m") = core_interface::kDefaultHnswNeighborCnt, py::arg("ef_construction") = core_interface::kDefaultHnswEfConstruction, - py::arg("quantize_type") = QuantizeType::UNDEFINED) + py::arg("quantize_type") = QuantizeType::UNDEFINED, + py::arg("use_contiguous_memory") = false) .def_property_readonly( "m", &HnswIndexParams::m, "int: Maximum number of neighbors per node in upper layers.") .def_property_readonly( "ef_construction", &HnswIndexParams::ef_construction, "int: Candidate list size during index construction.") + .def_property_readonly( + "use_contiguous_memory", &HnswIndexParams::use_contiguous_memory, + "bool: Whether to allocate a single contiguous memory arena for " + "all HNSW graph nodes. Improves cache locality and search " + "throughput at the cost of peak memory usage. Defaults to False.") .def( "to_dict", [](const HnswIndexParams &self) -> py::dict { @@ -353,6 +362,7 @@ Examples: dict["ef_construction"] = self.ef_construction(); dict["quantize_type"] = quantize_type_to_string(self.quantize_type()); + dict["use_contiguous_memory"] = self.use_contiguous_memory(); return dict; }, "Convert to dictionary with all fields") @@ -365,19 +375,22 @@ Examples: ", \"ef_construction\":" + std::to_string(self.ef_construction()) + ", \"quantize_type\":" + - quantize_type_to_string(self.quantize_type()) + "}"; + quantize_type_to_string(self.quantize_type()) + + ", \"use_contiguous_memory\":" + + (self.use_contiguous_memory() ? "true" : "false") + "}"; }) .def(py::pickle( [](const HnswIndexParams &self) { return py::make_tuple(self.metric_type(), self.m(), - self.ef_construction(), self.quantize_type()); + self.ef_construction(), self.quantize_type(), + self.use_contiguous_memory()); }, [](py::tuple t) { - if (t.size() != 4) + if (t.size() != 5) throw std::runtime_error("Invalid state for HnswIndexParams"); return std::make_shared( t[0].cast(), t[1].cast(), t[2].cast(), - t[3].cast()); + t[3].cast(), t[4].cast()); })); // binding hnsw rabitq index params @@ -479,6 +492,141 @@ Examples: t[3].cast(), t[4].cast(), t[5].cast()); })); + // binding vamana index params + py::class_> + vamana_params(m, "VamanaIndexParam", R"pbdoc( +Parameters for configuring a Vamana (DiskANN) index. + +Vamana is a single-layer graph-based approximate nearest neighbor search +index originally proposed in the DiskANN paper. This class encapsulates +its construction hyperparameters. + +Attributes: + metric_type (MetricType): Distance metric used for similarity computation. + Default is ``MetricType.IP`` (inner product). + max_degree (int): Maximum out-degree (R) of every node in the Vamana + graph. Higher values improve recall but increase memory usage and + construction time. Default is 64. + search_list_size (int): Size of the dynamic candidate list during graph + construction (analogous to HNSW's ef_construction). Larger values + yield better graph quality at the cost of slower build time. + Default is 100. + alpha (float): Pruning factor used by Vamana's RobustPrune. Values > 1.0 + keep more long-range edges and improve recall on hard datasets. + Default is 1.2. + saturate_graph (bool): If True, force every node to reach max_degree + neighbors during construction. Default is False. + use_contiguous_memory (bool): If True, allocate a single contiguous + memory arena for all graph nodes, improving cache locality and + search throughput at the cost of peak memory usage. Default is + False. + use_id_map (bool): Reserved flag for engine-level id remapping; the + db layer always supplies consecutive ids so this is currently + ignored by the engine. Default is False. + quantize_type (QuantizeType): Optional quantization type for vector + compression (e.g., FP16, INT8). Default is ``QuantizeType.UNDEFINED`` + to disable quantization. + +Examples: + >>> from zvec.typing import MetricType, QuantizeType + >>> params = VamanaIndexParam( + ... metric_type=MetricType.COSINE, + ... max_degree=64, + ... search_list_size=128, + ... alpha=1.2, + ... quantize_type=QuantizeType.INT8, + ... ) +)pbdoc"); + vamana_params + .def(py::init(), + py::arg("metric_type") = MetricType::IP, + py::arg("max_degree") = core_interface::kDefaultVamanaMaxDegree, + py::arg("search_list_size") = + core_interface::kDefaultVamanaSearchListSize, + py::arg("alpha") = core_interface::kDefaultVamanaAlpha, + py::arg("saturate_graph") = + core_interface::kDefaultVamanaSaturateGraph, + py::arg("use_contiguous_memory") = false, + py::arg("use_id_map") = false, + py::arg("quantize_type") = QuantizeType::UNDEFINED) + .def_property_readonly( + "max_degree", &VamanaIndexParams::max_degree, + "int: Maximum out-degree (R) of every node in the Vamana graph.") + .def_property_readonly( + "search_list_size", &VamanaIndexParams::search_list_size, + "int: Candidate list size during Vamana graph construction.") + .def_property_readonly("alpha", &VamanaIndexParams::alpha, + "float: Vamana RobustPrune alpha factor.") + .def_property_readonly( + "saturate_graph", &VamanaIndexParams::saturate_graph, + "bool: Whether to saturate every node to max_degree neighbors.") + .def_property_readonly( + "use_contiguous_memory", &VamanaIndexParams::use_contiguous_memory, + "bool: Whether to allocate a single contiguous memory arena for " + "all Vamana graph nodes. Improves cache locality and search " + "throughput at the cost of peak memory usage. Defaults to False.") + .def_property_readonly( + "use_id_map", &VamanaIndexParams::use_id_map, + "bool: Reserved flag for engine-level id remapping. Currently " + "ignored by the engine because the db layer always supplies " + "consecutive ids.") + .def( + "to_dict", + [](const VamanaIndexParams &self) -> py::dict { + py::dict dict; + dict["type"] = index_type_to_string(self.type()); + dict["metric_type"] = metric_type_to_string(self.metric_type()); + dict["max_degree"] = self.max_degree(); + dict["search_list_size"] = self.search_list_size(); + dict["alpha"] = self.alpha(); + dict["saturate_graph"] = self.saturate_graph(); + dict["use_contiguous_memory"] = self.use_contiguous_memory(); + dict["use_id_map"] = self.use_id_map(); + dict["quantize_type"] = + quantize_type_to_string(self.quantize_type()); + return dict; + }, + "Convert to dictionary with all fields") + .def("__repr__", + [](const VamanaIndexParams &self) -> std::string { + return "{" + "\"type\":\"" + + index_type_to_string(self.type()) + + "\", \"metric_type\":\"" + + metric_type_to_string(self.metric_type()) + + "\", \"max_degree\":" + std::to_string(self.max_degree()) + + ", \"search_list_size\":" + + std::to_string(self.search_list_size()) + + ", \"alpha\":" + std::to_string(self.alpha()) + + ", \"saturate_graph\":" + + std::string(self.saturate_graph() ? "true" : "false") + + ", \"use_contiguous_memory\":" + + std::string(self.use_contiguous_memory() ? "true" + : "false") + + ", \"use_id_map\":" + + std::string(self.use_id_map() ? "true" : "false") + + ", \"quantize_type\":\"" + + quantize_type_to_string(self.quantize_type()) + "\"}"; + }) + .def(py::pickle( + [](const VamanaIndexParams &self) { + return py::make_tuple(self.metric_type(), self.max_degree(), + self.search_list_size(), self.alpha(), + self.saturate_graph(), + self.use_contiguous_memory(), + self.use_id_map(), self.quantize_type()); + }, + [](py::tuple t) { + if (t.size() != 8) + throw std::runtime_error("Invalid state for VamanaIndexParams"); + return std::make_shared( + t[0].cast(), t[1].cast(), t[2].cast(), + t[3].cast(), t[4].cast(), t[5].cast(), + t[6].cast(), t[7].cast()); + })); + // FlatIndexParams py::class_> @@ -884,6 +1032,76 @@ Args: obj->set_is_using_refiner(t[3].cast()); return obj; })); + + // binding vamana query params + py::class_> + vamana_query_params(m, "VamanaQueryParam", R"pbdoc( +Query parameters for the Vamana (DiskANN) index. + +Controls the trade-off between search speed and accuracy via the +``ef_search`` parameter, which sets the size of the dynamic candidate list +explored during search. + +Attributes: + type (IndexType): Always ``IndexType.VAMANA``. + ef_search (int): Size of the dynamic candidate list during Vamana + search. Larger values improve recall but slow down search. + Default is 200. + radius (float): Search radius for range queries. Default is 0.0. + is_linear (bool): Force linear search. Default is False. + is_using_refiner (bool, optional): Whether to use refiner for the query. + Default is False. + +Examples: + >>> params = VamanaQueryParam(ef_search=200) + >>> print(params.ef_search) + 200 +)pbdoc"); + vamana_query_params + .def(py::init(), + py::arg("ef_search") = core_interface::kDefaultVamanaEfSearch, + py::arg("radius") = 0.0f, py::arg("is_linear") = false, + py::arg("is_using_refiner") = false, + R"pbdoc( +Constructs a VamanaQueryParam instance. + +Args: + ef_search (int, optional): Search-time candidate list size. + Higher values improve accuracy. Defaults to 200. + radius (float, optional): Search radius for range queries. Default is 0.0. + is_linear (bool, optional): Force linear search. Default is False. + is_using_refiner (bool, optional): Whether to use refiner for the query. + Default is False. +)pbdoc") + .def_property_readonly( + "ef_search", + [](const VamanaQueryParams &self) -> int { return self.ef_search(); }, + "int: Size of the dynamic candidate list during Vamana search.") + .def("__repr__", + [](const VamanaQueryParams &self) -> std::string { + return "{" + "\"type\":\"" + + index_type_to_string(self.type()) + + "\", \"ef_search\":" + std::to_string(self.ef_search()) + + ", \"radius\":" + std::to_string(self.radius()) + + ", \"is_linear\":" + std::to_string(self.is_linear()) + + ", \"is_using_refiner\":" + + std::to_string(self.is_using_refiner()) + "}"; + }) + .def(py::pickle( + [](const VamanaQueryParams &self) { + return py::make_tuple(self.ef_search(), self.radius(), + self.is_linear(), self.is_using_refiner()); + }, + [](py::tuple t) { + if (t.size() != 4) + throw std::runtime_error("Invalid state for VamanaQueryParams"); + auto obj = std::make_shared(t[0].cast()); + obj->set_radius(t[1].cast()); + obj->set_is_linear(t[2].cast()); + obj->set_is_using_refiner(t[3].cast()); + return obj; + })); } void ZVecPyParams::bind_options(py::module_ &m) { // binding collection options diff --git a/src/binding/python/model/python_collection.cc b/src/binding/python/model/python_collection.cc index 1ea32ac..2ced832 100644 --- a/src/binding/python/model/python_collection.cc +++ b/src/binding/python/model/python_collection.cc @@ -211,7 +211,19 @@ void ZVecPyCollection::bind_dql_methods( const auto result = self.Fetch(pks); // return DocPtrMap return unwrap_expected(result); - }); + }) + .def( + "_debug_hnsw_storage_mode", + [](const Collection &self, const std::string &column_name) { + const auto result = self.DebugGetHnswStorageMode(column_name); + return unwrap_expected(result); + }, + py::arg("column_name"), + "Debug-only: returns the storage mode of the HNSW entity on the " + "given vector column. One of 'mmap', 'buffer_pool', 'contiguous'. " + "Raises KeyError if no HNSW index exists on the column, or " + "ValueError if the column's index is not an HNSW index. Intended " + "for introspection and testing only; not part of the stable API."); } } // namespace zvec \ No newline at end of file diff --git a/src/binding/python/typing/python_type.cc b/src/binding/python/typing/python_type.cc index bb50034..e6b9bd7 100644 --- a/src/binding/python/typing/python_type.cc +++ b/src/binding/python/typing/python_type.cc @@ -99,6 +99,7 @@ Examples: .value("HNSW_RABITQ", IndexType::HNSW_RABITQ) .value("IVF", IndexType::IVF) .value("FLAT", IndexType::FLAT) + .value("VAMANA", IndexType::VAMANA) .value("INVERT", IndexType::INVERT); } diff --git a/src/core/algorithm/CMakeLists.txt b/src/core/algorithm/CMakeLists.txt index 20a4590..9a30e20 100644 --- a/src/core/algorithm/CMakeLists.txt +++ b/src/core/algorithm/CMakeLists.txt @@ -7,6 +7,7 @@ cc_directory(flat_sparse) cc_directory(ivf) cc_directory(hnsw) cc_directory(hnsw_sparse) +cc_directory(vamana) if(RABITQ_SUPPORTED) message(STATUS "BUILD RABITQ") cc_directory(hnsw_rabitq) diff --git a/src/core/algorithm/hnsw/hnsw_algorithm.cc b/src/core/algorithm/hnsw/hnsw_algorithm.cc index fa553f5..75fe5a9 100644 --- a/src/core/algorithm/hnsw/hnsw_algorithm.cc +++ b/src/core/algorithm/hnsw/hnsw_algorithm.cc @@ -12,28 +12,15 @@ // See the License for the specific language governing permissions and // limitations under the License. #include "hnsw_algorithm.h" -#include -#include -#include -#include namespace zvec { namespace core { -HnswAlgorithm::HnswAlgorithm(HnswEntity &entity) - : entity_(entity), - mt_(std::chrono::system_clock::now().time_since_epoch().count()), - lock_pool_(kLockCnt) {} - -int HnswAlgorithm::cleanup() { - return 0; -} - -int HnswAlgorithm::add_node(node_id_t id, level_t level, HnswContext *ctx) { +template +int HnswAlgorithm::add_node(node_id_t id, level_t level, + HnswContext *ctx) { spin_lock_.lock(); - // std::cout << "id: " << id << ", level: " << level << std::endl; - auto cur_max_level = entity_.cur_max_level(); auto entry_point = entity_.entry_point(); if (ailego_unlikely(entry_point == kInvalidNodeId)) { @@ -54,7 +41,7 @@ int HnswAlgorithm::add_node(node_id_t id, level_t level, HnswContext *ctx) { } level_t cur_level = cur_max_level; - dist_t dist = ctx->dist_calculator()(entry_point); + dist_t dist = ctx->dist_calculator().batch_dist(entry_point); for (; cur_level > level; --cur_level) { select_entry_point(cur_level, &entry_point, &dist, ctx); } @@ -81,7 +68,8 @@ int HnswAlgorithm::add_node(node_id_t id, level_t level, HnswContext *ctx) { return 0; } -int HnswAlgorithm::search(HnswContext *ctx) const { +template +int HnswAlgorithm::search(HnswContext *ctx) const { spin_lock_.lock(); auto maxLevel = entity_.cur_max_level(); auto entry_point = entity_.entry_point(); @@ -107,13 +95,15 @@ int HnswAlgorithm::search(HnswContext *ctx) const { return 0; } -//! select_entry_point on hnsw level, ef = 1 -void HnswAlgorithm::select_entry_point(level_t level, node_id_t *entry_point, - dist_t *dist, HnswContext *ctx) const { - auto &entity = ctx->get_entity(); +template +void HnswAlgorithm::select_entry_point(level_t level, + node_id_t *entry_point, + dist_t *dist, + HnswContext *ctx) const { + const auto &entity = static_cast(ctx->get_entity()); HnswDistCalculator &dc = ctx->dist_calculator(); while (true) { - const Neighbors neighbors = entity.get_neighbors(level, *entry_point); + const auto neighbors = entity.get_neighbors_typed(level, *entry_point); if (ailego_unlikely(ctx->debugging())) { (*ctx->mutable_stats_get_neighbors())++; } @@ -122,8 +112,8 @@ void HnswAlgorithm::select_entry_point(level_t level, node_id_t *entry_point, break; } - std::vector neighbor_vec_blocks; - int ret = entity.get_vector(&neighbors[0], size, neighbor_vec_blocks); + std::vector neighbor_vec_blocks; + int ret = entity.get_vector_typed(&neighbors[0], size, neighbor_vec_blocks); if (ailego_unlikely(ctx->debugging())) { (*ctx->mutable_stats_get_vector())++; } @@ -159,8 +149,10 @@ void HnswAlgorithm::select_entry_point(level_t level, node_id_t *entry_point, return; } -void HnswAlgorithm::add_neighbors(node_id_t id, level_t level, - TopkHeap &topk_heap, HnswContext *ctx) { +template +void HnswAlgorithm::add_neighbors(node_id_t id, level_t level, + TopkHeap &topk_heap, + HnswContext *ctx) { if (ailego_unlikely(topk_heap.size() == 0)) { return; } @@ -178,16 +170,20 @@ void HnswAlgorithm::add_neighbors(node_id_t id, level_t level, return; } -void HnswAlgorithm::search_neighbors(level_t level, node_id_t *entry_point, - dist_t *dist, TopkHeap &topk, - HnswContext *ctx) const { - const auto &entity = ctx->get_entity(); +template +void HnswAlgorithm::search_neighbors(level_t level, + node_id_t *entry_point, + dist_t *dist, TopkHeap &topk, + HnswContext *ctx) const { + const auto &entity = static_cast(ctx->get_entity()); HnswDistCalculator &dc = ctx->dist_calculator(); VisitFilter &visit = ctx->visit_filter(); CandidateHeap &candidates = ctx->candidates(); std::function filter = [](node_id_t) { return false; }; if (ctx->filter().is_valid()) { - filter = [&](node_id_t id) { return ctx->filter()(entity.get_key(id)); }; + filter = [&](node_id_t id) { + return ctx->filter()(entity.get_key_typed(id)); + }; } candidates.clear(); @@ -208,7 +204,7 @@ void HnswAlgorithm::search_neighbors(level_t level, node_id_t *entry_point, } candidates.pop(); - const Neighbors neighbors = entity.get_neighbors(level, main_node); + const auto neighbors = entity.get_neighbors_typed(level, main_node); ailego_prefetch(neighbors.data); if (ailego_unlikely(ctx->debugging())) { (*ctx->mutable_stats_get_neighbors())++; @@ -231,8 +227,9 @@ void HnswAlgorithm::search_neighbors(level_t level, node_id_t *entry_point, continue; } - std::vector neighbor_vec_blocks; - int ret = entity.get_vector(neighbor_ids.data(), size, neighbor_vec_blocks); + std::vector neighbor_vec_blocks; + int ret = + entity.get_vector_typed(neighbor_ids.data(), size, neighbor_vec_blocks); if (ailego_unlikely(ctx->debugging())) { (*ctx->mutable_stats_get_vector())++; } @@ -278,15 +275,16 @@ void HnswAlgorithm::search_neighbors(level_t level, node_id_t *entry_point, return; } -void HnswAlgorithm::expand_neighbors_by_group(TopkHeap &topk, - HnswContext *ctx) const { +template +void HnswAlgorithm::expand_neighbors_by_group( + TopkHeap &topk, HnswContext *ctx) const { if (!ctx->group_by().is_valid()) { return; } - const auto &entity = ctx->get_entity(); + const auto &entity = static_cast(ctx->get_entity()); std::function group_by = [&](node_id_t id) { - return ctx->group_by()(entity.get_key(id)); + return ctx->group_by()(entity.get_key_typed(id)); }; // devide into groups @@ -312,7 +310,9 @@ void HnswAlgorithm::expand_neighbors_by_group(TopkHeap &topk, std::function filter = [](node_id_t) { return false; }; if (ctx->filter().is_valid()) { - filter = [&](node_id_t id) { return ctx->filter()(entity.get_key(id)); }; + filter = [&](node_id_t id) { + return ctx->filter()(entity.get_key_typed(id)); + }; } // refill to get enough groups @@ -332,7 +332,7 @@ void HnswAlgorithm::expand_neighbors_by_group(TopkHeap &topk, node_id_t main_node = top->first; candidates.pop(); - const Neighbors neighbors = entity.get_neighbors(0, main_node); + const auto neighbors = entity.get_neighbors_typed(0, main_node); if (ailego_unlikely(ctx->debugging())) { (*ctx->mutable_stats_get_neighbors())++; } @@ -354,9 +354,9 @@ void HnswAlgorithm::expand_neighbors_by_group(TopkHeap &topk, continue; } - std::vector neighbor_vec_blocks; - int ret = - entity.get_vector(neighbor_ids.data(), size, neighbor_vec_blocks); + std::vector neighbor_vec_blocks; + int ret = entity.get_vector_typed(neighbor_ids.data(), size, + neighbor_vec_blocks); if (ailego_unlikely(ctx->debugging())) { (*ctx->mutable_stats_get_vector())++; } @@ -364,14 +364,16 @@ void HnswAlgorithm::expand_neighbors_by_group(TopkHeap &topk, break; } - static constexpr node_id_t PREFETCH_STEP = 2; + std::vector dists(size); + std::vector neighbor_vecs(size); + for (uint32_t i = 0; i < size; ++i) { + neighbor_vecs[i] = neighbor_vec_blocks[i].data(); + } + dc.batch_dist(neighbor_vecs.data(), size, dists.data()); + for (uint32_t i = 0; i < size; ++i) { node_id_t node = neighbor_ids[i]; - node_id_t prefetch_id = i + PREFETCH_STEP; - if (prefetch_id < size) { - ailego_prefetch(neighbor_vec_blocks[prefetch_id].data()); - } - dist_t cur_dist = dc.dist(neighbor_vec_blocks[i].data()); + dist_t cur_dist = dists[i]; if (!filter(node)) { std::string group_id = group_by(node); @@ -393,8 +395,10 @@ void HnswAlgorithm::expand_neighbors_by_group(TopkHeap &topk, } // end if } -void HnswAlgorithm::update_neighbors(HnswDistCalculator &dc, node_id_t id, - level_t level, TopkHeap &topk_heap) { +template +void HnswAlgorithm::update_neighbors(HnswDistCalculator &dc, + node_id_t id, level_t level, + TopkHeap &topk_heap) { topk_heap.sort(); uint32_t max_neighbor_cnt = entity_.neighbor_cnt(level); @@ -455,10 +459,10 @@ void HnswAlgorithm::update_neighbors(HnswDistCalculator &dc, node_id_t id, return; } -void HnswAlgorithm::reverse_update_neighbors(HnswDistCalculator &dc, - node_id_t id, level_t level, - node_id_t link_id, dist_t dist, - TopkHeap &update_heap) { +template +void HnswAlgorithm::reverse_update_neighbors( + HnswDistCalculator &dc, node_id_t id, level_t level, node_id_t link_id, + dist_t dist, TopkHeap &update_heap) { const size_t max_neighbor_cnt = entity_.neighbor_cnt(level); uint32_t lock_idx = id & kLockMask; @@ -516,5 +520,10 @@ void HnswAlgorithm::reverse_update_neighbors(HnswDistCalculator &dc, return; } +// Explicit template instantiation +template class HnswAlgorithm; +template class HnswAlgorithm; +template class HnswAlgorithm; + } // namespace core } // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_algorithm.h b/src/core/algorithm/hnsw/hnsw_algorithm.h index 886d870..ea73cd7 100644 --- a/src/core/algorithm/hnsw/hnsw_algorithm.h +++ b/src/core/algorithm/hnsw/hnsw_algorithm.h @@ -14,41 +14,64 @@ #pragma once #include +#include +#include +#include #include #include "hnsw_context.h" #include "hnsw_dist_calculator.h" #include "hnsw_entity.h" +#include "hnsw_streamer_entity.h" namespace zvec { namespace core { -//! hnsw graph algorithm implement -class HnswAlgorithm { +//! Non-template base class for HnswAlgorithm +class HnswAlgorithmBase { public: - typedef std::unique_ptr UPointer; + typedef std::unique_ptr UPointer; + virtual ~HnswAlgorithmBase() = default; + + virtual int cleanup() = 0; + virtual int add_node(node_id_t id, level_t level, HnswContext *ctx) = 0; + virtual int search(HnswContext *ctx) const = 0; + virtual int init() = 0; + virtual uint32_t get_random_level() const = 0; +}; + +//! hnsw graph algorithm implement, templated on EntityType +template +class HnswAlgorithm : public HnswAlgorithmBase { public: + using MemBlockType = typename EntityType::MemoryBlock; + //! Constructor - explicit HnswAlgorithm(HnswEntity &entity); + explicit HnswAlgorithm(EntityType &entity) + : entity_(entity), + mt_(std::chrono::system_clock::now().time_since_epoch().count()), + lock_pool_(kLockCnt) {} //! Destructor - ~HnswAlgorithm() = default; + ~HnswAlgorithm() override = default; //! Cleanup HnswAlgorithm - int cleanup(); + int cleanup() override { + return 0; + } //! Add a node to hnsw graph //! @id: the node unique id //! @level: a node will be add to graph in each level [0, level] //! return 0 on success, or errCode in failure - int add_node(node_id_t id, level_t level, HnswContext *ctx); + int add_node(node_id_t id, level_t level, HnswContext *ctx) override; //! do knn search in graph //! return 0 on success, or errCode in failure. results saved in ctx - int search(HnswContext *ctx) const; + int search(HnswContext *ctx) const override; //! Initiate HnswAlgorithm - int init() { + int init() override { level_probas_.clear(); double level_mult = 1 / std::log(static_cast(entity_.scaling_factor())); @@ -67,7 +90,7 @@ class HnswAlgorithm { //! Generate a random level //! return graph level - uint32_t get_random_level() const { + uint32_t get_random_level() const override { // gen rand float (0, 1) double f = mt_() / static_cast(mt_.max()); for (size_t level = 0; level < level_probas_.size(); level++) { @@ -116,7 +139,7 @@ class HnswAlgorithm { static constexpr uint32_t kLockCnt{1U << 8}; static constexpr uint32_t kLockMask{kLockCnt - 1U}; - HnswEntity &entity_; + EntityType &entity_; mutable std::mt19937 mt_{}; std::vector level_probas_{}; diff --git a/src/core/algorithm/hnsw/hnsw_chunk.h b/src/core/algorithm/hnsw/hnsw_chunk.h index 0483c7b..5d96d37 100644 --- a/src/core/algorithm/hnsw/hnsw_chunk.h +++ b/src/core/algorithm/hnsw/hnsw_chunk.h @@ -42,6 +42,7 @@ class ChunkBroker { CHUNK_TYPE_UPPER_NEIGHBOR = 4, CHUNK_TYPE_NEIGHBOR_INDEX = 5, CHUNK_TYPE_SPARSE_NODE = 6, + CHUNK_TYPE_NEIGHBOR_DIST = 7, // Vamana: per-node neighbor distances CHUNK_TYPE_MAX = 8 }; static constexpr size_t kDefaultChunkSeqId = 0UL; diff --git a/src/core/algorithm/hnsw/hnsw_entity.h b/src/core/algorithm/hnsw/hnsw_entity.h index 74b959e..c8d052f 100644 --- a/src/core/algorithm/hnsw/hnsw_entity.h +++ b/src/core/algorithm/hnsw/hnsw_entity.h @@ -171,6 +171,142 @@ struct Neighbors { IndexStorage::MemoryBlock neighbor_block; }; +//! Lightweight MemoryBlock for mmap mode: zero-cost construction/destruction +struct MmapMemoryBlock { + MmapMemoryBlock() = default; + explicit MmapMemoryBlock(void *data) : data_(data) {} + + MmapMemoryBlock(const MmapMemoryBlock &) = default; + MmapMemoryBlock &operator=(const MmapMemoryBlock &) = default; + MmapMemoryBlock(MmapMemoryBlock &&) = default; + MmapMemoryBlock &operator=(MmapMemoryBlock &&) = default; + ~MmapMemoryBlock() = default; + + const void *data() const { + return data_; + } + + void reset(void *data) { + data_ = data; + } + + void *data_{nullptr}; +}; + +//! Lightweight MemoryBlock for buffer pool mode: release on destruction +struct BufferPoolMemoryBlock { + BufferPoolMemoryBlock() = default; + + BufferPoolMemoryBlock(ailego::VecBufferPoolHandle *handle, size_t block_id, + void *data) + : buffer_pool_handle_(handle), buffer_block_id_(block_id), data_(data) {} + + BufferPoolMemoryBlock(const BufferPoolMemoryBlock &rhs) + : buffer_pool_handle_(rhs.buffer_pool_handle_), + buffer_block_id_(rhs.buffer_block_id_), + data_(rhs.data_) { + if (buffer_pool_handle_) { + buffer_pool_handle_->acquire_one(buffer_block_id_); + } + } + + BufferPoolMemoryBlock &operator=(const BufferPoolMemoryBlock &rhs) { + if (this != &rhs) { + release(); + buffer_pool_handle_ = rhs.buffer_pool_handle_; + buffer_block_id_ = rhs.buffer_block_id_; + data_ = rhs.data_; + if (buffer_pool_handle_) { + buffer_pool_handle_->acquire_one(buffer_block_id_); + } + } + return *this; + } + + BufferPoolMemoryBlock(BufferPoolMemoryBlock &&rhs) noexcept + : buffer_pool_handle_(rhs.buffer_pool_handle_), + buffer_block_id_(rhs.buffer_block_id_), + data_(rhs.data_) { + rhs.buffer_pool_handle_ = nullptr; + rhs.data_ = nullptr; + } + + BufferPoolMemoryBlock &operator=(BufferPoolMemoryBlock &&rhs) noexcept { + if (this != &rhs) { + release(); + buffer_pool_handle_ = rhs.buffer_pool_handle_; + buffer_block_id_ = rhs.buffer_block_id_; + data_ = rhs.data_; + rhs.buffer_pool_handle_ = nullptr; + rhs.data_ = nullptr; + } + return *this; + } + + ~BufferPoolMemoryBlock() { + release(); + } + + const void *data() const { + return data_; + } + + void reset(ailego::VecBufferPoolHandle *handle, size_t block_id, void *data) { + release(); + buffer_pool_handle_ = handle; + buffer_block_id_ = block_id; + data_ = data; + } + + private: + void release() { + if (buffer_pool_handle_) { + buffer_pool_handle_->release_one(buffer_block_id_); + buffer_pool_handle_ = nullptr; + } + data_ = nullptr; + } + + ailego::VecBufferPoolHandle *buffer_pool_handle_{nullptr}; + size_t buffer_block_id_{0}; + void *data_{nullptr}; +}; + +//! Typed Neighbors: holds a typed MemoryBlock to avoid runtime branching +template +struct NeighborsT { + NeighborsT() : cnt{0}, data{nullptr} {} + + NeighborsT(uint32_t cnt_in, const node_id_t *data_in) + : cnt{cnt_in}, data{data_in} {} + + explicit NeighborsT(const MemBlockType &mem_block) + : neighbor_block{mem_block} { + auto hd = reinterpret_cast(neighbor_block.data()); + cnt = hd->neighbor_cnt; + data = hd->neighbors; + } + + explicit NeighborsT(MemBlockType &&mem_block) + : neighbor_block{std::move(mem_block)} { + auto hd = reinterpret_cast(neighbor_block.data()); + cnt = hd->neighbor_cnt; + data = hd->neighbors; + } + + size_t size(void) const { + return cnt; + } + + const node_id_t &operator[](size_t idx) const { + return data[idx]; + } + + uint32_t cnt; + const node_id_t *data; + MemBlockType neighbor_block; +}; + //! level 0 neighbors offset struct GraphNeighborMeta { GraphNeighborMeta(size_t o, size_t cnt) : offset(o), neighbor_cnt(cnt) {} diff --git a/src/core/algorithm/hnsw/hnsw_params.h b/src/core/algorithm/hnsw/hnsw_params.h index a7b7c92..4caa148 100644 --- a/src/core/algorithm/hnsw/hnsw_params.h +++ b/src/core/algorithm/hnsw/hnsw_params.h @@ -108,5 +108,8 @@ static const std::string PARAM_HNSW_REDUCER_INDEX_NAME( static const std::string PARAM_HNSW_REDUCER_EFCONSTRUCTION( "proxima.hnsw.reducer.efconstruction"); +static const std::string PARAM_HNSW_STREAMER_USE_CONTIGUOUS_MEMORY( + "proxima.hnsw.streamer.use_contiguous_memory"); + } // namespace core } // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_streamer.cc b/src/core/algorithm/hnsw/hnsw_streamer.cc index 82319e2..935cae5 100644 --- a/src/core/algorithm/hnsw/hnsw_streamer.cc +++ b/src/core/algorithm/hnsw/hnsw_streamer.cc @@ -25,7 +25,7 @@ namespace zvec { namespace core { -HnswStreamer::HnswStreamer() : entity_(stats_) {} +HnswStreamer::HnswStreamer() = default; HnswStreamer::~HnswStreamer() { if (state_ == STATE_INITED) { @@ -68,7 +68,8 @@ int HnswStreamer::init(const IndexMeta &imeta, const ailego::Params ¶ms) { params.get(PARAM_HNSW_STREAMER_FORCE_PADDING_RESULT_ENABLE, &force_padding_topk_enabled_); params.get(PARAM_HNSW_STREAMER_USE_ID_MAP, &use_id_map_); - entity_.set_use_key_info_map(use_id_map_); + params.get(PARAM_HNSW_STREAMER_USE_CONTIGUOUS_MEMORY, + &use_contiguous_memory_); params.get(PARAM_HNSW_STREAMER_DOCS_SOFT_LIMIT, &docs_soft_limit_); if (docs_soft_limit_ > 0 && docs_soft_limit_ > docs_hard_limit_) { @@ -143,6 +144,7 @@ int HnswStreamer::init(const IndexMeta &imeta, const ailego::Params ¶ms) { if (prune_cnt == 0UL) { prune_cnt = upper_max_neighbor_cnt_; } + prune_cnt_ = prune_cnt; if (chunk_size_ == 0UL) { chunk_size_ = HnswEntity::kDefaultChunkSize; } @@ -152,47 +154,21 @@ int HnswStreamer::init(const IndexMeta &imeta, const ailego::Params ¶ms) { return IndexError_InvalidArgument; } - entity_.set_ef_construction(ef_construction_); - entity_.set_upper_neighbor_cnt(upper_max_neighbor_cnt_); - entity_.set_l0_neighbor_cnt(l0_max_neighbor_cnt_); - entity_.set_scaling_factor(scaling_factor_); - entity_.set_prune_cnt(prune_cnt); - - entity_.set_vector_size(meta_.element_size()); - - entity_.set_chunk_size(chunk_size_); - entity_.set_filter_same_key(filter_same_key_); - entity_.set_get_vector(get_vector_enabled_); - entity_.set_min_neighbor_cnt(min_neighbor_cnt_); - - int ret = entity_.init(docs_hard_limit_); - if (ret != 0) { - LOG_ERROR("Hnsw entity init failed for %s", IndexError::What(ret)); - return ret; - } - LOG_DEBUG( "Init params: maxIndexSize=%zu docsHardLimit=%zu docsSoftLimit=%zu " "efConstruction=%u ef=%u upperMaxNeighborCnt=%u l0MaxNeighborCnt=%u " "scalingFactor=%u maxScanRatio=%.3f minScanLimit=%zu maxScanLimit=%zu " "bfEnabled=%d bruteFoceThreshold=%zu bfNegativeProbability=%.5f " - "checkCrcEnabled=%d pruneSize=%zu vectorSize=%u chunkSize=%zu " + "checkCrcEnabled=%d pruneSize=%u vectorSize=%u chunkSize=%zu " "filterSameKey=%u getVectorEnabled=%u minNeighborCount=%u " "forcePadding=%u ", max_index_size_, docs_hard_limit_, docs_soft_limit_, ef_construction_, ef_, upper_max_neighbor_cnt_, l0_max_neighbor_cnt_, scaling_factor_, max_scan_ratio_, min_scan_limit_, max_scan_limit_, bf_enabled_, - bruteforce_threshold_, bf_negative_prob_, check_crc_enabled_, prune_cnt, + bruteforce_threshold_, bf_negative_prob_, check_crc_enabled_, prune_cnt_, meta_.element_size(), chunk_size_, filter_same_key_, get_vector_enabled_, min_neighbor_cnt_, force_padding_topk_enabled_); - alg_ = HnswAlgorithm::UPointer(new HnswAlgorithm(entity_)); - - ret = alg_->init(); - if (ret != 0) { - return ret; - } - state_ = STATE_INITED; return 0; @@ -208,7 +184,9 @@ int HnswStreamer::cleanup(void) { meta_.clear(); metric_.reset(); stats_.clear(); - entity_.cleanup(); + if (entity_) { + entity_->cleanup(); + } if (alg_) { alg_->cleanup(); @@ -237,6 +215,26 @@ int HnswStreamer::cleanup(void) { return 0; } +int HnswStreamer::setup_entity() { + entity_->set_use_key_info_map(use_id_map_); + entity_->set_ef_construction(ef_construction_); + entity_->set_upper_neighbor_cnt(upper_max_neighbor_cnt_); + entity_->set_l0_neighbor_cnt(l0_max_neighbor_cnt_); + entity_->set_scaling_factor(scaling_factor_); + entity_->set_prune_cnt(prune_cnt_); + entity_->set_vector_size(meta_.element_size()); + entity_->set_chunk_size(chunk_size_); + entity_->set_filter_same_key(filter_same_key_); + entity_->set_get_vector(get_vector_enabled_); + entity_->set_min_neighbor_cnt(min_neighbor_cnt_); + + int ret = entity_->init(docs_hard_limit_); + if (ret != 0) { + LOG_ERROR("Hnsw entity init failed for %s", IndexError::What(ret)); + } + return ret; +} + int HnswStreamer::open(IndexStorage::Pointer stg) { LOG_INFO("HnswStreamer open"); @@ -244,15 +242,36 @@ int HnswStreamer::open(IndexStorage::Pointer stg) { LOG_ERROR("Open storage failed, init streamer first!"); return IndexError_NoReady; } - int ret = entity_.open(std::move(stg), max_index_size_, check_crc_enabled_); + + // Create entity based on storage type + switch (stg->memory_block_type()) { + case IndexStorage::MemoryBlock::MBT_BUFFERPOOL: { + entity_ = std::make_unique(stats_); + break; + } + default: { + if (use_contiguous_memory_) { + entity_ = std::make_unique(stats_); + } else { + entity_ = std::make_unique(stats_); + } + break; + } + } + int ret = setup_entity(); + if (ret != 0) { + return ret; + } + + ret = entity_->open(std::move(stg), max_index_size_, check_crc_enabled_); if (ret != 0) { return ret; } IndexMeta index_meta; - ret = entity_.get_index_meta(&index_meta); + ret = entity_->get_index_meta(&index_meta); if (ret == IndexError_NoExist) { // Set IndexMeta for the new index - ret = entity_.set_index_meta(meta_); + ret = entity_->set_index_meta(meta_); if (ret != 0) { LOG_ERROR("Failed to set index meta for %s", IndexError::What(ret)); return ret; @@ -307,6 +326,36 @@ int HnswStreamer::open(IndexStorage::Pointer stg) { search_batch_distance_ = metric_->query_metric()->batch_distance(); } + // Create algorithm based on entity storage mode + switch (entity_->storage_mode()) { + case HnswStorageMode::kBufferPool: + alg_ = HnswAlgorithmBase::UPointer( + new HnswAlgorithm( + static_cast(*entity_))); + break; + case HnswStorageMode::kContiguous: { + auto &contiguous_entity = + static_cast(*entity_); + int build_ret = contiguous_entity.build_contiguous_memory(); + if (build_ret != 0) { + LOG_ERROR("Failed to build contiguous memory, ret=%d", build_ret); + return build_ret; + } + alg_ = HnswAlgorithmBase::UPointer( + new HnswAlgorithm(contiguous_entity)); + break; + } + default: + alg_ = + HnswAlgorithmBase::UPointer(new HnswAlgorithm( + static_cast(*entity_))); + break; + } + ret = alg_->init(); + if (ret != 0) { + return ret; + } + state_ = STATE_OPENED; magic_ = IndexContext::GenerateMagic(); @@ -318,8 +367,8 @@ int HnswStreamer::close(void) { stats_.clear(); meta_.set_metric(metric_->name(), 0, metric_->params()); - entity_.set_index_meta(meta_); - int ret = entity_.close(); + entity_->set_index_meta(meta_); + int ret = entity_->close(); if (ret != 0) { return ret; } @@ -332,8 +381,8 @@ int HnswStreamer::flush(uint64_t checkpoint) { LOG_INFO("HnswStreamer flush checkpoint=%zu", (size_t)checkpoint); meta_.set_metric(metric_->name(), 0, metric_->params()); - entity_.set_index_meta(meta_); - return entity_.flush(checkpoint); + entity_->set_index_meta(meta_); + return entity_->flush(checkpoint); } int HnswStreamer::dump(const IndexDumper::Pointer &dumper) { @@ -347,7 +396,7 @@ int HnswStreamer::dump(const IndexDumper::Pointer &dumper) { LOG_ERROR("Failed to serialize meta into dumper."); return ret; } - return entity_.dump(dumper); + return entity_->dump(dumper); } IndexStreamer::Context::Pointer HnswStreamer::create_context(void) const { @@ -356,7 +405,7 @@ IndexStreamer::Context::Pointer HnswStreamer::create_context(void) const { return Context::Pointer(); } - HnswEntity::Pointer entity = entity_.clone(); + HnswEntity::Pointer entity = entity_->clone(); if (ailego_unlikely(!entity)) { LOG_ERROR("CreateContext clone init failed"); return Context::Pointer(); @@ -387,9 +436,9 @@ IndexStreamer::Context::Pointer HnswStreamer::create_context(void) const { if (meta_.streamer_params().get(PARAM_HNSW_STREAMER_ESTIMATE_DOC_COUNT, &estimate_doc_count)) { LOG_DEBUG("HnswStreamer doc_count[%zu] estimate[%zu]", - (size_t)entity_.doc_cnt(), (size_t)estimate_doc_count); + (size_t)entity_->doc_cnt(), (size_t)estimate_doc_count); } - ctx->check_need_adjuct_ctx(std::max(entity_.doc_cnt(), estimate_doc_count)); + ctx->check_need_adjuct_ctx(std::max(entity_->doc_cnt(), estimate_doc_count)); return Context::Pointer(ctx); } @@ -397,7 +446,7 @@ IndexStreamer::Context::Pointer HnswStreamer::create_context(void) const { IndexProvider::Pointer HnswStreamer::create_provider(void) const { LOG_DEBUG("HnswStreamer create provider"); - auto entity = entity_.clone(); + auto entity = entity_->clone(); if (ailego_unlikely(!entity)) { LOG_ERROR("Clone HnswEntity failed"); return nullptr; @@ -407,7 +456,7 @@ IndexProvider::Pointer HnswStreamer::create_provider(void) const { } int HnswStreamer::update_context(HnswContext *ctx) const { - const HnswEntity::Pointer entity = entity_.clone(); + const HnswEntity::Pointer entity = entity_->clone(); if (!entity) { LOG_ERROR("Failed to clone search context entity"); return IndexError_Runtime; @@ -442,15 +491,15 @@ int HnswStreamer::add_with_id_impl(uint32_t id, const void *query, } } - if (ailego_unlikely(entity_.doc_cnt() >= docs_soft_limit_)) { - if (entity_.doc_cnt() >= docs_hard_limit_) { - LOG_ERROR("Current docs %u exceed [%s]", entity_.doc_cnt(), + if (ailego_unlikely(entity_->doc_cnt() >= docs_soft_limit_)) { + if (entity_->doc_cnt() >= docs_hard_limit_) { + LOG_ERROR("Current docs %u exceed [%s]", entity_->doc_cnt(), PARAM_HNSW_STREAMER_DOCS_HARD_LIMIT.c_str()); const std::lock_guard lk(mutex_); (*stats_.mutable_discarded_count())++; return IndexError_IndexFull; } else { - LOG_WARN("Current docs %u exceed [%s]", entity_.doc_cnt(), + LOG_WARN("Current docs %u exceed [%s]", entity_->doc_cnt(), PARAM_HNSW_STREAMER_DOCS_SOFT_LIMIT.c_str()); } } @@ -464,7 +513,7 @@ int HnswStreamer::add_with_id_impl(uint32_t id, const void *query, ctx->clear(); ctx->update_dist_caculator_distance(add_distance_, add_batch_distance_); ctx->reset_query(query); - ctx->check_need_adjuct_ctx(entity_.doc_cnt()); + ctx->check_need_adjuct_ctx(entity_->doc_cnt()); if (metric_->support_train()) { const std::lock_guard lk(mutex_); @@ -477,7 +526,7 @@ int HnswStreamer::add_with_id_impl(uint32_t id, const void *query, } level_t level = alg_->get_random_level(); - ret = entity_.add_vector_with_id(level, id, query); + ret = entity_->add_vector_with_id(level, id, query); if (ailego_unlikely(ret != 0)) { LOG_ERROR("Hnsw streamer add vector failed"); (*stats_.mutable_discarded_count())++; @@ -522,15 +571,15 @@ int HnswStreamer::add_impl(uint64_t pkey, const void *query, } } - if (ailego_unlikely(entity_.doc_cnt() >= docs_soft_limit_)) { - if (entity_.doc_cnt() >= docs_hard_limit_) { - LOG_ERROR("Current docs %u exceed [%s]", entity_.doc_cnt(), + if (ailego_unlikely(entity_->doc_cnt() >= docs_soft_limit_)) { + if (entity_->doc_cnt() >= docs_hard_limit_) { + LOG_ERROR("Current docs %u exceed [%s]", entity_->doc_cnt(), PARAM_HNSW_STREAMER_DOCS_HARD_LIMIT.c_str()); const std::lock_guard lk(mutex_); (*stats_.mutable_discarded_count())++; return IndexError_IndexFull; } else { - LOG_WARN("Current docs %u exceed [%s]", entity_.doc_cnt(), + LOG_WARN("Current docs %u exceed [%s]", entity_->doc_cnt(), PARAM_HNSW_STREAMER_DOCS_SOFT_LIMIT.c_str()); } } @@ -544,7 +593,7 @@ int HnswStreamer::add_impl(uint64_t pkey, const void *query, ctx->clear(); ctx->update_dist_caculator_distance(add_distance_, add_batch_distance_); ctx->reset_query(query); - ctx->check_need_adjuct_ctx(entity_.doc_cnt()); + ctx->check_need_adjuct_ctx(entity_->doc_cnt()); if (metric_->support_train()) { const std::lock_guard lk(mutex_); @@ -558,7 +607,7 @@ int HnswStreamer::add_impl(uint64_t pkey, const void *query, level_t level = alg_->get_random_level(); node_id_t id; - ret = entity_.add_vector(level, pkey, query, &id); + ret = entity_->add_vector(level, pkey, query, &id); if (ailego_unlikely(ret != 0)) { LOG_ERROR("Hnsw streamer add vector failed"); (*stats_.mutable_discarded_count())++; @@ -601,7 +650,7 @@ int HnswStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, return IndexError_Cast; } - if (entity_.doc_cnt() <= ctx->get_bruteforce_threshold()) { + if (entity_->doc_cnt() <= ctx->get_bruteforce_threshold()) { return search_bf_impl(query, qmeta, count, context); } @@ -616,7 +665,7 @@ int HnswStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, ctx->clear(); ctx->update_dist_caculator_distance(search_distance_, search_batch_distance_); ctx->resize_results(count); - ctx->check_need_adjuct_ctx(entity_.doc_cnt()); + ctx->check_need_adjuct_ctx(entity_->doc_cnt()); for (size_t q = 0; q < count; ++q) { ctx->reset_query(query); ret = alg_->search(ctx); @@ -636,11 +685,11 @@ int HnswStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, } void HnswStreamer::print_debug_info() { - for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) { - if (entity_.get_key(id) == kInvalidKey) { + for (node_id_t id = 0; id < entity_->doc_cnt(); ++id) { + if (entity_->get_key(id) == kInvalidKey) { continue; } - Neighbors neighbours = entity_.get_neighbors(0, id); + Neighbors neighbours = entity_->get_neighbors(0, id); std::cout << "node: " << id << "; "; if (neighbours.size() == 0) std::cout << std::endl; for (uint32_t i = 0; i < neighbours.size(); ++i) { @@ -654,7 +703,7 @@ void HnswStreamer::print_debug_info() { } } - // entity_.print_key_map(); + // entity_->print_key_map(); } int HnswStreamer::search_bf_impl( @@ -694,19 +743,19 @@ int HnswStreamer::search_bf_impl( } std::function group_by = [&](node_id_t id) { - return ctx->group_by()(entity_.get_key(id)); + return ctx->group_by()(entity_->get_key(id)); }; for (size_t q = 0; q < count; ++q) { ctx->reset_query(query); ctx->group_topk_heaps().clear(); - for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) { - if (entity_.get_key(id) == kInvalidKey) { + for (node_id_t id = 0; id < entity_->doc_cnt(); ++id) { + if (entity_->get_key(id) == kInvalidKey) { continue; } - if (!ctx->filter().is_valid() || !ctx->filter()(entity_.get_key(id))) { + if (!ctx->filter().is_valid() || !ctx->filter()(entity_->get_key(id))) { dist_t dist = ctx->dist_calculator().batch_dist(id); std::string group_id = group_by(id); @@ -728,12 +777,12 @@ int HnswStreamer::search_bf_impl( for (size_t q = 0; q < count; ++q) { ctx->reset_query(query); topk.clear(); - for (node_id_t id = 0; id < entity_.doc_cnt(); ++id) { - if (entity_.get_key(id) == kInvalidKey) { + for (node_id_t id = 0; id < entity_->doc_cnt(); ++id) { + if (entity_->get_key(id) == kInvalidKey) { continue; } - if (!filter.is_valid() || !filter(entity_.get_key(id))) { + if (!filter.is_valid() || !filter(entity_->get_key(id))) { dist_t dist = ctx->dist_calculator().batch_dist(id); topk.emplace(id, dist); } @@ -788,7 +837,7 @@ int HnswStreamer::search_bf_by_p_keys_impl( } std::function group_by = [&](node_id_t id) { - return ctx->group_by()(entity_.get_key(id)); + return ctx->group_by()(entity_->get_key(id)); }; for (size_t q = 0; q < count; ++q) { @@ -798,7 +847,7 @@ int HnswStreamer::search_bf_by_p_keys_impl( for (size_t idx = 0; idx < p_keys[q].size(); ++idx) { uint64_t pk = p_keys[q][idx]; if (!ctx->filter().is_valid() || !ctx->filter()(pk)) { - node_id_t id = entity_.get_id(pk); + node_id_t id = entity_->get_id(pk); if (id != kInvalidNodeId) { dist_t dist = ctx->dist_calculator().batch_dist(id); std::string group_id = group_by(id); @@ -824,7 +873,7 @@ int HnswStreamer::search_bf_by_p_keys_impl( for (size_t idx = 0; idx < p_keys[q].size(); ++idx) { key_t pk = p_keys[q][idx]; if (!filter.is_valid() || !filter(pk)) { - node_id_t id = entity_.get_id(pk); + node_id_t id = entity_->get_id(pk); if (id != kInvalidNodeId) { dist_t dist = ctx->dist_calculator().batch_dist(id); topk.emplace(id, dist); diff --git a/src/core/algorithm/hnsw/hnsw_streamer.h b/src/core/algorithm/hnsw/hnsw_streamer.h index b81106d..3f4511a 100644 --- a/src/core/algorithm/hnsw/hnsw_streamer.h +++ b/src/core/algorithm/hnsw/hnsw_streamer.h @@ -31,6 +31,17 @@ class HnswStreamer : public IndexStreamer { HnswStreamer(const HnswStreamer &streamer) = delete; HnswStreamer &operator=(const HnswStreamer &streamer) = delete; + public: + //! Retrieve the storage mode of the underlying entity. Returns + //! HnswStorageMode::kMmap when the entity has not been initialized yet. + //! Intended for introspection and debug/testing usage. + HnswStorageMode storage_mode() const { + if (!entity_) { + return HnswStorageMode::kMmap; + } + return entity_->storage_mode(); + } + protected: //! Initialize Streamer virtual int init(const IndexMeta &imeta, @@ -88,22 +99,22 @@ class HnswStreamer : public IndexStreamer { //! Fetch vector by key virtual const void *get_vector(uint64_t key) const override { - return entity_.get_vector_by_key(key); + return entity_->get_vector_by_key(key); } virtual int get_vector(const uint64_t key, IndexStorage::MemoryBlock &block) const override { - return entity_.get_vector_by_key(key, block); + return entity_->get_vector_by_key(key, block); } //! Fetch vector by id virtual const void *get_vector_by_id(uint32_t id) const override { - return entity_.get_vector(id); + return entity_->get_vector(id); } virtual int get_vector_by_id( const uint32_t id, IndexStorage::MemoryBlock &block) const override { - return entity_.get_vector(id, block); + return entity_->get_vector(id, block); } //! Open index from file path @@ -159,6 +170,9 @@ class HnswStreamer : public IndexStreamer { } private: + //! Configure and initialize the entity with saved parameters + int setup_entity(); + //! To share ctx across streamer/searcher, we need to update the context for //! current streamer/searcher int update_context(HnswContext *ctx) const; @@ -181,8 +195,8 @@ class HnswStreamer : public IndexStreamer { } }; - HnswStreamerEntity entity_; - HnswAlgorithm::UPointer alg_; + std::unique_ptr entity_; + HnswAlgorithmBase::UPointer alg_; IndexMeta meta_{}; IndexMetric::Pointer metric_{}; @@ -200,6 +214,7 @@ class HnswStreamer : public IndexStreamer { size_t docs_hard_limit_{HnswEntity::kDefaultDocsHardLimit}; size_t docs_soft_limit_{0UL}; uint32_t min_neighbor_cnt_{0u}; + uint32_t prune_cnt_{0u}; uint32_t upper_max_neighbor_cnt_{HnswEntity::kDefaultUpperMaxNeighborCnt}; uint32_t l0_max_neighbor_cnt_{HnswEntity::kDefaultL0MaxNeighborCnt}; uint32_t ef_{HnswEntity::kDefaultEf}; @@ -219,6 +234,7 @@ class HnswStreamer : public IndexStreamer { bool get_vector_enabled_{false}; bool force_padding_topk_enabled_{false}; bool use_id_map_{true}; + bool use_contiguous_memory_{false}; //! avoid add vector while dumping index ailego::SharedMutex shared_mutex_{}; diff --git a/src/core/algorithm/hnsw/hnsw_streamer_entity.cc b/src/core/algorithm/hnsw/hnsw_streamer_entity.cc index 4eef527..acc9bee 100644 --- a/src/core/algorithm/hnsw/hnsw_streamer_entity.cc +++ b/src/core/algorithm/hnsw/hnsw_streamer_entity.cc @@ -13,6 +13,9 @@ // limitations under the License. #include "hnsw_streamer_entity.h" +#if defined(__linux__) || defined(__APPLE__) +#include +#endif #include // #define DEBUG_PRINT @@ -773,5 +776,214 @@ const HnswEntity::Pointer HnswStreamerEntity::clone() const { return HnswEntity::Pointer(entity); } +const HnswEntity::Pointer HnswMmapStreamerEntity::clone() const { + std::vector node_chunks; + node_chunks.reserve(node_chunks_.size()); + for (size_t i = 0UL; i < node_chunks_.size(); ++i) { + node_chunks.emplace_back(node_chunks_[i]->clone()); + if (ailego_unlikely(!node_chunks[i])) { + LOG_ERROR("HnswMmapStreamerEntity get chunk failed in clone"); + return HnswEntity::Pointer(); + } + } + + std::vector upper_neighbor_chunks; + upper_neighbor_chunks.reserve(upper_neighbor_chunks_.size()); + for (size_t i = 0UL; i < upper_neighbor_chunks_.size(); ++i) { + upper_neighbor_chunks.emplace_back(upper_neighbor_chunks_[i]->clone()); + if (ailego_unlikely(!upper_neighbor_chunks[i])) { + LOG_ERROR("HnswMmapStreamerEntity get chunk failed in clone"); + return HnswEntity::Pointer(); + } + } + + auto *entity = new (std::nothrow) HnswMmapStreamerEntity( + stats_, header(), chunk_size_, node_index_mask_bits_, + upper_neighbor_mask_bits_, filter_same_key_, get_vector_enabled_, + upper_neighbor_index_, keys_map_lock_, keys_map_, use_key_info_map_, + std::move(node_chunks), std::move(upper_neighbor_chunks), broker_, + nullptr, nullptr); + if (ailego_unlikely(!entity)) { + LOG_ERROR("HnswMmapStreamerEntity new failed"); + } + return HnswEntity::Pointer(entity); +} + +const HnswEntity::Pointer HnswContiguousStreamerEntity::clone() const { + std::vector node_chunks; + node_chunks.reserve(node_chunks_.size()); + for (size_t i = 0UL; i < node_chunks_.size(); ++i) { + node_chunks.emplace_back(node_chunks_[i]->clone()); + if (ailego_unlikely(!node_chunks[i])) { + LOG_ERROR("HnswContiguousStreamerEntity get chunk failed in clone"); + return HnswEntity::Pointer(); + } + } + + std::vector upper_neighbor_chunks; + upper_neighbor_chunks.reserve(upper_neighbor_chunks_.size()); + for (size_t i = 0UL; i < upper_neighbor_chunks_.size(); ++i) { + upper_neighbor_chunks.emplace_back(upper_neighbor_chunks_[i]->clone()); + if (ailego_unlikely(!upper_neighbor_chunks[i])) { + LOG_ERROR("HnswContiguousStreamerEntity get chunk failed in clone"); + return HnswEntity::Pointer(); + } + } + + auto *entity = new (std::nothrow) HnswContiguousStreamerEntity( + stats_, header(), chunk_size_, node_index_mask_bits_, + upper_neighbor_mask_bits_, filter_same_key_, get_vector_enabled_, + upper_neighbor_index_, keys_map_lock_, keys_map_, use_key_info_map_, + std::move(node_chunks), std::move(upper_neighbor_chunks), broker_, + nullptr, nullptr); + if (ailego_unlikely(!entity)) { + LOG_ERROR("HnswContiguousStreamerEntity new failed"); + return HnswEntity::Pointer(); + } + + // Share contiguous memory with the clone (zero-copy) + entity->node_memory_ = node_memory_; + entity->node_base_ = node_base_; + entity->upper_neighbor_memory_ = upper_neighbor_memory_; + entity->upper_neighbor_base_ = upper_neighbor_base_; + entity->upper_chunk_offsets_ = upper_chunk_offsets_; + + return HnswEntity::Pointer(entity); +} + +// ============================================================================ +// HnswContiguousStreamerEntity implementation +// ============================================================================ + +char *HnswContiguousStreamerEntity::allocate_contiguous(size_t size) { + if (size == 0) { + return nullptr; + } +#if defined(__linux__) + // Use mmap with MAP_ANONYMOUS for contiguous memory + void *ptr = ::mmap(nullptr, size, PROT_READ | PROT_WRITE, + MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + if (ptr == MAP_FAILED) { + LOG_ERROR("mmap failed for contiguous memory, size=%zu", size); + return nullptr; + } + // Request transparent huge pages + ::madvise(ptr, size, MADV_HUGEPAGE); + return static_cast(ptr); +#elif defined(__APPLE__) + // macOS: use mmap with MAP_ANONYMOUS + void *ptr = ::mmap(nullptr, size, PROT_READ | PROT_WRITE, + MAP_PRIVATE | MAP_ANON, -1, 0); + if (ptr == MAP_FAILED) { + LOG_ERROR("mmap failed for contiguous memory, size=%zu", size); + return nullptr; + } + return static_cast(ptr); +#elif defined(_WIN32) + void *ptr = ::_aligned_malloc(size, ailego::MemoryHelper::PageSize()); + if (!ptr) { + LOG_ERROR("_aligned_malloc failed for contiguous memory, size=%zu", size); + return nullptr; + } + return static_cast(ptr); +#else + // Fallback: aligned allocation + void *ptr = std::aligned_alloc(ailego::MemoryHelper::PageSize(), size); + if (!ptr) { + LOG_ERROR("aligned_alloc failed for contiguous memory, size=%zu", size); + return nullptr; + } + return static_cast(ptr); +#endif +} + +int HnswContiguousStreamerEntity::build_contiguous_memory() { + node_memory_.reset(); + node_base_ = nullptr; + upper_neighbor_memory_.reset(); + upper_neighbor_base_ = nullptr; + upper_chunk_offsets_.clear(); + + const uint32_t total_docs = doc_cnt(); + if (total_docs == 0) { + return 0; + } + + // --- Build contiguous node memory --- + const size_t per_node = node_size(); + const size_t total_node_data = static_cast(total_docs) * per_node; + size_t node_memory_size = AlignHugePageSize(total_node_data); + char *raw_node = allocate_contiguous(node_memory_size); + if (!raw_node) { + return IndexError_Runtime; + } + node_memory_.reset(raw_node, ContiguousDeleter{node_memory_size}); + node_base_ = raw_node; + + // Copy node data from chunks into contiguous memory + // Each chunk holds node_cnt_per_chunk nodes, laid out at offset + // (id & mask) * node_size within the chunk. + const auto &chunks = node_chunks_; + const uint32_t nodes_per_chunk = 1U << node_index_mask_bits_; + for (size_t chunk_idx = 0; chunk_idx < chunks.size(); ++chunk_idx) { + const void *chunk_data = nullptr; + size_t data_size = chunks[chunk_idx]->data_size(); + chunks[chunk_idx]->read(0, &chunk_data, data_size); + + // Number of nodes in this chunk + uint32_t base_id = chunk_idx * nodes_per_chunk; + uint32_t count_in_chunk = std::min(nodes_per_chunk, total_docs - base_id); + + // Copy each node's data + const char *src = static_cast(chunk_data); + char *dst = node_base_ + static_cast(base_id) * per_node; + std::memcpy(dst, src, static_cast(count_in_chunk) * per_node); + } + + // --- Build contiguous upper neighbor memory --- + const auto &upper_chunks = upper_neighbor_chunks_; + if (upper_chunks.empty()) { + return 0; + } + + // Sync all upper neighbor chunks + sync_upper_neighbor_chunks(upper_chunks.size() - 1); + + // Calculate cumulative offsets and total size + upper_chunk_offsets_.resize(upper_chunks.size()); + size_t total_upper_size = 0; + for (size_t i = 0; i < upper_chunks.size(); ++i) { + upper_chunk_offsets_[i] = total_upper_size; + total_upper_size += upper_chunks[i]->data_size(); + } + + size_t upper_memory_size = AlignHugePageSize(total_upper_size); + char *raw_upper = allocate_contiguous(upper_memory_size); + if (!raw_upper) { + node_memory_.reset(); + node_base_ = nullptr; + return IndexError_Runtime; + } + upper_neighbor_memory_.reset(raw_upper, ContiguousDeleter{upper_memory_size}); + upper_neighbor_base_ = raw_upper; + + // Copy upper neighbor data from chunks into contiguous memory + for (size_t i = 0; i < upper_chunks.size(); ++i) { + const void *chunk_data = nullptr; + size_t data_size = upper_chunks[i]->data_size(); + upper_chunks[i]->read(0, &chunk_data, data_size); + std::memcpy(upper_neighbor_base_ + upper_chunk_offsets_[i], chunk_data, + data_size); + } + + LOG_INFO( + "Built contiguous memory: node_size=%zu upper_neighbor_size=%zu " + "total_docs=%u node_chunks=%zu upper_chunks=%zu", + node_memory_size, upper_memory_size, total_docs, chunks.size(), + upper_chunks.size()); + + return 0; +} + } // namespace core } // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_streamer_entity.h b/src/core/algorithm/hnsw/hnsw_streamer_entity.h index a357062..3dc6c96 100644 --- a/src/core/algorithm/hnsw/hnsw_streamer_entity.h +++ b/src/core/algorithm/hnsw/hnsw_streamer_entity.h @@ -15,6 +15,11 @@ #pragma once #include +#include +#include +#if defined(__linux__) || defined(__APPLE__) +#include +#endif #include #include #include @@ -28,6 +33,10 @@ namespace zvec { namespace core { + +//! Storage mode for HnswStreamerEntity +enum class HnswStorageMode { kMmap = 0, kBufferPool = 1, kContiguous = 2 }; + //! HnswStreamerEntity manage vector data, pkey, and node's neighbors class HnswStreamerEntity : public HnswEntity { public: @@ -84,6 +93,11 @@ class HnswStreamerEntity : public HnswEntity { virtual void update_ep_and_level(node_id_t ep, level_t level) override; + //! Get the storage mode of this entity + virtual HnswStorageMode storage_mode() const { + return HnswStorageMode::kMmap; + } + void set_use_key_info_map(bool use_id_map) { use_key_info_map_ = use_id_map; LOG_DEBUG("use_key_info_map_: %d", (int)use_key_info_map_); @@ -178,6 +192,22 @@ class HnswStreamerEntity : public HnswEntity { std::cout << "key map ends" << std::endl; } + //! Typed get_neighbors: returns NeighborsT without runtime + //! branching on MemoryBlock type. MmapMemoryBlock specialization uses + //! pointer-based read; BufferPoolMemoryBlock uses MemoryBlock-based read. + template + inline NeighborsT get_neighbors_typed(level_t level, + node_id_t id) const; + + //! Typed batch get_vector: fills vector without runtime branching + template + inline int get_vector_typed(const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const; + + //! Typed get_key: reads key using typed MemBlock + template + inline key_t get_key_typed(node_id_t id) const; + //! Get l0 neighbors size inline size_t neighbors_size() const { return sizeof(NeighborsHeader) + l0_neighbor_cnt() * sizeof(node_id_t); @@ -189,7 +219,7 @@ class HnswStreamerEntity : public HnswEntity { } - private: + protected: union UpperNeighborIndexMeta { struct { uint32_t level : 4; @@ -200,6 +230,7 @@ class HnswStreamerEntity : public HnswEntity { uint32_t data; }; + protected: template using HashMap = google::dense_hash_map>; template @@ -214,7 +245,7 @@ class HnswStreamerEntity : public HnswEntity { using NIHashMap = HnswIndexHashMap; using NIHashMapPointer = std::shared_ptr; - //! Private construct, only be called by clone method + //! Clone construct, used by clone method in subclasses HnswStreamerEntity(IndexStreamer::Stats &stats, const HNSWHeader &hd, size_t chunk_size, uint32_t node_index_mask_bits, uint32_t upper_neighbor_mask_bits, bool filter_same_key, @@ -480,14 +511,23 @@ class HnswStreamerEntity : public HnswEntity { return 0; } + protected: + //! Expose sync_chunks for subclass use + inline void sync_node_chunks(size_t idx) const { + sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, idx, &node_chunks_); + } + inline void sync_upper_neighbor_chunks(size_t idx) const { + sync_chunks(ChunkBroker::CHUNK_TYPE_UPPER_NEIGHBOR, idx, + &upper_neighbor_chunks_); + } + private: HnswStreamerEntity(const HnswStreamerEntity &) = delete; HnswStreamerEntity &operator=(const HnswStreamerEntity &) = delete; static constexpr uint64_t kUpperHashMemoryInflateRatio = 2.0f; - private: + protected: IndexStreamer::Stats &stats_; - HNSWHeader header_{}; std::mutex mutex_{}; size_t max_index_size_{0UL}; uint32_t chunk_size_{kDefaultChunkSize}; @@ -536,5 +576,434 @@ class HnswStreamerEntity : public HnswEntity { ChunkBroker::Pointer broker_{}; // chunk broker }; +// --- Template specializations for typed MemoryBlock access --- + +//! MmapMemoryBlock specialization: uses pointer-based Chunk::read +template <> +inline NeighborsT +HnswStreamerEntity::get_neighbors_typed(level_t level, + node_id_t id) const { + Chunk *chunk = nullptr; + size_t offset = 0UL; + size_t nbr_size = neighbor_size_; + if (level == 0UL) { + uint32_t chunk_idx = id >> node_index_mask_bits_; + offset = + (id & node_index_mask_) * node_size() + vector_size() + sizeof(key_t); + sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_); + ailego_assert_with(chunk_idx < node_chunks_.size(), "invalid chunk idx"); + chunk = node_chunks_[chunk_idx].get(); + } else { + auto p = get_upper_neighbor_chunk_loc(level, id); + chunk = upper_neighbor_chunks_[p.first].get(); + offset = p.second; + nbr_size = upper_neighbor_size_; + } + ailego_assert_with(offset < chunk->data_size(), "invalid chunk offset"); + const void *ptr = nullptr; + size_t ret = chunk->read(offset, &ptr, nbr_size); + if (ailego_unlikely(ret != nbr_size)) { + LOG_ERROR("Read neighbor header failed, ret=%zu", ret); + return NeighborsT(); + } + MmapMemoryBlock block(const_cast(ptr)); + return NeighborsT(std::move(block)); +} + +//! BufferPoolMemoryBlock specialization: uses MemoryBlock-based Chunk::read +template <> +inline NeighborsT +HnswStreamerEntity::get_neighbors_typed( + level_t level, node_id_t id) const { + Chunk *chunk = nullptr; + size_t offset = 0UL; + size_t nbr_size = neighbor_size_; + if (level == 0UL) { + uint32_t chunk_idx = id >> node_index_mask_bits_; + offset = + (id & node_index_mask_) * node_size() + vector_size() + sizeof(key_t); + sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_); + ailego_assert_with(chunk_idx < node_chunks_.size(), "invalid chunk idx"); + chunk = node_chunks_[chunk_idx].get(); + } else { + auto p = get_upper_neighbor_chunk_loc(level, id); + chunk = upper_neighbor_chunks_[p.first].get(); + offset = p.second; + nbr_size = upper_neighbor_size_; + } + ailego_assert_with(offset < chunk->data_size(), "invalid chunk offset"); + IndexStorage::MemoryBlock mem_block; + size_t ret = chunk->read(offset, mem_block, nbr_size); + if (ailego_unlikely(ret != nbr_size)) { + LOG_ERROR("Read neighbor header failed, ret=%zu", ret); + return NeighborsT(); + } + BufferPoolMemoryBlock block(mem_block.buffer_pool_handle_, + mem_block.buffer_block_id_, mem_block.data_); + mem_block.buffer_pool_handle_ = nullptr; + return NeighborsT(std::move(block)); +} + +//! MmapMemoryBlock specialization for batch get_vector +template <> +inline int HnswStreamerEntity::get_vector_typed( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + vec_blocks.resize(count); + for (auto i = 0U; i < count; ++i) { + auto loc = get_vector_chunk_loc(ids[i]); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + size_t read_size = vector_size(); + const void *ptr = nullptr; + size_t ret = node_chunks_[loc.first]->read(loc.second, &ptr, read_size); + if (ailego_unlikely(ret != read_size)) { + LOG_ERROR("Read vector failed, offset=%u, read size=%zu, ret=%zu", + loc.second, read_size, ret); + return IndexError_ReadData; + } + vec_blocks[i].reset(const_cast(ptr)); + } + return 0; +} + +//! BufferPoolMemoryBlock specialization for batch get_vector +template <> +inline int HnswStreamerEntity::get_vector_typed( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + vec_blocks.resize(count); + for (auto i = 0U; i < count; ++i) { + auto loc = get_vector_chunk_loc(ids[i]); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + size_t read_size = vector_size(); + IndexStorage::MemoryBlock mem_block; + size_t ret = + node_chunks_[loc.first]->read(loc.second, mem_block, read_size); + if (ailego_unlikely(ret != read_size)) { + LOG_ERROR("Read vector failed, offset=%u, read size=%zu, ret=%zu", + loc.second, read_size, ret); + return IndexError_ReadData; + } + vec_blocks[i] = + BufferPoolMemoryBlock(mem_block.buffer_pool_handle_, + mem_block.buffer_block_id_, mem_block.data_); + mem_block.buffer_pool_handle_ = nullptr; + } + return 0; +} + +//! MmapMemoryBlock specialization for get_key +template <> +inline key_t HnswStreamerEntity::get_key_typed( + node_id_t id) const { + if (!use_key_info_map_) { + return id; + } + auto loc = get_key_chunk_loc(id); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + const void *ptr = nullptr; + size_t ret = node_chunks_[loc.first]->read(loc.second, &ptr, sizeof(key_t)); + if (ailego_unlikely(ret != sizeof(key_t))) { + LOG_ERROR("Read key failed, ret=%zu", ret); + return kInvalidKey; + } + return *reinterpret_cast(ptr); +} + +//! BufferPoolMemoryBlock specialization for get_key +template <> +inline key_t HnswStreamerEntity::get_key_typed( + node_id_t id) const { + if (!use_key_info_map_) { + return id; + } + auto loc = get_key_chunk_loc(id); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + ailego_assert_with(loc.second < node_chunks_[loc.first]->data_size(), + "invalid chunk offset"); + IndexStorage::MemoryBlock key_block; + size_t ret = + node_chunks_[loc.first]->read(loc.second, key_block, sizeof(key_t)); + if (ailego_unlikely(ret != sizeof(key_t))) { + LOG_ERROR("Read key failed, ret=%zu", ret); + return kInvalidKey; + } + return *reinterpret_cast(key_block.data()); +} + +//! Typed entity subclass for mmap mode. +//! Caches chunk base addresses to eliminate virtual function calls on the +//! search hot path. For mmap mode, chunk data is memory-mapped at init time, +//! so we can directly compute pointers via base_addr + offset. +class HnswMmapStreamerEntity : public HnswStreamerEntity { + public: + using MemoryBlock = MmapMemoryBlock; + using TypedNeighbors = NeighborsT; + + using HnswStreamerEntity::HnswStreamerEntity; + + HnswStorageMode storage_mode() const override { + return HnswStorageMode::kMmap; + } + + //! Override clone to return correct subclass type, so that + //! static_cast in the algorithm is safe. + const HnswEntity::Pointer clone() const override; + + inline TypedNeighbors get_neighbors_typed(level_t level, node_id_t id) const { + if (level == 0UL) { + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = + (id & node_index_mask_) * node_size() + vector_size() + sizeof(key_t); + const char *base = get_node_chunk_base(chunk_idx); + MmapMemoryBlock block(const_cast(base + offset)); + return TypedNeighbors(std::move(block)); + } + // Upper level: use index to locate chunk and offset + auto it = upper_neighbor_index_->find(id); + ailego_assert_abort(it != upper_neighbor_index_->end(), + "Get upper neighbor header failed"); + auto meta = reinterpret_cast(&it->second); + uint32_t chunk_idx = (meta->index) >> upper_neighbor_mask_bits_; + uint32_t offset = (((meta->index) & upper_neighbor_mask_) + level - 1) * + upper_neighbor_size_; + const char *base = get_upper_neighbor_chunk_base(chunk_idx); + MmapMemoryBlock block(const_cast(base + offset)); + return TypedNeighbors(std::move(block)); + } + + inline int get_vector_typed(const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + vec_blocks.resize(count); + for (auto i = 0U; i < count; ++i) { + uint32_t chunk_idx = ids[i] >> node_index_mask_bits_; + uint32_t offset = (ids[i] & node_index_mask_) * node_size(); + const char *base = get_node_chunk_base(chunk_idx); + vec_blocks[i].reset(const_cast(base + offset)); + } + return 0; + } + + inline key_t get_key_typed(node_id_t id) const { + if (!use_key_info_map_) { + return id; + } + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = (id & node_index_mask_) * node_size() + vector_size(); + const char *base = get_node_chunk_base(chunk_idx); + return *reinterpret_cast(base + offset); + } + + private: + //! Get cached base address for a node chunk, syncing if needed + inline const char *get_node_chunk_base(uint32_t chunk_idx) const { + if (ailego_unlikely(chunk_idx >= node_chunk_bases_.size())) { + sync_node_chunk_bases(chunk_idx); + } + return node_chunk_bases_[chunk_idx]; + } + + //! Get cached base address for an upper neighbor chunk, syncing if needed + inline const char *get_upper_neighbor_chunk_base(uint32_t chunk_idx) const { + if (ailego_unlikely(chunk_idx >= upper_neighbor_chunk_bases_.size())) { + sync_upper_neighbor_chunk_bases(chunk_idx); + } + return upper_neighbor_chunk_bases_[chunk_idx]; + } + + //! Sync node chunk base addresses up to the given index + void sync_node_chunk_bases(uint32_t chunk_idx) const { + sync_node_chunks(chunk_idx); + const auto &chunks = node_chunks_; + for (size_t i = node_chunk_bases_.size(); i <= chunk_idx; ++i) { + const void *ptr = nullptr; + chunks[i]->read(0, &ptr, 1); + node_chunk_bases_.push_back(static_cast(ptr)); + } + } + + //! Sync upper neighbor chunk base addresses up to the given index + void sync_upper_neighbor_chunk_bases(uint32_t chunk_idx) const { + sync_upper_neighbor_chunks(chunk_idx); + const auto &chunks = upper_neighbor_chunks_; + for (size_t i = upper_neighbor_chunk_bases_.size(); i <= chunk_idx; ++i) { + const void *ptr = nullptr; + chunks[i]->read(0, &ptr, 1); + upper_neighbor_chunk_bases_.push_back(static_cast(ptr)); + } + } + + mutable std::vector node_chunk_bases_{}; + mutable std::vector upper_neighbor_chunk_bases_{}; +}; + +//! Typed entity subclass for buffer pool mode. +class HnswBufferPoolStreamerEntity : public HnswStreamerEntity { + public: + using MemoryBlock = BufferPoolMemoryBlock; + using TypedNeighbors = NeighborsT; + + using HnswStreamerEntity::HnswStreamerEntity; + + HnswStorageMode storage_mode() const override { + return HnswStorageMode::kBufferPool; + } + + inline TypedNeighbors get_neighbors_typed(level_t level, node_id_t id) const { + return HnswStreamerEntity::get_neighbors_typed(level, + id); + } + + inline int get_vector_typed( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + return HnswStreamerEntity::get_vector_typed( + ids, count, vec_blocks); + } + + inline key_t get_key_typed(node_id_t id) const { + return HnswStreamerEntity::get_key_typed(id); + } +}; + +//! Typed entity subclass for contiguous memory mode. +//! Allocates contiguous memory (with hugepage/THP support) and copies all +//! chunk data into it. Access is via a single base pointer + offset, +//! eliminating chunk-level indirection and maximizing memory locality. +class HnswContiguousStreamerEntity : public HnswMmapStreamerEntity { + public: + using HnswMmapStreamerEntity::HnswMmapStreamerEntity; + + HnswStorageMode storage_mode() const override { + return HnswStorageMode::kContiguous; + } + + //! Override clone to return correct subclass type. + //! Cloned entity shares contiguous memory via shared_ptr. + const HnswEntity::Pointer clone() const override; + + ~HnswContiguousStreamerEntity() = default; + + //! Build contiguous memory from chunks after open. + //! Must be called after the entity is fully opened and all chunks are loaded. + int build_contiguous_memory(); + + //! Degrade to mmap mode by releasing contiguous memory and falling back + //! to chunk-based access. + void degrade_to_mmap() { + node_memory_.reset(); + node_base_ = nullptr; + upper_neighbor_memory_.reset(); + upper_neighbor_base_ = nullptr; + upper_chunk_offsets_.clear(); + LOG_INFO("HNSW contiguous entity degraded to mmap mode for insertion"); + } + + bool is_contiguous() const { + return node_base_ != nullptr; + } + + int add_vector(level_t level, key_t key, const void *vec, + node_id_t *id) override { + if (ailego_unlikely(is_contiguous())) degrade_to_mmap(); + return HnswMmapStreamerEntity::add_vector(level, key, vec, id); + } + + int add_vector_with_id(level_t level, node_id_t id, + const void *vec) override { + if (ailego_unlikely(is_contiguous())) degrade_to_mmap(); + return HnswMmapStreamerEntity::add_vector_with_id(level, id, vec); + } + + inline TypedNeighbors get_neighbors_typed(level_t level, node_id_t id) const { + if (ailego_likely(node_base_ != nullptr)) { + if (level == 0UL) { + const char *ptr = node_base_ + static_cast(id) * node_size() + + vector_size() + sizeof(key_t); + MmapMemoryBlock block(const_cast(ptr)); + return TypedNeighbors(std::move(block)); + } + // Upper level: use index to locate global offset + auto it = upper_neighbor_index_->find(id); + ailego_assert_abort(it != upper_neighbor_index_->end(), + "Get upper neighbor header failed"); + auto meta = reinterpret_cast(&it->second); + uint32_t chunk_idx = (meta->index) >> upper_neighbor_mask_bits_; + uint32_t local_idx = (meta->index) & upper_neighbor_mask_; + size_t global_offset = + upper_chunk_offsets_[chunk_idx] + + static_cast(local_idx + level - 1) * upper_neighbor_size_; + const char *ptr = upper_neighbor_base_ + global_offset; + MmapMemoryBlock block(const_cast(ptr)); + return TypedNeighbors(std::move(block)); + } + return HnswMmapStreamerEntity::get_neighbors_typed(level, id); + } + + inline int get_vector_typed(const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + if (ailego_likely(node_base_ != nullptr)) { + vec_blocks.resize(count); + for (auto i = 0U; i < count; ++i) { + const char *ptr = + node_base_ + static_cast(ids[i]) * node_size(); + vec_blocks[i].reset(const_cast(ptr)); + } + return 0; + } + return HnswMmapStreamerEntity::get_vector_typed(ids, count, vec_blocks); + } + + inline key_t get_key_typed(node_id_t id) const { + if (ailego_likely(node_base_ != nullptr)) { + if (!use_key_info_map_) { + return id; + } + const char *ptr = + node_base_ + static_cast(id) * node_size() + vector_size(); + return *reinterpret_cast(ptr); + } + return HnswMmapStreamerEntity::get_key_typed(id); + } + + protected: + //! Custom deleter for contiguous memory (munmap / _aligned_free / free) + //! Used by shared_ptr to properly release mmap'd memory. + struct ContiguousDeleter { + size_t size; + void operator()(char *ptr) const { + if (!ptr) return; +#if defined(__linux__) || defined(__APPLE__) + ::munmap(ptr, size); +#elif defined(_WIN32) + ::_aligned_free(ptr); +#else + std::free(ptr); +#endif + } + }; + + //! Shared ownership of contiguous memory (enables zero-copy clone) + std::shared_ptr node_memory_{}; + std::shared_ptr upper_neighbor_memory_{}; + + //! Raw pointers for hot-path access (derived from shared_ptr) + char *node_base_{nullptr}; + char *upper_neighbor_base_{nullptr}; + + //! Cumulative offsets for each upper neighbor chunk in contiguous memory + std::vector upper_chunk_offsets_{}; + + private: + //! Allocate contiguous memory with hugepage/THP support + static char *allocate_contiguous(size_t size); +}; + } // namespace core } // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/vamana/CMakeLists.txt b/src/core/algorithm/vamana/CMakeLists.txt new file mode 100644 index 0000000..8e5bbda --- /dev/null +++ b/src/core/algorithm/vamana/CMakeLists.txt @@ -0,0 +1,11 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +cc_library( + NAME core_knn_vamana + STATIC SHARED STRICT ALWAYS_LINK + SRCS *.cc + LIBS core_framework core_knn_hnsw sparsehash + INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm + VERSION "${PROXIMA_ZVEC_VERSION}" + ) diff --git a/src/core/algorithm/vamana/vamana_algorithm.cc b/src/core/algorithm/vamana/vamana_algorithm.cc new file mode 100644 index 0000000..dbbff34 --- /dev/null +++ b/src/core/algorithm/vamana/vamana_algorithm.cc @@ -0,0 +1,465 @@ +// Copyright 2025-present the zvec 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 "vamana_algorithm.h" + +namespace zvec { +namespace core { + +// ============================================================================ +// add_node: Insert a new node into the Vamana graph. +// +// Algorithm (from DiskANN / Vamana paper): +// 1. GreedySearch from entry_point to find candidate neighbors +// 2. RobustPrune to select diverse neighbors for the new node +// 3. Update the new node's neighbor list +// 4. For each new neighbor, add reverse link; if over-degree, RobustPrune +// 5. If this is the first node, set it as entry point +// ============================================================================ +template +int VamanaAlgorithm::add_node(node_id_t id, VamanaContext *ctx) { + // Lazily initialize distance storage on first insert + entity_.ensure_dist_storage(); + + spin_lock_.lock(); + auto entry_point = entity_.entry_point(); + + if (ailego_unlikely(entry_point == kInvalidNodeId)) { + entity_.update_entry_point(id); + spin_lock_.unlock(); + return 0; + } + spin_lock_.unlock(); + + // Step 1: GreedySearch to find candidate neighbors + uint32_t search_list_size = entity_.search_list_size(); + ctx->topk_heap().clear(); + ctx->topk_heap().limit(search_list_size); + ctx->dist_calculator().clear_compare_cnt(); + + // Set query to the new node's vector + const void *query_vec = entity_.get_vector(id); + if (ailego_unlikely(query_vec == nullptr)) { + LOG_ERROR("Failed to get vector for node %u", id); + return IndexError_ReadData; + } + ctx->reset_query(query_vec); + + greedy_search(entry_point, ctx); + + // Step 2: RobustPrune to select diverse neighbors + auto &topk_heap = ctx->topk_heap(); + robust_prune(id, topk_heap, entity_.alpha(), entity_.max_degree(), ctx); + // Copy result before reverse updates (which also call robust_prune) + auto pruned_neighbors = ctx->prune_result(); + + // Step 3: Update the new node's neighbor list and distances + entity_.update_neighbors(id, pruned_neighbors); + entity_.update_neighbor_dists(id, pruned_neighbors); + + // Step 4: Reverse-link updates + update_neighbors_and_reverse_links(id, pruned_neighbors, ctx); + + return 0; +} + +// ============================================================================ +// search: Greedy search for approximate nearest neighbors. +// ============================================================================ +template +int VamanaAlgorithm::search(VamanaContext *ctx) const { + spin_lock_.lock(); + auto entry_point = entity_.entry_point(); + spin_lock_.unlock(); + + if (ailego_unlikely(entry_point == kInvalidNodeId)) { + return 0; + } + + auto &topk_heap = ctx->topk_heap(); + topk_heap.clear(); + + // Use ef (query-time parameter) instead of entity.search_list_size() + // (build-time L parameter). search_list_size controls construction; + // ef controls search quality and is user-configurable at query time. + uint32_t ef_search = std::max(static_cast(ctx->topk()), ctx->ef()); + topk_heap.limit(ef_search); + + greedy_search(entry_point, ctx); + + return 0; +} + +// ============================================================================ +// greedy_search: Beam search from entry_point. +// +// Maintains a candidate min-heap (ordered by distance) and a visited set. +// At each step, pops the closest unvisited candidate, expands its neighbors, +// and adds unvisited neighbors to both the candidate heap and the topk heap. +// Stops when the closest candidate is farther than the worst in topk, or +// when the scan limit is reached. +// ============================================================================ +template +void VamanaAlgorithm::greedy_search(node_id_t entry_point, + VamanaContext *ctx) const { + const auto &entity = static_cast(ctx->get_entity()); + VamanaDistCalculator &dc = ctx->dist_calculator(); + VisitFilter &visit = ctx->visit_filter(); + CandidateHeap &candidates = ctx->candidates(); + auto &topk_heap = ctx->topk_heap(); + + const IndexFilter &index_filter = + static_cast(ctx)->filter(); + std::function filter = [](node_id_t) { return false; }; + if (index_filter.is_valid()) { + filter = [&](node_id_t id) { + return index_filter(entity.get_key_typed(id)); + }; + } + + candidates.clear(); + visit.clear(); + + // Initialize with entry point using batch_dist (single-element batch). + // We must NOT use dc.dist(entry_point) here because dist() calls + // distance_() which is squared_euclidean_int8_distance (sign/abs trick, + // expects two raw int8 inputs), but query_ has been preprocessed by + // reset_query (+128 shift to uint8). batch_dist() correctly calls + // batch_distance_() which expects the preprocessed uint8 query. + dist_t entry_dist = dc.batch_dist(entry_point); + if (ailego_unlikely(dc.error())) { + return; + } + visit.set_visited(entry_point); + if (!filter(entry_point)) { + topk_heap.emplace(entry_point, entry_dist); + } + candidates.emplace(entry_point, entry_dist); + + // Pre-allocate temporary vectors outside the hot loop to avoid + // per-iteration heap allocations. Sized to max_degree initially; + // resized inside the loop if actual neighbor count exceeds this. + uint32_t buf_capacity = entity.max_degree(); + std::vector neighbor_ids(buf_capacity); + std::vector neighbor_vec_blocks; + neighbor_vec_blocks.reserve(buf_capacity); + std::vector dists(buf_capacity); + std::vector neighbor_vecs(buf_capacity); + + while (!candidates.empty() && !ctx->reach_scan_limit()) { + auto top = candidates.begin(); + node_id_t current_node = top->first; + dist_t current_dist = top->second; + + // Early termination: if the closest candidate is worse than the worst + // result in a full topk heap, we won't find anything better. + if (topk_heap.full() && current_dist > topk_heap[0].second) { + break; + } + + candidates.pop(); + + // Expand neighbors using typed access (no virtual dispatch) + const auto neighbors = entity.get_neighbors_typed(current_node); + ailego_prefetch(neighbors.data); + + // Resize buffers if this node has more neighbors than expected + if (neighbors.size() > buf_capacity) { + buf_capacity = neighbors.size(); + neighbor_ids.resize(buf_capacity); + dists.resize(buf_capacity); + neighbor_vecs.resize(buf_capacity); + } + + // Collect unvisited neighbors (reuse pre-allocated buffer) + uint32_t unvisited_count = 0; + for (uint32_t i = 0; i < neighbors.size(); ++i) { + node_id_t node = neighbors[i]; + if (visit.visited(node)) continue; + visit.set_visited(node); + neighbor_ids[unvisited_count++] = node; + } + if (unvisited_count == 0) continue; + + // Batch fetch vectors using typed access (reuse pre-allocated buffer) + neighbor_vec_blocks.clear(); + int ret = entity.get_vector_typed(neighbor_ids.data(), unvisited_count, + neighbor_vec_blocks); + if (ailego_unlikely(ret != 0)) break; + + // Prefetch for better cache performance + static constexpr uint32_t PREFETCH_BATCH = 2; + static constexpr uint32_t PREFETCH_STEP = 2; + for (uint32_t i = 0; + i < std::min(PREFETCH_BATCH * PREFETCH_STEP, unvisited_count); ++i) { + ailego_prefetch(neighbor_vec_blocks[i].data()); + } + + // Batch distance computation (reuse pre-allocated buffers) + for (uint32_t i = 0; i < unvisited_count; ++i) { + neighbor_vecs[i] = neighbor_vec_blocks[i].data(); + } + dc.batch_dist(neighbor_vecs.data(), unvisited_count, dists.data()); + + // Update candidates and topk. + // Unlike vanilla DiskANN which inserts all unvisited neighbors into + // the candidate queue unconditionally, we apply an early-pruning + // optimization: a neighbor is only inserted into the candidate queue + // (and topk_heap) if it could potentially improve the final results, + // i.e. either the topk heap is not yet full, or the neighbor is closer + // than the current worst result. This avoids expanding clearly + // unpromising branches and reduces the candidate queue size. + for (uint32_t i = 0; i < unvisited_count; ++i) { + node_id_t node = neighbor_ids[i]; + dist_t node_dist = dists[i]; + if ((!topk_heap.full()) || node_dist < topk_heap[0].second) { + candidates.emplace(node, node_dist); + if (!filter(node)) { + topk_heap.emplace(node, node_dist); + } + } + } + } +} + +// ============================================================================ +// robust_prune: Select up to max_degree diverse neighbors from candidates. +// +// Faithfully follows DiskANN's occlude_list algorithm: +// 1. Sort candidates by distance (ascending) +// 2. Truncate to max_occlusion_size (DiskANN's maxc parameter) +// 3. Multi-round alpha sweep: cur_alpha starts at 1.0, increments by *1.2 +// each round until reaching alpha. This progressively relaxes the +// occlusion criterion. +// 4. For each candidate, compute occlude_factor as: +// max over all selected neighbors p of: dist(query, candidate) / dist(p, +// candidate) +// If occlude_factor > cur_alpha, the candidate is occluded in this round. +// 5. After all rounds, if _saturate_graph and alpha > 1, fill remaining +// slots with any un-selected candidates. +// ============================================================================ +template +void VamanaAlgorithm::robust_prune(node_id_t id, + TopkHeap &candidates, + float alpha, uint32_t max_degree, + VamanaContext *ctx) const { + auto &result = ctx->prune_result(); + result.clear(); + + if (candidates.size() == 0) return; + + // Sort candidates by distance (ascending — closest first) + candidates.sort(); + + VamanaDistCalculator &dc = ctx->dist_calculator(); + size_t n = candidates.size(); + + // Truncate to max_occlusion_size (DiskANN's maxc parameter) + size_t maxc = entity_.max_occlusion_size(); + if (maxc > 0 && n > maxc) { + n = maxc; + } + + // Pre-cache all candidate vectors at once + auto &vec_cache = ctx->prune_vec_cache(); + vec_cache.resize(n); + for (size_t i = 0; i < n; ++i) { + vec_cache[i] = entity_.get_vector(candidates[i].first); + } + + // occlude_factor: tracks the maximum occlusion ratio for each candidate + // (DiskANN: occlude_factor[t] = max over selected p of + // dist_to_query/dist_to_p) + auto &occlude_factor = ctx->prune_occlude_factor(); + occlude_factor.assign(n, 0.0f); + + // Pre-allocated buffers for batch distance computation + auto &batch_vecs = ctx->batch_vecs_buf(); + auto &batch_dists = ctx->batch_dists_buf(); + auto &batch_indices = ctx->batch_indices_buf(); + batch_vecs.resize(n); + batch_dists.resize(n); + batch_indices.resize(n); + + // Multi-round alpha sweep (DiskANN: cur_alpha starts at 1, increments *1.2) + float cur_alpha = 1.0f; + while (cur_alpha <= alpha + 1e-6f && result.size() < max_degree) { + for (size_t i = 0; i < n && result.size() < max_degree; ++i) { + if (occlude_factor[i] > cur_alpha) { + continue; + } + + // Mark as consumed so it won't be reconsidered + occlude_factor[i] = std::numeric_limits::max(); + + // Skip self-loops + if (candidates[i].first == id) continue; + + const void *selected_vec = vec_cache[i]; + if (ailego_unlikely(selected_vec == nullptr)) continue; + + // Add this candidate as a neighbor + node_id_t candidate_id = candidates[i].first; + dist_t candidate_dist = candidates[i].second; + result.emplace_back(candidate_id, candidate_dist); + + // Update occlude_factor for remaining candidates + // Collect candidates that haven't been consumed yet + uint32_t batch_count = 0; + for (size_t j = i + 1; j < n; ++j) { + if (occlude_factor[j] > alpha) continue; // already fully occluded + if (ailego_unlikely(vec_cache[j] == nullptr)) continue; + batch_vecs[batch_count] = vec_cache[j]; + batch_indices[batch_count] = static_cast(j); + batch_count++; + } + + if (batch_count > 0) { + // Batch compute distances from selected candidate to remaining + dc.batch_dist_pair(selected_vec, batch_vecs.data(), batch_count, + batch_dists.data()); + + // DiskANN (L2/Cosine): + // occlude_factor[t] = max(occlude_factor[t], dist_to_query / + // dist_between) + // where dist_to_query = candidates[j].second (distance from query to j) + // dist_between = batch_dists[k] (distance from selected to j) + for (uint32_t k = 0; k < batch_count; ++k) { + uint32_t j = batch_indices[k]; + float dist_selected_to_candidate = batch_dists[k]; + if (dist_selected_to_candidate == 0.0f) { + occlude_factor[j] = std::numeric_limits::max(); + } else { + occlude_factor[j] = + std::max(occlude_factor[j], + candidates[j].second / dist_selected_to_candidate); + } + } + } + } + cur_alpha *= 1.2f; + } + + // Saturate graph: if enabled and alpha > 1, fill remaining slots with any + // un-selected candidates. This improves graph connectivity (better recall) + // at the cost of slightly more distance computations during search. + // Configurable via proxima.vamana.streamer.saturate_graph (default: false, + // matching DiskANN's default behavior). + if (entity_.saturate_graph() && alpha > 1.0f) { + for (size_t i = 0; i < n && result.size() < max_degree; ++i) { + if (candidates[i].first == id) continue; + bool already_selected = false; + for (const auto &r : result) { + if (r.first == candidates[i].first) { + already_selected = true; + break; + } + } + if (!already_selected) { + result.emplace_back(candidates[i].first, candidates[i].second); + } + } + } +} + +// ============================================================================ +// update_neighbors_and_reverse_links: For each new neighbor of `id`, +// add a reverse link from neighbor back to `id`. If the neighbor's degree +// exceeds max_degree, prune it using RobustPrune. +// ============================================================================ +template +void VamanaAlgorithm::update_neighbors_and_reverse_links( + node_id_t id, + const std::vector> &new_neighbors, + VamanaContext *ctx) { + for (const auto &[neighbor_id, dist] : new_neighbors) { + reverse_update_neighbor(id, neighbor_id, dist, ctx); + } +} + +// ============================================================================ +// reverse_update_neighbor: Add `id` as a neighbor of `neighbor_id`. +// If neighbor_id already has max_degree neighbors, collect all neighbors +// + the new one into a candidate set and RobustPrune. +// ============================================================================ +template +void VamanaAlgorithm::reverse_update_neighbor(node_id_t id, + node_id_t neighbor_id, + dist_t dist, + VamanaContext *ctx) { + std::lock_guard lock(lock_pool_[neighbor_id & kLockMask]); + + const Neighbors current_neighbors = entity_.get_neighbors(neighbor_id); + uint32_t current_size = current_neighbors.size(); + uint32_t max_deg = entity_.max_degree(); + + // Check if `id` is already a neighbor + for (uint32_t i = 0; i < current_size; ++i) { + if (current_neighbors[i] == id) return; + } + + if (current_size < max_deg) { + // Simply append and record distance + entity_.add_neighbor(neighbor_id, current_size, id); + entity_.set_neighbor_dist(neighbor_id, current_size, dist); + return; + } + + // Need to prune: collect current neighbors + new node into candidates + VamanaDistCalculator &dc = ctx->dist_calculator(); + + // Reuse update_heap from context instead of creating a new TopkHeap each time + TopkHeap &prune_candidates = ctx->update_heap(); + prune_candidates.clear(); + prune_candidates.limit(max_deg + 1); + + // Add existing neighbors — use cached distances when available + const dist_t *cached_dists = entity_.get_neighbor_dists(neighbor_id); + if (cached_dists != nullptr) { + // Fast path: read distances from storage, no recomputation needed + for (uint32_t i = 0; i < current_size; ++i) { + prune_candidates.emplace(current_neighbors[i], cached_dists[i]); + } + } else { + // Fallback: compute distances (first time or dist storage not loaded) + const void *neighbor_vec = entity_.get_vector(neighbor_id); + if (ailego_unlikely(neighbor_vec == nullptr)) return; + for (uint32_t i = 0; i < current_size; ++i) { + node_id_t nbr = current_neighbors[i]; + const void *nbr_vec = entity_.get_vector(nbr); + if (ailego_unlikely(nbr_vec == nullptr)) continue; + dist_t nbr_dist = dc.dist(neighbor_vec, nbr_vec); + prune_candidates.emplace(nbr, nbr_dist); + } + } + + // Add the new reverse link + prune_candidates.emplace(id, dist); + + // RobustPrune from neighbor_id's perspective + robust_prune(neighbor_id, prune_candidates, entity_.alpha(), max_deg, ctx); + + // Update neighbor_id's neighbor list and distances + const auto &prune_result = ctx->prune_result(); + entity_.update_neighbors(neighbor_id, prune_result); + entity_.update_neighbor_dists(neighbor_id, prune_result); +} + +// Explicit template instantiation for all entity types +template class VamanaAlgorithm; +template class VamanaAlgorithm; +template class VamanaAlgorithm; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/vamana/vamana_algorithm.h b/src/core/algorithm/vamana/vamana_algorithm.h new file mode 100644 index 0000000..c5b2ce9 --- /dev/null +++ b/src/core/algorithm/vamana/vamana_algorithm.h @@ -0,0 +1,115 @@ +// Copyright 2025-present the zvec 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 +#include +#include +#include +#include "vamana_context.h" +#include "vamana_dist_calculator.h" +#include "vamana_streamer_entity.h" + +namespace zvec { +namespace core { + +// Non-template base class providing a type-erased interface so that +// VamanaStreamer can hold a pointer without knowing the EntityType. +class VamanaAlgorithmBase { + public: + typedef std::unique_ptr UPointer; + + virtual ~VamanaAlgorithmBase() = default; + + virtual int cleanup() = 0; + + // Insert a new node into the Vamana graph. + // The node's vector must already be stored in the entity. + virtual int add_node(node_id_t id, VamanaContext *ctx) = 0; + + // Greedy search: find approximate nearest neighbors. + virtual int search(VamanaContext *ctx) const = 0; + + virtual int init() = 0; +}; + +// Vamana graph algorithm, templated on EntityType for hot-path optimization. +// EntityType should be VamanaMmapStreamerEntity, +// VamanaBufferPoolStreamerEntity, or VamanaContiguousStreamerEntity. +// +// Core operations: +// - GreedySearch: beam search from entry point, expanding best candidates +// - RobustPrune: select diverse neighbors using alpha-based pruning +// - add_node: insert + prune + reverse-link update +template +class VamanaAlgorithm : public VamanaAlgorithmBase { + public: + using MemBlockType = typename EntityType::MemoryBlock; + + explicit VamanaAlgorithm(EntityType &entity) + : entity_(entity), lock_pool_(kLockCnt) {} + + ~VamanaAlgorithm() override = default; + + int cleanup() override { + return 0; + } + + int init() override { + return 0; + } + + // Insert node `id` into the graph. Its vector must already be in the entity. + int add_node(node_id_t id, VamanaContext *ctx) override; + + // Greedy search from entry point. Results are stored in ctx->topk_heap(). + int search(VamanaContext *ctx) const override; + + private: + // GreedySearch: starting from entry_point, greedily expand the closest + // unvisited candidate until the search list is exhausted or scan limit + // is reached. Results accumulate in topk_heap. + void greedy_search(node_id_t entry_point, VamanaContext *ctx) const; + + // RobustPrune: given a candidate set (topk_heap), select up to max_degree + // diverse neighbors using alpha-based distance comparison. + // Result is stored in ctx->prune_result(). + void robust_prune(node_id_t id, TopkHeap &candidates, float alpha, + uint32_t max_degree, VamanaContext *ctx) const; + + // Update node's neighbors and handle reverse links. + void update_neighbors_and_reverse_links( + node_id_t id, + const std::vector> &new_neighbors, + VamanaContext *ctx); + + // Check if adding `id` as a reverse neighbor of `neighbor_id` requires + // pruning, and if so, prune neighbor_id's neighbor list. + void reverse_update_neighbor(node_id_t id, node_id_t neighbor_id, dist_t dist, + VamanaContext *ctx); + + private: + VamanaAlgorithm(const VamanaAlgorithm &) = delete; + VamanaAlgorithm &operator=(const VamanaAlgorithm &) = delete; + + static constexpr uint32_t kLockCnt{1U << 8}; + static constexpr uint32_t kLockMask{kLockCnt - 1U}; + + EntityType &entity_; + mutable ailego::SpinMutex spin_lock_{}; + std::vector lock_pool_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/vamana/vamana_context.cc b/src/core/algorithm/vamana/vamana_context.cc new file mode 100644 index 0000000..a3832e1 --- /dev/null +++ b/src/core/algorithm/vamana/vamana_context.cc @@ -0,0 +1,159 @@ +// Copyright 2025-present the zvec 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 "vamana_context.h" +#include +#include "vamana_params.h" + +namespace zvec { +namespace core { + +VamanaContext::VamanaContext(size_t dimension, + const IndexMetric::Pointer &metric, + const VamanaEntity::Pointer &entity) + : IndexContext(metric), + entity_(entity), + dc_(entity.get(), metric, dimension), + metric_(metric) {} + +VamanaContext::VamanaContext(const IndexMetric::Pointer &metric, + const VamanaEntity::Pointer &entity) + : IndexContext(metric), + entity_(entity), + dc_(entity.get(), metric), + metric_(metric) {} + +VamanaContext::~VamanaContext() { + visit_filter_.destroy(); +} + +int VamanaContext::init(ContextType type) { + int ret; + uint32_t doc_cnt; + + type_ = type; + results_.resize(1); + topk_heap_.limit(std::max(topk_, ef_)); + update_heap_.limit(entity_->max_degree()); + + switch (type) { + case kBuilderContext: + ret = visit_filter_.init(VisitFilter::ByteMap, entity_->doc_cnt(), + max_scan_num_, filter_negative_prob_); + if (ret != 0) { + LOG_ERROR("Create visit filter failed, mode %d", filter_mode_); + return ret; + } + candidates_.limit(max_scan_num_); + break; + + case kSearcherContext: + ret = visit_filter_.init(filter_mode_, entity_->doc_cnt(), max_scan_num_, + filter_negative_prob_); + if (ret != 0) { + LOG_ERROR("Create visit filter failed, mode %d", filter_mode_); + return ret; + } + candidates_.limit(max_scan_num_); + break; + + case kStreamerContext: + doc_cnt = entity_->doc_cnt(); + max_scan_num_ = compute_max_scan_num(doc_cnt); + reserve_max_doc_cnt_ = doc_cnt + compute_reserve_cnt(doc_cnt); + ret = visit_filter_.init(filter_mode_, reserve_max_doc_cnt_, + max_scan_num_, filter_negative_prob_); + if (ret != 0) { + LOG_ERROR("Create visit filter failed, mode %d", filter_mode_); + return ret; + } + candidates_.limit(max_scan_num_); + check_need_adjuct_ctx(); + break; + + default: + break; + } + + return 0; +} + +int VamanaContext::update_context(ContextType type, const IndexMeta &meta, + const IndexMetric::Pointer &metric, + const VamanaEntity::Pointer &entity, + uint32_t magic_num) { + if (magic_ == magic_num) { + return 0; + } + type_ = type; + entity_ = entity; + metric_ = metric; + magic_ = magic_num; + dc_.update(entity.get(), metric, meta.dimension()); + return 0; +} + +int VamanaContext::update(const ailego::Params ¶ms) { + uint32_t ef = ef_; + params.get(PARAM_VAMANA_STREAMER_EF, &ef); + ef_ = ef; + topk_heap_.limit(std::max(topk_, ef_)); + return 0; +} + +void VamanaContext::topk_to_result(uint32_t idx) { + if (force_padding_topk_ && !topk_heap_.full() && + topk_heap_.size() < entity_->doc_cnt()) { + this->fill_random_to_topk_full(); + } + if (ailego_unlikely(topk_heap_.size() == 0)) { + return; + } + + ailego_assert_with(idx < results_.size(), "invalid idx"); + int size = std::min(topk_, static_cast(topk_heap_.size())); + topk_heap_.sort(); + results_[idx].clear(); + + for (int i = 0; i < size; ++i) { + auto score = topk_heap_[i].second; + if (score > this->threshold()) { + break; + } + node_id_t id = topk_heap_[i].first; + if (fetch_vector_) { + results_[idx].emplace_back(entity_->get_key(id), score, id, + entity_->get_vector(id)); + } else { + results_[idx].emplace_back(entity_->get_key(id), score, id); + } + } +} + +void VamanaContext::fill_random_to_topk_full() { + std::mt19937 rng(42); + uint32_t doc_cnt = entity_->doc_cnt(); + uint32_t max_attempts = doc_cnt * 2; + uint32_t attempts = 0; + while (!topk_heap_.full() && doc_cnt > 0 && attempts < max_attempts) { + node_id_t random_id = rng() % doc_cnt; + if (entity_->get_key(random_id) != kInvalidKey) { + dist_t random_dist = dc_.dist(random_id); + topk_heap_.emplace_back(random_id, random_dist); + } + ++attempts; + } +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/vamana/vamana_context.h b/src/core/algorithm/vamana/vamana_context.h new file mode 100644 index 0000000..87e2903 --- /dev/null +++ b/src/core/algorithm/vamana/vamana_context.h @@ -0,0 +1,315 @@ +// Copyright 2025-present the zvec 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 +#include "utility/visit_filter.h" +#include "vamana_dist_calculator.h" +#include "vamana_entity.h" + +namespace zvec { +namespace core { + +class VamanaContext : public IndexContext { + public: + typedef std::unique_ptr Pointer; + + enum ContextType { + kUnknownContext = 0, + kSearcherContext = 1, + kBuilderContext = 2, + kStreamerContext = 3 + }; + + VamanaContext(size_t dimension, const IndexMetric::Pointer &metric, + const VamanaEntity::Pointer &entity); + + VamanaContext(const IndexMetric::Pointer &metric, + const VamanaEntity::Pointer &entity); + + virtual ~VamanaContext(); + + virtual void set_topk(uint32_t val) override { + topk_ = val; + topk_heap_.limit(std::max(val, ef_)); + } + + virtual const IndexDocumentList &result(void) const override { + return results_[0]; + } + + virtual const IndexDocumentList &result(size_t idx) const override { + return results_[idx]; + } + + virtual IndexDocumentList *mutable_result(size_t idx) override { + ailego_assert_with(idx < results_.size(), "invalid idx"); + return &results_[idx]; + } + + virtual uint32_t magic(void) const override { + return magic_; + } + + virtual void set_debug_mode(bool enable) override { + debug_mode_ = enable; + } + virtual bool debug_mode(void) const override { + return debug_mode_; + } + + virtual std::string debug_string(void) const override { + char buf[4096]; + size_t size = snprintf(buf, sizeof(buf), "scan_cnt=%zu", get_scan_num()); + return std::string(buf, size); + } + + virtual int update(const ailego::Params ¶ms) override; + + int init(ContextType type); + + int update_context(ContextType type, const IndexMeta &meta, + const IndexMetric::Pointer &metric, + const VamanaEntity::Pointer &entity, uint32_t magic_num); + + inline const VamanaEntity &get_entity() const { + return *entity_; + } + + inline void resize_results(size_t size) { + results_.resize(size); + } + + inline void topk_to_result() { + topk_to_result(0); + } + + void topk_to_result(uint32_t idx); + + inline void reset_query(const void *query) { + if (auto query_preprocess_func = index_metric_->get_query_preprocess_func(); + query_preprocess_func != nullptr) { + size_t dim = dc_.dimension(); + preprocess_buffer_.resize(dim); + memcpy(preprocess_buffer_.data(), query, dim); + query_preprocess_func(preprocess_buffer_.data(), dim); + query = preprocess_buffer_.data(); + } + dc_.reset_query(query); + dc_.clear_compare_cnt(); + } + + inline VamanaDistCalculator &dist_calculator() { + return dc_; + } + inline TopkHeap &topk_heap() { + return topk_heap_; + } + inline TopkHeap &update_heap() { + return update_heap_; + } + inline VisitFilter &visit_filter() { + return visit_filter_; + } + inline CandidateHeap &candidates() { + return candidates_; + } + + // Pre-allocated buffers for robust_prune optimization + inline std::vector &prune_vec_cache() { + return prune_vec_cache_; + } + inline std::vector &prune_active() { + return prune_active_; + } + inline std::vector &prune_occlude_factor() { + return prune_occlude_factor_; + } + inline std::vector> &prune_result() { + return prune_result_; + } + inline std::vector &batch_vecs_buf() { + return batch_vecs_buf_; + } + inline std::vector &batch_dists_buf() { + return batch_dists_buf_; + } + inline std::vector &batch_indices_buf() { + return batch_indices_buf_; + } + + inline void set_max_scan_num(uint32_t max_scan_num) { + max_scan_num_ = max_scan_num; + } + inline void set_ef(uint32_t v) { + ef_ = v; + } + + inline uint32_t ef() const { + return ef_; + } + inline void set_max_scan_ratio(float v) { + max_scan_ratio_ = v; + } + virtual void set_magic(uint32_t v) { + magic_ = v; + } + virtual void set_force_padding_topk(bool v) { + force_padding_topk_ = v; + } + + void set_bruteforce_threshold(uint32_t v) override { + bruteforce_threshold_ = v; + } + inline uint32_t get_bruteforce_threshold() const { + return bruteforce_threshold_; + } + + void set_fetch_vector(bool v) override { + fetch_vector_ = v; + } + bool fetch_vector() const override { + return fetch_vector_; + } + + void set_max_scan_limit(size_t v) { + max_scan_limit_ = v; + } + void set_min_scan_limit(size_t v) { + min_scan_limit_ = v; + } + + void set_filter_mode(VisitFilter::Mode mode) { + filter_mode_ = mode; + } + void set_filter_negative_probability(float prob) { + filter_negative_prob_ = prob; + } + + void reset(void) override { + dc_.clear(); + for (auto &it : results_) { + it.clear(); + } + IndexContext::reset_filter(); + IndexContext::reset_threshold(); + IndexContext::set_fetch_vector(false); + } + + inline void check_need_adjuct_ctx(void) { + check_need_adjuct_ctx(entity_->doc_cnt()); + } + + inline void check_need_adjuct_ctx(uint32_t doc_cnt) { + if (ailego_unlikely(doc_cnt + kTriggerReserveCnt > reserve_max_doc_cnt_)) { + while (doc_cnt + kTriggerReserveCnt > reserve_max_doc_cnt_) { + reserve_max_doc_cnt_ = + reserve_max_doc_cnt_ + compute_reserve_cnt(reserve_max_doc_cnt_); + } + uint32_t max_scan_cnt = compute_max_scan_num(reserve_max_doc_cnt_); + max_scan_num_ = max_scan_cnt; + visit_filter_.reset(reserve_max_doc_cnt_, max_scan_cnt); + candidates_.clear(); + candidates_.limit(max_scan_num_); + } + } + + inline size_t get_scan_num() const { + return dc_.compare_cnt(); + } + + inline uint64_t reach_scan_limit() const { + return dc_.compare_cnt() >= max_scan_num_; + } + + inline bool error() const { + return dc_.error(); + } + + inline void clear() { + dc_.clear(); + for (auto &it : results_) { + it.clear(); + } + } + + inline uint32_t topk() const override { + return topk_; + } + + inline void update_dist_caculator_distance( + const IndexMetric::MatrixDistance &distance, + const IndexMetric::MatrixBatchDistance &batch_distance) { + dc_.update_distance(distance, batch_distance); + } + + private: + void fill_random_to_topk_full(void); + + inline size_t compute_reserve_cnt(uint32_t cur_doc) const { + if (cur_doc > kMaxReserveDocCnt) return kMaxReserveDocCnt; + if (cur_doc < kMinReserveDocCnt) return kMinReserveDocCnt; + return cur_doc; + } + + inline uint32_t compute_max_scan_num(uint32_t max_doc_cnt) const { + uint32_t max_scan = max_doc_cnt * max_scan_ratio_; + if (max_scan < min_scan_limit_) max_scan = min_scan_limit_; + if (max_scan > max_scan_limit_) max_scan = max_scan_limit_; + return max_scan; + } + + constexpr static uint32_t kTriggerReserveCnt = 4096UL; + constexpr static uint32_t kMinReserveDocCnt = 4096UL; + constexpr static uint32_t kMaxReserveDocCnt = 128 * 1024UL; + + VamanaEntity::Pointer entity_; + VamanaDistCalculator dc_; + IndexMetric::Pointer metric_; + + bool debug_mode_{false}; + bool force_padding_topk_{false}; + uint32_t max_scan_num_{0}; + uint32_t reserve_max_doc_cnt_{kMinReserveDocCnt}; + uint32_t topk_{0}; + uint32_t ef_{VamanaEntity::kDefaultEf}; + float max_scan_ratio_{VamanaEntity::kDefaultScanRatio}; + size_t max_scan_limit_{VamanaEntity::kDefaultMaxScanLimit}; + size_t min_scan_limit_{VamanaEntity::kDefaultMinScanLimit}; + uint32_t magic_{0U}; + std::vector results_{}; + TopkHeap topk_heap_{}; + TopkHeap update_heap_{}; + CandidateHeap candidates_{}; + VisitFilter visit_filter_{}; + uint32_t bruteforce_threshold_{}; + bool fetch_vector_{false}; + uint32_t type_{kUnknownContext}; + std::string preprocess_buffer_; + + // Pre-allocated buffers for robust_prune optimization + std::vector prune_vec_cache_; + std::vector prune_active_; + std::vector prune_occlude_factor_; + std::vector> prune_result_; + std::vector batch_vecs_buf_; + std::vector batch_dists_buf_; + std::vector batch_indices_buf_; + + VisitFilter::Mode filter_mode_{VisitFilter::ByteMap}; + float filter_negative_prob_{VamanaEntity::kDefaultBFNegativeProbability}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/vamana/vamana_dist_calculator.h b/src/core/algorithm/vamana/vamana_dist_calculator.h new file mode 100644 index 0000000..c78f638 --- /dev/null +++ b/src/core/algorithm/vamana/vamana_dist_calculator.h @@ -0,0 +1,197 @@ +// Copyright 2025-present the zvec 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 +#include "vamana_entity.h" + +namespace zvec { +namespace core { + +class VamanaDistCalculator { + public: + typedef std::shared_ptr Pointer; + + VamanaDistCalculator(const VamanaEntity *entity, + const IndexMetric::Pointer &metric, uint32_t dim) + : entity_(entity), + distance_(metric->distance()), + batch_distance_(metric->batch_distance()), + query_(nullptr), + dim_(dim), + compare_cnt_(0) {} + + VamanaDistCalculator(const VamanaEntity *entity, + const IndexMetric::Pointer &metric, uint32_t dim, + const void *query) + : entity_(entity), + distance_(metric->distance()), + batch_distance_(metric->batch_distance()), + query_(query), + dim_(dim), + compare_cnt_(0) {} + + VamanaDistCalculator(const VamanaEntity *entity, + const IndexMetric::Pointer &metric) + : entity_(entity), + distance_(metric->distance()), + batch_distance_(metric->batch_distance()), + query_(nullptr), + dim_(0), + compare_cnt_(0) {} + + void update(const VamanaEntity *entity, const IndexMetric::Pointer &metric) { + entity_ = entity; + distance_ = metric->distance(); + batch_distance_ = metric->batch_distance(); + } + + void update(const VamanaEntity *entity, const IndexMetric::Pointer &metric, + uint32_t dim) { + entity_ = entity; + distance_ = metric->distance(); + batch_distance_ = metric->batch_distance(); + dim_ = dim; + } + + inline void update_distance( + const IndexMetric::MatrixDistance &distance, + const IndexMetric::MatrixBatchDistance &batch_distance) { + distance_ = distance; + batch_distance_ = batch_distance; + } + + inline void reset_query(const void *query) { + error_ = false; + query_ = query; + } + + inline dist_t dist(const void *vec_lhs, const void *vec_rhs) { + if (ailego_unlikely(vec_lhs == nullptr || vec_rhs == nullptr)) { + LOG_ERROR("Nullptr of dense vector"); + error_ = true; + return 0.0f; + } + float score{0.0f}; + distance_(vec_lhs, vec_rhs, dim_, &score); + return score; + } + + inline dist_t dist(const void *vec) { + compare_cnt_++; + return dist(vec, query_); + } + + inline dist_t dist(node_id_t id) { + compare_cnt_++; + const void *feat = entity_->get_vector(id); + if (ailego_unlikely(feat == nullptr)) { + LOG_ERROR("Get nullptr vector, id=%u", id); + error_ = true; + return 0.0f; + } + return dist(feat, query_); + } + + inline dist_t dist(node_id_t lhs, node_id_t rhs) { + compare_cnt_++; + const void *feat = entity_->get_vector(lhs); + const void *query = entity_->get_vector(rhs); + if (ailego_unlikely(feat == nullptr || query == nullptr)) { + LOG_ERROR("Get nullptr vector"); + error_ = true; + return 0.0f; + } + return dist(feat, query); + } + + inline void batch_dist(const void **vecs, uint32_t count, float *dists) { + compare_cnt_ += count; + batch_distance_(vecs, query_, count, dim_, dists); + } + + // Single-node batch distance: compute distance between query and a stored + // node using batch_distance_. Consistent with HnswDistCalculator::batch_dist. + inline dist_t batch_dist(node_id_t id) { + compare_cnt_++; + const void *feat = entity_->get_vector(id); + if (ailego_unlikely(feat == nullptr)) { + LOG_ERROR("Get nullptr vector, id=%u", id); + error_ = true; + return 0.0f; + } + dist_t score = 0; + batch_distance_(&feat, query_, 1, dim_, &score); + return score; + } + + // Batch distance computation between a base vector and multiple target + // vectors. Does NOT use query_ and does NOT increment compare_cnt. Used for + // inter-candidate distance computation in robust_prune. + // + // Uses the single distance function (distance_) in a loop rather than + // batch_distance_, because batch_distance_ (turbo AVX512-VNNI) expects + // the second argument to be a preprocessed uint8 query (+128 shift), + // while base_vec here is a raw int8 stored vector. The single distance + // function (AVX2 sign/abs trick) correctly handles two raw int8 inputs. + inline void batch_dist_pair(const void *base_vec, const void **vecs, + uint32_t count, float *dists) { + for (uint32_t i = 0; i < count; ++i) { + distance_(base_vec, vecs[i], dim_, &dists[i]); + } + } + + dist_t operator()(const void *vec) { + return dist(vec); + } + dist_t operator()(node_id_t i) { + return dist(i); + } + dist_t operator()(node_id_t lhs, node_id_t rhs) { + return dist(lhs, rhs); + } + + inline void clear() { + compare_cnt_ = 0; + error_ = false; + } + + inline void clear_compare_cnt() { + compare_cnt_ = 0; + } + inline bool error() const { + return error_; + } + inline uint32_t compare_cnt() const { + return compare_cnt_; + } + inline uint32_t dimension() const { + return dim_; + } + + private: + VamanaDistCalculator(const VamanaDistCalculator &) = delete; + VamanaDistCalculator &operator=(const VamanaDistCalculator &) = delete; + + const VamanaEntity *entity_; + IndexMetric::MatrixDistance distance_; + IndexMetric::MatrixBatchDistance batch_distance_; + const void *query_; + uint32_t dim_; + uint32_t compare_cnt_; + bool error_{false}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/vamana/vamana_entity.cc b/src/core/algorithm/vamana/vamana_entity.cc new file mode 100644 index 0000000..4114dc6 --- /dev/null +++ b/src/core/algorithm/vamana/vamana_entity.cc @@ -0,0 +1,203 @@ +// Copyright 2025-present the zvec 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 "vamana_entity.h" +#include + +namespace zvec { +namespace core { + +const std::string VamanaEntity::kGraphHeaderSegmentId = "vamana.graph.header"; +const std::string VamanaEntity::kGraphFeaturesSegmentId = + "vamana.graph.features"; +const std::string VamanaEntity::kGraphKeysSegmentId = "vamana.graph.keys"; +const std::string VamanaEntity::kGraphNeighborsSegmentId = + "vamana.graph.neighbors"; +const std::string VamanaEntity::kGraphOffsetsSegmentId = "vamana.graph.offsets"; +const std::string VamanaEntity::kGraphMappingSegmentId = "vamana.graph.mapping"; +const std::string VamanaEntity::kGraphNeighborDistsSegmentId = + "vamana.graph.neighbor_dists"; + +int VamanaEntity::CalcAndAddPadding(const IndexDumper::Pointer &dumper, + size_t data_size, size_t *padding_size) { + *padding_size = AlignSize(data_size) - data_size; + if (*padding_size == 0) { + return 0; + } + std::string padding(*padding_size, '\0'); + if (dumper->write(padding.data(), *padding_size) != *padding_size) { + LOG_ERROR("Append padding failed, size %lu", *padding_size); + return IndexError_WriteData; + } + return 0; +} + +int64_t VamanaEntity::dump_segment(const IndexDumper::Pointer &dumper, + const std::string &segment_id, + const void *data, size_t size) const { + size_t len = dumper->write(data, size); + if (len != size) { + LOG_ERROR("Dump segment %s data failed, expect: %lu, actual: %lu", + segment_id.c_str(), size, len); + return IndexError_WriteData; + } + + size_t padding_size = AlignSize(size) - size; + if (padding_size > 0) { + std::string padding(padding_size, '\0'); + if (dumper->write(padding.data(), padding_size) != padding_size) { + LOG_ERROR("Append padding failed, size %lu", padding_size); + return IndexError_WriteData; + } + } + + uint32_t crc = ailego::Crc32c::Hash(data, size); + int ret = dumper->append(segment_id, size, padding_size, crc); + if (ret != 0) { + LOG_ERROR("Dump segment %s meta failed, ret=%d", segment_id.c_str(), ret); + return ret; + } + + return len + padding_size; +} + +int64_t VamanaEntity::dump_header(const IndexDumper::Pointer &dumper, + const VamanaHeader &hd) const { + return dump_segment(dumper, kGraphHeaderSegmentId, &hd.graph, hd.graph.size); +} + +int64_t VamanaEntity::dump_vectors( + const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping) const { + size_t total_size = doc_cnt() * vector_size(); + std::vector buffer(total_size); + + for (node_id_t i = 0; i < doc_cnt(); ++i) { + node_id_t old_id = reorder_mapping[i]; + const void *vec = get_vector(old_id); + if (vec == nullptr) { + LOG_ERROR("Get vector failed for node %u", old_id); + return IndexError_ReadData; + } + memcpy(buffer.data() + static_cast(i) * vector_size(), vec, + vector_size()); + } + + return dump_segment(dumper, kGraphFeaturesSegmentId, buffer.data(), + total_size); +} + +int64_t VamanaEntity::dump_neighbors( + const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping, + const std::vector &neighbor_mapping) const { + size_t nbr_size = neighbors_size(); + size_t total_size = doc_cnt() * nbr_size; + std::vector buffer(total_size, 0); + + for (node_id_t i = 0; i < doc_cnt(); ++i) { + node_id_t old_id = reorder_mapping[i]; + const Neighbors nbrs = get_neighbors(old_id); + + auto *hd = reinterpret_cast( + buffer.data() + static_cast(i) * nbr_size); + hd->neighbor_cnt = nbrs.size(); + for (uint32_t j = 0; j < nbrs.size(); ++j) { + hd->neighbors[j] = neighbor_mapping[nbrs[j]]; + } + } + + return dump_segment(dumper, kGraphNeighborsSegmentId, buffer.data(), + total_size); +} + +int64_t VamanaEntity::dump_neighbor_dists( + const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping) const { + if (!dist_storage_loaded()) { + // No distance data to dump — this is fine for read-only indices + return 0; + } + + uint32_t max_deg = static_cast(max_degree()); + size_t dist_entry = max_deg * sizeof(dist_t); + size_t total_size = doc_cnt() * dist_entry; + std::vector buffer(total_size, 0); + + for (node_id_t i = 0; i < doc_cnt(); ++i) { + node_id_t old_id = reorder_mapping[i]; + const dist_t *dists = get_neighbor_dists(old_id); + if (dists != nullptr) { + memcpy(buffer.data() + static_cast(i) * dist_entry, dists, + dist_entry); + } + } + + return dump_segment(dumper, kGraphNeighborDistsSegmentId, buffer.data(), + total_size); +} + +int64_t VamanaEntity::dump_mapping_segment(const IndexDumper::Pointer &dumper, + const key_t *keys) const { + size_t total_size = doc_cnt() * sizeof(key_t); + return dump_segment(dumper, kGraphMappingSegmentId, keys, total_size); +} + +void VamanaEntity::reshuffle_vectors(std::vector *n2o_mapping, + std::vector *o2n_mapping, + key_t *keys) const { + uint32_t count = doc_cnt(); + n2o_mapping->resize(count); + o2n_mapping->resize(count); + + // Simple identity mapping for now (can be optimized with BFS traversal) + for (uint32_t i = 0; i < count; ++i) { + (*n2o_mapping)[i] = i; + (*o2n_mapping)[i] = i; + keys[i] = get_key(i); + } +} + +int64_t VamanaEntity::dump_segments(const IndexDumper::Pointer &dumper, + key_t *keys) const { + std::vector n2o_mapping; + std::vector o2n_mapping; + reshuffle_vectors(&n2o_mapping, &o2n_mapping, keys); + + // Remap entry point + VamanaHeader dump_header_copy = header_; + if (dump_header_copy.graph.entry_point != kInvalidNodeId) { + dump_header_copy.graph.entry_point = + o2n_mapping[dump_header_copy.graph.entry_point]; + } + + int64_t hd_size = dump_header(dumper, dump_header_copy); + if (hd_size < 0) return hd_size; + + int64_t vec_size = dump_vectors(dumper, n2o_mapping); + if (vec_size < 0) return vec_size; + + int64_t nbr_size = dump_neighbors(dumper, n2o_mapping, o2n_mapping); + if (nbr_size < 0) return nbr_size; + + int64_t map_size = dump_mapping_segment(dumper, keys); + if (map_size < 0) return map_size; + + int64_t dist_size = dump_neighbor_dists(dumper, n2o_mapping); + if (dist_size < 0) return dist_size; + + return hd_size + vec_size + nbr_size + map_size + dist_size; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/vamana/vamana_entity.h b/src/core/algorithm/vamana/vamana_entity.h new file mode 100644 index 0000000..fca4159 --- /dev/null +++ b/src/core/algorithm/vamana/vamana_entity.h @@ -0,0 +1,352 @@ +// Copyright 2025-present the zvec 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 +#include +#include +#include +#include +#include +#include +#include + +// Reuse typed MemoryBlock and NeighborsT from hnsw_entity.h +#include "algorithm/hnsw/hnsw_entity.h" + +namespace zvec { +namespace core { + +// Vamana graph header — single-layer graph (no hierarchical levels) +struct VamanaGraphHeader { + uint32_t size; + uint32_t version; + uint32_t graph_type; + uint32_t doc_count; + uint32_t vector_size; + uint32_t node_size; + uint32_t max_degree; // R: maximum out-degree + uint32_t search_list_size; // L: search list size for construction + uint32_t max_occlusion_size; // C: max candidate size for RobustPrune + uint32_t entry_point; // medoid node id + uint32_t options; + uint32_t reserved_pad; + float alpha; // alpha parameter for RobustPrune + uint8_t reserved_[4076]; +}; + +static_assert(sizeof(VamanaGraphHeader) % 32 == 0, + "VamanaGraphHeader must be aligned with 32 bytes"); + +struct VamanaHeader { + VamanaHeader() { + clear(); + } + + VamanaHeader(const VamanaHeader &header) { + memcpy(this, &header, sizeof(header)); + } + + VamanaHeader &operator=(const VamanaHeader &header) { + memcpy(this, &header, sizeof(header)); + return *this; + } + + void inline reset() { + graph.doc_count = 0U; + graph.entry_point = kInvalidNodeId; + } + + void inline clear() { + memset(this, 0, sizeof(VamanaHeader)); + graph.entry_point = kInvalidNodeId; + graph.size = sizeof(VamanaGraphHeader); + graph.alpha = 1.2f; + } + + size_t max_degree() const { + return graph.max_degree; + } + size_t vector_size() const { + return graph.vector_size; + } + size_t search_list_size() const { + return graph.search_list_size; + } + size_t max_occlusion_size() const { + return graph.max_occlusion_size; + } + float alpha() const { + return graph.alpha; + } + node_id_t entry_point() const { + return graph.entry_point; + } + node_id_t doc_cnt() const { + return graph.doc_count; + } + + VamanaGraphHeader graph; +}; + +// VamanaEntity: base class for Vamana graph data management +class VamanaEntity { + public: + VamanaEntity() {} + VamanaEntity(const VamanaHeader &hd) { + header_ = hd; + } + virtual ~VamanaEntity() {} + + typedef std::shared_ptr Pointer; + + // Options bit flags (stored in VamanaGraphHeader::options) + static constexpr uint32_t kOptionSaturateGraph = 1U << 0; + + // Default constants + static constexpr uint32_t kDefaultMaxDegree = 64; + static constexpr uint32_t kDefaultSearchListSize = 100; + static constexpr uint32_t kDefaultMaxOcclusionSize = 750; + static constexpr float kDefaultAlpha = 1.2f; + static constexpr bool kDefaultSaturateGraph = false; + static constexpr uint32_t kDefaultEf = 200; + static constexpr float kDefaultScanRatio = 0.1f; + static constexpr uint32_t kDefaultBruteForceThreshold = 1000U; + static constexpr uint32_t kDefaultDocsHardLimit = 1 << 30U; + static constexpr float kDefaultDocsSoftLimitRatio = 0.9f; + static constexpr size_t kMaxChunkSize = 0xFFFFFFFF; + static constexpr size_t kDefaultChunkSize = 2UL * 1024UL * 1024UL; + static constexpr size_t kDefaultMaxChunkCnt = 50000UL; + static constexpr uint32_t kDefaultMinScanLimit = 10000; + static constexpr uint32_t kDefaultMaxScanLimit = + std::numeric_limits::max(); + static constexpr float kDefaultBFNegativeProbability = 0.001f; + + inline size_t max_degree() const { + return header_.graph.max_degree; + } + inline node_id_t *mutable_doc_cnt() { + return &header_.graph.doc_count; + } + inline node_id_t doc_cnt() const { + return header_.graph.doc_count; + } + inline float alpha() const { + return header_.graph.alpha; + } + inline size_t search_list_size() const { + return header_.graph.search_list_size; + } + inline size_t max_occlusion_size() const { + return header_.graph.max_occlusion_size; + } + inline node_id_t entry_point() const { + return header_.graph.entry_point; + } + inline size_t vector_size() const { + return header_.graph.vector_size; + } + inline size_t node_size() const { + return header_.graph.node_size; + } + + void set_vector_size(size_t size) { + header_.graph.vector_size = size; + } + void set_max_degree(uint32_t val) { + header_.graph.max_degree = val; + } + void set_search_list_size(uint32_t val) { + header_.graph.search_list_size = val; + } + void set_max_occlusion_size(uint32_t val) { + header_.graph.max_occlusion_size = val; + } + void set_alpha(float val) { + header_.graph.alpha = val; + } + + inline bool saturate_graph() const { + return (header_.graph.options & kOptionSaturateGraph) != 0; + } + void set_saturate_graph(bool val) { + if (val) { + header_.graph.options |= kOptionSaturateGraph; + } else { + header_.graph.options &= ~kOptionSaturateGraph; + } + } + + // Neighbor size: NeighborsHeader + max_degree * sizeof(node_id_t) + inline size_t neighbors_size() const { + return sizeof(NeighborsHeader) + max_degree() * sizeof(node_id_t); + } + + virtual void update_entry_point(node_id_t ep) { + header_.graph.entry_point = ep; + } + + virtual int cleanup() { + header_.clear(); + return 0; + } + + virtual const VamanaEntity::Pointer clone() const { + return VamanaEntity::Pointer(); + } + + // Pure virtual interface + virtual key_t get_key(node_id_t id) const = 0; + virtual const void *get_vector(node_id_t id) const = 0; + virtual int get_vector(const node_id_t id, + IndexStorage::MemoryBlock &block) const = 0; + virtual int get_vector(const node_id_t *ids, uint32_t count, + const void **vecs) const = 0; + virtual int get_vector( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const = 0; + virtual const Neighbors get_neighbors(node_id_t id) const = 0; + + virtual int add_vector(key_t key, const void *vec, node_id_t *id) { + return IndexError_NotImplemented; + } + virtual int add_vector_with_id(node_id_t id, const void *vec) { + return IndexError_NotImplemented; + } + virtual int update_neighbors( + node_id_t id, + const std::vector> &neighbors) { + return IndexError_NotImplemented; + } + virtual void add_neighbor(node_id_t id, uint32_t size, + node_id_t neighbor_id) {} + + // --- Neighbor distance storage (CSR-like, lazy-loaded) --- + // Each node has max_degree dist_t slots, the i-th slot stores the distance + // from this node to its i-th neighbor. Only allocated/loaded when needed + // (first write operation). Search-only paths never touch this data. + + // Ensure distance storage is allocated/loaded. Must be called before + // any get/set neighbor dist operations. Thread-safe (idempotent). + virtual int ensure_dist_storage() { + return 0; + } + + // Whether distance storage is currently loaded + virtual bool dist_storage_loaded() const { + return false; + } + + // Get pointer to the distance array for node `id`. + // Returns nullptr if dist storage is not loaded. + virtual const dist_t *get_neighbor_dists(node_id_t id) const { + return nullptr; + } + + // Update all neighbor distances for node `id` from a prune result. + virtual void update_neighbor_dists( + node_id_t id, + const std::vector> &neighbors) {} + + // Set the distance for the `idx`-th neighbor of node `id`. + virtual void set_neighbor_dist(node_id_t id, uint32_t idx, dist_t dist) {} + + virtual int dump(const IndexDumper::Pointer &dumper) { + return IndexError_NotImplemented; + } + + virtual const void *get_vector_by_key(uint64_t /*key*/) const { + return nullptr; + } + virtual int get_vector_by_key(const key_t /*key*/, + IndexStorage::MemoryBlock & /*block*/) const { + return IndexError_NotImplemented; + } + + static int CalcAndAddPadding(const IndexDumper::Pointer &dumper, + size_t data_size, size_t *padding_size); + + protected: + inline const VamanaHeader &header() const { + return header_; + } + inline VamanaHeader *mutable_header() { + return &header_; + } + inline size_t header_size() const { + return sizeof(header_); + } + + void set_node_size(size_t size) { + header_.graph.node_size = size; + } + + int64_t dump_segments(const IndexDumper::Pointer &dumper, key_t *keys) const; + + int64_t dump_segment(const IndexDumper::Pointer &dumper, + const std::string &segment_id, const void *data, + size_t size) const; + + int64_t dump_header(const IndexDumper::Pointer &dumper, + const VamanaHeader &hd) const; + + int64_t dump_vectors(const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping) const; + + int64_t dump_neighbors(const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping, + const std::vector &neighbor_mapping) const; + + int64_t dump_neighbor_dists( + const IndexDumper::Pointer &dumper, + const std::vector &reorder_mapping) const; + + int64_t dump_mapping_segment(const IndexDumper::Pointer &dumper, + const key_t *keys) const; + + void reshuffle_vectors(std::vector *n2o_mapping, + std::vector *o2n_mapping, + key_t *keys) const; + + static inline size_t AlignSize(size_t size) { + return (size + 0x1F) & (~0x1F); + } + + static inline size_t AlignPageSize(size_t size) { + size_t page_mask = ailego::MemoryHelper::PageSize() - 1; + return (size + page_mask) & (~page_mask); + } + + static inline size_t AlignHugePageSize(size_t size) { + size_t page_mask = ailego::MemoryHelper::HugePageSize() - 1; + return (size + page_mask) & (~page_mask); + } + + public: + const static std::string kGraphHeaderSegmentId; + const static std::string kGraphFeaturesSegmentId; + const static std::string kGraphKeysSegmentId; + const static std::string kGraphNeighborsSegmentId; + const static std::string kGraphOffsetsSegmentId; + const static std::string kGraphMappingSegmentId; + const static std::string kGraphNeighborDistsSegmentId; + + static constexpr uint32_t kRevision = 0U; + + protected: + VamanaHeader header_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/vamana/vamana_index_provider.h b/src/core/algorithm/vamana/vamana_index_provider.h new file mode 100644 index 0000000..ab70ba5 --- /dev/null +++ b/src/core/algorithm/vamana/vamana_index_provider.h @@ -0,0 +1,117 @@ +// Copyright 2025-present the zvec 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 +#include +#include +#include "vamana_entity.h" + +namespace zvec { +namespace core { + +class VamanaIndexProvider : public IndexProvider { + public: + VamanaIndexProvider(const IndexMeta &meta, + const VamanaEntity::Pointer &entity, + const std::string &owner) + : meta_(meta), entity_(entity), owner_class_(owner) {} + + VamanaIndexProvider(const VamanaIndexProvider &) = delete; + VamanaIndexProvider &operator=(const VamanaIndexProvider &) = delete; + + public: + IndexProvider::Iterator::Pointer create_iterator() override { + return VamanaIndexProvider::Iterator::Pointer(new (std::nothrow) + Iterator(entity_)); + } + + size_t count(void) const override { + return entity_->doc_cnt(); + } + + size_t dimension(void) const override { + return meta_.dimension(); + } + + IndexMeta::DataType data_type(void) const override { + return meta_.data_type(); + } + + size_t element_size(void) const override { + return meta_.element_size(); + } + + const void *get_vector(uint64_t key) const override { + return entity_->get_vector_by_key(key); + } + + int get_vector(const uint64_t key, + IndexStorage::MemoryBlock &block) const override { + return entity_->get_vector_by_key(key, block); + } + + const std::string &owner_class(void) const override { + return owner_class_; + } + + private: + class Iterator : public IndexProvider::Iterator { + public: + Iterator(const VamanaEntity::Pointer &entity) + : entity_(entity), cur_id_(0U) { + cur_id_ = get_next_valid_id(0); + } + + const void *data(void) const override { + return entity_->get_vector(cur_id_); + } + + bool is_valid(void) const override { + return cur_id_ < entity_->doc_cnt(); + } + + uint64_t key(void) const override { + return entity_->get_key(cur_id_); + } + + void next(void) override { + cur_id_ = get_next_valid_id(cur_id_ + 1); + } + + void reset(void) { + cur_id_ = get_next_valid_id(0); + } + + private: + node_id_t get_next_valid_id(node_id_t start_id) { + for (node_id_t i = start_id; i < entity_->doc_cnt(); i++) { + if (entity_->get_key(i) != kInvalidNodeId) { + return i; + } + } + return kInvalidNodeId; + } + + const VamanaEntity::Pointer entity_; + node_id_t cur_id_; + }; + + const IndexMeta &meta_; + const VamanaEntity::Pointer entity_; + const std::string owner_class_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/vamana/vamana_params.h b/src/core/algorithm/vamana/vamana_params.h new file mode 100644 index 0000000..68ffaf8 --- /dev/null +++ b/src/core/algorithm/vamana/vamana_params.h @@ -0,0 +1,79 @@ +// Copyright 2025-present the zvec 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 + +namespace zvec { +namespace core { + +// Builder parameters +static const std::string PARAM_VAMANA_BUILDER_THREAD_COUNT( + "proxima.vamana.builder.thread_count"); +static const std::string PARAM_VAMANA_BUILDER_MEMORY_QUOTA( + "proxima.vamana.builder.memory_quota"); +static const std::string PARAM_VAMANA_BUILDER_MAX_DEGREE( + "proxima.vamana.builder.max_degree"); +static const std::string PARAM_VAMANA_BUILDER_SEARCH_LIST_SIZE( + "proxima.vamana.builder.search_list_size"); +static const std::string PARAM_VAMANA_BUILDER_ALPHA( + "proxima.vamana.builder.alpha"); +static const std::string PARAM_VAMANA_BUILDER_MAX_OCCLUSION_SIZE( + "proxima.vamana.builder.max_occlusion_size"); + +// Searcher parameters +static const std::string PARAM_VAMANA_SEARCHER_SEARCH_LIST_SIZE( + "proxima.vamana.searcher.search_list_size"); +static const std::string PARAM_VAMANA_SEARCHER_BRUTE_FORCE_THRESHOLD( + "proxima.vamana.searcher.brute_force_threshold"); + +// Streamer parameters +static const std::string PARAM_VAMANA_STREAMER_MAX_DEGREE( + "proxima.vamana.streamer.max_degree"); +static const std::string PARAM_VAMANA_STREAMER_SEARCH_LIST_SIZE( + "proxima.vamana.streamer.search_list_size"); +static const std::string PARAM_VAMANA_STREAMER_ALPHA( + "proxima.vamana.streamer.alpha"); +static const std::string PARAM_VAMANA_STREAMER_MAX_OCCLUSION_SIZE( + "proxima.vamana.streamer.max_occlusion_size"); +static const std::string PARAM_VAMANA_STREAMER_EF("proxima.vamana.streamer.ef"); +static const std::string PARAM_VAMANA_STREAMER_BRUTE_FORCE_THRESHOLD( + "proxima.vamana.streamer.brute_force_threshold"); +static const std::string PARAM_VAMANA_STREAMER_MAX_SCAN_RATIO( + "proxima.vamana.streamer.max_scan_ratio"); +static const std::string PARAM_VAMANA_STREAMER_DOCS_HARD_LIMIT( + "proxima.vamana.streamer.docs_hard_limit"); +static const std::string PARAM_VAMANA_STREAMER_DOCS_SOFT_LIMIT( + "proxima.vamana.streamer.docs_soft_limit"); +static const std::string PARAM_VAMANA_STREAMER_MAX_INDEX_SIZE( + "proxima.vamana.streamer.max_index_size"); +static const std::string PARAM_VAMANA_STREAMER_CHUNK_SIZE( + "proxima.vamana.streamer.chunk_size"); +static const std::string PARAM_VAMANA_STREAMER_GET_VECTOR_ENABLE( + "proxima.vamana.streamer.get_vector_enable"); +static const std::string PARAM_VAMANA_STREAMER_FORCE_PADDING_RESULT_ENABLE( + "proxima.vamana.streamer.force_padding_result_enable"); +static const std::string PARAM_VAMANA_STREAMER_USE_ID_MAP( + "proxima.vamana.streamer.use_id_map"); +static const std::string PARAM_VAMANA_STREAMER_MAX_SCAN_LIMIT( + "proxima.vamana.streamer.max_scan_limit"); +static const std::string PARAM_VAMANA_STREAMER_MIN_SCAN_LIMIT( + "proxima.vamana.streamer.min_scan_limit"); +static const std::string PARAM_VAMANA_STREAMER_SATURATE_GRAPH( + "proxima.vamana.streamer.saturate_graph"); +static const std::string PARAM_VAMANA_STREAMER_USE_CONTIGUOUS_MEMORY( + "proxima.vamana.streamer.use_contiguous_memory"); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/vamana/vamana_streamer.cc b/src/core/algorithm/vamana/vamana_streamer.cc new file mode 100644 index 0000000..ae935eb --- /dev/null +++ b/src/core/algorithm/vamana/vamana_streamer.cc @@ -0,0 +1,714 @@ +// Copyright 2025-present the zvec 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 "vamana_streamer.h" +#include +#include +#include +#include "vamana_algorithm.h" +#include "vamana_context.h" +#include "vamana_dist_calculator.h" +#include "vamana_index_provider.h" + +namespace zvec { +namespace core { + +VamanaStreamer::VamanaStreamer() = default; + +VamanaStreamer::~VamanaStreamer() { + if (state_ == STATE_INITED) { + this->cleanup(); + } +} + +int VamanaStreamer::init(const IndexMeta &imeta, const ailego::Params ¶ms) { + meta_ = imeta; + meta_.set_streamer("VamanaStreamer", VamanaEntity::kRevision, params); + + params.get(PARAM_VAMANA_STREAMER_MAX_INDEX_SIZE, &max_index_size_); + params.get(PARAM_VAMANA_STREAMER_MAX_DEGREE, &max_degree_); + params.get(PARAM_VAMANA_STREAMER_SEARCH_LIST_SIZE, &search_list_size_); + params.get(PARAM_VAMANA_STREAMER_ALPHA, &alpha_); + params.get(PARAM_VAMANA_STREAMER_MAX_OCCLUSION_SIZE, &max_occlusion_size_); + params.get(PARAM_VAMANA_STREAMER_EF, &ef_); + params.get(PARAM_VAMANA_STREAMER_BRUTE_FORCE_THRESHOLD, + &bruteforce_threshold_); + params.get(PARAM_VAMANA_STREAMER_MAX_SCAN_RATIO, &max_scan_ratio_); + params.get(PARAM_VAMANA_STREAMER_MAX_SCAN_LIMIT, &max_scan_limit_); + params.get(PARAM_VAMANA_STREAMER_MIN_SCAN_LIMIT, &min_scan_limit_); + params.get(PARAM_VAMANA_STREAMER_CHUNK_SIZE, &chunk_size_); + params.get(PARAM_VAMANA_STREAMER_GET_VECTOR_ENABLE, &get_vector_enabled_); + params.get(PARAM_VAMANA_STREAMER_FORCE_PADDING_RESULT_ENABLE, + &force_padding_topk_enabled_); + params.get(PARAM_VAMANA_STREAMER_USE_ID_MAP, &use_id_map_); + params.get(PARAM_VAMANA_STREAMER_DOCS_HARD_LIMIT, &docs_hard_limit_); + params.get(PARAM_VAMANA_STREAMER_SATURATE_GRAPH, &saturate_graph_); + params.get(PARAM_VAMANA_STREAMER_USE_CONTIGUOUS_MEMORY, + &use_contiguous_memory_); + + size_t docs_soft_limit = 0; + params.get(PARAM_VAMANA_STREAMER_DOCS_SOFT_LIMIT, &docs_soft_limit); + if (docs_soft_limit > 0 && docs_soft_limit > docs_hard_limit_) { + LOG_ERROR("[%s] must be >= [%s]", + PARAM_VAMANA_STREAMER_DOCS_HARD_LIMIT.c_str(), + PARAM_VAMANA_STREAMER_DOCS_SOFT_LIMIT.c_str()); + return IndexError_InvalidArgument; + } else if (docs_soft_limit == 0UL) { + docs_soft_limit_ = + docs_hard_limit_ * VamanaEntity::kDefaultDocsSoftLimitRatio; + } else { + docs_soft_limit_ = docs_soft_limit; + } + + // Validate parameters + if (max_degree_ == 0U) max_degree_ = VamanaEntity::kDefaultMaxDegree; + if (search_list_size_ == 0U) + search_list_size_ = VamanaEntity::kDefaultSearchListSize; + if (max_occlusion_size_ == 0U) + max_occlusion_size_ = VamanaEntity::kDefaultMaxOcclusionSize; + if (alpha_ <= 0.0f) alpha_ = VamanaEntity::kDefaultAlpha; + if (ef_ == 0U) ef_ = VamanaEntity::kDefaultEf; + if (chunk_size_ == 0UL) chunk_size_ = VamanaEntity::kDefaultChunkSize; + if (chunk_size_ > VamanaEntity::kMaxChunkSize) { + LOG_ERROR("[%s] must be < %zu", PARAM_VAMANA_STREAMER_CHUNK_SIZE.c_str(), + VamanaEntity::kMaxChunkSize); + return IndexError_InvalidArgument; + } + if (max_scan_ratio_ <= 0.0f || max_scan_ratio_ > 1.0f) { + LOG_ERROR("[%s] must be in range (0.0f,1.0f]", + PARAM_VAMANA_STREAMER_MAX_SCAN_RATIO.c_str()); + return IndexError_InvalidArgument; + } + if (max_scan_limit_ < min_scan_limit_) { + LOG_ERROR("[%s] must be >= [%s]", + PARAM_VAMANA_STREAMER_MAX_SCAN_LIMIT.c_str(), + PARAM_VAMANA_STREAMER_MIN_SCAN_LIMIT.c_str()); + return IndexError_InvalidArgument; + } + + LOG_DEBUG( + "Vamana init params: maxIndexSize=%zu docsHardLimit=%zu " + "docsSoftLimit=%zu maxDegree=%u searchListSize=%u alpha=%.2f " + "maxOcclusionSize=%u ef=%u maxScanRatio=%.3f minScanLimit=%zu " + "maxScanLimit=%zu bruteForceThreshold=%zu chunkSize=%zu " + "getVectorEnabled=%u forcePadding=%u", + max_index_size_, docs_hard_limit_, docs_soft_limit_, max_degree_, + search_list_size_, alpha_, max_occlusion_size_, ef_, max_scan_ratio_, + min_scan_limit_, max_scan_limit_, bruteforce_threshold_, chunk_size_, + get_vector_enabled_, force_padding_topk_enabled_); + + state_ = STATE_INITED; + return 0; +} + +int VamanaStreamer::cleanup(void) { + if (state_ == STATE_OPENED) { + this->close(); + } + + LOG_INFO("VamanaStreamer cleanup"); + + meta_.clear(); + metric_.reset(); + stats_.clear(); + if (entity_) entity_->cleanup(); + if (alg_) alg_->cleanup(); + + max_index_size_ = 0UL; + docs_hard_limit_ = VamanaEntity::kDefaultDocsHardLimit; + docs_soft_limit_ = 0UL; + max_degree_ = VamanaEntity::kDefaultMaxDegree; + search_list_size_ = VamanaEntity::kDefaultSearchListSize; + max_occlusion_size_ = VamanaEntity::kDefaultMaxOcclusionSize; + alpha_ = VamanaEntity::kDefaultAlpha; + ef_ = VamanaEntity::kDefaultEf; + bruteforce_threshold_ = VamanaEntity::kDefaultBruteForceThreshold; + max_scan_limit_ = VamanaEntity::kDefaultMaxScanLimit; + min_scan_limit_ = VamanaEntity::kDefaultMinScanLimit; + chunk_size_ = VamanaEntity::kDefaultChunkSize; + max_scan_ratio_ = VamanaEntity::kDefaultScanRatio; + state_ = STATE_INIT; + check_crc_enabled_ = false; + get_vector_enabled_ = false; + + return 0; +} + +int VamanaStreamer::setup_entity() { + entity_->set_use_key_info_map(use_id_map_); + entity_->set_vector_size(meta_.element_size()); + entity_->set_chunk_size(chunk_size_); + entity_->set_get_vector(get_vector_enabled_); + + // Set Vamana-specific parameters via public setters + entity_->set_max_degree(max_degree_); + entity_->set_search_list_size(search_list_size_); + entity_->set_max_occlusion_size(max_occlusion_size_); + entity_->set_alpha(alpha_); + entity_->set_saturate_graph(saturate_graph_); + + int ret = entity_->init(docs_hard_limit_); + if (ret != 0) { + LOG_ERROR("Vamana entity init failed: %s", IndexError::What(ret)); + } + return ret; +} + +int VamanaStreamer::open(IndexStorage::Pointer stg) { + LOG_INFO("VamanaStreamer open"); + + if (ailego_unlikely(state_ != STATE_INITED)) { + LOG_ERROR("Open storage failed, init streamer first!"); + return IndexError_NoReady; + } + + // Create entity based on storage type + switch (stg->memory_block_type()) { + case IndexStorage::MemoryBlock::MBT_BUFFERPOOL: { + entity_ = std::make_unique(stats_); + break; + } + default: { + if (use_contiguous_memory_) { + entity_ = std::make_unique(stats_); + } else { + entity_ = std::make_unique(stats_); + } + break; + } + } + + auto cleanup_on_error = [this]() { + if (entity_) { + entity_->close(); + entity_.reset(); + } + alg_.reset(); + metric_.reset(); + }; + + int ret = setup_entity(); + if (ret != 0) { + cleanup_on_error(); + return ret; + } + + ret = entity_->open(std::move(stg), max_index_size_, check_crc_enabled_); + if (ret != 0) { + cleanup_on_error(); + return ret; + } + + // Handle IndexMeta + IndexMeta index_meta; + ret = entity_->get_index_meta(&index_meta); + if (ret == IndexError_NoExist) { + ret = entity_->set_index_meta(meta_); + if (ret != 0) { + LOG_ERROR("Failed to set index meta: %s", IndexError::What(ret)); + cleanup_on_error(); + return ret; + } + } else if (ret != 0) { + LOG_ERROR("Failed to get index meta: %s", IndexError::What(ret)); + cleanup_on_error(); + return ret; + } else { + if (index_meta.dimension() != meta_.dimension() || + index_meta.element_size() != meta_.element_size() || + index_meta.metric_name() != meta_.metric_name() || + index_meta.data_type() != meta_.data_type()) { + LOG_ERROR("IndexMeta mismatch from the previous in index"); + cleanup_on_error(); + return IndexError_Mismatch; + } + auto metric_params = index_meta.metric_params(); + metric_params.merge(meta_.metric_params()); + meta_.set_metric(index_meta.metric_name(), 0, metric_params); + } + + // Create metric + metric_ = IndexFactory::CreateMetric(meta_.metric_name()); + if (!metric_) { + LOG_ERROR("Failed to create metric %s", meta_.metric_name().c_str()); + cleanup_on_error(); + return IndexError_NoExist; + } + ret = metric_->init(meta_, meta_.metric_params()); + if (ret != 0) { + LOG_ERROR("Failed to init metric, ret=%d", ret); + cleanup_on_error(); + return ret; + } + if (!metric_->distance() || !metric_->batch_distance()) { + LOG_ERROR("Invalid metric distance functions"); + cleanup_on_error(); + return IndexError_InvalidArgument; + } + + add_distance_ = metric_->distance(); + add_batch_distance_ = metric_->batch_distance(); + search_distance_ = add_distance_; + search_batch_distance_ = add_batch_distance_; + + if (metric_->query_metric() && metric_->query_metric()->distance() && + metric_->query_metric()->batch_distance()) { + search_distance_ = metric_->query_metric()->distance(); + search_batch_distance_ = metric_->query_metric()->batch_distance(); + } + + // Create algorithm based on entity storage mode + switch (entity_->storage_mode()) { + case VamanaStorageMode::kBufferPool: + alg_ = VamanaAlgorithmBase::UPointer( + new VamanaAlgorithm( + static_cast(*entity_))); + break; + case VamanaStorageMode::kContiguous: { + auto &contiguous_entity = + static_cast(*entity_); + int build_ret = contiguous_entity.build_contiguous_memory(); + if (build_ret != 0) { + LOG_ERROR("Failed to build contiguous memory, ret=%d", build_ret); + cleanup_on_error(); + return build_ret; + } + alg_ = VamanaAlgorithmBase::UPointer( + new VamanaAlgorithm( + contiguous_entity)); + break; + } + default: + alg_ = VamanaAlgorithmBase::UPointer( + new VamanaAlgorithm( + static_cast(*entity_))); + break; + } + + ret = alg_->init(); + if (ret != 0) { + cleanup_on_error(); + return ret; + } + + state_ = STATE_OPENED; + magic_ = IndexContext::GenerateMagic(); + + return 0; +} + +int VamanaStreamer::close(void) { + LOG_INFO("VamanaStreamer close"); + + stats_.clear(); + meta_.set_metric(metric_->name(), 0, metric_->params()); + entity_->set_index_meta(meta_); + int ret = entity_->close(); + if (ret != 0) return ret; + state_ = STATE_INITED; + return 0; +} + +int VamanaStreamer::flush(uint64_t checkpoint) { + LOG_INFO("VamanaStreamer flush checkpoint=%zu", (size_t)checkpoint); + + meta_.set_metric(metric_->name(), 0, metric_->params()); + entity_->set_index_meta(meta_); + return entity_->flush(checkpoint); +} + +int VamanaStreamer::dump(const IndexDumper::Pointer &dumper) { + LOG_INFO("VamanaStreamer dump"); + + shared_mutex_.lock(); + AILEGO_DEFER([&]() { shared_mutex_.unlock(); }); + + meta_.set_searcher("VamanaSearcher", VamanaEntity::kRevision, + ailego::Params()); + + int ret = IndexHelper::SerializeToDumper(meta_, dumper.get()); + if (ret != 0) { + LOG_ERROR("Failed to serialize meta into dumper."); + return ret; + } + return entity_->dump(dumper); +} + +IndexStreamer::Context::Pointer VamanaStreamer::create_context(void) const { + if (ailego_unlikely(state_ != STATE_OPENED)) { + LOG_ERROR("Create context failed, open storage first!"); + return Context::Pointer(); + } + + VamanaEntity::Pointer entity = entity_->clone(); + if (ailego_unlikely(!entity)) { + LOG_ERROR("CreateContext clone failed"); + return Context::Pointer(); + } + + VamanaContext *ctx = + new (std::nothrow) VamanaContext(meta_.dimension(), metric_, entity); + if (ailego_unlikely(ctx == nullptr)) { + LOG_ERROR("Failed to new VamanaContext"); + return Context::Pointer(); + } + + ctx->set_ef(ef_); + ctx->set_max_scan_limit(max_scan_limit_); + ctx->set_min_scan_limit(min_scan_limit_); + ctx->set_max_scan_ratio(max_scan_ratio_); + ctx->set_magic(magic_); + ctx->set_force_padding_topk(force_padding_topk_enabled_); + ctx->set_bruteforce_threshold(bruteforce_threshold_); + + if (ailego_unlikely(ctx->init(VamanaContext::kStreamerContext) != 0)) { + LOG_ERROR("Init VamanaContext failed"); + delete ctx; + return Context::Pointer(); + } + + ctx->check_need_adjuct_ctx(entity_->doc_cnt()); + + return Context::Pointer(ctx); +} + +IndexProvider::Pointer VamanaStreamer::create_provider(void) const { + LOG_DEBUG("VamanaStreamer create provider"); + + auto entity = entity_->clone(); + if (ailego_unlikely(!entity)) { + LOG_ERROR("Clone VamanaEntity failed"); + return nullptr; + } + return IndexProvider::Pointer( + new VamanaIndexProvider(meta_, entity, "VamanaStreamer")); +} + +int VamanaStreamer::update_context(VamanaContext *ctx) const { + const VamanaEntity::Pointer entity = entity_->clone(); + if (!entity) { + LOG_ERROR("Failed to clone search context entity"); + return IndexError_Runtime; + } + ctx->set_max_scan_limit(max_scan_limit_); + ctx->set_min_scan_limit(min_scan_limit_); + ctx->set_max_scan_ratio(max_scan_ratio_); + ctx->set_bruteforce_threshold(bruteforce_threshold_); + return ctx->update_context(VamanaContext::kStreamerContext, meta_, metric_, + entity, magic_); +} + +int VamanaStreamer::add_impl(uint64_t pkey, const void *query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) { + int ret = check_params(query, qmeta); + if (ailego_unlikely(ret != 0)) return ret; + + VamanaContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to VamanaContext failed"); + return IndexError_Cast; + } + if (ctx->magic() != magic_) { + ret = update_context(ctx); + if (ret != 0) return ret; + } + + if (ailego_unlikely(entity_->doc_cnt() >= docs_soft_limit_)) { + if (entity_->doc_cnt() >= docs_hard_limit_) { + LOG_ERROR("Current docs %u exceed hard limit", entity_->doc_cnt()); + const std::lock_guard lk(mutex_); + (*stats_.mutable_discarded_count())++; + return IndexError_IndexFull; + } else { + LOG_WARN("Current docs %u exceed soft limit", entity_->doc_cnt()); + } + } + + if (ailego_unlikely(!shared_mutex_.try_lock_shared())) { + LOG_ERROR("Cannot add vector while dumping index"); + (*stats_.mutable_discarded_count())++; + return IndexError_Unsupported; + } + AILEGO_DEFER([&]() { shared_mutex_.unlock_shared(); }); + + ctx->clear(); + ctx->update_dist_caculator_distance(add_distance_, add_batch_distance_); + ctx->reset_query(query); + ctx->check_need_adjuct_ctx(entity_->doc_cnt()); + + if (metric_->support_train()) { + const std::lock_guard lk(mutex_); + ret = metric_->train(query, meta_.dimension()); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Vamana streamer metric train failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + } + + node_id_t id; + ret = entity_->add_vector(pkey, query, &id); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Vamana streamer add vector failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + + ret = alg_->add_node(id, ctx); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Vamana streamer add node failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + + if (ailego_unlikely(ctx->error())) { + (*stats_.mutable_discarded_count())++; + return IndexError_Runtime; + } + (*stats_.mutable_added_count())++; + + return 0; +} + +int VamanaStreamer::add_with_id_impl(uint32_t id, const void *query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) { + int ret = check_params(query, qmeta); + if (ailego_unlikely(ret != 0)) return ret; + + VamanaContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to VamanaContext failed"); + return IndexError_Cast; + } + if (ctx->magic() != magic_) { + ret = update_context(ctx); + if (ret != 0) return ret; + } + + if (ailego_unlikely(entity_->doc_cnt() >= docs_soft_limit_)) { + if (entity_->doc_cnt() >= docs_hard_limit_) { + LOG_ERROR("Current docs %u exceed hard limit", entity_->doc_cnt()); + const std::lock_guard lk(mutex_); + (*stats_.mutable_discarded_count())++; + return IndexError_IndexFull; + } else { + LOG_WARN("Current docs %u exceed soft limit", entity_->doc_cnt()); + } + } + + if (ailego_unlikely(!shared_mutex_.try_lock_shared())) { + LOG_ERROR("Cannot add vector while dumping index"); + (*stats_.mutable_discarded_count())++; + return IndexError_Unsupported; + } + AILEGO_DEFER([&]() { shared_mutex_.unlock_shared(); }); + + ctx->clear(); + ctx->update_dist_caculator_distance(add_distance_, add_batch_distance_); + ctx->reset_query(query); + ctx->check_need_adjuct_ctx(entity_->doc_cnt()); + + if (metric_->support_train()) { + const std::lock_guard lk(mutex_); + ret = metric_->train(query, meta_.dimension()); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Vamana streamer metric train failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + } + + ret = entity_->add_vector_with_id(id, query); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Vamana streamer add vector failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + + ret = alg_->add_node(id, ctx); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Vamana streamer add node failed"); + (*stats_.mutable_discarded_count())++; + return ret; + } + + if (ailego_unlikely(ctx->error())) { + (*stats_.mutable_discarded_count())++; + return IndexError_Runtime; + } + (*stats_.mutable_added_count())++; + + return 0; +} + +int VamanaStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const { + return search_impl(query, qmeta, 1, context); +} + +int VamanaStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, + Context::Pointer &context) const { + int ret = check_params(query, qmeta); + if (ailego_unlikely(ret != 0)) return ret; + + VamanaContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to VamanaContext failed"); + return IndexError_Cast; + } + + if (entity_->doc_cnt() <= ctx->get_bruteforce_threshold()) { + return search_bf_impl(query, qmeta, count, context); + } + + if (ctx->magic() != magic_) { + ret = update_context(ctx); + if (ret != 0) return ret; + } + + ctx->clear(); + ctx->update_dist_caculator_distance(search_distance_, search_batch_distance_); + ctx->resize_results(count); + ctx->check_need_adjuct_ctx(entity_->doc_cnt()); + + for (size_t q = 0; q < count; ++q) { + ctx->reset_query(query); + ret = alg_->search(ctx); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Vamana search failed"); + return ret; + } + ctx->topk_to_result(q); + query = static_cast(query) + qmeta.element_size(); + } + + if (ailego_unlikely(ctx->error())) return IndexError_Runtime; + return 0; +} + +void VamanaStreamer::print_debug_info() { + for (node_id_t id = 0; id < entity_->doc_cnt(); ++id) { + if (entity_->get_key(id) == kInvalidKey) continue; + Neighbors neighbours = entity_->get_neighbors(id); + std::cout << "node: " << id << "; "; + if (neighbours.size() == 0) { + std::cout << std::endl; + continue; + } + for (uint32_t i = 0; i < neighbours.size(); ++i) { + std::cout << neighbours[i]; + if (i == neighbours.size() - 1) { + std::cout << std::endl; + } else { + std::cout << ", "; + } + } + } +} + +int VamanaStreamer::search_bf_impl(const void *query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) const { + return search_bf_impl(query, qmeta, 1, context); +} + +int VamanaStreamer::search_bf_impl(const void *query, + const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const { + int ret = check_params(query, qmeta); + if (ailego_unlikely(ret != 0)) return ret; + + VamanaContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to VamanaContext failed"); + return IndexError_Cast; + } + if (ctx->magic() != magic_) { + ret = update_context(ctx); + if (ret != 0) return ret; + } + + ctx->clear(); + ctx->update_dist_caculator_distance(search_distance_, search_batch_distance_); + ctx->resize_results(count); + + const auto &filter = static_cast(ctx)->filter(); + auto &topk = ctx->topk_heap(); + + for (size_t q = 0; q < count; ++q) { + ctx->reset_query(query); + topk.clear(); + for (node_id_t id = 0; id < entity_->doc_cnt(); ++id) { + if (entity_->get_key(id) == kInvalidKey) continue; + if (!filter.is_valid() || !filter(entity_->get_key(id))) { + dist_t dist = ctx->dist_calculator().batch_dist(id); + topk.emplace(id, dist); + } + } + ctx->topk_to_result(q); + query = static_cast(query) + qmeta.element_size(); + } + + if (ailego_unlikely(ctx->error())) return IndexError_Runtime; + return 0; +} + +int VamanaStreamer::search_bf_by_p_keys_impl( + const void *query, const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const { + int ret = check_params(query, qmeta); + if (ailego_unlikely(ret != 0)) return ret; + + VamanaContext *ctx = dynamic_cast(context.get()); + ailego_do_if_false(ctx) { + LOG_ERROR("Cast context to VamanaContext failed"); + return IndexError_Cast; + } + if (ctx->magic() != magic_) { + ret = update_context(ctx); + if (ret != 0) return ret; + } + + ctx->clear(); + ctx->update_dist_caculator_distance(search_distance_, search_batch_distance_); + ctx->resize_results(count); + + auto &topk = ctx->topk_heap(); + + for (size_t q = 0; q < count; ++q) { + ctx->reset_query(query); + topk.clear(); + for (const auto &keys : p_keys) { + for (auto key : keys) { + node_id_t id = entity_->get_id(key); + if (id == kInvalidNodeId) continue; + dist_t dist = ctx->dist_calculator().batch_dist(id); + topk.emplace(id, dist); + } + } + ctx->topk_to_result(q); + query = static_cast(query) + qmeta.element_size(); + } + + if (ailego_unlikely(ctx->error())) return IndexError_Runtime; + return 0; +} + +INDEX_FACTORY_REGISTER_STREAMER(VamanaStreamer); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/vamana/vamana_streamer.h b/src/core/algorithm/vamana/vamana_streamer.h new file mode 100644 index 0000000..5520b37 --- /dev/null +++ b/src/core/algorithm/vamana/vamana_streamer.h @@ -0,0 +1,192 @@ +// Copyright 2025-present the zvec 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 +#include +#include "vamana_algorithm.h" +#include "vamana_streamer_entity.h" + +namespace zvec { +namespace core { + +class VamanaStreamer : public IndexStreamer { + public: + using ContextPointer = IndexStreamer::Context::Pointer; + + VamanaStreamer(void); + virtual ~VamanaStreamer(void); + + VamanaStreamer(const VamanaStreamer &) = delete; + VamanaStreamer &operator=(const VamanaStreamer &) = delete; + + protected: + virtual int init(const IndexMeta &imeta, + const ailego::Params ¶ms) override; + + virtual int cleanup(void) override; + + virtual Context::Pointer create_context(void) const override; + + virtual IndexProvider::Pointer create_provider(void) const override; + + virtual int add_impl(uint64_t pkey, const void *query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) override; + + virtual int add_with_id_impl(uint32_t id, const void *query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) override; + + virtual int search_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const override; + + virtual int search_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, + Context::Pointer &context) const override; + + virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const override; + + virtual int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, + Context::Pointer &context) const override; + + virtual int search_bf_by_p_keys_impl( + const void *query, const std::vector> &p_keys, + const IndexQueryMeta &qmeta, ContextPointer &context) const override { + return search_bf_by_p_keys_impl(query, p_keys, qmeta, 1, context); + } + + virtual int search_bf_by_p_keys_impl( + const void *query, const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + ContextPointer &context) const override; + + virtual const void *get_vector(uint64_t key) const override { + return entity_->get_vector_by_key(key); + } + + virtual int get_vector(const uint64_t key, + IndexStorage::MemoryBlock &block) const override { + return entity_->get_vector_by_key(key, block); + } + + virtual const void *get_vector_by_id(uint32_t id) const override { + return entity_->get_vector(id); + } + + virtual int get_vector_by_id( + const uint32_t id, IndexStorage::MemoryBlock &block) const override { + return entity_->get_vector(id, block); + } + + virtual int open(IndexStorage::Pointer stg) override; + + virtual int close(void) override; + + virtual int flush(uint64_t checkpoint) override; + + virtual int dump(const IndexDumper::Pointer &dumper) override; + + virtual const Stats &stats(void) const override { + return stats_; + } + + virtual const IndexMeta &meta(void) const override { + return meta_; + } + + virtual void print_debug_info() override; + + private: + inline int check_params(const void *query, + const IndexQueryMeta &qmeta) const { + if (ailego_unlikely(!query)) { + LOG_ERROR("null query"); + return IndexError_InvalidArgument; + } + if (ailego_unlikely(qmeta.dimension() != meta_.dimension() || + qmeta.data_type() != meta_.data_type() || + qmeta.element_size() != meta_.element_size())) { + LOG_ERROR("Unsupported query meta"); + return IndexError_Mismatch; + } + return 0; + } + + int setup_entity(); + int update_context(VamanaContext *ctx) const; + + private: + enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_OPENED = 2 }; + + class Stats : public IndexStreamer::Stats { + public: + void clear(void) { + set_revision_id(0u); + set_loaded_count(0u); + set_added_count(0u); + set_discarded_count(0u); + set_index_size(0u); + set_dumped_size(0u); + set_check_point(0u); + set_create_time(0u); + set_update_time(0u); + clear_attributes(); + } + }; + + std::unique_ptr entity_; + VamanaAlgorithmBase::UPointer alg_; + IndexMeta meta_{}; + IndexMetric::Pointer metric_{}; + + IndexMetric::MatrixDistance add_distance_{}; + IndexMetric::MatrixDistance search_distance_{}; + IndexMetric::MatrixBatchDistance add_batch_distance_{}; + IndexMetric::MatrixBatchDistance search_batch_distance_{}; + + Stats stats_{}; + std::mutex mutex_{}; + + size_t max_index_size_{0UL}; + size_t chunk_size_{VamanaEntity::kDefaultChunkSize}; + size_t docs_hard_limit_{VamanaEntity::kDefaultDocsHardLimit}; + size_t docs_soft_limit_{0UL}; + uint32_t max_degree_{VamanaEntity::kDefaultMaxDegree}; + uint32_t search_list_size_{VamanaEntity::kDefaultSearchListSize}; + uint32_t max_occlusion_size_{VamanaEntity::kDefaultMaxOcclusionSize}; + float alpha_{VamanaEntity::kDefaultAlpha}; + uint32_t ef_{VamanaEntity::kDefaultEf}; + size_t bruteforce_threshold_{VamanaEntity::kDefaultBruteForceThreshold}; + size_t max_scan_limit_{VamanaEntity::kDefaultMaxScanLimit}; + size_t min_scan_limit_{VamanaEntity::kDefaultMinScanLimit}; + float bf_negative_prob_{VamanaEntity::kDefaultBFNegativeProbability}; + float max_scan_ratio_{VamanaEntity::kDefaultScanRatio}; + + uint32_t magic_{0U}; + State state_{STATE_INIT}; + bool check_crc_enabled_{false}; + bool get_vector_enabled_{false}; + bool force_padding_topk_enabled_{false}; + bool use_id_map_{true}; + bool saturate_graph_{VamanaEntity::kDefaultSaturateGraph}; + bool use_contiguous_memory_{false}; + + ailego::SharedMutex shared_mutex_{}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/vamana/vamana_streamer_entity.cc b/src/core/algorithm/vamana/vamana_streamer_entity.cc new file mode 100644 index 0000000..ffe8e75 --- /dev/null +++ b/src/core/algorithm/vamana/vamana_streamer_entity.cc @@ -0,0 +1,773 @@ +// Copyright 2025-present the zvec 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 "vamana_streamer_entity.h" +#if defined(__linux__) || defined(__APPLE__) +#include +#endif +#include +#include +#include + +namespace zvec { +namespace core { + +VamanaStreamerEntity::VamanaStreamerEntity(IndexStreamer::Stats &stats) + : stats_(stats) { + keys_map_lock_ = std::make_shared(); + keys_map_ = std::make_shared>(); + keys_map_->set_empty_key(kInvalidKey); + broker_ = std::make_shared(stats); +} + +VamanaStreamerEntity::~VamanaStreamerEntity() {} + +int VamanaStreamerEntity::cleanup() { + node_chunks_.clear(); + if (keys_map_) { + keys_map_->clear(); + } + header_.clear(); + return 0; +} + +int VamanaStreamerEntity::init(size_t /*max_doc_cnt*/) { + // node_size = vector_size + key_size + neighbors_size + set_node_size(vector_size() + sizeof(key_t) + neighbors_size()); + neighbor_size_ = neighbors_size(); + return 0; +} + +key_t VamanaStreamerEntity::get_key(node_id_t id) const { + if (!use_key_info_map_) return id; + auto loc = get_key_chunk_loc(id); + if (ailego_unlikely(loc.first >= node_chunks_.size())) return kInvalidKey; + const void *ptr = nullptr; + size_t ret = node_chunks_[loc.first]->read(loc.second, &ptr, sizeof(key_t)); + if (ailego_unlikely(ret != sizeof(key_t))) { + LOG_ERROR("Read key failed, ret=%zu", ret); + return kInvalidKey; + } + return *reinterpret_cast(ptr); +} + +const void *VamanaStreamerEntity::get_vector(node_id_t id) const { + auto loc = get_vector_chunk_loc(id); + if (ailego_unlikely(loc.first >= node_chunks_.size())) return nullptr; + const void *ptr = nullptr; + size_t ret = node_chunks_[loc.first]->read(loc.second, &ptr, vector_size()); + if (ailego_unlikely(ret != vector_size())) { + LOG_ERROR("Read vector failed, ret=%zu", ret); + return nullptr; + } + return ptr; +} + +int VamanaStreamerEntity::get_vector(const node_id_t id, + IndexStorage::MemoryBlock &block) const { + auto loc = get_vector_chunk_loc(id); + if (ailego_unlikely(loc.first >= node_chunks_.size())) + return IndexError_NoExist; + size_t ret = node_chunks_[loc.first]->read(loc.second, block, vector_size()); + if (ailego_unlikely(ret != vector_size())) { + LOG_ERROR("Read vector failed, ret=%zu", ret); + return IndexError_ReadData; + } + return 0; +} + +int VamanaStreamerEntity::get_vector(const node_id_t *ids, uint32_t count, + const void **vecs) const { + for (uint32_t i = 0; i < count; ++i) { + vecs[i] = get_vector(ids[i]); + if (ailego_unlikely(vecs[i] == nullptr)) { + return IndexError_NoExist; + } + } + return 0; +} + +int VamanaStreamerEntity::get_vector( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + vec_blocks.resize(count); + for (uint32_t i = 0; i < count; ++i) { + int ret = get_vector(ids[i], vec_blocks[i]); + if (ailego_unlikely(ret != 0)) return ret; + } + return 0; +} + +const Neighbors VamanaStreamerEntity::get_neighbors(node_id_t id) const { + auto loc = get_neighbor_chunk_loc(id); + IndexStorage::MemoryBlock mem_block; + size_t ret = loc.first->read(loc.second, mem_block, neighbor_size_); + if (ailego_unlikely(ret != neighbor_size_)) { + LOG_ERROR("Read neighbor header failed, ret=%zu", ret); + return Neighbors(); + } + return Neighbors(mem_block); +} + +int VamanaStreamerEntity::add_vector(key_t key, const void *vec, + node_id_t *id) { + Chunk::Pointer node_chunk; + size_t chunk_offset = static_cast(-1); + + std::lock_guard lock(mutex_); + + node_id_t local_id = static_cast(doc_cnt()); + uint32_t chunk_index = node_chunks_.size() - 1U; + if (chunk_index == static_cast(-1) || + (node_chunks_[chunk_index]->data_size() >= + node_cnt_per_chunk_ * node_size())) { + if (ailego_unlikely(node_chunks_.capacity() == node_chunks_.size())) { + LOG_ERROR("add vector failed for no memory quota"); + return IndexError_IndexFull; + } + chunk_index++; + if (auto dret = ensure_dist_chunk_for(chunk_index); dret != 0) { + return dret; + } + auto p = broker_->alloc_chunk(ChunkBroker::CHUNK_TYPE_NODE, chunk_index, + chunk_size_); + if (ailego_unlikely(p.first != 0)) { + LOG_ERROR("Alloc data chunk failed"); + return p.first; + } + node_chunk = p.second; + chunk_offset = 0UL; + node_chunks_.emplace_back(node_chunk); + } else { + node_chunk = node_chunks_[chunk_index]; + chunk_offset = node_chunk->data_size(); + } + + // Write vector + size_t size = node_chunk->write(chunk_offset, vec, vector_size()); + if (ailego_unlikely(size != vector_size())) { + LOG_ERROR("Chunk write vec failed, ret=%zu", size); + return IndexError_WriteData; + } + // Write key + size = node_chunk->write(chunk_offset + vector_size(), &key, sizeof(key_t)); + if (ailego_unlikely(size != sizeof(key_t))) { + LOG_ERROR("Chunk write key failed, ret=%zu", size); + return IndexError_WriteData; + } + // Neighbors are initialized to zero by default (chunk is zero-filled) + + chunk_offset += node_size(); + if (ailego_unlikely(node_chunk->resize(chunk_offset) != chunk_offset)) { + LOG_ERROR("Chunk resize to %zu failed", chunk_offset); + return IndexError_Runtime; + } + + if (use_key_info_map_) { + keys_map_lock_->lock(); + (*keys_map_)[key] = local_id; + keys_map_lock_->unlock(); + } + + *mutable_doc_cnt() += 1; + broker_->mark_dirty(); + *id = local_id; + + return 0; +} + +int VamanaStreamerEntity::add_vector_with_id(node_id_t id, const void *vec) { + Chunk::Pointer node_chunk; + size_t chunk_offset = static_cast(-1); + key_t key = id; + + std::lock_guard lock(mutex_); + + auto func_get_node_chunk_and_offset = [&](node_id_t node_id) -> int { + uint32_t chunk_idx = node_id >> node_index_mask_bits_; + ailego_assert_with(chunk_idx <= node_chunks_.size(), "invalid chunk idx"); + if (chunk_idx == node_chunks_.size()) { + if (ailego_unlikely(node_chunks_.capacity() == node_chunks_.size())) { + LOG_ERROR("add vector failed for no memory quota"); + return IndexError_IndexFull; + } + if (auto dret = ensure_dist_chunk_for(chunk_idx); dret != 0) { + return dret; + } + auto p = broker_->alloc_chunk(ChunkBroker::CHUNK_TYPE_NODE, chunk_idx, + chunk_size_); + if (ailego_unlikely(p.first != 0)) { + LOG_ERROR("Alloc data chunk failed"); + return p.first; + } + node_chunk = p.second; + node_chunks_.emplace_back(node_chunk); + } + node_chunk = node_chunks_[chunk_idx]; + chunk_offset = (node_id & node_index_mask_) * node_size(); + return 0; + }; + + // Fill gaps with invalid keys + for (size_t start_id = doc_cnt(); start_id < id; ++start_id) { + if (auto ret = func_get_node_chunk_and_offset(start_id); ret != 0) { + return ret; + } + size_t size = node_chunk->write(chunk_offset + vector_size(), &kInvalidKey, + sizeof(key_t)); + if (ailego_unlikely(size != sizeof(key_t))) { + LOG_ERROR("Chunk write key failed, ret=%zu", size); + return IndexError_WriteData; + } + chunk_offset += node_size(); + if (ailego_unlikely(node_chunk->resize(chunk_offset) != chunk_offset)) { + LOG_ERROR("Chunk resize to %zu failed", chunk_offset); + return IndexError_Runtime; + } + } + + if (auto ret = func_get_node_chunk_and_offset(id); ret != 0) { + return ret; + } + + // Write vector + size_t size = node_chunk->write(chunk_offset, vec, vector_size()); + if (ailego_unlikely(size != vector_size())) { + LOG_ERROR("Chunk write vec failed, ret=%zu", size); + return IndexError_WriteData; + } + // Write key + size = node_chunk->write(chunk_offset + vector_size(), &key, sizeof(key_t)); + if (ailego_unlikely(size != sizeof(key_t))) { + LOG_ERROR("Chunk write key failed, ret=%zu", size); + return IndexError_WriteData; + } + + if (*mutable_doc_cnt() <= id) { + *mutable_doc_cnt() = id + 1; + chunk_offset += node_size(); + if (ailego_unlikely(node_chunk->resize(chunk_offset) != chunk_offset)) { + LOG_ERROR("Chunk resize to %zu failed", chunk_offset); + return IndexError_Runtime; + } + } + + if (use_key_info_map_) { + keys_map_lock_->lock(); + (*keys_map_)[key] = id; + keys_map_lock_->unlock(); + } + + broker_->mark_dirty(); + return 0; +} + +int VamanaStreamerEntity::update_neighbors( + node_id_t id, const std::vector> &neighbors) { + auto loc = get_neighbor_chunk_loc(id); + uint32_t count = std::min(static_cast(neighbors.size()), + static_cast(max_degree())); + + // Build neighbor data in a local buffer + size_t nbr_size = neighbors_size(); + std::vector buffer(nbr_size, 0); + auto *hd = reinterpret_cast(buffer.data()); + hd->neighbor_cnt = count; + for (uint32_t i = 0; i < count; ++i) { + hd->neighbors[i] = neighbors[i].first; + } + + size_t ret = loc.first->write(loc.second, buffer.data(), nbr_size); + if (ailego_unlikely(ret != nbr_size)) { + LOG_ERROR("Write neighbors failed, ret=%zu", ret); + return IndexError_WriteData; + } + return 0; +} + +void VamanaStreamerEntity::add_neighbor(node_id_t id, uint32_t size, + node_id_t neighbor_id) { + auto loc = get_neighbor_chunk_loc(id); + if (size >= max_degree()) return; + + // Read current neighbors + IndexStorage::MemoryBlock mem_block; + size_t ret = loc.first->read(loc.second, mem_block, neighbor_size_); + if (ailego_unlikely(ret != neighbor_size_)) { + LOG_ERROR("Read neighbor header failed, ret=%zu", ret); + return; + } + + // Copy to mutable buffer, update, and write back + std::vector buffer(neighbor_size_); + memcpy(buffer.data(), mem_block.data(), neighbor_size_); + auto *hd = reinterpret_cast(buffer.data()); + hd->neighbors[size] = neighbor_id; + hd->neighbor_cnt = size + 1; + + ret = loc.first->write(loc.second, buffer.data(), neighbor_size_); + if (ailego_unlikely(ret != neighbor_size_)) { + LOG_ERROR("Write neighbor failed, ret=%zu", ret); + } +} + +void VamanaStreamerEntity::update_entry_point(node_id_t ep) { + VamanaEntity::update_entry_point(ep); + flush_header(); +} + +int VamanaStreamerEntity::open(IndexStorage::Pointer stg, + uint64_t max_index_size, bool check_crc) { + std::lock_guard lock(mutex_); + bool huge_page = stg->isHugePage(); + + int ret = broker_->open(std::move(stg), chunk_size_, check_crc); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("Open index failed: %s", IndexError::What(ret)); + return ret; + } + + ret = init_chunk_params(max_index_size, huge_page); + if (ailego_unlikely(ret != 0)) { + LOG_ERROR("init_chunk_params failed: %s", IndexError::What(ret)); + return ret; + } + broker_->set_max_chunks_size(max_index_size_); + + // Init header + auto header_chunk = broker_->get_chunk(ChunkBroker::CHUNK_TYPE_HEADER, + ChunkBroker::kDefaultChunkSeqId); + if (!header_chunk) { + // Open empty index, create header + auto p = + broker_->alloc_chunk(ChunkBroker::CHUNK_TYPE_HEADER, + ChunkBroker::kDefaultChunkSeqId, header_size()); + if (ailego_unlikely(p.first != 0)) { + LOG_ERROR("Alloc header chunk failed"); + return p.first; + } + size_t size = p.second->write(0UL, &header(), header_size()); + if (ailego_unlikely(size != header_size())) { + LOG_ERROR("Write header chunk failed"); + return IndexError_WriteData; + } + return 0; + } + + // Open existing index + ret = init_chunks(header_chunk); + if (ailego_unlikely(ret != 0)) return ret; + + // Verify total docs + node_id_t total_vecs = 0; + if (!node_chunks_.empty()) { + size_t last_idx = node_chunks_.size() - 1; + if (node_chunks_[last_idx]->data_size() % node_size()) { + LOG_WARN("The index may be broken"); + return IndexError_InvalidFormat; + } + total_vecs = last_idx * node_cnt_per_chunk_ + + node_chunks_[last_idx]->data_size() / node_size(); + } + + LOG_INFO("Open Vamana index, maxDegree=%zu docCnt=%u totalVecs=%u", + max_degree(), doc_cnt(), total_vecs); + + if (doc_cnt() != total_vecs) { + LOG_WARN("Index closed abnormally, using totalVecs as curDocCnt"); + *mutable_doc_cnt() = total_vecs; + } + + // Rebuild key map + if (use_key_info_map_) { + for (node_id_t i = 0; i < doc_cnt(); ++i) { + key_t k = get_key(i); + if (k != kInvalidKey) { + (*keys_map_)[k] = i; + } + } + } + + stats_.set_loaded_count(doc_cnt()); + return 0; +} + +int VamanaStreamerEntity::init_chunks(const Chunk::Pointer &header_chunk) { + // Read header from chunk + const void *hd_ptr = nullptr; + size_t ret = header_chunk->read(0UL, &hd_ptr, header_size()); + if (ailego_unlikely(ret != header_size())) { + LOG_ERROR("Read header chunk failed"); + return IndexError_ReadData; + } + auto *hd = reinterpret_cast(hd_ptr); + + // Validate + if (vector_size() != hd->vector_size()) { + LOG_ERROR("vector size %zu mismatch index previous %zu", vector_size(), + hd->vector_size()); + return IndexError_Mismatch; + } + if (max_degree() != hd->max_degree()) { + LOG_ERROR("max_degree %zu mismatch index previous %zu", max_degree(), + hd->max_degree()); + return IndexError_Mismatch; + } + + *mutable_header() = *hd; + + // Load node chunks + size_t chunk_cnt = broker_->get_chunk_cnt(ChunkBroker::CHUNK_TYPE_NODE); + for (size_t i = 0; i < chunk_cnt; ++i) { + auto chunk = broker_->get_chunk(ChunkBroker::CHUNK_TYPE_NODE, i); + if (ailego_unlikely(!chunk)) { + LOG_ERROR("Get node chunk %zu failed", i); + return IndexError_ReadData; + } + node_chunks_.emplace_back(std::move(chunk)); + } + return 0; +} + +int VamanaStreamerEntity::close() { + LOG_DEBUG("close Vamana index"); + std::lock_guard lock(mutex_); + flush_header(); + mutable_header()->reset(); + keys_map_->clear(); + header_.clear(); + node_chunks_.clear(); + dist_chunks_.clear(); + dist_loaded_ = false; + return broker_->close(); +} + +int VamanaStreamerEntity::flush(uint64_t checkpoint) { + LOG_INFO("Flush Vamana index, curDocs=%u", doc_cnt()); + std::lock_guard lock(mutex_); + flush_header(); + return broker_->flush(checkpoint); +} + +int VamanaStreamerEntity::dump(const IndexDumper::Pointer &dumper) { + LOG_INFO("Dump Vamana index, curDocs=%u", doc_cnt()); + + std::vector keys(doc_cnt()); + auto ret = dump_segments(dumper, keys.data()); + if (ailego_unlikely(ret < 0)) { + return static_cast(ret); + } + *stats_.mutable_dumped_size() += ret; + return 0; +} + +const VamanaEntity::Pointer VamanaStreamerEntity::clone() const { + std::vector cloned_chunks; + cloned_chunks.reserve(node_chunks_.size()); + for (size_t i = 0; i < node_chunks_.size(); ++i) { + cloned_chunks.emplace_back(node_chunks_[i]->clone()); + if (ailego_unlikely(!cloned_chunks[i])) { + LOG_ERROR("VamanaStreamerEntity get chunk failed in clone"); + return VamanaEntity::Pointer(); + } + } + + auto *entity = new (std::nothrow) VamanaStreamerEntity( + stats_, header(), chunk_size_, node_index_mask_bits_, get_vector_enabled_, + use_key_info_map_, keys_map_lock_, keys_map_, std::move(cloned_chunks), + broker_); + if (ailego_unlikely(!entity)) { + LOG_ERROR("VamanaStreamerEntity new failed"); + } + return VamanaEntity::Pointer(entity); +} + +const VamanaEntity::Pointer VamanaMmapStreamerEntity::clone() const { + std::vector cloned_chunks; + cloned_chunks.reserve(node_chunks_.size()); + for (size_t i = 0; i < node_chunks_.size(); ++i) { + cloned_chunks.emplace_back(node_chunks_[i]->clone()); + if (ailego_unlikely(!cloned_chunks[i])) { + LOG_ERROR("VamanaMmapStreamerEntity get chunk failed in clone"); + return VamanaEntity::Pointer(); + } + } + + auto *entity = new (std::nothrow) VamanaMmapStreamerEntity( + stats_, header(), chunk_size_, node_index_mask_bits_, get_vector_enabled_, + use_key_info_map_, keys_map_lock_, keys_map_, std::move(cloned_chunks), + broker_); + if (ailego_unlikely(!entity)) { + LOG_ERROR("VamanaMmapStreamerEntity new failed"); + } + return VamanaEntity::Pointer(entity); +} + +const VamanaEntity::Pointer VamanaContiguousStreamerEntity::clone() const { + std::vector cloned_chunks; + cloned_chunks.reserve(node_chunks_.size()); + for (size_t i = 0; i < node_chunks_.size(); ++i) { + cloned_chunks.emplace_back(node_chunks_[i]->clone()); + if (ailego_unlikely(!cloned_chunks[i])) { + LOG_ERROR("VamanaContiguousStreamerEntity get chunk failed in clone"); + return VamanaEntity::Pointer(); + } + } + + auto *entity = new (std::nothrow) VamanaContiguousStreamerEntity( + stats_, header(), chunk_size_, node_index_mask_bits_, get_vector_enabled_, + use_key_info_map_, keys_map_lock_, keys_map_, std::move(cloned_chunks), + broker_); + if (ailego_unlikely(!entity)) { + LOG_ERROR("VamanaContiguousStreamerEntity new failed"); + return VamanaEntity::Pointer(); + } + + // Share contiguous memory with the clone (zero-copy) + entity->node_memory_ = node_memory_; + entity->node_base_ = node_base_; + + return VamanaEntity::Pointer(entity); +} + +// ============================================================================ +// VamanaContiguousStreamerEntity implementation +// ============================================================================ + +char *VamanaContiguousStreamerEntity::allocate_contiguous(size_t size) { + if (size == 0) return nullptr; +#if defined(__linux__) + void *ptr = ::mmap(nullptr, size, PROT_READ | PROT_WRITE, + MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + if (ptr == MAP_FAILED) { + LOG_ERROR("mmap failed for contiguous memory, size=%zu", size); + return nullptr; + } + ::madvise(ptr, size, MADV_HUGEPAGE); + return static_cast(ptr); +#elif defined(__APPLE__) + void *ptr = ::mmap(nullptr, size, PROT_READ | PROT_WRITE, + MAP_PRIVATE | MAP_ANON, -1, 0); + if (ptr == MAP_FAILED) { + LOG_ERROR("mmap failed for contiguous memory, size=%zu", size); + return nullptr; + } + return static_cast(ptr); +#elif defined(_WIN32) + void *ptr = ::_aligned_malloc(size, ailego::MemoryHelper::PageSize()); + if (!ptr) { + LOG_ERROR("_aligned_malloc failed for contiguous memory, size=%zu", size); + return nullptr; + } + return static_cast(ptr); +#else + void *ptr = std::aligned_alloc(ailego::MemoryHelper::PageSize(), size); + if (!ptr) { + LOG_ERROR("aligned_alloc failed, size=%zu", size); + return nullptr; + } + return static_cast(ptr); +#endif +} + +int VamanaContiguousStreamerEntity::build_contiguous_memory() { + node_memory_.reset(); + node_base_ = nullptr; + + const uint32_t total_docs = doc_cnt(); + if (total_docs == 0) return 0; + + const size_t per_node = node_size(); + const size_t total_node_data = static_cast(total_docs) * per_node; + size_t node_memory_size = AlignHugePageSize(total_node_data); + char *raw_node = allocate_contiguous(node_memory_size); + if (!raw_node) return IndexError_Runtime; + node_memory_.reset(raw_node, ContiguousDeleter{node_memory_size}); + node_base_ = raw_node; + + // Copy node data from chunks into contiguous memory + const auto &chunks = node_chunks_; + const uint32_t nodes_per_chunk = 1U << node_index_mask_bits_; + for (size_t chunk_idx = 0; chunk_idx < chunks.size(); ++chunk_idx) { + const void *chunk_data = nullptr; + size_t data_size = chunks[chunk_idx]->data_size(); + chunks[chunk_idx]->read(0, &chunk_data, data_size); + + uint32_t base_id = chunk_idx * nodes_per_chunk; + uint32_t count_in_chunk = std::min(nodes_per_chunk, total_docs - base_id); + + const char *src = static_cast(chunk_data); + char *dst = node_base_ + static_cast(base_id) * per_node; + std::memcpy(dst, src, static_cast(count_in_chunk) * per_node); + } + + LOG_INFO( + "Built Vamana contiguous memory: node_size=%zu total_docs=%u " + "node_chunks=%zu", + node_memory_size, total_docs, chunks.size()); + + return 0; +} + +// ============================================================================ +// Neighbor distance storage implementation (CSR-like, lazy-loaded) +// ============================================================================ + +int VamanaStreamerEntity::ensure_dist_storage() { + if (dist_loaded_) return 0; + + std::lock_guard lock(mutex_); + if (dist_loaded_) return 0; // double-check after lock + + dist_entry_size_ = static_cast(max_degree() * sizeof(dist_t)); + + // Calculate how many dist chunks we need for existing nodes + uint32_t total_docs = doc_cnt(); + if (total_docs == 0) { + dist_loaded_ = true; + return 0; + } + + // Check if dist chunks already exist in storage (reopened index) + size_t existing_dist_chunks = + broker_->get_chunk_cnt(ChunkBroker::CHUNK_TYPE_NEIGHBOR_DIST); + if (existing_dist_chunks > 0) { + // Load existing dist chunks + for (size_t i = 0; i < existing_dist_chunks; ++i) { + auto chunk = broker_->get_chunk(ChunkBroker::CHUNK_TYPE_NEIGHBOR_DIST, i); + if (ailego_unlikely(!chunk)) { + LOG_ERROR("Failed to load dist chunk %zu", i); + return IndexError_ReadData; + } + dist_chunks_.emplace_back(std::move(chunk)); + } + LOG_INFO("Loaded %zu existing dist chunks", existing_dist_chunks); + } else { + // Allocate new dist chunks for all existing nodes + int ret = alloc_dist_chunks_for_existing_nodes(); + if (ret != 0) return ret; + } + + dist_loaded_ = true; + return 0; +} + +int VamanaStreamerEntity::ensure_dist_chunk_for(uint32_t chunk_index) { + // No-op when dist storage is not active. + if (!dist_loaded_ || dist_entry_size_ == 0) return 0; + + // Idempotent: nothing to do if this dist chunk slot already exists and is + // populated. (Slots created by the placeholder loop below will hold + // nullptr and must still be (re-)allocated.) + if (chunk_index < dist_chunks_.size() && dist_chunks_[chunk_index]) { + return 0; + } + + uint32_t dist_chunk_data_size = node_cnt_per_chunk_ * dist_entry_size_; + uint32_t dist_chunk_size = AlignPageSize(dist_chunk_data_size); + auto dp = broker_->alloc_chunk(ChunkBroker::CHUNK_TYPE_NEIGHBOR_DIST, + chunk_index, dist_chunk_size); + if (ailego_unlikely(dp.first != 0)) { + LOG_ERROR("Alloc dist chunk %u failed", chunk_index); + return dp.first; + } + dp.second->resize(dist_chunk_data_size); + while (dist_chunks_.size() <= chunk_index) { + dist_chunks_.emplace_back(nullptr); + } + dist_chunks_[chunk_index] = std::move(dp.second); + return 0; +} + +int VamanaStreamerEntity::alloc_dist_chunks_for_existing_nodes() { + uint32_t total_docs = doc_cnt(); + if (total_docs == 0) return 0; + + // Calculate dist chunk size: same number of nodes per chunk as node chunks + uint32_t dist_chunk_data_size = node_cnt_per_chunk_ * dist_entry_size_; + uint32_t dist_chunk_size = AlignPageSize(dist_chunk_data_size); + + uint32_t num_chunks_needed = + (total_docs + node_cnt_per_chunk_ - 1) >> node_index_mask_bits_; + + for (uint32_t i = 0; i < num_chunks_needed; ++i) { + auto p = broker_->alloc_chunk(ChunkBroker::CHUNK_TYPE_NEIGHBOR_DIST, i, + dist_chunk_size); + if (ailego_unlikely(p.first != 0)) { + LOG_ERROR("Alloc dist chunk %u failed", i); + return p.first; + } + // Resize to cover all nodes in this chunk + uint32_t nodes_in_chunk = + std::min(node_cnt_per_chunk_, total_docs - i * node_cnt_per_chunk_); + size_t data_size = static_cast(nodes_in_chunk) * dist_entry_size_; + p.second->resize(data_size); + dist_chunks_.emplace_back(std::move(p.second)); + } + broker_->mark_dirty(); + LOG_INFO("Allocated %u dist chunks for %u existing nodes", num_chunks_needed, + total_docs); + return 0; +} + +const dist_t *VamanaStreamerEntity::get_neighbor_dists(node_id_t id) const { + if (!dist_loaded_) return nullptr; + + auto loc = get_dist_chunk_loc(id); + if (ailego_unlikely(loc.first >= dist_chunks_.size())) { + sync_dist_chunks(loc.first); + } + if (ailego_unlikely(loc.first >= dist_chunks_.size())) return nullptr; + + const void *ptr = nullptr; + size_t ret = + dist_chunks_[loc.first]->read(loc.second, &ptr, dist_entry_size_); + if (ailego_unlikely(ret != dist_entry_size_)) return nullptr; + return static_cast(ptr); +} + +void VamanaStreamerEntity::update_neighbor_dists( + node_id_t id, const std::vector> &neighbors) { + if (!dist_loaded_) return; + + auto loc = get_dist_chunk_loc(id); + // Dist chunk must have been pre-allocated by add_vector or + // ensure_dist_storage + if (ailego_unlikely(loc.first >= dist_chunks_.size() || + dist_chunks_[loc.first] == nullptr)) { + LOG_ERROR("Dist chunk %u not allocated for node %u", loc.first, id); + return; + } + + // Write distances: fill max_degree slots, zero-pad unused slots + uint32_t max_deg = static_cast(max_degree()); + std::vector dists(max_deg, 0.0f); + for (size_t i = 0; i < neighbors.size() && i < max_deg; ++i) { + dists[i] = neighbors[i].second; + } + dist_chunks_[loc.first]->write(loc.second, dists.data(), dist_entry_size_); +} + +void VamanaStreamerEntity::set_neighbor_dist(node_id_t id, uint32_t idx, + dist_t dist) { + if (!dist_loaded_) return; + + auto loc = get_dist_chunk_loc(id); + if (ailego_unlikely(loc.first >= dist_chunks_.size())) return; + + uint32_t offset = loc.second + idx * sizeof(dist_t); + dist_chunks_[loc.first]->write(offset, &dist, sizeof(dist_t)); +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/vamana/vamana_streamer_entity.h b/src/core/algorithm/vamana/vamana_streamer_entity.h new file mode 100644 index 0000000..ae29187 --- /dev/null +++ b/src/core/algorithm/vamana/vamana_streamer_entity.h @@ -0,0 +1,635 @@ +// Copyright 2025-present the zvec 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 +#include +#include +#if defined(__linux__) || defined(__APPLE__) +#include +#endif +#include +#include +#include +#include +#include "algorithm/hnsw/hnsw_chunk.h" +#include "algorithm/hnsw/hnsw_entity.h" // MmapMemoryBlock, BufferPoolMemoryBlock, NeighborsT +#include "vamana_entity.h" +#include "vamana_params.h" + +namespace zvec { +namespace core { + + +// Storage mode for VamanaStreamerEntity +enum class VamanaStorageMode { kMmap = 0, kBufferPool = 1, kContiguous = 2 }; + +// VamanaStreamerEntity manages vector data, primary keys, and neighbors +// for a single-layer Vamana graph in streaming (incremental) mode. +// Unlike HNSW, Vamana has no upper-level neighbors — only a single +// neighbor list per node. Node layout in chunk: +// [vector_data (vector_size) | key (sizeof(key_t)) | NeighborsHeader + +// neighbors (neighbors_size)] +class VamanaStreamerEntity : public VamanaEntity { + public: + // Virtual interface implementation + virtual int cleanup() override; + virtual const VamanaEntity::Pointer clone() const override; + virtual key_t get_key(node_id_t id) const override; + virtual const void *get_vector(node_id_t id) const override; + virtual int get_vector(const node_id_t id, + IndexStorage::MemoryBlock &block) const override; + virtual int get_vector(const node_id_t *ids, uint32_t count, + const void **vecs) const override; + virtual int get_vector( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const override; + virtual const Neighbors get_neighbors(node_id_t id) const override; + + virtual int add_vector(key_t key, const void *vec, node_id_t *id) override; + virtual int add_vector_with_id(node_id_t id, const void *vec) override; + virtual int update_neighbors( + node_id_t id, + const std::vector> &neighbors) override; + virtual void add_neighbor(node_id_t id, uint32_t size, + node_id_t neighbor_id) override; + virtual int dump(const IndexDumper::Pointer &dumper) override; + virtual void update_entry_point(node_id_t ep) override; + + // --- Neighbor distance storage --- + int ensure_dist_storage() override; + bool dist_storage_loaded() const override { + return dist_loaded_; + } + const dist_t *get_neighbor_dists(node_id_t id) const override; + void update_neighbor_dists( + node_id_t id, + const std::vector> &neighbors) override; + void set_neighbor_dist(node_id_t id, uint32_t idx, dist_t dist) override; + + virtual VamanaStorageMode storage_mode() const { + return VamanaStorageMode::kMmap; + } + + void set_use_key_info_map(bool use_id_map) { + use_key_info_map_ = use_id_map; + } + + public: + VamanaStreamerEntity(IndexStreamer::Stats &stats); + ~VamanaStreamerEntity(); + + virtual const void *get_vector_by_key(key_t key) const override { + auto id = get_id(key); + return id == kInvalidNodeId ? nullptr : get_vector(id); + } + + virtual int get_vector_by_key( + const key_t key, IndexStorage::MemoryBlock &block) const override { + auto id = get_id(key); + if (id != kInvalidNodeId) { + return get_vector(id, block); + } + return IndexError_InvalidArgument; + } + + int init(size_t max_doc_cnt); + int flush(uint64_t checkpoint); + int open(IndexStorage::Pointer stg, uint64_t max_index_size, bool check_crc); + int close(); + + int set_index_meta(const IndexMeta &meta) const { + return IndexHelper::SerializeToStorage(meta, broker_->storage().get()); + } + + int get_index_meta(IndexMeta *meta) const { + return IndexHelper::DeserializeFromStorage(broker_->storage().get(), meta); + } + + inline void set_chunk_size(size_t val) { + chunk_size_ = val; + } + inline void set_get_vector(bool val) { + get_vector_enabled_ = val; + } + + inline node_id_t get_id(key_t key) const { + if (use_key_info_map_) { + keys_map_lock_->lock_shared(); + auto it = keys_map_->find(key); + keys_map_lock_->unlock_shared(); + return it == keys_map_->end() ? kInvalidNodeId : it->second; + } + return key; + } + + // --- Typed access methods for hot-path optimization --- + // These are templated on MemBlock type to avoid runtime branching. + + template + inline NeighborsT get_neighbors_typed(node_id_t id) const; + + template + inline int get_vector_typed(const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const; + + template + inline key_t get_key_typed(node_id_t id) const; + + protected: + inline void sync_node_chunks(size_t idx) const { + sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, idx, &node_chunks_); + } + + protected: + template + using HashMap = google::dense_hash_map>; + template + using HashMapPointer = std::shared_ptr>; + + //! Clone constructor, used by clone method in subclasses + VamanaStreamerEntity(IndexStreamer::Stats &stats, const VamanaHeader &hd, + size_t chunk_size, uint32_t node_index_mask_bits, + bool get_vector_enabled, bool use_key_info_map, + std::shared_ptr &keys_map_lock, + const HashMapPointer &keys_map, + std::vector &&node_chunks, + const ChunkBroker::Pointer &broker) + : stats_(stats), + chunk_size_(chunk_size), + node_index_mask_bits_(node_index_mask_bits), + node_cnt_per_chunk_(1UL << node_index_mask_bits_), + node_index_mask_(node_cnt_per_chunk_ - 1), + get_vector_enabled_(get_vector_enabled), + use_key_info_map_(use_key_info_map), + keys_map_lock_(keys_map_lock), + keys_map_(keys_map), + node_chunks_(std::move(node_chunks)), + broker_(broker) { + *mutable_header() = hd; + neighbor_size_ = neighbors_size(); + } + + //! Lazy chunk synchronization: fetches chunks from broker when needed. + //! Each clone entity has its own node_chunks_ vector, so concurrent + //! search threads do not race with the writer's emplace_back. + void sync_chunks(ChunkBroker::CHUNK_TYPE type, size_t idx, + std::vector *chunks) const { + if (ailego_likely(idx < chunks->size())) { + return; + } + for (size_t i = chunks->size(); i <= idx; ++i) { + auto chunk = broker_->get_chunk(type, i); + ailego_assert_with(!!chunk, "get chunk failed"); + chunks->emplace_back(std::move(chunk)); + } + } + + inline std::pair get_vector_chunk_loc( + node_id_t id) const { + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = (id & node_index_mask_) * node_size(); + sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_); + return std::make_pair(chunk_idx, offset); + } + + inline std::pair get_key_chunk_loc(node_id_t id) const { + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = (id & node_index_mask_) * node_size() + vector_size(); + sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_); + return std::make_pair(chunk_idx, offset); + } + + inline std::pair get_neighbor_chunk_loc(node_id_t id) const { + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = + (id & node_index_mask_) * node_size() + vector_size() + sizeof(key_t); + sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_); + ailego_assert_abort(chunk_idx < node_chunks_.size(), "invalid chunk idx"); + return std::make_pair(node_chunks_[chunk_idx].get(), offset); + } + + // Get chunk location for neighbor distance data. + // Uses the same chunk indexing as node chunks but with dist_entry_size_. + inline std::pair get_dist_chunk_loc(node_id_t id) const { + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = (id & node_index_mask_) * dist_entry_size_; + return std::make_pair(chunk_idx, offset); + } + + void sync_dist_chunks(size_t idx) const { + sync_chunks(ChunkBroker::CHUNK_TYPE_NEIGHBOR_DIST, idx, &dist_chunks_); + } + + int ensure_dist_chunk_for(uint32_t chunk_index); + + int alloc_dist_chunks_for_existing_nodes(); + + size_t estimate_doc_capacity() const { + return node_chunks_.capacity() * node_cnt_per_chunk_; + } + + int init_chunk_params(size_t max_index_size, bool huge_page) { + node_cnt_per_chunk_ = std::max(1, chunk_size_ / node_size()); + node_index_mask_bits_ = std::ceil(std::log2(node_cnt_per_chunk_)); + node_cnt_per_chunk_ = 1UL << node_index_mask_bits_; + if (huge_page) { + chunk_size_ = AlignHugePageSize(node_cnt_per_chunk_ * node_size()); + } else { + chunk_size_ = AlignPageSize(node_cnt_per_chunk_ * node_size()); + } + node_index_mask_ = node_cnt_per_chunk_ - 1; + + if (max_index_size == 0UL) { + max_index_size_ = chunk_size_ * kDefaultMaxChunkCnt; + } else { + max_index_size_ = max_index_size; + } + + size_t max_node_chunk_cnt = + std::ceil(static_cast(max_index_size_) / chunk_size_); + node_chunks_.reserve(max_node_chunk_cnt); + + LOG_DEBUG( + "VamanaSettings: nodeSize=%zu chunkSize=%u nodeCntPerChunk=%u " + "maxChunkCnt=%zu maxIndexSize=%zu", + node_size(), chunk_size_, node_cnt_per_chunk_, max_node_chunk_cnt, + max_index_size_); + + return 0; + } + + int init_chunks(const Chunk::Pointer &header_chunk); + + int flush_header(void) { + if (!broker_->dirty()) { + return 0; + } + auto header_chunk = broker_->get_chunk(ChunkBroker::CHUNK_TYPE_HEADER, + ChunkBroker::kDefaultChunkSeqId); + if (ailego_unlikely(!header_chunk)) { + LOG_ERROR("get header chunk failed"); + return IndexError_Runtime; + } + size_t size = header_chunk->write(0UL, &header(), header_size()); + if (ailego_unlikely(size != header_size())) { + LOG_ERROR("Write header chunk failed"); + return IndexError_WriteData; + } + return 0; + } + + protected: + IndexStreamer::Stats &stats_; + std::mutex mutex_{}; + size_t max_index_size_{0UL}; + uint32_t chunk_size_{kDefaultChunkSize}; + uint32_t node_index_mask_bits_{0U}; + uint32_t node_cnt_per_chunk_{0U}; + uint32_t node_index_mask_{0U}; + uint32_t neighbor_size_{0U}; + bool get_vector_enabled_{false}; + bool use_key_info_map_{true}; + + mutable std::shared_ptr keys_map_lock_; + HashMapPointer keys_map_; + + ChunkBroker::Pointer broker_; + + mutable std::vector node_chunks_{}; + + private: + mutable std::vector dist_chunks_{}; + bool dist_loaded_{false}; + uint32_t dist_entry_size_{0}; // max_degree * sizeof(dist_t) +}; + +// --- Template specializations for typed MemoryBlock access --- + +template <> +inline NeighborsT +VamanaStreamerEntity::get_neighbors_typed(node_id_t id) const { + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = + (id & node_index_mask_) * node_size() + vector_size() + sizeof(key_t); + sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_); + ailego_assert_with(chunk_idx < node_chunks_.size(), "invalid chunk idx"); + ailego_assert_with(offset < node_chunks_[chunk_idx]->data_size(), + "invalid chunk offset"); + const void *ptr = nullptr; + size_t ret = node_chunks_[chunk_idx]->read(offset, &ptr, neighbor_size_); + if (ailego_unlikely(ret != neighbor_size_)) { + LOG_ERROR("Read neighbor header failed, ret=%zu", ret); + return NeighborsT(); + } + MmapMemoryBlock block(const_cast(ptr)); + return NeighborsT(std::move(block)); +} + +template <> +inline NeighborsT +VamanaStreamerEntity::get_neighbors_typed( + node_id_t id) const { + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = + (id & node_index_mask_) * node_size() + vector_size() + sizeof(key_t); + sync_chunks(ChunkBroker::CHUNK_TYPE_NODE, chunk_idx, &node_chunks_); + ailego_assert_with(chunk_idx < node_chunks_.size(), "invalid chunk idx"); + IndexStorage::MemoryBlock mem_block; + size_t ret = node_chunks_[chunk_idx]->read(offset, mem_block, neighbor_size_); + if (ailego_unlikely(ret != neighbor_size_)) { + LOG_ERROR("Read neighbor header failed, ret=%zu", ret); + return NeighborsT(); + } + BufferPoolMemoryBlock block(mem_block.buffer_pool_handle_, + mem_block.buffer_block_id_, mem_block.data_); + mem_block.buffer_pool_handle_ = nullptr; + return NeighborsT(std::move(block)); +} + +template <> +inline int VamanaStreamerEntity::get_vector_typed( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + vec_blocks.resize(count); + for (auto i = 0U; i < count; ++i) { + auto loc = get_vector_chunk_loc(ids[i]); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + const void *ptr = nullptr; + size_t ret = node_chunks_[loc.first]->read(loc.second, &ptr, vector_size()); + if (ailego_unlikely(ret != vector_size())) { + LOG_ERROR("Read vector failed, ret=%zu", ret); + return IndexError_ReadData; + } + vec_blocks[i].reset(const_cast(ptr)); + } + return 0; +} + +template <> +inline int VamanaStreamerEntity::get_vector_typed( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + vec_blocks.resize(count); + for (auto i = 0U; i < count; ++i) { + auto loc = get_vector_chunk_loc(ids[i]); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + IndexStorage::MemoryBlock mem_block; + size_t ret = + node_chunks_[loc.first]->read(loc.second, mem_block, vector_size()); + if (ailego_unlikely(ret != vector_size())) { + LOG_ERROR("Read vector failed, ret=%zu", ret); + return IndexError_ReadData; + } + vec_blocks[i] = + BufferPoolMemoryBlock(mem_block.buffer_pool_handle_, + mem_block.buffer_block_id_, mem_block.data_); + mem_block.buffer_pool_handle_ = nullptr; + } + return 0; +} + +template <> +inline key_t VamanaStreamerEntity::get_key_typed( + node_id_t id) const { + if (!use_key_info_map_) return id; + auto loc = get_key_chunk_loc(id); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + const void *ptr = nullptr; + size_t ret = node_chunks_[loc.first]->read(loc.second, &ptr, sizeof(key_t)); + if (ailego_unlikely(ret != sizeof(key_t))) { + LOG_ERROR("Read key failed, ret=%zu", ret); + return kInvalidKey; + } + return *reinterpret_cast(ptr); +} + +template <> +inline key_t VamanaStreamerEntity::get_key_typed( + node_id_t id) const { + if (!use_key_info_map_) return id; + auto loc = get_key_chunk_loc(id); + ailego_assert_with(loc.first < node_chunks_.size(), "invalid chunk idx"); + IndexStorage::MemoryBlock key_block; + size_t ret = + node_chunks_[loc.first]->read(loc.second, key_block, sizeof(key_t)); + if (ailego_unlikely(ret != sizeof(key_t))) { + LOG_ERROR("Read key failed, ret=%zu", ret); + return kInvalidKey; + } + return *reinterpret_cast(key_block.data()); +} + +// --- Typed entity subclass for mmap mode --- +// Caches chunk base addresses to eliminate virtual function calls on the +// search hot path. For mmap mode, chunk data is memory-mapped at init time, +// so we can directly compute pointers via base_addr + offset. +class VamanaMmapStreamerEntity : public VamanaStreamerEntity { + public: + using MemoryBlock = MmapMemoryBlock; + using TypedNeighbors = NeighborsT; + + using VamanaStreamerEntity::VamanaStreamerEntity; + + VamanaStorageMode storage_mode() const override { + return VamanaStorageMode::kMmap; + } + + //! Override clone to return correct subclass type, so that + //! static_cast in the algorithm is safe. + const VamanaEntity::Pointer clone() const override; + + inline TypedNeighbors get_neighbors_typed(node_id_t id) const { + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = + (id & node_index_mask_) * node_size() + vector_size() + sizeof(key_t); + const char *base = get_node_chunk_base(chunk_idx); + MmapMemoryBlock block(const_cast(base + offset)); + return TypedNeighbors(std::move(block)); + } + + inline int get_vector_typed(const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + vec_blocks.resize(count); + for (auto i = 0U; i < count; ++i) { + uint32_t chunk_idx = ids[i] >> node_index_mask_bits_; + uint32_t offset = (ids[i] & node_index_mask_) * node_size(); + const char *base = get_node_chunk_base(chunk_idx); + vec_blocks[i].reset(const_cast(base + offset)); + } + return 0; + } + + inline key_t get_key_typed(node_id_t id) const { + if (!use_key_info_map_) return id; + uint32_t chunk_idx = id >> node_index_mask_bits_; + uint32_t offset = (id & node_index_mask_) * node_size() + vector_size(); + const char *base = get_node_chunk_base(chunk_idx); + return *reinterpret_cast(base + offset); + } + + private: + inline const char *get_node_chunk_base(uint32_t chunk_idx) const { + if (ailego_unlikely(chunk_idx >= node_chunk_bases_.size())) { + sync_node_chunk_bases(chunk_idx); + } + return node_chunk_bases_[chunk_idx]; + } + + void sync_node_chunk_bases(uint32_t chunk_idx) const { + sync_node_chunks(chunk_idx); + const auto &chunks = node_chunks_; + for (size_t i = node_chunk_bases_.size(); i <= chunk_idx; ++i) { + const void *ptr = nullptr; + chunks[i]->read(0, &ptr, 1); + node_chunk_bases_.push_back(static_cast(ptr)); + } + } + + mutable std::vector node_chunk_bases_{}; +}; + +// --- Typed entity subclass for buffer pool mode --- +class VamanaBufferPoolStreamerEntity : public VamanaStreamerEntity { + public: + using MemoryBlock = BufferPoolMemoryBlock; + using TypedNeighbors = NeighborsT; + + using VamanaStreamerEntity::VamanaStreamerEntity; + + VamanaStorageMode storage_mode() const override { + return VamanaStorageMode::kBufferPool; + } + + inline TypedNeighbors get_neighbors_typed(node_id_t id) const { + return VamanaStreamerEntity::get_neighbors_typed(id); + } + + inline int get_vector_typed( + const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + return VamanaStreamerEntity::get_vector_typed( + ids, count, vec_blocks); + } + + inline key_t get_key_typed(node_id_t id) const { + return VamanaStreamerEntity::get_key_typed(id); + } +}; + +// --- Typed entity subclass for contiguous memory mode --- +// Allocates contiguous memory and copies all chunk data into it. +// Access is via a single base pointer + offset, eliminating chunk-level +// indirection and maximizing memory locality. +class VamanaContiguousStreamerEntity : public VamanaMmapStreamerEntity { + public: + using VamanaMmapStreamerEntity::VamanaMmapStreamerEntity; + + VamanaStorageMode storage_mode() const override { + return VamanaStorageMode::kContiguous; + } + + //! Override clone to return correct subclass type. + //! Cloned entity shares contiguous memory via shared_ptr. + const VamanaEntity::Pointer clone() const override; + + ~VamanaContiguousStreamerEntity() = default; + + // Build contiguous memory from chunks after open. + int build_contiguous_memory(); + + //! Degrade to mmap mode by releasing contiguous memory and falling back + //! to chunk-based access. + void degrade_to_mmap() { + node_memory_.reset(); + node_base_ = nullptr; + LOG_INFO("Vamana contiguous entity degraded to mmap mode for insertion"); + } + + bool is_contiguous() const { + return node_base_ != nullptr; + } + + int add_vector(key_t key, const void *vec, node_id_t *id) override { + if (ailego_unlikely(is_contiguous())) degrade_to_mmap(); + return VamanaMmapStreamerEntity::add_vector(key, vec, id); + } + + int add_vector_with_id(node_id_t id, const void *vec) override { + if (ailego_unlikely(is_contiguous())) degrade_to_mmap(); + return VamanaMmapStreamerEntity::add_vector_with_id(id, vec); + } + + inline TypedNeighbors get_neighbors_typed(node_id_t id) const { + if (ailego_likely(node_base_ != nullptr)) { + const char *ptr = node_base_ + static_cast(id) * node_size() + + vector_size() + sizeof(key_t); + MmapMemoryBlock block(const_cast(ptr)); + return TypedNeighbors(std::move(block)); + } + return VamanaMmapStreamerEntity::get_neighbors_typed(id); + } + + inline int get_vector_typed(const node_id_t *ids, uint32_t count, + std::vector &vec_blocks) const { + if (ailego_likely(node_base_ != nullptr)) { + vec_blocks.resize(count); + for (auto i = 0U; i < count; ++i) { + const char *ptr = + node_base_ + static_cast(ids[i]) * node_size(); + vec_blocks[i].reset(const_cast(ptr)); + } + return 0; + } + return VamanaMmapStreamerEntity::get_vector_typed(ids, count, vec_blocks); + } + + inline key_t get_key_typed(node_id_t id) const { + if (ailego_likely(node_base_ != nullptr)) { + if (!use_key_info_map_) return id; + const char *ptr = + node_base_ + static_cast(id) * node_size() + vector_size(); + return *reinterpret_cast(ptr); + } + return VamanaMmapStreamerEntity::get_key_typed(id); + } + + protected: + //! Custom deleter for contiguous memory (munmap / _aligned_free / free) + struct ContiguousDeleter { + size_t size; + void operator()(char *ptr) const { + if (!ptr) return; +#if defined(__linux__) || defined(__APPLE__) + ::munmap(ptr, size); +#elif defined(_WIN32) + ::_aligned_free(ptr); +#else + std::free(ptr); +#endif + } + }; + + //! Shared ownership of contiguous memory (enables zero-copy clone) + std::shared_ptr node_memory_{}; + + //! Raw pointer for hot-path access (derived from shared_ptr) + char *node_base_{nullptr}; + + private: + static char *allocate_contiguous(size_t size); +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/interface/index_factory.cc b/src/core/interface/index_factory.cc index 0d81572..1eb1604 100644 --- a/src/core/interface/index_factory.cc +++ b/src/core/interface/index_factory.cc @@ -47,6 +47,8 @@ Index::Pointer IndexFactory::CreateAndInitIndex(const BaseIndexParam ¶m) { ptr = std::make_shared(); } else if (param.index_type == IndexType::kHNSWRabitq) { ptr = std::make_shared(); + } else if (param.index_type == IndexType::kVamana) { + ptr = std::make_shared(); } else { LOG_ERROR("Unsupported index type: "); return nullptr; @@ -115,6 +117,14 @@ BaseIndexParam::Pointer IndexFactory::DeserializeIndexParamFromJson( } return param; } + case IndexType::kVamana: { + VamanaIndexParam::Pointer param = std::make_shared(); + if (!param->DeserializeFromJson(json_str)) { + LOG_ERROR("Failed to deserialize vamana index param"); + return nullptr; + } + return param; + } default: LOG_ERROR("Unsupported index type: %s", magic_enum::enum_name(index_type).data()); @@ -167,6 +177,11 @@ std::string IndexFactory::QueryParamSerializeToJson(const QueryParamType ¶m, json_obj.set("ef_search", ailego::JsonValue(param.ef_search)); } index_type = IndexType::kHNSWRabitq; + } else if constexpr (std::is_same_v) { + if (!omit_empty_value || param.ef_search != 0) { + json_obj.set("ef_search", ailego::JsonValue(param.ef_search)); + } + index_type = IndexType::kVamana; } json_obj.set("index_type", @@ -272,6 +287,17 @@ typename QueryParamType::Pointer IndexFactory::QueryParamDeserializeFromJson( return nullptr; } return param; + } else if (index_type == IndexType::kVamana) { + auto param = std::make_shared(); + if (!parse_common_fields(param)) { + return nullptr; + } + if (!extract_value_from_json(json_obj, "ef_search", param->ef_search, + tmp_json_value)) { + LOG_ERROR("Failed to deserialize ef_search"); + return nullptr; + } + return param; } else { LOG_ERROR("Unsupported index type: %s", magic_enum::enum_name(index_type).data()); @@ -301,6 +327,12 @@ typename QueryParamType::Pointer IndexFactory::QueryParamDeserializeFromJson( LOG_ERROR("Failed to deserialize ef_search"); return nullptr; } + } else if constexpr (std::is_same_v) { + if (!extract_value_from_json(json_obj, "ef_search", param->ef_search, + tmp_json_value)) { + LOG_ERROR("Failed to deserialize ef_search"); + return nullptr; + } } else { LOG_ERROR("Unsupported index type: %s", magic_enum::enum_name(index_type).data()); @@ -319,5 +351,9 @@ template HNSWQueryParam::Pointer IndexFactory::QueryParamDeserializeFromJson< HNSWQueryParam>(const std::string &json_str); template IVFQueryParam::Pointer IndexFactory::QueryParamDeserializeFromJson< IVFQueryParam>(const std::string &json_str); +template std::string IndexFactory::QueryParamSerializeToJson( + const VamanaQueryParam ¶m, bool omit_empty_value); +template VamanaQueryParam::Pointer IndexFactory::QueryParamDeserializeFromJson< + VamanaQueryParam>(const std::string &json_str); } // namespace zvec::core_interface diff --git a/src/core/interface/index_param.cc b/src/core/interface/index_param.cc index e21256b..962a78e 100644 --- a/src/core/interface/index_param.cc +++ b/src/core/interface/index_param.cc @@ -83,6 +83,10 @@ ailego::JsonObject HNSWIndexParam::SerializeToJsonObject( auto json_obj = BaseIndexParam::SerializeToJsonObject(omit_empty_value); json_obj.set("m", ailego::JsonValue(m)); json_obj.set("ef_construction", ailego::JsonValue(ef_construction)); + if (!omit_empty_value || use_contiguous_memory) { + json_obj.set("use_contiguous_memory", + ailego::JsonValue(use_contiguous_memory)); + } return json_obj; } @@ -137,6 +141,7 @@ bool HNSWIndexParam::DeserializeFromJsonObject( DESERIALIZE_VALUE_FIELD(json_obj, m); DESERIALIZE_VALUE_FIELD(json_obj, ef_construction); + DESERIALIZE_VALUE_FIELD(json_obj, use_contiguous_memory); return true; } @@ -174,6 +179,47 @@ ailego::JsonObject HNSWRabitqIndexParam::SerializeToJsonObject( return json_obj; } +ailego::JsonObject VamanaIndexParam::SerializeToJsonObject( + bool omit_empty_value) const { + auto json_obj = BaseIndexParam::SerializeToJsonObject(omit_empty_value); + json_obj.set("max_degree", ailego::JsonValue(max_degree)); + json_obj.set("search_list_size", ailego::JsonValue(search_list_size)); + json_obj.set("alpha", ailego::JsonValue(alpha)); + if (!omit_empty_value || + max_occlusion_size != static_cast(kDefaultVamanaMaxOcclusionSize)) { + json_obj.set("max_occlusion_size", ailego::JsonValue(max_occlusion_size)); + } + if (!omit_empty_value || saturate_graph) { + json_obj.set("saturate_graph", ailego::JsonValue(saturate_graph)); + } + if (!omit_empty_value || use_contiguous_memory) { + json_obj.set("use_contiguous_memory", + ailego::JsonValue(use_contiguous_memory)); + } + return json_obj; +} + +bool VamanaIndexParam::DeserializeFromJsonObject( + const ailego::JsonObject &json_obj) { + if (!BaseIndexParam::DeserializeFromJsonObject(json_obj)) { + return false; + } + + if (index_type != IndexType::kVamana) { + LOG_ERROR("index_type is not kVamana"); + return false; + } + + DESERIALIZE_VALUE_FIELD(json_obj, max_degree); + DESERIALIZE_VALUE_FIELD(json_obj, search_list_size); + DESERIALIZE_VALUE_FIELD(json_obj, alpha); + DESERIALIZE_VALUE_FIELD(json_obj, max_occlusion_size); + DESERIALIZE_VALUE_FIELD(json_obj, saturate_graph); + DESERIALIZE_VALUE_FIELD(json_obj, use_contiguous_memory); + + return true; +} + ailego::JsonObject QuantizerParam::SerializeToJsonObject( bool omit_empty_value) const { ailego::JsonObject json_obj; diff --git a/src/core/interface/indexes/hnsw_index.cc b/src/core/interface/indexes/hnsw_index.cc index 921db8f..816665b 100644 --- a/src/core/interface/indexes/hnsw_index.cc +++ b/src/core/interface/indexes/hnsw_index.cc @@ -16,10 +16,32 @@ #include #include #include "algorithm/hnsw/hnsw_params.h" +#include "algorithm/hnsw/hnsw_streamer.h" +#include "algorithm/hnsw/hnsw_streamer_entity.h" #include "algorithm/hnsw_sparse/hnsw_sparse_params.h" namespace zvec::core_interface { +std::string HNSWIndex::storage_mode() const { + if (!streamer_) { + return ""; + } + auto *hnsw_streamer = dynamic_cast(streamer_.get()); + if (!hnsw_streamer) { + // e.g. sparse branch uses HnswSparseStreamer which is a different type + return ""; + } + switch (hnsw_streamer->storage_mode()) { + case core::HnswStorageMode::kMmap: + return "mmap"; + case core::HnswStorageMode::kBufferPool: + return "buffer_pool"; + case core::HnswStorageMode::kContiguous: + return "contiguous"; + } + return ""; +} + int HNSWIndex::CreateAndInitStreamer(const BaseIndexParam ¶m) { param_ = dynamic_cast(param); @@ -57,6 +79,8 @@ int HNSWIndex::CreateAndInitStreamer(const BaseIndexParam ¶m) { kDefaultHnswEfSearch); proxima_index_params_.set(core::PARAM_HNSW_STREAMER_USE_ID_MAP, param_.use_id_map); + proxima_index_params_.set(core::PARAM_HNSW_STREAMER_USE_CONTIGUOUS_MEMORY, + param_.use_contiguous_memory); streamer_ = core::IndexFactory::CreateStreamer("HnswStreamer"); } diff --git a/src/core/interface/indexes/vamana_index.cc b/src/core/interface/indexes/vamana_index.cc new file mode 100644 index 0000000..297cf37 --- /dev/null +++ b/src/core/interface/indexes/vamana_index.cc @@ -0,0 +1,105 @@ +// Copyright 2025-present the zvec 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 "algorithm/vamana/vamana_params.h" + +namespace zvec::core_interface { + +int VamanaIndex::CreateAndInitStreamer(const BaseIndexParam ¶m) { + param_ = dynamic_cast(param); + + // Validate parameters + param_.max_degree = std::max(5, std::min(256, param_.max_degree)); + param_.search_list_size = + std::max(10, std::min(2048, param_.search_list_size)); + if (param_.alpha <= 0.0f) param_.alpha = kDefaultVamanaAlpha; + + proxima_index_params_.set(core::PARAM_VAMANA_STREAMER_MAX_DEGREE, + static_cast(param_.max_degree)); + proxima_index_params_.set(core::PARAM_VAMANA_STREAMER_SEARCH_LIST_SIZE, + static_cast(param_.search_list_size)); + proxima_index_params_.set(core::PARAM_VAMANA_STREAMER_ALPHA, param_.alpha); + proxima_index_params_.set(core::PARAM_VAMANA_STREAMER_MAX_OCCLUSION_SIZE, + static_cast(param_.max_occlusion_size)); + proxima_index_params_.set(core::PARAM_VAMANA_STREAMER_SATURATE_GRAPH, + param_.saturate_graph); + proxima_index_params_.set(core::PARAM_VAMANA_STREAMER_GET_VECTOR_ENABLE, + true); + proxima_index_params_.set(core::PARAM_VAMANA_STREAMER_EF, + kDefaultVamanaEfSearch); + proxima_index_params_.set(core::PARAM_VAMANA_STREAMER_USE_ID_MAP, + param_.use_id_map); + proxima_index_params_.set(core::PARAM_VAMANA_STREAMER_USE_CONTIGUOUS_MEMORY, + param_.use_contiguous_memory); + + streamer_ = core::IndexFactory::CreateStreamer("VamanaStreamer"); + + if (ailego_unlikely(!streamer_)) { + LOG_ERROR("Failed to create VamanaStreamer"); + return core::IndexError_Runtime; + } + if (ailego_unlikely( + streamer_->init(proxima_index_meta_, proxima_index_params_) != 0)) { + LOG_ERROR("Failed to init VamanaStreamer"); + return core::IndexError_Runtime; + } + return 0; +} + +int VamanaIndex::_prepare_for_search( + const VectorData & /*vector_data*/, + const BaseIndexQueryParam::Pointer &search_param, + core::IndexContext::Pointer &context) { + const auto &vamana_search_param = + std::dynamic_pointer_cast(search_param); + + if (ailego_unlikely(!vamana_search_param)) { + LOG_ERROR("Invalid search param type, expected VamanaQueryParam"); + return core::IndexError_Runtime; + } + + if (vamana_search_param->ef_search == 0 || + vamana_search_param->ef_search > 2048) { + LOG_ERROR( + "ef_search must be greater than 0 and less than or equal to 2048."); + return core::IndexError_Runtime; + } + + context->set_topk(vamana_search_param->topk); + context->set_fetch_vector(vamana_search_param->fetch_vector); + if (vamana_search_param->filter) { + context->set_filter(std::move(*vamana_search_param->filter)); + } + if (vamana_search_param->radius > 0.0f) { + context->set_threshold(vamana_search_param->radius); + } + ailego::Params params; + const uint32_t real_search_ef = + std::max(1u, std::min(2048u, vamana_search_param->ef_search)); + params.set(core::PARAM_VAMANA_STREAMER_EF, real_search_ef); + context->update(params); + return 0; +} + +int VamanaIndex::_get_coarse_search_topk( + const BaseIndexQueryParam::Pointer &search_param) { + const auto &vamana_search_param = + std::dynamic_pointer_cast(search_param); + return std::max(search_param->topk, vamana_search_param->ef_search); +} + +} // namespace zvec::core_interface diff --git a/src/core/utility/buffer_storage.cc b/src/core/utility/buffer_storage.cc index aa2fd19..62d442a 100644 --- a/src/core/utility/buffer_storage.cc +++ b/src/core/utility/buffer_storage.cc @@ -173,6 +173,11 @@ class BufferStorage : public IndexStorage { this->cleanup(); } + //! Retrieve the memory block type of this storage + MemoryBlock::MemoryBlockType memory_block_type(void) const override { + return MemoryBlock::MBT_BUFFERPOOL; + } + //! Initialize storage int init(const ailego::Params ¶ms) override { return 0; diff --git a/src/db/collection.cc b/src/db/collection.cc index e75492b..7cae6fc 100644 --- a/src/db/collection.cc +++ b/src/db/collection.cc @@ -35,6 +35,7 @@ #include "db/common/file_helper.h" #include "db/common/profiler.h" #include "db/common/typedef.h" +#include "db/index/column/vector_column/vector_column_indexer.h" #include "db/index/common/delete_store.h" #include "db/index/common/id_map.h" #include "db/index/common/index_filter.h" @@ -43,6 +44,7 @@ #include "db/index/segment/segment_helper.h" #include "db/index/segment/segment_manager.h" #include "db/sqlengine/sqlengine.h" +#include "zvec/core/interface/index.h" namespace zvec { @@ -121,6 +123,9 @@ class CollectionImpl : public Collection { Result Fetch(const std::vector &pks) const override; + Result DebugGetHnswStorageMode( + const std::string &column_name) const override; + private: void prepare_schema(); @@ -1636,6 +1641,50 @@ Result CollectionImpl::Fetch( return results; } +Result CollectionImpl::DebugGetHnswStorageMode( + const std::string &column_name) const { + std::shared_lock lock(schema_handle_mtx_); + + CHECK_DESTROY_RETURN_STATUS_EXPECTED(destroyed_, false); + + // Try all segments (including the writing one). The first segment that has + // a fully-built HNSW index wins; if only a building segment exists we still + // surface its current storage mode so that tests can observe the entity + // type right after Open(). + auto segments = get_all_segments(); + + for (const auto &segment : segments) { + if (!segment) { + continue; + } + auto indexers = segment->get_vector_indexer(column_name); + for (const auto &indexer : indexers) { + if (!indexer) { + continue; + } + auto index = indexer->debug_get_index(); + if (!index) { + continue; + } + auto *hnsw_index = dynamic_cast(index.get()); + if (!hnsw_index) { + return tl::make_unexpected(Status::InvalidArgument( + "Column '", column_name, + "' does not have an HNSW index (or index is sparse)")); + } + auto mode = hnsw_index->storage_mode(); + if (mode.empty()) { + // streamer not initialized yet; skip and look at other segments + continue; + } + return mode; + } + } + + return tl::make_unexpected( + Status::NotFound("No HNSW index found for column '", column_name, "'")); +} + Status CollectionImpl::recovery() { if (!FileHelper::DirectoryExists(path_.c_str())) { return Status::InvalidArgument("collection path{", path_, "} not exist."); diff --git a/src/db/index/column/vector_column/engine_helper.hpp b/src/db/index/column/vector_column/engine_helper.hpp index e7e323f..0901f0d 100644 --- a/src/db/index/column/vector_column/engine_helper.hpp +++ b/src/db/index/column/vector_column/engine_helper.hpp @@ -202,6 +202,26 @@ class ProximaEngineHelper { } return std::move(ivf_query_param); } + + case IndexType::VAMANA: { + auto vamana_query_param_result = + _build_common_query_param( + query_params); + if (!vamana_query_param_result.has_value()) { + return tl::make_unexpected(Status::InvalidArgument( + "failed to build query param: " + + vamana_query_param_result.error().message())); + } + auto &vamana_query_param = vamana_query_param_result.value(); + if (query_params.query_params) { + auto db_vamana_query_params = dynamic_cast( + query_params.query_params.get()); + vamana_query_param->ef_search = + static_cast(db_vamana_query_params->ef_search()); + } + return std::move(vamana_query_param); + } + default: return tl::make_unexpected(Status::InvalidArgument("not supported")); } @@ -346,6 +366,8 @@ class ProximaEngineHelper { index_param_builder->WithM(db_index_params->m()); index_param_builder->WithEFConstruction( db_index_params->ef_construction()); + index_param_builder->WithUseContiguousMemory( + db_index_params->use_contiguous_memory()); return index_param_builder->Build(); } @@ -395,6 +417,37 @@ class ProximaEngineHelper { return index_param_builder->Build(); } + case IndexType::VAMANA: { + auto index_param_builder_result = + _build_common_index_param( + field_schema); + if (!index_param_builder_result.has_value()) { + return tl::make_unexpected(Status::InvalidArgument( + "failed to build index param: " + + index_param_builder_result.error().message())); + } + auto index_param_builder = index_param_builder_result.value(); + + auto db_index_params = dynamic_cast( + field_schema.index_params().get()); + index_param_builder->WithMaxDegree(db_index_params->max_degree()); + index_param_builder->WithSearchListSize( + db_index_params->search_list_size()); + index_param_builder->WithAlpha(db_index_params->alpha()); + index_param_builder->WithSaturateGraph( + db_index_params->saturate_graph()); + index_param_builder->WithUseContiguousMemory( + db_index_params->use_contiguous_memory()); + // db_index_params->use_id_map() is intentionally ignored here: + // db ensures id is consecutive (see _build_common_index_param), so + // the engine-level use_id_map is forced to false in the common + // builder. The flag is preserved on the db-side params for schema + // round-trip / introspection only. + + return index_param_builder->Build(); + } + default: return tl::make_unexpected(Status::InvalidArgument("not supported")); } diff --git a/src/db/index/column/vector_column/vector_column_indexer.h b/src/db/index/column/vector_column/vector_column_indexer.h index 0006080..f3b2286 100644 --- a/src/db/index/column/vector_column/vector_column_indexer.h +++ b/src/db/index/column/vector_column/vector_column_indexer.h @@ -105,6 +105,12 @@ class VectorColumnIndexer { return index->GetDocCount(); } + //! Debug-only accessor for the underlying core_interface Index. + //! Intended for introspection/testing; not part of the stable API. + core_interface::Index::Pointer debug_get_index() const { + return index; + } + // for ut protected: VectorColumnIndexer() = default; diff --git a/src/db/index/common/proto_converter.cc b/src/db/index/common/proto_converter.cc index 46eb93f..d58dc18 100644 --- a/src/db/index/common/proto_converter.cc +++ b/src/db/index/common/proto_converter.cc @@ -21,7 +21,8 @@ HnswIndexParams::OPtr ProtoConverter::FromPb( auto params = std::make_shared( MetricTypeCodeBook::Get(params_pb.base().metric_type()), params_pb.m(), params_pb.ef_construction(), - QuantizeTypeCodeBook::Get(params_pb.base().quantize_type())); + QuantizeTypeCodeBook::Get(params_pb.base().quantize_type()), + params_pb.use_contiguous_memory()); return params; } @@ -34,6 +35,7 @@ proto::HnswIndexParams ProtoConverter::ToPb(const HnswIndexParams *params) { QuantizeTypeCodeBook::Get(params->quantize_type())); params_pb.set_ef_construction(params->ef_construction()); params_pb.set_m(params->m()); + params_pb.set_use_contiguous_memory(params->use_contiguous_memory()); return params_pb; } @@ -101,6 +103,32 @@ proto::IVFIndexParams ProtoConverter::ToPb(const IVFIndexParams *params) { return params_pb; } +// VamanaIndexParams +VamanaIndexParams::OPtr ProtoConverter::FromPb( + const proto::VamanaIndexParams ¶ms_pb) { + return std::make_shared( + MetricTypeCodeBook::Get(params_pb.base().metric_type()), + params_pb.max_degree(), params_pb.search_list_size(), params_pb.alpha(), + params_pb.saturate_graph(), params_pb.use_contiguous_memory(), + params_pb.use_id_map(), + QuantizeTypeCodeBook::Get(params_pb.base().quantize_type())); +} + +proto::VamanaIndexParams ProtoConverter::ToPb(const VamanaIndexParams *params) { + proto::VamanaIndexParams params_pb; + params_pb.mutable_base()->set_metric_type( + MetricTypeCodeBook::Get(params->metric_type())); + params_pb.mutable_base()->set_quantize_type( + QuantizeTypeCodeBook::Get(params->quantize_type())); + params_pb.set_max_degree(params->max_degree()); + params_pb.set_search_list_size(params->search_list_size()); + params_pb.set_alpha(params->alpha()); + params_pb.set_saturate_graph(params->saturate_graph()); + params_pb.set_use_contiguous_memory(params->use_contiguous_memory()); + params_pb.set_use_id_map(params->use_id_map()); + return params_pb; +} + // InvertIndexParams InvertIndexParams::OPtr ProtoConverter::FromPb( const proto::InvertIndexParams ¶ms_pb) { @@ -185,6 +213,8 @@ IndexParams::Ptr ProtoConverter::FromPb(const proto::IndexParams ¶ms_pb) { return ProtoConverter::FromPb(params_pb.flat()); } else if (params_pb.has_hnsw_rabitq()) { return ProtoConverter::FromPb(params_pb.hnsw_rabitq()); + } else if (params_pb.has_vamana()) { + return ProtoConverter::FromPb(params_pb.vamana()); } return nullptr; @@ -246,6 +276,15 @@ proto::IndexParams ProtoConverter::ToPb(const IndexParams *params) { params_pb.mutable_hnsw_rabitq()->CopyFrom( ProtoConverter::ToPb(hnsw_rabitq_params)); } + break; + } + case IndexType::VAMANA: { + auto vamana_params = dynamic_cast(params); + if (vamana_params) { + params_pb.mutable_vamana()->CopyFrom( + ProtoConverter::ToPb(vamana_params)); + } + break; } default: break; diff --git a/src/db/index/common/proto_converter.h b/src/db/index/common/proto_converter.h index ad96007..362f950 100644 --- a/src/db/index/common/proto_converter.h +++ b/src/db/index/common/proto_converter.h @@ -38,6 +38,11 @@ struct ProtoConverter { static IVFIndexParams::OPtr FromPb(const proto::IVFIndexParams ¶ms_pb); static proto::IVFIndexParams ToPb(const IVFIndexParams *params); + // VamanaIndexParams + static VamanaIndexParams::OPtr FromPb( + const proto::VamanaIndexParams ¶ms_pb); + static proto::VamanaIndexParams ToPb(const VamanaIndexParams *params); + // InvertIndexParams static InvertIndexParams::OPtr FromPb( const proto::InvertIndexParams ¶ms_pb); diff --git a/src/db/index/common/schema.cc b/src/db/index/common/schema.cc index 971de61..1236f5f 100644 --- a/src/db/index/common/schema.cc +++ b/src/db/index/common/schema.cc @@ -54,7 +54,8 @@ std::unordered_set support_sparse_vector_type = { }; std::unordered_set support_dense_vector_index = { - IndexType::FLAT, IndexType::HNSW, IndexType::HNSW_RABITQ, IndexType::IVF}; + IndexType::FLAT, IndexType::HNSW, IndexType::HNSW_RABITQ, IndexType::IVF, + IndexType::VAMANA}; std::unordered_set support_sparse_vector_index = {IndexType::FLAT, IndexType::HNSW}; diff --git a/src/db/index/common/type_helper.h b/src/db/index/common/type_helper.h index 33d2ee3..02b7c0b 100644 --- a/src/db/index/common/type_helper.h +++ b/src/db/index/common/type_helper.h @@ -33,6 +33,8 @@ struct IndexTypeCodeBook { return IndexType::FLAT; case proto::IT_IVF: return IndexType::IVF; + case proto::IT_VAMANA: + return IndexType::VAMANA; case proto::IT_INVERT: return IndexType::INVERT; default: @@ -52,6 +54,8 @@ struct IndexTypeCodeBook { return proto::IT_FLAT; case IndexType::IVF: return proto::IT_IVF; + case IndexType::VAMANA: + return proto::IT_VAMANA; case IndexType::INVERT: return proto::IT_INVERT; default: @@ -71,6 +75,8 @@ struct IndexTypeCodeBook { return "FLAT"; case IndexType::IVF: return "IVF"; + case IndexType::VAMANA: + return "VAMANA"; case IndexType::INVERT: return "INVERT"; default: diff --git a/src/db/proto/zvec.proto b/src/db/proto/zvec.proto index 3c9d333..197914c 100644 --- a/src/db/proto/zvec.proto +++ b/src/db/proto/zvec.proto @@ -58,6 +58,8 @@ enum IndexType { IT_FLAT = 3; // Proxima HNSW RABITQ Index IT_HNSW_RABITQ = 4; + // Proxima Vamana (DiskANN) Index + IT_VAMANA = 5; // Invert Index IT_INVERT = 10; }; @@ -90,6 +92,10 @@ message HnswIndexParams { BaseIndexParams base = 1; int32 m = 2; int32 ef_construction = 3; + // When enabled, the HNSW streamer allocates a single contiguous memory + // arena for all graph nodes, which improves cache locality / search + // throughput at the cost of peak memory usage. Defaults to false. + bool use_contiguous_memory = 4; } message HnswRabitqIndexParams { @@ -112,6 +118,19 @@ message IVFIndexParams { bool use_soar = 4; } +message VamanaIndexParams { + BaseIndexParams base = 1; + int32 max_degree = 2; + int32 search_list_size = 3; + float alpha = 4; + bool saturate_graph = 5; + // When enabled, the Vamana streamer allocates a single contiguous memory + // arena for all graph nodes, which improves cache locality / search + // throughput at the cost of peak memory usage. Defaults to false. + bool use_contiguous_memory = 6; + bool use_id_map = 7; +} + message IndexParams { oneof params { InvertIndexParams invert = 1; @@ -119,6 +138,7 @@ message IndexParams { FlatIndexParams flat = 3; IVFIndexParams ivf = 4; HnswRabitqIndexParams hnsw_rabitq = 5; + VamanaIndexParams vamana = 6; }; }; diff --git a/src/include/zvec/core/framework/index_storage.h b/src/include/zvec/core/framework/index_storage.h index a4bec2c..ac1052e 100644 --- a/src/include/zvec/core/framework/index_storage.h +++ b/src/include/zvec/core/framework/index_storage.h @@ -273,6 +273,11 @@ class IndexStorage : public IndexModule { virtual bool isHugePage(void) const { return false; } + + //! Retrieve the memory block type of this storage + virtual MemoryBlock::MemoryBlockType memory_block_type(void) const { + return MemoryBlock::MBT_MMAP; + } }; } // namespace core diff --git a/src/include/zvec/core/interface/constants.h b/src/include/zvec/core/interface/constants.h index 1bc61dc..5b6f516 100644 --- a/src/include/zvec/core/interface/constants.h +++ b/src/include/zvec/core/interface/constants.h @@ -23,6 +23,13 @@ constexpr static uint32_t kDefaultHnswNeighborCnt = 50; constexpr static uint32_t kDefaultHnswEfSearch = 300; +constexpr static uint32_t kDefaultVamanaMaxDegree = 64; +constexpr static uint32_t kDefaultVamanaSearchListSize = 100; +constexpr static float kDefaultVamanaAlpha = 1.2f; +constexpr static uint32_t kDefaultVamanaEfSearch = 200; +constexpr static uint32_t kDefaultVamanaMaxOcclusionSize = 750; +constexpr static bool kDefaultVamanaSaturateGraph = false; + constexpr const uint32_t kDefaultRabitqTotalBits = 7; constexpr const uint32_t kDefaultRabitqNumClusters = 16; diff --git a/src/include/zvec/core/interface/index.h b/src/include/zvec/core/interface/index.h index 8634e39..a4f0e82 100644 --- a/src/include/zvec/core/interface/index.h +++ b/src/include/zvec/core/interface/index.h @@ -285,6 +285,13 @@ class HNSWIndex : public Index { public: HNSWIndex() = default; + //! Retrieve the storage mode of the underlying HNSW streamer entity. + //! Returns a string among {"mmap", "buffer_pool", "contiguous"}. + //! Intended for introspection and debug/testing usage. Returns empty + //! string when the streamer has not been initialized or is of an + //! unexpected type (e.g. the sparse branch). + std::string storage_mode() const; + protected: virtual int CreateAndInitStreamer(const BaseIndexParam ¶m) override; @@ -294,11 +301,27 @@ class HNSWIndex : public Index { int _get_coarse_search_topk( const BaseIndexQueryParam::Pointer &search_param) override; - private: HNSWIndexParam param_{}; }; +class VamanaIndex : public Index { + public: + VamanaIndex() = default; + + protected: + virtual int CreateAndInitStreamer(const BaseIndexParam ¶m) override; + + virtual int _prepare_for_search( + const VectorData &query, const BaseIndexQueryParam::Pointer &search_param, + core::IndexContext::Pointer &context) override; + int _get_coarse_search_topk( + const BaseIndexQueryParam::Pointer &search_param) override; + + private: + VamanaIndexParam param_{}; +}; + class HNSWRabitqIndex : public Index { public: HNSWRabitqIndex() = default; diff --git a/src/include/zvec/core/interface/index_param.h b/src/include/zvec/core/interface/index_param.h index 0d7bf30..48c6b84 100644 --- a/src/include/zvec/core/interface/index_param.h +++ b/src/include/zvec/core/interface/index_param.h @@ -63,6 +63,7 @@ enum class IndexType { kIVF, // it's actual a two-layer index kHNSW, kHNSWRabitq, + kVamana, }; enum class IVFSearchMethod { kBF, kHNSW }; @@ -309,6 +310,7 @@ struct HNSWIndexParam : public BaseIndexParam { using Pointer = std::shared_ptr; int m = kDefaultHnswNeighborCnt; int ef_construction = kDefaultHnswEfConstruction; + bool use_contiguous_memory = false; // Constructors with delegation HNSWIndexParam() : BaseIndexParam(IndexType::kHNSW) {} @@ -329,6 +331,46 @@ struct HNSWIndexParam : public BaseIndexParam { bool omit_empty_value = false) const override; }; +struct VamanaIndexParam : public BaseIndexParam { + using Pointer = std::shared_ptr; + int max_degree = kDefaultVamanaMaxDegree; + int search_list_size = kDefaultVamanaSearchListSize; + float alpha = kDefaultVamanaAlpha; + int max_occlusion_size = kDefaultVamanaMaxOcclusionSize; + bool saturate_graph = kDefaultVamanaSaturateGraph; + bool use_contiguous_memory = false; + + VamanaIndexParam() : BaseIndexParam(IndexType::kVamana) {} + + VamanaIndexParam(int max_degree, int search_list_size, float alpha) + : BaseIndexParam(IndexType::kVamana), + max_degree(max_degree), + search_list_size(search_list_size), + alpha(alpha) {} + + VamanaIndexParam(MetricType metric, int dim, int max_degree, + int search_list_size, float alpha) + : BaseIndexParam(IndexType::kVamana, metric, dim), + max_degree(max_degree), + search_list_size(search_list_size), + alpha(alpha) {} + + protected: + bool DeserializeFromJsonObject(const ailego::JsonObject &json_obj) override; + ailego::JsonObject SerializeToJsonObject( + bool omit_empty_value = false) const override; +}; + +struct VamanaQueryParam : public BaseIndexQueryParam { + using Pointer = std::shared_ptr; + + uint32_t ef_search = kDefaultVamanaEfSearch; + + BaseIndexQueryParam::Pointer Clone() const override { + return std::make_shared(*this); + } +}; + struct HNSWRabitqIndexParam : public BaseIndexParam { using Pointer = std::shared_ptr; diff --git a/src/include/zvec/core/interface/index_param_builders.h b/src/include/zvec/core/interface/index_param_builders.h index e22ecb3..5dcafbf 100644 --- a/src/include/zvec/core/interface/index_param_builders.h +++ b/src/include/zvec/core/interface/index_param_builders.h @@ -33,7 +33,7 @@ template class BaseIndexParamBuilder { // : public // std::enable_shared_from_this public: - BaseIndexParamBuilder() : param(std::make_shared()) {}; + BaseIndexParamBuilder() : param(std::make_shared()) {} virtual ~BaseIndexParamBuilder() = default; ActualIndexParamBuilderType &WithVersion(int version) { @@ -145,6 +145,10 @@ class HNSWIndexParamBuilder param->ef_construction = ef_construction; return *this; } + HNSWIndexParamBuilder &WithUseContiguousMemory(bool use_contiguous_memory) { + param->use_contiguous_memory = use_contiguous_memory; + return *this; + } std::shared_ptr Build() override { return param; @@ -191,6 +195,40 @@ class HNSWRabitqIndexParamBuilder } }; +class VamanaIndexParamBuilder + : public BaseIndexParamBuilder { + public: + VamanaIndexParamBuilder() = default; + VamanaIndexParamBuilder &WithMaxDegree(int max_degree) { + param->max_degree = max_degree; + return *this; + } + VamanaIndexParamBuilder &WithSearchListSize(int search_list_size) { + param->search_list_size = search_list_size; + return *this; + } + VamanaIndexParamBuilder &WithAlpha(float alpha) { + param->alpha = alpha; + return *this; + } + VamanaIndexParamBuilder &WithMaxOcclusionSize(int max_occlusion_size) { + param->max_occlusion_size = max_occlusion_size; + return *this; + } + VamanaIndexParamBuilder &WithSaturateGraph(bool saturate_graph) { + param->saturate_graph = saturate_graph; + return *this; + } + VamanaIndexParamBuilder &WithUseContiguousMemory(bool use_contiguous_memory) { + param->use_contiguous_memory = use_contiguous_memory; + return *this; + } + + std::shared_ptr Build() override { + return param; + } +}; + // class CompositeIndexParamBuilder : public // BaseIndexParamBuilder // { public: @@ -349,6 +387,21 @@ class HNSWRabitqQueryParamBuilder } }; +// Vamana builder (adds ef_search field) +class VamanaQueryParamBuilder + : public BaseIndexQueryParamBuilder { + public: + VamanaQueryParamBuilder &with_ef_search(int ef_search) { + m_param.ef_search = ef_search; + return *this; + } + + VamanaQueryParam::Pointer build() { + return std::make_shared(std::move(m_param)); + } +}; + // Example Usage: // // First, build the required nested params // auto nested_hnsw = HNSWQueryParamBuilder().with_ef_search(64).build(); diff --git a/src/include/zvec/db/collection.h b/src/include/zvec/db/collection.h index 57c2319..010ba36 100644 --- a/src/include/zvec/db/collection.h +++ b/src/include/zvec/db/collection.h @@ -103,6 +103,15 @@ class Collection { virtual Result Fetch( const std::vector &pks) const = 0; + + public: + //! Debug-only: retrieve the storage mode string of an HNSW index on the + //! given vector column. Returns one of {"mmap", "buffer_pool", + //! "contiguous"}. Returns an error Status when the column does not exist, + //! has no index, or the index is not an HNSW index. Intended for + //! introspection and testing; not part of the stable public API. + virtual Result DebugGetHnswStorageMode( + const std::string &column_name) const = 0; }; } // namespace zvec \ No newline at end of file diff --git a/src/include/zvec/db/index_params.h b/src/include/zvec/db/index_params.h index fcccf08..5f6faff 100644 --- a/src/include/zvec/db/index_params.h +++ b/src/include/zvec/db/index_params.h @@ -46,7 +46,8 @@ class IndexParams { bool is_vector_index_type() const { return type_ == IndexType::FLAT || type_ == IndexType::HNSW || - type_ == IndexType::HNSW_RABITQ || type_ == IndexType::IVF; + type_ == IndexType::HNSW_RABITQ || type_ == IndexType::IVF || + type_ == IndexType::VAMANA; } IndexType type() const { @@ -156,17 +157,20 @@ class HnswIndexParams : public VectorIndexParams { HnswIndexParams( MetricType metric_type, int m = core_interface::kDefaultHnswNeighborCnt, int ef_construction = core_interface::kDefaultHnswEfConstruction, - QuantizeType quantize_type = QuantizeType::UNDEFINED) + QuantizeType quantize_type = QuantizeType::UNDEFINED, + bool use_contiguous_memory = false) : VectorIndexParams(IndexType::HNSW, metric_type, quantize_type), m_(m), - ef_construction_(ef_construction) {} + ef_construction_(ef_construction), + use_contiguous_memory_(use_contiguous_memory) {} using OPtr = std::shared_ptr; public: Ptr clone() const override { return std::make_shared(metric_type_, m_, ef_construction_, - quantize_type_); + quantize_type_, + use_contiguous_memory_); } std::string to_string() const override { @@ -174,7 +178,8 @@ class HnswIndexParams : public VectorIndexParams { metric_type_, quantize_type_); std::ostringstream oss; oss << base_str << ",m:" << m_ << ",ef_construction:" << ef_construction_ - << "}"; + << ",use_contiguous_memory:" + << (use_contiguous_memory_ ? "true" : "false") << "}"; return oss.str(); } @@ -186,7 +191,9 @@ class HnswIndexParams : public VectorIndexParams { ef_construction_ == static_cast(other).ef_construction_ && quantize_type() == - static_cast(other).quantize_type(); + static_cast(other).quantize_type() && + use_contiguous_memory_ == static_cast(other) + .use_contiguous_memory_; } void set_m(int m) { @@ -202,9 +209,21 @@ class HnswIndexParams : public VectorIndexParams { return ef_construction_; } + void set_use_contiguous_memory(bool use_contiguous_memory) { + use_contiguous_memory_ = use_contiguous_memory; + } + bool use_contiguous_memory() const { + return use_contiguous_memory_; + } + protected: int m_; int ef_construction_; + // When enabled, HNSW streamer allocates a single contiguous memory arena + // for all graph nodes, improving cache locality and search throughput at + // the cost of peak memory usage. Defaults to false for backward + // compatibility. + bool use_contiguous_memory_{false}; }; class HnswRabitqIndexParams : public VectorIndexParams { @@ -428,4 +447,115 @@ class IVFIndexParams : public VectorIndexParams { bool use_soar_; }; +/* + * Vector: Vamana index params + */ +class VamanaIndexParams : public VectorIndexParams { + public: + VamanaIndexParams( + MetricType metric_type, + int max_degree = core_interface::kDefaultVamanaMaxDegree, + int search_list_size = core_interface::kDefaultVamanaSearchListSize, + float alpha = core_interface::kDefaultVamanaAlpha, + bool saturate_graph = core_interface::kDefaultVamanaSaturateGraph, + bool use_contiguous_memory = false, bool use_id_map = false, + QuantizeType quantize_type = QuantizeType::UNDEFINED) + : VectorIndexParams(IndexType::VAMANA, metric_type, quantize_type), + max_degree_(max_degree), + search_list_size_(search_list_size), + alpha_(alpha), + saturate_graph_(saturate_graph), + use_contiguous_memory_(use_contiguous_memory), + use_id_map_(use_id_map) {} + + using OPtr = std::shared_ptr; + + public: + Ptr clone() const override { + return std::make_shared( + metric_type_, max_degree_, search_list_size_, alpha_, saturate_graph_, + use_contiguous_memory_, use_id_map_, quantize_type_); + } + + std::string to_string() const override { + auto base_str = vector_index_params_to_string("VamanaIndexParams", + metric_type_, quantize_type_); + std::ostringstream oss; + oss << base_str << ",max_degree:" << max_degree_ + << ",search_list_size:" << search_list_size_ << ",alpha:" << alpha_ + << ",saturate_graph:" << (saturate_graph_ ? "true" : "false") + << ",use_contiguous_memory:" + << (use_contiguous_memory_ ? "true" : "false") + << ",use_id_map:" << (use_id_map_ ? "true" : "false") << "}"; + return oss.str(); + } + + bool operator==(const IndexParams &other) const override { + if (type() != other.type()) { + return false; + } + auto &rhs = static_cast(other); + return metric_type() == rhs.metric_type() && + quantize_type() == rhs.quantize_type() && + max_degree_ == rhs.max_degree_ && + search_list_size_ == rhs.search_list_size_ && alpha_ == rhs.alpha_ && + saturate_graph_ == rhs.saturate_graph_ && + use_contiguous_memory_ == rhs.use_contiguous_memory_ && + use_id_map_ == rhs.use_id_map_; + } + + int max_degree() const { + return max_degree_; + } + void set_max_degree(int max_degree) { + max_degree_ = max_degree; + } + + int search_list_size() const { + return search_list_size_; + } + void set_search_list_size(int search_list_size) { + search_list_size_ = search_list_size; + } + + float alpha() const { + return alpha_; + } + void set_alpha(float alpha) { + alpha_ = alpha; + } + + bool saturate_graph() const { + return saturate_graph_; + } + void set_saturate_graph(bool saturate_graph) { + saturate_graph_ = saturate_graph; + } + + bool use_contiguous_memory() const { + return use_contiguous_memory_; + } + void set_use_contiguous_memory(bool use_contiguous_memory) { + use_contiguous_memory_ = use_contiguous_memory; + } + + bool use_id_map() const { + return use_id_map_; + } + void set_use_id_map(bool use_id_map) { + use_id_map_ = use_id_map; + } + + private: + int max_degree_; + int search_list_size_; + float alpha_; + bool saturate_graph_; + // When enabled, Vamana streamer allocates a single contiguous memory arena + // for all graph nodes, improving cache locality and search throughput at + // the cost of peak memory usage. + bool use_contiguous_memory_; + bool use_id_map_; +}; + } // namespace zvec \ No newline at end of file diff --git a/src/include/zvec/db/query_params.h b/src/include/zvec/db/query_params.h index ba62dab..fc06672 100644 --- a/src/include/zvec/db/query_params.h +++ b/src/include/zvec/db/query_params.h @@ -172,4 +172,29 @@ class FlatQueryParams : public QueryParams { float scale_factor_{10}; }; +class VamanaQueryParams : public QueryParams { + public: + VamanaQueryParams(int ef_search = core_interface::kDefaultVamanaEfSearch, + float radius = 0.0f, bool is_linear = false, + bool is_using_refiner = false) + : QueryParams(IndexType::VAMANA), ef_search_(ef_search) { + set_radius(radius); + set_is_linear(is_linear); + set_is_using_refiner(is_using_refiner); + } + + virtual ~VamanaQueryParams() = default; + + int ef_search() const { + return ef_search_; + } + + void set_ef_search(int ef_search) { + ef_search_ = ef_search; + } + + private: + int ef_search_; +}; + } // namespace zvec \ No newline at end of file diff --git a/src/include/zvec/db/type.h b/src/include/zvec/db/type.h index 1578f81..31b8850 100644 --- a/src/include/zvec/db/type.h +++ b/src/include/zvec/db/type.h @@ -26,6 +26,7 @@ enum class IndexType : uint32_t { IVF = 2, FLAT = 3, HNSW_RABITQ = 4, + VAMANA = 5, INVERT = 10, }; diff --git a/src/turbo/avx512_vnni/record_quantized_int8/squared_euclidean.cc b/src/turbo/avx512_vnni/record_quantized_int8/squared_euclidean.cc index 4bfba53..16590af 100644 --- a/src/turbo/avx512_vnni/record_quantized_int8/squared_euclidean.cc +++ b/src/turbo/avx512_vnni/record_quantized_int8/squared_euclidean.cc @@ -82,36 +82,71 @@ void squared_euclidean_int8_batch_distance(const void *const *vectors, if (original_dim <= 0) { return; } - - internal::ip_int8_batch_avx512_vnni(vectors, query, n, original_dim, - distances); + static constexpr size_t batch_size = 12; + static constexpr size_t prefetch_step = 2; + size_t i = 0; + float *dist_ptr = distances; + const int8_t *const *data_ptrs_ptr = + reinterpret_cast(vectors); const float *q_tail = reinterpret_cast( reinterpret_cast(query) + original_dim); - float qa = q_tail[0]; - float qb = q_tail[1]; - float qs = q_tail[2]; - float qs2 = q_tail[3]; + float qA = q_tail[0]; + float qB = q_tail[1]; + float qS = q_tail[2]; + float qS2 = q_tail[3]; + const float sum = qA * qS; + const float sum2 = qA * qA * qS2; - const float sum = qa * qs; - const float sum2 = qa * qa * qs2; - for (size_t i = 0; i < n; ++i) { + for (; i + batch_size <= n; i += batch_size) { + std::array prefetch_ptrs; + std::array ip_dists; + for (size_t j = 0; j < batch_size; ++j) { + if (i + j + batch_size * prefetch_step < n) { + prefetch_ptrs[j] = vectors[i + j + batch_size * prefetch_step]; + } else { + prefetch_ptrs[j] = nullptr; + } + } + internal::ip_int8_batch_avx512_vnni_impl( + query, &vectors[i], prefetch_ptrs, original_dim, ip_dists.data()); + for (size_t j = 0; j < batch_size; ++j) { + const float *m_tail = reinterpret_cast( + reinterpret_cast(data_ptrs_ptr[j]) + original_dim); + float mA = m_tail[0]; + float mB = m_tail[1]; + float mS = m_tail[2]; + float mS2 = m_tail[3]; + int int8_sum = reinterpret_cast(m_tail)[4]; + float result = ip_dists[j]; + result -= 128.0f * static_cast(int8_sum); + result = mA * mA * mS2 + sum2 - 2 * mA * qA * result + + (mB - qB) * (mB - qB) * original_dim + + 2 * (mB - qB) * (mS * mA - sum); + dist_ptr[j] = result; + } + dist_ptr += batch_size; + data_ptrs_ptr += batch_size; + } + for (; i < n; ++i) { + std::array prefetch_ptrs{nullptr}; + float ip_dist; + internal::ip_int8_batch_avx512_vnni_impl<1>( + query, &vectors[i], prefetch_ptrs, original_dim, &ip_dist); const float *m_tail = reinterpret_cast( - reinterpret_cast(vectors[i]) + original_dim); - float ma = m_tail[0]; - float mb = m_tail[1]; - float ms = m_tail[2]; - float ms2 = m_tail[3]; - // Correct for the +128 shift applied to the query during preprocessing: - // dpbusd computes sum(uint8_query[i] * int8_data[i]) - // = sum((int8_query[i] + 128) * int8_data[i]) - // = true_ip + 128 * sum(int8_data[i]) - // int8_sum is stored as the 5th int-sized field after the 4 floats. + reinterpret_cast(data_ptrs_ptr[0]) + original_dim); + float mA = m_tail[0]; + float mB = m_tail[1]; + float mS = m_tail[2]; + float mS2 = m_tail[3]; int int8_sum = reinterpret_cast(m_tail)[4]; - float &result = distances[i]; + float result = ip_dist; result -= 128.0f * static_cast(int8_sum); - result = ma * ma * ms2 + sum2 - 2 * ma * qa * result + - (mb - qb) * (mb - qb) * original_dim + - 2 * (mb - qb) * (ms * ma - sum); + result = mA * mA * mS2 + sum2 - 2 * mA * qA * result + + (mB - qB) * (mB - qB) * original_dim + + 2 * (mB - qB) * (mS * mA - sum); + *dist_ptr = result; + data_ptrs_ptr += 1; + dist_ptr += 1; } #else (void)vectors; diff --git a/tests/core/algorithm/CMakeLists.txt b/tests/core/algorithm/CMakeLists.txt index 9ef1ec2..36a1f3e 100644 --- a/tests/core/algorithm/CMakeLists.txt +++ b/tests/core/algorithm/CMakeLists.txt @@ -7,6 +7,7 @@ cc_directories(flat_sparse) cc_directories(ivf) cc_directories(hnsw) cc_directories(hnsw_sparse) +cc_directories(vamana) if(RABITQ_SUPPORTED) cc_directories(hnsw_rabitq) endif() diff --git a/tests/core/algorithm/hnsw/hnsw_streamer_test.cc b/tests/core/algorithm/hnsw/hnsw_streamer_test.cc index 3257494..2966b33 100644 --- a/tests/core/algorithm/hnsw/hnsw_streamer_test.cc +++ b/tests/core/algorithm/hnsw/hnsw_streamer_test.cc @@ -3772,6 +3772,213 @@ TEST_F(HnswStreamerTest, TestBasicRefiner) { #endif +TEST_F(HnswStreamerTest, TestContiguousMemorySearch) { + // Build index with mmap mode + auto storage = IndexFactory::CreateStorage("MMapFileStorage"); + ASSERT_NE(nullptr, storage); + ailego::Params stg_params; + ASSERT_EQ(0, storage->init(stg_params)); + ASSERT_EQ(0, storage->open(dir_ + "TestContiguous.index", true)); + + { + auto builder = IndexFactory::CreateStreamer("HnswStreamer"); + ASSERT_NE(nullptr, builder); + ailego::Params build_params; + build_params.set(PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT, 16U); + build_params.set(PARAM_HNSW_STREAMER_SCALING_FACTOR, 5U); + build_params.set(PARAM_HNSW_STREAMER_EFCONSTRUCTION, 32U); + build_params.set(PARAM_HNSW_STREAMER_EF, 16U); + build_params.set(PARAM_HNSW_STREAMER_BRUTE_FORCE_THRESHOLD, 2000U); + ASSERT_EQ(0, builder->init(*index_meta_ptr_, build_params)); + ASSERT_EQ(0, builder->open(storage)); + + auto ctx = builder->create_context(); + ASSERT_TRUE(!!ctx); + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); + NumericalVector vec(dim); + size_t cnt = 3000UL; + for (size_t i = 0; i < cnt; i++) { + for (size_t j = 0; j < dim; ++j) { + vec[j] = static_cast(i); + } + ASSERT_EQ(0, builder->add_impl(i, vec.data(), qmeta, ctx)); + } + ASSERT_EQ(0, builder->flush(0UL)); + ASSERT_EQ(0, builder->close()); + } + + // Re-open with contiguous memory mode + auto searcher = IndexFactory::CreateStreamer("HnswStreamer"); + ASSERT_NE(nullptr, searcher); + ailego::Params search_params; + search_params.set(PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT, 16U); + search_params.set(PARAM_HNSW_STREAMER_SCALING_FACTOR, 5U); + search_params.set(PARAM_HNSW_STREAMER_EFCONSTRUCTION, 32U); + search_params.set(PARAM_HNSW_STREAMER_EF, 16U); + search_params.set(PARAM_HNSW_STREAMER_BRUTE_FORCE_THRESHOLD, 2000U); + search_params.set(PARAM_HNSW_STREAMER_USE_CONTIGUOUS_MEMORY, true); + ASSERT_EQ(0, searcher->init(*index_meta_ptr_, search_params)); + ASSERT_EQ(0, searcher->open(storage)); + + size_t cnt = 3000UL; + size_t topk = 50; + NumericalVector vec(dim); + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); + auto linearCtx = searcher->create_context(); + auto knnCtx = searcher->create_context(); + linearCtx->set_topk(topk); + knnCtx->set_topk(topk); + int totalHits = 0; + int totalCnts = 0; + for (size_t i = 0; i < cnt; i++) { + for (size_t j = 0; j < dim; ++j) { + vec[j] = static_cast(i) + 0.1f; + } + ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, knnCtx)); + ASSERT_EQ(0, searcher->search_bf_impl(vec.data(), qmeta, linearCtx)); + auto &knnResult = knnCtx->result(); + ASSERT_EQ(topk, knnResult.size()); + auto &linearResult = linearCtx->result(); + ASSERT_EQ(topk, linearResult.size()); + ASSERT_EQ(i, linearResult[0].key()); + for (size_t k = 0; k < topk; ++k) { + totalCnts++; + for (size_t j = 0; j < topk; ++j) { + if (linearResult[j].key() == knnResult[k].key()) { + totalHits++; + break; + } + } + } + } + float recall = totalHits * 1.0f / totalCnts; + EXPECT_GT(recall, 0.90f); +} + +TEST_F(HnswStreamerTest, TestContiguousMultiThreadSearch) { + constexpr size_t dim_mt = 32; + IndexMeta meta(IndexMeta::DataType::DT_FP32, dim_mt); + meta.set_metric("SquaredEuclidean", 0, ailego::Params()); + + // Build with mmap mode + auto storage = IndexFactory::CreateStorage("MMapFileStorage"); + ASSERT_NE(nullptr, storage); + ailego::Params stg_params; + ASSERT_EQ(0, storage->init(stg_params)); + ASSERT_EQ(0, storage->open(dir_ + "TestContiguousMT", true)); + + { + ailego::Params build_params; + build_params.set(PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT, 128); + build_params.set(PARAM_HNSW_STREAMER_SCALING_FACTOR, 10); + build_params.set(PARAM_HNSW_STREAMER_EFCONSTRUCTION, 64); + build_params.set(PARAM_HNSW_STREAMER_MAX_INDEX_SIZE, 30 * 1024 * 1024U); + build_params.set(PARAM_HNSW_STREAMER_BRUTE_FORCE_THRESHOLD, 1000U); + build_params.set(PARAM_HNSW_STREAMER_EF, 32); + build_params.set(PARAM_HNSW_STREAMER_GET_VECTOR_ENABLE, true); + + auto builder = IndexFactory::CreateStreamer("HnswStreamer"); + ASSERT_NE(nullptr, builder); + ASSERT_EQ(0, builder->init(meta, build_params)); + ASSERT_EQ(0, builder->open(storage)); + + auto addVector = [&builder, dim_mt](int baseKey, size_t addCnt) { + NumericalVector vec(dim_mt); + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim_mt); + size_t succAdd = 0; + auto ctx = builder->create_context(); + for (size_t i = 0; i < addCnt; i++) { + for (size_t j = 0; j < dim_mt; ++j) { + vec[j] = static_cast(i + baseKey); + } + succAdd += !builder->add_impl(baseKey + i, vec.data(), qmeta, ctx); + } + builder->flush(0UL); + return succAdd; + }; + auto t1 = std::async(std::launch::async, addVector, 0, 1000); + auto t2 = std::async(std::launch::async, addVector, 1000, 1000); + auto t3 = std::async(std::launch::async, addVector, 2000, 1000); + ASSERT_EQ(1000U, t1.get()); + ASSERT_EQ(1000U, t2.get()); + ASSERT_EQ(1000U, t3.get()); + ASSERT_EQ(0, builder->close()); + } + + // Re-open with contiguous memory + ailego::Params search_params; + search_params.set(PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT, 128); + search_params.set(PARAM_HNSW_STREAMER_SCALING_FACTOR, 10); + search_params.set(PARAM_HNSW_STREAMER_EFCONSTRUCTION, 64); + search_params.set(PARAM_HNSW_STREAMER_MAX_INDEX_SIZE, 30 * 1024 * 1024U); + search_params.set(PARAM_HNSW_STREAMER_BRUTE_FORCE_THRESHOLD, 1000U); + search_params.set(PARAM_HNSW_STREAMER_EF, 32); + search_params.set(PARAM_HNSW_STREAMER_GET_VECTOR_ENABLE, true); + search_params.set(PARAM_HNSW_STREAMER_USE_CONTIGUOUS_MEMORY, true); + + auto searcher = IndexFactory::CreateStreamer("HnswStreamer"); + ASSERT_NE(nullptr, searcher); + ASSERT_EQ(0, searcher->init(meta, search_params)); + ASSERT_EQ(0, searcher->open(storage)); + + // Verify data via provider + auto provider = searcher->create_provider(); + auto iter = provider->create_iterator(); + ASSERT_TRUE(!!iter); + size_t total = 0; + while (iter->is_valid()) { + float *data = (float *)iter->data(); + for (size_t d = 0; d < dim_mt; ++d) { + ASSERT_FLOAT_EQ(static_cast(iter->key()), data[d]); + } + total++; + iter->next(); + } + ASSERT_EQ(3000, total); + + // Multi-thread search on contiguous memory + size_t topk = 100; + size_t cnt = 3000; + auto knnSearch = [&]() { + NumericalVector vec(dim_mt); + auto linearCtx = searcher->create_context(); + auto knnCtx = searcher->create_context(); + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim_mt); + linearCtx->set_topk(topk); + knnCtx->set_topk(topk); + size_t totalCnts = 0; + size_t totalHits = 0; + for (size_t i = 0; i < cnt; i += 1) { + for (size_t j = 0; j < dim_mt; ++j) { + vec[j] = static_cast(i) + 0.1f; + } + ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, knnCtx)); + ASSERT_EQ(0, searcher->search_bf_impl(vec.data(), qmeta, linearCtx)); + auto &knnResult = knnCtx->result(); + ASSERT_EQ(topk, knnResult.size()); + auto &linearResult = linearCtx->result(); + ASSERT_EQ(topk, linearResult.size()); + ASSERT_EQ(i, linearResult[0].key()); + for (size_t k = 0; k < topk; ++k) { + totalCnts++; + for (size_t j = 0; j < topk; ++j) { + if (linearResult[j].key() == knnResult[k].key()) { + totalHits++; + break; + } + } + } + } + ASSERT_TRUE((totalHits * 1.0f / totalCnts) > 0.80f); + }; + auto s1 = std::async(std::launch::async, knnSearch); + auto s2 = std::async(std::launch::async, knnSearch); + auto s3 = std::async(std::launch::async, knnSearch); + s1.wait(); + s2.wait(); + s3.wait(); +} + } // namespace core } // namespace zvec diff --git a/tests/core/algorithm/vamana/CMakeLists.txt b/tests/core/algorithm/vamana/CMakeLists.txt new file mode 100644 index 0000000..43d35b0 --- /dev/null +++ b/tests/core/algorithm/vamana/CMakeLists.txt @@ -0,0 +1,14 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) + +file(GLOB_RECURSE ALL_TEST_SRCS *_test.cc) + +foreach(CC_SRCS ${ALL_TEST_SRCS}) + get_filename_component(CC_TARGET ${CC_SRCS} NAME_WE) + cc_gtest( + NAME ${CC_TARGET} + STRICT + LIBS zvec_ailego core_framework core_utility core_metric core_quantizer core_knn_vamana core_knn_hnsw core_knn_flat + SRCS ${CC_SRCS} + INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm/vamana + ) +endforeach() diff --git a/tests/core/algorithm/vamana/vamana_streamer_test.cc b/tests/core/algorithm/vamana/vamana_streamer_test.cc new file mode 100644 index 0000000..645ce43 --- /dev/null +++ b/tests/core/algorithm/vamana/vamana_streamer_test.cc @@ -0,0 +1,721 @@ +// Copyright 2025-present the zvec 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 "vamana_streamer.h" +#include +#include +#ifndef _MSC_VER +#include +#include +#endif +#include +#include +#include +#include +#include +#include "tests/test_util.h" + +#if defined(__GNUC__) || defined(__GNUG__) +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wunused-result" +#endif + +using namespace std; +using namespace testing; +using namespace zvec::ailego; + +namespace zvec { +namespace core { + +constexpr size_t kDim = 16; + +class VamanaStreamerTest : public testing::Test { + protected: + void SetUp(void); + void TearDown(void); + + IndexStreamer::Pointer CreateVamanaStreamer( + const ailego::Params &extra_params = ailego::Params()); + + static std::string dir_; + static shared_ptr index_meta_ptr_; +}; + +std::string VamanaStreamerTest::dir_("vamana_streamer_test_dir/"); +shared_ptr VamanaStreamerTest::index_meta_ptr_; + +void VamanaStreamerTest::SetUp(void) { + index_meta_ptr_.reset(new (nothrow) + IndexMeta(IndexMeta::DataType::DT_FP32, kDim)); + index_meta_ptr_->set_metric("SquaredEuclidean", 0, ailego::Params()); + + zvec::test_util::RemoveTestPath(dir_); +} + +void VamanaStreamerTest::TearDown(void) { + zvec::test_util::RemoveTestPath(dir_); +} + +IndexStreamer::Pointer VamanaStreamerTest::CreateVamanaStreamer( + const ailego::Params &extra_params) { + auto streamer = IndexFactory::CreateStreamer("VamanaStreamer"); + if (!streamer) return nullptr; + + ailego::Params params; + params.set(PARAM_VAMANA_STREAMER_MAX_DEGREE, 32U); + params.set(PARAM_VAMANA_STREAMER_SEARCH_LIST_SIZE, 100U); + params.set(PARAM_VAMANA_STREAMER_ALPHA, 1.2f); + params.set(PARAM_VAMANA_STREAMER_EF, 64U); + params.set(PARAM_VAMANA_STREAMER_BRUTE_FORCE_THRESHOLD, 500U); + params.merge(extra_params); + + if (streamer->init(*index_meta_ptr_, params) != 0) { + return nullptr; + } + return streamer; +} + +TEST_F(VamanaStreamerTest, TestAddVector) { + auto streamer = CreateVamanaStreamer(); + ASSERT_NE(nullptr, streamer); + + auto storage = IndexFactory::CreateStorage("MMapFileStorage"); + ASSERT_NE(nullptr, storage); + ailego::Params stg_params; + ASSERT_EQ(0, storage->init(stg_params)); + ASSERT_EQ(0, storage->open(dir_ + "TestAddVector", true)); + ASSERT_EQ(0, streamer->open(storage)); + + auto ctx = streamer->create_context(); + ASSERT_TRUE(!!ctx); + + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, kDim); + for (size_t i = 0; i < 1000UL; i++) { + NumericalVector vec(kDim); + for (size_t j = 0; j < kDim; ++j) { + vec[j] = static_cast(i); + } + ASSERT_EQ(0, streamer->add_impl(i, vec.data(), qmeta, ctx)); + } + + streamer->flush(0UL); + streamer.reset(); +} + +TEST_F(VamanaStreamerTest, TestLinearSearch) { + auto streamer = CreateVamanaStreamer(); + ASSERT_NE(nullptr, streamer); + + auto storage = IndexFactory::CreateStorage("MMapFileStorage"); + ASSERT_NE(nullptr, storage); + ailego::Params stg_params; + ASSERT_EQ(0, storage->init(stg_params)); + ASSERT_EQ(0, storage->open(dir_ + "TestLinearSearch.index", true)); + ASSERT_EQ(0, streamer->open(storage)); + + size_t cnt = 5000UL; + auto ctx = streamer->create_context(); + ASSERT_TRUE(!!ctx); + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, kDim); + NumericalVector vec(kDim); + for (size_t i = 0; i < cnt; i++) { + for (size_t j = 0; j < kDim; ++j) { + vec[j] = static_cast(i); + } + ASSERT_EQ(0, streamer->add_impl(i, vec.data(), qmeta, ctx)); + } + + size_t topk = 3; + for (size_t i = 0; i < cnt; i += 1) { + for (size_t j = 0; j < kDim; ++j) { + vec[j] = static_cast(i); + } + ctx->set_topk(1U); + ASSERT_EQ(0, streamer->search_bf_impl(vec.data(), qmeta, ctx)); + auto &result1 = ctx->result(); + ASSERT_EQ(1UL, result1.size()); + ASSERT_EQ(i, result1[0].key()); + + for (size_t j = 0; j < kDim; ++j) { + vec[j] = static_cast(i) + 0.1f; + } + ctx->set_topk(topk); + ASSERT_EQ(0, streamer->search_bf_impl(vec.data(), qmeta, ctx)); + auto &result2 = ctx->result(); + ASSERT_EQ(topk, result2.size()); + ASSERT_EQ(i, result2[0].key()); + ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key()); + ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key()); + } +} + +TEST_F(VamanaStreamerTest, TestKnnSearch) { + auto streamer = CreateVamanaStreamer(); + ASSERT_NE(nullptr, streamer); + + ailego::Params stg_params; + auto storage = IndexFactory::CreateStorage("MMapFileStorage"); + ASSERT_EQ(0, storage->init(stg_params)); + ASSERT_EQ(0, storage->open(dir_ + "TestKnnSearch.index", true)); + ASSERT_EQ(0, streamer->open(storage)); + + NumericalVector vec(kDim); + size_t cnt = 5000U; + auto ctx = streamer->create_context(); + ASSERT_TRUE(!!ctx); + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, kDim); + for (size_t i = 0; i < cnt; i++) { + for (size_t j = 0; j < kDim; ++j) { + vec[j] = static_cast(i); + } + ASSERT_EQ(0, streamer->add_impl(i, vec.data(), qmeta, ctx)); + } + + auto linearCtx = streamer->create_context(); + auto knnCtx = streamer->create_context(); + size_t topk = 100; + linearCtx->set_topk(topk); + knnCtx->set_topk(topk); + int totalHits = 0; + int totalCnts = 0; + int topk1Hits = 0; + for (size_t i = 0; i < cnt; i++) { + for (size_t j = 0; j < kDim; ++j) { + vec[j] = static_cast(i) + 0.1f; + } + ASSERT_EQ(0, streamer->search_impl(vec.data(), qmeta, knnCtx)); + ASSERT_EQ(0, streamer->search_bf_impl(vec.data(), qmeta, linearCtx)); + + auto &knnResult = knnCtx->result(); + ASSERT_EQ(topk, knnResult.size()); + topk1Hits += i == knnResult[0].key(); + + auto &linearResult = linearCtx->result(); + ASSERT_EQ(topk, linearResult.size()); + ASSERT_EQ(i, linearResult[0].key()); + + for (size_t k = 0; k < topk; ++k) { + totalCnts++; + for (size_t j = 0; j < topk; ++j) { + if (linearResult[j].key() == knnResult[k].key()) { + totalHits++; + break; + } + } + } + } + float recall = totalHits * 1.0f / totalCnts; + float topk1Recall = topk1Hits * 1.0f / cnt; + EXPECT_GT(recall, 0.90f); + EXPECT_GT(topk1Recall, 0.95f); +} + +TEST_F(VamanaStreamerTest, TestOpenClose) { + auto streamer = CreateVamanaStreamer(); + ASSERT_NE(nullptr, streamer); + + constexpr size_t dim_large = 128; + IndexMeta meta(IndexMeta::DataType::DT_FP32, dim_large); + meta.set_metric("SquaredEuclidean", 0, ailego::Params()); + + ailego::Params params; + params.set(PARAM_VAMANA_STREAMER_MAX_DEGREE, 32U); + params.set(PARAM_VAMANA_STREAMER_SEARCH_LIST_SIZE, 100U); + params.set(PARAM_VAMANA_STREAMER_ALPHA, 1.2f); + + streamer = IndexFactory::CreateStreamer("VamanaStreamer"); + ASSERT_NE(nullptr, streamer); + ASSERT_EQ(0, streamer->init(meta, params)); + + auto storage = IndexFactory::CreateStorage("MMapFileStorage"); + ASSERT_NE(nullptr, storage); + ailego::Params stg_params; + ASSERT_EQ(0, storage->init(stg_params)); + ASSERT_EQ(0, storage->open(dir_ + "TestOpenClose.index", true)); + ASSERT_EQ(0, streamer->open(storage)); + + size_t testCnt = 200; + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim_large); + auto ctx = streamer->create_context(); + ASSERT_TRUE(!!ctx); + for (size_t i = 0; i < testCnt; i++) { + std::vector vec(dim_large); + for (size_t d = 0; d < dim_large; ++d) { + vec[d] = static_cast(i); + } + ASSERT_EQ(0, streamer->add_impl(i, vec.data(), qmeta, ctx)); + } + + ASSERT_EQ(0, streamer->flush(0UL)); + ASSERT_EQ(0, streamer->close()); + + // Re-open and verify data + ASSERT_EQ(0, streamer->open(storage)); + auto provider = streamer->create_provider(); + auto iter = provider->create_iterator(); + ASSERT_TRUE(!!iter); + size_t total = 0; + while (iter->is_valid()) { + float *data = (float *)iter->data(); + for (size_t d = 0; d < dim_large; ++d) { + ASSERT_FLOAT_EQ(static_cast(iter->key()), data[d]); + } + total++; + iter->next(); + } + ASSERT_EQ(testCnt, total); +} + +TEST_F(VamanaStreamerTest, TestKnnMultiThread) { + constexpr size_t dim = 32; + IndexMeta meta(IndexMeta::DataType::DT_FP32, dim); + meta.set_metric("SquaredEuclidean", 0, ailego::Params()); + + ailego::Params params; + params.set(PARAM_VAMANA_STREAMER_MAX_DEGREE, 64U); + params.set(PARAM_VAMANA_STREAMER_SEARCH_LIST_SIZE, 500U); + params.set(PARAM_VAMANA_STREAMER_ALPHA, 1.2f); + params.set(PARAM_VAMANA_STREAMER_EF, 200U); + params.set(PARAM_VAMANA_STREAMER_BRUTE_FORCE_THRESHOLD, 1000U); + params.set(PARAM_VAMANA_STREAMER_MAX_INDEX_SIZE, 30U * 1024U * 1024U); + params.set(PARAM_VAMANA_STREAMER_GET_VECTOR_ENABLE, true); + + auto streamer = IndexFactory::CreateStreamer("VamanaStreamer"); + ASSERT_NE(nullptr, streamer); + ASSERT_EQ(0, streamer->init(meta, params)); + + auto storage = IndexFactory::CreateStorage("MMapFileStorage"); + ASSERT_NE(nullptr, storage); + ailego::Params stg_params; + ASSERT_EQ(0, storage->init(stg_params)); + ASSERT_EQ(0, storage->open(dir_ + "TestKnnMultiThread", true)); + ASSERT_EQ(0, streamer->open(storage)); + + auto addVector = [&streamer, dim](int baseKey, size_t addCnt) { + NumericalVector vec(dim); + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); + size_t succAdd = 0; + auto ctx = streamer->create_context(); + for (size_t i = 0; i < addCnt; i++) { + for (size_t j = 0; j < dim; ++j) { + vec[j] = static_cast(i + baseKey); + } + succAdd += !streamer->add_impl(baseKey + i, vec.data(), qmeta, ctx); + } + streamer->flush(0UL); + return succAdd; + }; + auto t1 = std::async(std::launch::async, addVector, 0, 1000); + auto t2 = std::async(std::launch::async, addVector, 1000, 1000); + auto t3 = std::async(std::launch::async, addVector, 2000, 1000); + ASSERT_EQ(1000U, t1.get()); + ASSERT_EQ(1000U, t2.get()); + ASSERT_EQ(1000U, t3.get()); + streamer->close(); + + // Verify data + ASSERT_EQ(0, streamer->open(storage)); + auto provider = streamer->create_provider(); + auto iter = provider->create_iterator(); + ASSERT_TRUE(!!iter); + size_t total = 0; + uint64_t minKey = 10000; + uint64_t maxKey = 0; + while (iter->is_valid()) { + float *data = (float *)iter->data(); + for (size_t d = 0; d < dim; ++d) { + ASSERT_FLOAT_EQ(static_cast(iter->key()), data[d]); + } + total++; + minKey = std::min(minKey, iter->key()); + maxKey = std::max(maxKey, iter->key()); + iter->next(); + } + ASSERT_EQ(3000, total); + ASSERT_EQ(0, minKey); + ASSERT_EQ(2999, maxKey); + + // Multi-thread search + size_t topk = 100; + size_t cnt = 3000; + auto knnSearch = [&]() { + NumericalVector vec(dim); + auto linearCtx = streamer->create_context(); + auto knnCtx = streamer->create_context(); + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); + linearCtx->set_topk(topk); + knnCtx->set_topk(topk); + size_t totalCnts = 0; + size_t totalHits = 0; + for (size_t i = 0; i < cnt; i += 1) { + for (size_t j = 0; j < dim; ++j) { + vec[j] = static_cast(i) + 0.1f; + } + ASSERT_EQ(0, streamer->search_impl(vec.data(), qmeta, knnCtx)); + ASSERT_EQ(0, streamer->search_bf_impl(vec.data(), qmeta, linearCtx)); + auto &knnResult = knnCtx->result(); + ASSERT_EQ(topk, knnResult.size()); + auto &linearResult = linearCtx->result(); + ASSERT_EQ(topk, linearResult.size()); + ASSERT_EQ(i, linearResult[0].key()); + for (size_t k = 0; k < topk; ++k) { + totalCnts++; + for (size_t j = 0; j < topk; ++j) { + if (linearResult[j].key() == knnResult[k].key()) { + totalHits++; + break; + } + } + } + } + ASSERT_TRUE((totalHits * 1.0f / totalCnts) > 0.80f); + }; + auto s1 = std::async(std::launch::async, knnSearch); + auto s2 = std::async(std::launch::async, knnSearch); + auto s3 = std::async(std::launch::async, knnSearch); + s1.wait(); + s2.wait(); + s3.wait(); +} + +TEST_F(VamanaStreamerTest, TestContiguousMemory) { + ailego::Params extra; + extra.set(PARAM_VAMANA_STREAMER_USE_CONTIGUOUS_MEMORY, true); + extra.set(PARAM_VAMANA_STREAMER_BRUTE_FORCE_THRESHOLD, 2000U); + auto streamer = CreateVamanaStreamer(extra); + ASSERT_NE(nullptr, streamer); + + auto storage = IndexFactory::CreateStorage("MMapFileStorage"); + ASSERT_NE(nullptr, storage); + ailego::Params stg_params; + ASSERT_EQ(0, storage->init(stg_params)); + ASSERT_EQ(0, storage->open(dir_ + "TestContiguous.index", true)); + + // First build with default mmap mode + { + auto builder_streamer = CreateVamanaStreamer(); + ASSERT_NE(nullptr, builder_streamer); + ASSERT_EQ(0, builder_streamer->open(storage)); + auto ctx = builder_streamer->create_context(); + ASSERT_TRUE(!!ctx); + + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, kDim); + NumericalVector vec(kDim); + size_t cnt = 3000UL; + for (size_t i = 0; i < cnt; i++) { + for (size_t j = 0; j < kDim; ++j) { + vec[j] = static_cast(i); + } + ASSERT_EQ(0, builder_streamer->add_impl(i, vec.data(), qmeta, ctx)); + } + ASSERT_EQ(0, builder_streamer->flush(0UL)); + ASSERT_EQ(0, builder_streamer->close()); + } + + // Re-open with contiguous memory mode for search + ASSERT_EQ(0, streamer->open(storage)); + + size_t cnt = 3000UL; + size_t topk = 50; + NumericalVector vec(kDim); + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, kDim); + auto linearCtx = streamer->create_context(); + auto knnCtx = streamer->create_context(); + linearCtx->set_topk(topk); + knnCtx->set_topk(topk); + int totalHits = 0; + int totalCnts = 0; + for (size_t i = 0; i < cnt; i++) { + for (size_t j = 0; j < kDim; ++j) { + vec[j] = static_cast(i) + 0.1f; + } + ASSERT_EQ(0, streamer->search_impl(vec.data(), qmeta, knnCtx)); + ASSERT_EQ(0, streamer->search_bf_impl(vec.data(), qmeta, linearCtx)); + auto &knnResult = knnCtx->result(); + ASSERT_EQ(topk, knnResult.size()); + auto &linearResult = linearCtx->result(); + ASSERT_EQ(topk, linearResult.size()); + ASSERT_EQ(i, linearResult[0].key()); + for (size_t k = 0; k < topk; ++k) { + totalCnts++; + for (size_t j = 0; j < topk; ++j) { + if (linearResult[j].key() == knnResult[k].key()) { + totalHits++; + break; + } + } + } + } + float recall = totalHits * 1.0f / totalCnts; + EXPECT_GT(recall, 0.90f); +} + +TEST_F(VamanaStreamerTest, TestContiguousMultiThreadSearch) { + constexpr size_t dim = 32; + IndexMeta meta(IndexMeta::DataType::DT_FP32, dim); + meta.set_metric("SquaredEuclidean", 0, ailego::Params()); + + // Build with mmap mode + auto storage = IndexFactory::CreateStorage("MMapFileStorage"); + ASSERT_NE(nullptr, storage); + ailego::Params stg_params; + ASSERT_EQ(0, storage->init(stg_params)); + ASSERT_EQ(0, storage->open(dir_ + "TestContiguousMT", true)); + + { + ailego::Params build_params; + build_params.set(PARAM_VAMANA_STREAMER_MAX_DEGREE, 64U); + build_params.set(PARAM_VAMANA_STREAMER_SEARCH_LIST_SIZE, 128U); + build_params.set(PARAM_VAMANA_STREAMER_ALPHA, 1.2f); + build_params.set(PARAM_VAMANA_STREAMER_EF, 64U); + build_params.set(PARAM_VAMANA_STREAMER_MAX_INDEX_SIZE, 30U * 1024U * 1024U); + build_params.set(PARAM_VAMANA_STREAMER_GET_VECTOR_ENABLE, true); + + auto builder = IndexFactory::CreateStreamer("VamanaStreamer"); + ASSERT_NE(nullptr, builder); + ASSERT_EQ(0, builder->init(meta, build_params)); + ASSERT_EQ(0, builder->open(storage)); + + auto ctx = builder->create_context(); + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); + NumericalVector vec(dim); + for (size_t i = 0; i < 3000; i++) { + for (size_t j = 0; j < dim; ++j) { + vec[j] = static_cast(i); + } + ASSERT_EQ(0, builder->add_impl(i, vec.data(), qmeta, ctx)); + } + ASSERT_EQ(0, builder->flush(0UL)); + ASSERT_EQ(0, builder->close()); + } + + // Re-open with contiguous memory + ailego::Params search_params; + search_params.set(PARAM_VAMANA_STREAMER_MAX_DEGREE, 64U); + search_params.set(PARAM_VAMANA_STREAMER_SEARCH_LIST_SIZE, 128U); + search_params.set(PARAM_VAMANA_STREAMER_ALPHA, 1.2f); + search_params.set(PARAM_VAMANA_STREAMER_EF, 64U); + search_params.set(PARAM_VAMANA_STREAMER_MAX_INDEX_SIZE, 30U * 1024U * 1024U); + search_params.set(PARAM_VAMANA_STREAMER_GET_VECTOR_ENABLE, true); + search_params.set(PARAM_VAMANA_STREAMER_USE_CONTIGUOUS_MEMORY, true); + + auto searcher = IndexFactory::CreateStreamer("VamanaStreamer"); + ASSERT_NE(nullptr, searcher); + ASSERT_EQ(0, searcher->init(meta, search_params)); + ASSERT_EQ(0, searcher->open(storage)); + + size_t topk = 50; + size_t cnt = 3000; + auto knnSearch = [&]() { + NumericalVector vec(dim); + auto linearCtx = searcher->create_context(); + auto knnCtx = searcher->create_context(); + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); + linearCtx->set_topk(topk); + knnCtx->set_topk(topk); + size_t totalCnts = 0; + size_t totalHits = 0; + for (size_t i = 0; i < cnt; i++) { + for (size_t j = 0; j < dim; ++j) { + vec[j] = static_cast(i) + 0.1f; + } + ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, knnCtx)); + ASSERT_EQ(0, searcher->search_bf_impl(vec.data(), qmeta, linearCtx)); + auto &knnResult = knnCtx->result(); + ASSERT_EQ(topk, knnResult.size()); + auto &linearResult = linearCtx->result(); + ASSERT_EQ(topk, linearResult.size()); + ASSERT_EQ(i, linearResult[0].key()); + for (size_t k = 0; k < topk; ++k) { + totalCnts++; + for (size_t j = 0; j < topk; ++j) { + if (linearResult[j].key() == knnResult[k].key()) { + totalHits++; + break; + } + } + } + } + ASSERT_TRUE((totalHits * 1.0f / totalCnts) > 0.80f); + }; + auto s1 = std::async(std::launch::async, knnSearch); + auto s2 = std::async(std::launch::async, knnSearch); + auto s3 = std::async(std::launch::async, knnSearch); + s1.wait(); + s2.wait(); + s3.wait(); +} + +TEST_F(VamanaStreamerTest, TestProvider) { + auto streamer = CreateVamanaStreamer(); + ASSERT_NE(nullptr, streamer); + + auto storage = IndexFactory::CreateStorage("MMapFileStorage"); + ASSERT_NE(nullptr, storage); + ailego::Params stg_params; + ASSERT_EQ(0, storage->init(stg_params)); + ASSERT_EQ(0, storage->open(dir_ + "TestProvider", true)); + ASSERT_EQ(0, streamer->open(storage)); + + size_t cnt = 500; + auto ctx = streamer->create_context(); + ASSERT_TRUE(!!ctx); + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, kDim); + NumericalVector vec(kDim); + for (size_t i = 0; i < cnt; i++) { + for (size_t j = 0; j < kDim; ++j) { + vec[j] = static_cast(i); + } + ASSERT_EQ(0, streamer->add_impl(i, vec.data(), qmeta, ctx)); + } + ASSERT_EQ(0, streamer->flush(0UL)); + + auto provider = streamer->create_provider(); + ASSERT_NE(nullptr, provider); + auto iter = provider->create_iterator(); + ASSERT_TRUE(!!iter); + size_t total = 0; + while (iter->is_valid()) { + ASSERT_NE(nullptr, iter->data()); + float *data = (float *)iter->data(); + for (size_t d = 0; d < kDim; ++d) { + ASSERT_FLOAT_EQ(static_cast(iter->key()), data[d]); + } + total++; + iter->next(); + } + ASSERT_EQ(cnt, total); +} + +TEST_F(VamanaStreamerTest, TestAddAndSearch) { + auto streamer = CreateVamanaStreamer(); + ASSERT_NE(nullptr, streamer); + + auto storage = IndexFactory::CreateStorage("MMapFileStorage"); + ASSERT_NE(nullptr, storage); + ailego::Params stg_params; + ASSERT_EQ(0, storage->init(stg_params)); + ASSERT_EQ(0, storage->open(dir_ + "TestAddAndSearch.index", true)); + ASSERT_EQ(0, streamer->open(storage)); + + NumericalVector vec(kDim); + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, kDim); + auto ctx = streamer->create_context(); + ASSERT_TRUE(!!ctx); + + // Add and search interleaved + for (size_t batch = 0; batch < 5; batch++) { + size_t base = batch * 200; + for (size_t i = 0; i < 200; i++) { + for (size_t j = 0; j < kDim; ++j) { + vec[j] = static_cast(base + i); + } + ASSERT_EQ(0, streamer->add_impl(base + i, vec.data(), qmeta, ctx)); + } + + // Search for recently added vectors + size_t current_cnt = (batch + 1) * 200; + size_t topk = std::min(current_cnt, (size_t)10); + auto searchCtx = streamer->create_context(); + searchCtx->set_topk(topk); + for (size_t j = 0; j < kDim; ++j) { + vec[j] = static_cast(base); + } + ASSERT_EQ(0, streamer->search_bf_impl(vec.data(), qmeta, searchCtx)); + auto &result = searchCtx->result(); + ASSERT_EQ(topk, result.size()); + ASSERT_EQ(base, result[0].key()); + } +} + +TEST_F(VamanaStreamerTest, TestKnnConcurrentAddAndSearch) { + constexpr size_t dim = 32; + IndexMeta meta(IndexMeta::DataType::DT_FP32, dim); + meta.set_metric("SquaredEuclidean", 0, ailego::Params()); + + ailego::Params params; + params.set(PARAM_VAMANA_STREAMER_MAX_DEGREE, 64U); + params.set(PARAM_VAMANA_STREAMER_SEARCH_LIST_SIZE, 128U); + params.set(PARAM_VAMANA_STREAMER_ALPHA, 1.2f); + params.set(PARAM_VAMANA_STREAMER_EF, 64U); + params.set(PARAM_VAMANA_STREAMER_BRUTE_FORCE_THRESHOLD, 500U); + params.set(PARAM_VAMANA_STREAMER_MAX_INDEX_SIZE, 30U * 1024U * 1024U); + params.set(PARAM_VAMANA_STREAMER_GET_VECTOR_ENABLE, true); + + auto streamer = IndexFactory::CreateStreamer("VamanaStreamer"); + ASSERT_NE(nullptr, streamer); + ASSERT_EQ(0, streamer->init(meta, params)); + + auto storage = IndexFactory::CreateStorage("MMapFileStorage"); + ASSERT_NE(nullptr, storage); + ailego::Params stg_params; + ASSERT_EQ(0, storage->init(stg_params)); + ASSERT_EQ(0, storage->open(dir_ + "TestConcurrentAddSearch", true)); + ASSERT_EQ(0, streamer->open(storage)); + + // First add some base data + { + auto ctx = streamer->create_context(); + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); + NumericalVector vec(dim); + for (size_t i = 0; i < 2000; i++) { + for (size_t j = 0; j < dim; ++j) { + vec[j] = static_cast(i); + } + ASSERT_EQ(0, streamer->add_impl(i, vec.data(), qmeta, ctx)); + } + } + + std::atomic stop_search{false}; + + // Concurrent add + auto addFuture = std::async(std::launch::async, [&]() { + auto ctx = streamer->create_context(); + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); + NumericalVector vec(dim); + for (size_t i = 2000; i < 3000; i++) { + for (size_t j = 0; j < dim; ++j) { + vec[j] = static_cast(i); + } + streamer->add_impl(i, vec.data(), qmeta, ctx); + } + stop_search.store(true); + }); + + // Concurrent search + auto searchFuture = std::async(std::launch::async, [&]() { + auto ctx = streamer->create_context(); + IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); + NumericalVector vec(dim); + ctx->set_topk(10); + while (!stop_search.load()) { + for (size_t j = 0; j < dim; ++j) { + vec[j] = 100.1f; + } + int ret = streamer->search_impl(vec.data(), qmeta, ctx); + ASSERT_EQ(0, ret); + auto &result = ctx->result(); + ASSERT_GT(result.size(), 0UL); + } + }); + + addFuture.wait(); + searchFuture.wait(); +} + +} // namespace core +} // namespace zvec + +#if defined(__GNUC__) || defined(__GNUG__) +#pragma GCC diagnostic pop +#endif diff --git a/tests/core/interface/CMakeLists.txt b/tests/core/interface/CMakeLists.txt index 62e7bc9..83bf8d9 100644 --- a/tests/core/interface/CMakeLists.txt +++ b/tests/core/interface/CMakeLists.txt @@ -8,7 +8,7 @@ foreach(CC_SRCS ${ALL_TEST_SRCS}) NAME ${CC_TARGET} STRICT LIBS zvec_ailego core_framework core_metric core_interface core_knn_flat core_utility core_quantizer sparsehash core_knn_hnsw core_mix_reducer - core_knn_flat_sparse core_knn_hnsw_sparse core_knn_ivf core_knn_hnsw_rabitq + core_knn_flat_sparse core_knn_hnsw_sparse core_knn_ivf core_knn_hnsw_rabitq core_knn_vamana SRCS ${CC_SRCS} INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm ) diff --git a/tests/core/interface/index_interface_test.cc b/tests/core/interface/index_interface_test.cc index 4016d31..0f590ea 100644 --- a/tests/core/interface/index_interface_test.cc +++ b/tests/core/interface/index_interface_test.cc @@ -153,6 +153,38 @@ TEST(IndexInterface, General) { .WithQuantizerParam(QuantizerParam(QuantizerType::kFP16)) .Build(), IVFQueryParamBuilder().with_topk(10).with_fetch_vector(true).build()); + + func(VamanaIndexParamBuilder() + .WithMetricType(MetricType::kInnerProduct) + .WithDataType(DataType::DT_FP32) + .WithDimension(kDimension) + .WithIsSparse(false) + .WithMaxDegree(32) + .WithSearchListSize(100) + .WithAlpha(1.2f) + .Build(), + VamanaQueryParamBuilder() + .with_topk(10) + .with_fetch_vector(true) + .with_ef_search(50) + .build()); + + // Vamana with topk > ef_search to exercise _get_coarse_search_topk branch + // that picks max(topk, ef_search). + func(VamanaIndexParamBuilder() + .WithMetricType(MetricType::kInnerProduct) + .WithDataType(DataType::DT_FP32) + .WithDimension(kDimension) + .WithIsSparse(false) + .WithMaxDegree(32) + .WithSearchListSize(100) + .WithAlpha(1.2f) + .Build(), + VamanaQueryParamBuilder() + .with_topk(100) + .with_fetch_vector(true) + .with_ef_search(10) + .build()); } TEST(IndexInterface, BufferGeneral) { @@ -651,6 +683,116 @@ TEST(IndexInterface, Serialize) { ASSERT_TRUE(IndexFactory::QueryParamSerializeToJson(*deserialized_param) == IndexFactory::QueryParamSerializeToJson(*param)); } + + { + std::cout << "\n\n----vamana index----" << std::endl; + auto param = VamanaIndexParamBuilder() + .WithMetricType(MetricType::kInnerProduct) + .WithDataType(DataType::DT_FP32) + .WithDimension(64) + .WithIsSparse(false) + .WithMaxDegree(32) + .WithSearchListSize(100) + .WithAlpha(1.2f) + .Build(); + + std::cout << "vamana index -- omit=true: " << param->SerializeToJson(true) + << std::endl; + std::cout << "vamana index -- omit=false: " << param->SerializeToJson() + << std::endl; + + auto deserialized_param = + IndexFactory::DeserializeIndexParamFromJson(param->SerializeToJson()); + ASSERT_NE(nullptr, deserialized_param.get()); + + std::cout << "serialize then de then se:" + << deserialized_param->SerializeToJson() << std::endl; + + ASSERT_TRUE(deserialized_param->SerializeToJson() == + param->SerializeToJson()); + ASSERT_TRUE(deserialized_param->SerializeToJson(true) == + param->SerializeToJson(true)); + } + + { + std::cout << "\n\n----hnsw index with use_contiguous_memory----" + << std::endl; + auto param = std::make_shared(); + param->metric_type = MetricType::kL2sq; + param->data_type = DataType::DT_FP32; + param->dimension = 64; + param->use_contiguous_memory = true; + + auto json_str = param->SerializeToJson(); + std::cout << "hnsw contiguous -- json: " << json_str << std::endl; + ASSERT_TRUE(json_str.find("use_contiguous_memory") != std::string::npos); + + auto deserialized_param = + IndexFactory::DeserializeIndexParamFromJson(json_str); + ASSERT_NE(nullptr, deserialized_param.get()); + auto hnsw_param = + std::dynamic_pointer_cast(deserialized_param); + ASSERT_NE(nullptr, hnsw_param.get()); + ASSERT_TRUE(hnsw_param->use_contiguous_memory); + + ASSERT_TRUE(deserialized_param->SerializeToJson() == json_str); + } + + { + std::cout << "\n\n----vamana index with use_contiguous_memory----" + << std::endl; + auto param = std::make_shared(); + param->metric_type = MetricType::kL2sq; + param->data_type = DataType::DT_FP32; + param->dimension = 64; + param->max_degree = 48; + param->search_list_size = 200; + param->alpha = 1.5f; + param->use_contiguous_memory = true; + + auto json_str = param->SerializeToJson(); + std::cout << "vamana contiguous -- json: " << json_str << std::endl; + ASSERT_TRUE(json_str.find("use_contiguous_memory") != std::string::npos); + + auto deserialized_param = + IndexFactory::DeserializeIndexParamFromJson(json_str); + ASSERT_NE(nullptr, deserialized_param.get()); + auto vamana_param = + std::dynamic_pointer_cast(deserialized_param); + ASSERT_NE(nullptr, vamana_param.get()); + ASSERT_TRUE(vamana_param->use_contiguous_memory); + ASSERT_EQ(48, vamana_param->max_degree); + ASSERT_EQ(200, vamana_param->search_list_size); + ASSERT_FLOAT_EQ(1.5f, vamana_param->alpha); + + ASSERT_TRUE(deserialized_param->SerializeToJson() == json_str); + } + + { + std::cout << "\n\n----vamana query----" << std::endl; + auto param = VamanaQueryParamBuilder() + .with_topk(10) + .with_fetch_vector(true) + .with_ef_search(50) + .build(); + std::cout << "vamana query -- omit=true: " + << IndexFactory::QueryParamSerializeToJson(*param, true) + << std::endl; + std::cout << "vamana query -- omit=false: " + << IndexFactory::QueryParamSerializeToJson(*param) << std::endl; + + auto deserialized_param = + IndexFactory::QueryParamDeserializeFromJson( + IndexFactory::QueryParamSerializeToJson(*param)); + ASSERT_NE(nullptr, deserialized_param.get()); + + std::cout << "serialize then de then se:" + << IndexFactory::QueryParamSerializeToJson(*deserialized_param) + << std::endl; + + ASSERT_TRUE(IndexFactory::QueryParamSerializeToJson(*deserialized_param) == + IndexFactory::QueryParamSerializeToJson(*param)); + } } TEST(IndexInterface, Failure) { @@ -888,6 +1030,110 @@ TEST(IndexInterface, Failure) { zvec::test_util::RemoveTestFiles("test2.index"); zvec::test_util::RemoveTestFiles("test3.index"); } + + // Test Vamana search with ef_search == 0 (invalid, ef_search must be > 0) + { + auto param = VamanaIndexParamBuilder() + .WithMetricType(MetricType::kInnerProduct) + .WithDataType(DataType::DT_FP32) + .WithDimension(64) + .WithIsSparse(false) + .WithMaxDegree(32) + .WithSearchListSize(100) + .WithAlpha(1.2f) + .Build(); + auto index = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, index); + + index->Open("test.index", {StorageOptions::StorageType::kMMAP, true}); + + std::vector vector(64, 1.0f); + VectorData vector_data{DenseVector{vector.data()}}; + ASSERT_EQ(0, index->Add(vector_data, 1)); + + VectorData query{DenseVector{vector.data()}}; + auto query_param = VamanaQueryParamBuilder() + .with_topk(10) + .with_fetch_vector(false) + .with_ef_search(0) + .build(); + SearchResult result; + int ret = index->Search(query, query_param, &result); + ASSERT_NE(0, ret); + + index->Close(); + zvec::test_util::RemoveTestFiles("test.index"); + } + + // Test Vamana search with ef_search > 2048 (invalid upper bound) + { + auto param = VamanaIndexParamBuilder() + .WithMetricType(MetricType::kInnerProduct) + .WithDataType(DataType::DT_FP32) + .WithDimension(64) + .WithIsSparse(false) + .WithMaxDegree(32) + .WithSearchListSize(100) + .WithAlpha(1.2f) + .Build(); + auto index = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, index); + + index->Open("test.index", {StorageOptions::StorageType::kMMAP, true}); + + std::vector vector(64, 1.0f); + VectorData vector_data{DenseVector{vector.data()}}; + ASSERT_EQ(0, index->Add(vector_data, 1)); + + VectorData query{DenseVector{vector.data()}}; + auto query_param = VamanaQueryParamBuilder() + .with_topk(10) + .with_fetch_vector(false) + .with_ef_search(4096) + .build(); + SearchResult result; + int ret = index->Search(query, query_param, &result); + ASSERT_NE(0, ret); + + index->Close(); + zvec::test_util::RemoveTestFiles("test.index"); + } + + // Test Vamana search with wrong query param type (HNSWQueryParam instead of + // VamanaQueryParam) + { + auto param = VamanaIndexParamBuilder() + .WithMetricType(MetricType::kInnerProduct) + .WithDataType(DataType::DT_FP32) + .WithDimension(64) + .WithIsSparse(false) + .WithMaxDegree(32) + .WithSearchListSize(100) + .WithAlpha(1.2f) + .Build(); + auto index = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, index); + + index->Open("test.index", {StorageOptions::StorageType::kMMAP, true}); + + std::vector vector(64, 1.0f); + VectorData vector_data{DenseVector{vector.data()}}; + ASSERT_EQ(0, index->Add(vector_data, 1)); + + VectorData query{DenseVector{vector.data()}}; + // Intentionally pass an HNSWQueryParam to a Vamana index + auto wrong_query_param = HNSWQueryParamBuilder() + .with_topk(10) + .with_fetch_vector(false) + .with_ef_search(50) + .build(); + SearchResult result; + int ret = index->Search(query, wrong_query_param, &result); + ASSERT_NE(0, ret); + + index->Close(); + zvec::test_util::RemoveTestFiles("test.index"); + } } TEST(IndexInterface, SerializeFailure) { @@ -1172,6 +1418,22 @@ TEST(IndexInterface, Score) { .build(), MetricType::kInnerProduct); + dense_func(VamanaIndexParamBuilder() + .WithMetricType(MetricType::kInnerProduct) + .WithDataType(DataType::DT_FP32) + .WithDimension(kDimension) + .WithIsSparse(false) + .WithMaxDegree(32) + .WithSearchListSize(100) + .WithAlpha(1.2f) + .Build(), + VamanaQueryParamBuilder() + .with_topk(kTopk) + .with_fetch_vector(true) + .with_ef_search(50) + .build(), + MetricType::kInnerProduct); + LOG_INFO("Test DenseVector, MetricType::kInnerProduct, QuantizerType::kFP16"); dense_func( FlatIndexParamBuilder() @@ -1455,6 +1717,135 @@ TEST(IndexInterface, HNSWRabitqGeneral) { } #endif +// Verify that enabling use_contiguous_memory on HNSW / Vamana index params at +// the interface layer is correctly propagated to the underlying streamer and +// yields a working build -> close -> reopen-for-search pipeline. This guards +// the interface -> streamer param binding introduced for contiguous memory +// mode. +TEST(IndexInterface, ContiguousMemoryEndToEnd) { + constexpr uint32_t kDimension = 32; + constexpr uint32_t kNumDocs = 500; + constexpr int kTopk = 10; + const std::string index_name{"test_contiguous.index"}; + + // build_then_search builds an index from scratch (with use_contiguous_memory + // possibly enabled), closes it, then reopens with the same params and runs a + // search for each inserted vector, asserting top-1 is itself. + auto build_then_search = [&](const BaseIndexParam::Pointer ¶m, + const BaseIndexQueryParam::Pointer &query_param) { + zvec::test_util::RemoveTestFiles(index_name); + + // Phase 1: build & persist. + { + auto index = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, index); + ASSERT_EQ(0, index->Open(index_name, + {StorageOptions::StorageType::kMMAP, true})); + + std::vector vec(kDimension); + for (uint32_t i = 0; i < kNumDocs; ++i) { + for (uint32_t d = 0; d < kDimension; ++d) { + vec[d] = static_cast(i); + } + VectorData data{DenseVector{vec.data()}}; + ASSERT_EQ(0, index->Add(data, i)); + } + ASSERT_EQ(0, index->Train()); + ASSERT_EQ(0, index->Close()); + } + + // Phase 2: reopen with same params (contiguous memory takes effect here) + // and search. + { + auto index = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, index); + ASSERT_EQ(0, index->Open(index_name, + {StorageOptions::StorageType::kMMAP, false})); + + std::vector q(kDimension); + for (uint32_t i = 0; i < kNumDocs; i += 50) { + for (uint32_t d = 0; d < kDimension; ++d) { + q[d] = static_cast(i); + } + VectorData query{DenseVector{q.data()}}; + SearchResult result; + ASSERT_EQ(0, index->Search(query, query_param, &result)); + ASSERT_GT(result.doc_list_.size(), 0UL); + ASSERT_EQ(i, result.doc_list_[0].key()); + } + ASSERT_EQ(0, index->Close()); + } + + zvec::test_util::RemoveTestFiles(index_name); + }; + + // HNSW + use_contiguous_memory=true + build_then_search(HNSWIndexParamBuilder() + .WithMetricType(MetricType::kL2sq) + .WithDataType(DataType::DT_FP32) + .WithDimension(kDimension) + .WithIsSparse(false) + .WithM(16) + .WithEFConstruction(64) + .WithUseContiguousMemory(true) + .Build(), + HNSWQueryParamBuilder() + .with_topk(kTopk) + .with_fetch_vector(false) + .with_ef_search(64) + .build()); + + // HNSW + use_contiguous_memory=false (baseline, same harness) + build_then_search(HNSWIndexParamBuilder() + .WithMetricType(MetricType::kL2sq) + .WithDataType(DataType::DT_FP32) + .WithDimension(kDimension) + .WithIsSparse(false) + .WithM(16) + .WithEFConstruction(64) + .WithUseContiguousMemory(false) + .Build(), + HNSWQueryParamBuilder() + .with_topk(kTopk) + .with_fetch_vector(false) + .with_ef_search(64) + .build()); + + // Vamana + use_contiguous_memory=true + build_then_search(VamanaIndexParamBuilder() + .WithMetricType(MetricType::kL2sq) + .WithDataType(DataType::DT_FP32) + .WithDimension(kDimension) + .WithIsSparse(false) + .WithMaxDegree(32) + .WithSearchListSize(100) + .WithAlpha(1.2f) + .WithUseContiguousMemory(true) + .Build(), + VamanaQueryParamBuilder() + .with_topk(kTopk) + .with_fetch_vector(false) + .with_ef_search(64) + .build()); + + // Vamana + use_contiguous_memory=false (baseline, same harness) + build_then_search(VamanaIndexParamBuilder() + .WithMetricType(MetricType::kL2sq) + .WithDataType(DataType::DT_FP32) + .WithDimension(kDimension) + .WithIsSparse(false) + .WithMaxDegree(32) + .WithSearchListSize(100) + .WithAlpha(1.2f) + .WithUseContiguousMemory(false) + .Build(), + VamanaQueryParamBuilder() + .with_topk(kTopk) + .with_fetch_vector(false) + .with_ef_search(64) + .build()); +} + #if defined(__GNUC__) || defined(__GNUG__) #pragma GCC diagnostic pop #endif \ No newline at end of file diff --git a/tools/core/CMakeLists.txt b/tools/core/CMakeLists.txt index c36b264..cd33267 100644 --- a/tools/core/CMakeLists.txt +++ b/tools/core/CMakeLists.txt @@ -14,7 +14,7 @@ cc_binary( STRICT PACKED SRCS local_builder.cc INCS ${PROJECT_ROOT_DIR}/src/core/ - LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_cluster core_knn_ivf core_interface + LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_vamana core_knn_cluster core_knn_ivf core_interface ) cc_binary( @@ -22,7 +22,7 @@ cc_binary( STRICT PACKED SRCS recall.cc INCS ${PROJECT_ROOT_DIR}/src/core/ - LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_cluster core_knn_ivf roaring core_interface + LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_vamana core_knn_cluster core_knn_ivf roaring core_interface ) cc_binary( @@ -30,7 +30,7 @@ cc_binary( STRICT PACKED SRCS bench.cc INCS ${PROJECT_ROOT_DIR}/src/core/ - LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_cluster core_knn_ivf roaring core_interface + LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_vamana core_knn_cluster core_knn_ivf roaring core_interface ) @@ -39,7 +39,7 @@ cc_binary( STRICT PACKED SRCS recall_original.cc flow.cc INCS ${PROJECT_ROOT_DIR}/src/core/ - LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_cluster core_knn_ivf roaring core_interface + LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_vamana core_knn_cluster core_knn_ivf roaring core_interface ) cc_binary( @@ -47,7 +47,7 @@ cc_binary( STRICT PACKED SRCS bench_original.cc flow.cc INCS ${PROJECT_ROOT_DIR}/src/core/ - LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_cluster core_knn_ivf roaring core_interface + LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_vamana core_knn_cluster core_knn_ivf roaring core_interface ) cc_binary( @@ -55,5 +55,5 @@ cc_binary( STRICT PACKED SRCS local_builder_original.cc INCS ${PROJECT_ROOT_DIR}/src/core/ - LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_cluster core_knn_ivf core_interface + LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_vamana core_knn_cluster core_knn_ivf core_interface )