alibaba--zvec/python/tests/test_collection_group_by.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.
from __future__ import annotations

import pytest
import zvec
from zvec import (
    Collection,
    CollectionOption,
    DataType,
    Doc,
    FieldSchema,
    FlatIndexParam,
    Fts,
    GroupResult,
    HnswIndexParam,
    HnswQueryParam,
    InvertIndexParam,
    Query,
    VectorSchema,
)

# ==================== Constants ====================

GB_DIMENSION = 4
GB_NUM_DOCS = 12
GB_NUM_GROUPS = 3
GB_TOPK_PER_GROUP = 2


# ==================== Fixtures ====================


@pytest.fixture(scope="session")
def group_by_collection_schema():
    """Collection schema for group-by end-to-end tests.

    Mirrors the data layout in ``vector_column_indexer_test.cc``:
    a 4-dimensional dense vector and a scalar ``group_id`` field used
    for grouping.
    """
    return zvec.CollectionSchema(
        name="test_group_by_collection",
        fields=[
            FieldSchema(
                "id",
                DataType.INT64,
                nullable=False,
                index_param=InvertIndexParam(enable_range_optimization=True),
            ),
            FieldSchema(
                "group_id",
                DataType.INT64,
                nullable=False,
                index_param=InvertIndexParam(),
            ),
        ],
        vectors=[
            VectorSchema(
                "dense",
                DataType.VECTOR_FP32,
                dimension=GB_DIMENSION,
                index_param=HnswIndexParam(),
            ),
            VectorSchema(
                "dense_flat",
                DataType.VECTOR_FP32,
                dimension=GB_DIMENSION,
                index_param=FlatIndexParam(),
            ),
        ],
    )


@pytest.fixture(scope="session")
def collection_option():
    return CollectionOption(read_only=False, enable_mmap=True)


@pytest.fixture
def group_by_docs():
    """Generate docs matching the C++ GroupByIndexerTest fixture.

    Doc ``i`` has vector ``[i, i, i, i]`` and ``group_id = i % 3``.
    """
    return [
        Doc(
            id=f"{i}",
            fields={"id": i, "group_id": i % GB_NUM_GROUPS},
            vectors={
                "dense": [float(i)] * GB_DIMENSION,
                "dense_flat": [float(i)] * GB_DIMENSION,
            },
        )
        for i in range(GB_NUM_DOCS)
    ]


@pytest.fixture(scope="function")
def group_by_collection(
    tmp_path_factory, group_by_collection_schema, collection_option
) -> Collection:
    """Function-scoped fixture: creates and opens a collection for group-by tests."""
    temp_dir = tmp_path_factory.mktemp("zvec_group_by")
    collection_path = temp_dir / "test_group_by_collection"

    coll = zvec.create_and_open(
        path=str(collection_path),
        schema=group_by_collection_schema,
        option=collection_option,
    )

    assert coll is not None, "Failed to create and open group-by collection"
    assert coll.path == str(collection_path)
    assert coll.schema.name == group_by_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 group_by_collection_with_docs(
    group_by_collection: Collection, group_by_docs
) -> Collection:
    """Setup: insert group-by fixture docs."""
    assert group_by_collection.stats.doc_count == 0
    result = group_by_collection.insert(group_by_docs)
    assert len(result) == len(group_by_docs)
    for item in result:
        assert item.ok()
    assert group_by_collection.stats.doc_count == len(group_by_docs)

    yield group_by_collection

    # Teardown
    group_by_collection.delete([doc.id for doc in group_by_docs])


# ==================== Helpers ====================


def _assert_grouped_results(results, num_groups, topk_per_group, query_value):
    """Validate group-by result structure and ordering.

    Each returned group must:
      - contain only docs whose ``group_id`` matches ``group_by_value``
      - have at most ``topk_per_group`` docs
      - have docs sorted by descending score
    """
    assert len(results) == num_groups, (
        f"Expected {num_groups} groups, got {len(results)}"
    )

    group_values = set()
    for group in results:
        assert isinstance(group, GroupResult)
        group_value = int(group.group_by_value)
        group_values.add(group_value)
        docs = group.docs
        assert 1 <= len(docs) <= topk_per_group

        for doc in docs:
            assert int(doc.field("group_id")) == group_value

        scores = [doc.score for doc in docs]
        assert scores == sorted(scores, reverse=True), (
            "Docs must be sorted by score desc"
        )

        # Score sanity: for query [1,1,1,1] and vector [i,i,i,i],
        # IP score is 4 * i.
        for doc in docs:
            doc_id = int(doc.field("id"))
            expected_score = float(doc_id * sum(query_value))
            assert abs(doc.score - expected_score) < 0.1

    assert group_values == set(range(num_groups))


# ==================== Tests ====================


@pytest.mark.usefixtures("group_by_collection_with_docs")
class TestGroupBySearch:
    def test_group_by_defaults(self, group_by_collection: Collection):
        results = group_by_collection.group_by_query(
            Query(field_name="dense_flat", vector=[1.0] * GB_DIMENSION),
            group_by_field_name="group_id",
        )
        assert len(results) == 2
        assert all(1 <= len(group.docs) <= 3 for group in results)

    def test_group_by_hnsw(self, group_by_collection: Collection):
        """Group-by search over an HNSW index."""
        query_vector = [1.0] * GB_DIMENSION
        results = group_by_collection.group_by_query(
            Query(
                field_name="dense", vector=query_vector, param=HnswQueryParam(ef=300)
            ),
            group_by_field_name="group_id",
            group_count=GB_NUM_GROUPS,
            topk_per_group=GB_TOPK_PER_GROUP,
        )
        _assert_grouped_results(results, GB_NUM_GROUPS, GB_TOPK_PER_GROUP, query_vector)

    def test_group_by_flat(self, group_by_collection: Collection):
        """Group-by search over a FLAT index."""
        query_vector = [1.0] * GB_DIMENSION
        results = group_by_collection.group_by_query(
            Query(field_name="dense_flat", vector=query_vector),
            group_by_field_name="group_id",
            group_count=GB_NUM_GROUPS,
            topk_per_group=GB_TOPK_PER_GROUP,
        )
        _assert_grouped_results(results, GB_NUM_GROUPS, GB_TOPK_PER_GROUP, query_vector)

    def test_group_by_with_filter(self, group_by_collection: Collection):
        """Group-by search with a scalar filter."""
        query_vector = [1.0] * GB_DIMENSION
        results = group_by_collection.group_by_query(
            Query(field_name="dense_flat", vector=query_vector),
            group_by_field_name="group_id",
            group_count=GB_NUM_GROUPS,
            topk_per_group=GB_TOPK_PER_GROUP,
            filter="id < 6",
        )
        # Only docs 0..5 are visible; every group still has at least one doc.
        assert len(results) == GB_NUM_GROUPS
        for group in results:
            for doc in group.docs:
                assert int(doc.field("id")) < 6

    def test_group_by_include_vector(self, group_by_collection: Collection):
        """Group-by search returns original vectors when requested."""
        query_vector = [1.0] * GB_DIMENSION
        results = group_by_collection.group_by_query(
            Query(field_name="dense_flat", vector=query_vector),
            group_by_field_name="group_id",
            group_count=GB_NUM_GROUPS,
            topk_per_group=GB_TOPK_PER_GROUP,
            include_vector=True,
        )
        assert len(results) == GB_NUM_GROUPS
        for group in results:
            for doc in group.docs:
                vec = doc.vector("dense_flat")
                doc_id = int(doc.field("id"))
                assert vec == pytest.approx([float(doc_id)] * GB_DIMENSION, abs=1e-5)

    def test_group_by_output_fields(self, group_by_collection: Collection):
        """Group-by search honors scalar output field selection."""
        query_vector = [1.0] * GB_DIMENSION
        results = group_by_collection.group_by_query(
            Query(field_name="dense_flat", vector=query_vector),
            group_by_field_name="group_id",
            group_count=GB_NUM_GROUPS,
            topk_per_group=GB_TOPK_PER_GROUP,
            output_fields=["group_id"],
        )
        assert len(results) == GB_NUM_GROUPS
        for group in results:
            for doc in group.docs:
                assert doc.has_field("group_id")

    def test_group_by_invalid_field(self, group_by_collection: Collection):
        """Group-by with a non-existent vector field raises an error."""
        with pytest.raises(ValueError):
            group_by_collection.group_by_query(
                Query(field_name="nonexistent", vector=[1.0] * GB_DIMENSION),
                group_by_field_name="group_id",
            )

    def test_group_by_query_by_id(self, group_by_collection: Collection):
        results = group_by_collection.group_by_query(
            Query(field_name="dense_flat", id="11"),
            group_by_field_name="group_id",
            group_count=GB_NUM_GROUPS,
            topk_per_group=GB_TOPK_PER_GROUP,
        )
        assert len(results) == GB_NUM_GROUPS

    @pytest.mark.parametrize(
        ("query", "error"),
        [
            (Query(field_name="content", fts=Fts(match_string="text")), "FTS"),
            (Query(field_name="dense_flat"), "vector or document id"),
        ],
    )
    def test_group_by_rejects_unsupported_query(
        self, group_by_collection: Collection, query: Query, error: str
    ):
        with pytest.raises(ValueError, match=error):
            group_by_collection.group_by_query(query, "group_id")

    @pytest.mark.parametrize(
        ("kwargs", "error"),
        [
            ({"group_by_field_name": ""}, "group_by_field_name"),
            ({"group_by_field_name": "group_id", "group_count": 0}, "group_count"),
            (
                {"group_by_field_name": "group_id", "topk_per_group": 0},
                "topk_per_group",
            ),
        ],
    )
    def test_group_by_rejects_invalid_group_params(
        self, group_by_collection: Collection, kwargs: dict, error: str
    ):
        with pytest.raises(ValueError, match=error):
            group_by_collection.group_by_query(
                Query(field_name="dense_flat", vector=[1.0] * GB_DIMENSION),
                **kwargs,
            )


class TestGroupByEmptyCollection:
    def test_group_by_empty_collection(self, group_by_collection: Collection):
        """Group-by on an empty collection returns an empty list."""
        results = group_by_collection.group_by_query(
            Query(field_name="dense_flat", vector=[1.0] * GB_DIMENSION),
            group_by_field_name="group_id",
            group_count=GB_NUM_GROUPS,
            topk_per_group=GB_TOPK_PER_GROUP,
        )
        assert results == []


def test_group_by_public_api_exports():
    assert zvec.GroupResult is GroupResult
    assert not hasattr(zvec, "GroupByQuery")
    assert not hasattr(Collection, "groupby_query")