From 775462d3f77dcfe67d788a1d0dcb263be06aaa4b Mon Sep 17 00:00:00 2001
From: "junjie.jiang" <junjie.jiang@zilliz.com>
Date: Wed, 14 Dec 2022 19:42:56 +0800
Subject: [PATCH] update

Signed-off-by: junjie.jiang <junjie.jiang@zilliz.com>
---
 __init__.py      |  5 +++
 milvus_client.py | 90 ++++++++++++++++++++++++++++++++++++++++++++++++
 requirements.txt |  1 +
 3 files changed, 96 insertions(+)
 create mode 100644 __init__.py
 create mode 100644 milvus_client.py
 create mode 100644 requirements.txt

diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..675ca86
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,5 @@
+from .milvus_client import MilvusClientls
+
+
+def milvus_client(*args, **kwargs):
+    return MilvusClient(*args, **kwargs)
diff --git a/milvus_client.py b/milvus_client.py
new file mode 100644
index 0000000..2ff8056
--- /dev/null
+++ b/milvus_client.py
@@ -0,0 +1,90 @@
+from pymilvus import connections, Collection
+from towhee.operator import PyOperator
+import uuid
+
+
+class Milvus(PyOperator):
+    """
+    Search for embedding vectors in Milvus. Note that the Milvus collection has data before searching,
+
+    Args:
+        collection (`str`):
+            The collection name.
+        kwargs
+            The kwargs with collection.search, refer to https://milvus.io/docs/v2.0.x/search.md#Prepare-search-parameters.
+            And the `anns_field` defaults to the vector field name, `limit` defaults to 10, and `metric_type` in `param` defaults to 'L2'
+            if there has no index(FLAT), and for default index `param`:
+                IVF_FLAT: {"params": {"nprobe": 10}},
+                IVF_SQ8: {"params": {"nprobe": 10}},
+                IVF_PQ: {"params": {"nprobe": 10}},
+                HNSW: {"params": {"ef": 10}},
+                IVF_HNSW: {"params": {"nprobe": 10, "ef": 10}},
+                RHNSW_FLAT: {"params": {"ef": 10}},
+                RHNSW_SQ: {"params": {"ef": 10}},
+                RHNSW_PQ: {"params": {"ef": 10}},
+                ANNOY: {"params": {"search_k": 10}}.
+    """
+
+    def __init__(self, host: str = 'localhost', port: int = 19530, collection_name: str = None, **kwargs):
+        """
+        Get an existing collection.
+        """
+        self._host = host
+        self._port = port
+        self._collection_name = collection_name
+        self._connect_name = uuid.uuid4().hex
+        connections.connect(alias=self._connect_name, host=self._host, port=self._port)
+        self._collection = Collection(self._collection_name, using=self._connect_name)
+
+        self.kwargs = kwargs
+        if 'anns_field' not in self.kwargs:
+            fields_schema = self._collection.schema.fields
+            for schema in fields_schema:
+                if schema.dtype in (101, 100):
+                    self.kwargs['anns_field'] = schema.name
+
+        if 'limit' not in self.kwargs:
+            self.kwargs['limit'] = 10
+
+        index_params = {
+            'IVF_FLAT': {'params': {'nprobe': 10}},
+            'IVF_SQ8': {'params': {'nprobe': 10}},
+            'IVF_PQ': {'params': {'nprobe': 10}},
+            'HNSW': {'params': {'ef': 10}},
+            'RHNSW_FLAT': {'params': {'ef': 10}},
+            'RHNSW_SQ': {'params': {'ef': 10}},
+            'RHNSW_PQ': {'params': {'ef': 10}},
+            'IVF_HNSW': {'params': {'nprobe': 10, 'ef': 10}},
+            'ANNOY': {'params': {'search_k': 10}}
+        }
+
+        if 'param' not in self.kwargs:
+            if len(self._collection.indexes) != 0:
+                index_type = self._collection.indexes[0].params['index_type']
+                self.kwargs['param'] = index_params[index_type]
+            else:
+                self.kwargs['param'] = index_params['IVF_FLAT']
+            if 'metric_type' in self.kwargs:
+                self.kwargs['param']['metric_type'] = self.kwargs['metric_type']
+            else:
+                self.kwargs['param']['metric_type'] = 'L2'
+
+    def __call__(self, query: list):
+        milvus_result = self._collection.search(
+            data=[query],
+            **self.kwargs
+        )
+
+        result = []
+        for re in milvus_result:
+            row = []
+            for hit in re:
+                row.extend([hit.id, hit.score])
+                if 'output_fields' in self.kwargs:
+                    for k in self.kwargs['output_fields']:
+                        row.append(hit.entity._row_data[k])
+                result.append(row)
+        return result
+
+    def __del__(self):
+        connections.disconnect(self._connect_name)
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..e0472f0
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1 @@
+pymilvus