edits

Author: timgraham

django_mongodb_backend/features.py

@@ -634,9 +634,10 @@ def supports_atlas_search(self):
             # An existing collection must be used on MongoDB 6, otherwise
             # the operation will not error when unsupported.
             self.connection.get_collection("django_migrations").list_search_indexes()
-        except OperationFailure as exc:
-            if "$listSearchIndexes stage is only allowed on MongoDB Atlas" in str(exc):
-                return False
-            raise
+        except OperationFailure:
+            # It would be best to check the error message or error code to, but
+            # they vary across MongoDB versions. Example: "$listSearchIndexes
+            # stage is only allowed on MongoDB Atlas".
+            return False
         else:
             return True

django_mongodb_backend/indexes.py

@@ -131,7 +131,7 @@ def check(self, model, connection):
     def search_index_data_types(self, db_type):
         """
         Map a model field's type to search index type.
-        Reference: https://www.mongodb.com/docs/atlas/atlas-search/define-field-mappings/#data-types
+        https://www.mongodb.com/docs/atlas/atlas-search/define-field-mappings/#data-types
         """
         if db_type in {"double", "int", "long"}:
             return "number"
@@ -163,17 +163,11 @@ def get_pymongo_index_model(
 
 class VectorSearchIndex(SearchIndex):
     suffix = "vsi"
-    VALID_SIMILARITIES = frozenset(("euclidean", "cosine", "dotProduct"))
+    VALID_SIMILARITIES = frozenset(("cosine", "dotProduct", "euclidean"))
     VALID_FIELD_TYPES = frozenset(("boolean", "date", "number", "objectId", "string", "uuid"))
     _error_id_prefix = "django_mongodb_backend.indexes.VectorSearchIndex"
 
-    def __init__(
-        self,
-        *,
-        fields=(),
-        similarities="cosine",
-        name=None,
-    ):
+    def __init__(self, *, fields=(), similarities="cosine", name=None):
         super().__init__(fields=fields, name=name)
         self.similarities = similarities
         self._multiple_similarities = isinstance(similarities, tuple | list)
@@ -191,11 +185,11 @@ def __init__(
 
     def check(self, model, connection):
         errors = super().check(model, connection)
-        expected_similarities = 0
+        num_arrayfields = 0
         for field_name, _ in self.fields_orders:
             field = model._meta.get_field(field_name)
             if isinstance(field, ArrayField):
-                expected_similarities += 1
+                num_arrayfields += 1
                 try:
                     int(field.size)
                 except (ValueError, TypeError):
@@ -219,7 +213,6 @@ def check(self, model, connection):
                     )
             else:
                 search_type = self.search_index_data_types(field.db_type(connection))
-                # Validate allowed search types.
                 if search_type not in self.VALID_FIELD_TYPES:
                     errors.append(
                         Error(
@@ -230,34 +223,33 @@ def check(self, model, connection):
                             hint=f"Allowed types are {', '.join(sorted(self.VALID_FIELD_TYPES))}.",
                         )
                     )
-        if self._multiple_similarities and expected_similarities != len(self.similarities):
+        if self._multiple_similarities and num_arrayfields != len(self.similarities):
             errors.append(
                 Error(
                     f"VectorSearchIndex requires the same number of similarities "
                     f"and vector fields; {model._meta.object_name} has "
-                    f"{expected_similarities} ArrayField(s) but similarities "
+                    f"{num_arrayfields} ArrayField(s) but similarities "
                     f"has {len(self.similarities)} element(s).",
                     obj=model,
                     id=f"{self._error_id_prefix}.E005",
                 )
             )
-        # There isn't any vector.
-        if expected_similarities == 0:
+        if num_arrayfields == 0:
             errors.append(
                 Error(
-                    "VectorSearchIndex requires at least one field containing vector data "
-                    "(e.g., an ArrayField(FloatField, size=10)). "
-                    "If you're aiming to perform search operations on other data types, "
-                    "consider using SearchIndex instead.",
+                    "VectorSearchIndex requires at least one ArrayField to " "store vector data.",
                     obj=model,
                     id=f"{self._error_id_prefix}.E006",
+                    hint="If you want to perform search operations without vectors, "
+                    "use SearchIndex instead.",
                 )
             )
         return errors
 
     def deconstruct(self):
         path, args, kwargs = super().deconstruct()
-        kwargs["similarities"] = self.similarities
+        if self.similarities != "cosine":
+            kwargs["similarities"] = self.similarities
         return path, args, kwargs
 
     def get_pymongo_index_model(

docs/source/ref/models/indexes.rst

@@ -38,14 +38,17 @@ model has multiple indexes).
 A subclass of :class:`SearchIndex` that creates a :doc:`vector search index
 <atlas:atlas-vector-search/vector-search-type>` on the given field(s).
 
-Each index should references at least one vector field: an :class:`.ArrayField`
+The index should reference at least one vector field: an :class:`.ArrayField`
 with a :attr:`~.ArrayField.base_field` of :class:`~django.db.models.FloatField`
-or :class:`~django.db.models.IntegerField`.
+or :class:`~django.db.models.IntegerField`. It cannot reference an
+:class:`.ArrayField` of any other type. Each :class:`.ArrayField` must have a
+:attr:`~.ArrayField.size`.
 
 It may also have other fields to filter on, provided the field stores
 ``boolean``, ``date``, ``objectId``, ``numeric``, ``string``, or ``uuid``.
 
-Available values for ``similarities`` are ``"euclidean"``, ``"cosine"``, and
-``"dotProduct"``. You can provide this value either a string, in which case
-that value will be applied to all vector fields, or a list or tuple of values
-with a similarity corresponding to each vector field.
+Available values for ``similarities`` are ``"cosine"``, ``"dotProduct"``, and
+``"euclidean"`` (see :ref:`atlas:avs-similarity-functions`). You can provide
+this value either a string, in which case that value will be applied to all
+vector fields, or a list or tuple of values with a similarity corresponding to
+each vector field.

tests/indexes_/test_checks.py

@@ -41,10 +41,10 @@ class Meta:
     def test_vector_search_requires_atlas_search_support(self):
         class Article(models.Model):
             title = models.CharField(max_length=10)
-            vector_float = ArrayField(models.FloatField(), size=10)
+            vector = ArrayField(models.FloatField(), size=10)
 
             class Meta:
-                indexes = [VectorSearchIndex(fields=["title", "vector_float"])]
+                indexes = [VectorSearchIndex(fields=["title", "vector"])]
 
         errors = checks.run_checks(app_configs=self.apps.get_app_configs(), databases={"default"})
         self.assertEqual(
@@ -109,10 +109,10 @@ class Meta:
     def test_unsupported_type(self):
         class Article(models.Model):
             data = models.JSONField()
-            vector_float = ArrayField(models.FloatField(), size=10)
+            vector = ArrayField(models.FloatField(), size=10)
 
             class Meta:
-                indexes = [VectorSearchIndex(fields=["data", "vector_float"])]
+                indexes = [VectorSearchIndex(fields=["data", "vector"])]
 
         errors = checks.run_checks(app_configs=self.apps.get_app_configs(), databases={"default"})
         self.assertEqual(
@@ -129,12 +129,12 @@ class Meta:
 
     def test_invalid_number_similarity_function_singular(self):
         class Article(models.Model):
-            vector_data = ArrayField(models.FloatField(), size=10)
+            vector = ArrayField(models.FloatField(), size=10)
 
             class Meta:
                 indexes = [
                     VectorSearchIndex(
-                        fields=["vector_data"],
+                        fields=["vector"],
                         similarities=["dotProduct", "cosine"],
                     )
                 ]
@@ -182,10 +182,10 @@ class Meta:
 
     def test_simple(self):
         class Article(models.Model):
-            vector_data = ArrayField(models.FloatField(), size=10)
+            vector = ArrayField(models.FloatField(), size=10)
 
             class Meta:
-                indexes = [VectorSearchIndex(fields=["vector_data"])]
+                indexes = [VectorSearchIndex(fields=["vector"])]
 
         errors = checks.run_checks(app_configs=self.apps.get_app_configs(), databases={"default"})
         self.assertEqual(errors, [])
@@ -234,11 +234,11 @@ class Meta:
             errors,
             [
                 checks.Error(
-                    "VectorSearchIndex requires at least one field containing vector data "
-                    "(e.g., an ArrayField(FloatField, size=10)). If you're aiming to perform "
-                    "search operations on other data types, consider using SearchIndex instead.",
+                    "VectorSearchIndex requires at least one ArrayField to " "store vector data.",
                     id="django_mongodb_backend.indexes.VectorSearchIndex.E006",
                     obj=NoSearchVectorModel,
+                    hint="If you want to perform search operations without vectors, "
+                    "use SearchIndex instead.",
                 ),
             ],
         )

tests/indexes_/test_search_indexes.py

@@ -12,30 +12,32 @@
 @mock.patch.object(connection.features, "supports_atlas_search", False)
 class UnsupportedSearchIndexesTests(TestCase):
     def test_search_index_not_created(self):
-        with connection.schema_editor() as editor:
-            index = SearchIndex(name="recent_test_idx", fields=["number"])
-            with self.assertNumQueries(0):
-                editor.add_index(index=index, model=SearchIndexTestModel)
+        index = SearchIndex(name="recent_test_idx", fields=["number"])
+        with connection.schema_editor() as editor, self.assertNumQueries(0):
+            editor.add_index(index=index, model=SearchIndexTestModel)
         self.assertNotIn(
             index.name,
             connection.introspection.get_constraints(
                 cursor=None,
                 table_name=SearchIndexTestModel._meta.db_table,
             ),
         )
+        with connection.schema_editor() as editor, self.assertNumQueries(0):
+            editor.remove_index(index=index, model=SearchIndexTestModel)
 
     def test_vector_index_not_created(self):
-        with connection.schema_editor() as editor:
-            index = VectorSearchIndex(name="recent_test_idx", fields=["number"])
-            with self.assertNumQueries(0):
-                editor.add_index(index=index, model=SearchIndexTestModel)
+        index = VectorSearchIndex(name="recent_test_idx", fields=["number"])
+        with connection.schema_editor() as editor, self.assertNumQueries(0):
+            editor.add_index(index=index, model=SearchIndexTestModel)
         self.assertNotIn(
             index.name,
             connection.introspection.get_constraints(
                 cursor=None,
                 table_name=SearchIndexTestModel._meta.db_table,
             ),
         )
+        with connection.schema_editor() as editor, self.assertNumQueries(0):
+            editor.remove_index(index=index, model=SearchIndexTestModel)
 
 
 class SearchIndexTests(SimpleTestCase):
@@ -68,6 +70,7 @@ def test_no_extra_kargs(self):
     def test_deconstruct(self):
         index = VectorSearchIndex(name="recent_test_idx", fields=["number"])
         name, args, kwargs = index.deconstruct()
+        self.assertEqual(kwargs, {"name": "recent_test_idx", "fields": ["number"]})
         new = VectorSearchIndex(*args, **kwargs)
         self.assertEqual(new.similarities, index.similarities)
 
@@ -255,3 +258,89 @@ def test_multiple_fields(self):
         finally:
             with connection.schema_editor() as editor:
                 editor.remove_index(index=index, model=SearchIndexTestModel)
+
+    def test_similarities_value(self):
+        index = VectorSearchIndex(
+            name="recent_test_idx",
+            fields=["vector_float", "vector_integer"],
+            similarities="euclidean",
+        )
+        with connection.schema_editor() as editor:
+            editor.add_index(index=index, model=SearchIndexTestModel)
+        try:
+            index_info = connection.introspection.get_constraints(
+                cursor=None,
+                table_name=SearchIndexTestModel._meta.db_table,
+            )
+            expected_options = {
+                "latestDefinition": {
+                    "fields": [
+                        {
+                            "numDimensions": 10,
+                            "path": "vector_float",
+                            "similarity": "euclidean",
+                            "type": "vector",
+                        },
+                        {
+                            "numDimensions": 10,
+                            "path": "vector_integer",
+                            "similarity": "euclidean",
+                            "type": "vector",
+                        },
+                    ]
+                },
+                "latestVersion": 0,
+                "name": "recent_test_idx",
+                "queryable": False,
+                "type": "vectorSearch",
+            }
+            self.assertCountEqual(index_info[index.name]["columns"], index.fields)
+            index_info[index.name]["options"].pop("id")
+            index_info[index.name]["options"].pop("status")
+            self.assertEqual(index_info[index.name]["options"], expected_options)
+        finally:
+            with connection.schema_editor() as editor:
+                editor.remove_index(index=index, model=SearchIndexTestModel)
+
+    def test_similarities_list(self):
+        index = VectorSearchIndex(
+            name="recent_test_idx",
+            fields=["vector_float", "vector_integer"],
+            similarities=["cosine", "euclidean"],
+        )
+        with connection.schema_editor() as editor:
+            editor.add_index(index=index, model=SearchIndexTestModel)
+        try:
+            index_info = connection.introspection.get_constraints(
+                cursor=None,
+                table_name=SearchIndexTestModel._meta.db_table,
+            )
+            expected_options = {
+                "latestDefinition": {
+                    "fields": [
+                        {
+                            "numDimensions": 10,
+                            "path": "vector_float",
+                            "similarity": "cosine",
+                            "type": "vector",
+                        },
+                        {
+                            "numDimensions": 10,
+                            "path": "vector_integer",
+                            "similarity": "euclidean",
+                            "type": "vector",
+                        },
+                    ]
+                },
+                "latestVersion": 0,
+                "name": "recent_test_idx",
+                "queryable": False,
+                "type": "vectorSearch",
+            }
+            self.assertCountEqual(index_info[index.name]["columns"], index.fields)
+            index_info[index.name]["options"].pop("id")
+            index_info[index.name]["options"].pop("status")
+            self.assertEqual(index_info[index.name]["options"], expected_options)
+        finally:
+            with connection.schema_editor() as editor:
+                editor.remove_index(index=index, model=SearchIndexTestModel)