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=True) 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)) else: self.set_collection(collection_name) return "OK" except Exception as e: LOGGER.error("Failed to load data to Milvus: {}".format(e)) sys.exit(1)
def test_cmp(self, raw_dict_binary_vector): import copy field1 = FieldSchema.construct_from_dict(raw_dict_binary_vector) field2 = FieldSchema.construct_from_dict(raw_dict_binary_vector) assert field1 == field2 dict1 = copy.deepcopy(raw_dict_binary_vector) dict1["name"] = dict1["name"] + "_" field3 = FieldSchema.construct_from_dict(dict1) assert field1 != field3
def gen_binary_schema(): binary_fields = [ FieldSchema(name="int64", dtype=DataType.INT64, is_primary=True), FieldSchema(name="double", dtype=DataType.DOUBLE), FieldSchema(name=default_binary_vec_field_name, dtype=DataType.BINARY_VECTOR, dim=default_dim) ] default_schema = CollectionSchema(fields=binary_fields, description="test collection") return default_schema
def gen_default_fields_with_primary_key_1(): default_fields = [ FieldSchema(name="int64", dtype=DataType.INT64, is_primary=True), FieldSchema(name="double", dtype=DataType.DOUBLE), FieldSchema(name=default_float_vec_field_name, dtype=DataType.FLOAT_VECTOR, dim=default_dim) ] default_schema = CollectionSchema(fields=default_fields, description="test collection") return default_schema
def test_to_dict(self, raw_dict_norm, raw_dict_float_vector, raw_dict_binary_vector): fields = [] dicts = [raw_dict_norm, raw_dict_float_vector, raw_dict_binary_vector] fields.append(FieldSchema.construct_from_dict(raw_dict_norm)) fields.append(FieldSchema.construct_from_dict(raw_dict_float_vector)) fields.append(FieldSchema.construct_from_dict(raw_dict_binary_vector)) for i, f in enumerate(fields): target = f.to_dict() assert target == dicts[i] assert target is not dicts[i]
def gen_schema(): from pymilvus_orm.schema import CollectionSchema, FieldSchema fields = [ FieldSchema(gen_field_name(), DataType.INT64, is_primary=True, auto_id=False), FieldSchema(gen_field_name(), DataType.FLOAT), FieldSchema(gen_field_name(), DataType.FLOAT_VECTOR, dim=default_dim) ] collection_schema = CollectionSchema(fields) return collection_schema
def gen_default_fields(auto_id=True): default_fields = [ FieldSchema(name="count", dtype=DataType.INT64, is_primary=True), FieldSchema(name="float", dtype=DataType.FLOAT), FieldSchema(name=default_float_vec_field_name, dtype=DataType.FLOAT_VECTOR, dim=default_dim) ] default_schema = CollectionSchema( fields=default_fields, description="test collection", segment_row_limit=default_segment_row_limit, auto_id=False) return default_schema
def test_collection_by_DataFrame(self): from pymilvus_orm import Collection, connections from pymilvus_orm.schema import FieldSchema, CollectionSchema from pymilvus_orm.types import DataType fields = [ FieldSchema("int64", DataType.INT64), FieldSchema("float", DataType.FLOAT), FieldSchema("float_vector", DataType.FLOAT_VECTOR, dim=128) ] collection_schema = CollectionSchema(fields, primary_field="int64") connections.connect() collection = Collection(name=gen_collection_name(), schema=collection_schema) collection.drop()
def test_constructor_from_float_dict(self, raw_dict_float_vector): field = FieldSchema.construct_from_dict(raw_dict_float_vector) assert field.dtype == DataType.FLOAT_VECTOR assert field.description == raw_dict_float_vector['description'] assert field.is_primary == False assert field.name == raw_dict_float_vector['name'] assert field.dim == raw_dict_float_vector['params']['dim']
def gen_float_field(name=ct.default_float_field, is_primary=False, description=ct.default_desc): float_field = FieldSchema(name=name, dtype=DataType.FLOAT, description=description, is_primary=is_primary) return float_field
def gen_int64_field(name=ct.default_int64_field, is_primary=False, description=ct.default_desc): int64_field = FieldSchema(name=name, dtype=DataType.INT64, description=description, is_primary=is_primary) return int64_field
def test_constructor_from_norm_dict(self, raw_dict_norm): field = FieldSchema.construct_from_dict(raw_dict_norm) assert field.dtype == DataType.INT64 assert field.description == raw_dict_norm['description'] assert field.is_primary == False assert field.name == raw_dict_norm['name'] assert field.dim is None assert field.dummy is None
def gen_binary_vec_field(is_primary=False, dim=default_dim, description=binary_vec_field_desc): binary_vec_field = FieldSchema(name=default_binary_vec_field_name, dtype=DataType.BINARY_VECTOR, description=description, dim=dim, is_primary=is_primary) return binary_vec_field
def gen_float_vec_field(is_primary=False, dim=default_dim, description=float_vec_field_desc): float_vec_field = FieldSchema(name=default_float_vec_field_name, dtype=DataType.FLOAT_VECTOR, description=description, dim=dim, is_primary=is_primary) return float_vec_field
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 gen_all_type_fields(): fields = [] for k, v in DataType.__members__.items(): field = FieldSchema(name=k.lower(), dtype=v) fields.append(field) return fields