Replace decimal in vector index for integer

Author: WaVEV

django_mongodb_backend/indexes.py

@@ -3,7 +3,7 @@
 
 from django.core.checks import Error, Warning
 from django.db import NotSupportedError
-from django.db.models import DecimalField, FloatField, Index
+from django.db.models import FloatField, Index, IntegerField
 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
@@ -200,12 +200,12 @@ def check(self, model, connection):
                             id=f"{self._error_id_prefix}.E002",
                         )
                     )
-                if not isinstance(field_.base_field, FloatField | DecimalField):
+                if not isinstance(field_.base_field, FloatField | IntegerField):
                     errors.append(
                         Error(
                             "VectorSearchIndex requires the base field of "
                             f"ArrayField '{field_.name}' to be FloatField or "
-                            "DecimalField but is "
+                            "IntegerField but is "
                             f"{field_.base_field.get_internal_type()}.",
                             obj=model,
                             id=f"{self._error_id_prefix}.E003",

tests/indexes_/models.py

@@ -15,4 +15,4 @@ class Article(models.Model):
     data = models.JSONField()
     embedded = EmbeddedModelField(Data)
     created_at = models.DateTimeField(auto_now=True)
-    description_semantic = ArrayField(models.DecimalField(decimal_places=3, max_digits=10), size=10)
+    description_semantic = ArrayField(models.IntegerField(), size=10)

tests/indexes_/test_checks.py

@@ -97,7 +97,7 @@ class Meta:
             [
                 checks.Error(
                     "VectorSearchIndex requires the base field of ArrayField "
-                    "'title_embedded' to be FloatField or DecimalField but is CharField.",
+                    "'title_embedded' to be FloatField or IntegerField but is CharField.",
                     id="django_mongodb_backend.indexes.VectorSearchIndex.E003",
                     obj=Article,
                 )
@@ -125,7 +125,7 @@ class Meta:
 
     def test_invalid_number_similarity_function_singular(self):
         class Article(models.Model):
-            vector_data = ArrayField(models.DecimalField(), size=10)
+            vector_data = ArrayField(models.FloatField(), size=10)
 
             class Meta:
                 indexes = [
@@ -150,8 +150,8 @@ class Meta:
 
     def test_invalid_number_similarity_function_plural(self):
         class Article(models.Model):
-            vector1 = ArrayField(models.DecimalField(), size=10)
-            vector2 = ArrayField(models.DecimalField(), size=10)
+            vector1 = ArrayField(models.FloatField(), size=10)
+            vector2 = ArrayField(models.FloatField(), size=10)
 
             class Meta:
                 indexes = [
@@ -176,7 +176,7 @@ class Meta:
 
     def test_simple(self):
         class Article(models.Model):
-            vector_data = ArrayField(models.DecimalField(), size=10)
+            vector_data = ArrayField(models.FloatField(), size=10)
 
             class Meta:
                 indexes = [VectorSearchIndex(fields=["vector_data"])]