logo
Browse Source

update

Signed-off-by: junjie.jiang <junjie.jiang@zilliz.com>
main
junjie.jiang 2 years ago
parent
commit
c7a0b70fa9
  1. 4
      __init__.py
  2. 44
      milvus_client.py
  3. 1
      requirements.txt

4
__init__.py

@ -0,0 +1,4 @@
from .milvus_client import MilvusClient
def milvus_client(*args, **kwargs):
return MilvusClient(*args, **kwargs)

44
milvus_client.py

@ -0,0 +1,44 @@
import uuid
import logging
from towhee.operator import PyOperator
from pymilvus import connections, Collection
logger = logging.getLogger()
class MilvusClient(PyOperator):
"""
Milvus ANN index class.
"""
def __init__(self, host: str, port: int, collection_name: str):
"""
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)
def __call__(self, *data):
"""
Insert data to Milvus.
Args:
data (`list`):
The data to insert into milvus.
Returns:
A MutationResult object contains `insert_count` represents how many and a `primary_keys` of primary keys.
"""
vectors = []
for v in data:
vectors.append(v if isinstance(v, list) else [v])
mr = self._collection.insert(vectors)
if mr.err_count > 0:
raise RuntimeError("Insert to milvus failed")
return None

1
requirements.txt

@ -0,0 +1 @@
pymilvus
Loading…
Cancel
Save