Changed vector type to return Vector object instead of NumPy array - closes #99

Author: ankane

CHANGELOG.md

@@ -1,6 +1,7 @@
 ## 0.5.0 (unreleased)
 
 - Added type hints
+- Changed `vector` type to return `Vector` object instead of NumPy array
 - Removed `utils` package (use top-level `pgvector` package instead)
 - Removed re-exported classes (use top-level `pgvector` package instead)
 - Dropped support for Python < 3.10

pgvector/django/vector.py

@@ -1,6 +1,5 @@
 from django import forms
 from django.db.models import Field
-import numpy as np
 from typing import Any
 from .. import Vector
 
@@ -25,30 +24,23 @@ 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) -> np.ndarray | None:
+    def from_db_value(self, value: Any, expression: Any, connection: Any) -> Vector | None:
         return Vector._from_db(value)
 
-    def to_python(self, value: Any) -> np.ndarray | None:
-        if isinstance(value, list):
-            return np.array(value, dtype=np.float32)
-        return Vector._from_db(value)
+    def to_python(self, value: Any) -> Vector | None:
+        if value is None or isinstance(value, Vector):
+            return value
+        elif isinstance(value, str):
+            return Vector._from_db(value)
+        else:
+            return Vector(value)
 
     def get_prep_value(self, value: Any) -> str | None:
         return Vector._to_db(value)
 
     def value_to_string(self, obj: Any) -> str | None:
         return self.get_prep_value(self.value_from_object(obj))
 
-    def validate(self, value: Any, model_instance: Any) -> None:
-        if isinstance(value, np.ndarray):
-            value = value.tolist()
-        super().validate(value, model_instance)
-
-    def run_validators(self, value: Any) -> None:
-        if isinstance(value, np.ndarray):
-            value = value.tolist()
-        super().run_validators(value)
-
     def formfield(self, form_class: Any = None, choices_form_class: Any = None, **kwargs: Any) -> forms.Field:
         return super().formfield(
             form_class=VectorFormField if form_class is None else form_class,
@@ -59,19 +51,14 @@ def formfield(self, form_class: Any = None, choices_form_class: Any = None, **kw
 
 class VectorWidget(forms.TextInput):
     def format_value(self, value: Any) -> str | None:
-        if isinstance(value, np.ndarray):
-            value = value.tolist()
+        if isinstance(value, Vector):
+            value = value.to_list()
         return super().format_value(value)
 
 
 class VectorFormField(forms.CharField):
     widget = VectorWidget
 
-    def has_changed(self, initial: Any, data: Any) -> bool:
-        if isinstance(initial, np.ndarray):
-            initial = initial.tolist()
-        return super().has_changed(initial, data)
-
     def to_python(self, value: Any) -> Any:
         if isinstance(value, str) and value == '':
             return None

pgvector/peewee/vector.py

@@ -1,4 +1,3 @@
-import numpy as np
 from peewee import Expression, Field
 from typing import Any
 from .. import Vector
@@ -17,7 +16,7 @@ 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) -> np.ndarray | None:
+    def python_value(self, value: Any) -> Vector | None:
         return Vector._from_db(value)
 
     def _distance(self, op: str, vector: object) -> Expression:

pgvector/psycopg/vector.py

@@ -31,7 +31,7 @@ class VectorLoader(Loader):
 
     format = Format.TEXT
 
-    def load(self, data: Buffer) -> np.ndarray | None:
+    def load(self, data: Buffer) -> Vector | None:
         if isinstance(data, memoryview):
             data = bytes(data)
         return Vector._from_db(data.decode('utf8'))
@@ -41,7 +41,7 @@ class VectorBinaryLoader(VectorLoader):
 
     format = Format.BINARY
 
-    def load(self, data: Buffer) -> np.ndarray | None:
+    def load(self, data: Buffer) -> Vector | None:
         if isinstance(data, (bytearray, memoryview)):
             data = bytes(data)
         return Vector._from_db_binary(data)

pgvector/psycopg2/vector.py

@@ -11,7 +11,7 @@ def getquoted(self) -> bytes:
         return adapt(Vector._to_db(self._value)).getquoted()
 
 
-def cast_vector(value: str | None, cur: cursor) -> np.ndarray | None:
+def cast_vector(value: str | None, cur: cursor) -> Vector | None:
     return Vector._from_db(value)
 
 

pgvector/sqlalchemy/vector.py

@@ -1,4 +1,3 @@
-import numpy as np
 from sqlalchemy.dialects.postgresql.base import ischema_names
 from sqlalchemy.types import UserDefinedType, Float, String
 from sqlalchemy import Dialect, Operators
@@ -32,7 +31,7 @@ def process(value: Any) -> Any:
         return process
 
     def result_processor(self, dialect: Dialect, coltype: Any) -> Any:
-        def process(value: Any) -> np.ndarray | None:
+        def process(value: Any) -> Vector | None:
             return Vector._from_db(value)
         return process
 

pgvector/vector.py

@@ -82,15 +82,15 @@ def _to_db_binary(cls, value: object) -> bytes | None:
         return value.to_binary()
 
     @classmethod
-    def _from_db(cls, value: str | np.ndarray | None) -> np.ndarray | None:
-        if value is None or isinstance(value, np.ndarray):
+    def _from_db(cls, value: str | Vector | None) -> Vector | None:
+        if value is None or isinstance(value, Vector):
             return value
 
-        return cls.from_text(value).to_numpy().astype(np.float32)
+        return cls.from_text(value)
 
     @classmethod
-    def _from_db_binary(cls, value: bytes | np.ndarray | None) -> np.ndarray | None:
-        if value is None or isinstance(value, np.ndarray):
+    def _from_db_binary(cls, value: bytes | Vector | None) -> Vector | None:
+        if value is None or isinstance(value, Vector):
             return value
 
-        return cls.from_binary(value).to_numpy().astype(np.float32)
+        return cls.from_binary(value)

tests/test_asyncpg.py

@@ -20,9 +20,8 @@ async def test_vector(self):
         await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), (NULL)", embedding, embedding2)
 
         res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id")
-        assert np.array_equal(res[0]['embedding'], embedding.to_numpy())
-        assert res[0]['embedding'].dtype == np.float32
-        assert np.array_equal(res[1]['embedding'], embedding2)
+        assert res[0]['embedding'] == embedding
+        assert res[1]['embedding'] == Vector(embedding2)
         assert res[2]['embedding'] is None
 
         # ensures binary format is correct
@@ -116,10 +115,8 @@ async def test_vector_array(self):
         await conn.execute("INSERT INTO asyncpg_items (embeddings) VALUES (ARRAY[$1, $2]::vector[])", embeddings2[0], embeddings2[1])
 
         res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id")
-        assert np.array_equal(res[0]['embeddings'][0], embeddings[0].to_numpy())
-        assert np.array_equal(res[0]['embeddings'][1], embeddings[1].to_numpy())
-        assert np.array_equal(res[1]['embeddings'][0], embeddings2[0])
-        assert np.array_equal(res[1]['embeddings'][1], embeddings2[1])
+        assert res[0]['embeddings'] == embeddings
+        assert res[1]['embeddings'] == [Vector(e) for e in embeddings2]
 
         await conn.close()
 
@@ -140,7 +137,6 @@ async def init(conn):
             await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), (NULL)", embedding, embedding2)
 
             res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id")
-            assert np.array_equal(res[0]['embedding'], embedding.to_numpy())
-            assert res[0]['embedding'].dtype == np.float32
-            assert np.array_equal(res[1]['embedding'], embedding2)
+            assert res[0]['embedding'] == embedding
+            assert res[1]['embedding'] == Vector(embedding2)
             assert res[2]['embedding'] is None

tests/test_django.py

@@ -14,7 +14,7 @@
 import numpy as np
 import os
 import pgvector.django
-from pgvector import HalfVector, SparseVector
+from pgvector import HalfVector, SparseVector, Vector
 from pgvector.django import VectorExtension, VectorField, HalfVectorField, BitField, SparseVectorField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, JaccardDistance
 from unittest import mock
 
@@ -167,12 +167,11 @@ def setup_method(self):
     def test_vector(self):
         Item(id=1, embedding=[1, 2, 3]).save()
         item = Item.objects.get(pk=1)
-        assert np.array_equal(item.embedding, [1, 2, 3])
-        assert item.embedding.dtype == np.float32
+        assert item.embedding == Vector([1, 2, 3])
 
     def test_vector_l2_distance(self):
         create_items()
-        distance = L2Distance('embedding', [1, 1, 1])
+        distance = L2Distance('embedding', Vector([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)]
@@ -295,15 +294,15 @@ def test_vector_avg(self):
         Item(embedding=[1, 2, 3]).save()
         Item(embedding=[4, 5, 6]).save()
         avg = Item.objects.aggregate(Avg('embedding'))['embedding__avg']
-        assert np.array_equal(avg, [2.5, 3.5, 4.5])
+        assert avg == Vector([2.5, 3.5, 4.5])
 
     def test_vector_sum(self):
         sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum']
         assert sum is None
         Item(embedding=[1, 2, 3]).save()
         Item(embedding=[4, 5, 6]).save()
         sum = Item.objects.aggregate(Sum('embedding'))['embedding__sum']
-        assert np.array_equal(sum, [5, 7, 9])
+        assert sum == Vector([5, 7, 9])
 
     def test_halfvec_avg(self):
         avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg']
@@ -349,7 +348,7 @@ def test_vector_form_save(self):
         assert form.has_changed()
         assert form.is_valid()
         assert form.save()
-        assert np.array_equal(Item.objects.get(pk=1).embedding, [4, 5, 6])
+        assert Item.objects.get(pk=1).embedding == Vector([4, 5, 6])
 
     def test_vector_form_save_missing(self):
         Item(id=1).save()
@@ -467,8 +466,7 @@ def test_vector_array(self):
 
             # this fails if the driver does not cast arrays
             item = Item.objects.get(pk=1)
-            assert np.array_equal(item.embeddings[0], [1, 2, 3])
-            assert np.array_equal(item.embeddings[1], [4, 5, 6])
+            assert item.embeddings == [Vector([1, 2, 3]), Vector([4, 5, 6])]
 
     def test_double_array(self):
         Item(id=1, double_embedding=[1, 1, 1]).save()

tests/test_peewee.py

@@ -1,7 +1,7 @@
 from math import sqrt
 import numpy as np
 from peewee import Model, PostgresqlDatabase, fn
-from pgvector import HalfVector, SparseVector
+from pgvector import HalfVector, SparseVector, Vector
 from pgvector.peewee import VectorField, HalfVectorField, FixedBitField, SparseVectorField
 
 db = PostgresqlDatabase('pgvector_python_test')
@@ -43,8 +43,7 @@ def setup_method(self):
     def test_vector(self):
         Item.create(id=1, embedding=[1, 2, 3])
         item = Item.get_by_id(1)
-        assert np.array_equal(item.embedding, [1, 2, 3])
-        assert item.embedding.dtype == np.float32
+        assert item.embedding == Vector([1, 2, 3])
 
     def test_vector_l2_distance(self):
         create_items()
@@ -170,15 +169,15 @@ def test_vector_avg(self):
         Item.create(embedding=[1, 2, 3])
         Item.create(embedding=[4, 5, 6])
         avg = Item.select(fn.avg(Item.embedding).coerce(True)).scalar()
-        assert np.array_equal(avg, [2.5, 3.5, 4.5])
+        assert avg == Vector([2.5, 3.5, 4.5])
 
     def test_vector_sum(self):
         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 np.array_equal(sum, [5, 7, 9])
+        assert sum == Vector([5, 7, 9])
 
     def test_halfvec_avg(self):
         avg = Item.select(fn.avg(Item.half_embedding).coerce(True)).scalar()