asyncpg-vector 0.1.0__tar.gz → 0.1.1__tar.gz

{asyncpg_vector-0.1.0 → asyncpg_vector-0.1.1}/LICENSE

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2025 Levente Hunyadi
+Copyright (c) 2025-2026 Levente Hunyadi
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

{asyncpg_vector-0.1.0 → asyncpg_vector-0.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: asyncpg_vector
-Version: 0.1.0
+Version: 0.1.1
 Summary: PostgreSQL vector support for asyncpg
 Author-email: Levente Hunyadi <hunyadi@gmail.com>
 Maintainer-email: Levente Hunyadi <hunyadi@gmail.com>

{asyncpg_vector-0.1.0 → asyncpg_vector-0.1.1}/asyncpg_vector/__init__.py

@@ -1,18 +1,19 @@
 """
 PostgreSQL vector support for asyncpg.
 
-Registers data types `vector` and `halfvec` from the PostgreSQL extension `vector` to the asynchronous PostgreSQL
-client `asyncpg`, and marshals vector data to and from PostgreSQL database tables.
+Registers data types `vector`, `halfvec` and `sparsevec` from the PostgreSQL extension `vector` to the asynchronous
+PostgreSQL client `asyncpg`, and marshals vector data to and from PostgreSQL database tables.
+
+Copyright 2025-2026, Levente Hunyadi
 
 :see: https://github.com/hunyadi/asyncpg_vector
 """
 
 import base64
-import struct
 import sys
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
-from struct import pack
+from struct import pack, unpack
 from typing import ClassVar
 
 import asyncpg
@@ -27,9 +28,9 @@ if sys.version_info >= (3, 12):
 else:
     from typing_extensions import override
 
-__version__ = "0.1.0"
+__version__ = "0.1.1"
 __author__ = "Levente Hunyadi"
-__copyright__ = "Copyright 2025, Levente Hunyadi"
+__copyright__ = "Copyright 2025-2026, Levente Hunyadi"
 __license__ = "MIT"
 __maintainer__ = "Levente Hunyadi"
 __status__ = "Production"
@@ -38,32 +39,10 @@ __status__ = "Production"
 class BasicVector(ABC):
     "Base class for PostgreSQL vector types."
 
-    __slots__ = ("_data",)
-
-    _data: bytes
-
-    def __init__(self, data: bytes) -> None:
-        self._data = data
-
-    def __eq__(self, value: object) -> bool:
-        if not isinstance(value, self.__class__):
-            return False
-
-        return self._data == value._data
-
-    def __repr__(self) -> str:
-        return f"{type(self).__name__}({self._data!r})"
-
+    @abstractmethod
     def size(self) -> int:
         "Number of vector dimensions."
 
-        return len(self._data) // self.bytes_per_item()
-
-    @classmethod
-    @abstractmethod
-    def bytes_per_item(cls) -> int:
-        "Number of bytes per item in the vector."
-
         ...
 
     @abstractmethod
@@ -88,22 +67,21 @@ class BasicVector(ABC):
         """
 
         decoded_bytes = base64.b64decode(base64_encoded)
-        float_tuple = struct.unpack(f"{len(decoded_bytes) // 4}f", decoded_bytes)
+        float_tuple = unpack(f"{len(decoded_bytes) // 4}f", decoded_bytes)
         return cls.from_float_list(float_tuple)
 
+    @abstractmethod
     def to_database_binary(self) -> bytes:
         "Writes the data into the PostgreSQL data transfer representation."
 
-        return pack(">HH", self.size(), 0) + self._data
+        ...
 
     @classmethod
+    @abstractmethod
     def from_database_binary(cls, data: bytes) -> Self:
         "Reads the data from the PostgreSQL data transfer representation."
 
-        size, _unused = struct.unpack(">HH", data[0:4])
-        if len(data) != 4 + cls.bytes_per_item() * size:
-            raise ValueError(f"expected size: {cls.bytes_per_item()} * {size}; got {len(data) - 4} bytes")
-        return cls(data[4:])
+        ...
 
     @classmethod
     def _to_database_binary(cls, value: Self | list[float] | None) -> bytes | None:
@@ -123,7 +101,7 @@ class BasicVector(ABC):
 
                     obj = cls.from_float_list(value)
                 else:
-                    obj = cls(bytes())
+                    obj = cls()
             case _:
                 raise ValueError(f"unsupported type: {type(value).__name__}")
 
@@ -139,7 +117,56 @@ class BasicVector(ABC):
         return cls.from_database_binary(value)
 
 
-class HalfVector(BasicVector):
+class DenseVector(BasicVector):
+    "Base class for PostgreSQL dense vector types."
+
+    __slots__ = ("_data",)
+
+    _data: bytes
+
+    def __init__(self, data: bytes | None = None) -> None:
+        if data is not None:
+            self._data = data
+        else:
+            self._data = bytes()
+
+    def __eq__(self, value: object) -> bool:
+        if not isinstance(value, self.__class__):
+            return False
+
+        return self._data == value._data
+
+    def __str__(self) -> str:
+        return f"{type(self).__name__}(size={self.size()})"
+
+    def __repr__(self) -> str:
+        return f"{type(self).__name__}({self._data!r})"
+
+    @override
+    def size(self) -> int:
+        return len(self._data) // self.bytes_per_item()
+
+    @classmethod
+    @abstractmethod
+    def bytes_per_item(cls) -> int:
+        "Number of bytes per item in the vector."
+
+        ...
+
+    @override
+    def to_database_binary(self) -> bytes:
+        return pack(">HH", self.size(), 0) + self._data
+
+    @override
+    @classmethod
+    def from_database_binary(cls, data: bytes) -> Self:
+        size, _unused = unpack(">HH", data[0:4])
+        if len(data) != 4 + cls.bytes_per_item() * size:
+            raise ValueError(f"expected size: {cls.bytes_per_item()} * {size}; got {len(data) - 4} bytes")
+        return cls(data[4:])
+
+
+class HalfVector(DenseVector):
     "Implements the PostgreSQL extension type `halfvec`."
 
     type_name: ClassVar[str] = "halfvec"
@@ -152,15 +179,15 @@ class HalfVector(BasicVector):
 
     @override
     def to_float_list(self) -> list[float]:
-        return list(struct.unpack(f">{self.size()}e", self._data))
+        return list(unpack(f">{self.size()}e", self._data))
 
     @override
     @classmethod
     def from_float_list(cls, vec: Sequence[float]) -> Self:
-        return cls(struct.pack(f">{len(vec)}e", *vec))
+        return cls(pack(f">{len(vec)}e", *vec))
 
 
-class Vector(BasicVector):
+class Vector(DenseVector):
     "Implements the PostgreSQL extension type `vector`."
 
     type_name: ClassVar[str] = "vector"
@@ -173,12 +200,12 @@ class Vector(BasicVector):
 
     @override
     def to_float_list(self) -> list[float]:
-        return list(struct.unpack(f">{self.size()}f", self._data))
+        return list(unpack(f">{self.size()}f", self._data))
 
     @override
     @classmethod
     def from_float_list(cls, vec: Sequence[float]) -> Self:
-        return cls(struct.pack(f">{len(vec)}f", *vec))
+        return cls(pack(f">{len(vec)}f", *vec))
 
     @override
     @classmethod
@@ -186,8 +213,83 @@ class Vector(BasicVector):
         return cls(base64.b64decode(base64_encoded))
 
 
+class SparseVector(BasicVector):
+    "Implements the PostgreSQL extension type `sparsevec`."
+
+    __slots__ = ("_size", "_indices", "_values")
+
+    _size: int
+    _indices: bytes
+    _values: bytes
+
+    type_name: ClassVar[str] = "sparsevec"
+    cosine_similarity: ClassVar[str] = "sparsevec_cosine_ops"
+
+    def __init__(self, size: int | None = None, indices: bytes | None = None, values: bytes | None = None) -> None:
+        if size is not None and indices is not None and values is not None:
+            self._size = size
+            self._indices = indices
+            self._values = values
+            if len(self._indices) != len(self._values):
+                raise ValueError(f"expected: `indices` (of len {len(self._indices)}) and `values` (of len {len(self._values)}) to match in size")
+        elif size is None and indices is None and values is None:
+            self._size = 0
+            self._indices = bytes()
+            self._values = bytes()
+        else:
+            raise ValueError("expected: either all of `size`, `indices` and `values`, or neither")
+
+    def __eq__(self, value: object) -> bool:
+        if not isinstance(value, self.__class__):
+            return False
+
+        return self._size == value._size and self._indices == value._indices and self._values == value._values
+
+    def __str__(self) -> str:
+        return f"{type(self).__name__}(size={self.size()}, nnz={self.nnz()})"
+
+    def __repr__(self) -> str:
+        return f"{type(self).__name__}(size={self.size()}, indices={self._indices!r}, values={self._values!r})"
+
+    @override
+    def size(self) -> int:
+        return self._size
+
+    def nnz(self) -> int:
+        return len(self._indices) // 4
+
+    @override
+    def to_float_list(self) -> list[float]:
+        count = self.nnz()
+        indices = unpack(f">{count}i", self._indices)
+        values = unpack(f">{count}f", self._values)
+        items = [0.0 for _ in range(self._size)]
+        for index, value in zip(indices, values, strict=True):
+            items[index] = value
+        return items
+
+    @override
+    @classmethod
+    def from_float_list(cls, vec: Sequence[float]) -> Self:
+        indices = [index for index, value in enumerate(vec) if value != 0]
+        values = [vec[index] for index in indices]
+        count = len(indices)
+        return cls(len(vec), pack(f">{count}i", *indices), pack(f">{count}f", *values))
+
+    @override
+    def to_database_binary(self) -> bytes:
+        return pack(">iii", self._size, self.nnz(), 0) + self._indices + self._values
+
+    @override
+    @classmethod
+    def from_database_binary(cls, data: bytes) -> Self:
+        size, count, _ = unpack(">iii", data[0:12])
+        return cls(size, data[12 : 12 + 4 * count], data[12 + 4 * count :])
+
+
 async def register_vector(conn: asyncpg.Connection, schema: str = "public") -> None:
     "Registers `vector` extension types with Python module `asyncpg`."
 
     await conn.set_type_codec("vector", schema=schema, encoder=Vector._to_database_binary, decoder=Vector._from_database_binary, format="binary")  # pyright: ignore[reportPrivateUsage]
     await conn.set_type_codec("halfvec", schema=schema, encoder=HalfVector._to_database_binary, decoder=HalfVector._from_database_binary, format="binary")  # pyright: ignore[reportPrivateUsage]
+    await conn.set_type_codec("sparsevec", schema=schema, encoder=SparseVector._to_database_binary, decoder=SparseVector._from_database_binary, format="binary")  # pyright: ignore[reportPrivateUsage]

{asyncpg_vector-0.1.0 → asyncpg_vector-0.1.1}/asyncpg_vector.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: asyncpg_vector
-Version: 0.1.0
+Version: 0.1.1
 Summary: PostgreSQL vector support for asyncpg
 Author-email: Levente Hunyadi <hunyadi@gmail.com>
 Maintainer-email: Levente Hunyadi <hunyadi@gmail.com>

{asyncpg_vector-0.1.0 → asyncpg_vector-0.1.1}/asyncpg_vector.egg-info/SOURCES.txt

@@ -9,4 +9,5 @@ asyncpg_vector.egg-info/SOURCES.txt
 asyncpg_vector.egg-info/dependency_links.txt
 asyncpg_vector.egg-info/requires.txt
 asyncpg_vector.egg-info/top_level.txt
-asyncpg_vector.egg-info/zip-safe
+asyncpg_vector.egg-info/zip-safe
+test/test_vector.py

asyncpg_vector-0.1.1/test/test_vector.py

@@ -0,0 +1,118 @@
+import unittest
+from collections.abc import AsyncIterator, Sequence
+from contextlib import asynccontextmanager
+from random import random
+from struct import pack, unpack
+
+import asyncpg
+
+from asyncpg_vector import HalfVector, SparseVector, Vector, register_vector
+
+
+@asynccontextmanager
+async def get_connection() -> AsyncIterator[asyncpg.Connection]:
+    conn = await asyncpg.connect(host="localhost", port=5432, user="postgres", password="postgres")
+    try:
+        yield conn
+    finally:
+        await conn.close()
+
+
+def to_float16(values: Sequence[float]) -> list[float]:
+    "Truncates a list of floating-point values to 16-bit width."
+
+    return list(unpack(f"{len(values)}e", pack(f"{len(values)}e", *values)))
+
+
+def to_float32(values: Sequence[float]) -> list[float]:
+    "Truncates a list of floating-point values to 32-bit width."
+
+    return list(unpack(f"{len(values)}f", pack(f"{len(values)}f", *values)))
+
+
+def random_dense() -> list[float]:
+    "Generates a dense random vector."
+
+    return [random() for _ in range(1536)]
+
+
+def random_sparse() -> list[float]:
+    "Generates a sparse random vector with many zeros."
+
+    return [(random() if random() > 0.9 else 0.0) for _ in range(1536)]
+
+
+class TestVector(unittest.IsolatedAsyncioTestCase):
+    async def test_types(self) -> None:
+        f32d_vector = to_float32(random_dense())
+        f16d_vector = to_float16(random_dense())
+        f32s_vector = to_float32(random_sparse())
+
+        # empty vectors
+        self.assertEqual(Vector().size(), 0)
+        self.assertEqual(HalfVector().size(), 0)
+        self.assertEqual(SparseVector().size(), 0)
+
+        # str() yields compact representation
+        self.assertLess(len(str(Vector.from_float_list(f32d_vector))), 64)
+        self.assertLess(len(str(HalfVector.from_float_list(f16d_vector))), 64)
+        self.assertLess(len(str(SparseVector.from_float_list(f32s_vector))), 64)
+
+        # repr() yields full representation
+        self.assertGreater(len(repr(Vector.from_float_list(f32d_vector))), 1000)
+        self.assertGreater(len(repr(HalfVector.from_float_list(f16d_vector))), 1000)
+        self.assertGreater(len(repr(SparseVector.from_float_list(f32s_vector))), 1000)
+
+        # round trip for list[float]
+        self.assertEqual(Vector.from_float_list(f32d_vector).to_float_list(), f32d_vector)
+        self.assertEqual(HalfVector.from_float_list(f16d_vector).to_float_list(), f16d_vector)
+        self.assertEqual(SparseVector.from_float_list(f32s_vector).to_float_list(), f32s_vector)
+
+    async def test_connection(self) -> None:
+        create_sql = """
+        --sql
+        CREATE EXTENSION IF NOT EXISTS vector;
+
+        --sql
+        CREATE TEMPORARY TABLE vector_types(
+            id bigint GENERATED ALWAYS AS IDENTITY,
+            embedding vector(1536) NOT NULL,
+            half_embedding halfvec(1536) NOT NULL,
+            sparse_embedding sparsevec(1536) NOT NULL,
+            CONSTRAINT pk_vector_types PRIMARY KEY (id)
+        );
+        """
+
+        insert_sql = """
+        --sql
+        INSERT INTO vector_types (embedding, half_embedding, sparse_embedding)
+        VALUES ($1, $2, $3);
+        """
+
+        select_sql = """
+        --sql
+        SELECT embedding, half_embedding, sparse_embedding
+        FROM vector_types
+        ORDER BY id;
+        """
+
+        async with get_connection() as conn:
+            await conn.execute(create_sql)
+            await register_vector(conn)
+
+            records = [
+                (
+                    Vector.from_float_list(to_float32(random_dense())),
+                    HalfVector.from_float_list(to_float16(random_dense())),
+                    SparseVector.from_float_list(to_float32(random_sparse())),
+                )
+                for _ in range(1)
+            ]
+            await conn.executemany(insert_sql, records)
+            rows = await conn.fetch(select_sql)
+            for row, record in zip(rows, records, strict=True):
+                self.assertEqual(tuple(row), record)
+
+
+if __name__ == "__main__":
+    unittest.main()