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()