Skip to content

Commit

Permalink
feat: support qdrant
Browse files Browse the repository at this point in the history
Signed-off-by: Keming <[email protected]>
  • Loading branch information
kemingy committed Jan 31, 2024
1 parent 48b4906 commit 024db9b
Show file tree
Hide file tree
Showing 8 changed files with 90 additions and 6 deletions.
4 changes: 2 additions & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ 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:

Expand All @@ -24,7 +24,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
Expand Down
Original file line number Diff line number Diff line change
@@ -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
Expand Down
Original file line number Diff line number Diff line change
@@ -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
Expand Down
17 changes: 17 additions & 0 deletions docker/compose.qdrant.yaml
Original file line number Diff line number Diff line change
@@ -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"
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ dependencies = [
"tqdm~=4.66",
"httpx~=0.25",
"psycopg[binary]~=3.1",
"qdrant-client~=1.7.1",
]
[project.optional-dependencies]
dev = [
Expand Down
2 changes: 2 additions & 0 deletions vector_bench/client/__init__.py
Original file line number Diff line number Diff line change
@@ -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
2 changes: 1 addition & 1 deletion vector_bench/client/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
64 changes: 64 additions & 0 deletions vector_bench/client/qdrant.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,64 @@
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)
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
]

0 comments on commit 024db9b

Please sign in to comment.