Merge branch 'pgvector:master' into master

Cdingram · web-flow · commit a9007a7f95bb · 2024-05-16T16:50:45.000-06:00
diff --git a/pgvector/django/__init__.py b/pgvector/django/__init__.py
@@ -4,6 +4,19 @@
 from .indexes import IvfflatIndex, HnswIndex
 from .sparsevec import SparsevecField
 from .vector import VectorField
-from ..utils import SparseVec
+from ..utils import HalfVec, SparseVec
 
-__all__ = ['VectorExtension', 'VectorField', 'HalfvecField', 'SparsevecField', 'IvfflatIndex', 'HnswIndex', 'L2Distance', 'MaxInnerProduct', 'CosineDistance', 'L1Distance']
+__all__ = [
+    'VectorExtension',
+    'VectorField',
+    'HalfvecField',
+    'SparsevecField',
+    'IvfflatIndex',
+    'HnswIndex',
+    'L2Distance',
+    'MaxInnerProduct',
+    'CosineDistance',
+    'L1Distance',
+    'HalfVec',
+    'SparseVec'
+]
diff --git a/pgvector/django/functions.py b/pgvector/django/functions.py
@@ -1,13 +1,18 @@
 from django.db.models import FloatField, Func, Value
-from ..utils import Vector
+from ..utils import Vector, HalfVec, SparseVec
 
 
 class DistanceBase(Func):
     output_field = FloatField()
 
     def __init__(self, expression, vector, **extra):
         if not hasattr(vector, 'resolve_expression'):
-            vector = Value(Vector.to_db(vector))
+            if isinstance(vector, HalfVec):
+                vector = Value(HalfVec.to_db(vector))
+            elif isinstance(vector, SparseVec):
+                vector = Value(SparseVec.to_db(vector))
+            else:
+                vector = Value(Vector.to_db(vector))
         super().__init__(expression, vector, **extra)
 
 
diff --git a/tests/test_django.py b/tests/test_django.py
@@ -8,7 +8,7 @@
 from math import sqrt
 import numpy as np
 import pgvector.django
-from pgvector.django import VectorExtension, VectorField, HalfvecField, SparsevecField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance, SparseVec
+from pgvector.django import VectorExtension, VectorField, HalfvecField, SparsevecField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HalfVec, SparseVec
 from unittest import mock
 
 settings.configure(
@@ -92,7 +92,7 @@ def create_items():
         [1, 1, 2]
     ]
     for i, v in enumerate(vectors):
-        item = Item(id=i + 1, embedding=v)
+        item = Item(id=i + 1, embedding=v, half_embedding=v, sparse_embedding=SparseVec.from_dense(v))
         item.save()
 
 
@@ -156,6 +156,20 @@ def test_l1_distance(self):
         assert [v.id for v in items] == [1, 3, 2]
         assert [v.distance for v in items] == [0, 1, 3]
 
+    def test_halfvec_l2_distance(self):
+        create_items()
+        distance = L2Distance('half_embedding', HalfVec([1, 1, 1]))
+        items = Item.objects.annotate(distance=distance).order_by(distance)
+        assert [v.id for v in items] == [1, 3, 2]
+        assert [v.distance for v in items] == [0, 1, sqrt(3)]
+
+    def test_sparsevec_l2_distance(self):
+        create_items()
+        distance = L2Distance('sparse_embedding', SparseVec.from_dense([1, 1, 1]))
+        items = Item.objects.annotate(distance=distance).order_by(distance)
+        assert [v.id for v in items] == [1, 3, 2]
+        assert [v.distance for v in items] == [0, 1, sqrt(3)]
+
     def test_filter(self):
         create_items()
         distance = L2Distance('embedding', [1, 1, 1])
