mirror of
https://github.com/qdrant/qdrant.git
synced 2026-09-25 15:37:42 -05:00
45 lines
1.3 KiB
Python
45 lines
1.3 KiB
Python
"""Execute a batch through QueryBatchRequest and compare with individual queries."""
|
|
|
|
from tempfile import TemporaryDirectory
|
|
|
|
from qdrant_edge import (
|
|
Distance,
|
|
EdgeConfig,
|
|
EdgeShard,
|
|
EdgeVectorParams,
|
|
Point,
|
|
Query,
|
|
QueryBatchRequest,
|
|
QueryRequest,
|
|
UpdateOperation,
|
|
)
|
|
|
|
|
|
with TemporaryDirectory() as path:
|
|
shard = EdgeShard.create(
|
|
path, EdgeConfig(vectors=EdgeVectorParams(size=2, distance=Distance.Dot))
|
|
)
|
|
shard.update(
|
|
UpdateOperation.upsert_points(
|
|
[Point(1, [1.0, 0.0]), Point(2, [0.0, 1.0]), Point(3, [0.5, 0.5])]
|
|
)
|
|
)
|
|
request = QueryBatchRequest(
|
|
queries=[
|
|
QueryRequest(limit=1, query=Query.Nearest([1.0, 0.0])),
|
|
QueryRequest(limit=2, query=Query.Nearest([0.0, 1.0])),
|
|
]
|
|
)
|
|
batches = shard.query_batch(request)
|
|
assert len(batches) == len(request.queries)
|
|
for query, actual in zip(request.queries, batches):
|
|
expected = shard.query(query)
|
|
assert [(p.id, p.score) for p in actual] == [
|
|
(p.id, p.score) for p in expected
|
|
]
|
|
assert shard.query_batch(QueryBatchRequest(queries=[])) == []
|
|
assert repr(request).startswith("QueryBatchRequest(queries=")
|
|
shard.close()
|
|
|
|
print("QueryBatchRequest: batched results match individual queries")
|