alibaba--zvec/python/tests/test_collection_diskann.py

# 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 index.

DiskAnn must be built for Linux (x86_64/ARM64) or macOS ARM64. Other
platforms are skipped wholesale.

If the prerequisite fails the whole module is skipped so the rest of the
test suite is not affected. macOS uses synchronous pread.
"""

from __future__ import annotations

import math
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", "aarch64", "arm64")
        )
        or (sys.platform == "darwin" and platform.machine() in ("aarch64", "arm64"))
    ),
    reason="DiskAnn is supported on Linux (x86_64/ARM64) and macOS ARM64",
)

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()