Add vector search test.

Author: WaVEV

django_mongodb_backend/compiler.py

@@ -248,12 +248,13 @@ def _build_aggregation_pipeline(self, ids, group):
                 pipeline.append({"$unset": "_id"})
         return pipeline
 
-    def _compound_searches_queries(self, searches):
+    def _compound_searches_queries(self, searches, search_replacements):
         if not searches:
             return []
         if len(searches) > 1:
             raise ValueError("Cannot perform more than one search operation.")
-        return [searches[0], {"$addFields": {"__search_expr.search1": {"$meta": "searchScore"}}}]
+        score_function = "searchScore" if "$search" in searches[0] else "vectorSearchScore"
+        return [searches[0], {"$addFields": {"__search_expr.search1": {"$meta": score_function}}}]
 
     def pre_sql_setup(self, with_col_aliases=False):
         extra_select, order_by, group_by = super().pre_sql_setup(with_col_aliases=with_col_aliases)
@@ -262,7 +263,7 @@ def pre_sql_setup(self, with_col_aliases=False):
         )
         group, group_replacements = self._prepare_annotations_for_aggregation_pipeline(order_by)
         all_replacements = {**search_replacements, **group_replacements}
-        self.search_pipeline = self._compound_searches_queries(searches)
+        self.search_pipeline = self._compound_searches_queries(searches, search_replacements)
         # query.group_by is either:
         # - None: no GROUP BY
         # - True: group by select fields

django_mongodb_backend/expressions/builtins.py

@@ -598,61 +598,6 @@ def get_search_fields(self):
         return needed_fields
 
 
-class SearchVector(SearchExpression):
-    def __init__(
-        self,
-        path,
-        query_vector,
-        limit,
-        num_candidates=None,
-        exact=None,
-        filter=None,
-    ):
-        self.path = path
-        self.query_vector = query_vector
-        self.limit = limit
-        self.num_candidates = num_candidates
-        self.exact = exact
-        self.filter = filter
-        super().__init__()
-
-    def __invert__(self):
-        return ValueError("SearchVector cannot be negated")
-
-    def __and__(self, other):
-        raise NotSupportedError("SearchVector cannot be combined")
-
-    def __rand__(self, other):
-        raise NotSupportedError("SearchVector cannot be combined")
-
-    def __or__(self, other):
-        raise NotSupportedError("SearchVector cannot be combined")
-
-    def __ror__(self, other):
-        raise NotSupportedError("SearchVector cannot be combined")
-
-    def get_search_fields(self):
-        return {self.path}
-
-    def _get_query_index(self, field, compiler):
-        return "default"
-
-    def as_mql(self, compiler, connection):
-        params = {
-            "index": self._get_query_index(self.get_search_fields()),
-            "path": self.path,
-            "queryVector": self.query_vector,
-            "limit": self.limit,
-        }
-        if self.num_candidates is not None:
-            params["numCandidates"] = self.num_candidates
-        if self.exact is not None:
-            params["exact"] = self.exact
-        if self.filter is not None:
-            params["filter"] = self.filter
-        return {"$vectorSearch": params}
-
-
 class CompoundExpression(SearchExpression):
     def __init__(
         self,
@@ -716,22 +661,76 @@ def resolve(node, negated=False):
         if operator == Operator.OR:
             return CompoundExpression(should=[lhs_compound, rhs_compound], minimum_should_match=1)
         if operator == Operator.AND:
-            # NOTE: we can't just do the code below, think about this case (A | B) & (C | D)
-            # return CompoundExpression(
-            #     must=lhs_compound.must + rhs_compound.must,
-            #     must_not=lhs_compound.must_not + rhs_compound.must_not,
-            #     should=lhs_compound.should + rhs_compound.should,
-            #     filter=lhs_compound.filter + rhs_compound.filter,
-            # )
             return CompoundExpression(must=[lhs_compound, rhs_compound])
-        # not operator
         return lhs_compound
 
     def as_mql(self, compiler, connection):
         expression = self.resolve(self)
         return expression.as_mql(compiler, connection)
 
 
+class SearchVector(SearchExpression):
+    def __init__(
+        self,
+        path,
+        query_vector,
+        limit,
+        num_candidates=None,
+        exact=None,
+        filter=None,
+    ):
+        self.path = path
+        self.query_vector = query_vector
+        self.limit = limit
+        self.num_candidates = num_candidates
+        self.exact = exact
+        self.filter = filter
+        super().__init__()
+
+    def __invert__(self):
+        return ValueError("SearchVector cannot be negated")
+
+    def __and__(self, other):
+        raise NotSupportedError("SearchVector cannot be combined")
+
+    def __rand__(self, other):
+        raise NotSupportedError("SearchVector cannot be combined")
+
+    def __or__(self, other):
+        raise NotSupportedError("SearchVector cannot be combined")
+
+    def __ror__(self, other):
+        raise NotSupportedError("SearchVector cannot be combined")
+
+    def get_search_fields(self):
+        return {self.path}
+
+    def _get_query_index(self, fields, compiler):
+        for search_indexes in compiler.collection.list_search_indexes():
+            if search_indexes["type"] == "vectorSearch":
+                index_field = {
+                    field["path"] for field in search_indexes["latestDefinition"]["fields"]
+                }
+                if fields.issubset(index_field):
+                    return search_indexes["name"]
+        return "default"
+
+    def as_mql(self, compiler, connection):
+        params = {
+            "index": self._get_query_index(self.get_search_fields(), compiler),
+            "path": self.path,
+            "queryVector": self.query_vector,
+            "limit": self.limit,
+        }
+        if self.num_candidates is not None:
+            params["numCandidates"] = self.num_candidates
+        if self.exact is not None:
+            params["exact"] = self.exact
+        if self.filter is not None:
+            params["filter"] = self.filter
+        return {"$vectorSearch": params}
+
+
 class SearchScoreOption:
     """Class to mutate scoring on a search operation"""
 

tests/queries_/models.py

@@ -1,6 +1,6 @@
 from django.db import models
 
-from django_mongodb_backend.fields import ObjectIdAutoField, ObjectIdField
+from django_mongodb_backend.fields import ArrayField, ObjectIdAutoField, ObjectIdField
 
 
 class Author(models.Model):
@@ -60,3 +60,4 @@ class Article(models.Model):
     number = models.IntegerField()
     body = models.TextField()
     location = models.JSONField(null=True)
+    plot_embedding = ArrayField(models.FloatField(), size=3)

tests/queries_/test_search.py

@@ -17,6 +17,7 @@
     SearchRange,
     SearchRegex,
     SearchText,
+    SearchVector,
     SearchWildcard,
 )
 
@@ -29,9 +30,9 @@ def _get_collection(model):
         return connection.database.get_collection(model._meta.db_table)
 
     @staticmethod
-    def create_search_index(model, index_name, definition):
+    def create_search_index(model, index_name, definition, type="search"):
         collection = CreateIndexMixin._get_collection(model)
-        idx = SearchIndexModel(definition=definition, name=index_name)
+        idx = SearchIndexModel(definition=definition, name=index_name, type=type)
         collection.create_search_index(idx)
 
 
@@ -365,3 +366,49 @@ def test_compound_operations(self):
         )
         qs = Article.objects.annotate(score=expr)
         self.assertCountEqual(qs, [self.mars_mission, self.exoplanet])
+
+
+class SearchVectorTest(TestCase, CreateIndexMixin):
+    @classmethod
+    def setUpTestData(cls):
+        cls.create_search_index(
+            Article,
+            "vector_index",
+            {
+                "fields": [
+                    {
+                        "type": "vector",
+                        "path": "plot_embedding",
+                        "numDimensions": 3,
+                        "similarity": "cosine",
+                        "quantization": "scalar",
+                    }
+                ]
+            },
+            type="vectorSearch",
+        )
+
+        cls.mars = Article.objects.create(
+            headline="Mars landing",
+            number=1,
+            body="The rover has landed on Mars",
+            plot_embedding=[0.1, 0.2, 0.3],
+        )
+        Article.objects.create(
+            headline="Cooking tips",
+            number=2,
+            body="This article is about pasta",
+            plot_embedding=[0.9, 0.8, 0.7],
+        )
+        time.sleep(1)
+
+    def test_vector_search(self):
+        vector_query = [0.1, 0.2, 0.3]
+        expr = SearchVector(
+            path="plot_embedding",
+            query_vector=vector_query,
+            num_candidates=5,
+            limit=2,
+        )
+        qs = Article.objects.annotate(score=expr).order_by("-score")
+        self.assertEqual(qs.first(), self.mars)