Spaces:
Build error
Build error
File size: 3,342 Bytes
3932407 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
import pathlib
from .utils import *
from .assertions import assert_http_ok
N_PEERS = 5
N_SHARDS = 4
N_REPLICA = 2
def test_collection_sharding(tmp_path: pathlib.Path):
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")
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))
# 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 = requests.get(f"{uri}/collections")
assert_http_ok(r)
assert len(r.json()["result"]["collections"]) == 0
# Create collection in first peer
r = 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)
# 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 = 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]}
]
})
assert_http_ok(r)
# Check that 'search' returns the same results on all peers
for uri in peer_api_uris:
r = requests.post(
f"{uri}/collections/test_collection/points/search", json={
"vector": [0.2, 0.1, 0.9, 0.7],
"top": 3,
}
)
assert_http_ok(r)
assert r.json()["result"][0]["id"] == 4
assert r.json()["result"][1]["id"] == 1
assert r.json()["result"][2]["id"] == 3
|