Changed vector type to return Vector class instead of NumPy array [skip ci]

Author: ankane

pgvector/django/vector.py

@@ -1,6 +1,5 @@
 from django import forms
 from django.db.models import Field
-import numpy as np
 from .. import Vector
 
 
@@ -28,45 +27,33 @@ def from_db_value(self, value, expression, connection):
         return Vector._from_db(value)
 
     def to_python(self, value):
-        if isinstance(value, list):
-            return np.array(value, dtype=np.float32)
-        return Vector._from_db(value)
+        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):
         return Vector._to_db(value)
 
     def value_to_string(self, obj):
         return self.get_prep_value(self.value_from_object(obj))
 
-    def validate(self, value, model_instance):
-        if isinstance(value, np.ndarray):
-            value = value.tolist()
-        super().validate(value, model_instance)
-
-    def run_validators(self, value):
-        if isinstance(value, np.ndarray):
-            value = value.tolist()
-        super().run_validators(value)
-
     def formfield(self, **kwargs):
         return super().formfield(form_class=VectorFormField, **kwargs)
 
 
 class VectorWidget(forms.TextInput):
     def format_value(self, value):
-        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, data):
-        if isinstance(initial, np.ndarray):
-            initial = initial.tolist()
-        return super().has_changed(initial, data)
-
     def to_python(self, value):
         if isinstance(value, str) and value == '':
             return None

pgvector/vector.py

@@ -70,14 +70,14 @@ def _to_db_binary(cls, value):
 
     @classmethod
     def _from_db(cls, value):
-        if value is None or isinstance(value, np.ndarray):
+        if value is None or isinstance(value, cls):
             return value
 
-        return cls.from_text(value).to_numpy().astype(np.float32)
+        return cls.from_text(value)
 
     @classmethod
     def _from_db_binary(cls, value):
-        if value is None or isinstance(value, np.ndarray):
+        if value is None or isinstance(value, cls):
             return value
 
-        return cls.from_binary(value).to_numpy().astype(np.float32)
+        return cls.from_binary(value)

tests/test_asyncpg.py

@@ -1,6 +1,6 @@
 import asyncpg
 import numpy as np
-from pgvector import SparseVector
+from pgvector import Vector, HalfVector, SparseVector
 from pgvector.asyncpg import register_vector
 import pytest
 
@@ -15,12 +15,11 @@ async def test_vector(self):
 
         await register_vector(conn)
 
-        embedding = np.array([1.5, 2, 3])
+        embedding = Vector([1.5, 2, 3])
         await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding)
 
         res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id")
-        assert np.array_equal(res[0]['embedding'], embedding)
-        assert res[0]['embedding'].dtype == np.float32
+        assert res[0]['embedding'] == embedding
         assert res[1]['embedding'] is None
 
         # ensures binary format is correct
@@ -38,11 +37,11 @@ async def test_halfvec(self):
 
         await register_vector(conn)
 
-        embedding = [1.5, 2, 3]
+        embedding = HalfVector([1.5, 2, 3])
         await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding)
 
         res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id")
-        assert res[0]['embedding'].to_list() == [1.5, 2, 3]
+        assert res[0]['embedding'] == embedding
         assert res[1]['embedding'] is None
 
         # ensures binary format is correct
@@ -87,7 +86,7 @@ async def test_sparsevec(self):
         await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding)
 
         res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id")
-        assert res[0]['embedding'].to_list() == [1.5, 2, 3]
+        assert res[0]['embedding'] == embedding
         assert res[1]['embedding'] is None
 
         # ensures binary format is correct
@@ -105,12 +104,12 @@ async def test_vector_array(self):
 
         await register_vector(conn)
 
-        embeddings = [np.array([1.5, 2, 3]), np.array([4.5, 5, 6])]
+        embeddings = [Vector([1.5, 2, 3]), Vector([4.5, 5, 6])]
         await conn.execute("INSERT INTO asyncpg_items (embeddings) VALUES (ARRAY[$1, $2]::vector[])", embeddings[0], embeddings[1])
 
         res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id")
-        assert np.array_equal(res[0]['embeddings'][0], embeddings[0])
-        assert np.array_equal(res[0]['embeddings'][1], embeddings[1])
+        assert res[0]['embeddings'][0] == embeddings[0]
+        assert res[0]['embeddings'][1] == embeddings[1]
 
         await conn.close()
 
@@ -126,10 +125,9 @@ async def init(conn):
             await conn.execute('DROP TABLE IF EXISTS asyncpg_items')
             await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding vector(3))')
 
-            embedding = np.array([1.5, 2, 3])
+            embedding = Vector([1.5, 2, 3])
             await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding)
 
             res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id")
-            assert np.array_equal(res[0]['embedding'], embedding)
-            assert res[0]['embedding'].dtype == np.float32
+            assert res[0]['embedding'] == embedding
             assert res[1]['embedding'] is None

tests/test_django.py

@@ -12,7 +12,7 @@
 import numpy as np
 import os
 import pgvector.django
-from pgvector import HalfVector, SparseVector
+from pgvector import Vector, HalfVector, SparseVector
 from pgvector.django import VectorExtension, VectorField, HalfVectorField, BitField, SparseVectorField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, JaccardDistance
 from unittest import mock
 
@@ -165,12 +165,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, np.array([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)]
@@ -293,31 +292,31 @@ 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, np.array([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, np.array([5, 7, 9]))
+        assert sum == Vector([5, 7, 9])
 
     def test_halfvec_avg(self):
         avg = Item.objects.aggregate(Avg('half_embedding'))['half_embedding__avg']
         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']
-        assert avg.to_list() == [2.5, 3.5, 4.5]
+        assert avg == HalfVector([2.5, 3.5, 4.5])
 
     def test_halfvec_sum(self):
         sum = Item.objects.aggregate(Sum('half_embedding'))['half_embedding__sum']
         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']
-        assert sum.to_list() == [5, 7, 9]
+        assert sum == HalfVector([5, 7, 9])
 
     def test_serialization(self):
         create_items()
@@ -347,7 +346,7 @@ def test_vector_form_save(self):
         assert form.has_changed()
         assert form.is_valid()
         assert form.save()
-        assert [4, 5, 6] == Item.objects.get(pk=1).embedding.tolist()
+        assert [4, 5, 6] == Item.objects.get(pk=1).embedding.to_list()
 
     def test_vector_form_save_missing(self):
         Item(id=1).save()
@@ -465,8 +464,8 @@ def test_vector_array(self):
 
             # this fails if the driver does not cast arrays
             item = Item.objects.get(pk=1)
-            assert item.embeddings[0].tolist() == [1, 2, 3]
-            assert item.embeddings[1].tolist() == [4, 5, 6]
+            assert item.embeddings[0].to_list() == [1, 2, 3]
+            assert item.embeddings[1].to_list() == [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 SparseVector
+from pgvector import Vector, HalfVector, SparseVector
 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, np.array([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,31 +169,31 @@ 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, np.array([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, np.array([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()
         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.to_list() == [2.5, 3.5, 4.5]
+        assert avg == HalfVector([2.5, 3.5, 4.5])
 
     def test_halfvec_sum(self):
         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.to_list() == [5, 7, 9]
+        assert sum == HalfVector([5, 7, 9])
 
     def test_get_or_create(self):
         Item.get_or_create(id=1, defaults={'embedding': [1, 2, 3]})