-
Notifications
You must be signed in to change notification settings - Fork 28
/
customize_schema_auto_id.py
70 lines (53 loc) · 2.49 KB
/
customize_schema_auto_id.py
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
import numpy as np
from pymilvus import (
MilvusClient,
DataType
)
fmt = "\n=== {:30} ===\n"
dim = 8
collection_name = "hello_milvus"
milvus_client = MilvusClient(uri='./local_test.db')
has_collection = milvus_client.has_collection(collection_name, timeout=5)
if has_collection:
milvus_client.drop_collection(collection_name)
schema = milvus_client.create_schema(enable_dynamic_field=True, auto_id=True)
schema.add_field("id", DataType.INT64, is_primary=True)
schema.add_field("embeddings", DataType.FLOAT_VECTOR, dim=dim)
schema.add_field("title", DataType.VARCHAR, max_length=64)
index_params = milvus_client.prepare_index_params()
index_params.add_index(field_name = "embeddings", metric_type="L2")
milvus_client.create_collection(collection_name, schema=schema, index_params=index_params, consistency_level="Strong")
print(fmt.format(" all collections "))
print(milvus_client.list_collections())
print(fmt.format(f"schema of collection {collection_name}"))
print(milvus_client.describe_collection(collection_name))
rng = np.random.default_rng(seed=19530)
rows = [
{"embeddings": rng.random((1, dim))[0], "a": 100, "title": "t1"},
{"embeddings": rng.random((1, dim))[0], "b": 200, "title": "t2"},
{"embeddings": rng.random((1, dim))[0], "c": 300, "title": "t3"},
{"embeddings": rng.random((1, dim))[0], "d": 400, "title": "t4"},
{"embeddings": rng.random((1, dim))[0], "e": 500, "title": "t5"},
{"embeddings": rng.random((1, dim))[0], "f": 600, "title": "t6"},
]
print(fmt.format("Start inserting entities"))
insert_result = milvus_client.insert(collection_name, rows)
print(fmt.format("Inserting entities done"))
print(insert_result)
print(fmt.format("Start load collection "))
milvus_client.load_collection(collection_name)
print(fmt.format("Start query by specifying primary keys"))
query_results = milvus_client.query(collection_name, ids=insert_result['ids'][0])
print(query_results[0])
print(fmt.format("Start query by specifying filtering expression"))
query_results = milvus_client.query(collection_name, filter= "f == 600 or title == 't2'")
for ret in query_results:
print(ret)
rng = np.random.default_rng(seed=19530)
vectors_to_search = rng.random((1, dim))
print(fmt.format(f"Start search with retrieve serveral fields."))
result = milvus_client.search(collection_name, vectors_to_search, limit=3, output_fields=["pk", "a", "b"])
for hits in result:
for hit in hits:
print(f"hit: {hit}")
milvus_client.drop_collection(collection_name)