diff --git a/README.md b/README.md index 5b61e56..d0849b8 100644 --- a/README.md +++ b/README.md @@ -4,12 +4,17 @@ Supported databases/extensions: - [x] [`pgvecto.rs`](https://github.com/tensorchord/pgvecto.rs) - [x] [`pgvector`](https://github.com/pgvector/pgvector) -- [ ] [`qdrant`](https://github.com/qdrant/qdrant/) +- [x] [`qdrant`](https://github.com/qdrant/qdrant/) Supported datasets: - [x] random generated - [x] GIST 960 +- [x] GLOVE +- [x] Deep Image +- [x] LAION + +For more information, check the [source.py](./vector_bench/dataset/source.py). ## Installation @@ -24,7 +29,7 @@ pip install vector_bench Run the docker compose file under [`server`](server/) folder. ```base -cd server/pgvecto.rs && docker compose up -d +docker compose -f docker/compose.${DB_NAME}.yaml up -d ``` ### Client diff --git a/server/pgvecto.rs/compose.yaml b/docker/compose.pgvecto_rs.yaml similarity index 88% rename from server/pgvecto.rs/compose.yaml rename to docker/compose.pgvecto_rs.yaml index f1d7ee1..fef771c 100644 --- a/server/pgvecto.rs/compose.yaml +++ b/docker/compose.pgvecto_rs.yaml @@ -1,7 +1,7 @@ services: - pgvector: + pgvectors: image: tensorchord/pgvecto-rs:pg15-v0.1.13 - container_name: pgvector + container_name: pgvectors environment: - POSTGRES_USER=postgres - POSTGRES_PASSWORD=password diff --git a/server/pgvector/compose.yaml b/docker/compose.pgvector.yaml similarity index 89% rename from server/pgvector/compose.yaml rename to docker/compose.pgvector.yaml index fd30a66..85d13d3 100644 --- a/server/pgvector/compose.yaml +++ b/docker/compose.pgvector.yaml @@ -1,6 +1,6 @@ services: pgvector: - image: ankane/pgvector:v0.5.1 + image: pgvector/pgvector:0.6.0-pg15 container_name: pgvector environment: - POSTGRES_USER=postgres diff --git a/docker/compose.qdrant.yaml b/docker/compose.qdrant.yaml new file mode 100644 index 0000000..d1260ec --- /dev/null +++ b/docker/compose.qdrant.yaml @@ -0,0 +1,17 @@ +services: + qdrant: + image: qdrant/qdrant:v1.7.4 + container_name: qdrant + ports: + - "6333:6333" + logging: + driver: "json-file" + options: + max-file: "1" + max-size: "10m" + deploy: + resources: + limits: + cpus: "8" + reservations: + cpus: "4" diff --git a/pyproject.toml b/pyproject.toml index 17bf012..8ed845c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,6 +16,7 @@ dependencies = [ "tqdm~=4.66", "httpx~=0.25", "psycopg[binary]~=3.1", + "qdrant-client~=1.7.1", ] [project.optional-dependencies] dev = [ diff --git a/vector_bench/client/__init__.py b/vector_bench/client/__init__.py index 8101650..31acedc 100644 --- a/vector_bench/client/__init__.py +++ b/vector_bench/client/__init__.py @@ -1,8 +1,10 @@ from vector_bench.client.pgvecto_rs import PgVectorsClient from vector_bench.client.pgvector import PgvectorClient +from vector_bench.client.qdrant import QdrantVectorClient from vector_bench.spec import EnumSelector class DataBaseClient(EnumSelector): PGVECTO_RS = PgVectorsClient PGVECTOR = PgvectorClient + QDRANT = QdrantVectorClient diff --git a/vector_bench/client/base.py b/vector_bench/client/base.py index d163f35..a4f668a 100644 --- a/vector_bench/client/base.py +++ b/vector_bench/client/base.py @@ -11,7 +11,7 @@ def insert_batch(self, records: list[Record]): pass @abc.abstractmethod - def query(self, vector: list[float], top_k: int = 10): + def query(self, vector: list[float], top_k: int = 10) -> list[Record]: pass @abc.abstractclassmethod diff --git a/vector_bench/client/pgvecto_rs.py b/vector_bench/client/pgvecto_rs.py index 74224d1..d142d6a 100644 --- a/vector_bench/client/pgvecto_rs.py +++ b/vector_bench/client/pgvecto_rs.py @@ -16,8 +16,8 @@ class VectorDumper(Dumper): def dump(self, obj): if isinstance(obj, np.ndarray): - obj = f"[{','.join(map(str, obj))}]" - return str(obj).replace(" ", "") + return f"[{','.join(map(str, obj))}]".encode() + return str(obj).replace(" ", "").encode() class VectorLoader(Loader): @@ -152,19 +152,6 @@ def indexing(self): conn.execute(self.sql_create_index) conn.commit() - async def insert(self, record: Record): - async with await psycopg.AsyncConnection.connect(self.url) as conn: - register_vector_async(conn) - await conn.execute( - self.sql_insert, - ( - record.id, - record.vector, - Jsonb(record.metadata or {}, dumps=msgspec.json.encode), - ), - ) - await conn.commit() - def insert_batch(self, records: list[Record]): with psycopg.connect(self.url) as conn: register_vector(conn) diff --git a/vector_bench/client/pgvector.py b/vector_bench/client/pgvector.py index e66490f..0e93d74 100644 --- a/vector_bench/client/pgvector.py +++ b/vector_bench/client/pgvector.py @@ -69,7 +69,7 @@ def to_db_binary(value): class VectorDumper(Dumper): def dump(self, obj): - return to_db(obj).encode("utf8") + return to_db(obj).encode() class VectorBinaryDumper(VectorDumper): @@ -83,7 +83,7 @@ class VectorLoader(Loader): def load(self, data): if isinstance(data, memoryview): data = bytes(data) - return from_db(data.decode("utf8")) + return from_db(data.decode()) class VectorBinaryLoader(VectorLoader): diff --git a/vector_bench/client/qdrant.py b/vector_bench/client/qdrant.py new file mode 100644 index 0000000..d9f449f --- /dev/null +++ b/vector_bench/client/qdrant.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +from qdrant_client import QdrantClient +from qdrant_client.models import Distance as QdrantDistance +from qdrant_client.models import PointStruct, ScoredPoint, VectorParams + +from vector_bench.client.base import BaseClient +from vector_bench.spec import DatabaseConfig, Distance, Record + +DISTANCE_TO_QDRANT = { + Distance.COSINE: QdrantDistance.COSINE, + Distance.EUCLIDEAN: QdrantDistance.EUCLID, + Distance.DOT_PRODUCT: QdrantDistance.DOT, +} + + +class QdrantVectorClient(BaseClient): + dim: int + url: str + table: str + distance: Distance + + @classmethod + def from_config(cls, config: DatabaseConfig) -> QdrantVectorClient: + cls.dim = config.vector_dim + cls.url = config.url + cls.table = f"{config.table}_qdrant" + cls.distance = config.distance + + cls = QdrantVectorClient() + cls.init_db() + return cls + + def init_db(self): + self.client = QdrantClient(url=self.url) + collections_response = self.client.get_collections() + for collection in collections_response.collections: + if collection.name == self.table: + # already exists, return + return + + self.client.create_collection( + collection_name=self.table, + vectors_config=VectorParams( + size=self.dim, + distance=DISTANCE_TO_QDRANT[self.distance.__func__], + ), + ) + + def insert_batch(self, records: list[Record]): + self.client.upsert( + collection_name=self.table, + points=[ + PointStruct( + id=record.id, vector=record.vector.tolist(), payload=record.metadata + ) + for record in records + ], + ) + + def query(self, vector: list[float], top_k: int = 10) -> list[Record]: + points: list[ScoredPoint] = self.client.search( + collection_name=self.table, + query_vector=vector, + limit=top_k, + ) + return [ + Record(id=point.id, vector=point.vector, metadata=point.payload) + for point in points + ] diff --git a/vector_bench/dataset/source.py b/vector_bench/dataset/source.py index 4b16b60..1397d9b 100644 --- a/vector_bench/dataset/source.py +++ b/vector_bench/dataset/source.py @@ -11,6 +11,42 @@ link="https://ann-benchmarks.com/gist-960-euclidean.hdf5", ) +GLOVE_25_COSINE = DatasetConfig( + vector_dim=25, + num=1_200_000, + distance=Distance.COSINE, + type=FileType.H5, + path="datasets/glove-25-angular.hdf5", + link="https://ann-benchmarks.com/glove-25-angular.hdf5", +) + +GLOVE_100_COSINE = DatasetConfig( + vector_dim=100, + num=1_200_000, + distance=Distance.COSINE, + type=FileType.H5, + path="datasets/glove-100-angular.hdf5", + link="https://ann-benchmarks.com/glove-100-angular.hdf5", +) + +DEEP_96_COSINE = DatasetConfig( + vector_dim=96, + num=10_000_000, + distance=Distance.COSINE, + type=FileType.H5, + path="datasets/deep-image-96-angular.hdf5", + link="https://ann-benchmarks.com/deep-image-96-angular.hdf5", +) + +LAION_768_DOT_PRODUCT = DatasetConfig( + vector_dim=512, + num=5_000_000, + distance=Distance.DOT_PRODUCT, + type=FileType.H5, + path="datasets/laion-768-ip.hdf5", + link="https://myscale-datasets.s3.ap-southeast-1.amazonaws.com/laion-5m-test-ip.hdf5", +) + RANDOM_128_L2 = DatasetConfig( vector_dim=128, num=100_000, @@ -23,4 +59,8 @@ class DataSource(EnumSelector): GIST_960_L2 = GIST_960_L2 + GLOVE_25_COSINE = GLOVE_25_COSINE + GLOVE_100_COSINE = GLOVE_100_COSINE + DEEP_96_COSINE = DEEP_96_COSINE + LAION_768_DOT_PRODUCT = LAION_768_DOT_PRODUCT RANDOM_128_L2 = RANDOM_128_L2