alibaba--zvec/python/zvec/model/collection.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 warnings
from typing import Optional, Union, overload

from zvec._zvec import _Collection

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 .param import (
    AddColumnOption,
    AlterColumnOption,
    CollectionOption,
    FlatIndexParam,
    FtsIndexParam,
    HnswIndexParam,
    HnswRabitqIndexParam,
    IndexOption,
    InvertIndexParam,
    IVFIndexParam,
    OptimizeOption,
)
from .param.query import Query
from .schema import CollectionSchema, CollectionStats, FieldSchema

__all__ = ["Collection"]


class Collection:
    """Represents an opened collection in Zvec.

    A `Collection` provides methods for data definition (DDL), data manipulation (DML),
    and querying (DQL). It is obtained via `create_and_open()` or `open()`.

    This class is not meant to be instantiated directly; use factory functions instead.
    """

    def __init__(self, obj: _Collection):
        self._obj = obj
        self._schema = None
        self._querier = None

    @classmethod
    def _from_core(cls, core_collection: _Collection) -> Collection:
        if not core_collection:
            raise ValueError("Collection is None")
        inst = cls.__new__(cls)
        inst._obj = core_collection
        schema = CollectionSchema._from_core(core_collection.Schema())
        inst._schema = schema
        inst._querier = QueryExecutor(schema)
        return inst

    @property
    def path(self) -> str:
        """str: The filesystem path of the collection."""
        return self._obj.Path()

    @property
    def option(self) -> CollectionOption:
        """CollectionOption: The options used to open the collection."""
        return self._obj.Options()

    @property
    def schema(self) -> CollectionSchema:
        """CollectionSchema: The schema defining the structure of the collection."""
        return self._schema

    @property
    def stats(self) -> CollectionStats:
        """CollectionStats: Runtime statistics about the collection (e.g., doc count, size)."""
        return self._obj.Stats()

    # ========== Collection DDL Methods ==========
    def destroy(self) -> None:
        """Permanently delete the collection from disk.

        Warning:
            This operation is irreversible. All data will be lost.
        """
        self._obj.Destroy()

    def flush(self) -> None:
        """Force all pending writes to disk.

        Ensures durability of recent inserts/updates.
        """
        self._obj.Flush()

    # ========== Index DDL Methods ==========
    def create_index(
        self,
        field_name: str,
        index_param: Union[
            HnswIndexParam,
            HnswRabitqIndexParam,
            IVFIndexParam,
            FlatIndexParam,
            InvertIndexParam,
            FtsIndexParam,
        ],
        option: IndexOption = IndexOption(),
    ) -> None:
        """Create an index on a field.

        Vector index types (HNSW, IVF, FLAT) can only be applied to vector fields.
        Inverted index (`InvertIndexParam`) is for scalar fields.
        FTS index (`FtsIndexParam`) is for full-text search on STRING fields.

        Args:
            field_name (str): Name of the field to index.
            index_param (Union[HnswIndexParam, HnswRabitqIndexParam, IVFIndexParam, FlatIndexParam, InvertIndexParam, FtsIndexParam]):
                Index configuration.
            option (Optional[IndexOption], optional): Index creation options.
                Defaults to ``IndexOption()``.

        """
        self._obj.CreateIndex(field_name, index_param, option)
        self._schema = CollectionSchema._from_core(self._obj.Schema())
        self._querier._schema = self._schema

    def drop_index(self, field_name: str) -> None:
        """Remove the index from a field.

        Args:
            field_name (str): Name of the indexed field.
        """
        self._obj.DropIndex(field_name)
        self._schema = CollectionSchema._from_core(self._obj.Schema())
        self._querier._schema = self._schema

    def optimize(self, option: OptimizeOption = OptimizeOption()) -> None:
        """Optimize the collection (e.g., merge segments, rebuild index).

        Args:
            option (Optional[OptimizeOption], optional): Optimization options.
                Defaults to ``OptimizeOption()``.
        """
        self._obj.Optimize(option)

    # ========== COLUMN DDL Methods ==========
    def add_column(
        self,
        field_schema: FieldSchema,
        expression: str = "",
        option: AddColumnOption = AddColumnOption(),
    ) -> None:
        """Add a new column to the collection.

        The column is populated using the provided expression (e.g., SQL-like formula).

        Args:
            field_schema (FieldSchema): Schema definition for the new column.
            expression (str): Expression to compute values for existing documents.
            option (Optional[AddColumnOption], optional): Options for the operation.
                Defaults to ``AddColumnOption()``.
        """
        self._obj.AddColumn(field_schema._get_object(), expression, option)
        self._schema = CollectionSchema._from_core(self._obj.Schema())
        self._querier._schema = self._schema

    def drop_column(self, field_name: str) -> None:
        """Remove a column from the collection.

        Args:
            field_name (str): Name of the column to drop.
        """
        self._obj.DropColumn(field_name)
        self._schema = CollectionSchema._from_core(self._obj.Schema())
        self._querier._schema = self._schema

    def alter_column(
        self,
        old_name: str,
        new_name: Optional[str] = None,
        field_schema: Optional[FieldSchema] = None,
        option: AlterColumnOption = AlterColumnOption(),
    ) -> None:
        """Rename a column, update its schema.

        This method supports three atomic operations:
          1. Rename only (when `field_schema` is None).
          2. Modify schema only (when `new_name` is None or empty string).

        Args:
            old_name (str): The current name of the column to be altered.
            new_name (Optional[str]): The new name for the column.
                - If provided and non-empty, the column will be renamed.
                - If `None` or empty string, no rename occurs.
            field_schema (Optional[FieldSchema]): The new schema definition.
                - If provided, the column's type, dimension, or other properties will be updated.
                - If `None`, only renaming (if requested) is performed.
            option (AlterColumnOption, optional): Options controlling the alteration behavior.
                Defaults to ``AlterColumnOption()``.

        **Limitation**: This operation **only supports scalar numeric columns**. such as:
        - `DOUBLE`, `FLOAT`,
        - `INT32`, `INT64`, `UINT32`, `UINT64`

        Note:
            - Schema modification may trigger data migration or index rebuild.

        Examples:
            >>> # Rename column only
            >>> results = collection.alter_column(old_name="id", new_name="doc_id")

            >>> # Modify schema only
            >>> new_schema = FieldSchema(name="doc_id", dtype=DataType.INT64)
            >>> collection.alter_column("id", field_schema=new_schema)
        """
        self._obj.AlterColumn(
            old_name,
            new_name or "",
            field_schema._get_object() if field_schema else None,
            option,
        )
        self._schema = CollectionSchema._from_core(self._obj.Schema())
        self._querier._schema = self._schema

    # ========== Collection DDL Methods ==========
    @overload
    def insert(self, docs: Doc) -> Status:
        pass

    @overload
    def insert(self, docs: list[Doc]) -> list[Status]:
        pass

    def insert(self, docs: Union[Doc, list[Doc]]) -> Union[Status, list[Status]]:
        """Insert new documents into the collection.

        Documents must have unique IDs and conform to the schema.

        Args:
            docs (Union[Doc, list[Doc]]): One or more documents to insert.

        Returns:
            Union[Status, list[Status]]: If a single Doc was given, returns its Status;
            if a list was given, returns a list of Status objects.
        """
        is_single = isinstance(docs, Doc)
        doc_list = [docs] if is_single else docs
        results = self._obj.Insert(
            [convert_to_cpp_doc(doc, self.schema) for doc in doc_list]
        )
        return results[0] if is_single else results

    @overload
    def upsert(self, docs: Doc) -> Status:
        pass

    @overload
    def upsert(self, docs: list[Doc]) -> list[Status]:
        pass

    def upsert(self, docs: Union[Doc, list[Doc]]) -> Union[Status, list[Status]]:
        """Insert new documents or update existing ones by ID.

        Args:
            docs (Union[Doc, list[Doc]]): Documents to upsert.

        Returns:
            Union[Status, list[Status]]: If a single Doc was given, returns its Status;
            if a list was given, returns a list of Status objects.
        """
        is_single = isinstance(docs, Doc)
        doc_list = [docs] if is_single else docs
        results = self._obj.Upsert(
            [convert_to_cpp_doc(doc, self.schema) for doc in doc_list]
        )
        return results[0] if is_single else results

    @overload
    def update(self, docs: Doc) -> Status:
        pass

    @overload
    def update(self, docs: list[Doc]) -> list[Status]:
        pass

    def update(self, docs: Union[Doc, list[Doc]]) -> Union[Status, list[Status]]:
        """Update existing documents by ID.

        Only specified fields are updated; others remain unchanged.

        Args:
            docs (Union[Doc, list[Doc]]): Documents containing updated fields.

        Returns:
            Union[Status, list[Status]]: If a single Doc was given, returns its Status;
            if a list was given, returns a list of Status objects.
        """
        is_single = isinstance(docs, Doc)
        doc_list = [docs] if is_single else docs
        results = self._obj.Update(
            [convert_to_cpp_doc(doc, self.schema) for doc in doc_list]
        )
        return results[0] if is_single else results

    @overload
    def delete(self, ids: str) -> Status:
        pass

    @overload
    def delete(self, ids: list[str]) -> list[Status]:
        pass

    def delete(self, ids: Union[str, list[str]]) -> Union[Status, list[Status]]:
        """Delete documents by ID.

        Args:
            ids (Union[str, list[str]]): One or more document IDs to delete.

        Returns:
            Union[Status, list[Status]]: If a single id was given, returns its Status;
            if a list was given, returns a list of Status objects.
        """
        is_single = isinstance(ids, str)
        id_list = [ids] if isinstance(ids, str) else ids
        results = self._obj.Delete(id_list)
        return results[0] if is_single else results

    def delete_by_filter(self, filter: str) -> None:
        """Delete documents matching a filter expression.

        Args:
            filter (str): Boolean expression (e.g., ``"age > 30"``).
        """
        self._obj.DeleteByFilter(filter)

    # ========== Collection DQL-fetch Methods ==========
    def fetch(
        self,
        ids: Union[str, list[str]],
        *,
        output_fields: Optional[list[str]] = None,
        include_vector: bool = True,
    ) -> dict[str, Doc]:
        """Retrieve documents by ID.

        Args:
            ids (Union[str, list[str]]): Document IDs to fetch.
            output_fields (Optional[list[str]], optional): Scalar fields to
                include. If None, all fields are returned. Defaults to None.
            include_vector (bool, optional): Whether to include vector data in
                results. Defaults to True.

        Returns:
            dict[str, Doc]: Mapping from ID to document. Missing IDs are omitted.
        """
        ids = [ids] if isinstance(ids, str) else ids
        docs = self._obj.Fetch(ids, output_fields, include_vector)
        return {
            doc_id: py_doc
            for doc_id, core_doc in docs.items()
            if (py_doc := convert_to_py_doc(core_doc, self.schema)) is not None
        }

    # ========== Collection DQL-Query Methods ==========

    def query(
        self,
        queries: Optional[Union[Query, list[Query]]] = None,
        *,
        vectors: Optional[Union[Query, list[Query]]] = None,
        topk: int = 10,
        filter: Optional[str] = None,
        include_vector: bool = False,
        output_fields: Optional[list[str]] = None,
        reranker: Optional[ReRanker] = None,
    ) -> DocList:
        """Perform vector similarity search with optional filtering and re-ranking.

        At least one `Query` must be provided via `queries`.

        Args:
            queries (Optional[Union[Query, list[Query]]], optional):
                One or more vector queries. Defaults to None.
            vectors (Optional[Union[Query, list[Query]]], optional):
                Deprecated. Use `queries` instead.
            topk (int, optional): Number of nearest neighbors to return.
                Defaults to 10.
            filter (Optional[str], optional): Boolean expression to pre-filter candidates.
                Defaults to None.
            include_vector (bool, optional): Whether to include vector data in results.
                Defaults to False.
            output_fields (Optional[list[str]], optional): Scalar fields to include.
                If None, all fields are returned. Defaults to None.
            reranker (Optional[ReRanker], optional): Re-ranker to refine results.
                Defaults to None.

        Returns:
            DocList: Top-k matching documents, sorted by relevance score.

        Examples:
            >>> from zvec import Query
            >>> results = collection.query(
            ...     queries=Query(field_name="embedding", vector=[0.1, 0.2]),
            ...     topk=5,
            ...     filter="category == 'tech'",
            ...     output_fields=["title", "url"]
            ... )
        """
        if vectors is not None:
            warnings.warn(
                "The 'vectors' parameter is deprecated and will be removed in a future version. "
                "Use 'queries' instead.",
                DeprecationWarning,
                stacklevel=2,
            )
            if queries is not None:
                raise ValueError("Cannot specify both 'queries' and 'vectors'.")
            queries = vectors

        ctx = QueryContext(
            topk=topk,
            filter=filter,
            queries=[queries] if isinstance(queries, Query) else queries,
            include_vector=include_vector,
            output_fields=output_fields,
            reranker=reranker,
        )
        return self._querier.execute(ctx, self._obj)