pyvastbase 0.2.2__py3-none-any.whl

pyvastbase/async_impl/collection.py

@@ -0,0 +1,1367 @@
+"""
+Vastbase Vector Database SDK - Async Collection Module
+
+Async implementation of Collection class.
+Fully compatible with the sync API - same method names, parameters, and return types.
+
+All methods are async and return awaitables. Use 'await' to call them.
+
+Example:
+    from pyvastbase.async import AsyncCollection, async_connect
+    
+    await async_connect(host="localhost", database="mydb")
+    collection = AsyncCollection("my_vectors")
+    await collection.insert([{"id": 1, "embedding": [0.1] * 128}])
+    
+    results = await collection.search(data=[[0.15] * 128], limit=10)
+"""
+
+import json
+import functools
+from typing import Optional, List, Dict, Any, Union, TYPE_CHECKING, Iterator
+
+from ..core.constants import DistanceType, IndexType
+from ..core.collection_core import CollectionCore
+from ..core.exceptions import (
+    CollectionError, CollectionNotExistsError,
+    SchemaError, ParamError, DataError, VectorDimensionError
+)
+from ..core.mutation_result import MutationResult
+from ..operations import (
+    build_delete_sql,
+    build_upsert_clauses,
+    classify_insert_columns,
+    build_columns_query_sql,
+    build_existence_check_sql,
+    build_pk_query_sql,
+    parse_column_row,
+    validate_and_normalize_queries,
+)
+from ..core.schema import CollectionSchema, FieldSchema, DataType
+from ..index.index_builder import IndexParams, IndexBuilder
+from ..search.search_builder import SearchParams, SearchQueryBuilder, SearchResult
+from ..utils import (
+    normalize_vector, validate_vector, ensure_vector_dim,
+    batch_iterator,
+    format_entity, extract_ids, to_list_dict,
+)
+
+if TYPE_CHECKING:
+    from .connections import AsyncVastbaseConnection
+
+from ..executor.async_impl import AsyncExecutor, AsyncTransaction
+
+
+class QueryIterator:
+    """
+    Async iterator for query results in batches.
+    
+    Example:
+        # Iterate over results in batches
+        async for batch in collection.query(batch_size=100):
+            await process(batch)
+    """
+    
+    def __init__(
+        self,
+        collection: 'AsyncCollection',
+        filter_expr: str = "",
+        batch_size: int = 100,
+        output_fields: Optional[List[str]] = None,
+        limit: Optional[int] = None,
+        offset: Optional[int] = None
+    ):
+        self._collection = collection
+        self._filter_expr = filter_expr
+        self._batch_size = batch_size
+        self._output_fields = output_fields
+        self._limit = limit
+        self._offset = offset or 0
+        self._buffer: List[Dict[str, Any]] = []
+        self._exhausted = False
+    
+    @property
+    def batch_size(self) -> int:
+        """Get batch size (for test compatibility)"""
+        return self._batch_size
+    
+    def __aiter__(self) -> 'QueryIterator':
+        return self
+    
+    async def __anext__(self) -> List[Dict[str, Any]]:
+        if self._exhausted:
+            raise StopAsyncIteration
+        
+        # If buffer is empty, fetch next batch
+        if not self._buffer:
+            remaining = self._limit - self._offset if self._limit else None
+            fetch_limit = min(self._batch_size, remaining) if remaining else self._batch_size
+            
+            rows = await self._collection._fetch_batch(
+                filter_expr=self._filter_expr,
+                output_fields=self._output_fields,
+                limit=fetch_limit,
+                offset=self._offset
+            )
+            
+            if not rows:
+                self._exhausted = True
+                raise StopAsyncIteration
+            
+            self._buffer = rows
+        
+        # Return batch from buffer
+        batch = self._buffer[:self._batch_size]
+        self._buffer = self._buffer[self._batch_size:]
+        self._offset += len(batch)
+        
+        return batch
+
+
+# ---------------------------------------------------------------------------
+# Decorator — eliminates connection boilerplate from every public method
+# ---------------------------------------------------------------------------
+
+def _with_connection(method=None, *, schema: bool = False):
+    """Ensure a live connection (and optionally a loaded schema) before the method runs.
+
+    Supports both ``@_with_connection`` and ``@_with_connection(schema=True)``.
+
+    The async version calls sync ``_get_connection()`` and awaits
+    ``_ensure_connection()`` and (when *schema* is True) ``_ensure_schema_loaded()``.
+    """
+    def deco(m):
+        @functools.wraps(m)
+        async def wrapper(self, *args, **kwargs):
+            self._get_connection()
+            await self._ensure_connection()
+            if schema:
+                await self._ensure_schema_loaded()
+            return await m(self, *args, **kwargs)
+        return wrapper
+
+    if method is not None:
+        return deco(method)
+    return deco
+
+
+class AsyncCollection:
+    """
+    Async version of Collection class.
+    
+    Provides the same API as Collection, but all database operations are async.
+    Use 'await' to call methods.
+    
+    Example:
+        from pyvastbase.async import AsyncCollection, async_connect
+        
+        await async_connect(host="localhost", database="mydb")
+        collection = AsyncCollection("my_vectors")
+        
+        # Insert data
+        await collection.insert([
+            {"id": 1, "embedding": [0.1] * 128, "text": "hello"},
+            {"id": 2, "embedding": [0.2] * 128, "text": "world"}
+        ])
+        
+        # Search
+        results = await collection.search(
+            data=[[0.15] * 128],
+            anns_field="embedding",
+            limit=10
+        )
+        
+        # Query with async iterator
+        async for batch in collection.query(batch_size=100):
+            print(batch)
+    """
+    
+    def __init__(
+        self,
+        name: str,
+        schema: Optional[CollectionSchema] = None,
+        using: str = "default",
+        timeout: Optional[float] = None,
+        _connection_provider=None,
+        **kwargs
+    ):
+        self._name = name
+        self._schema = schema
+        self._using = using
+        self._timeout = timeout
+        self._connection_provider = _connection_provider
+        self._extra_params = kwargs
+
+        self._conn: Optional['AsyncVastbaseConnection'] = None
+        self._core_obj: Optional[CollectionCore] = None
+        self._executor: Optional[AsyncExecutor] = None
+    
+    # =========================================================================
+    # Properties
+    # =========================================================================
+    
+    @property
+    def schema(self) -> CollectionSchema:
+        """Get collection schema (synchronous, cached)"""
+        if self._schema is None:
+            raise CollectionError(
+                "Schema not loaded. Use 'await collection.schema_async' or "
+                "call 'await collection.load_schema()' first."
+            )
+        return self._schema
+    
+    @property
+    async def schema_async(self) -> CollectionSchema:
+        """Get collection schema (async, auto-loads if needed)"""
+        if self._schema is None:
+            await self._load_schema_async()
+        return self._schema
+    
+    @property
+    def name(self) -> str:
+        """Get collection name"""
+        return self._name
+    
+    @property
+    def description(self) -> str:
+        """Get collection description"""
+        return self.schema.description
+    
+    @property
+    async def num_entities(self) -> int:
+        """Get number of entities in collection (async property)"""
+        try:
+            rows = await self._executor.execute(
+                f'SELECT COUNT(*) AS cnt FROM "{self._name}"', []
+            )
+            if rows:
+                return rows[0].get("cnt", next(iter(rows[0].values())))
+            return 0
+        except Exception as e:
+            raise CollectionError(f"Failed to get entity count: {e}")
+
+    @property
+    async def is_empty(self) -> bool:
+        """Check if collection is empty"""
+        count = await self.num_entities
+        return count == 0
+
+    async def index_params_async(self) -> Dict[str, Any]:
+        """Get current index parameters (async version)"""
+        try:
+            rows = await self._executor.execute(
+                """
+                SELECT indexname, indexdef
+                FROM pg_indexes
+                WHERE tablename = %s
+                """,
+                [self._name]
+            )
+            return {row["indexname"]: row["indexdef"] for row in rows}
+        except Exception as e:
+            raise CollectionError(f"Failed to get index params: {e}")
+    
+    # =========================================================================
+    # Connection Management
+    # =========================================================================
+    
+    def _get_connection(self) -> 'AsyncVastbaseConnection':
+        """Get database connection"""
+        if self._conn is None:
+            if self._connection_provider:
+                self._conn = self._connection_provider.get_connection(self._using)
+            else:
+                from .connections import AsyncConnections
+                self._conn = AsyncConnections.get_connection(self._using)
+        return self._conn
+    
+    async def _ensure_connection(self) -> None:
+        """Ensure connection is established"""
+        self._conn = self._get_connection()
+        if not self._conn.is_connected:
+            await self._conn.connect()
+        if self._executor is None:
+            self._executor = AsyncExecutor(self._conn)
+
+    @property
+    def _core(self) -> CollectionCore:
+        """Lazy-initialized CollectionCore with the AsyncExecutor.
+
+        Used for delegating business logic to the shared CollectionCore layer.
+        Schema must be loaded before this property is accessed.
+        """
+        if self._core_obj is None:
+            self._core_obj = CollectionCore(
+                self._schema, self._executor, name=self._name
+            )
+        return self._core_obj
+
+    async def _ensure_schema_loaded(self) -> None:
+        """Ensure schema is loaded before delegating to core."""
+        if self._schema is None:
+            await self._load_schema_async()
+
+    def _get_core(self) -> CollectionCore:
+        """Get CollectionCore with FakeExecutor for validation-only operations.
+
+        Unlike ``self._core`` (which uses the real AsyncExecutor), this method
+        returns a core backed by a FakeExecutor, safe for calling
+        ``_process_vectors_for_insert`` / ``_validate_insert_data`` without
+        executing real SQL.
+        """
+        if self._schema is None:
+            raise CollectionError(
+                "Schema not loaded. Use 'await collection.schema_async' or "
+                "call 'await collection.load_schema()' first."
+            )
+        from ..core.executor import FakeExecutor
+
+        return CollectionCore(self._schema, FakeExecutor(), name=self._name)
+    
+    # =========================================================================
+    # Schema Operations
+    # =========================================================================
+    
+    @_with_connection
+    async def load_schema(self) -> CollectionSchema:
+        """Load schema from database"""
+        await self._load_schema_async()
+        return self._schema
+
+    async def _load_schema_async(self) -> None:
+        """Load schema from database (async)"""
+        try:
+            # Check existence first
+            check_sql = build_existence_check_sql()
+            rows = await self._executor.execute(check_sql, [self._name])
+            if not rows:
+                raise CollectionNotExistsError(self._name)
+
+            # Primary keys
+            pk_rows = await self._executor.execute(
+                build_pk_query_sql(), [self._name]
+            )
+            pk_columns = {row["column_name"] for row in pk_rows}
+
+            # Columns with atttypmod dimensions
+            field_rows = await self._executor.execute(
+                build_columns_query_sql(), [self._name]
+            )
+
+            fields = [parse_column_row(row, pk_columns) for row in field_rows]
+            self._schema = CollectionSchema(name=self._name, fields=fields)
+
+        except Exception as e:
+            if isinstance(e, CollectionNotExistsError):
+                raise
+            raise CollectionError(f"Failed to load schema: {e}")
+
+    async def _check_schema_version(self, conn) -> None:
+        """Verify server version supports all vector types in the schema.
+
+        Uses VectorTypeRegistry to look up the minimum Vastbase version
+        required for each vector type in the schema.
+
+        The ``conn`` argument is the AsyncVastbaseConnection, which provides
+        ``require_version``.
+        """
+        from ..core.vector_types import get_registry
+
+        registry = get_registry()
+        checked: set = set()
+        for f in self._schema.fields:
+            if f.is_vector_field():
+                min_ver = registry.get_min_version(f.dtype)
+                if min_ver is not None and min_ver not in checked:
+                    checked.add(min_ver)
+                    ver_str = ".".join(str(v) for v in min_ver)
+                    await conn.require_version(
+                        (min_ver[0], min_ver[1]),
+                        min_ver[2],
+                        min_ver[3],
+                        f"Vector type '{f.dtype}' requires Vastbase {ver_str}",
+                    )
+
+    # =========================================================================
+    # DDL Operations
+    # =========================================================================
+    
+    @_with_connection
+    async def create(
+        self,
+        timeout: Optional[float] = None,
+        **kwargs
+    ) -> None:
+        """
+        Create the collection in the database.
+
+        Args:
+            timeout: Operation timeout in seconds
+        """
+        if self._schema is None:
+            raise SchemaError("Schema is required to create collection")
+
+        # Version-gating: check if schema uses newer vector types
+        await self._check_schema_version(self._conn)
+
+        try:
+            sql = self._schema.to_sql(if_not_exists=True)
+            await self._executor.execute(sql, [])
+        except Exception as e:
+            raise CollectionError(f"Failed to create collection: {e}")
+
+    @_with_connection
+    async def drop(self, timeout: Optional[float] = None, **kwargs) -> None:
+        """Drop the collection from the database."""
+        from ..operations.ddl import build_drop_table_sql
+        try:
+            sql, params = build_drop_table_sql(self._name)
+            await self._executor.execute(sql, params)
+        except Exception as e:
+            raise CollectionError(f"Failed to drop collection: {e}")
+
+    @_with_connection
+    async def exists(self) -> bool:
+        """Check if collection exists"""
+        try:
+            from ..operations.schema_reflect import build_existence_check_sql
+            sql = build_existence_check_sql()
+            rows = await self._executor.execute(sql, {"table_name": self._name})
+            return len(rows) > 0
+        except Exception:
+            return False
+
+    @_with_connection
+    async def truncate(self, timeout: Optional[float] = None, **kwargs) -> None:
+        """Truncate (empty) the collection."""
+        from ..operations.ddl import build_truncate_table_sql
+        try:
+            sql, params = build_truncate_table_sql(self._name)
+            await self._executor.execute(sql, params)
+        except Exception as e:
+            raise CollectionError(f"Failed to truncate collection: {e}")
+
+    @_with_connection
+    async def refresh_collection(self, timeout: Optional[float] = None) -> None:
+        """Refresh collection statistics (ANALYZE)."""
+        from ..operations.ddl import build_refresh_collection_sql
+        try:
+            sql, params = build_refresh_collection_sql(self._name)
+            await self._executor.execute(sql, params)
+        except Exception as e:
+            raise CollectionError(f"Failed to refresh collection: {e}")
+
+    @_with_connection
+    async def has_collection_revision(
+        self,
+        revision: Optional[int] = None
+    ) -> Union[bool, Dict[str, Any]]:
+        """Check collection revision or get revision info."""
+        from ..operations.ddl import build_has_revision_sql, parse_has_revision_result
+        try:
+            sql, params = build_has_revision_sql(self._name)
+            rows = await self._executor.execute(sql, params)
+            result = parse_has_revision_result(rows, revision)
+            if isinstance(result, dict):
+                result["collection_name"] = self._name
+            return result
+
+        except Exception as e:
+            if revision is not None:
+                return False
+            raise CollectionError(f"Failed to get collection revision: {e}")
+    
+    # =========================================================================
+    # Index Operations
+    # =========================================================================
+
+    @_with_connection
+    async def create_index(
+        self,
+        field_name: str,
+        index_params: Optional[IndexParams] = None,
+        timeout: Optional[float] = None,
+        **kwargs
+    ) -> None:
+        """
+        Create index on a field.
+
+        Args:
+            field_name: Field name to index
+            index_params: Index parameters
+            timeout: Operation timeout in seconds
+        """
+        if index_params is None:
+            index_params = IndexParams()
+
+        # Get field info (load schema if needed)
+        if self._schema is None:
+            await self._load_schema_async()
+
+        field = self._schema.get_field(field_name)
+        if field is None:
+            raise CollectionError(f"Field '{field_name}' not found")
+
+        # Generate index name
+        index_name = f"idx_{self._name}_{field_name}"
+
+        # Build index using builder
+        builder = IndexBuilder(self._name)
+        builder.set_index_name(index_name)
+        builder.set_column(field_name)
+        builder.set_params(index_params)
+
+        try:
+            sql = builder.build()
+            await self._executor.execute(sql, [])
+            return index_name
+        except Exception as e:
+            raise CollectionError(f"Failed to create index: {e}")
+
+    @_with_connection
+    async def drop_index(
+        self,
+        field_name: str,
+        timeout: Optional[float] = None,
+        **kwargs
+    ) -> None:
+        """Drop index on a field."""
+        from ..operations.index import build_drop_index_sql
+        index_name = f"idx_{self._name}_{field_name}"
+        try:
+            sql, params = build_drop_index_sql(index_name)
+            await self._executor.execute(sql, params)
+        except Exception as e:
+            raise CollectionError(f"Failed to drop index: {e}")
+
+    @_with_connection
+    async def has_index(self, field_name: Optional[str] = None) -> bool:
+        """Check if index exists on field."""
+        from ..operations.index import build_has_index_sql
+        try:
+            index_name = f"idx_{self._name}_{field_name}" if field_name else None
+            sql, params = build_has_index_sql(self._name, index_name)
+            rows = await self._executor.execute(sql, params)
+            return len(rows) > 0
+        except Exception:
+            return False
+
+    @_with_connection
+    async def list_indexes(self) -> List[str]:
+        """List all indexes of the collection."""
+        from ..operations.index import build_list_indexes_sql, parse_list_indexes_result
+        try:
+            sql, params = build_list_indexes_sql(self._name)
+            rows = await self._executor.execute(sql, params)
+            return parse_list_indexes_result(rows)
+        except Exception as e:
+            raise CollectionError(f"Failed to list indexes: {e}")
+
+    @_with_connection
+    async def get_index(self, field_name: Optional[str] = None) -> Dict[str, Any]:
+        """Get index information."""
+        from ..operations.index import build_get_index_sql, parse_get_index_result
+        try:
+            sql, params = build_get_index_sql(self._name, field_name)
+            rows = await self._executor.execute(sql, params)
+            return parse_get_index_result(rows, self._name)
+        except Exception as e:
+            raise CollectionError(f"Failed to get index: {e}")
+    
+    # =========================================================================
+    # Data Operations - Helpers
+    # =========================================================================
+    # Data Operations - Insert
+    # =========================================================================
+    
+    @_with_connection(schema=True)
+    async def insert(
+        self,
+        data: Union[List[Dict[str, Any]], Dict[str, Any]],
+        partition_name: Optional[str] = None,
+        timeout: Optional[float] = None,
+        normalize_vectors: bool = False,
+        **kwargs
+    ) -> MutationResult:
+        """
+        Insert data into collection.
+
+        Args:
+            data: List of records or single record
+            timeout: Operation timeout in seconds
+            normalize_vectors: Whether to L2-normalize vector fields before insert
+
+        Returns:
+            MutationResult with inserted IDs (backward compatible with List[int])
+
+        Example:
+            >>> result = await collection.insert([{"id": 1, "embedding": [0.1] * 128}])
+            >>> # Access as list (backward compatible)
+            >>> ids = list(result)
+            >>> # Access as MutationResult
+            >>> print(result.insert_count)
+        """
+        # Normalize to list
+        if isinstance(data, dict):
+            data = [data]
+
+        if not data:
+            return MutationResult(primary_keys=[])
+
+        # Deep copy to avoid mutating original data
+        data = json.loads(json.dumps(data))
+
+        # Process vector fields: validate and optionally normalize
+        core = self._get_core()
+        core._process_vectors_for_insert(data, normalize=normalize_vectors)
+
+        # Validate against schema
+        core._validate_insert_data(data)
+
+        # Build insert statement
+        columns = list(data[0].keys())
+        from ..operations.insert import build_insert_col_names, build_insert_row_sql
+        from ..operations.result import parse_insert_result
+        from ..operations.insert import build_value_parts
+        col_names = build_insert_col_names(columns)
+        pk_field = self._schema.get_primary_field()
+        field_map = {f.name: f for f in self._schema.fields}
+
+        try:
+            ids = []
+            for row in data:
+                value_parts, params = build_value_parts(row, columns, field_map)
+                sql = build_insert_row_sql(self._name, col_names, pk_field.name, value_parts)
+                rows = await self._executor.execute(sql, params)
+                pk_col = pk_field.name
+                row_ids = [r[pk_col] for r in rows]
+                ids.extend(row_ids)
+
+            return parse_insert_result(ids)
+        except Exception as e:
+            await self._conn.rollback()
+            raise DataError(f"Failed to insert data: {e}")
+
+    @_with_connection(schema=True)
+    async def batch_insert(
+        self,
+        data: Union[List[Dict[str, Any]], Dict[str, Any]],
+        batch_size: int = 100,
+        timeout: Optional[float] = None,
+        normalize_vectors: bool = False,
+        **kwargs
+    ) -> MutationResult:
+        """
+        Insert data in batches for better performance with large datasets.
+
+        Args:
+            data: List of records or single record
+            batch_size: Number of records per batch
+            timeout: Operation timeout in seconds
+            normalize_vectors: Whether to L2-normalize vector fields before insert
+
+        Returns:
+            MutationResult with all inserted IDs (backward compatible with List[int])
+
+        Example:
+            >>> result = await collection.batch_insert(large_data, batch_size=1000)
+            >>> print(result.insert_count)
+            >>> print(result.primary_keys[:5])
+        """
+        # Normalize to list
+        if isinstance(data, dict):
+            data = [data]
+
+        if not data:
+            return MutationResult(primary_keys=[])
+
+        all_ids = []
+
+        # Use batch_iterator utility to process data in batches
+        for batch in batch_iterator(data, batch_size):
+            # Process vectors (validate, normalize)
+            batch_copy = json.loads(json.dumps(list(batch)))
+            core = self._get_core()
+            core._process_vectors_for_insert(batch_copy, normalize=normalize_vectors)
+
+            # Validate batch
+            core._validate_insert_data(batch_copy)
+
+            # Build batch insert SQL
+            columns = list(batch_copy[0].keys())
+            from ..operations.insert import build_insert_col_names, build_insert_row_sql
+            from ..operations.insert import build_value_parts
+            col_names = build_insert_col_names(columns)
+            field_map = {f.name: f for f in self._schema.fields}
+            pk_field = self._schema.get_primary_field()
+
+            try:
+                batch_ids = []
+                for row in batch_copy:
+                    value_parts, params = build_value_parts(row, columns, field_map)
+                    sql = build_insert_row_sql(self._name, col_names, pk_field.name, value_parts)
+                    rows = await self._executor.execute(sql, params)
+                    pk_col = pk_field.name
+                    batch_ids.extend(r[pk_col] for r in rows)
+                all_ids.extend(batch_ids)
+            except Exception as e:
+                await self._conn.rollback()
+                raise DataError(f"Batch insert failed: {e}")
+
+        from ..operations.result import parse_insert_result
+        return parse_insert_result(all_ids)
+    
+    # =========================================================================
+    # Data Operations - Delete/Upsert
+    # =========================================================================
+    
+    @_with_connection(schema=True)
+    async def delete(
+        self,
+        expr: Optional[str] = None,
+        pks: Optional[List[Any]] = None,
+        partition_name: Optional[str] = None,
+        timeout: Optional[float] = None,
+        **kwargs
+    ) -> MutationResult:
+        """
+        Delete entities from collection.
+
+        Args:
+            expr: Filter expression (SQL WHERE clause)
+            pks: List of primary key values to delete
+            timeout: Operation timeout in seconds
+
+        Returns:
+            MutationResult with delete count (backward compatible with int)
+
+        Example:
+            >>> result = await collection.delete(pks=[1, 2, 3])
+            >>> # Access as int (backward compatible)
+            >>> count = len(result)
+            >>> # Access as MutationResult
+            >>> print(result.delete_count)
+        """
+        conn = self._get_connection()
+        await self._ensure_connection()
+
+        if self._schema is None:
+            await self._load_schema_async()
+
+        try:
+            select_sql, delete_sql, params = build_delete_sql(
+                self._name,
+                self._schema.get_primary_field().name,
+                expr=expr if expr else None,
+                pks=pks if pks else None,
+            )
+        except ValueError as e:
+            raise ParamError(str(e))
+
+        try:
+            from ..operations.result import parse_delete_result
+            if select_sql:
+                pk_name = self._schema.get_primary_field().name
+                rows = await self._executor.execute(select_sql, [])
+                deleted_pks = [row[pk_name] for row in rows]
+            else:
+                deleted_pks = list(pks)
+
+            await self._executor.execute(delete_sql, params if params else [])
+            return parse_delete_result(deleted_pks)
+        except Exception as e:
+            await self._conn.rollback()
+            raise DataError(f"Failed to delete data: {e}")
+
+    @_with_connection(schema=True)
+    async def upsert(
+        self,
+        data: Union[List[Dict[str, Any]], Dict[str, Any]],
+        partition_name: Optional[str] = None,
+        timeout: Optional[float] = None,
+        **kwargs
+    ) -> MutationResult:
+        """
+        Upsert data (insert or update on conflict).
+        
+        Args:
+            data: Data to upsert
+            timeout: Operation timeout in seconds
+            
+        Returns:
+            MutationResult with upserted IDs (backward compatible with List[int])
+            
+        Example:
+            >>> result = await collection.upsert(data)
+            >>> print(result.insert_count)
+            >>> print(result.primary_keys[:5])
+        """
+        if isinstance(data, dict):
+            data = [data]
+
+        if not data:
+            return MutationResult(primary_keys=[])
+
+        pk_field = self._schema.get_primary_field()
+        if not pk_field:
+            raise CollectionError("Upsert requires a primary key")
+
+        # Validate data
+        self._get_core()._validate_insert_data(data)
+
+        columns = list(data[0].keys())
+        from ..operations.upsert import build_upsert_clauses
+        from ..operations.insert import build_insert_col_names, build_upsert_row_sql
+        from ..operations.insert import build_value_parts
+        from ..operations.result import parse_upsert_result
+        col_names, on_conflict = build_upsert_clauses(columns, pk_field.name)
+        field_map = {f.name: f for f in self._schema.fields}
+
+        try:
+            ids = []
+            for row in data:
+                value_parts, params = build_value_parts(row, columns, field_map)
+                sql = build_upsert_row_sql(self._name, col_names, value_parts, on_conflict)
+                rows = await self._executor.execute(sql, params)
+                pk_col = pk_field.name
+                ids.extend(r[pk_col] for r in rows)
+
+            return parse_upsert_result(ids)
+        except Exception as e:
+            await self._conn.rollback()
+            raise DataError(f"Failed to upsert data: {e}")
+
+    # =========================================================================
+    # Search Operations
+    # =========================================================================
+    
+    @_with_connection(schema=True)
+    async def search(
+        self,
+        data: Union[List[float], List[List[float]]],
+        anns_field: Optional[str] = None,
+        param: Optional[Dict[str, Any]] = None,
+        limit: int = 10,
+        expr: Optional[str] = None,
+        partition_names: Optional[List[str]] = None,
+        output_fields: Optional[List[str]] = None,
+        timeout: Optional[float] = None,
+        round_decimal: int = -1,
+        ranker: Optional[Any] = None,
+        highlighter: Optional[Any] = None,
+        normalize_vectors: bool = False,
+        **kwargs
+    ) -> SearchResult:
+        """
+        Search for similar vectors.
+
+        Args:
+            data: Query vector(s)
+            anns_field: Name of vector field to search
+            param: Search parameters
+            limit: Number of results to return
+            expr: Filter expression
+            output_fields: Fields to return in results
+            timeout: Operation timeout in seconds
+            normalize_vectors: Whether to normalize query vectors before search
+            
+        Returns:
+            SearchResult object
+        """
+        if ranker is not None:
+            raise ParamError("ranker is not supported by Vastbase (PostgreSQL-based). "
+                             "Vastbase does not support custom ranking strategies.")
+        if highlighter is not None:
+            raise ParamError("highlighter is not supported by Vastbase. "
+                             "Use bm25_search_highlight utility for BM25 highlighting.")
+
+        # Get vector field
+        if anns_field is None:
+            vector_fields = self._schema.vector_field_names
+            if not vector_fields:
+                raise CollectionError("No vector field found in schema")
+            anns_field = vector_fields[0]
+        
+        # Get field info for dimension validation
+        field = self._schema.get_field(anns_field)
+        if field is None:
+            raise CollectionError(f"Vector field '{anns_field}' not found")
+        
+        expected_dim = field.dim
+        
+        # Get metric type from schema
+        metric_type = field.metric_type if field and field.metric_type else DistanceType.L2
+        
+        # Normalize queries (ensure correct dimension)
+        queries = data if isinstance(data[0], list) else [data]
+        
+        processed_queries = validate_and_normalize_queries(
+            queries,
+            expected_dim,
+            normalize=normalize_vectors,
+        )
+
+        # Build search parameters
+        search_params = SearchParams(
+            anns_field=anns_field,
+            top_k=limit,
+            metric_type=metric_type
+        )
+
+        if param:
+            if "ef" in param:
+                search_params.hnsw_ef_search = param["ef"]
+            if "metric_type" in param:
+                search_params.metric_type = DistanceType(param["metric_type"])
+
+        # Build query using SearchQueryBuilder
+        builder = SearchQueryBuilder(self._name)
+        builder.set_vector_field(anns_field)
+        builder.set_queries(processed_queries)
+        builder.set_params(search_params)
+        builder.set_limit(limit)
+
+        if expr:
+            builder.set_filter(expr)
+
+        if output_fields:
+            builder.set_output_fields(output_fields)
+
+        # Build and execute query
+        build_result = builder.build()
+
+        try:
+            if isinstance(build_result, tuple) and len(build_result) == 3:
+                set_clauses, main_sql, params = build_result
+                async with self._executor.transaction() as tx:
+                    for set_sql in set_clauses:
+                        await tx.execute(set_sql, [])
+                    rows = await tx.execute(
+                        main_sql, list(params[0]) if params else []
+                    )
+            else:
+                sql, params = build_result
+                if params:
+                    rows = await self._executor.execute(sql, list(params[0]))
+                else:
+                    rows = await self._executor.execute(sql, [])
+
+            # Parse results
+            from ..operations.search_result import parse_search_rows
+            pk_field = self._schema.get_primary_field()
+            query_results = parse_search_rows(rows, pk_field.name if pk_field else None, output_fields)
+
+            results = SearchResult()
+            results.num_queries = len(processed_queries)
+            results.queries = processed_queries
+            results.add_result(0, query_results)
+            return results
+
+        except Exception as e:
+            raise CollectionError(f"Search failed: {e}")
+    
+    async def hybrid_search(
+        self,
+        reqs: Optional[List[Any]] = None,
+        rerank: Optional[Any] = None,
+        limit: int = 10,
+        partition_names: Optional[List[str]] = None,
+        output_fields: Optional[List[str]] = None,
+        timeout: Optional[float] = None,
+        round_decimal: int = -1,
+        ranker: Optional[Any] = None,
+        parallel: bool = True,
+        **kwargs
+    ) -> SearchResult:
+        """
+        Multi-way vector recall with reranking (pymilvus-compatible, async).
+
+        Issues independent ANN search SQL per AnnSearchRequest, merges
+        results using the specified reranking strategy.
+
+        Args:
+            reqs: List of AnnSearchRequest objects
+            rerank: BaseRanker instance (RRFRanker, WeightedRRFRanker, WeightedRanker)
+            limit: Max merged results
+            output_fields: Fields to include
+            parallel: Execute searches concurrently via asyncio.gather (default True)
+        """
+        from ..core.exceptions import ParamError
+        from ..search.search_builder import BaseRanker
+
+        if not reqs:
+            raise ParamError("reqs must not be empty")
+        if not isinstance(rerank, BaseRanker):
+            raise ParamError(f"rerank must be a BaseRanker instance, got {type(rerank).__name__}")
+
+        pk_field = self._schema.get_primary_field()
+        all_results: list[tuple[list, DistanceType]] = []
+
+        async def _execute_single_req(req):
+            field = self._schema.get_field(req.anns_field)
+            if field is None:
+                raise ParamError(f"Field '{req.anns_field}' not found in schema")
+
+            metric_str = req.param.get("metric_type", "l2")
+            try:
+                metric_type = DistanceType(metric_str.lower())
+            except ValueError:
+                raise ParamError(f"Unknown metric_type '{metric_str}'")
+
+            params = SearchParams(
+                anns_field=req.anns_field,
+                top_k=req.limit,
+                metric_type=metric_type,
+            )
+            if "ef" in req.param:
+                params.hnsw_ef_search = req.param["ef"]
+
+            query_vec = req.data if req.data and isinstance(req.data[0], list) else [req.data]
+
+            builder = SearchQueryBuilder(table_name=self._name)
+            builder.set_vector_field(req.anns_field)
+            builder.set_queries(query_vec)
+            builder.set_params(params)
+            builder.set_limit(req.limit)
+            if output_fields:
+                builder.set_output_fields(output_fields)
+
+            build_result = builder.build()
+            if isinstance(build_result, tuple) and len(build_result) == 3:
+                set_clauses, main_sql, sql_params = build_result
+                async with self._executor.transaction() as tx:
+                    for set_sql in set_clauses:
+                        await tx.execute(set_sql, [])
+                    rows = await tx.execute(
+                        main_sql, list(sql_params[0]) if sql_params else []
+                    )
+            else:
+                sql, sql_params = build_result
+                rows = await self._executor.execute(sql, list(sql_params[0]) if sql_params else [])
+            items = parse_search_rows(rows, pk_field.name if pk_field else None, output_fields)
+            return items, metric_type
+
+        if parallel and len(reqs) > 1:
+            import asyncio
+            pairs = await asyncio.gather(*[_execute_single_req(req) for req in reqs])
+            all_results = list(pairs)
+        else:
+            for req in reqs:
+                items, metric_type = await _execute_single_req(req)
+                all_results.append((items, metric_type))
+
+        merged = rerank.merge(all_results, limit)
+
+        results = SearchResult()
+        results.num_queries = 1
+        results.add_result(0, merged)
+        return results
+
+    async def hybrid_ann_search(
+        self,
+        data: Union[List[float], List[List[float]]],
+        anns_field: Optional[str] = None,
+        param: Optional[Dict[str, Any]] = None,
+        limit: int = 10,
+        filter: Optional[Union[Dict[str, Any], str]] = None,
+        output_fields: Optional[List[str]] = None,
+        timeout: Optional[float] = None,
+        normalize_vectors: bool = False,
+        **kwargs
+    ) -> SearchResult:
+        """
+        Hybrid ANN search: vector similarity search with scalar field filtering.
+
+        Vastbase-native vector + scalar hybrid search using HybridANN index.
+
+        Args:
+            data: Query vector(s) - single vector or list of vectors
+            anns_field: Name of vector field to search (auto-detected if None)
+            param: Search parameters (ef, metric_type, etc.)
+            limit: Maximum number of results
+            filter: Filter expression - can be:
+                - dict: MongoDB-style filter {"field": {"$op": value}}
+                - str: Raw SQL WHERE clause
+            output_fields: Fields to return in results
+            timeout: Operation timeout in seconds
+            normalize_vectors: Whether to L2-normalize query vectors
+
+        Returns:
+            SearchResult object
+
+        Supported Filter Operators:
+            - Comparison: $eq, $ne, $gt, $gte, $lt, $lte
+            - Set: $in, $nin
+            - Range: $between
+            - String: $like, $ilike
+            - Logical: $and, $or
+
+        Example:
+            # Simple filter
+            results = await collection.hybrid_search(
+                data=[0.1] * 128,
+                filter={"category": {"$eq": "electronics"}},
+                limit=10
+            )
+
+            # Complex filter with AND/OR
+            results = await collection.hybrid_search(
+                data=[0.1] * 128,
+                filter={
+                    "$and": [
+                        {"price": {"$lte": 1000}},
+                        {"category": {"$in": ["electronics", "books"]}}
+                    ]
+                },
+                limit=20
+            )
+        """
+        from ..utils.filter_utils import parse_filter_expression
+
+        # Get vector field
+        if anns_field is None:
+            vector_fields = self._schema.vector_field_names
+            if not vector_fields:
+                raise CollectionError("No vector field found in schema")
+            anns_field = vector_fields[0]
+
+        # Get field info for dimension validation
+        field = self._schema.get_field(anns_field)
+        if field is None:
+            raise CollectionError(f"Vector field '{anns_field}' not found")
+
+        expected_dim = field.dim
+
+        # Get metric type from schema
+        metric_type = field.metric_type if field and field.metric_type else DistanceType.L2
+
+        # Normalize queries (ensure correct dimension)
+        queries = data if isinstance(data[0], list) else [data]
+
+        processed_queries = validate_and_normalize_queries(
+            queries,
+            expected_dim,
+            normalize=normalize_vectors,
+        )
+
+        # Parse filter expression to SQL
+        where_clause = ""
+        filter_params = []
+
+        if filter is not None:
+            where_clause, filter_params = parse_filter_expression(
+                filter, param_style="dollar"
+            )
+
+        # Build search parameters
+        search_params = SearchParams(
+            anns_field=anns_field,
+            top_k=limit,
+            metric_type=metric_type
+        )
+
+        if param:
+            if "ef" in param:
+                search_params.hnsw_ef_search = param["ef"]
+            if "metric_type" in param:
+                search_params.metric_type = DistanceType(param["metric_type"])
+
+        # Build query using SearchQueryBuilder
+        builder = SearchQueryBuilder(self._name)
+        builder.set_vector_field(anns_field)
+        builder.set_queries(processed_queries)
+        builder.set_limit(limit)
+        builder.set_params(search_params)
+
+        if where_clause:
+            builder.set_filter(where_clause)
+
+        if output_fields:
+            builder.set_output_fields(output_fields)
+
+        # Build and execute query
+        build_result = builder.build()
+
+        try:
+            if isinstance(build_result, tuple) and len(build_result) == 3:
+                set_clauses, main_sql, params = build_result
+                final_params = filter_params + (params[0] if params else [])
+                async with self._executor.transaction() as tx:
+                    for set_sql in set_clauses:
+                        await tx.execute(set_sql, [])
+                    rows = await tx.execute(
+                        main_sql, list(final_params) if final_params else []
+                    )
+            else:
+                sql, params = build_result
+                final_params = filter_params + (params[0] if params else [])
+                if final_params:
+                    rows = await self._executor.execute(sql, list(final_params))
+                else:
+                    rows = await self._executor.execute(sql, [])
+
+            # Parse results
+            from ..operations.search_result import parse_search_rows
+            pk_field = self._schema.get_primary_field()
+            query_results = parse_search_rows(rows, pk_field.name if pk_field else None, output_fields)
+
+            results = SearchResult()
+            results.num_queries = len(processed_queries)
+            results.queries = processed_queries
+            results.add_result(0, query_results)
+            return results
+
+        except Exception as e:
+            raise CollectionError(f"Hybrid search failed: {e}")
+    
+    async def _fetch_batch(
+        self,
+        filter_expr: str,
+        output_fields: Optional[List[str]] = None,
+        limit: Optional[int] = None,
+        offset: int = 0
+    ) -> List[Dict[str, Any]]:
+        """Internal method to fetch a batch of results"""
+        await self._ensure_connection()
+        from ..operations.query import build_query_sql, parse_query_result
+        sql, params = build_query_sql(self._name, filter_expr, output_fields, limit, offset if offset > 0 else None)
+        try:
+            rows = await self._executor.execute(sql, params)
+            return parse_query_result(rows, output_fields)
+        except Exception as e:
+            raise CollectionError(f"Query failed: {e}")
+    
+    @_with_connection(schema=True)
+    async def query(
+        self,
+        expr: Optional[str] = None,
+        output_fields: Optional[List[str]] = None,
+        limit: Optional[int] = None,
+        offset: Optional[int] = None,
+        partition_names: Optional[List[str]] = None,
+        timeout: Optional[float] = None,
+        batch_size: Optional[int] = None,
+        **kwargs
+    ) -> Union[List[Dict[str, Any]], QueryIterator]:
+        """
+        Query entities with filters.
+
+        Args:
+            expr: Filter expression (SQL WHERE clause)
+            output_fields: Fields to return
+            limit: Maximum results
+            offset: Offset for pagination
+            timeout: Operation timeout in seconds
+            batch_size: If specified, returns an async iterator that yields batches
+
+        Returns:
+            List of matching records, or QueryIterator if batch_size specified
+        """
+        # If batch_size specified, return async iterator
+        if batch_size is not None and batch_size > 0:
+            return QueryIterator(
+                collection=self,
+                filter_expr=expr or "",
+                batch_size=batch_size,
+                output_fields=output_fields,
+                limit=limit,
+                offset=offset
+            )
+
+        # Build and execute query
+        from ..operations.query import build_query_sql, parse_query_result
+        sql, params = build_query_sql(self._name, expr, output_fields, limit, offset)
+        try:
+            rows = await self._executor.execute(sql, params)
+            return parse_query_result(rows, output_fields)
+        except Exception as e:
+            raise CollectionError(f"Query failed: {e}")
+
+    @_with_connection(schema=True)
+    async def query_iterator(
+        self,
+        expr: Optional[str] = None,
+        output_fields: Optional[List[str]] = None,
+        limit: Optional[int] = None,
+        offset: Optional[int] = None,
+        batch_size: int = 1000,
+        partition_names: Optional[List[str]] = None,
+        timeout: Optional[float] = None,
+        **kwargs
+    ) -> QueryIterator:
+        """Return a QueryIterator for paginated result access."""
+        results = await self._fetch_batch(
+            filter_expr=expr or "",
+            output_fields=output_fields,
+            limit=limit,
+            offset=offset or 0,
+        )
+        return QueryIterator(
+            data=results,
+            batch_size=batch_size,
+            output_fields=output_fields,
+        )
+
+    @_with_connection(schema=True)
+    async def get(
+        self,
+        ids: Union[List[Any], Any],
+        output_fields: Optional[List[str]] = None,
+        partition_names: Optional[List[str]] = None,
+        timeout: Optional[float] = None,
+        **kwargs
+    ) -> List[Dict[str, Any]]:
+        """
+        Get entities by primary key.
+
+        Args:
+            ids: Primary key value(s)
+            output_fields: Fields to return
+            timeout: Operation timeout in seconds
+
+        Returns:
+            List of matching records
+        """
+        if not isinstance(ids, list):
+            ids = [ids]
+
+        pk_field = self._schema.get_primary_field()
+        if not pk_field:
+            raise CollectionError("Get requires a primary key")
+
+        from ..operations.query import build_get_sql, parse_query_result
+        sql, params = build_get_sql(self._name, pk_field.name, ids, output_fields)
+        try:
+            rows = await self._executor.execute(sql, params)
+            return parse_query_result(rows, output_fields)
+        except Exception as e:
+            raise CollectionError(f"Failed to get entities: {e}")
+
+    async def get_entity_by_id(
+        self,
+        ids: Union[List[Any], Any],
+        output_fields: Optional[List[str]] = None,
+        timeout: Optional[float] = None,
+        **kwargs
+    ) -> List[Dict[str, Any]]:
+        """
+        Get entities by primary key (PyMilvus API compatibility).
+
+        Deprecated: use get() instead.
+
+        Args:
+            ids: Primary key value(s)
+            output_fields: Fields to return
+            timeout: Operation timeout in seconds
+
+        Returns:
+            List of matching records
+        """
+        import warnings
+        warnings.warn(
+            "get_entity_by_id() is deprecated, use get() instead",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return await self.get(ids=ids, output_fields=output_fields, timeout=timeout, **kwargs)
+