Support milvus local

Signed-off-by: junjie.jiang <[email protected]>

Author: junjiejiangjjj

pymilvus/client/types.py

@@ -763,7 +763,6 @@ def groups(self):
 
 class ResourceGroupInfo:
     def __init__(self, resource_group: Any) -> None:
-
         self._name = resource_group.name
         self._capacity = resource_group.capacity
         self._num_available_node = resource_group.num_available_node

pymilvus/orm/connections.py

@@ -15,6 +15,7 @@
 import threading
 import time
 from typing import Callable, Tuple, Union
+import pathlib
 from urllib import parse
 
 from pymilvus.client.check import is_legal_address, is_legal_host, is_legal_port
@@ -357,6 +358,30 @@ def connect(
             >>> connections.connect("test", host="localhost", port="19530")
         """
 
+        if kwargs.get("uri") and parse.urlparse(kwargs["uri"]).scheme.lower() not in [
+            "unix",
+            "http",
+            "https",
+            "tcp",
+        ]:
+            # start and connect milvuslite
+            if kwargs["uri"].endswith("/"):
+                raise ConnectionConfigException(
+                    message=f"Open local milvus failed, {kwargs['uri']} is not a local file path"
+                )
+            parent_path = pathlib.Path(kwargs["uri"]).parent
+            if not parent_path.is_dir():
+                raise ConnectionConfigException(
+                    message=f"Open local milvus failed, dir: {str(parent_path)} not exists"
+                )
+
+            from milvus_lite.server_manager import server_manager_instance
+
+            local_uri = server_manager_instance.start_and_get_uri(kwargs["uri"])
+            if local_uri is None:
+                raise ConnectionConfigException(message="Open local milvus failed")
+            kwargs["uri"] = local_uri
+
         # kwargs_copy is used for auto reconnect
         kwargs_copy = copy.deepcopy(kwargs)
         kwargs_copy["user"] = user

pyproject.toml

@@ -29,6 +29,7 @@ dependencies=[
     "pyarrow>=12.0.0",
     "azure-storage-blob",
     "scipy",
+    "milvus-lite>=2.4.0",
 ]
 
 classifiers=[

requirements.txt

@@ -38,3 +38,4 @@ black
 requests
 minio
 azure-storage-blob
+milvus-lite>=2.4.0

tests/test_milvus_lite.py

@@ -0,0 +1,56 @@
+import os
+from tempfile import TemporaryDirectory
+import numpy as np
+
+from pymilvus.milvus_client import MilvusClient
+
+
+class TestMilvusLite:
+
+    def test_milvus_lite(self):
+        with TemporaryDirectory(dir='./') as root:
+            db_file = os.path.join(root, 'test.db')
+            client = MilvusClient(db_file)
+            client.create_collection(
+                collection_name="demo_collection",
+                dimension=3
+            )
+
+            # Text strings to search from.
+            docs = [
+                "Artificial intelligence was founded as an academic discipline in 1956.",
+                "Alan Turing was the first person to conduct substantial research in AI.",
+                "Born in Maida Vale, London, Turing was raised in southern England.",
+            ]
+
+            vectors = [[np.random.uniform(-1, 1) for _ in range(3) ] for _ in range(len(docs))]
+            data = [{"id": i, "vector": vectors[i], "text": docs[i], "subject": "history"} for i in range(len(vectors))]
+            res = client.insert(
+                collection_name="demo_collection",
+                data=data
+            )
+            assert res["insert_count"] == 3
+
+            res = client.search(
+                    collection_name="demo_collection",
+                    data=[vectors[0]],
+                    filter="subject == 'history'",
+                    limit=2,
+                    output_fields=["text", "subject"],
+                )
+            assert len(res[0]) == 2
+
+            # a query that retrieves all entities matching filter expressions.
+            res = client.query(
+                    collection_name="demo_collection",
+                    filter="subject == 'history'",
+                    output_fields=["text", "subject"],
+                )
+            assert len(res) == 3
+
+            # delete
+            res = client.delete(
+                    collection_name="demo_collection",
+                    filter="subject == 'history'",
+                )
+            assert len(res) == 3