Changed vector and halfvec types to return list for Django, SQLAlchemy, SQLModel, and Peewee

Author: ankane

CHANGELOG.md

@@ -1,7 +1,8 @@
 ## 1.0.0 (unreleased)
 
 - Added experimental support for type hints
-- Changed `vector` type to return `Vector` object instead of NumPy array
+- Changed `vector` and `halfvec` types to return list for Django, SQLAlchemy, SQLModel, and Peewee
+- Changed `vector` type to return `Vector` object for Psycopg 3, Psycopg 2, asyncpg, and pg8000
 - Removed `utils` package (use top-level `pgvector` package instead)
 - Removed re-exported classes (use top-level `pgvector` package instead)
 - Removed dependency on NumPy

pgvector/django/halfvec.py

@@ -24,16 +24,18 @@ def db_type(self, connection: Any) -> str:
             return 'halfvec'
         return 'halfvec(%d)' % self.dimensions
 
-    def from_db_value(self, value: Any, expression: Any, connection: Any) -> HalfVector | None:
-        return HalfVector._from_db(value)
+    def from_db_value(self, value: Any, expression: Any, connection: Any) -> list[float] | None:
+        value = HalfVector._from_db(value)
+        return None if value is None else value.to_list()
 
-    def to_python(self, value: Any) -> HalfVector | None:
-        if value is None or isinstance(value, HalfVector):
+    def to_python(self, value: Any) -> list[float] | None:
+        if value is None or isinstance(value, list):
             return value
-        elif isinstance(value, str):
-            return HalfVector.from_text(value)
-        else:
-            return HalfVector(value)
+        if isinstance(value, str):
+            value = HalfVector.from_text(value)
+        if not isinstance(value, HalfVector):
+            value = HalfVector(value)
+        return value.to_list()
 
     def get_prep_value(self, value: Any) -> str | None:
         return HalfVector._to_db(value)
@@ -50,16 +52,7 @@ def formfield(self, form_class: Any = None, choices_form_class: Any = None, **kw
         )
 
 
-class HalfVectorWidget(forms.TextInput):
-    def format_value(self, value: Any) -> Any:
-        if isinstance(value, HalfVector):
-            value = value.to_list()
-        return super().format_value(value)
-
-
 class HalfVectorFormField(forms.CharField):
-    widget = HalfVectorWidget
-
     def to_python(self, value: Any) -> Any:
         if isinstance(value, str) and value == '':
             return None

pgvector/django/vector.py

@@ -24,16 +24,18 @@ def db_type(self, connection: Any) -> str:
             return 'vector'
         return 'vector(%d)' % self.dimensions
 
-    def from_db_value(self, value: Any, expression: Any, connection: Any) -> Vector | None:
-        return Vector._from_db(value)
+    def from_db_value(self, value: Any, expression: Any, connection: Any) -> list[float] | None:
+        value = Vector._from_db(value)
+        return None if value is None else value.to_list()
 
-    def to_python(self, value: Any) -> Vector | None:
-        if value is None or isinstance(value, Vector):
+    def to_python(self, value: Any) -> list[float] | None:
+        if value is None or isinstance(value, list):
             return value
-        elif isinstance(value, str):
-            return Vector.from_text(value)
-        else:
-            return Vector(value)
+        if isinstance(value, str):
+            value = Vector.from_text(value)
+        if not isinstance(value, Vector):
+            value = Vector(value)
+        return value.to_list()
 
     def get_prep_value(self, value: Any) -> str | None:
         return Vector._to_db(value)
@@ -50,16 +52,7 @@ def formfield(self, form_class: Any = None, choices_form_class: Any = None, **kw
         )
 
 
-class VectorWidget(forms.TextInput):
-    def format_value(self, value: Any) -> Any:
-        if isinstance(value, Vector):
-            value = value.to_list()
-        return super().format_value(value)
-
-
 class VectorFormField(forms.CharField):
-    widget = VectorWidget
-
     def to_python(self, value: Any) -> Any:
         if isinstance(value, str) and value == '':
             return None

pgvector/peewee/halfvec.py

@@ -16,8 +16,9 @@ def get_modifiers(self) -> list[int] | None:
     def db_value(self, value: object) -> str | None:
         return HalfVector._to_db(value)
 
-    def python_value(self, value: Any) -> HalfVector | None:
-        return HalfVector._from_db(value)
+    def python_value(self, value: Any) -> list[float] | None:
+        value = HalfVector._from_db(value)
+        return None if value is None else value.to_list()
 
     def _distance(self, op: str, vector: object) -> Expression:
         return Expression(lhs=self, op=op, rhs=self.to_value(vector))

pgvector/peewee/vector.py

@@ -16,8 +16,9 @@ def get_modifiers(self) -> list[int] | None:
     def db_value(self, value: object) -> str | None:
         return Vector._to_db(value)
 
-    def python_value(self, value: Any) -> Vector | None:
-        return Vector._from_db(value)
+    def python_value(self, value: Any) -> list[float] | None:
+        value = Vector._from_db(value)
+        return None if value is None else value.to_list()
 
     def _distance(self, op: str, vector: object) -> Expression:
         return Expression(lhs=self, op=op, rhs=self.to_value(vector))

pgvector/sqlalchemy/halfvec.py

@@ -31,8 +31,9 @@ def process(value: Any) -> Any:
         return process
 
     def result_processor(self, dialect: Dialect, coltype: Any) -> Any:
-        def process(value: Any) -> HalfVector | None:
-            return HalfVector._from_db(value)
+        def process(value: Any) -> list[float] | None:
+            value = HalfVector._from_db(value)
+            return None if value is None else value.to_list()
         return process
 
     class Comparator(TypeEngine.Comparator[Any]):

pgvector/sqlalchemy/vector.py

@@ -31,8 +31,9 @@ def process(value: Any) -> Any:
         return process
 
     def result_processor(self, dialect: Dialect, coltype: Any) -> Any:
-        def process(value: Any) -> Vector | None:
-            return Vector._from_db(value)
+        def process(value: Any) -> list[float] | None:
+            value = Vector._from_db(value)
+            return None if value is None else value.to_list()
         return process
 
     class Comparator(TypeEngine.Comparator[Any]):

tests/test_django.py

@@ -11,7 +11,7 @@
 from math import sqrt
 import os
 import pgvector.django
-from pgvector import HalfVector, SparseVector, Vector
+from pgvector import SparseVector
 from pgvector.django import VectorExtension, VectorField, HalfVectorField, BitField, SparseVectorField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, JaccardDistance
 from unittest import mock
 
@@ -164,11 +164,11 @@ def setup_method(self) -> None:
     def test_vector(self) -> None:
         Item(id=1, embedding=[1, 2, 3]).save()
         item = Item.objects.get(pk=1)  # type: ignore
-        assert item.embedding == Vector([1, 2, 3])
+        assert item.embedding == [1, 2, 3]
 
     def test_vector_l2_distance(self) -> None:
         create_items()
-        distance = L2Distance('embedding', Vector([1, 1, 1]))
+        distance = L2Distance('embedding', [1, 1, 1])
         items = Item.objects.annotate(distance=distance).order_by(distance)  # type: ignore
         assert [v.id for v in items] == [1, 3, 2]
         assert [v.distance for v in items] == [0, 1, sqrt(3)]
@@ -210,32 +210,32 @@ def test_vector_l1_distance(self) -> None:
     def test_halfvec(self) -> None:
         Item(id=1, half_embedding=[1, 2, 3]).save()
         item = Item.objects.get(pk=1)  # type: ignore
-        assert item.half_embedding == HalfVector([1, 2, 3])
+        assert item.half_embedding == [1, 2, 3]
 
     def test_halfvec_l2_distance(self) -> None:
         create_items()
-        distance = L2Distance('half_embedding', HalfVector([1, 1, 1]))
+        distance = L2Distance('half_embedding', [1, 1, 1])
         items = Item.objects.annotate(distance=distance).order_by(distance)  # type: ignore
         assert [v.id for v in items] == [1, 3, 2]
         assert [v.distance for v in items] == [0, 1, sqrt(3)]
 
     def test_halfvec_max_inner_product(self) -> None:
         create_items()
-        distance = MaxInnerProduct('half_embedding', HalfVector([1, 1, 1]))
+        distance = MaxInnerProduct('half_embedding', [1, 1, 1])
         items = Item.objects.annotate(distance=distance).order_by(distance)  # type: ignore
         assert [v.id for v in items] == [2, 3, 1]
         assert [v.distance for v in items] == [-6, -4, -3]
 
     def test_halfvec_cosine_distance(self) -> None:
         create_items()
-        distance = CosineDistance('half_embedding', HalfVector([1, 1, 1]))
+        distance = CosineDistance('half_embedding', [1, 1, 1])
         items = Item.objects.annotate(distance=distance).order_by(distance)  # type: ignore
         assert [v.id for v in items] == [1, 2, 3]
         assert [v.distance for v in items] == [0, 0, 0.05719095841793653]
 
     def test_halfvec_l1_distance(self) -> None:
         create_items()
-        distance = L1Distance('half_embedding', HalfVector([1, 1, 1]))
+        distance = L1Distance('half_embedding', [1, 1, 1])
         items = Item.objects.annotate(distance=distance).order_by(distance)  # type: ignore
         assert [v.id for v in items] == [1, 3, 2]
         assert [v.distance for v in items] == [0, 1, 3]
@@ -304,31 +304,31 @@ def test_vector_avg(self) -> None:
         Item(embedding=[1, 2, 3]).save()
         Item(embedding=[4, 5, 6]).save()
         avg = Item.objects.aggregate(Avg('embedding'))['embedding__avg']  # type: ignore
-        assert avg == Vector([2.5, 3.5, 4.5])
+        assert avg == [2.5, 3.5, 4.5]
 
     def test_vector_sum(self) -> None:
         sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum']  # type: ignore
         assert sum is None
         Item(embedding=[1, 2, 3]).save()
         Item(embedding=[4, 5, 6]).save()
         sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum']  # type: ignore
-        assert sum == Vector([5, 7, 9])
+        assert sum == [5, 7, 9]
 
     def test_halfvec_avg(self) -> None:
         avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg']  # type: ignore
         assert avg is None
         Item(half_embedding=[1, 2, 3]).save()
         Item(half_embedding=[4, 5, 6]).save()
         avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg']  # type: ignore
-        assert avg == HalfVector([2.5, 3.5, 4.5])
+        assert avg == [2.5, 3.5, 4.5]
 
     def test_halfvec_sum(self) -> None:
         sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum']  # type: ignore
         assert sum is None
         Item(half_embedding=[1, 2, 3]).save()
         Item(half_embedding=[4, 5, 6]).save()
         sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum']  # type: ignore
-        assert sum == HalfVector([5, 7, 9])
+        assert sum == [5, 7, 9]
 
     def test_serialization(self) -> None:
         create_items()
@@ -358,7 +358,7 @@ def test_vector_form_save(self) -> None:
         assert form.has_changed()
         assert form.is_valid()
         assert form.save()
-        assert Item.objects.get(pk=1).embedding == Vector([4, 5, 6])  # type: ignore
+        assert Item.objects.get(pk=1).embedding == [4, 5, 6]  # type: ignore
 
     def test_vector_form_save_missing(self) -> None:
         Item(id=1).save()
@@ -386,7 +386,7 @@ def test_halfvec_form_save(self) -> None:
         assert form.has_changed()
         assert form.is_valid()
         assert form.save()
-        assert Item.objects.get(pk=1).half_embedding == HalfVector([4, 5, 6])  # type: ignore
+        assert Item.objects.get(pk=1).half_embedding == [4, 5, 6]  # type: ignore
 
     def test_halfvec_form_save_missing(self) -> None:
         Item(id=1).save()
@@ -468,15 +468,15 @@ def test_missing(self) -> None:
         assert Item.objects.first().sparse_embedding is None  # type: ignore
 
     def test_vector_array(self) -> None:
-        Item(id=1, embeddings=[Vector([1, 2, 3]), Vector([4, 5, 6])]).save()
+        Item(id=1, embeddings=[[1, 2, 3], [4, 5, 6]]).save()
 
         with connection.cursor() as cursor:
             from pgvector.psycopg import register_vector
             register_vector(cursor.connection)
 
             # this fails if the driver does not cast arrays
             item = Item.objects.get(pk=1)  # type: ignore
-            assert item.embeddings == [Vector([1, 2, 3]), Vector([4, 5, 6])]
+            assert item.embeddings == [[1, 2, 3], [4, 5, 6]]
 
     def test_double_array(self) -> None:
         Item(id=1, double_embedding=[1, 1, 1]).save()

tests/test_peewee.py

@@ -1,6 +1,6 @@
 from math import sqrt
 from peewee import Model, PostgresqlDatabase, fn
-from pgvector import HalfVector, SparseVector, Vector
+from pgvector import SparseVector, Vector
 from pgvector.peewee import VectorField, HalfVectorField, FixedBitField, SparseVectorField
 from typing import Any
 
@@ -44,7 +44,7 @@ def setup_method(self) -> None:
     def test_vector(self) -> None:
         Item.create(id=1, embedding=[1, 2, 3])
         item = Item.get_by_id(1)
-        assert item.embedding == Vector([1, 2, 3])
+        assert item.embedding == [1, 2, 3]
 
     def test_vector_l2_distance(self) -> None:
         create_items()
@@ -77,7 +77,7 @@ def test_vector_l1_distance(self) -> None:
     def test_halfvec(self) -> None:
         Item.create(id=1, half_embedding=[1, 2, 3])
         item = Item.get_by_id(1)
-        assert item.half_embedding == HalfVector([1, 2, 3])
+        assert item.half_embedding == [1, 2, 3]
 
     def test_halfvec_l2_distance(self) -> None:
         create_items()
@@ -170,31 +170,31 @@ def test_vector_avg(self) -> None:
         Item.create(embedding=[1, 2, 3])
         Item.create(embedding=[4, 5, 6])
         avg = Item.select(fn.avg(Item.embedding).coerce(True)).scalar()
-        assert avg == Vector([2.5, 3.5, 4.5])
+        assert avg == [2.5, 3.5, 4.5]
 
     def test_vector_sum(self) -> None:
         sum = Item.select(fn.sum(Item.embedding).coerce(True)).scalar()
         assert sum is None
         Item.create(embedding=[1, 2, 3])
         Item.create(embedding=[4, 5, 6])
         sum = Item.select(fn.sum(Item.embedding).coerce(True)).scalar()
-        assert sum == Vector([5, 7, 9])
+        assert sum == [5, 7, 9]
 
     def test_halfvec_avg(self) -> None:
         avg = Item.select(fn.avg(Item.half_embedding).coerce(True)).scalar()
         assert avg is None
         Item.create(half_embedding=[1, 2, 3])
         Item.create(half_embedding=[4, 5, 6])
         avg = Item.select(fn.avg(Item.half_embedding).coerce(True)).scalar()
-        assert avg == HalfVector([2.5, 3.5, 4.5])
+        assert avg == [2.5, 3.5, 4.5]
 
     def test_halfvec_sum(self) -> None:
         sum = Item.select(fn.sum(Item.half_embedding).coerce(True)).scalar()
         assert sum is None
         Item.create(half_embedding=[1, 2, 3])
         Item.create(half_embedding=[4, 5, 6])
         sum = Item.select(fn.sum(Item.half_embedding).coerce(True)).scalar()
-        assert sum == HalfVector([5, 7, 9])
+        assert sum == [5, 7, 9]
 
     def test_get_or_create(self) -> None:
         Item.get_or_create(id=1, defaults={'embedding': [1, 2, 3]})