|
| 1 | +import time |
| 2 | +import numpy as np |
| 3 | +from pymilvus import ( |
| 4 | + MilvusClient, |
| 5 | + DataType, |
| 6 | + Function, |
| 7 | + FunctionType, |
| 8 | + AnnSearchRequest, |
| 9 | +) |
| 10 | + |
| 11 | +fmt = "\n=== {:30} ===\n" |
| 12 | +dim = 8 |
| 13 | +collection_name = "hello_milvus" |
| 14 | +milvus_client = MilvusClient("http://localhost:19530") |
| 15 | + |
| 16 | +has_collection = milvus_client.has_collection(collection_name, timeout=5) |
| 17 | +if has_collection: |
| 18 | + milvus_client.drop_collection(collection_name) |
| 19 | + |
| 20 | +schema = milvus_client.create_schema(enable_dynamic_field=False, auto_id=True) |
| 21 | +schema.add_field("id", DataType.INT64, is_primary=True) |
| 22 | +schema.add_field("embeddings", DataType.FLOAT_VECTOR, dim=dim) |
| 23 | +schema.add_field("ts", DataType.INT64) |
| 24 | + |
| 25 | + |
| 26 | +index_params = milvus_client.prepare_index_params() |
| 27 | +index_params.add_index(field_name = "embeddings", index_type="FLAT", metric_type="L2") |
| 28 | +milvus_client.create_collection(collection_name, schema=schema, index_params=index_params, consistency_level="Strong") |
| 29 | + |
| 30 | +print(fmt.format(" all collections ")) |
| 31 | +print(milvus_client.list_collections()) |
| 32 | + |
| 33 | +print(fmt.format(f"schema of collection {collection_name}")) |
| 34 | +print(milvus_client.describe_collection(collection_name)) |
| 35 | + |
| 36 | +rng = np.random.default_rng(seed=19530) |
| 37 | +rows = [ |
| 38 | + {"embeddings": rng.random((1, dim))[0], "ts": 100}, |
| 39 | + {"embeddings": rng.random((1, dim))[0], "ts": 200}, |
| 40 | + {"embeddings": rng.random((1, dim))[0], "ts": 300}, |
| 41 | + {"embeddings": rng.random((1, dim))[0], "ts": 400}, |
| 42 | + {"embeddings": rng.random((1, dim))[0], "ts": 500}, |
| 43 | + {"embeddings": rng.random((1, dim))[0], "ts": 600}, |
| 44 | +] |
| 45 | + |
| 46 | +print(fmt.format("Start inserting entities")) |
| 47 | +insert_result = milvus_client.insert(collection_name, rows) |
| 48 | +print(fmt.format("Inserting entities done")) |
| 49 | +print(insert_result) |
| 50 | + |
| 51 | + |
| 52 | +print(fmt.format("Start load collection ")) |
| 53 | +milvus_client.load_collection(collection_name) |
| 54 | + |
| 55 | +rng = np.random.default_rng(seed=19530) |
| 56 | +vectors_to_search = rng.random((1, dim)) |
| 57 | + |
| 58 | +ranker = Function( |
| 59 | + name="rerank_fn", |
| 60 | + input_field_names=["ts"], |
| 61 | + function_type=FunctionType.RERANK, |
| 62 | + params={ |
| 63 | + "reranker": "decay", |
| 64 | + "function": "exp", |
| 65 | + "origin": 0, |
| 66 | + "offset": 200, |
| 67 | + "decay": 0.9, |
| 68 | + "scale": 100 |
| 69 | + } |
| 70 | +) |
| 71 | + |
| 72 | +print(fmt.format(f"Start search with retrieve serveral fields.")) |
| 73 | +result = milvus_client.search(collection_name, vectors_to_search, limit=3, output_fields=["*"], ranker=ranker) |
| 74 | +for hits in result: |
| 75 | + for hit in hits: |
| 76 | + print(f"hit: {hit}") |
| 77 | + |
| 78 | +vectors_to_search = rng.random((1, dim)) |
| 79 | +search_param = { |
| 80 | + "data": vectors_to_search, |
| 81 | + "anns_field": "embeddings", |
| 82 | + "param": {"metric_type": "L2"}, |
| 83 | + "limit": 3, |
| 84 | +} |
| 85 | +req = AnnSearchRequest(**search_param) |
| 86 | + |
| 87 | +hybrid_res = milvus_client.hybrid_search(collection_name, [req, req], ranker=ranker, limit=3, output_fields=["ts"]) |
| 88 | +for hits in hybrid_res: |
| 89 | + for hit in hits: |
| 90 | + print(f" hybrid search hit: {hit}") |
| 91 | + |
| 92 | +milvus_client.drop_collection(collection_name) |
0 commit comments