chore: import upstream snapshot with attribution

2026-07-13 12:31:17 +08:00
commit 498b235461
5446 changed files with 2748612 additions and 0 deletions
@@ -0,0 +1,201 @@
+import threading
+import h5py
+import numpy as np
+import time
+import sys
+import copy
+from pathlib import Path
+from loguru import logger
+import pymilvus
+from pymilvus import (
+    connections,
+    FieldSchema, CollectionSchema, DataType,
+    Collection, utility
+)
+
+pymilvus_version = pymilvus.__version__
+
+
+all_index_types = ["IVF_FLAT", "IVF_SQ8", "HNSW"]
+default_index_params = [{"nlist": 128}, {"nlist": 128}, {"M": 48, "efConstruction": 200}]
+index_params_map = dict(zip(all_index_types, default_index_params))
+
+
+def gen_index_params(index_type, metric_type="L2"):
+    default_index = {"index_type": "IVF_FLAT", "params": {"nlist": 128}, "metric_type": metric_type}
+    index = copy.deepcopy(default_index)
+    index["index_type"] = index_type
+    index["params"] = index_params_map[index_type]
+    if index_type in ["BIN_FLAT", "BIN_IVF_FLAT"]:
+        index["metric_type"] = "HAMMING"
+    return index
+
+
+def gen_search_param(index_type, metric_type="L2"):
+    search_params = []
+    if index_type in ["FLAT", "IVF_FLAT", "IVF_SQ8", "IVF_PQ"]:
+        for nprobe in [10]:
+            ivf_search_params = {"metric_type": metric_type, "params": {"nprobe": nprobe}}
+            search_params.append(ivf_search_params)
+    elif index_type in ["BIN_FLAT", "BIN_IVF_FLAT"]:
+        for nprobe in [10]:
+            bin_search_params = {"metric_type": "HAMMING", "params": {"nprobe": nprobe}}
+            search_params.append(bin_search_params)
+    elif index_type in ["HNSW"]:
+        for ef in [150]:
+            hnsw_search_param = {"metric_type": metric_type, "params": {"ef": ef}}
+            search_params.append(hnsw_search_param)
+    elif index_type == "ANNOY":
+        for search_k in [1000]:
+            annoy_search_param = {"metric_type": metric_type, "params": {"search_k": search_k}}
+            search_params.append(annoy_search_param)
+    else:
+        logger.info("Invalid index_type.")
+        raise Exception("Invalid index_type.")
+    return search_params[0]
+
+
+def read_benchmark_hdf5(file_path):
+
+    f = h5py.File(file_path, 'r')
+    train = np.array(f["train"])
+    test = np.array(f["test"])
+    neighbors = np.array(f["neighbors"])
+    f.close()
+    return train, test, neighbors
+
+
+dim = 128
+TIMEOUT = 200
+
+
+def milvus_recall_test(host='127.0.0.1', index_type="HNSW"):
+    logger.info(f"recall test for index type {index_type}")
+    file_path = f"{str(Path(__file__).absolute().parent.parent.parent)}/assets/ann_hdf5/sift-128-euclidean.hdf5"
+    train, test, neighbors = read_benchmark_hdf5(file_path)
+    connections.connect(host=host, port="19530")
+    default_fields = [
+        FieldSchema(name="int64", dtype=DataType.INT64, is_primary=True),
+        FieldSchema(name="float", dtype=DataType.FLOAT),
+        FieldSchema(name="varchar", dtype=DataType.VARCHAR, max_length=65535),
+        FieldSchema(name="float_vector", dtype=DataType.FLOAT_VECTOR, dim=dim)
+    ]
+    default_schema = CollectionSchema(
+        fields=default_fields, description="test collection")
+
+    name = f"sift_128_euclidean_{index_type}"
+    logger.info(f"Create collection {name}")
+    collection = Collection(name=name, schema=default_schema)
+    nb = len(train)
+    batch_size = 50000
+    epoch = int(nb / batch_size)
+    t0 = time.time()
+    for i in range(epoch):
+        logger.info(f"epoch: {i}")
+        start = i * batch_size
+        end = (i + 1) * batch_size
+        if end > nb:
+            end = nb
+        data = [
+            [i for i in range(start, end)],
+            [np.float32(i) for i in range(start, end)],
+            [str(i) for i in range(start, end)],
+            train[start:end]
+        ]
+        collection.insert(data)
+    t1 = time.time()
+    logger.info(f"Insert {nb} vectors cost {t1 - t0:.4f} seconds")
+
+    t0 = time.time()
+    logger.info(f"Get collection entities...")
+    if pymilvus_version >= "2.2.0":
+        collection.flush()
+    else:
+        collection.num_entities
+    logger.info(collection.num_entities)
+    t1 = time.time()
+    logger.info(f"Get collection entities cost {t1 - t0:.4f} seconds")
+
+    # create index
+    default_index = gen_index_params(index_type)
+    logger.info(f"Create index...")
+    t0 = time.time()
+    collection.create_index(field_name="float_vector",
+                            index_params=default_index)
+    t1 = time.time()
+    logger.info(f"Create index cost {t1 - t0:.4f} seconds")
+
+    # load collection
+    replica_number = 1
+    logger.info(f"load collection...")
+    t0 = time.time()
+    collection.load(replica_number=replica_number)
+    t1 = time.time()
+    logger.info(f"load collection cost {t1 - t0:.4f} seconds")
+    res = utility.get_query_segment_info(name)
+    cnt = 0
+    logger.info(f"segments info: {res}")
+    for segment in res:
+        cnt += segment.num_rows
+    assert cnt == collection.num_entities
+    logger.info(f"wait for loading complete...")
+    time.sleep(30)
+    res = utility.get_query_segment_info(name)
+    logger.info(f"segments info: {res}")
+
+    # search
+    topK = 100
+    nq = 10000
+    current_search_params = gen_search_param(index_type)
+
+    # define output_fields of search result
+    for i in range(3):
+        t0 = time.time()
+        logger.info(f"Search...")
+        res = collection.search(
+            test[:nq], "float_vector", current_search_params, topK, output_fields=["int64"], timeout=TIMEOUT
+        )
+        t1 = time.time()
+        logger.info(f"search cost  {t1 - t0:.4f} seconds")
+        result_ids = []
+        for hits in res:
+            result_id = []
+            for hit in hits:
+                result_id.append(hit.entity.get("int64"))
+            result_ids.append(result_id)
+
+        # calculate recall
+        true_ids = neighbors[:nq, :topK]
+        sum_radio = 0.0
+        logger.info(f"Calculate recall...")
+        for index, item in enumerate(result_ids):
+            # tmp = set(item).intersection(set(flat_id_list[index]))
+            assert len(item) == len(true_ids[index])
+            tmp = set(true_ids[index]).intersection(set(item))
+            sum_radio = sum_radio + len(tmp) / len(item)
+        recall = round(sum_radio / len(result_ids), 6)
+        logger.info(f"recall={recall}")
+        if index_type in ["IVF_PQ", "ANNOY"]:
+            assert recall >= 0.6, f"recall={recall} < 0.6"
+        else:
+            assert 0.95 <= recall < 1.0, f"recall is {recall}, less than 0.95, greater than or equal to 1.0"
+    # query
+    expr = "int64 in [2,4,6,8]"
+    output_fields = ["int64", "float"]
+    res = collection.query(expr, output_fields, timeout=TIMEOUT)
+    sorted_res = sorted(res, key=lambda k: k['int64'])
+    for r in sorted_res:
+        logger.info(r)
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser(description='config for recall test')
+    parser.add_argument('--host', type=str,
+                        default="127.0.0.1", help='milvus server ip')
+    args = parser.parse_args()
+    host = args.host
+    tasks = []
+    for index_type in ["HNSW"]:
+        milvus_recall_test(host, index_type)
+