mirror of
https://github.com/qdrant/qdrant.git
synced 2026-08-02 16:10:58 -05:00
424bc2c265
* Enable async_scorer in integration tests Enable `async_scorer=true` for all openapi and shell-based integration tests by setting the environment variable in `integration-tests.sh`. This exercises the io_uring vector storage path (on Linux) across all vector configurations tested there, with graceful fallback if io_uring is unavailable on the CI runner. Made-with: Cursor * Scope async_scorer to test_points_search consensus test only Revert the blanket async_scorer=true from integration-tests.sh to keep the normal mmap path covered. Instead, enable it via extra_env in just the test_points_search consensus test, so we exercise the io_uring path without losing coverage of the non-io_uring path on Linux. Made-with: Cursor * Parametrize test_points_search for sync/async scorer modes Use pytest.mark.parametrize to run test_points_search twice: once with the default scorer (sync_scorer) and once with async_scorer=true. This exercises both the normal mmap and io_uring paths in a single test. Made-with: Cursor --------- Co-authored-by: Cursor Agent <agent@cursor.com>
196 lines
6.6 KiB
Python
196 lines
6.6 KiB
Python
import pathlib
|
|
|
|
import pytest
|
|
|
|
from .utils import *
|
|
from .assertions import assert_http_ok
|
|
|
|
N_PEERS = 5
|
|
N_SHARDS = 4
|
|
N_REPLICA = 2
|
|
|
|
|
|
@pytest.mark.parametrize("async_scorer", [False, True], ids=["sync_scorer", "async_scorer"])
|
|
def test_points_search(tmp_path: pathlib.Path, async_scorer: bool):
|
|
extra_env = {}
|
|
if async_scorer:
|
|
extra_env["QDRANT__STORAGE__PERFORMANCE__ASYNC_SCORER"] = "true"
|
|
|
|
assert_project_root()
|
|
peer_dirs = make_peer_folders(tmp_path, N_PEERS)
|
|
|
|
# Gathers REST API uris
|
|
peer_api_uris = []
|
|
|
|
# Start bootstrap
|
|
(bootstrap_api_uri, bootstrap_uri) = start_first_peer(
|
|
peer_dirs[0], "peer_0_0.log", extra_env=extra_env)
|
|
peer_api_uris.append(bootstrap_api_uri)
|
|
|
|
# Wait for leader
|
|
leader = wait_peer_added(bootstrap_api_uri)
|
|
|
|
# Start other peers
|
|
for i in range(1, len(peer_dirs)):
|
|
peer_api_uris.append(start_peer(
|
|
peer_dirs[i], f"peer_0_{i}.log", bootstrap_uri, extra_env=extra_env))
|
|
|
|
# Wait for cluster
|
|
wait_for_uniform_cluster_status(peer_api_uris, leader)
|
|
|
|
# Check that there are no collections on all peers
|
|
for uri in peer_api_uris:
|
|
r_batch = requests.get(f"{uri}/collections")
|
|
assert_http_ok(r_batch)
|
|
assert len(r_batch.json()["result"]["collections"]) == 0
|
|
|
|
# Create collection in first peer
|
|
r_batch = requests.put(
|
|
f"{peer_api_uris[0]}/collections/test_collection", json={
|
|
"vectors": {
|
|
"size": 4,
|
|
"distance": "Dot"
|
|
},
|
|
"shard_number": N_SHARDS,
|
|
"replication_factor": N_REPLICA,
|
|
})
|
|
assert_http_ok(r_batch)
|
|
|
|
# Check that it exists on all peers
|
|
wait_collection_exists_and_active_on_all_peers(collection_name="test_collection", peer_api_uris=peer_api_uris)
|
|
|
|
# Check collection's cluster info
|
|
collection_cluster_info = get_collection_cluster_info(peer_api_uris[0], "test_collection")
|
|
assert collection_cluster_info["shard_count"] == N_SHARDS
|
|
|
|
# Create points in first peer's collection
|
|
r_batch = requests.put(
|
|
f"{peer_api_uris[0]}/collections/test_collection/points?wait=true", json={
|
|
"points": [
|
|
{
|
|
"id": 1,
|
|
"vector": [0.05, 0.61, 0.76, 0.74],
|
|
"payload": {
|
|
"city": "Berlin",
|
|
"country": "Germany",
|
|
"count": 1000000,
|
|
"square": 12.5,
|
|
"coords": {"lat": 1.0, "lon": 2.0}
|
|
}
|
|
},
|
|
{"id": 2, "vector": [0.19, 0.81, 0.75, 0.11],
|
|
"payload": {"city": ["Berlin", "London"]}},
|
|
{"id": 3, "vector": [0.36, 0.55, 0.47, 0.94],
|
|
"payload": {"city": ["Berlin", "Moscow"]}},
|
|
{"id": 4, "vector": [0.18, 0.01, 0.85, 0.80],
|
|
"payload": {"city": ["London", "Moscow"]}},
|
|
{"id": 5, "vector": [0.24, 0.18, 0.22, 0.44],
|
|
"payload": {"count": [0]}},
|
|
{"id": 6, "vector": [0.35, 0.08, 0.11, 0.44]},
|
|
{"id": 7, "vector": [0.45, 0.07, 0.21, 0.04]},
|
|
{"id": 8, "vector": [0.75, 0.18, 0.91, 0.48]},
|
|
{"id": 9, "vector": [0.30, 0.01, 0.1, 0.12]},
|
|
{"id": 10, "vector": [0.95, 0.8, 0.17, 0.19]}
|
|
]
|
|
})
|
|
assert_http_ok(r_batch)
|
|
|
|
# Check that 'search' & `search_batch` return the same results on all peers
|
|
q = {
|
|
"vector": [0.2, 0.1, 0.9, 0.7],
|
|
"top": 3,
|
|
"offset": 1,
|
|
"with_vector": True,
|
|
"with_payload": True,
|
|
"score_threshold": 0.1
|
|
}
|
|
|
|
# Capture result from first peer
|
|
r_init_search = requests.post(
|
|
f"{peer_api_uris[0]}/collections/test_collection/points/search", json=q
|
|
).json()["result"]
|
|
|
|
for uri in peer_api_uris:
|
|
r_search = requests.post(
|
|
f"{uri}/collections/test_collection/points/search", json=q
|
|
)
|
|
assert_http_ok(r_search)
|
|
|
|
r_batch = requests.post(
|
|
f"{uri}/collections/test_collection/points/search/batch", json={
|
|
"searches": [q]
|
|
}
|
|
)
|
|
assert_http_ok(r_batch)
|
|
|
|
# only one search in the batch
|
|
assert len(r_batch.json()["result"]) == 1
|
|
# assert same number of results
|
|
assert len(r_search.json()["result"]) == len(r_batch.json()["result"][0])
|
|
# assert stable across peers
|
|
assert r_search.json()["result"] == r_init_search
|
|
# search equivalent to single batch
|
|
assert r_search.json()["result"] == r_batch.json()["result"][0]
|
|
|
|
# Check that `search_batch` return the same results on all peers for duplicated searches
|
|
for uri in peer_api_uris:
|
|
r_batch = requests.post(
|
|
f"{uri}/collections/test_collection/points/search/batch", json={
|
|
"searches": [q, q, q, q]
|
|
}
|
|
)
|
|
assert_http_ok(r_batch)
|
|
|
|
# assert num searches
|
|
assert len(r_batch.json()["result"]) == 4
|
|
# assert the search limit
|
|
assert len(r_batch.json()["result"][0]) == 3
|
|
assert len(r_batch.json()["result"][1]) == 3
|
|
assert len(r_batch.json()["result"][2]) == 3
|
|
assert len(r_batch.json()["result"][3]) == 3
|
|
|
|
assert r_batch.json()["result"] == [r_init_search, r_init_search, r_init_search, r_init_search]
|
|
|
|
# Check that `search_batch` return the same results on all peers compared to multiple searches
|
|
q1 = {
|
|
"vector": [0.2, 0.1, 0.9, 0.7],
|
|
"top": 3,
|
|
"offset": 1,
|
|
"with_vector": True
|
|
}
|
|
q2 = {
|
|
"vector": [0.1, 0.2, 0.9, 0.7],
|
|
"top": 5,
|
|
"offset": 3,
|
|
"with_payload": True,
|
|
}
|
|
q3 = {
|
|
"vector": [0.2, 0.1, 0.7, 0.9],
|
|
"top": 10,
|
|
"score_threshold": 1.1
|
|
}
|
|
for uri in peer_api_uris:
|
|
r_batch = requests.post(
|
|
f"{uri}/collections/test_collection/points/search/batch", json={
|
|
"searches": [q1, q2, q3]
|
|
}
|
|
)
|
|
assert_http_ok(r_batch)
|
|
|
|
r_search_1 = requests.post(
|
|
f"{uri}/collections/test_collection/points/search", json=q1
|
|
)
|
|
assert_http_ok(r_search_1)
|
|
r_search_2 = requests.post(
|
|
f"{uri}/collections/test_collection/points/search", json=q2
|
|
)
|
|
assert_http_ok(r_search_2)
|
|
|
|
r_search_3 = requests.post(
|
|
f"{uri}/collections/test_collection/points/search", json=q3
|
|
)
|
|
assert_http_ok(r_search_3)
|
|
|
|
accumulated = [r_search_1.json()["result"], r_search_2.json()["result"], r_search_3.json()["result"]]
|
|
assert accumulated == r_batch.json()["result"]
|