feat(python): expose group-by search to Python API (#561)

Co-authored-by: jiliang.ljl <jiliang.ljl@alibaba-inc.com>

python/tests/test_collection_fts.py

@@ -172,7 +172,9 @@ class TestFtsOnlyCollectionLifecycle:
 class TestFtsOnlyCollectionQueryValidation:
     def test_vector_query_rejected(self, fts_collection: Collection):
         """Vector query on a no-vector collection must raise."""
-        with pytest.raises(ValueError, match="No vector field found"):
+        with pytest.raises(
+            ValueError, match="Vector field 'content' not found in schema"
+        ):
             fts_collection.query(
                 queries=Query(field_name="content", vector=[0.1, 0.2, 0.3]),
                 topk=5,
@@ -181,7 +183,9 @@ class TestFtsOnlyCollectionQueryValidation:
     def test_id_query_rejected(self, fts_collection: Collection):
         """ID-based query on a no-vector collection must raise."""
         fts_collection.insert(_make_docs()[:1])
-        with pytest.raises(ValueError, match="No vector field found"):
+        with pytest.raises(
+            ValueError, match="Vector field 'content' not found in schema"
+        ):
             fts_collection.query(
                 queries=Query(field_name="content", id="pk_0"),
                 topk=5,

python/tests/test_collection_group_by.py

@@ -0,0 +1,348 @@
+# 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")

python/zvec/__init__.py

@@ -95,7 +95,7 @@ from .model import schema as schema
 
 # —— Core data structures ——
 from .model.collection import Collection
-from .model.doc import Doc, DocList
+from .model.doc import Doc, DocList, GroupResult
 
 # —— Query & index parameters ——
 # —— FTS params (C++ binding) ——
@@ -161,6 +161,7 @@ __all__ = [
     "VectorSchema",
     "CollectionStats",
     # Parameters
+    "GroupResult",
     "Query",
     "VectorQuery",
     "Fts",

python/zvec/__init__.pyi

@@ -11,7 +11,7 @@ from .extension import ReRanker, RrfReRanker, WeightedReRanker
 from .extension.embedding import DenseEmbeddingFunction
 from .model import param, schema
 from .model.collection import Collection
-from .model.doc import Doc, DocList
+from .model.doc import Doc, DocList, GroupResult
 from .model.param import (
     AddColumnOption,
     AlterColumnOption,
@@ -66,6 +66,7 @@ __all__: list = [
     "Fts",
     "FtsIndexParam",
     "FtsQueryParam",
+    "GroupResult",
     "HnswIndexParam",
     "HnswQueryParam",
     "HnswRabitqIndexParam",
@@ -132,7 +133,7 @@ class _Collection:
         include_vector: bool = True,
     ) -> dict[str, _Doc]: ...
     def Flush(self) -> None: ...
-    def GroupByQuery(self, arg0: ...) -> list[...]: ...
+    def GroupByQuery(self, arg0: param._GroupByVectorQuery) -> list[_GroupResult]: ...
     def Insert(self, arg0: collections.abc.Sequence[_Doc]) -> list[typing.Status]: ...
     def Optimize(self, arg0: param.OptimizeOption) -> None: ...
     def Options(self) -> param.CollectionOption: ...
@@ -165,6 +166,12 @@ class _Doc:
     def set_pk(self, arg0: str) -> None: ...
     def set_score(self, arg0: typing.SupportsFloat) -> None: ...
 
+class _GroupResult:
+    @property
+    def docs(self) -> list[_Doc]: ...
+    @property
+    def group_by_value(self) -> str: ...
+
 class _DocOp:
     """
     Members:

python/zvec/executor/query_executor.py

@@ -223,6 +223,31 @@ class QueryExecutor:
             fts.match_string = query.fts.match_string or ""
             search_query.fts = fts
 
+    def set_query_vector(
+        self, query: Query, search_query, collection: _Collection
+    ) -> None:
+        """Resolve a Query vector and set it on a native query object."""
+        vector_schema = self._schema.vector(query.field_name)
+        if vector_schema is None:
+            raise ValueError(f"Vector field '{query.field_name}' not found in schema")
+
+        if query.has_vector():
+            vec_data = query.vector
+        elif query.has_id():
+            fetched = collection.Fetch([query.id])
+            doc = next(iter(fetched.values()), None)
+            if not doc:
+                raise ValueError(f"Document with id '{query.id}' not found")
+            vec_data = doc.get_any(vector_schema.name, vector_schema.data_type)
+        else:
+            raise ValueError("Group by query requires a vector or document id")
+
+        target_dtype = DTYPE_MAP.get(vector_schema.data_type.value)
+        search_query.set_vector(
+            vector_schema._get_object(),
+            convert_to_numpy(vec_data, target_dtype) if target_dtype else vec_data,
+        )
+
     def _build_search_query(
         self, ctx: QueryContext, query: Query, collection: _Collection
     ) -> _SearchQuery:
@@ -235,32 +260,6 @@ class QueryExecutor:
         # set FTS query if provided
         self._apply_fts(query, search_query)
 
-        vector_schema = None
         if query.has_vector() or query.has_id():
-            vector_schema = (
-                self._schema.vector(query.field_name)
-                if query
-                else self._schema.vectors[0]
-            )
-
-            if vector_schema is None:
-                raise ValueError("No vector field found")
-
-        # set vector
-        if query.has_vector():
-            vec_data = query.vector
-        elif query.has_id():
-            fetched = collection.Fetch([query.id])
-            doc = next(iter(fetched.values()), None)
-            if not doc:
-                raise ValueError(f"Document with id '{query.id}' not found")
-            vec_data = doc.get_any(vector_schema.name, vector_schema.data_type)
-        else:
-            return search_query
-
-        target_dtype = DTYPE_MAP.get(vector_schema.data_type.value)
-        search_query.set_vector(
-            vector_schema._get_object(),
-            convert_to_numpy(vec_data, target_dtype) if target_dtype else vec_data,
-        )
+            self.set_query_vector(query, search_query, collection)
         return search_query

python/zvec/model/__init__.py

@@ -14,7 +14,7 @@
 from __future__ import annotations
 
 from .collection import Collection
-from .doc import Doc
+from .doc import Doc, GroupResult
 from .param.query import Fts, Query, VectorQuery
 from .schema.collection_schema import CollectionSchema
 from .schema.field_schema import FieldSchema
@@ -25,6 +25,7 @@ __all__ = [
     "Doc",
     "FieldSchema",
     "Fts",
+    "GroupResult",
     "Query",
     "VectorQuery",
 ]

python/zvec/model/collection.py

@@ -17,12 +17,13 @@ import warnings
 from typing import Optional, Union, overload
 
 from zvec._zvec import _Collection
+from zvec._zvec.param import _GroupByVectorQuery
 
 from ..executor import QueryContext, QueryExecutor
 from ..extension import ReRanker
 from ..typing import Status
 from .convert import convert_to_cpp_doc, convert_to_py_doc
-from .doc import Doc, DocList
+from .doc import Doc, DocList, GroupResult
 from .param import (
     AddColumnOption,
     AlterColumnOption,
@@ -42,6 +43,11 @@ from .schema import CollectionSchema, CollectionStats, FieldSchema
 __all__ = ["Collection"]
 
 
+def _require_positive_integer(value, name: str) -> None:
+    if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
+        raise ValueError(f"{name} must be a positive integer")
+
+
 class Collection:
     """Represents an opened collection in Zvec.
 
@@ -437,3 +443,80 @@ class Collection:
             reranker=reranker,
         )
         return self._querier.execute(ctx, self._obj)
+
+    def group_by_query(
+        self,
+        query: Query,
+        group_by_field_name: str,
+        group_count=2,
+        topk_per_group=3,
+        *,
+        filter: Optional[str] = None,
+        include_vector: bool = False,
+        output_fields: Optional[list[str]] = None,
+    ) -> list[GroupResult]:
+        """Perform group-by vector search.
+
+        Groups results by a scalar field value, returning the top-k documents
+        within each group, ordered by similarity score.
+
+        Accepts a single vector Query. Full-text search is not supported.
+
+        Args:
+            query (Query): Vector query.
+            group_by_field_name (str): Scalar field used to group results.
+            group_count (int): Maximum number of groups to return.
+            topk_per_group (int): Maximum number of documents in each group.
+            filter (Optional[str]): Boolean expression used to filter candidates.
+            include_vector (bool): Whether returned documents include vectors.
+            output_fields (Optional[list[str]]): Scalar fields to return.
+
+        Returns:
+            list[GroupResult]: Grouped documents sorted by score.
+
+        Examples:
+            >>> results = collection.group_by_query(
+            ...     zvec.Query(
+            ...         field_name="embedding",
+            ...         vector=[0.1, 0.2, 0.3],
+            ...         param=zvec.HnswQueryParam(ef=300),
+            ...     ),
+            ...     group_by_field_name="category",
+            ...     group_count=5,
+            ...     topk_per_group=3,
+            ... )
+            >>> for group in results:
+            ...     print(group.group_by_value, len(group.docs))
+        """
+        query._validate()
+        if query.has_fts():
+            raise ValueError("Group by query does not support FTS")
+        if not query.has_vector() and not query.has_id():
+            raise ValueError("Group by query requires a vector or document id")
+        if not group_by_field_name:
+            raise ValueError("group_by_field_name cannot be empty")
+        _require_positive_integer(group_count, "group_count")
+        _require_positive_integer(topk_per_group, "topk_per_group")
+        cpp_query = _GroupByVectorQuery()
+        cpp_query.field_name = query.field_name
+        cpp_query.group_by_field_name = group_by_field_name
+        cpp_query.group_count = group_count
+        cpp_query.topk_per_group = topk_per_group
+        cpp_query.include_vector = include_vector
+        if filter:
+            cpp_query.filter = filter
+        if output_fields is not None:
+            cpp_query.output_fields = output_fields
+        if query.param:
+            cpp_query.query_params = query.param
+        self._querier.set_query_vector(query, cpp_query, self._obj)
+
+        raw_results = self._obj.GroupByQuery(cpp_query)
+
+        return [
+            GroupResult(
+                group_by_value=group.group_by_value,
+                docs=[convert_to_py_doc(doc, self.schema) for doc in group.docs],
+            )
+            for group in raw_results
+        ]

python/zvec/model/doc.py

@@ -14,6 +14,7 @@
 from __future__ import annotations
 
 import json
+from dataclasses import dataclass
 from typing import Any, Optional
 
 from ..common import VectorType
@@ -21,6 +22,7 @@ from ..common import VectorType
 __all__ = [
     "Doc",
     "DocList",
+    "GroupResult",
 ]
 
 
@@ -176,3 +178,16 @@ class Doc:
 
 #: Type alias for query results: a list of documents returned by a single query route.
 DocList = list[Doc]
+
+
+@dataclass(frozen=True)
+class GroupResult:
+    """A group value and its matching documents.
+
+    Attributes:
+        group_by_value (str): String representation of the grouped scalar field value.
+        docs (list[Doc]): Matching documents in similarity order.
+    """
+
+    group_by_value: str
+    docs: list[Doc]