Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[2.4]Support milvuslite #2075

Merged
merged 1 commit into from
May 8, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 25 additions & 0 deletions pymilvus/orm/connections.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@

import copy
import logging
import pathlib
import threading
import time
from typing import Callable, Tuple, Union
Expand Down Expand Up @@ -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: {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
Expand Down
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ dependencies=[
"pyarrow>=12.0.0",
"azure-storage-blob",
"scipy",
"milvus_lite>=2.4.0,<2.5.0",
]

classifiers=[
Expand Down
1 change: 1 addition & 0 deletions requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -38,3 +38,4 @@ black
requests
minio
azure-storage-blob
milvus_lite>=2.4.0,<2.5.0
56 changes: 56 additions & 0 deletions tests/test_milvus_lite.py
Original file line number Diff line number Diff line change
@@ -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