PYCBC-1680: Add support for FTS Vector Search pre-filtering

Changes
=======
* Added prefilter option to `VectorQuery`
* Updated unit tests to confirm search JSON is encoded correctly

Results
=======
All tests pass

Change-Id: I5107074fea06b4486fce150dd868de0079db5804
Reviewed-on: https://review.couchbase.org/c/couchbase-python-client/+/226340
Reviewed-by: Dimitris Christodoulou <[email protected]>
Tested-by: Build Bot <[email protected]>

Author: thejcfactor

couchbase/logic/search.py

@@ -969,6 +969,8 @@ def encode_vector_search(self) -> Optional[List[Dict[str, Any]]]:
                 encoded_query['vector_base64'] = query.vector_base64
             if query.boost is not None:
                 encoded_query['boost'] = query.boost
+            if query.prefilter is not None:
+                encoded_query['filter'] = query.prefilter.encodable
             encoded_queries.append(encoded_query)
 
         return encoded_queries

couchbase/logic/vector_search.py

@@ -1,14 +1,18 @@
 from __future__ import annotations
 
 from enum import Enum
-from typing import (List,
+from typing import (TYPE_CHECKING,
+                    List,
                     Optional,
                     Union)
 
 from couchbase._utils import is_null_or_empty
 from couchbase.exceptions import InvalidArgumentException
 from couchbase.options import VectorSearchOptions
 
+if TYPE_CHECKING:
+    from couchbase.logic.search_queries import SearchQuery
+
 
 class VectorQueryCombination(Enum):
     """ Specifies how multiple vector searches are combined.
@@ -31,6 +35,7 @@ class VectorQuery:
         vector (Union[List[float], str]): The vector to use in the query.
         num_candidates (int, optional): Specifies the number of results returned. If provided, must be greater or equal to 1.
         boost (float, optional): Add boost to query.
+        prefilter (`~couchbase.search.SearchQuery`, optional): Specifies a pre-filter to use for the vector query.
 
     Raises:
         :class:`~couchbase.exceptions.InvalidArgumentException`: If the vector is not provided.
@@ -46,18 +51,21 @@ def __init__(self,
                  vector,  # type: Union[List[float], str]
                  num_candidates=None,  # type: Optional[int]
                  boost=None,  # type: Optional[float]
+                 prefilter=None,  # type: Optional[SearchQuery]
                  ):
         if is_null_or_empty(field_name):
             raise InvalidArgumentException('Must provide a field name.')
         self._field_name = field_name
         self._vector = None
         self._vector_base64 = None
         self._validate_and_set_vector(vector)
-        self._num_candidates = self._boost = None
+        self._num_candidates = self._boost = self._prefilter = None
         if num_candidates is not None:
             self.num_candidates = num_candidates
         if boost is not None:
             self.boost = boost
+        if prefilter is not None:
+            self.prefilter = prefilter
 
     @property
     def boost(self) -> Optional[float]:
@@ -98,6 +106,23 @@ def num_candidates(self,
             raise InvalidArgumentException('num_candidates must be >= 1.')
         self._num_candidates = value
 
+    @property
+    def prefilter(self) -> Optional[SearchQuery]:
+        """
+        Optional[SearchQuery]: Returns vector query's prefilter query, if it exists.
+        """
+        return self._prefilter
+
+    @prefilter.setter
+    def prefilter(self,
+                  value  # type: SearchQuery
+                  ):
+        # avoid circular import
+        from couchbase.logic.search_queries import SearchQuery
+        if not isinstance(value, SearchQuery):
+            raise InvalidArgumentException('prefilter must be a SearchQuery.')
+        self._prefilter = value
+
     @property
     def vector(self) -> Optional[List[float]]:
         """
@@ -138,6 +163,7 @@ def create(cls,
                vector,  # type: Union[List[float], str]
                num_candidates=None,  # type: Optional[int]
                boost=None,  # type: Optional[float]
+               prefilter=None,  # type: Optional[SearchQuery]
                ) -> VectorQuery:
         """ Creates a :class:`~couchbase.vector_search.VectorQuery`.
 
@@ -146,6 +172,7 @@ def create(cls,
             vector (Union[List[float], str]): The vector to use in the query.
             num_candidates (int, optional): Specifies the number of results returned. If provided, must be greater or equal to 1.
             boost (float, optional): Add boost to query.
+            prefilter (`~couchbase.search.SearchQuery`, optional): Specifies a pre-filter to use for the vector query.
 
         Raises:
             :class:`~couchbase.exceptions.InvalidArgumentException`: If the vector is not provided.
@@ -155,7 +182,7 @@ def create(cls,
         Returns:
             :class:`~couchbase.vector_search.VectorQuery`: The created vector query.
         """  # noqa: E501
-        return cls(field_name, vector, num_candidates=num_candidates, boost=boost)
+        return cls(field_name, vector, num_candidates=num_candidates, boost=boost, prefilter=prefilter)
 
 
 class VectorSearch:

couchbase/tests/search_params_t.py

@@ -955,6 +955,7 @@ class VectorSearchParamTestSuite:
         'test_vector_query_invalid_num_candidates',
         'test_vector_query_invalid_vector',
         'test_vector_search',
+        'test_vector_search_with_prefilter',
         'test_vector_search_base64',
         'test_vector_search_invalid',
         'test_vector_search_multiple_queries'
@@ -1047,6 +1048,43 @@ def test_vector_search(self, cb_env):
         encoded_q = cb_env.get_encoded_query(search_query)
         assert exp_json == encoded_q
 
+    def test_vector_search_with_prefilter(self, cb_env):
+        exp_json = {
+            'query': {'match_none': None},
+            'index_name': cb_env.TEST_INDEX_NAME,
+            'metrics': True,
+            'show_request': False,
+            'vector_search': [
+                {
+                    'field': 'vector_field',
+                    'vector': self.TEST_VECTOR,
+                    'k': 3,
+                    'filter': {
+                        'match': 'salty beers',
+                        'analyzer': 'analyzer',
+                        'boost': 1.5,
+                        'field': 'field',
+                        'fuzziness': 1234,
+                        'prefix_length': 4,
+                        'operator': 'or'
+                    }
+                }
+            ]
+        }
+
+        q = search.MatchQuery('salty beers', boost=1.5, analyzer='analyzer',
+                              field='field', fuzziness=1234, prefix_length=4, match_operator=MatchOperator.OR)
+        vector_search = VectorSearch.from_vector_query(VectorQuery('vector_field',
+                                                                   self.TEST_VECTOR,
+                                                                   prefilter=q))
+        req = SearchRequest.create(vector_search)
+        search_query = search.SearchQueryBuilder.create_search_query_from_request(
+            cb_env.TEST_INDEX_NAME,
+            req
+        )
+        encoded_q = cb_env.get_encoded_query(search_query)
+        assert exp_json == encoded_q
+
     def test_vector_search_base64(self, cb_env):
         exp_json = {
             'query': {'match_none': None},