Merge branch 'master' into return-vector

Author: ankane

.github/workflows/build.yml

@@ -24,3 +24,6 @@ jobs:
           make
           sudo make install
       - run: pytest
+
+      - run: pip install "SQLAlchemy<2" -U
+      - run: pytest tests/test_sqlalchemy.py

pgvector/bit.py

@@ -7,14 +7,13 @@ def __init__(self, value):
         if isinstance(value, str):
             self._value = self.from_text(value)._value
         else:
-            # TODO change in 0.4.0
-            # TODO raise if dtype not bool or uint8
-            # if isinstance(value, np.ndarray) and value.dtype == np.uint8:
-            #     value = np.unpackbits(value)
-            # else:
-            #     value = np.asarray(value, dtype=bool)
-
-            value = np.asarray(value, dtype=bool)
+            if isinstance(value, np.ndarray):
+                if value.dtype == np.uint8:
+                    value = np.unpackbits(value).astype(bool)
+                elif value.dtype != np.bool:
+                    raise ValueError('expected dtype to be bool or uint8')
+            else:
+                value = np.asarray(value, dtype=bool)
 
             if value.ndim != 1:
                 raise ValueError('expected ndim to be 1')

tests/test_asyncpg.py

@@ -1,6 +1,6 @@
 import asyncpg
 import numpy as np
-from pgvector import Vector, HalfVector, SparseVector
+from pgvector import HalfVector, SparseVector, Vector
 from pgvector.asyncpg import register_vector
 import pytest
 
@@ -16,11 +16,13 @@ async def test_vector(self):
         await register_vector(conn)
 
         embedding = Vector([1.5, 2, 3])
-        await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding)
+        embedding2 = np.array([4.5, 5, 6])
+        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 res[0]['embedding'] == embedding
-        assert res[1]['embedding'] is None
+        assert res[1]['embedding'] == Vector(embedding2)
+        assert res[2]['embedding'] is None
 
         # ensures binary format is correct
         text_res = await conn.fetch("SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1")
@@ -38,11 +40,13 @@ async def test_halfvec(self):
         await register_vector(conn)
 
         embedding = HalfVector([1.5, 2, 3])
-        await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding)
+        embedding2 = [4.5, 5, 6]
+        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 res[0]['embedding'] == embedding
-        assert res[1]['embedding'] is None
+        assert res[1]['embedding'] == HalfVector(embedding2)
+        assert res[2]['embedding'] is None
 
         # ensures binary format is correct
         text_res = await conn.fetch("SELECT embedding::text FROM asyncpg_items ORDER BY id LIMIT 1")
@@ -105,11 +109,14 @@ async def test_vector_array(self):
         await register_vector(conn)
 
         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])
+        await conn.execute("INSERT INTO asyncpg_items (embeddings) VALUES ($1)", embeddings)
+
+        embeddings2 = [np.array([1.5, 2, 3]), np.array([4.5, 5, 6])]
+        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 res[0]['embeddings'][0] == embeddings[0]
-        assert res[0]['embeddings'][1] == embeddings[1]
+        assert res[0]['embeddings'] == embeddings
+        assert res[1]['embeddings'] == [Vector(e) for e in embeddings2]
 
         await conn.close()
 
@@ -126,8 +133,10 @@ async def init(conn):
             await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding vector(3))')
 
             embedding = Vector([1.5, 2, 3])
-            await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding)
+            embedding2 = np.array([1.5, 2, 3])
+            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 res[0]['embedding'] == embedding
-            assert res[1]['embedding'] is None
+            assert res[1]['embedding'] == Vector(embedding2)
+            assert res[2]['embedding'] is None

tests/test_bit.py

@@ -15,9 +15,13 @@ def test_str(self):
 
     def test_ndarray_uint8(self):
         arr = np.array([254, 7, 0], dtype=np.uint8)
-        # TODO change in 0.4.0
-        # assert Bit(arr).to_text() == '111111100000011100000000'
-        assert Bit(arr).to_text() == '110'
+        assert Bit(arr).to_text() == '111111100000011100000000'
+
+    def test_ndarray_uint16(self):
+        arr = np.array([254, 7, 0], dtype=np.uint16)
+        with pytest.raises(ValueError) as error:
+            Bit(arr)
+        assert str(error.value) == 'expected dtype to be bool or uint8'
 
     def test_ndarray_same_object(self):
         arr = np.array([True, False, True])

tests/test_django.py

@@ -198,7 +198,7 @@ def test_vector_l1_distance(self):
     def test_halfvec(self):
         Item(id=1, half_embedding=[1, 2, 3]).save()
         item = Item.objects.get(pk=1)
-        assert item.half_embedding.to_list() == [1, 2, 3]
+        assert item.half_embedding == HalfVector([1, 2, 3])
 
     def test_halfvec_l2_distance(self):
         create_items()
@@ -250,7 +250,7 @@ def test_bit_jaccard_distance(self):
     def test_sparsevec(self):
         Item(id=1, sparse_embedding=SparseVector([1, 2, 3])).save()
         item = Item.objects.get(pk=1)
-        assert item.sparse_embedding.to_list() == [1, 2, 3]
+        assert item.sparse_embedding == SparseVector([1, 2, 3])
 
     def test_sparsevec_l2_distance(self):
         create_items()
@@ -346,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.to_list()
+        assert Item.objects.get(pk=1).embedding == Vector([4, 5, 6])
 
     def test_vector_form_save_missing(self):
         Item(id=1).save()
@@ -374,7 +374,7 @@ def test_halfvec_form_save(self):
         assert form.has_changed()
         assert form.is_valid()
         assert form.save()
-        assert [4, 5, 6] == Item.objects.get(pk=1).half_embedding.to_list()
+        assert Item.objects.get(pk=1).half_embedding == HalfVector([4, 5, 6])
 
     def test_halfvec_form_save_missing(self):
         Item(id=1).save()
@@ -431,7 +431,7 @@ def test_sparsevec_form_save(self):
         assert form.has_changed()
         assert form.is_valid()
         assert form.save()
-        assert [4, 5, 6] == Item.objects.get(pk=1).sparse_embedding.to_list()
+        assert Item.objects.get(pk=1).sparse_embedding == SparseVector([4, 5, 6])
 
     def test_sparesevec_form_save_missing(self):
         Item(id=1).save()
@@ -464,8 +464,7 @@ def test_vector_array(self):
 
             # this fails if the driver does not cast arrays
             item = Item.objects.get(pk=1)
-            assert item.embeddings[0].to_list() == [1, 2, 3]
-            assert item.embeddings[1].to_list() == [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_half_vector.py

@@ -1,6 +1,7 @@
 import numpy as np
 from pgvector import HalfVector
 import pytest
+from struct import pack
 
 
 class TestHalfVector:
@@ -44,3 +45,15 @@ def test_equality(self):
 
     def test_dimensions(self):
         assert HalfVector([1, 2, 3]).dimensions() == 3
+
+    def test_from_text(self):
+        vec = HalfVector.from_text('[1.5,2,3]')
+        assert vec.to_list() == [1.5, 2, 3]
+        assert np.array_equal(vec.to_numpy(), [1.5, 2, 3])
+
+    def test_from_binary(self):
+        data = pack('>HH3e', 3, 0, 1.5, 2, 3)
+        vec = HalfVector.from_binary(data)
+        assert vec.to_list() == [1.5, 2, 3]
+        assert np.array_equal(vec.to_numpy(), [1.5, 2, 3])
+        assert vec.to_binary() == data

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 Vector, HalfVector, SparseVector
+from pgvector import HalfVector, SparseVector, Vector
 from pgvector.peewee import VectorField, HalfVectorField, FixedBitField, SparseVectorField
 
 db = PostgresqlDatabase('pgvector_python_test')
@@ -76,7 +76,7 @@ def test_vector_l1_distance(self):
     def test_halfvec(self):
         Item.create(id=1, half_embedding=[1, 2, 3])
         item = Item.get_by_id(1)
-        assert item.half_embedding.to_list() == [1, 2, 3]
+        assert item.half_embedding == HalfVector([1, 2, 3])
 
     def test_halfvec_l2_distance(self):
         create_items()
@@ -128,7 +128,7 @@ def test_bit_jaccard_distance(self):
     def test_sparsevec(self):
         Item.create(id=1, sparse_embedding=[1, 2, 3])
         item = Item.get_by_id(1)
-        assert item.sparse_embedding.to_list() == [1, 2, 3]
+        assert item.sparse_embedding == SparseVector([1, 2, 3])
 
     def test_sparsevec_l2_distance(self):
         create_items()
@@ -219,5 +219,5 @@ class Meta:
         # fails with column "embeddings" is of type vector[] but expression is of type text[]
         # ExtItem.create(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])])
         # item = ExtItem.get_by_id(1)
-        # assert np.array_equal(item.embeddings[0], np.array([1, 2, 3]))
-        # assert np.array_equal(item.embeddings[1], np.array([4, 5, 6]))
+        # assert np.array_equal(item.embeddings[0], [1, 2, 3])
+        # assert np.array_equal(item.embeddings[1], [4, 5, 6])

tests/test_psycopg.py

@@ -68,19 +68,17 @@ def test_halfvec(self):
         conn.execute('INSERT INTO psycopg_items (half_embedding) VALUES (%s)', (embedding,))
 
         res = conn.execute('SELECT half_embedding FROM psycopg_items ORDER BY id').fetchone()[0]
-        assert res.to_list() == [1.5, 2, 3]
+        assert res == HalfVector([1.5, 2, 3])
 
     def test_halfvec_binary_format(self):
         embedding = HalfVector([1.5, 2, 3])
         res = conn.execute('SELECT %b::halfvec', (embedding,), binary=True).fetchone()[0]
-        assert res.to_list() == [1.5, 2, 3]
-        assert np.array_equal(res.to_numpy(), np.array([1.5, 2, 3]))
+        assert res == HalfVector([1.5, 2, 3])
 
     def test_halfvec_text_format(self):
         embedding = HalfVector([1.5, 2, 3])
         res = conn.execute('SELECT %t::halfvec', (embedding,)).fetchone()[0]
-        assert res.to_list() == [1.5, 2, 3]
-        assert np.array_equal(res.to_numpy(), np.array([1.5, 2, 3]))
+        assert res == HalfVector([1.5, 2, 3])
 
     def test_bit(self):
         embedding = Bit([True, False, True])
@@ -105,25 +103,17 @@ def test_sparsevec(self):
         conn.execute('INSERT INTO psycopg_items (sparse_embedding) VALUES (%s)', (embedding,))
 
         res = conn.execute('SELECT sparse_embedding FROM psycopg_items ORDER BY id').fetchone()[0]
-        assert res.to_list() == [1.5, 2, 3]
+        assert res == SparseVector([1.5, 2, 3])
 
     def test_sparsevec_binary_format(self):
         embedding = SparseVector([1.5, 0, 2, 0, 3, 0])
         res = conn.execute('SELECT %b::sparsevec', (embedding,), binary=True).fetchone()[0]
-        assert res.dimensions() == 6
-        assert res.indices() == [0, 2, 4]
-        assert res.values() == [1.5, 2, 3]
-        assert res.to_list() == [1.5, 0, 2, 0, 3, 0]
-        assert np.array_equal(res.to_numpy(), np.array([1.5, 0, 2, 0, 3, 0]))
+        assert res == embedding
 
     def test_sparsevec_text_format(self):
         embedding = SparseVector([1.5, 0, 2, 0, 3, 0])
         res = conn.execute('SELECT %t::sparsevec', (embedding,)).fetchone()[0]
-        assert res.dimensions() == 6
-        assert res.indices() == [0, 2, 4]
-        assert res.values() == [1.5, 2, 3]
-        assert res.to_list() == [1.5, 0, 2, 0, 3, 0]
-        assert np.array_equal(res.to_numpy(), np.array([1.5, 0, 2, 0, 3, 0]))
+        assert res == embedding
 
     def test_text_copy_from(self):
         embedding = np.array([1.5, 2, 3])
@@ -161,8 +151,8 @@ def test_binary_copy_to(self):
         cur = conn.cursor()
         with cur.copy("COPY psycopg_items (embedding, half_embedding) TO STDOUT WITH (FORMAT BINARY)") as copy:
             for row in copy.rows():
-                assert Vector.from_binary(row[0]).to_list() == [1.5, 2, 3]
-                assert HalfVector.from_binary(row[1]).to_list() == [1.5, 2, 3]
+                assert Vector.from_binary(row[0]) == embedding
+                assert HalfVector.from_binary(row[1]) == half_embedding
 
     def test_binary_copy_to_set_types(self):
         embedding = Vector([1.5, 2, 3])

tests/test_psycopg2.py

@@ -80,17 +80,17 @@ def test_halfvec_array(self):
 
         cur.execute('SELECT half_embeddings FROM psycopg2_items ORDER BY id')
         res = cur.fetchone()
-        assert res[0][0].to_list() == [1.5, 2, 3]
-        assert res[0][1].to_list() == [4.5, 5, 6]
+        assert res[0][0] == HalfVector([1.5, 2, 3])
+        assert res[0][1] == HalfVector([4.5, 5, 6])
 
     def test_sparsevec_array(self):
         embeddings = [SparseVector([1.5, 2, 3]), SparseVector([4.5, 5, 6])]
         cur.execute('INSERT INTO psycopg2_items (sparse_embeddings) VALUES (%s::sparsevec[])', (embeddings,))
 
         cur.execute('SELECT sparse_embeddings FROM psycopg2_items ORDER BY id')
         res = cur.fetchone()
-        assert res[0][0].to_list() == [1.5, 2, 3]
-        assert res[0][1].to_list() == [4.5, 5, 6]
+        assert res[0][0] == SparseVector([1.5, 2, 3])
+        assert res[0][1] == SparseVector([4.5, 5, 6])
 
     def test_cursor_factory(self):
         for cursor_factory in [DictCursor, RealDictCursor, NamedTupleCursor]:

tests/test_sparse_vector.py

@@ -2,13 +2,14 @@
 from pgvector import SparseVector
 import pytest
 from scipy.sparse import coo_array
+from struct import pack
 
 
 class TestSparseVector:
     def test_list(self):
         vec = SparseVector([1, 0, 2, 0, 3, 0])
         assert vec.to_list() == [1, 0, 2, 0, 3, 0]
-        assert vec.to_numpy().tolist() == [1, 0, 2, 0, 3, 0]
+        assert np.array_equal(vec.to_numpy(), [1, 0, 2, 0, 3, 0])
         assert vec.indices() == [0, 2, 4]
 
     def test_list_dimensions(self):
@@ -56,6 +57,7 @@ def test_equality(self):
         assert SparseVector([1, 0, 2, 0, 3, 0]) == SparseVector([1, 0, 2, 0, 3, 0])
         assert SparseVector([1, 0, 2, 0, 3, 0]) != SparseVector([1, 0, 2, 0, 3, 1])
         assert SparseVector([1, 0, 2, 0, 3, 0]) == SparseVector({2: 2, 4: 3, 0: 1, 3: 0}, 6)
+        assert SparseVector({}, 1) != SparseVector({}, 2)
 
     def test_dimensions(self):
         assert SparseVector([1, 0, 2, 0, 3, 0]).dimensions() == 6
@@ -67,8 +69,26 @@ def test_values(self):
         assert SparseVector([1, 0, 2, 0, 3, 0]).values() == [1, 2, 3]
 
     def test_to_coo(self):
-        assert SparseVector([1, 0, 2, 0, 3, 0]).to_coo().toarray().tolist() == [[1, 0, 2, 0, 3, 0]]
+        assert np.array_equal(SparseVector([1, 0, 2, 0, 3, 0]).to_coo().toarray(), [[1, 0, 2, 0, 3, 0]])
 
     def test_zero_vector_text(self):
         vec = SparseVector({}, 3)
         assert vec.to_list() == SparseVector.from_text(vec.to_text()).to_list()
+
+    def test_from_text(self):
+        vec = SparseVector.from_text('{1:1.5,3:2,5:3}/6')
+        assert vec.dimensions() == 6
+        assert vec.indices() == [0, 2, 4]
+        assert vec.values() == [1.5, 2, 3]
+        assert vec.to_list() == [1.5, 0, 2, 0, 3, 0]
+        assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0])
+
+    def test_from_binary(self):
+        data = pack('>iii3i3f', 6, 3, 0, 0, 2, 4, 1.5, 2, 3)
+        vec = SparseVector.from_binary(data)
+        assert vec.dimensions() == 6
+        assert vec.indices() == [0, 2, 4]
+        assert vec.values() == [1.5, 2, 3]
+        assert vec.to_list() == [1.5, 0, 2, 0, 3, 0]
+        assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0])
+        assert vec.to_binary() == data

tests/test_sqlalchemy.py

@@ -1,7 +1,7 @@
 import asyncpg
 import numpy as np
 import os
-from pgvector import Vector, HalfVector, SparseVector
+from pgvector import HalfVector, SparseVector, Vector
 from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, avg, sum
 import pytest
 from sqlalchemy import create_engine, event, insert, inspect, select, text, MetaData, Table, Column, Index, Integer, ARRAY
@@ -539,8 +539,7 @@ def test_vector_array(self, engine):
 
             # this fails if the driver does not cast arrays
             item = session.get(Item, 1)
-            assert item.embeddings[0] == Vector([1, 2, 3])
-            assert item.embeddings[1] == Vector([4, 5, 6])
+            assert item.embeddings == [Vector([1, 2, 3]), Vector([4, 5, 6])]
 
     def test_halfvec_array(self, engine):
         with Session(engine) as session:
@@ -637,7 +636,10 @@ async def test_vector_array(self, engine):
             async with session.begin():
                 session.add(Item(id=1, embeddings=[Vector([1, 2, 3]), Vector([4, 5, 6])]))
                 item = await session.get(Item, 1)
-                assert item.embeddings[0] == Vector([1, 2, 3])
-                assert item.embeddings[1] == Vector([4, 5, 6])
+                assert item.embeddings == [Vector([1, 2, 3]), Vector([4, 5, 6])]
+
+                session.add(Item(id=2, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])]))
+                item = await session.get(Item, 2)
+                assert item.embeddings == [Vector([1, 2, 3]), Vector([4, 5, 6])]
 
         await engine.dispose()