def create_collection(self, collection_name):
try:
if not self.has_collection(collection_name):
field1 = FieldSchema(name="id",
dtype=DataType.INT64,
descrition="int64",
is_primary=True,
auto_id=False)
field2 = FieldSchema(name="embedding",
dtype=DataType.FLOAT_VECTOR,
descrition="float vector",
dim=VECTOR_DIMENSION,
is_primary=False)
schema = CollectionSchema(fields=[field1, field2],
description="collection description")
self.collection = Collection(name=collection_name,
schema=schema)
LOGGER.debug("Create Milvus collection: {}".format(
self.collection))
return "OK"
else:
self.collection = Collection(collection_name)
return "collection {} exists".format(collection_name)
except Exception as e:
LOGGER.error("Failed to load data to Milvus: {}".format(e))
sys.exit(1)
def test_collection_only_name():
name = gen_unique_str()
collection_temp = Collection(name=name, schema=gen_default_fields())
collection = Collection(name=name)
data = gen_float_data(default_nb)
collection.insert(data)
collection.load()
assert collection.is_empty is False
assert collection.num_entities == default_nb
collection.drop()
def test_partition():
connections.connect(alias="default")
print("create collection")
collection = Collection(name=gen_unique_str(), schema=gen_default_fields())
print("create partition")
partition = Partition(collection, name=gen_unique_str())
print(list_collections())
assert has_partition(collection.name, partition.name) == True
data = gen_data(default_nb)
print("insert data to partition")
partition.insert(data)
assert partition.is_empty is False
assert partition.num_entities == default_nb
print("load partition")
partition.load()
topK = 5
search_params = {"metric_type": "L2", "params": {"nprobe": 10}}
print("search partition")
res = partition.search(data[2][-2:], "float_vector", search_params, topK,
"count > 100")
for hits in res:
for hit in hits:
print(hit)
print("release partition")
partition.release()
print("drop partition")
partition.drop()
print("drop collection")
collection.drop()
def test_create_index_binary_vector():
collection = Collection(name=gen_unique_str(), schema=gen_binary_schema())
data = gen_binary_data(default_nb)
collection.insert(data)
collection.create_index(field_name=default_binary_vec_field_name,
index_params=default_binary_index)
assert len(collection.indexes) != 0
collection.drop()
def set_collection(self, collection_name):
try:
if self.has_collection(collection_name):
self.collection = Collection(name=collection_name)
else:
raise Exception("There has no collection named:{}".format(collection_name))
except Exception as e:
LOGGER.error("Failed to load data to Milvus: {}".format(e))
sys.exit(1)
def test_collection_with_dataframe():
data = gen_dataframe(default_nb)
collection = Collection.construct_from_dataframe(name=gen_unique_str(),
dataframe=data,
primary_field="int64")
collection.load()
assert collection.is_empty is False
assert collection.num_entities == default_nb
collection.drop()
def test_specify_primary_key():
data = gen_float_data(default_nb)
collection = Collection(name=gen_unique_str(),
data=data,
schema=gen_default_fields_with_primary_key_1())
for index_param in gen_simple_index():
collection.create_index(field_name=default_float_vec_field_name,
index_params=index_param)
assert len(collection.indexes) != 0
collection.drop()
collection2 = Collection(name=gen_unique_str(),
data=data,
schema=gen_default_fields_with_primary_key_2())
for index_param in gen_simple_index():
collection2.create_index(field_name=default_float_vec_field_name,
index_params=index_param)
assert len(collection2.indexes) != 0
collection2.drop()
def test_create_index_float_vector():
data = gen_float_data(default_nb)
collection = Collection(name=gen_unique_str(),
data=data,
schema=gen_default_fields())
for index_param in gen_simple_index():
collection.create_index(field_name=default_float_vec_field_name,
index_params=index_param)
assert len(collection.indexes) != 0
collection.drop()
def insert(self, collection_name, vectors):
try:
self.create_collection(collection_name)
self.collection = Collection(name=collection_name)
data = [vectors]
mr = self.collection.insert(data)
ids = mr.primary_keys
self.collection.load()
LOGGER.debug(
"Insert vectors to Milvus in collection: {} with {} rows".
format(collection_name, len(vectors)))
return ids
except Exception as e:
LOGGER.error("Failed to load data to Milvus: {}".format(e))
sys.exit(1)
def create_collection(self, collection_name, collection_param):
try:
if not self.has_collection(collection_name):
field1 = FieldSchema(name="id",
dtype=DataType.INT64,
descrition="int64",
is_primary=True,
auto_id=False)
field2 = FieldSchema(name="embedding",
dtype=collection_param['data_type'],
descrition="float vector",
dim=collection_param['dimension'],
is_primary=False)
schema = CollectionSchema(fields=[field1, field2],
description="collection description")
self.collection = Collection(name=collection_name,
schema=schema)
self.logger.debug("Create Milvus collection: {}".format(
self.collection))
except Exception as e:
self.logger.error(e)
sys.exit(1)
def test_create_collection():
collection = Collection(name=gen_unique_str(), schema=gen_default_fields())
assert collection.is_empty is True
assert collection.num_entities == 0
collection.drop()
