Add atlas vector search index

WaVEV · WaVEV · commit d55cab218e4d · 2025-03-08T00:49:58.000-03:00
diff --git a/django_mongodb_backend/indexes.py b/django_mongodb_backend/indexes.py
@@ -1,11 +1,26 @@
+import itertools
+
 from django.db import NotSupportedError
-from django.db.models import Index
+from django.db.models import (
+    BooleanField,
+    CharField,
+    DateField,
+    DateTimeField,
+    DecimalField,
+    FloatField,
+    Index,
+    IntegerField,
+    TextField,
+    UUIDField,
+)
 from django.db.models.lookups import BuiltinLookup
 from django.db.models.sql.query import Query
 from django.db.models.sql.where import AND, XOR, WhereNode
 from pymongo import ASCENDING, DESCENDING
 from pymongo.operations import IndexModel, SearchIndexModel
 
+from django_mongodb_backend.fields import ArrayField, ObjectIdAutoField, ObjectIdField
+
 from .query_utils import process_rhs
 
 MONGO_INDEX_OPERATORS = {
@@ -122,12 +137,81 @@ def create_mongodb_index(
         for field_name, _ in self.fields_orders:
             field_ = model._meta.get_field(field_name)
             type_ = connection.mongo_data_types[field_.get_internal_type()]
-            fields[field_name] = {"type": type_}
+            field_path = column_prefix + model._meta.get_field(field_name).column
+            fields[field_path] = {"type": type_}
         return SearchIndexModel(
             definition={"mappings": {"dynamic": False, "fields": fields}}, name=self.name
         )
 
 
+class AtlasVectorSearchIndex(Index):
+    suffix = "atlas_vector_search"
+    ALLOWED_SIMILARITY_FUNCTIONS = ("euclidean", "cosine", "dotProduct")
+
+    def __init__(self, *expressions, similarities="cosine", **kwargs):
+        super().__init__(*expressions, **kwargs)
+        # validate the similarities types
+        if isinstance(similarities, str):
+            self._check_similarity_functions([similarities])
+        else:
+            self._check_similarity_functions(similarities)
+        self.similarities = similarities
+
+    def _check_similarity_functions(self, similarities):
+        for func in similarities:
+            if func not in self.ALLOWED_SIMILARITY_FUNCTIONS:
+                raise ValueError(
+                    f"{func} isn't a valid similarity function, options "
+                    f"'are {','.join(self.ALLOWED_SIMILARITY_FUNCTIONS)}"
+                )
+
+    def create_mongodb_index(
+        self, model, schema_editor, connection=None, field=None, unique=False, column_prefix=""
+    ):
+        similarities = (
+            itertools.cycle([self.similarities])
+            if isinstance(self.similarities, str)
+            else iter(self.similarities)
+        )
+        fields = []
+        for field_name, _ in self.fields_orders:
+            field_ = model._meta.get_field(field_name)
+            field_path = column_prefix + model._meta.get_field(field_name).column
+            mappings = {"path": field_path}
+            if isinstance(field_, ArrayField):
+                try:
+                    vector_size = int(field_.size)
+                except (ValueError, TypeError) as err:
+                    raise ValueError("Atlas vector search requires fixed size.") from err
+                if not isinstance(field_.base_field, FloatField | DecimalField):
+                    raise ValueError("Base type must be Float or Decimal.")
+                mappings.update(
+                    {
+                        "type": "vector",
+                        "numDimensions": vector_size,
+                        "similarity": next(similarities),
+                    }
+                )
+            elif isinstance(
+                field_,
+                BooleanField
+                | IntegerField
+                | DateField
+                | DateTimeField
+                | CharField
+                | TextField
+                | UUIDField
+                | ObjectIdField
+                | ObjectIdAutoField,
+            ):
+                mappings["type"] = "filter"
+            else:
+                field_type = field_.get_internal_type()
+                raise ValueError(f"Unsupported filter of type {field_type}.")
+            fields.append(mappings)
+        return SearchIndexModel(definition={"fields": fields}, name=self.name, type="vectorSearch")
+
+
 def register_indexes():
     BuiltinLookup.as_mql_idx = builtin_lookup_idx
     Index._get_condition_mql = _get_condition_mql
diff --git a/django_mongodb_backend/introspection.py b/django_mongodb_backend/introspection.py
@@ -2,7 +2,7 @@
 from django.db.models import Index
 from pymongo import ASCENDING, DESCENDING
 
-from django_mongodb_backend.indexes import AtlasSearchIndex
+from django_mongodb_backend.indexes import AtlasSearchIndex, AtlasVectorSearchIndex
 
 
 class DatabaseIntrospection(BaseDatabaseIntrospection):
@@ -37,7 +37,14 @@ def _get_atlas_index_info(self, table_name):
         constraints = {}
         indexes = self.connection.get_collection(table_name).list_search_indexes()
         for details in indexes:
-            columns = list(details["latestDefinition"]["mappings"].get("fields", {}).keys())
+            if details["type"] == "vectorSearch":
+                columns = [field["path"] for field in details["latestDefinition"]["fields"]]
+                type_ = AtlasVectorSearchIndex.suffix
+                options = details
+            else:
+                columns = list(details["latestDefinition"]["mappings"].get("fields", {}).keys())
+                options = details["latestDefinition"]["mappings"]
+                type_ = AtlasSearchIndex.suffix
             constraints[details["name"]] = {
                 "check": False,
                 "columns": columns,
@@ -46,9 +53,9 @@ def _get_atlas_index_info(self, table_name):
                 "index": True,
                 "orders": [],
                 "primary_key": False,
-                "type": AtlasSearchIndex.suffix,
+                "type": type_,
                 "unique": False,
-                "options": details["latestDefinition"]["mappings"],
+                "options": options,
             }
         return constraints
 
diff --git a/django_mongodb_backend/schema.py b/django_mongodb_backend/schema.py
@@ -4,7 +4,7 @@
 from django.db.models import Index, UniqueConstraint
 from pymongo.operations import IndexModel, SearchIndexModel
 
-from django_mongodb_backend.indexes import AtlasSearchIndex
+from django_mongodb_backend.indexes import AtlasSearchIndex, AtlasVectorSearchIndex
 
 from .fields import EmbeddedModelField
 from .query import wrap_database_errors
@@ -310,7 +310,7 @@ def _(self, index: Index, model):
         return self.get_collection(model._meta.db_table).drop_index(index.name)
 
     @_remove_index.register
-    def _(self, index: AtlasSearchIndex, model):
+    def _(self, index: AtlasSearchIndex | AtlasVectorSearchIndex, model):
         return self.get_collection(model._meta.db_table).drop_search_index(index.name)
 
     @ignore_embedded_models
diff --git a/tests/indexes_/models.py b/tests/indexes_/models.py
@@ -1,6 +1,6 @@
 from django.db import models
 
-from django_mongodb_backend.fields import EmbeddedModelField
+from django_mongodb_backend.fields import ArrayField, EmbeddedModelField
 from django_mongodb_backend.models import EmbeddedModel
 
 
@@ -15,3 +15,7 @@ class Article(models.Model):
     data = models.JSONField()
     embedded = EmbeddedModelField(Data)
     auto_now = models.DateTimeField(auto_now=True)
+    title_embedded = ArrayField(models.FloatField(), size=10)
+    description_embedded = ArrayField(models.FloatField(), size=10)
+    number_list = ArrayField(models.FloatField())
+    name_list = ArrayField(models.CharField(max_length=30), size=10)
diff --git a/tests/indexes_/test_atlas_indexes.py b/tests/indexes_/test_atlas_indexes.py
@@ -4,7 +4,7 @@
 from django.db import connection
 from django.test import TestCase
 
-from django_mongodb_backend.indexes import AtlasSearchIndex
+from django_mongodb_backend.indexes import AtlasSearchIndex, AtlasVectorSearchIndex
 
 from .models import Article, Data
 
@@ -32,7 +32,7 @@ def assertAddRemoveIndex(self, editor, model, index):
             ),
         )
 
-    def test_simple_atlas_index(self):
+    def test_simple(self):
         with connection.schema_editor() as editor:
             index = AtlasSearchIndex(
                 name="recent_article_idx",
@@ -41,7 +41,7 @@ def test_simple_atlas_index(self):
             editor.add_index(index=index, model=Article)
             self.assertAddRemoveIndex(editor, Article, index)
 
-    def test_multiple_fields_atlas_index(self):
+    def test_multiple_fields(self):
         with connection.schema_editor() as editor:
             index = AtlasSearchIndex(
                 name="recent_article_idx",
@@ -106,6 +106,112 @@ def setUpTestData(cls):
                 data="{json: i}",
                 embedded=Data(integer=i),
                 auto_now=datetime.datetime.now(),
+                title_embedded=[0.1] * 10,
+                description_embedded=[2.5] * 10,
+                number_list=[2] * i,
+                name_list=[f"name_{i}"] * 10,
+            )
+            for i in range(5)
+        ]
+        cls.objs = Article.objects.bulk_create(articles)
+
+
+class AtlasSearchIndexTests(TestCase):
+    # Schema editor is used to  create the index to test that it works.
+    # available_apps = ["indexes"]
+    available_apps = None  # could be removed?
+
+    def assertAddRemoveIndex(self, editor, model, index):
+        editor.add_index(index=index, model=model)
+        self.assertIn(
+            index.name,
+            connection.introspection.get_constraints(
+                cursor=None,
+                table_name=model._meta.db_table,
+            ),
+        )
+        editor.remove_index(index=index, model=model)
+        self.assertNotIn(
+            index.name,
+            connection.introspection.get_constraints(
+                cursor=None,
+                table_name=model._meta.db_table,
+            ),
+        )
+
+    def test_simple_atlas_vector_search(self):
+        with connection.schema_editor() as editor:
+            index = AtlasVectorSearchIndex(
+                name="recent_article_idx",
+                fields=["number"],
+            )
+            editor.add_index(index=index, model=Article)
+            self.assertAddRemoveIndex(editor, Article, index)
+
+    def test_multiple_fields(self):
+        with connection.schema_editor() as editor:
+            index = AtlasVectorSearchIndex(
+                name="recent_article_idx",
+                fields=["headline", "number", "body", "description_embedded"],
+            )
+            editor.add_index(index=index, model=Article)
+            index_info = connection.introspection.get_constraints(
+                cursor=None,
+                table_name=Article._meta.db_table,
+            )
+            expected_options = {
+                "latestDefinition": {
+                    "fields": [
+                        {"path": "headline", "type": "filter"},
+                        {"path": "number", "type": "filter"},
+                        {"path": "body", "type": "filter"},
+                        {
+                            "numDimensions": 10,
+                            "path": "description_embedded",
+                            "similarity": "cosine",
+                            "type": "vector",
+                        },
+                    ]
+                },
+                "latestVersion": 0,
+                "name": "recent_article_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)
+            self.assertAddRemoveIndex(editor, Article, index)
+
+    def test_field_not_exists(self):
+        index = AtlasVectorSearchIndex(
+            name="recent_article_idx",
+            fields=["headline", "number1", "title_embedded"],
+        )
+        with connection.schema_editor() as editor:
+            msg = "Article has no field named 'number1'"
+            with self.assertRaisesMessage(
+                FieldDoesNotExist, msg
+            ), connection.schema_editor() as editor:
+                editor.add_index(index=index, model=Article)
+
+
+class AtlasSearchIndexTestsWithData(AtlasSearchIndexTests):
+    @classmethod
+    def setUpTestData(cls):
+        articles = [
+            Article(
+                headline=f"Title {i}",
+                number=i,
+                body=f"body {i}",
+                data="{json: i}",
+                embedded=Data(integer=i),
+                auto_now=datetime.datetime.now(),
+                title_embedded=[0.1] * 10,
+                description_embedded=[2.5] * 10,
+                number_list=[2] * i,
+                name_list=[f"name_{i}"] * 10,
             )
             for i in range(5)
         ]
