Support milvuslite

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

pymilvus/orm/connections.py

@@ -12,6 +12,7 @@
 
 import copy
 import logging
+import pathlib
 import threading
 import time
 from typing import Callable, Tuple, Union
@@ -240,11 +241,7 @@ def add_connection(self, **kwargs):
             self._alias[alias] = alias_config
 
     def __get_full_address(
-        self,
-        address: str = "",
-        uri: str = "",
-        host: str = "",
-        port: str = "",
+        self, address: str = "", uri: str = "", host: str = "", port: str = "",
     ) -> (str, parse.ParseResult):
         if address != "":
             if not is_legal_address(address):
@@ -357,6 +354,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: {parent_path} is 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,<2.5.0",
 ]
 
 classifiers=[

requirements.txt

@@ -38,3 +38,4 @@ black
 requests
 minio
 azure-storage-blob
+milvus_lite>=2.4.0,<2.5.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