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