feat: Add QueryHint support for database optimization (#177)

Add QueryHint support for OceanBase database optimization with parallel
execution and query timeout hints.

## Key Features:
- QueryHint dataclass with parallel and query_timeout parameters
- SQL hint generation for OceanBase database optimization
- Added to collection.get(), query(), and hybrid_search() methods
- Comprehensive test coverage across all database modes
- Full documentation with usage examples

## Usage:
from pyseekdb.client.query_types import QueryHint
query_hint = QueryHint(parallel=4, query_timeout=5.0)
collection.get(ids=['1', '2'], query_hint=query_hint)

Made with [Cursor](https://cursor.com)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Introduced QueryHint as a client-side option to tune
collection.query(), get(), and hybrid_search() (parallelism, timeout,
vector index control). When provided, hint text is integrated into
generated queries and can affect execution plans.
* **Documentation**
  * Docstrings and examples updated to show QueryHint usage.
* **Tests**
* Added integration tests covering QueryHint scenarios for query, get,
and hybrid_search.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: hnwyllmm

src/pyseekdb/client/__init__.py

@@ -46,6 +46,7 @@
     get_default_embedding_function,
     register_embedding_function,
 )
+from .query_types import QueryHint
 from .schema import Schema
 from .sparse_embedding_function import (
     SparseEmbeddingFunction,
@@ -172,6 +173,7 @@ def __getattr__(name: str) -> Any:
     "K",
     "Ngram2FulltextIndexConfig",
     "NgramFulltextIndexConfig",
+    "QueryHint",
     "RemoteServerClient",
     "Schema",
     "SeekdbEmbeddedClient",

src/pyseekdb/client/client_base.py

@@ -38,14 +38,15 @@
 )
 from .filters import FilterBuilder
 from .meta_info import CollectionFieldNames, CollectionNames
+from .query_types import QueryHint
 from .schema import Schema, SparseVectorIndexConfig
 from .sparse_embedding_function import (
     SparseEmbeddingFunction,
     SparseEmbeddingFunctionRegistry,
     SparseVector,
     _sparse_vector_to_sql,
 )
-from .sql_utils import is_query_sql
+from .sql_utils import _query_hint_to_sql, is_query_sql
 from .types import K as FieldKey
 from .version import Version
 
@@ -2588,6 +2589,7 @@ def _collection_query(  # noqa: C901
         where_document: dict[str, Any] | None = None,
         include: list[str] | None = None,
         query_key: FieldKey | None = None,
+        query_hint: QueryHint | None = None,
         **kwargs,
     ) -> dict[str, Any]:
         """
@@ -2647,6 +2649,7 @@ def _collection_query(  # noqa: C901
                 include=include,
                 sparse_config=sparse_config,
                 collection_name=collection_name,
+                query_hint=query_hint,
                 **kwargs,
             )
 
@@ -2725,12 +2728,15 @@ def _collection_query(  # noqa: C901
             # Convert vector to string format for SQL
             vector_str = _embedding_to_hexstring(query_vector)
 
+            # Build query hint
+            hint_sql = _query_hint_to_sql(query_hint, table_name=table_name)
+
             # Build SQL query with vector distance calculation
             # Reference: SELECT id, vec FROM t2 ORDER BY l2_distance(vec, '[0.1, 0.2, 0.3]') APPROXIMATE LIMIT 5;
             # Need to include distance in SELECT for result processing
             # Use the appropriate distance function based on the index configuration
             sql = f"""
-                SELECT {select_clause},
+                SELECT {hint_sql} {select_clause},
                        {distance_func}(embedding, {vector_str}) AS distance
                 FROM `{table_name}`
                 {where_clause}
@@ -2806,6 +2812,7 @@ def _collection_query_sparse(  # noqa: C901
         include: list[str] | None = None,
         sparse_config=None,
         collection_name: str = "",
+        query_hint=None,
         **kwargs,
     ) -> dict[str, Any]:
         """
@@ -2858,6 +2865,8 @@ def _collection_query_sparse(  # noqa: C901
         if not sparse_query_vectors:
             raise ValueError("No sparse query vectors resolved.")
 
+        hint_sql = _query_hint_to_sql(query_hint, table_name=table_name)
+
         # Normalize include fields
         include_fields = self._normalize_include_fields(include)
 
@@ -2885,14 +2894,14 @@ def _collection_query_sparse(  # noqa: C901
 
             # Build SQL query with sparse vector distance calculation
             sql = f"""
-                SELECT {select_clause},
+                SELECT {hint_sql} {select_clause},
                        {distance_func}(sparse_embedding, {sv_sql}) AS distance
                 FROM `{table_name}`
                 {where_clause}
                 ORDER BY {distance_func}(sparse_embedding, {sv_sql})
                 APPROXIMATE
                 LIMIT %s
-            """
+            """.strip()
 
             # Execute query
             query_params = [*params, n_results]
@@ -2960,6 +2969,7 @@ def _collection_get(  # noqa: C901
         limit: int | None = None,
         offset: int | None = None,
         include: list[str] | None = None,
+        query_hint: QueryHint | None = None,
         **kwargs,
     ) -> dict[str, Any]:
         """
@@ -2974,6 +2984,7 @@ def _collection_get(  # noqa: C901
             limit: Maximum number of results (optional)
             offset: Number of results to skip (optional)
             include: Fields to include in results (optional)
+            query_hint: Query optimization hints for database execution (optional)
             **kwargs: Additional parameters
 
         Returns:
@@ -3015,9 +3026,12 @@ def _collection_get(  # noqa: C901
         # Build WHERE clause from filters
         where_clause, params = self._build_where_clause(where, where_document, id_list)
 
+        # Build query hint
+        hint_sql = _query_hint_to_sql(query_hint, table_name=table_name)
+
         # Build SQL query
         sql = f"""
-            SELECT {select_clause}
+            SELECT {hint_sql} {select_clause}
             FROM `{table_name}`
             {where_clause}
             LIMIT %s OFFSET %s
@@ -3073,6 +3087,7 @@ def _collection_hybrid_search(
         rank: dict[str, Any] | None = None,
         n_results: int = 10,
         include: list[str] | None = None,
+        query_hint: QueryHint | None = None,
         dimension: int | None = None,
         **kwargs,
     ) -> dict[str, Any]:
@@ -3171,6 +3186,12 @@ def _collection_hybrid_search(
             # Remove any surrounding quotes if present
             query_sql = query_sql.strip().strip("'\"")
 
+        # Add query hint to the generated SQL
+        hint_sql = _query_hint_to_sql(query_hint, table_name=table_name)
+        if hint_sql and query_sql.upper().startswith("SELECT"):
+            # Insert hint after SELECT keyword
+            query_sql = f"SELECT {hint_sql} {query_sql[len('SELECT') :]}"
+
         logger.debug(f"Executing query SQL: {query_sql}")
 
         # Execute the returned SQL query

src/pyseekdb/client/collection.py

@@ -13,6 +13,7 @@
 if TYPE_CHECKING:
     from .embedding_function import Documents as EmbeddingDocuments
     from .embedding_function import EmbeddingFunction
+    from .query_types import QueryHint
     from .schema import SparseVectorIndexConfig
     from .sparse_embedding_function import SparseEmbeddingFunction
 
@@ -346,6 +347,7 @@ def query(
         where_document: dict[str, Any] | None = None,
         include: list[str] | None = None,
         query_key: Any | None = None,
+        query_hint: "QueryHint | None" = None,
         **kwargs,
     ) -> dict[str, Any]:
         """
@@ -370,6 +372,7 @@ def query(
             query_key: Specify which index to query. Default is None (dense vector).
                        Use ``K.SPARSE_EMBEDDING`` (or ``"#sparse_embedding"``)
                        to query using sparse vector index.
+            query_hint: Query optimization hints for database execution (optional)
             **kwargs: Additional parameters
 
         Returns:
@@ -399,6 +402,15 @@ def query(
                 query_key=K.SPARSE_EMBEDDING,
                 n_results=5
             )
+
+            # Query with query hint
+            from pyseekdb.client.query_types import QueryHint
+            results = collection.query(
+                query_texts=["machine learning"],
+                n_results=5,
+                where={"score": {"$gte": 90}},
+                query_hint=QueryHint(parallel=8, query_timeout=10.0)
+            )
         """
         return self._client._collection_query(
             collection_id=self._id,
@@ -409,6 +421,7 @@ def query(
             where=where,
             where_document=where_document,
             include=include,
+            query_hint=query_hint,
             embedding_function=self._embedding_function,
             distance=self._distance,
             query_key=query_key,
@@ -424,6 +437,7 @@ def get(
         limit: int | None = None,
         offset: int | None = None,
         include: list[str] | None = None,
+        query_hint: "QueryHint | None" = None,
         **kwargs,
     ) -> dict[str, Any]:
         """
@@ -436,6 +450,7 @@ def get(
             limit: Maximum number of results to return (optional)
             offset: Number of results to skip (optional)
             include: Fields to include in results, e.g., ["metadatas", "documents", "embeddings"] (optional)
+            query_hint: Query optimization hints for database execution (optional)
             **kwargs: Additional parameters
 
         Returns:
@@ -469,6 +484,14 @@ def get(
 
             # Get all data
             results = collection.get(limit=100)
+
+            # Get with query hint
+            from pyseekdb.client.query_types import QueryHint
+            results = collection.get(
+                where={"category": "AI"},
+                limit=10,
+                query_hint=QueryHint(parallel=4, query_timeout=5.0)
+            )
         """
         return self._client._collection_get(
             collection_id=self._id,
@@ -479,6 +502,7 @@ def get(
             limit=limit,
             offset=offset,
             include=include,
+            query_hint=query_hint,
             **kwargs,
         )
 
@@ -489,6 +513,7 @@ def hybrid_search(
         rank: dict[str, Any] | None = None,
         n_results: int = 10,
         include: list[str] | None = None,
+        query_hint: "QueryHint | None" = None,
         **kwargs,
     ) -> dict[str, Any]:
         """
@@ -507,6 +532,7 @@ def hybrid_search(
             rank: Ranking configuration dict (e.g., {"rrf": {"rank_window_size": 60, "rank_constant": 60}})
             n_results: Final number of results to return after ranking (default: 10)
             include: Fields to include in results (e.g., ["documents", "metadatas", "embeddings"])
+            query_hint: Query optimization hints for database execution (optional)
             **kwargs: Additional parameters
 
         Returns:
@@ -537,6 +563,22 @@ def hybrid_search(
             # results["ids"][0] contains IDs for the hybrid search
             # results["documents"][0] contains documents for the hybrid search
             # results["distances"][0] contains distances for the hybrid search
+
+            # Hybrid search with query hint
+            from pyseekdb.client.query_types import QueryHint
+            results = collection.hybrid_search(
+                query={
+                    "where_document": {"$contains": "AI"},
+                    "n_results": 8
+                },
+                knn={
+                    "query_texts": ["artificial intelligence"],
+                    "n_results": 8
+                },
+                rank={"rrf": {"rank_window_size": 60}},
+                n_results=10,
+                query_hint=QueryHint(parallel=6, query_timeout=15.0)
+            )
         """
         # When no query/knn provided, return only ids/distances by default
         if include is None and not query and not knn:
@@ -550,6 +592,7 @@ def hybrid_search(
             rank=rank,
             n_results=n_results,
             include=include,
+            query_hint=query_hint,
             embedding_function=self._embedding_function,
             dimension=self._dimension,
             **kwargs,

src/pyseekdb/client/query_types.py

@@ -0,0 +1,32 @@
+"""
+Query-related type definitions and data classes for pyseekdb client.
+
+This module contains types and classes used for query operations,
+such as query hints for database optimization.
+"""
+
+from dataclasses import dataclass
+
+
+@dataclass
+class QueryHint:
+    """
+    Query hint for database optimization
+
+    Args:
+        parallel: Number of parallel execution threads (optional)
+        query_timeout: Query timeout in seconds (optional, converted to microseconds for OceanBase)
+        vector_index: Whether to add vector index hint (optional).
+    """
+
+    parallel: int | None = None
+    query_timeout: float | None = None
+    vector_index: bool | None = None
+
+    def __post_init__(self):
+        if self.parallel is not None and self.parallel <= 0:
+            raise ValueError(f"parallel must be positive, got {self.parallel}")
+        if self.query_timeout is not None and self.query_timeout <= 0:
+            raise ValueError(f"query_timeout must be positive, got {self.query_timeout}")
+        if self.vector_index is not None and not isinstance(self.vector_index, bool):
+            raise ValueError(f"vector_index must be a boolean, got {type(self.vector_index)}")

src/pyseekdb/client/sql_utils.py

@@ -9,6 +9,8 @@
 
 from pymysql.converters import escape_string
 
+from .query_types import QueryHint
+
 
 def escape_percent_for_sql(value: str) -> str:
     """
@@ -61,3 +63,36 @@ def render_sql_with_params(sql: str, params: Sequence[Any]) -> str:
         rendered_parts.append(replacement)
         rendered_parts.append(part)
     return "".join(rendered_parts)
+
+
+def _query_hint_to_sql(query_hint: QueryHint | None, table_name: str | None = None) -> str:
+    """
+    Convert QueryHint to SQL HINT string for OceanBase/seekdb.
+
+    Args:
+        query_hint: QueryHint object containing hint parameters
+
+    Returns:
+        SQL HINT string in format "/*+ hint1(value1) hint2(value2) */" or empty string if no hints
+    """
+    if query_hint is None:
+        return ""
+
+    hints = []
+
+    if query_hint.parallel is not None:
+        hints.append(f"parallel({query_hint.parallel})")
+
+    if query_hint.query_timeout is not None:
+        # Convert seconds to microseconds for OceanBase
+        timeout_microseconds = int(query_hint.query_timeout * 1_000_000)
+        hints.append(f"query_timeout({timeout_microseconds})")
+
+    if query_hint.vector_index is True and table_name:
+        hints.append(f"INDEX({table_name}, idx_vec)")
+
+    if not hints:
+        return ""
+
+    hint_str = " ".join(hints)
+    return f"/*+ {hint_str} */"