feat: add diskann index (#369)
This commit is contained in:
@@ -1,9 +1,60 @@
|
||||
import pytest
|
||||
import logging
|
||||
import platform
|
||||
|
||||
DISKANN_SUPPORTED = platform.system() == "Linux" and platform.machine() in (
|
||||
"x86_64",
|
||||
"AMD64",
|
||||
"i686",
|
||||
"i386",
|
||||
)
|
||||
|
||||
from typing import Any, Generator
|
||||
from zvec.typing import DataType, StatusCode, MetricType, QuantizeType
|
||||
import zvec
|
||||
|
||||
|
||||
# Cache the DiskAnn plugin preload status so we pay the load cost once per
|
||||
# test session. The plugin normally auto-loads on first DiskAnn use, but we
|
||||
# preload it explicitly here so a missing libaio / misplaced plugin .so
|
||||
# surfaces as a clear pytest skip instead of a confusing
|
||||
# "Create vector column indexer failed" deep inside the collection code path.
|
||||
_DISKANN_PRELOAD_REASON: str | None = None
|
||||
_DISKANN_PRELOAD_DONE: bool = False
|
||||
|
||||
|
||||
def _ensure_diskann_runtime_or_reason() -> str | None:
|
||||
"""Preload the DiskAnn plugin and return None on success or a human-readable
|
||||
skip reason on failure. Idempotent across calls."""
|
||||
global _DISKANN_PRELOAD_DONE, _DISKANN_PRELOAD_REASON
|
||||
if _DISKANN_PRELOAD_DONE:
|
||||
return _DISKANN_PRELOAD_REASON
|
||||
_DISKANN_PRELOAD_DONE = True
|
||||
|
||||
if not DISKANN_SUPPORTED:
|
||||
_DISKANN_PRELOAD_REASON = "DiskAnn only supported on Linux x86_64"
|
||||
return _DISKANN_PRELOAD_REASON
|
||||
|
||||
if not zvec.is_libaio_available():
|
||||
_DISKANN_PRELOAD_REASON = (
|
||||
"libaio is not available on this host; DiskAnn cannot run. "
|
||||
"Install libaio1 (or libaio1t64 on Ubuntu 24.04+) and retry."
|
||||
)
|
||||
return _DISKANN_PRELOAD_REASON
|
||||
|
||||
status = zvec.load_diskann_plugin()
|
||||
if status != zvec.DISKANN_PLUGIN_OK:
|
||||
_DISKANN_PRELOAD_REASON = (
|
||||
f"Failed to load DiskAnn plugin (status={status}); "
|
||||
"check that libzvec_diskann_plugin.so is installed alongside "
|
||||
"_zvec.so in the Python site-packages directory."
|
||||
)
|
||||
return _DISKANN_PRELOAD_REASON
|
||||
|
||||
_DISKANN_PRELOAD_REASON = None
|
||||
return None
|
||||
|
||||
|
||||
from zvec import (
|
||||
CollectionOption,
|
||||
InvertIndexParam,
|
||||
@@ -97,6 +148,15 @@ def full_schema_new(request) -> CollectionSchema:
|
||||
else:
|
||||
nullable, has_index, vector_index = True, False, HnswIndexParam()
|
||||
|
||||
# Skip DiskAnn tests on unsupported platforms or when the runtime cannot
|
||||
# be brought up (missing libaio, plugin .so not installed, etc.).
|
||||
from zvec.model.param import DiskAnnIndexParam
|
||||
|
||||
if isinstance(vector_index, DiskAnnIndexParam):
|
||||
skip_reason = _ensure_diskann_runtime_or_reason()
|
||||
if skip_reason is not None:
|
||||
pytest.skip(skip_reason)
|
||||
|
||||
scalar_index_param = None
|
||||
vector_index_param = None
|
||||
if has_index:
|
||||
|
||||
@@ -22,8 +22,10 @@ from zvec.model.param import (
|
||||
HnswIndexParam,
|
||||
FlatIndexParam,
|
||||
IVFIndexParam,
|
||||
DiskAnnIndexParam,
|
||||
HnswQueryParam,
|
||||
IVFQueryParam,
|
||||
DiskAnnQueryParam,
|
||||
)
|
||||
|
||||
from zvec.model.schema import FieldSchema, VectorSchema
|
||||
@@ -179,10 +181,24 @@ def get_ground_truth_map(collection, test_docs, query_vectors_map, metric_type,
|
||||
for field_name, query_vectors in query_vectors_map.items():
|
||||
ground_truth_map[field_name] = {}
|
||||
|
||||
# Support per-field metric type: metric_type can be a dict mapping
|
||||
# field_name -> MetricType, or a single MetricType applied to all fields.
|
||||
if isinstance(metric_type, dict):
|
||||
field_metric = metric_type.get(field_name, MetricType.IP)
|
||||
else:
|
||||
field_metric = metric_type
|
||||
|
||||
for i, query_vector in enumerate(query_vectors):
|
||||
# Get the ground truth for this query
|
||||
relevant_doc_ids_scores = get_ground_truth_for_vector_query(
|
||||
collection, query_vector, field_name, test_docs, i, metric_type, k, True
|
||||
collection,
|
||||
query_vector,
|
||||
field_name,
|
||||
test_docs,
|
||||
i,
|
||||
field_metric,
|
||||
k,
|
||||
True,
|
||||
)
|
||||
ground_truth_map[field_name][i] = relevant_doc_ids_scores
|
||||
|
||||
@@ -292,6 +308,7 @@ class TestRecall:
|
||||
[
|
||||
(True, True, HnswIndexParam()),
|
||||
(False, True, IVFIndexParam()),
|
||||
(False, True, DiskAnnIndexParam()),
|
||||
(False, True, FlatIndexParam()), # ——ok
|
||||
(
|
||||
True,
|
||||
@@ -371,6 +388,19 @@ class TestRecall:
|
||||
use_soar=False,
|
||||
),
|
||||
),
|
||||
(
|
||||
True,
|
||||
True,
|
||||
DiskAnnIndexParam(
|
||||
metric_type=MetricType.IP,
|
||||
max_degree=32,
|
||||
),
|
||||
),
|
||||
(
|
||||
True,
|
||||
True,
|
||||
DiskAnnIndexParam(metric_type=MetricType.L2, max_degree=32),
|
||||
),
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
@@ -388,10 +418,16 @@ class TestRecall:
|
||||
):
|
||||
full_schema_params = request.getfixturevalue("full_schema_new")
|
||||
|
||||
# Build per-field metric type map so ground truth uses each field's
|
||||
# actual index metric (fields may fall back to HnswIndexParam/IP).
|
||||
field_metric_map = {}
|
||||
for vector_para in full_schema_params.vectors:
|
||||
if vector_para.name == "vector_fp32_field":
|
||||
metric_type = vector_para.index_param.metric_type
|
||||
break
|
||||
if vector_para.index_param is not None:
|
||||
field_metric_map[vector_para.name] = vector_para.index_param.metric_type
|
||||
else:
|
||||
field_metric_map[vector_para.name] = MetricType.IP
|
||||
|
||||
metric_type = field_metric_map.get("vector_fp32_field", MetricType.IP)
|
||||
|
||||
multiple_docs = [
|
||||
generate_doc_recall(i, full_collection_new.schema) for i in range(doc_num)
|
||||
@@ -438,9 +474,13 @@ class TestRecall:
|
||||
multiple_docs[i].vectors[field_name] for i in range(query_num)
|
||||
]
|
||||
|
||||
# Get ground truth mapping
|
||||
# Get ground truth mapping (pass per-field metric map)
|
||||
ground_truth_map = get_ground_truth_map(
|
||||
full_collection_new, multiple_docs, query_vectors_map, metric_type, top_k
|
||||
full_collection_new,
|
||||
multiple_docs,
|
||||
query_vectors_map,
|
||||
field_metric_map,
|
||||
top_k,
|
||||
)
|
||||
|
||||
# Validate ground truth mapping structure
|
||||
@@ -479,8 +519,8 @@ class TestRecall:
|
||||
|
||||
print("(recall_at_k_stats:\n")
|
||||
print(recall_at_k_stats)
|
||||
print("metric_type:")
|
||||
print(metric_type)
|
||||
print("field_metric_map:")
|
||||
print(field_metric_map)
|
||||
# Print Recall@K statistics
|
||||
print(f"Recall@{top_k} using Ground Truth:")
|
||||
for field_name, stats in recall_at_k_stats.items():
|
||||
@@ -552,7 +592,21 @@ class TestRecall:
|
||||
use_soar=True,
|
||||
),
|
||||
),
|
||||
# (True, True, IVFIndexParam(metric_type=MetricType.COSINE, n_list=150, n_iters=15, use_soar=False, )),
|
||||
(
|
||||
True,
|
||||
True,
|
||||
DiskAnnIndexParam(metric_type=MetricType.IP, max_degree=32),
|
||||
),
|
||||
(
|
||||
True,
|
||||
True,
|
||||
DiskAnnIndexParam(metric_type=MetricType.L2, max_degree=32),
|
||||
),
|
||||
(
|
||||
True,
|
||||
True,
|
||||
DiskAnnIndexParam(metric_type=MetricType.COSINE, max_degree=32),
|
||||
),
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
@@ -571,10 +625,16 @@ class TestRecall:
|
||||
):
|
||||
full_schema_params = request.getfixturevalue("full_schema_new")
|
||||
|
||||
# Build per-field metric type map so ground truth uses each field's
|
||||
# actual index metric (fields may fall back to HnswIndexParam/IP).
|
||||
field_metric_map = {}
|
||||
for vector_para in full_schema_params.vectors:
|
||||
if vector_para.name == "vector_fp32_field":
|
||||
metric_type = vector_para.index_param.metric_type
|
||||
break
|
||||
if vector_para.index_param is not None:
|
||||
field_metric_map[vector_para.name] = vector_para.index_param.metric_type
|
||||
else:
|
||||
field_metric_map[vector_para.name] = MetricType.IP
|
||||
|
||||
metric_type = field_metric_map.get("vector_fp32_field", MetricType.IP)
|
||||
|
||||
multiple_docs = [
|
||||
generate_doc_recall(i, full_collection_new.schema) for i in range(doc_num)
|
||||
@@ -621,9 +681,13 @@ class TestRecall:
|
||||
multiple_docs[i].vectors[field_name] for i in range(query_num)
|
||||
]
|
||||
|
||||
# Get ground truth mapping
|
||||
# Get ground truth mapping (pass per-field metric map)
|
||||
ground_truth_map = get_ground_truth_map(
|
||||
full_collection_new, multiple_docs, query_vectors_map, metric_type, top_k
|
||||
full_collection_new,
|
||||
multiple_docs,
|
||||
query_vectors_map,
|
||||
field_metric_map,
|
||||
top_k,
|
||||
)
|
||||
|
||||
# Validate ground truth mapping structure
|
||||
@@ -662,8 +726,8 @@ class TestRecall:
|
||||
|
||||
print("(recall_at_k_stats:\n")
|
||||
print(recall_at_k_stats)
|
||||
print("metric_type:")
|
||||
print(metric_type)
|
||||
print("field_metric_map:")
|
||||
print(field_metric_map)
|
||||
# Print Recall@K statistics
|
||||
print(f"Recall@{top_k} using Ground Truth:")
|
||||
for field_name, stats in recall_at_k_stats.items():
|
||||
|
||||
@@ -0,0 +1,600 @@
|
||||
# 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.
|
||||
"""End-to-end collection tests for the DiskAnn index.
|
||||
|
||||
Mirrors ``test_collection_hnsw_rabitq.py`` but targets the DiskAnn plugin.
|
||||
|
||||
Two platform-level prerequisites are enforced at module import time:
|
||||
|
||||
1. DiskAnn is currently built only for Linux x86_64 — other platforms are
|
||||
skipped wholesale.
|
||||
2. The DiskAnn backend lives in a *runtime-loaded* plugin
|
||||
(``libzvec_diskann_plugin.so``). It must be loaded with ``RTLD_GLOBAL |
|
||||
RTLD_NOW`` BEFORE ``import zvec`` so that the plugin's ``IndexFactory``
|
||||
singleton is unified with the one inside ``_zvec.so``. After ``import
|
||||
zvec`` we must also call ``zvec.load_diskann_plugin()`` exactly once.
|
||||
|
||||
If either prerequisite fails the whole module is skipped so the rest of the
|
||||
test-suite is not affected.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import os
|
||||
import platform
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Platform gating (must happen BEFORE we touch zvec).
|
||||
# --------------------------------------------------------------------------- #
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not (sys.platform == "linux" and platform.machine() in ("x86_64", "AMD64")),
|
||||
reason="DiskAnn plugin is only supported on Linux x86_64",
|
||||
)
|
||||
|
||||
# Promote all symbols in subsequently-loaded DSOs to the global namespace and
|
||||
# resolve relocations eagerly. This is REQUIRED so the DiskAnn plugin can see
|
||||
# the ``IndexFactory`` singleton that lives in ``_zvec.so`` and vice versa.
|
||||
# See: DiskAnn RTLD_GLOBAL + RTLD_NOW Requirement.
|
||||
if sys.platform == "linux":
|
||||
sys.setdlopenflags(sys.getdlopenflags() | os.RTLD_GLOBAL | os.RTLD_NOW)
|
||||
|
||||
import zvec # noqa: E402
|
||||
|
||||
from zvec import ( # noqa: E402
|
||||
Collection,
|
||||
CollectionOption,
|
||||
DataType,
|
||||
DiskAnnIndexParam,
|
||||
DiskAnnQueryParam,
|
||||
Doc,
|
||||
FieldSchema,
|
||||
MetricType,
|
||||
Query,
|
||||
VectorSchema,
|
||||
)
|
||||
from zvec.typing import QuantizeType # noqa: E402
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def diskann_collection_schema():
|
||||
"""Create a collection schema with a DiskAnn index."""
|
||||
return zvec.CollectionSchema(
|
||||
name="test_diskann_collection",
|
||||
fields=[
|
||||
FieldSchema("id", DataType.INT64, nullable=False),
|
||||
FieldSchema("name", DataType.STRING, nullable=False),
|
||||
],
|
||||
vectors=[
|
||||
VectorSchema(
|
||||
"embedding",
|
||||
DataType.VECTOR_FP32,
|
||||
dimension=128,
|
||||
index_param=DiskAnnIndexParam(
|
||||
metric_type=MetricType.L2,
|
||||
max_degree=64,
|
||||
list_size=100,
|
||||
pq_chunk_num=0,
|
||||
quantize_type=QuantizeType.UNDEFINED,
|
||||
),
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def collection_option():
|
||||
"""Create collection options."""
|
||||
return CollectionOption(read_only=False, enable_mmap=True)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def single_doc():
|
||||
"""Create a single document for testing."""
|
||||
return Doc(
|
||||
id="0",
|
||||
fields={"id": 0, "name": "test_doc_0"},
|
||||
vectors={"embedding": [0.1 + i * 0.01 for i in range(128)]},
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def multiple_docs():
|
||||
"""Create multiple documents for testing."""
|
||||
return [
|
||||
Doc(
|
||||
id=f"{i}",
|
||||
fields={"id": i, "name": f"test_doc_{i}"},
|
||||
vectors={"embedding": [i * 0.1 + j * 0.01 for j in range(128)]},
|
||||
)
|
||||
for i in range(1, 101)
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def diskann_collection(
|
||||
tmp_path_factory, diskann_collection_schema, collection_option
|
||||
) -> Collection:
|
||||
"""
|
||||
Function-scoped fixture: creates and opens a collection with DiskAnn index.
|
||||
"""
|
||||
temp_dir = tmp_path_factory.mktemp("zvec_diskann")
|
||||
collection_path = temp_dir / "test_diskann_collection"
|
||||
|
||||
coll = zvec.create_and_open(
|
||||
path=str(collection_path),
|
||||
schema=diskann_collection_schema,
|
||||
option=collection_option,
|
||||
)
|
||||
|
||||
assert coll is not None, "Failed to create and open DiskAnn collection"
|
||||
assert coll.path == str(collection_path)
|
||||
assert coll.schema.name == diskann_collection_schema.name
|
||||
|
||||
try:
|
||||
yield coll
|
||||
finally:
|
||||
if hasattr(coll, "destroy") and coll is not None:
|
||||
try:
|
||||
coll.destroy()
|
||||
except Exception as e:
|
||||
print(f"Warning: failed to destroy collection: {e}")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def collection_with_single_doc(
|
||||
diskann_collection: Collection, single_doc: Doc
|
||||
) -> Collection:
|
||||
"""Setup: insert single doc into collection."""
|
||||
assert diskann_collection.stats.doc_count == 0
|
||||
result = diskann_collection.insert(single_doc)
|
||||
assert bool(result)
|
||||
assert result.ok()
|
||||
assert diskann_collection.stats.doc_count == 1
|
||||
|
||||
yield diskann_collection
|
||||
|
||||
# Teardown: delete single doc
|
||||
diskann_collection.delete(single_doc.id)
|
||||
assert diskann_collection.stats.doc_count == 0
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def collection_with_multiple_docs(
|
||||
diskann_collection: Collection, multiple_docs: list[Doc]
|
||||
) -> Collection:
|
||||
"""Setup: insert multiple docs into collection."""
|
||||
assert diskann_collection.stats.doc_count == 0
|
||||
result = diskann_collection.insert(multiple_docs)
|
||||
assert len(result) == len(multiple_docs)
|
||||
for item in result:
|
||||
assert item.ok()
|
||||
assert diskann_collection.stats.doc_count == len(multiple_docs)
|
||||
|
||||
yield diskann_collection
|
||||
|
||||
# Teardown: delete multiple docs
|
||||
diskann_collection.delete([doc.id for doc in multiple_docs])
|
||||
|
||||
|
||||
# ==================== Tests ====================
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("diskann_collection")
|
||||
class TestDiskAnnCollectionCreation:
|
||||
"""Test DiskAnn collection creation and schema validation."""
|
||||
|
||||
def test_collection_creation(
|
||||
self, diskann_collection: Collection, diskann_collection_schema
|
||||
):
|
||||
"""Test that collection is created with correct schema."""
|
||||
assert diskann_collection is not None
|
||||
assert diskann_collection.schema.name == diskann_collection_schema.name
|
||||
assert len(diskann_collection.schema.fields) == len(
|
||||
diskann_collection_schema.fields
|
||||
)
|
||||
assert len(diskann_collection.schema.vectors) == len(
|
||||
diskann_collection_schema.vectors
|
||||
)
|
||||
|
||||
def test_vector_schema_validation(self, diskann_collection: Collection):
|
||||
"""Test that vector schema has correct DiskAnn configuration."""
|
||||
vector_schema = diskann_collection.schema.vector("embedding")
|
||||
assert vector_schema is not None
|
||||
assert vector_schema.name == "embedding"
|
||||
assert vector_schema.data_type == DataType.VECTOR_FP32
|
||||
assert vector_schema.dimension == 128
|
||||
|
||||
index_param = vector_schema.index_param
|
||||
assert index_param is not None
|
||||
assert index_param.metric_type == MetricType.L2
|
||||
assert index_param.max_degree == 64
|
||||
assert index_param.list_size == 100
|
||||
assert index_param.pq_chunk_num == 0
|
||||
|
||||
def test_collection_stats(self, diskann_collection: Collection):
|
||||
"""Test initial collection statistics."""
|
||||
stats = diskann_collection.stats
|
||||
assert stats is not None
|
||||
assert stats.doc_count == 0
|
||||
assert len(stats.index_completeness) == 1
|
||||
assert stats.index_completeness["embedding"] == 1
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("diskann_collection")
|
||||
class TestDiskAnnCollectionInsert:
|
||||
"""Test document insertion into DiskAnn collection."""
|
||||
|
||||
def test_insert_single_doc(self, diskann_collection: Collection, single_doc: Doc):
|
||||
"""Test inserting a single document."""
|
||||
result = diskann_collection.insert(single_doc)
|
||||
assert bool(result)
|
||||
assert result.ok()
|
||||
|
||||
stats = diskann_collection.stats
|
||||
assert stats is not None
|
||||
assert stats.doc_count == 1
|
||||
|
||||
def test_insert_multiple_docs(
|
||||
self, diskann_collection: Collection, multiple_docs: list[Doc]
|
||||
):
|
||||
"""Test inserting multiple documents."""
|
||||
result = diskann_collection.insert(multiple_docs)
|
||||
assert len(result) == len(multiple_docs)
|
||||
for item in result:
|
||||
assert item.ok()
|
||||
|
||||
stats = diskann_collection.stats
|
||||
assert stats is not None
|
||||
assert stats.doc_count == len(multiple_docs)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("diskann_collection")
|
||||
class TestDiskAnnCollectionFetch:
|
||||
"""Test document fetching from DiskAnn collection."""
|
||||
|
||||
def test_fetch_single_doc(
|
||||
self, collection_with_single_doc: Collection, single_doc: Doc
|
||||
):
|
||||
"""Test fetching a single document by ID."""
|
||||
result = collection_with_single_doc.fetch(ids=[single_doc.id])
|
||||
assert bool(result)
|
||||
assert single_doc.id in result.keys()
|
||||
|
||||
doc = result[single_doc.id]
|
||||
assert doc is not None
|
||||
assert doc.id == single_doc.id
|
||||
assert doc.field("id") == single_doc.field("id")
|
||||
assert doc.field("name") == single_doc.field("name")
|
||||
|
||||
def test_fetch_multiple_docs(
|
||||
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
|
||||
):
|
||||
"""Test fetching multiple documents by IDs."""
|
||||
ids = [doc.id for doc in multiple_docs[:10]]
|
||||
result = collection_with_multiple_docs.fetch(ids=ids)
|
||||
assert bool(result)
|
||||
assert len(result) == len(ids)
|
||||
|
||||
for doc_id in ids:
|
||||
assert doc_id in result
|
||||
doc = result[doc_id]
|
||||
assert doc is not None
|
||||
assert doc.id == doc_id
|
||||
|
||||
def test_fetch_nonexistent_doc(self, collection_with_single_doc: Collection):
|
||||
"""Test fetching a non-existent document."""
|
||||
result = collection_with_single_doc.fetch(ids=["nonexistent_id"])
|
||||
assert len(result) == 0
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("diskann_collection")
|
||||
class TestDiskAnnCollectionQuery:
|
||||
"""Test vector search queries on DiskAnn collection."""
|
||||
|
||||
def test_query_by_vector(
|
||||
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
|
||||
):
|
||||
"""Test querying by vector with DiskAnn index."""
|
||||
query_vector = multiple_docs[0].vector("embedding")
|
||||
query = Query(
|
||||
field_name="embedding",
|
||||
vector=query_vector,
|
||||
param=DiskAnnQueryParam(list_size=100),
|
||||
)
|
||||
|
||||
result = collection_with_multiple_docs.query(queries=query, topk=10)
|
||||
assert len(result) > 0
|
||||
assert len(result) <= 10
|
||||
|
||||
# First result should be the query document itself (or very close)
|
||||
first_doc = result[0]
|
||||
assert first_doc is not None
|
||||
assert first_doc.id is not None
|
||||
|
||||
def test_query_by_id(
|
||||
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
|
||||
):
|
||||
"""Test querying by document ID with DiskAnn index."""
|
||||
query = Query(
|
||||
field_name="embedding",
|
||||
id=multiple_docs[0].id,
|
||||
param=DiskAnnQueryParam(list_size=100),
|
||||
)
|
||||
|
||||
result = collection_with_multiple_docs.query(queries=query, topk=10)
|
||||
assert len(result) > 0
|
||||
assert len(result) <= 10
|
||||
|
||||
def test_query_with_different_list_size(
|
||||
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
|
||||
):
|
||||
"""Test querying with different list_size parameter values."""
|
||||
query_vector = multiple_docs[0].vector("embedding")
|
||||
|
||||
# Test with list_size=50
|
||||
query_small = Query(
|
||||
field_name="embedding",
|
||||
vector=query_vector,
|
||||
param=DiskAnnQueryParam(list_size=50),
|
||||
)
|
||||
result_small = collection_with_multiple_docs.query(queries=query_small, topk=10)
|
||||
assert len(result_small) > 0
|
||||
|
||||
# Test with list_size=200
|
||||
query_large = Query(
|
||||
field_name="embedding",
|
||||
vector=query_vector,
|
||||
param=DiskAnnQueryParam(list_size=200),
|
||||
)
|
||||
result_large = collection_with_multiple_docs.query(queries=query_large, topk=10)
|
||||
assert len(result_large) > 0
|
||||
|
||||
def test_query_with_topk(
|
||||
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
|
||||
):
|
||||
"""Test querying with different topk values."""
|
||||
query_vector = multiple_docs[0].vector("embedding")
|
||||
query = Query(
|
||||
field_name="embedding",
|
||||
vector=query_vector,
|
||||
param=DiskAnnQueryParam(list_size=100),
|
||||
)
|
||||
|
||||
# Test topk=5
|
||||
result_5 = collection_with_multiple_docs.query(queries=query, topk=5)
|
||||
assert len(result_5) <= 5
|
||||
|
||||
# Test topk=20
|
||||
result_20 = collection_with_multiple_docs.query(queries=query, topk=20)
|
||||
assert len(result_20) <= 20
|
||||
|
||||
def test_query_with_filter(
|
||||
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
|
||||
):
|
||||
"""Test querying with filter conditions."""
|
||||
query_vector = multiple_docs[0].vector("embedding")
|
||||
query = Query(
|
||||
field_name="embedding",
|
||||
vector=query_vector,
|
||||
param=DiskAnnQueryParam(list_size=100),
|
||||
)
|
||||
|
||||
# Query with id filter
|
||||
result = collection_with_multiple_docs.query(
|
||||
queries=query, topk=10, filter="id < 50"
|
||||
)
|
||||
assert len(result) > 0
|
||||
for doc in result:
|
||||
assert doc.field("id") < 50
|
||||
|
||||
def test_query_with_output_fields(
|
||||
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
|
||||
):
|
||||
"""Test querying with specific output fields."""
|
||||
query_vector = multiple_docs[0].vector("embedding")
|
||||
query = Query(
|
||||
field_name="embedding",
|
||||
vector=query_vector,
|
||||
param=DiskAnnQueryParam(list_size=100),
|
||||
)
|
||||
|
||||
result = collection_with_multiple_docs.query(
|
||||
queries=query, topk=10, output_fields=["id", "name"]
|
||||
)
|
||||
assert len(result) > 0
|
||||
|
||||
first_doc = result[0]
|
||||
assert "id" in first_doc.field_names()
|
||||
assert "name" in first_doc.field_names()
|
||||
|
||||
def test_query_with_include_vector(
|
||||
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
|
||||
):
|
||||
"""Test querying with vector data included in results."""
|
||||
query_vector = multiple_docs[0].vector("embedding")
|
||||
query = Query(
|
||||
field_name="embedding",
|
||||
vector=query_vector,
|
||||
param=DiskAnnQueryParam(list_size=100),
|
||||
)
|
||||
|
||||
result = collection_with_multiple_docs.query(
|
||||
queries=query, topk=10, include_vector=True
|
||||
)
|
||||
assert len(result) > 0
|
||||
|
||||
first_doc = result[0]
|
||||
assert first_doc.vector("embedding") is not None
|
||||
assert len(first_doc.vector("embedding")) == 128
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("diskann_collection")
|
||||
class TestDiskAnnCollectionUpdate:
|
||||
"""Test document update in DiskAnn collection."""
|
||||
|
||||
def test_update_doc_fields(
|
||||
self, collection_with_single_doc: Collection, single_doc: Doc
|
||||
):
|
||||
"""Test updating document fields."""
|
||||
updated_doc = Doc(
|
||||
id=single_doc.id,
|
||||
fields={"id": single_doc.field("id"), "name": "updated_name"},
|
||||
)
|
||||
|
||||
result = collection_with_single_doc.update(updated_doc)
|
||||
assert bool(result)
|
||||
assert result.ok()
|
||||
|
||||
# Verify update
|
||||
fetched = collection_with_single_doc.fetch(ids=[single_doc.id])
|
||||
assert single_doc.id in fetched
|
||||
doc = fetched[single_doc.id]
|
||||
assert doc.field("name") == "updated_name"
|
||||
|
||||
def test_update_doc_vector(
|
||||
self, collection_with_single_doc: Collection, single_doc: Doc
|
||||
):
|
||||
"""Test updating document vector."""
|
||||
new_vector = [0.5 + i * 0.01 for i in range(128)]
|
||||
updated_doc = Doc(
|
||||
id=single_doc.id,
|
||||
vectors={"embedding": new_vector},
|
||||
)
|
||||
|
||||
result = collection_with_single_doc.update(updated_doc)
|
||||
assert bool(result)
|
||||
assert result.ok()
|
||||
|
||||
# Verify update
|
||||
fetched = collection_with_single_doc.fetch(
|
||||
ids=[single_doc.id],
|
||||
)
|
||||
assert single_doc.id in fetched
|
||||
doc = fetched[single_doc.id]
|
||||
assert doc.vector("embedding") is not None
|
||||
embedding = doc.vector("embedding")
|
||||
assert len(embedding) == 128
|
||||
# Verify vector values are approximately equal (float comparison)
|
||||
for i in range(128):
|
||||
assert math.isclose(embedding[i], new_vector[i], rel_tol=1e-5)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("diskann_collection")
|
||||
class TestDiskAnnCollectionDelete:
|
||||
"""Test document deletion from DiskAnn collection."""
|
||||
|
||||
def test_delete_single_doc(
|
||||
self, collection_with_single_doc: Collection, single_doc: Doc
|
||||
):
|
||||
"""Test deleting a single document."""
|
||||
result = collection_with_single_doc.delete(single_doc.id)
|
||||
assert bool(result)
|
||||
assert result.ok()
|
||||
|
||||
stats = collection_with_single_doc.stats
|
||||
assert stats.doc_count == 0
|
||||
|
||||
def test_delete_multiple_docs(
|
||||
self, collection_with_multiple_docs: Collection, multiple_docs: list[Doc]
|
||||
):
|
||||
"""Test deleting multiple documents."""
|
||||
ids_to_delete = [doc.id for doc in multiple_docs[:10]]
|
||||
result = collection_with_multiple_docs.delete(ids_to_delete)
|
||||
assert len(result) == len(ids_to_delete)
|
||||
for item in result:
|
||||
assert item.ok()
|
||||
|
||||
stats = collection_with_multiple_docs.stats
|
||||
assert stats.doc_count == len(multiple_docs) - len(ids_to_delete)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("diskann_collection")
|
||||
class TestDiskAnnCollectionOptimizeAndReopen:
|
||||
"""Test collection optimize and reopen functionality."""
|
||||
|
||||
def test_optimize_close_reopen_and_query(
|
||||
self,
|
||||
tmp_path_factory,
|
||||
diskann_collection_schema,
|
||||
collection_option,
|
||||
multiple_docs: list[Doc],
|
||||
):
|
||||
"""Test inserting 100 docs, optimize, close, reopen and query."""
|
||||
# Create collection and insert 100 documents
|
||||
temp_dir = tmp_path_factory.mktemp("zvec_diskann_optimize")
|
||||
collection_path = temp_dir / "test_optimize_collection"
|
||||
|
||||
coll = zvec.create_and_open(
|
||||
path=str(collection_path),
|
||||
schema=diskann_collection_schema,
|
||||
option=collection_option,
|
||||
)
|
||||
|
||||
assert coll is not None
|
||||
assert coll.stats.doc_count == 0
|
||||
|
||||
# Insert 100 documents
|
||||
result = coll.insert(multiple_docs)
|
||||
assert len(result) == len(multiple_docs)
|
||||
for item in result:
|
||||
assert item.ok()
|
||||
assert coll.stats.doc_count == len(multiple_docs)
|
||||
|
||||
# Call optimize
|
||||
from zvec import OptimizeOption
|
||||
|
||||
coll.optimize(option=OptimizeOption())
|
||||
|
||||
# Verify data is still accessible after optimize
|
||||
query_vector = multiple_docs[0].vector("embedding")
|
||||
query = Query(
|
||||
field_name="embedding",
|
||||
vector=query_vector,
|
||||
param=DiskAnnQueryParam(list_size=100),
|
||||
)
|
||||
result_before_close = coll.query(queries=query, topk=10)
|
||||
assert len(result_before_close) > 0
|
||||
|
||||
# Close collection (destroy will close it)
|
||||
collection_path_str = str(collection_path)
|
||||
del coll
|
||||
|
||||
# Reopen collection
|
||||
reopened_coll = zvec.open(path=collection_path_str, option=collection_option)
|
||||
assert reopened_coll is not None
|
||||
assert reopened_coll.stats.doc_count == len(multiple_docs)
|
||||
|
||||
# Execute query on reopened collection
|
||||
query_after_reopen = Query(
|
||||
field_name="embedding",
|
||||
vector=query_vector,
|
||||
param=DiskAnnQueryParam(list_size=100),
|
||||
)
|
||||
result_after_reopen = reopened_coll.query(queries=query_after_reopen, topk=10)
|
||||
assert len(result_after_reopen) > 0
|
||||
assert len(result_after_reopen) <= 10
|
||||
|
||||
# Verify query results are valid
|
||||
first_doc = result_after_reopen[0]
|
||||
assert first_doc is not None
|
||||
assert first_doc.id is not None
|
||||
assert first_doc.field("id") is not None
|
||||
assert first_doc.field("name") is not None
|
||||
|
||||
# Cleanup
|
||||
reopened_coll.destroy()
|
||||
Reference in New Issue
Block a user