Adds support for HYBRID_POLICY on KNN queries with filters

Author: bsbodden

redisvl/query/query.py

@@ -188,6 +188,8 @@ def __init__(
         dialect: int = 2,
         sort_by: Optional[str] = None,
         in_order: bool = False,
+        hybrid_policy: Optional[str] = None,
+        batch_size: Optional[int] = None,
     ):
         """A query for running a vector search along with an optional filter
         expression.
@@ -224,6 +226,8 @@ def __init__(
         self._vector_field_name = vector_field_name
         self._dtype = dtype
         self._num_results = num_results
+        self._hybrid_policy: Optional[str] = None
+        self._batch_size: Optional[int] = None
         self.set_filter(filter_expression)
         query_string = self._build_query_string()
 
@@ -246,12 +250,89 @@ def __init__(
         if in_order:
             self.in_order()
 
+        if hybrid_policy is not None:
+            self.set_hybrid_policy(hybrid_policy)
+
+        if batch_size is not None:
+            self.set_batch_size(batch_size)
+
     def _build_query_string(self) -> str:
         """Build the full query string for vector search with optional filtering."""
         filter_expression = self._filter_expression
         if isinstance(filter_expression, FilterExpression):
             filter_expression = str(filter_expression)
-        return f"{filter_expression}=>[KNN {self._num_results} @{self._vector_field_name} ${self.VECTOR_PARAM} AS {self.DISTANCE_ID}]"
+
+        # Base KNN query
+        knn_query = (
+            f"KNN {self._num_results} @{self._vector_field_name} ${self.VECTOR_PARAM}"
+        )
+
+        # Add hybrid policy parameters if specified
+        if self._hybrid_policy:
+            knn_query += f" HYBRID_POLICY {self._hybrid_policy}"
+
+            # Add batch size if specified and using BATCHES policy
+            if self._hybrid_policy == "BATCHES" and self._batch_size:
+                knn_query += f" BATCH_SIZE {self._batch_size}"
+
+        # Add distance field alias
+        knn_query += f" AS {self.DISTANCE_ID}"
+
+        return f"{filter_expression}=>[{knn_query}]"
+
+    def set_hybrid_policy(self, hybrid_policy: str):
+        """Set the hybrid policy for the query.
+
+        Args:
+            hybrid_policy (str): The hybrid policy to use. Options are "BATCHES"
+                                or "ADHOC_BF".
+
+        Raises:
+            ValueError: If hybrid_policy is not one of the valid options
+        """
+        if hybrid_policy not in {"BATCHES", "ADHOC_BF"}:
+            raise ValueError("hybrid_policy must be one of {'BATCHES', 'ADHOC_BF'}")
+        self._hybrid_policy = hybrid_policy
+
+        # Reset the query string
+        self._query_string = self._build_query_string()
+
+    def set_batch_size(self, batch_size: int):
+        """Set the batch size for the query.
+
+        Args:
+            batch_size (int): The batch size to use when hybrid_policy is "BATCHES".
+
+        Raises:
+            TypeError: If batch_size is not an integer
+            ValueError: If batch_size is not positive
+        """
+        if not isinstance(batch_size, int):
+            raise TypeError("batch_size must be an integer")
+        if batch_size <= 0:
+            raise ValueError("batch_size must be positive")
+        self._batch_size = batch_size
+
+        # Reset the query string
+        self._query_string = self._build_query_string()
+
+    @property
+    def hybrid_policy(self) -> Optional[str]:
+        """Return the hybrid policy for the query.
+
+        Returns:
+            Optional[str]: The hybrid policy for the query.
+        """
+        return self._hybrid_policy
+
+    @property
+    def batch_size(self) -> Optional[int]:
+        """Return the batch size for the query.
+
+        Returns:
+            Optional[int]: The batch size for the query.
+        """
+        return self._batch_size
 
     @property
     def params(self) -> Dict[str, Any]:
@@ -265,7 +346,9 @@ def params(self) -> Dict[str, Any]:
         else:
             vector = array_to_buffer(self._vector, dtype=self._dtype)
 
-        return {self.VECTOR_PARAM: vector}
+        params = {self.VECTOR_PARAM: vector}
+
+        return params
 
 
 class VectorRangeQuery(BaseVectorQuery, BaseQuery):
@@ -279,6 +362,7 @@ def __init__(
         filter_expression: Optional[Union[str, FilterExpression]] = None,
         dtype: str = "float32",
         distance_threshold: float = 0.2,
+        epsilon: Optional[float] = None,
         num_results: int = 10,
         return_score: bool = True,
         dialect: int = 2,
@@ -301,6 +385,11 @@ def __init__(
             distance_threshold (str, float): The threshold for vector distance.
                 A smaller threshold indicates a stricter semantic search.
                 Defaults to 0.2.
+            epsilon (Optional[float]): The relative factor for vector range queries,
+                setting boundaries for candidates within radius * (1 + epsilon).
+                This controls how extensive the search is beyond the specified radius.
+                Higher values increase recall at the expense of performance.
+                Defaults to None, which uses the index-defined epsilon (typically 0.01).
             num_results (int): The MAX number of results to return.
                 Defaults to 10.
             return_score (bool, optional): Whether to return the vector
@@ -324,6 +413,11 @@ def __init__(
         self._vector_field_name = vector_field_name
         self._dtype = dtype
         self._num_results = num_results
+        self._epsilon: Optional[float] = None
+
+        if epsilon is not None:
+            self.set_epsilon(epsilon)
+
         self.set_distance_threshold(distance_threshold)
         self.set_filter(filter_expression)
         query_string = self._build_query_string()
@@ -349,15 +443,20 @@ def __init__(
 
     def _build_query_string(self) -> str:
         """Build the full query string for vector range queries with optional filtering"""
+        range_params = ""
+        if self._epsilon is not None:
+            range_params = f"; $EPSILON: {self._epsilon}"
+
         base_query = f"@{self._vector_field_name}:[VECTOR_RANGE ${self.DISTANCE_THRESHOLD_PARAM} ${self.VECTOR_PARAM}]"
+        attr_section = f"=>{{$yield_distance_as: {self.DISTANCE_ID}{range_params}}}"
 
         filter_expression = self._filter_expression
         if isinstance(filter_expression, FilterExpression):
             filter_expression = str(filter_expression)
 
         if filter_expression == "*":
-            return f"{base_query}=>{{$yield_distance_as: {self.DISTANCE_ID}}}"
-        return f"({base_query}=>{{$yield_distance_as: {self.DISTANCE_ID}}} {filter_expression})"
+            return f"{base_query}{attr_section}"
+        return f"({base_query}{attr_section} {filter_expression})"
 
     def set_distance_threshold(self, distance_threshold: float):
         """Set the distance threshold for the query.
@@ -369,6 +468,35 @@ def set_distance_threshold(self, distance_threshold: float):
             raise TypeError("distance_threshold must be of type int or float")
         self._distance_threshold = distance_threshold
 
+    def set_epsilon(self, epsilon: float):
+        """Set the epsilon parameter for the range query.
+
+        Args:
+            epsilon (float): The relative factor for vector range queries,
+                setting boundaries for candidates within radius * (1 + epsilon).
+
+        Raises:
+            TypeError: If epsilon is not a float or int
+            ValueError: If epsilon is negative
+        """
+        if not isinstance(epsilon, (float, int)):
+            raise TypeError("epsilon must be of type float or int")
+        if epsilon < 0:
+            raise ValueError("epsilon must be non-negative")
+        self._epsilon = epsilon
+
+        # Reset the query string
+        self._query_string = self._build_query_string()
+
+    @property
+    def epsilon(self) -> Optional[float]:
+        """Return the epsilon for the query.
+
+        Returns:
+            Optional[float]: The epsilon for the query, or None if not set.
+        """
+        return self._epsilon
+
     @property
     def distance_threshold(self) -> float:
         """Return the distance threshold for the query.
@@ -390,11 +518,14 @@ def params(self) -> Dict[str, Any]:
         else:
             vector_param = array_to_buffer(self._vector, dtype=self._dtype)
 
-        return {
+        # Only include the necessary parameters, not EPSILON
+        params = {
             self.VECTOR_PARAM: vector_param,
             self.DISTANCE_THRESHOLD_PARAM: self._distance_threshold,
         }
 
+        return params
+
 
 class RangeQuery(VectorRangeQuery):
     # keep for backwards compatibility

tests/integration/test_query.py

@@ -4,6 +4,7 @@
 from redisvl.index import SearchIndex
 from redisvl.query import CountQuery, FilterQuery, RangeQuery, VectorQuery
 from redisvl.query.filter import FilterExpression, Geo, GeoRadius, Num, Tag, Text
+from redisvl.query.query import VectorRangeQuery
 from redisvl.redis.utils import array_to_buffer
 
 # TODO expand to multiple schema types and sync + async
@@ -466,3 +467,93 @@ def test_query_with_chunk_number_zero():
     assert (
         str(filter_conditions) == expected_query_str
     ), "Query with chunk_number zero is incorrect"
+
+
+def test_hybrid_policy_batches_mode(index, vector_query):
+    """Test vector query with BATCHES hybrid policy."""
+    # Create a filter
+    t = Tag("credit_score") == "high"
+
+    # Set hybrid policy to BATCHES
+    vector_query.set_hybrid_policy("BATCHES")
+    vector_query.set_batch_size(2)
+
+    # Set the filter
+    vector_query.set_filter(t)
+
+    # Check query string
+    assert "HYBRID_POLICY BATCHES BATCH_SIZE 2" in str(vector_query)
+
+    # Execute query
+    results = index.query(vector_query)
+
+    # Check results - should have filtered to "high" credit scores
+    assert len(results) > 0
+    for result in results:
+        assert result["credit_score"] == "high"
+
+
+def test_hybrid_policy_adhoc_bf_mode(index, vector_query):
+    """Test vector query with ADHOC_BF hybrid policy."""
+    # Create a filter
+    t = Tag("credit_score") == "high"
+
+    # Set hybrid policy to ADHOC_BF
+    vector_query.set_hybrid_policy("ADHOC_BF")
+
+    # Set the filter
+    vector_query.set_filter(t)
+
+    # Check query string
+    assert "HYBRID_POLICY ADHOC_BF" in str(vector_query)
+
+    # Execute query
+    results = index.query(vector_query)
+
+    # Check results - should have filtered to "high" credit scores
+    assert len(results) > 0
+    for result in results:
+        assert result["credit_score"] == "high"
+
+
+def test_range_query_with_epsilon(index):
+    """Test vector range query with epsilon parameter."""
+    # Create a range query with default epsilon
+    default_query = VectorRangeQuery(
+        vector=[0.1, 0.1, 0.5],
+        vector_field_name="user_embedding",
+        return_fields=["user", "credit_score", "age", "job"],
+        distance_threshold=0.2,
+    )
+
+    # Create a range query with custom epsilon
+    custom_query = VectorRangeQuery(
+        vector=[0.1, 0.1, 0.5],
+        vector_field_name="user_embedding",
+        return_fields=["user", "credit_score", "age", "job"],
+        distance_threshold=0.2,
+        epsilon=0.05,  # More strict than the default
+    )
+
+    # Verify query string contains epsilon attribute
+    query_string = str(custom_query)
+    assert "$EPSILON: 0.05" in query_string
+
+    # Verify epsilon property is set
+    assert custom_query.epsilon == 0.05
+
+    # Test setting epsilon
+    custom_query.set_epsilon(0.1)
+    assert custom_query.epsilon == 0.1
+    assert "$EPSILON: 0.1" in str(custom_query)
+
+    # Verify the parameters don't include epsilon
+    assert "EPSILON" not in custom_query.params
+
+    # Execute basic query without epsilon to ensure functionality
+    results = index.query(default_query)
+
+    # Check results
+    assert len(results) > 0
+    for result in results:
+        assert float(result["vector_distance"]) <= 0.2