import json import logging import random import time import uuid from pathlib import Path import numpy as np import pyarrow as pa import pyarrow.parquet as pq import pytest from base.client_base import TestcaseBase from common import common_func as cf from common import common_type as ct from common.bulk_insert_data import ( DataField as df, ) from common.bulk_insert_data import ( gen_utf8_string, prepare_bulk_insert_json_files, prepare_bulk_insert_new_json_files, prepare_bulk_insert_numpy_files, prepare_bulk_insert_parquet_files, ) from common.common_params import DefaultVectorIndexParams, DefaultVectorSearchParams from common.common_type import CaseLabel, CheckTasks from common.minio_comm import copy_files_to_minio from faker import Faker from pymilvus import CollectionSchema, DataType, FieldSchema, Function, FunctionType from pymilvus.bulk_writer import BulkFileType, RemoteBulkWriter from utils.util_log import test_log as log fake = Faker() default_vec_only_fields = [df.vec_field] default_multi_fields = [df.vec_field, df.int_field, df.string_field, df.bool_field, df.float_field, df.array_int_field] default_vec_n_int_fields = [df.vec_field, df.int_field, df.array_int_field] # milvus_ns = "chaos-testing" base_dir = "/tmp/bulk_insert_data" def entity_suffix(entities): if entities // 1000000 > 0: suffix = f"{entities // 1000000}m" elif entities // 1000 > 0: suffix = f"{entities // 1000}k" else: suffix = f"{entities}" return suffix class TestcaseBaseBulkInsert(TestcaseBase): @pytest.fixture(scope="function", autouse=True) def init_minio_client(self, minio_host, minio_bucket): Path("/tmp/bulk_insert_data").mkdir(parents=True, exist_ok=True) minio_port = "9000" self.minio_endpoint = f"{minio_host}:{minio_port}" self.bucket_name = minio_bucket class TestBulkInsertNullableVector(TestcaseBaseBulkInsert): nullable_vector_field = df.float_vec_field @staticmethod def _float_vector(seed, dim): return [float(seed + i) / dim for i in range(dim)] @staticmethod def _int8_vector(seed, dim): return [((seed + i) % 255) - 128 for i in range(dim)] @staticmethod def _json_vector_value(vector_type, seed, dim): if vector_type in [DataType.FLOAT_VECTOR, DataType.FLOAT16_VECTOR, DataType.BFLOAT16_VECTOR]: return TestBulkInsertNullableVector._float_vector(seed, dim) if vector_type == DataType.BINARY_VECTOR: return [(seed + i) % 256 for i in range(dim // 8)] if vector_type == DataType.SPARSE_FLOAT_VECTOR: return {0: float(seed + 1), 1: float(seed + 2)} if vector_type == DataType.INT8_VECTOR: return TestBulkInsertNullableVector._int8_vector(seed, dim) raise ValueError(f"unsupported vector type: {vector_type}") @staticmethod def _search_vector_value(vector_type, seed, dim): if vector_type == DataType.BINARY_VECTOR: return cf.gen_binary_vectors(1, dim)[1] if vector_type == DataType.SPARSE_FLOAT_VECTOR: return cf.gen_sparse_vectors(1, dim) if vector_type == DataType.FLOAT16_VECTOR: return cf.gen_vectors(1, dim, vector_data_type=DataType.FLOAT16_VECTOR) if vector_type == DataType.BFLOAT16_VECTOR: return cf.gen_vectors(1, dim, vector_data_type=DataType.BFLOAT16_VECTOR) if vector_type == DataType.INT8_VECTOR: return [np.array(TestBulkInsertNullableVector._int8_vector(seed, dim), dtype=np.int8)] return [TestBulkInsertNullableVector._float_vector(seed, dim)] @staticmethod def _index_params(vector_type): if vector_type == DataType.BINARY_VECTOR: return ct.default_binary_index if vector_type == DataType.SPARSE_FLOAT_VECTOR: return ct.default_sparse_inverted_index if vector_type == DataType.INT8_VECTOR: return {"index_type": "HNSW", "metric_type": "COSINE", "params": {"M": 16, "efConstruction": 100}} return ct.default_index @staticmethod def _search_params(vector_type): if vector_type == DataType.BINARY_VECTOR: return ct.default_search_binary_params if vector_type == DataType.SPARSE_FLOAT_VECTOR: return ct.default_sparse_search_params if vector_type == DataType.INT8_VECTOR: return {"metric_type": "COSINE", "params": {"ef": 64}} return ct.default_search_params @staticmethod def _vector_field(vector_type, name, dim, nullable=True): if vector_type == DataType.FLOAT_VECTOR: return cf.gen_float_vec_field(name=name, dim=dim, nullable=nullable) if vector_type == DataType.BINARY_VECTOR: return cf.gen_binary_vec_field(name=name, dim=dim, nullable=nullable) if vector_type == DataType.FLOAT16_VECTOR: return cf.gen_float16_vec_field(name=name, dim=dim, nullable=nullable) if vector_type == DataType.BFLOAT16_VECTOR: return cf.gen_bfloat16_vec_field(name=name, dim=dim, nullable=nullable) if vector_type == DataType.SPARSE_FLOAT_VECTOR: return cf.gen_sparse_vec_field(name=name, nullable=nullable) if vector_type == DataType.INT8_VECTOR: return cf.gen_int8_vec_field(name=name, dim=dim, nullable=nullable) raise ValueError(f"unsupported vector type: {vector_type}") @staticmethod def _write_json_file_to_minio(minio_endpoint, bucket_name, rows, file_num=1): prefix = f"nullable-vector-json-{uuid.uuid4()}" files = [] rows_per_file = max(1, (len(rows) + file_num - 1) // file_num) for idx in range(file_num): batch_rows = rows[idx * rows_per_file : (idx + 1) * rows_per_file] if not batch_rows: continue relative_file = f"{prefix}/data_{idx}.json" local_file = Path(base_dir) / relative_file local_file.parent.mkdir(parents=True, exist_ok=True) with open(local_file, "w") as f: json.dump(batch_rows, f) files.append(relative_file) copy_files_to_minio( host=minio_endpoint, r_source=base_dir, files=files, bucket_name=bucket_name, force=True, ) return files @staticmethod def _write_parquet_file_to_minio(minio_endpoint, bucket_name, rows, vector_field): prefix = f"nullable-vector-parquet-{uuid.uuid4()}" relative_file = f"{prefix}/data.parquet" local_file = Path(base_dir) / relative_file local_file.parent.mkdir(parents=True, exist_ok=True) columns = { df.pk_field: pa.array([row[df.pk_field] for row in rows], type=pa.int64()), df.int_field: pa.array([row.get(df.int_field) for row in rows], type=pa.int64()), vector_field: pa.array([row.get(vector_field) for row in rows], type=pa.list_(pa.float32())), } if any(df.string_field in row for row in rows): columns[df.string_field] = pa.array([row.get(df.string_field) for row in rows], type=pa.string()) table = pa.table(columns) pq.write_table(table, local_file) copy_files_to_minio( host=minio_endpoint, r_source=base_dir, files=[relative_file], bucket_name=bucket_name, force=True, ) return [relative_file] def _create_nullable_float_vector_collection( self, c_name, dim, nullable=True, nullable_scalar=False, include_string_scalar=False ): fields = [ cf.gen_int64_field(name=df.pk_field, is_primary=True, auto_id=False), cf.gen_int64_field(name=df.int_field, nullable=nullable_scalar), cf.gen_float_vec_field(name=self.nullable_vector_field, dim=dim, nullable=nullable), ] if include_string_scalar: fields.insert(2, cf.gen_string_field(name=df.string_field, nullable=nullable_scalar)) schema = cf.gen_collection_schema(fields=fields, auto_id=False) self.collection_wrap.init_collection(c_name, schema=schema) return schema def _import_files(self, c_name, files, timeout=300, **kwargs): task_id, _ = self.utility_wrap.do_bulk_insert(collection_name=c_name, files=files, **kwargs) logging.info(f"bulk insert task ids:{task_id}") success, states = self.utility_wrap.wait_for_bulk_insert_tasks_completed(task_ids=[task_id], timeout=timeout) log.info(f"bulk insert state:{success} with states:{states}") return success, states def _assert_nullable_float_vector_data(self, c_name, entities, null_ids, non_null_ids, dim): assert self.collection_wrap.num_entities == entities self.collection_wrap.create_index(field_name=self.nullable_vector_field, index_params=ct.default_index) self.utility_wrap.wait_for_index_building_complete(c_name, timeout=300) self.collection_wrap.load() results, _ = self.collection_wrap.query( expr=f"{df.pk_field} >= 0", output_fields=[df.pk_field, self.nullable_vector_field], ) assert len(results) == entities rows_by_id = {row[df.pk_field]: row for row in results} for pk in null_ids: assert rows_by_id[pk][self.nullable_vector_field] is None for pk in non_null_ids: vector = rows_by_id[pk][self.nullable_vector_field] assert vector is not None assert len(vector) == dim search_res, _ = self.collection_wrap.search( [self._float_vector(100, dim)], self.nullable_vector_field, param=ct.default_search_params, limit=len(non_null_ids), ) returned_ids = set(search_res[0].ids) assert returned_ids assert len(returned_ids) == len(non_null_ids) assert returned_ids.issubset(set(non_null_ids)) assert returned_ids.isdisjoint(set(null_ids)) @pytest.mark.tags(CaseLabel.L0) @pytest.mark.parametrize("dim", [32]) @pytest.mark.parametrize("entities", [60]) def test_bulk_insert_json_nullable_float_vector_null_and_omit(self, dim, entities): """ target: verify JSON import supports nullable vector/scalar fields and multiple vector columns method: import row-based JSON with mixed valid/null/omitted dense and sparse vectors plus nullable scalars expected: import succeeds; query returns None for NULL fields; each search skips NULL vectors on its field """ self._connect() c_name = cf.gen_unique_str("bulk_insert_nullable_vector") fields = [ cf.gen_int64_field(name=df.pk_field, is_primary=True, auto_id=False), cf.gen_int64_field(name=df.int_field, nullable=True), cf.gen_string_field(name=df.string_field, nullable=True), cf.gen_float_vec_field(name=self.nullable_vector_field, dim=dim, nullable=True), cf.gen_sparse_vec_field(name=df.sparse_vec_field, nullable=True), ] schema = cf.gen_collection_schema(fields=fields, auto_id=False) self.collection_wrap.init_collection(c_name, schema=schema) rows = [] float_null_ids = [] float_non_null_ids = [] sparse_null_ids = [] sparse_non_null_ids = [] int_null_ids = [] string_null_ids = [] for i in range(entities): row = { df.pk_field: i, df.int_field: None if i % 5 == 0 else i, df.string_field: None if i % 4 == 0 else f"string_{i}", } if row[df.int_field] is None: int_null_ids.append(i) if row[df.string_field] is None: string_null_ids.append(i) if i % 4 == 0: row[self.nullable_vector_field] = None float_null_ids.append(i) elif i % 4 == 1: float_null_ids.append(i) else: row[self.nullable_vector_field] = self._float_vector(i, dim) float_non_null_ids.append(i) if i % 3 == 0: row[df.sparse_vec_field] = None sparse_null_ids.append(i) elif i % 3 == 1: sparse_null_ids.append(i) else: row[df.sparse_vec_field] = cf.gen_sparse_vectors(1, dim)[0] sparse_non_null_ids.append(i) rows.append(row) files = self._write_json_file_to_minio(self.minio_endpoint, self.bucket_name, rows) success, _ = self._import_files(c_name, files) assert success assert self.collection_wrap.num_entities == entities dense_index_name = "nullable_dense_idx" sparse_index_name = "nullable_sparse_idx" self.collection_wrap.create_index( field_name=self.nullable_vector_field, index_params=ct.default_index, index_name=dense_index_name ) self.collection_wrap.create_index( field_name=df.sparse_vec_field, index_params=ct.default_sparse_inverted_index, index_name=sparse_index_name ) self.utility_wrap.wait_for_index_building_complete(c_name, index_name=dense_index_name, timeout=300) self.utility_wrap.wait_for_index_building_complete(c_name, index_name=sparse_index_name, timeout=300) self.collection_wrap.load() results, _ = self.collection_wrap.query( expr=f"{df.pk_field} >= 0", output_fields=[df.pk_field, df.int_field, df.string_field, self.nullable_vector_field, df.sparse_vec_field], ) assert len(results) == entities rows_by_id = {row[df.pk_field]: row for row in results} for pk in int_null_ids: assert rows_by_id[pk][df.int_field] is None for pk in string_null_ids: assert rows_by_id[pk][df.string_field] is None for pk in float_null_ids: assert rows_by_id[pk][self.nullable_vector_field] is None for pk in float_non_null_ids: assert rows_by_id[pk][self.nullable_vector_field] is not None for pk in sparse_null_ids: assert rows_by_id[pk][df.sparse_vec_field] is None for pk in sparse_non_null_ids: assert rows_by_id[pk][df.sparse_vec_field] is not None search_res, _ = self.collection_wrap.search( [self._float_vector(100, dim)], self.nullable_vector_field, param=ct.default_search_params, limit=len(float_non_null_ids), ) returned_ids = set(search_res[0].ids) assert returned_ids assert len(returned_ids) == len(float_non_null_ids) assert returned_ids.issubset(set(float_non_null_ids)) assert returned_ids.isdisjoint(set(float_null_ids)) search_res, _ = self.collection_wrap.search( cf.gen_sparse_vectors(1, dim), df.sparse_vec_field, param=ct.default_sparse_search_params, limit=len(sparse_non_null_ids), ) returned_ids = set(search_res[0].ids) assert returned_ids assert len(returned_ids) == len(sparse_non_null_ids) assert returned_ids.issubset(set(sparse_non_null_ids)) assert returned_ids.isdisjoint(set(sparse_null_ids)) @pytest.mark.tags(CaseLabel.L1) @pytest.mark.parametrize("dim", [32]) @pytest.mark.parametrize("entities", [12]) def test_bulk_insert_parquet_nullable_float_vector(self, dim, entities): """ target: verify Parquet import supports nullable vector through Arrow List validity bitmap method: import a Parquet file whose nullable vector column contains valid lists and nulls expected: import succeeds; query returns None for NULL vectors; search skips NULL vectors """ self._connect() c_name = cf.gen_unique_str("bulk_insert_nullable_vector") self._create_nullable_float_vector_collection(c_name, dim, nullable_scalar=True, include_string_scalar=True) rows = [] null_ids = [] non_null_ids = [] int_null_ids = [] string_null_ids = [] for i in range(entities): row = { df.pk_field: i, df.int_field: None if i % 4 == 0 else i, df.string_field: None if i % 5 == 0 else f"string_{i}", } if row[df.int_field] is None: int_null_ids.append(i) if row[df.string_field] is None: string_null_ids.append(i) if i % 3 == 0: row[self.nullable_vector_field] = None null_ids.append(i) else: row[self.nullable_vector_field] = self._float_vector(i, dim) non_null_ids.append(i) rows.append(row) files = self._write_parquet_file_to_minio( self.minio_endpoint, self.bucket_name, rows, self.nullable_vector_field ) success, _ = self._import_files(c_name, files) assert success self._assert_nullable_float_vector_data(c_name, entities, null_ids, non_null_ids, dim) results, _ = self.collection_wrap.query( expr=f"{df.pk_field} >= 0", output_fields=[df.pk_field, df.int_field, df.string_field], ) rows_by_id = {row[df.pk_field]: row for row in results} for pk in int_null_ids: assert rows_by_id[pk][df.int_field] is None for pk in string_null_ids: assert rows_by_id[pk][df.string_field] is None @pytest.mark.tags(CaseLabel.L1) @pytest.mark.parametrize( "file_type", [ BulkFileType.JSON, BulkFileType.CSV, BulkFileType.PARQUET, ], ) @pytest.mark.parametrize("dim", [32]) @pytest.mark.parametrize("entities", [60]) def test_bulk_writer_nullable_float_vector(self, file_type, dim, entities): """ target: verify BulkWriter can write nullable vector values for import method: write JSON/CSV/Parquet files with RemoteBulkWriter, then import generated batches expected: import succeeds; query returns None for NULL vectors; search skips NULL vectors """ self._connect() c_name = cf.gen_unique_str("bulk_writer_nullable_vector") schema = self._create_nullable_float_vector_collection(c_name, dim) null_ids = [] non_null_ids = [] with RemoteBulkWriter( schema=schema, remote_path="bulk_data", connect_param=RemoteBulkWriter.ConnectParam( bucket_name=self.bucket_name, endpoint=self.minio_endpoint, access_key="minioadmin", secret_key="minioadmin", ), file_type=file_type, chunk_size=2048 if file_type == BulkFileType.PARQUET else 1024 * 1024 * 1024, ) as remote_writer: for i in range(entities): vector = None if i % 3 == 0 else self._float_vector(i, dim) if vector is None: null_ids.append(i) else: non_null_ids.append(i) remote_writer.append_row( { df.pk_field: i, df.int_field: i, self.nullable_vector_field: vector, } ) remote_writer.commit() files = remote_writer.batch_files if file_type == BulkFileType.PARQUET: assert len(files) > 1 for files_in_batch in files: success, _ = self._import_files(c_name, files_in_batch) assert success self._assert_nullable_float_vector_data(c_name, entities, null_ids, non_null_ids, dim) @pytest.mark.tags(CaseLabel.L2) @pytest.mark.parametrize("dim", [32]) @pytest.mark.parametrize("entities", [6]) def test_bulk_insert_non_nullable_float_vector_with_null_failed(self, dim, entities): """ target: verify non-nullable vector fields still reject NULL values in import method: import JSON data with null values for a non-nullable vector field expected: import task fails with a field/value validation error """ self._connect() c_name = cf.gen_unique_str("bulk_insert_non_nullable_vector") self._create_nullable_float_vector_collection(c_name, dim, nullable=False) rows = [] for i in range(entities): rows.append( { df.pk_field: i, df.int_field: i, self.nullable_vector_field: None if i == 0 else self._float_vector(i, dim), } ) files = self._write_json_file_to_minio(self.minio_endpoint, self.bucket_name, rows) success, states = self._import_files(c_name, files) assert not success for state in states.values(): assert state.state_name in ["Failed", "Failed and cleaned"] failed_reason = state.infos.get("failed_reason", "") assert f"expected type 'FloatVector' for field '{self.nullable_vector_field}'" in failed_reason assert "got type ''" in failed_reason @pytest.mark.tags(CaseLabel.L2) @pytest.mark.parametrize( "vector_type,vector_field", [ (DataType.FLOAT_VECTOR, df.float_vec_field), (DataType.BINARY_VECTOR, df.binary_vec_field), (DataType.FLOAT16_VECTOR, df.fp16_vec_field), (DataType.BFLOAT16_VECTOR, df.bf16_vec_field), (DataType.SPARSE_FLOAT_VECTOR, df.sparse_vec_field), (DataType.INT8_VECTOR, ct.default_int8_vec_field_name), ], ) @pytest.mark.parametrize("dim", [32]) @pytest.mark.parametrize("entities", [6]) def test_bulk_insert_json_nullable_vector_all_types(self, vector_type, vector_field, dim, entities): """ target: verify JSON import supports nullable values for every supported vector dtype method: import mixed NULL/non-NULL rows for each vector type, then query and search the vector field expected: import succeeds; query returns None for NULL vectors; search skips NULL vectors """ self._connect() c_name = cf.gen_unique_str("bulk_insert_nullable_vector_types") fields = [ cf.gen_int64_field(name=df.pk_field, is_primary=True, auto_id=False), cf.gen_int64_field(name=df.int_field), self._vector_field(vector_type, vector_field, dim, nullable=True), ] schema = cf.gen_collection_schema(fields=fields, auto_id=False) self.collection_wrap.init_collection(c_name, schema=schema) rows = [] null_ids = [] non_null_ids = [] for i in range(entities): row = {df.pk_field: i, df.int_field: i} if i % 2 == 0: row[vector_field] = None null_ids.append(i) else: row[vector_field] = self._json_vector_value(vector_type, i, dim) non_null_ids.append(i) rows.append(row) files = self._write_json_file_to_minio(self.minio_endpoint, self.bucket_name, rows) success, _ = self._import_files(c_name, files) assert success self.collection_wrap.create_index(field_name=vector_field, index_params=self._index_params(vector_type)) self.utility_wrap.wait_for_index_building_complete(c_name, timeout=300) self.collection_wrap.load() results, _ = self.collection_wrap.query( expr=f"{df.pk_field} >= 0", output_fields=[df.pk_field, vector_field], ) assert len(results) == entities rows_by_id = {row[df.pk_field]: row for row in results} for pk in null_ids: assert rows_by_id[pk][vector_field] is None for pk in non_null_ids: assert rows_by_id[pk][vector_field] is not None search_res, _ = self.collection_wrap.search( self._search_vector_value(vector_type, 100, dim), vector_field, param=self._search_params(vector_type), limit=len(non_null_ids), ) returned_ids = set(search_res[0].ids) assert returned_ids assert len(returned_ids) == len(non_null_ids) assert returned_ids.issubset(set(non_null_ids)) assert returned_ids.isdisjoint(set(null_ids)) class TestBulkInsert(TestcaseBaseBulkInsert): @pytest.mark.tags(CaseLabel.L0) @pytest.mark.parametrize("is_row_based", [True]) @pytest.mark.parametrize("auto_id", [True, False]) @pytest.mark.parametrize("dim", [128]) # 8, 128 @pytest.mark.parametrize("entities", [100]) # 100, 1000 def test_float_vector_only(self, is_row_based, auto_id, dim, entities): """ collection: auto_id, customized_id collection schema: [pk, float_vector] Steps: 1. create collection 2. import data 3. verify the data entities equal the import data 4. load the collection 5. verify search successfully 6. verify query successfully """ self._connect() c_name = cf.gen_unique_str("bulk_insert") fields = [ cf.gen_int64_field(name=df.pk_field, is_primary=True, auto_id=auto_id), cf.gen_float_vec_field(name=df.float_vec_field, dim=dim), ] data_fields = [f.name for f in fields if not f.to_dict().get("auto_id", False)] files = prepare_bulk_insert_new_json_files( minio_endpoint=self.minio_endpoint, bucket_name=self.bucket_name, is_row_based=is_row_based, rows=entities, dim=dim, auto_id=auto_id, data_fields=data_fields, force=True, ) schema = cf.gen_collection_schema(fields=fields) self.collection_wrap.init_collection(c_name, schema=schema) # import data t0 = time.time() task_id, _ = self.utility_wrap.do_bulk_insert( collection_name=c_name, partition_name=None, files=files, ) logging.info(f"bulk insert task id:{task_id}") success, _ = self.utility_wrap.wait_for_bulk_insert_tasks_completed(task_ids=[task_id], timeout=300) tt = time.time() - t0 log.info(f"bulk insert state:{success} in {tt}") assert success num_entities = self.collection_wrap.num_entities log.info(f" collection entities: {num_entities}") assert num_entities == entities # verify imported data is available for search index_params = ct.default_index self.collection_wrap.create_index(field_name=df.float_vec_field, index_params=index_params) time.sleep(2) self.utility_wrap.wait_for_index_building_complete(c_name, timeout=300) res, _ = self.utility_wrap.index_building_progress(c_name) log.info(f"index building progress: {res}") self.collection_wrap.load() self.collection_wrap.load(_refresh=True) log.info("wait for load finished and be ready for search") time.sleep(2) log.info(f"query seg info: {self.utility_wrap.get_query_segment_info(c_name)[0]}") nq = 2 topk = 2 search_data = cf.gen_vectors(nq, dim) search_params = ct.default_search_params res, _ = self.collection_wrap.search( search_data, df.float_vec_field, param=search_params, limit=topk, check_task=CheckTasks.check_search_results, check_items={"nq": nq, "limit": topk}, ) for hits in res: ids = hits.ids results, _ = self.collection_wrap.query(expr=f"{df.pk_field} in {ids}") assert len(results) == len(ids) @pytest.mark.tags(CaseLabel.L0) @pytest.mark.parametrize("is_row_based", [True]) @pytest.mark.parametrize("dim", [128]) # 8 @pytest.mark.parametrize("entities", [100]) # 100 def test_str_pk_float_vector_only(self, is_row_based, dim, entities): """ collection schema: [str_pk, float_vector] Steps: 1. create collection 2. import data 3. verify the data entities equal the import data 4. load the collection 5. verify search successfully 6. verify query successfully """ auto_id = False # no auto id for string_pk schema string_pk = True self._connect() c_name = cf.gen_unique_str("bulk_insert") fields = [ cf.gen_string_field(name=df.string_field, is_primary=True, auto_id=auto_id), cf.gen_float_vec_field(name=df.float_vec_field, dim=dim), ] schema = cf.gen_collection_schema(fields=fields) data_fields = [f.name for f in fields if not f.to_dict().get("auto_id", False)] files = prepare_bulk_insert_new_json_files( minio_endpoint=self.minio_endpoint, bucket_name=self.bucket_name, is_row_based=is_row_based, rows=entities, dim=dim, auto_id=auto_id, str_pk=string_pk, data_fields=data_fields, schema=schema, ) self.collection_wrap.init_collection(c_name, schema=schema) # import data t0 = time.time() task_id, _ = self.utility_wrap.do_bulk_insert(collection_name=c_name, files=files) logging.info(f"bulk insert task ids:{task_id}") completed, _ = self.utility_wrap.wait_for_bulk_insert_tasks_completed(task_ids=[task_id], timeout=300) tt = time.time() - t0 log.info(f"bulk insert state:{completed} in {tt}") assert completed num_entities = self.collection_wrap.num_entities log.info(f" collection entities: {num_entities}") assert num_entities == entities # verify imported data is available for search index_params = ct.default_index self.collection_wrap.create_index(field_name=df.float_vec_field, index_params=index_params) self.utility_wrap.wait_for_index_building_complete(c_name, timeout=300) res, _ = self.utility_wrap.index_building_progress(c_name) log.info(f"index building progress: {res}") self.collection_wrap.load() self.collection_wrap.load(_refresh=True) log.info("wait for load finished and be ready for search") time.sleep(2) log.info(f"query seg info: {self.utility_wrap.get_query_segment_info(c_name)[0]}") nq = 3 topk = 2 search_data = cf.gen_vectors(nq, dim) search_params = ct.default_search_params time.sleep(2) res, _ = self.collection_wrap.search( search_data, df.float_vec_field, param=search_params, limit=topk, check_task=CheckTasks.check_search_results, check_items={"nq": nq, "limit": topk}, ) for hits in res: ids = hits.ids expr = f"{df.string_field} in {ids}" expr = expr.replace("'", '"') results, _ = self.collection_wrap.query(expr=expr) assert len(results) == len(ids) @pytest.mark.tags(CaseLabel.L0) @pytest.mark.parametrize("is_row_based", [True]) @pytest.mark.parametrize("auto_id", [True, False]) @pytest.mark.parametrize("dim", [128]) @pytest.mark.parametrize("entities", [2000]) def test_partition_float_vector_int_scalar(self, is_row_based, auto_id, dim, entities): """ collection: customized partitions collection schema: [pk, float_vectors, int_scalar] 1. create collection and a partition 2. build index and load partition 3. import data into the partition 4. verify num entities 5. verify index status 6. verify search and query """ files = prepare_bulk_insert_json_files( minio_endpoint=self.minio_endpoint, bucket_name=self.bucket_name, is_row_based=is_row_based, rows=entities, dim=dim, auto_id=auto_id, data_fields=default_vec_n_int_fields, file_nums=1, ) self._connect() c_name = cf.gen_unique_str("bulk_insert") fields = [ cf.gen_int64_field(name=df.pk_field, is_primary=True), cf.gen_float_vec_field(name=df.vec_field, dim=dim), cf.gen_int32_field(name=df.int_field), cf.gen_array_field(name=df.array_int_field, element_type=DataType.INT32), ] schema = cf.gen_collection_schema(fields=fields, auto_id=auto_id) self.collection_wrap.init_collection(c_name, schema=schema) # create a partition p_name = cf.gen_unique_str("bulk_insert") m_partition, _ = self.collection_wrap.create_partition(partition_name=p_name) # build index before bulk insert index_params = ct.default_index self.collection_wrap.create_index(field_name=df.vec_field, index_params=index_params) # load before bulk insert self.collection_wrap.load(partition_names=[p_name]) # import data into the partition t0 = time.time() task_id, _ = self.utility_wrap.do_bulk_insert( collection_name=c_name, partition_name=p_name, files=files, ) logging.info(f"bulk insert task ids:{task_id}") success, state = self.utility_wrap.wait_for_bulk_insert_tasks_completed(task_ids=[task_id], timeout=300) tt = time.time() - t0 log.info(f"bulk insert state:{success} in {tt}") assert success assert m_partition.num_entities == entities assert self.collection_wrap.num_entities == entities log.debug(state) time.sleep(2) self.utility_wrap.wait_for_index_building_complete(c_name, timeout=300) res, _ = self.utility_wrap.index_building_progress(c_name) log.info(f"index building progress: {res}") log.info("wait for load finished and be ready for search") self.collection_wrap.load(_refresh=True) time.sleep(2) log.info(f"query seg info: {self.utility_wrap.get_query_segment_info(c_name)[0]}") nq = 10 topk = 5 search_data = cf.gen_vectors(nq, dim) search_params = ct.default_search_params res, _ = self.collection_wrap.search( search_data, df.vec_field, param=search_params, limit=topk, check_task=CheckTasks.check_search_results, check_items={"nq": nq, "limit": topk}, ) for hits in res: ids = hits.ids results, _ = self.collection_wrap.query(expr=f"{df.pk_field} in {ids}") assert len(results) == len(ids) @pytest.mark.tags(CaseLabel.L1) @pytest.mark.parametrize("is_row_based", [True]) @pytest.mark.parametrize("auto_id", [True, False]) @pytest.mark.parametrize("dim", [128]) @pytest.mark.parametrize("entities", [2000]) def test_binary_vector_json(self, is_row_based, auto_id, dim, entities): """ collection schema: [pk, binary_vector] Steps: 1. create collection 2. create index and load collection 3. import data 4. verify build status 5. verify the data entities 6. load collection 7. verify search successfully 6. verify query successfully """ float_vec = False files = prepare_bulk_insert_json_files( minio_endpoint=self.minio_endpoint, bucket_name=self.bucket_name, is_row_based=is_row_based, rows=entities, dim=dim, auto_id=auto_id, float_vector=float_vec, data_fields=default_vec_only_fields, ) self._connect() c_name = cf.gen_unique_str("bulk_insert") fields = [ cf.gen_int64_field(name=df.pk_field, is_primary=True), cf.gen_binary_vec_field(name=df.vec_field, dim=dim), ] schema = cf.gen_collection_schema(fields=fields, auto_id=auto_id) self.collection_wrap.init_collection(c_name, schema=schema) # build index before bulk insert binary_index_params = { "index_type": "BIN_IVF_FLAT", "metric_type": "JACCARD", "params": {"nlist": 64}, } self.collection_wrap.create_index(field_name=df.vec_field, index_params=binary_index_params) # load collection self.collection_wrap.load() # import data t0 = time.time() task_id, _ = self.utility_wrap.do_bulk_insert(collection_name=c_name, files=files) logging.info(f"bulk insert task ids:{task_id}") success, _ = self.utility_wrap.wait_for_bulk_insert_tasks_completed(task_ids=[task_id], timeout=300) tt = time.time() - t0 log.info(f"bulk insert state:{success} in {tt}") assert success time.sleep(2) self.utility_wrap.wait_for_index_building_complete(c_name, timeout=300) res, _ = self.utility_wrap.index_building_progress(c_name) log.info(f"index building progress: {res}") # verify num entities assert self.collection_wrap.num_entities == entities # verify search and query log.info("wait for load finished and be ready for search") self.collection_wrap.load(_refresh=True) time.sleep(2) search_data = cf.gen_binary_vectors(1, dim)[1] search_params = {"metric_type": "JACCARD", "params": {"nprobe": 10}} res, _ = self.collection_wrap.search( search_data, df.vec_field, param=search_params, limit=1, check_task=CheckTasks.check_search_results, check_items={"nq": 1, "limit": 1}, ) for hits in res: ids = hits.ids results, _ = self.collection_wrap.query(expr=f"{df.pk_field} in {ids}") assert len(results) == len(ids) @pytest.mark.tags(CaseLabel.L1) @pytest.mark.parametrize("insert_before_bulk_insert", [True, False]) def test_insert_before_or_after_bulk_insert(self, insert_before_bulk_insert): """ collection schema: [pk, float_vector] Steps: 1. create collection 2. create index and insert data or not 3. import data 4. insert data or not 5. verify the data entities 6. verify search and query """ bulk_insert_row = 500 direct_insert_row = 3000 dim = 128 self._connect() c_name = cf.gen_unique_str("bulk_insert") fields = [ cf.gen_int64_field(name=df.pk_field, is_primary=True), cf.gen_float_field(name=df.float_field), cf.gen_float_vec_field(name=df.float_vec_field, dim=dim), ] data = [ [i for i in range(direct_insert_row)], [np.float32(i) for i in range(direct_insert_row)], cf.gen_vectors(direct_insert_row, dim=dim), ] schema = cf.gen_collection_schema(fields=fields) self.collection_wrap.init_collection(c_name, schema=schema) # build index index_params = ct.default_index self.collection_wrap.create_index(field_name=df.float_vec_field, index_params=index_params) # load collection self.collection_wrap.load() if insert_before_bulk_insert: # insert data self.collection_wrap.insert(data) self.collection_wrap.num_entities files = prepare_bulk_insert_new_json_files( minio_endpoint=self.minio_endpoint, bucket_name=self.bucket_name, is_row_based=True, rows=bulk_insert_row, dim=dim, data_fields=[df.pk_field, df.float_field, df.float_vec_field], force=True, schema=schema, ) # import data t0 = time.time() task_id, _ = self.utility_wrap.do_bulk_insert(collection_name=c_name, files=files) logging.info(f"bulk insert task ids:{task_id}") success, states = self.utility_wrap.wait_for_bulk_insert_tasks_completed(task_ids=[task_id], timeout=300) tt = time.time() - t0 log.info(f"bulk insert state:{success} in {tt}") assert success if not insert_before_bulk_insert: # insert data self.collection_wrap.insert(data) self.collection_wrap.num_entities num_entities = self.collection_wrap.num_entities log.info(f"collection entities: {num_entities}") assert num_entities == bulk_insert_row + direct_insert_row # verify index time.sleep(2) self.utility_wrap.wait_for_index_building_complete(c_name, timeout=300) res, _ = self.utility_wrap.index_building_progress(c_name) log.info(f"index building progress: {res}") # verify search and query log.info("wait for load finished and be ready for search") self.collection_wrap.load(_refresh=True) time.sleep(2) nq = 3 topk = 10 search_data = cf.gen_vectors(nq, dim=dim) search_params = ct.default_search_params res, _ = self.collection_wrap.search( search_data, df.float_vec_field, param=search_params, limit=topk, check_task=CheckTasks.check_search_results, check_items={"nq": nq, "limit": topk}, ) for hits in res: ids = hits.ids expr = f"{df.pk_field} in {ids}" expr = expr.replace("'", '"') results, _ = self.collection_wrap.query(expr=expr) assert len(results) == len(ids) @pytest.mark.tags(CaseLabel.L1) @pytest.mark.parametrize("create_index_before_bulk_insert", [True, False]) @pytest.mark.parametrize("loaded_before_bulk_insert", [True, False]) def test_load_before_or_after_bulk_insert(self, loaded_before_bulk_insert, create_index_before_bulk_insert): """ collection schema: [pk, float_vector] Steps: 1. create collection 2. create index and load collection or not 3. import data 4. load collection or not 5. verify the data entities 5. verify the index status 6. verify search and query """ if loaded_before_bulk_insert and not create_index_before_bulk_insert: pytest.skip("can not load collection if index not created") files = prepare_bulk_insert_json_files( minio_endpoint=self.minio_endpoint, bucket_name=self.bucket_name, is_row_based=True, rows=500, dim=16, auto_id=True, data_fields=[df.vec_field], force=True, ) self._connect() c_name = cf.gen_unique_str("bulk_insert") fields = [ cf.gen_int64_field(name=df.pk_field, is_primary=True), cf.gen_float_vec_field(name=df.vec_field, dim=16), ] schema = cf.gen_collection_schema(fields=fields, auto_id=True) self.collection_wrap.init_collection(c_name, schema=schema) # build index index_params = ct.default_index self.collection_wrap.create_index(field_name=df.vec_field, index_params=index_params) if loaded_before_bulk_insert: # load collection self.collection_wrap.load() # import data t0 = time.time() task_id, _ = self.utility_wrap.do_bulk_insert(collection_name=c_name, files=files) logging.info(f"bulk insert task ids:{task_id}") success, states = self.utility_wrap.wait_for_bulk_insert_tasks_completed(task_ids=[task_id], timeout=300) tt = time.time() - t0 log.info(f"bulk insert state:{success} in {tt}") assert success if not loaded_before_bulk_insert: # load collection self.collection_wrap.load() num_entities = self.collection_wrap.num_entities log.info(f"collection entities: {num_entities}") assert num_entities == 500 time.sleep(2) self.utility_wrap.wait_for_index_building_complete(c_name, timeout=300) res, _ = self.utility_wrap.index_building_progress(c_name) log.info(f"index building progress: {res}") # verify search and query log.info("wait for load finished and be ready for search") self.collection_wrap.load(_refresh=True) time.sleep(2) nq = 3 topk = 10 search_data = cf.gen_vectors(nq, 16) search_params = ct.default_search_params res, _ = self.collection_wrap.search( search_data, df.vec_field, param=search_params, limit=topk, check_task=CheckTasks.check_search_results, check_items={"nq": nq, "limit": topk}, ) for hits in res: ids = hits.ids expr = f"{df.pk_field} in {ids}" expr = expr.replace("'", '"') results, _ = self.collection_wrap.query(expr=expr) assert len(results) == len(ids) @pytest.mark.tags(CaseLabel.L1) def test_index_load_before_bulk_insert(self): """ Steps: 1. create collection 2. create index and load collection 3. import data 4. verify """ enable_dynamic_field = True auto_id = True dim = 128 entities = 1000 fields = [ cf.gen_int64_field(name=df.pk_field, is_primary=True, auto_id=auto_id), cf.gen_int64_field(name=df.int_field), cf.gen_float_field(name=df.float_field), cf.gen_string_field(name=df.string_field), cf.gen_json_field(name=df.json_field), cf.gen_array_field(name=df.array_int_field, element_type=DataType.INT64), cf.gen_float_vec_field(name=df.float_vec_field, dim=dim), ] self._connect() c_name = cf.gen_unique_str("bulk_insert") schema = cf.gen_collection_schema(fields=fields, auto_id=auto_id, enable_dynamic_field=enable_dynamic_field) self.collection_wrap.init_collection(c_name, schema=schema) data_fields = [f.name for f in fields if not f.to_dict().get("auto_id", False)] files = prepare_bulk_insert_new_json_files( minio_endpoint=self.minio_endpoint, bucket_name=self.bucket_name, rows=entities, dim=dim, data_fields=data_fields, enable_dynamic_field=enable_dynamic_field, force=True, schema=schema, ) # create index and load before bulk insert scalar_field_list = [df.int_field, df.float_field, df.double_field, df.string_field] scalar_fields = [f.name for f in fields if f.name in scalar_field_list] float_vec_fields = [f.name for f in fields if "vec" in f.name and "float" in f.name] binary_vec_fields = [f.name for f in fields if "vec" in f.name and "binary" in f.name] for f in scalar_fields: self.collection_wrap.create_index(field_name=f, index_params={"index_type": "INVERTED"}) for f in float_vec_fields: self.collection_wrap.create_index(field_name=f, index_params=ct.default_index) for f in binary_vec_fields: self.collection_wrap.create_index(field_name=f, index_params=ct.default_binary_index) # add json path index for json field json_path_index_params_double = { "index_type": "INVERTED", "params": {"json_cast_type": "double", "json_path": f"{df.json_field}['number']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_double) json_path_index_params_varchar = { "index_type": "INVERTED", "params": {"json_cast_type": "VARCHAR", "json_path": f"{df.json_field}['address']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_varchar) json_path_index_params_bool = { "index_type": "INVERTED", "params": {"json_cast_type": "Bool", "json_path": f"{df.json_field}['name']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_bool) json_path_index_params_not_exist = { "index_type": "INVERTED", "params": {"json_cast_type": "Double", "json_path": f"{df.json_field}['not_exist']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_not_exist) self.collection_wrap.load() t0 = time.time() task_id, _ = self.utility_wrap.do_bulk_insert(collection_name=c_name, files=files) logging.info(f"bulk insert task ids:{task_id}") success, states = self.utility_wrap.wait_for_bulk_insert_tasks_completed(task_ids=[task_id], timeout=300) tt = time.time() - t0 log.info(f"bulk insert state:{success} in {tt} with states:{states}") assert success num_entities = self.collection_wrap.num_entities log.info(f" collection entities: {num_entities}") assert num_entities == entities # verify imported data is available for search log.info("wait for load finished and be ready for search") self.collection_wrap.load(_refresh=True) time.sleep(5) # log.info(f"query seg info: {self.utility_wrap.get_query_segment_info(c_name)[0]}") # query data for f in scalar_fields: if f == df.string_field: expr = f"{f} > '0'" else: expr = f"{f} > 0" res, result = self.collection_wrap.query(expr=expr, output_fields=["count(*)"]) log.info(f"query result: {res}") assert result # search data search_data = cf.gen_vectors(1, dim) search_params = ct.default_search_params for field_name in float_vec_fields: res, _ = self.collection_wrap.search( search_data, field_name, param=search_params, limit=1, output_fields=["*"], check_task=CheckTasks.check_search_results, check_items={"nq": 1, "limit": 1}, ) for hit in res: for r in hit: fields_from_search = r.fields.keys() for f in fields: assert f.name in fields_from_search if enable_dynamic_field: assert "name" in fields_from_search assert "address" in fields_from_search _, search_data = cf.gen_binary_vectors(1, dim) search_params = ct.default_search_binary_params for field_name in binary_vec_fields: res, _ = self.collection_wrap.search( search_data, field_name, param=search_params, limit=1, output_fields=["*"], check_task=CheckTasks.check_search_results, check_items={"nq": 1, "limit": 1}, ) for hit in res: for r in hit: fields_from_search = r.fields.keys() for f in fields: assert f.name in fields_from_search if enable_dynamic_field: assert "name" in fields_from_search assert "address" in fields_from_search res, _ = self.collection_wrap.query(expr=f"{df.json_field}['number'] >= 0", output_fields=[df.json_field]) assert len(res) == entities res, _ = self.collection_wrap.query(expr=f"{df.json_field}['number'] == 1", output_fields=[df.json_field]) assert len(res) == 1 @pytest.mark.tags(CaseLabel.L1) @pytest.mark.parametrize("auto_id", [True]) @pytest.mark.parametrize("dim", [128]) # 128 @pytest.mark.parametrize("entities", [2000]) @pytest.mark.parametrize("enable_dynamic_field", [True]) @pytest.mark.parametrize("enable_partition_key", [True, False]) @pytest.mark.parametrize("nullable", [True, False]) @pytest.mark.parametrize("add_field", [True, False]) def test_bulk_insert_all_field_with_new_json_format( self, auto_id, dim, entities, enable_dynamic_field, enable_partition_key, nullable, add_field ): """ collection schema 1: [pk, int64, float64, string float_vector] data file: vectors.npy and uid.npy, Steps: 1. create collection 2. import data 3. verify """ if enable_partition_key is True and nullable is True: pytest.skip("partition key field not support nullable") float_vec_field_dim = dim binary_vec_field_dim = ((dim + random.randint(-16, 32)) // 8) * 8 bf16_vec_field_dim = dim + random.randint(-16, 32) fp16_vec_field_dim = dim + random.randint(-16, 32) fields = [ cf.gen_int64_field(name=df.pk_field, is_primary=True, auto_id=auto_id), cf.gen_int64_field(name=df.int_field, nullable=nullable), cf.gen_float_field(name=df.float_field, nullable=nullable), cf.gen_string_field(name=df.string_field, is_partition_key=enable_partition_key, nullable=nullable), cf.gen_string_field(name=df.text_field, enable_analyzer=True, enable_match=True, nullable=nullable), cf.gen_json_field(name=df.json_field, nullable=nullable), cf.gen_array_field(name=df.array_int_field, element_type=DataType.INT64, nullable=nullable), cf.gen_array_field(name=df.array_float_field, element_type=DataType.FLOAT, nullable=nullable), cf.gen_array_field( name=df.array_string_field, element_type=DataType.VARCHAR, max_length=100, nullable=nullable ), cf.gen_array_field(name=df.array_bool_field, element_type=DataType.BOOL, nullable=nullable), cf.gen_geometry_field(name=df.geo_field), cf.gen_timestamptz_field(name=df.timestamp_field, nullable=nullable), cf.gen_float_vec_field(name=df.float_vec_field, dim=float_vec_field_dim), cf.gen_binary_vec_field(name=df.binary_vec_field, dim=binary_vec_field_dim), cf.gen_bfloat16_vec_field(name=df.bf16_vec_field, dim=bf16_vec_field_dim), cf.gen_float16_vec_field(name=df.fp16_vec_field, dim=fp16_vec_field_dim), ] data_fields = [f.name for f in fields if not f.to_dict().get("auto_id", False)] self._connect() c_name = cf.gen_unique_str("bulk_insert") schema = cf.gen_collection_schema(fields=fields, auto_id=auto_id, enable_dynamic_field=enable_dynamic_field) files = prepare_bulk_insert_new_json_files( minio_endpoint=self.minio_endpoint, bucket_name=self.bucket_name, rows=entities, dim=dim, data_fields=data_fields, enable_dynamic_field=enable_dynamic_field, force=True, schema=schema, ) self.collection_wrap.init_collection(c_name, schema=schema) if add_field: self._connect(enable_milvus_client_api=True) self.client.add_collection_field( collection_name=c_name, field_name=df.new_field, data_type=DataType.INT64, nullable=True ) # import data t0 = time.time() task_id, _ = self.utility_wrap.do_bulk_insert(collection_name=c_name, files=files) logging.info(f"bulk insert task ids:{task_id}") success, states = self.utility_wrap.wait_for_bulk_insert_tasks_completed(task_ids=[task_id], timeout=300) tt = time.time() - t0 log.info(f"bulk insert state:{success} in {tt} with states:{states}") assert success num_entities = self.collection_wrap.num_entities log.info(f" collection entities: {num_entities}") assert num_entities == entities # verify imported data is available for search index_params = ct.default_index float_vec_fields = [f.name for f in fields if "vec" in f.name and "float" in f.name] binary_vec_fields = [f.name for f in fields if "vec" in f.name and "binary" in f.name] for f in float_vec_fields: self.collection_wrap.create_index(field_name=f, index_params=index_params) for f in binary_vec_fields: self.collection_wrap.create_index(field_name=f, index_params=ct.default_binary_index) # add json path index for json field json_path_index_params_double = { "index_type": "INVERTED", "params": {"json_cast_type": "double", "json_path": f"{df.json_field}['number']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_double) json_path_index_params_varchar = { "index_type": "INVERTED", "params": {"json_cast_type": "VARCHAR", "json_path": f"{df.json_field}['address']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_varchar) json_path_index_params_bool = { "index_type": "INVERTED", "params": {"json_cast_type": "Bool", "json_path": f"{df.json_field}['name']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_bool) json_path_index_params_not_exist = { "index_type": "INVERTED", "params": {"json_cast_type": "Double", "json_path": f"{df.json_field}['not_exist']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_not_exist) self.collection_wrap.load() log.info("wait for load finished and be ready for search") time.sleep(2) # log.info(f"query seg info: {self.utility_wrap.get_query_segment_info(c_name)[0]}") for f in [df.float_vec_field, df.bf16_vec_field, df.fp16_vec_field]: vector_data_type = DataType.FLOAT_VECTOR if f == df.float_vec_field: dim = float_vec_field_dim vector_data_type = DataType.FLOAT_VECTOR elif f == df.bf16_vec_field: dim = bf16_vec_field_dim vector_data_type = DataType.BFLOAT16_VECTOR else: dim = fp16_vec_field_dim vector_data_type = DataType.FLOAT16_VECTOR search_data = cf.gen_vectors(1, dim, vector_data_type=vector_data_type) search_params = ct.default_search_params res, _ = self.collection_wrap.search( search_data, f, param=search_params, limit=1, output_fields=["*"], check_task=CheckTasks.check_search_results, check_items={"nq": 1, "limit": 1}, ) for hit in res: for r in hit: fields_from_search = r.fields.keys() for f in fields: assert f.name in fields_from_search if enable_dynamic_field: assert "name" in fields_from_search assert "address" in fields_from_search _, search_data = cf.gen_binary_vectors(1, binary_vec_field_dim) search_params = ct.default_search_binary_params for field_name in binary_vec_fields: res, _ = self.collection_wrap.search( search_data, field_name, param=search_params, limit=1, output_fields=["*"], check_task=CheckTasks.check_search_results, check_items={"nq": 1, "limit": 1}, ) for hit in res: for r in hit: fields_from_search = r.fields.keys() for f in fields: assert f.name in fields_from_search if enable_dynamic_field: assert "name" in fields_from_search assert "address" in fields_from_search # query data if not nullable: expr_field = df.string_field expr = f"{expr_field} >= '0'" else: res, _ = self.collection_wrap.query( expr=f"{df.string_field} >= '0'", output_fields=[df.string_field, df.int_field] ) assert len(res) == 0 expr_field = df.pk_field expr = f"{expr_field} >= 0" if add_field: res, _ = self.collection_wrap.query( expr=f"{df.new_field} is not null", output_fields=[df.string_field, df.int_field, df.new_field] ) assert len(res) == 0 res, _ = self.collection_wrap.query(expr=f"{expr}", output_fields=[expr_field, df.int_field]) assert len(res) == entities log.info(res) query_data = [r[expr_field] for r in res][: len(self.collection_wrap.partitions)] res, _ = self.collection_wrap.query(expr=f"{expr_field} in {query_data}", output_fields=[expr_field]) assert len(res) == len(query_data) res, _ = self.collection_wrap.query( expr=f"text_match({df.text_field}, 'milvus')", output_fields=[df.text_field] ) if nullable is False: assert len(res) == entities else: assert 0 < len(res) < entities if enable_partition_key: assert len(self.collection_wrap.partitions) > 1 res, _ = self.collection_wrap.query(expr=f"{df.json_field}['number'] >= 0", output_fields=[df.json_field]) if not nullable: assert len(res) == entities else: assert len(res) == 0 res, _ = self.collection_wrap.query(expr=f"{df.json_field}['number'] == 1", output_fields=[df.json_field]) if not nullable: assert len(res) == 1 else: assert len(res) == 0 @pytest.mark.tags(CaseLabel.L1) @pytest.mark.parametrize("auto_id", [True, False]) @pytest.mark.parametrize("dim", [128]) # 128 @pytest.mark.parametrize("entities", [2000]) @pytest.mark.parametrize("enable_dynamic_field", [True, False]) @pytest.mark.parametrize("enable_partition_key", [True, False]) @pytest.mark.parametrize("include_meta", [True, False]) @pytest.mark.parametrize("nullable", [True, False]) @pytest.mark.parametrize("add_field", [True, False]) def test_bulk_insert_all_field_with_numpy( self, auto_id, dim, entities, enable_dynamic_field, enable_partition_key, include_meta, nullable, add_field ): """ collection schema 1: [pk, int64, float64, string float_vector] data file: vectors.npy and uid.npy, note: numpy file is not supported for array field Steps: 1. create collection 2. import data 3. verify """ if enable_dynamic_field is False and include_meta is True: pytest.skip("include_meta only works with enable_dynamic_field") if nullable is True: pytest.skip("not support bulk insert numpy files in field which set nullable == true") float_vec_field_dim = dim binary_vec_field_dim = ((dim + random.randint(-16, 32)) // 8) * 8 bf16_vec_field_dim = dim + random.randint(-16, 32) fp16_vec_field_dim = dim + random.randint(-16, 32) fields = [ cf.gen_int64_field(name=df.pk_field, is_primary=True, auto_id=auto_id), cf.gen_int64_field(name=df.int_field, nullable=nullable), cf.gen_float_field(name=df.float_field), cf.gen_string_field(name=df.string_field, is_partition_key=enable_partition_key), cf.gen_string_field(name=df.text_field, enable_analyzer=True, enable_match=True, nullable=nullable), cf.gen_json_field(name=df.json_field), cf.gen_geometry_field(name=df.geo_field), cf.gen_float_vec_field(name=df.float_vec_field, dim=float_vec_field_dim), cf.gen_binary_vec_field(name=df.binary_vec_field, dim=binary_vec_field_dim), cf.gen_bfloat16_vec_field(name=df.bf16_vec_field, dim=bf16_vec_field_dim), cf.gen_float16_vec_field(name=df.fp16_vec_field, dim=fp16_vec_field_dim), ] data_fields = [f.name for f in fields if not f.to_dict().get("auto_id", False)] self._connect() c_name = cf.gen_unique_str("bulk_insert") schema = cf.gen_collection_schema(fields=fields, auto_id=auto_id, enable_dynamic_field=enable_dynamic_field) files = prepare_bulk_insert_numpy_files( minio_endpoint=self.minio_endpoint, bucket_name=self.bucket_name, rows=entities, dim=dim, data_fields=data_fields, enable_dynamic_field=enable_dynamic_field, force=True, schema=schema, ) self.collection_wrap.init_collection(c_name, schema=schema) if add_field: self._connect(enable_milvus_client_api=True) self.client.add_collection_field( collection_name=c_name, field_name=df.new_field, data_type=DataType.INT64, nullable=True ) # import data t0 = time.time() task_id, _ = self.utility_wrap.do_bulk_insert(collection_name=c_name, files=files) logging.info(f"bulk insert task ids:{task_id}") success, states = self.utility_wrap.wait_for_bulk_insert_tasks_completed(task_ids=[task_id], timeout=300) tt = time.time() - t0 log.info(f"bulk insert state:{success} in {tt} with states:{states}") assert success num_entities = self.collection_wrap.num_entities log.info(f" collection entities: {num_entities}") assert num_entities == entities # verify imported data is available for search index_params = ct.default_index float_vec_fields = [f.name for f in fields if "vec" in f.name and "float" in f.name] binary_vec_fields = [f.name for f in fields if "vec" in f.name and "binary" in f.name] for f in float_vec_fields: self.collection_wrap.create_index(field_name=f, index_params=index_params) # add json path index for json field json_path_index_params_double = { "index_type": "INVERTED", "params": {"json_cast_type": "double", "json_path": f"{df.json_field}['number']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_double) json_path_index_params_varchar = { "index_type": "INVERTED", "params": {"json_cast_type": "VARCHAR", "json_path": f"{df.json_field}['address']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_varchar) json_path_index_params_bool = { "index_type": "INVERTED", "params": {"json_cast_type": "Bool", "json_path": f"{df.json_field}['name']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_bool) json_path_index_params_not_exist = { "index_type": "INVERTED", "params": {"json_cast_type": "Double", "json_path": f"{df.json_field}['not_exist']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_not_exist) for f in binary_vec_fields: self.collection_wrap.create_index(field_name=f, index_params=ct.default_binary_index) self.collection_wrap.load() log.info("wait for load finished and be ready for search") time.sleep(2) # log.info(f"query seg info: {self.utility_wrap.get_query_segment_info(c_name)[0]}") for f in [df.float_vec_field, df.bf16_vec_field, df.fp16_vec_field]: vector_data_type = DataType.FLOAT_VECTOR if f == df.float_vec_field: dim = float_vec_field_dim vector_data_type = DataType.FLOAT_VECTOR elif f == df.bf16_vec_field: dim = bf16_vec_field_dim vector_data_type = DataType.BFLOAT16_VECTOR else: dim = fp16_vec_field_dim vector_data_type = DataType.FLOAT16_VECTOR search_data = cf.gen_vectors(1, dim, vector_data_type=vector_data_type) search_params = ct.default_search_params res, _ = self.collection_wrap.search( search_data, f, param=search_params, limit=1, output_fields=["*"], check_task=CheckTasks.check_search_results, check_items={"nq": 1, "limit": 1}, ) for hit in res: for r in hit: fields_from_search = r.fields.keys() for f in fields: assert f.name in fields_from_search if enable_dynamic_field: assert "name" in fields_from_search assert "address" in fields_from_search _, search_data = cf.gen_binary_vectors(1, binary_vec_field_dim) search_params = ct.default_search_binary_params for field_name in binary_vec_fields: res, _ = self.collection_wrap.search( search_data, field_name, param=search_params, limit=1, output_fields=["*"], check_task=CheckTasks.check_search_results, check_items={"nq": 1, "limit": 1}, ) for hit in res: for r in hit: fields_from_search = r.fields.keys() for f in fields: assert f.name in fields_from_search if enable_dynamic_field: assert "name" in fields_from_search assert "address" in fields_from_search # query data if add_field: res, _ = self.collection_wrap.query( expr=f"{df.new_field} is not null", output_fields=[df.string_field, df.int_field, df.new_field] ) assert len(res) == 0 res, _ = self.collection_wrap.query(expr=f"{df.string_field} >= '0'", output_fields=[df.string_field]) assert len(res) == entities query_data = [r[df.string_field] for r in res][: len(self.collection_wrap.partitions)] res, _ = self.collection_wrap.query(expr=f"{df.string_field} in {query_data}", output_fields=[df.string_field]) assert len(res) == len(query_data) res, _ = self.collection_wrap.query( expr=f"TEXT_MATCH({df.text_field}, 'milvus')", output_fields=[df.text_field] ) if nullable is False: assert len(res) == entities else: assert 0 < len(res) < entities if enable_partition_key: assert len(self.collection_wrap.partitions) > 1 res, _ = self.collection_wrap.query(expr=f"{df.json_field}['number'] >= 0", output_fields=[df.json_field]) assert len(res) == entities res, _ = self.collection_wrap.query(expr=f"{df.json_field}['number'] == 1", output_fields=[df.json_field]) assert len(res) == 1 @pytest.mark.tags(CaseLabel.L1) @pytest.mark.parametrize("auto_id", [True, False]) @pytest.mark.parametrize("dim", [128]) # 128 @pytest.mark.parametrize("entities", [2000]) @pytest.mark.parametrize("enable_dynamic_field", [True, False]) @pytest.mark.parametrize("enable_partition_key", [True, False]) @pytest.mark.parametrize("include_meta", [True, False]) @pytest.mark.parametrize("nullable", [True, False]) @pytest.mark.parametrize("add_field", [True, False]) def test_bulk_insert_all_field_with_parquet( self, auto_id, dim, entities, enable_dynamic_field, enable_partition_key, include_meta, nullable, add_field ): """ collection schema 1: [pk, int64, float64, string float_vector] data file: vectors.parquet and uid.parquet, Steps: 1. create collection 2. import data 3. verify """ if enable_dynamic_field is False and include_meta is True: pytest.skip("include_meta only works with enable_dynamic_field") if enable_partition_key is True and nullable is True: pytest.skip("partition key field not support nullable") float_vec_field_dim = dim binary_vec_field_dim = ((dim + random.randint(-16, 32)) // 8) * 8 bf16_vec_field_dim = dim + random.randint(-16, 32) fp16_vec_field_dim = dim + random.randint(-16, 32) fields = [ cf.gen_int64_field(name=df.pk_field, is_primary=True, auto_id=auto_id), cf.gen_int64_field(name=df.int_field, nullable=nullable), cf.gen_float_field(name=df.float_field, nullable=nullable), cf.gen_string_field(name=df.string_field, is_partition_key=enable_partition_key, nullable=nullable), cf.gen_string_field(name=df.text_field, enable_analyzer=True, enable_match=True, nullable=nullable), cf.gen_json_field(name=df.json_field, nullable=nullable), cf.gen_array_field(name=df.array_int_field, element_type=DataType.INT64, nullable=nullable), cf.gen_array_field(name=df.array_float_field, element_type=DataType.FLOAT, nullable=nullable), cf.gen_array_field( name=df.array_string_field, element_type=DataType.VARCHAR, max_length=100, nullable=nullable ), cf.gen_array_field(name=df.array_bool_field, element_type=DataType.BOOL, nullable=nullable), cf.gen_geometry_field(name=df.geo_field), cf.gen_timestamptz_field(name=df.timestamp_field, nullable=nullable), cf.gen_float_vec_field(name=df.float_vec_field, dim=float_vec_field_dim), cf.gen_binary_vec_field(name=df.binary_vec_field, dim=binary_vec_field_dim), cf.gen_bfloat16_vec_field(name=df.bf16_vec_field, dim=bf16_vec_field_dim), cf.gen_float16_vec_field(name=df.fp16_vec_field, dim=fp16_vec_field_dim), ] data_fields = [f.name for f in fields if not f.to_dict().get("auto_id", False)] self._connect() c_name = cf.gen_unique_str("bulk_insert") schema = cf.gen_collection_schema(fields=fields, auto_id=auto_id, enable_dynamic_field=enable_dynamic_field) files = prepare_bulk_insert_parquet_files( minio_endpoint=self.minio_endpoint, bucket_name=self.bucket_name, rows=entities, dim=dim, data_fields=data_fields, enable_dynamic_field=enable_dynamic_field, force=True, schema=schema, use_utf8_data=True, ) self.collection_wrap.init_collection(c_name, schema=schema) if add_field: self._connect(enable_milvus_client_api=True) self.client.add_collection_field( collection_name=c_name, field_name=df.new_field, data_type=DataType.INT64, nullable=True ) # import data t0 = time.time() task_id, _ = self.utility_wrap.do_bulk_insert(collection_name=c_name, files=files) logging.info(f"bulk insert task ids:{task_id}") success, states = self.utility_wrap.wait_for_bulk_insert_tasks_completed(task_ids=[task_id], timeout=300) tt = time.time() - t0 log.info(f"bulk insert state:{success} in {tt} with states:{states}") assert success num_entities = self.collection_wrap.num_entities log.info(f" collection entities: {num_entities}") assert num_entities == entities # verify imported data is available for search index_params = ct.default_index float_vec_fields = [f.name for f in fields if "vec" in f.name and "float" in f.name] binary_vec_fields = [f.name for f in fields if "vec" in f.name and "binary" in f.name] for f in float_vec_fields: self.collection_wrap.create_index(field_name=f, index_params=index_params) for f in binary_vec_fields: self.collection_wrap.create_index(field_name=f, index_params=ct.default_binary_index) # add json path index for json field json_path_index_params_double = { "index_type": "INVERTED", "params": {"json_cast_type": "double", "json_path": f"{df.json_field}['number']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_double) json_path_index_params_varchar = { "index_type": "INVERTED", "params": {"json_cast_type": "VARCHAR", "json_path": f"{df.json_field}['address']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_varchar) json_path_index_params_bool = { "index_type": "INVERTED", "params": {"json_cast_type": "Bool", "json_path": f"{df.json_field}['name']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_bool) json_path_index_params_not_exist = { "index_type": "INVERTED", "params": {"json_cast_type": "Double", "json_path": f"{df.json_field}['not_exist']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_not_exist) self.collection_wrap.load() log.info("wait for load finished and be ready for search") time.sleep(2) # log.info(f"query seg info: {self.utility_wrap.get_query_segment_info(c_name)[0]}") for f in [df.float_vec_field, df.bf16_vec_field, df.fp16_vec_field]: vector_data_type = DataType.FLOAT_VECTOR if f == df.float_vec_field: dim = float_vec_field_dim vector_data_type = DataType.FLOAT_VECTOR elif f == df.bf16_vec_field: dim = bf16_vec_field_dim vector_data_type = DataType.BFLOAT16_VECTOR else: dim = fp16_vec_field_dim vector_data_type = DataType.FLOAT16_VECTOR search_data = cf.gen_vectors(1, dim, vector_data_type=vector_data_type) search_params = ct.default_search_params res, _ = self.collection_wrap.search( search_data, f, param=search_params, limit=1, output_fields=["*"], check_task=CheckTasks.check_search_results, check_items={"nq": 1, "limit": 1}, ) for hit in res: for r in hit: fields_from_search = r.fields.keys() for f in fields: assert f.name in fields_from_search if enable_dynamic_field: assert "name" in fields_from_search assert "address" in fields_from_search _, search_data = cf.gen_binary_vectors(1, binary_vec_field_dim) search_params = ct.default_search_binary_params for field_name in binary_vec_fields: res, _ = self.collection_wrap.search( search_data, field_name, param=search_params, limit=1, output_fields=["*"], check_task=CheckTasks.check_search_results, check_items={"nq": 1, "limit": 1}, ) for hit in res: for r in hit: fields_from_search = r.fields.keys() for f in fields: assert f.name in fields_from_search if enable_dynamic_field: assert "name" in fields_from_search assert "address" in fields_from_search # query data if not nullable: expr_field = df.string_field expr = f"{expr_field} >= '0'" else: res, _ = self.collection_wrap.query(expr=f"{df.string_field} >= '0'", output_fields=[df.string_field]) assert len(res) == 0 expr_field = df.pk_field expr = f"{expr_field} >= 0" if add_field: res, _ = self.collection_wrap.query( expr=f"{df.new_field} is not null", output_fields=[df.string_field, df.int_field, df.new_field] ) assert len(res) == 0 res, _ = self.collection_wrap.query( expr=f"{expr}", output_fields=[df.pk_field, df.string_field, df.text_field, df.array_string_field, df.json_field], ) assert len(res) == entities if not nullable: sample = res[0] json_value = sample[df.json_field] if isinstance(json_value, str): json_value = json.loads(json_value) assert any(ord(ch) > 127 for ch in sample[df.string_field]) assert any(ord(ch) > 127 for ch in sample[df.text_field]) assert any(any(ord(ch) > 127 for ch in item) for item in sample[df.array_string_field]) assert any(ord(ch) > 127 for ch in json.dumps(json_value, ensure_ascii=False)) query_data = [r[expr_field] for r in res][: len(self.collection_wrap.partitions)] res, _ = self.collection_wrap.query(expr=f"{expr_field} in {query_data}", output_fields=[expr_field]) assert len(res) == len(query_data) res, _ = self.collection_wrap.query( expr=f"TEXT_MATCH({df.text_field}, 'milvus')", output_fields=[df.text_field] ) if not nullable: assert len(res) == entities else: assert 0 < len(res) < entities if enable_partition_key: assert len(self.collection_wrap.partitions) > 1 res, _ = self.collection_wrap.query(expr=f"{df.json_field}['number'] >= 0", output_fields=[df.json_field]) if not nullable: assert len(res) == entities else: assert len(res) == 0 res, _ = self.collection_wrap.query(expr=f"{df.json_field}['number'] == 1", output_fields=[df.json_field]) if not nullable: assert len(res) == 1 else: assert len(res) == 0 @pytest.mark.tags(CaseLabel.L1) @pytest.mark.parametrize("auto_id", [True, False]) @pytest.mark.parametrize("dim", [128]) # 128 @pytest.mark.parametrize("entities", [2000]) @pytest.mark.parametrize("enable_dynamic_field", [True, False]) @pytest.mark.parametrize("include_meta", [True, False]) @pytest.mark.parametrize("sparse_format", ["doc", "coo"]) def test_bulk_insert_sparse_vector_with_parquet( self, auto_id, dim, entities, enable_dynamic_field, include_meta, sparse_format ): """ collection schema 1: [pk, int64, float64, string float_vector] data file: vectors.parquet and uid.parquet, Steps: 1. create collection 2. import data 3. verify """ if enable_dynamic_field is False and include_meta is True: pytest.skip("include_meta only works with enable_dynamic_field") fields = [ cf.gen_int64_field(name=df.pk_field, is_primary=True, auto_id=auto_id), cf.gen_int64_field(name=df.int_field), cf.gen_float_field(name=df.float_field), cf.gen_string_field(name=df.string_field), cf.gen_json_field(name=df.json_field), cf.gen_array_field(name=df.array_int_field, element_type=DataType.INT64), cf.gen_array_field(name=df.array_float_field, element_type=DataType.FLOAT), cf.gen_array_field(name=df.array_string_field, element_type=DataType.VARCHAR, max_length=100), cf.gen_array_field(name=df.array_bool_field, element_type=DataType.BOOL), cf.gen_float_vec_field(name=df.float_vec_field, dim=dim), cf.gen_sparse_vec_field(name=df.sparse_vec_field), ] data_fields = [f.name for f in fields if not f.to_dict().get("auto_id", False)] self._connect() c_name = cf.gen_unique_str("bulk_insert") schema = cf.gen_collection_schema(fields=fields, auto_id=auto_id, enable_dynamic_field=enable_dynamic_field) files = prepare_bulk_insert_parquet_files( minio_endpoint=self.minio_endpoint, bucket_name=self.bucket_name, rows=entities, dim=dim, data_fields=data_fields, enable_dynamic_field=enable_dynamic_field, force=True, include_meta=include_meta, sparse_format=sparse_format, schema=schema, ) self.collection_wrap.init_collection(c_name, schema=schema) # import data t0 = time.time() task_id, _ = self.utility_wrap.do_bulk_insert(collection_name=c_name, files=files) logging.info(f"bulk insert task ids:{task_id}") success, states = self.utility_wrap.wait_for_bulk_insert_tasks_completed(task_ids=[task_id], timeout=300) tt = time.time() - t0 log.info(f"bulk insert state:{success} in {tt} with states:{states}") assert success num_entities = self.collection_wrap.num_entities log.info(f" collection entities: {num_entities}") assert num_entities == entities # verify imported data is available for search index_params = ct.default_index float_vec_fields = [f.name for f in fields if "vec" in f.name and "float" in f.name] sparse_vec_fields = [f.name for f in fields if "vec" in f.name and "sparse" in f.name] for f in float_vec_fields: self.collection_wrap.create_index(field_name=f, index_params=index_params) for f in sparse_vec_fields: self.collection_wrap.create_index(field_name=f, index_params=ct.default_sparse_inverted_index) self.collection_wrap.load() log.info("wait for load finished and be ready for search") time.sleep(2) # log.info(f"query seg info: {self.utility_wrap.get_query_segment_info(c_name)[0]}") search_data = cf.gen_vectors(1, dim) search_params = ct.default_search_params for field_name in float_vec_fields: res, _ = self.collection_wrap.search( search_data, field_name, param=search_params, limit=1, output_fields=["*"], check_task=CheckTasks.check_search_results, check_items={"nq": 1, "limit": 1}, ) for hit in res: for r in hit: fields_from_search = r.fields.keys() for f in fields: assert f.name in fields_from_search if enable_dynamic_field and include_meta: assert "name" in fields_from_search assert "address" in fields_from_search search_data = cf.gen_sparse_vectors(1, dim) search_params = ct.default_sparse_search_params for field_name in sparse_vec_fields: res, _ = self.collection_wrap.search( search_data, field_name, param=search_params, limit=1, output_fields=["*"], check_task=CheckTasks.check_search_results, check_items={"nq": 1, "limit": 1}, ) for hit in res: for r in hit: fields_from_search = r.fields.keys() for f in fields: assert f.name in fields_from_search if enable_dynamic_field and include_meta: assert "name" in fields_from_search assert "address" in fields_from_search @pytest.mark.tags(CaseLabel.L1) @pytest.mark.parametrize("auto_id", [True, False]) @pytest.mark.parametrize("dim", [128]) # 128 @pytest.mark.parametrize("entities", [2000]) @pytest.mark.parametrize("enable_dynamic_field", [True, False]) @pytest.mark.parametrize("include_meta", [True, False]) @pytest.mark.parametrize("sparse_format", ["doc", "coo"]) def test_bulk_insert_sparse_vector_with_json( self, auto_id, dim, entities, enable_dynamic_field, include_meta, sparse_format ): """ collection schema 1: [pk, int64, float64, string float_vector] data file: vectors.parquet and uid.parquet, Steps: 1. create collection 2. import data 3. verify """ if enable_dynamic_field is False and include_meta is True: pytest.skip("include_meta only works with enable_dynamic_field") fields = [ cf.gen_int64_field(name=df.pk_field, is_primary=True, auto_id=auto_id), cf.gen_int64_field(name=df.int_field), cf.gen_float_field(name=df.float_field), cf.gen_string_field(name=df.string_field), cf.gen_json_field(name=df.json_field), cf.gen_array_field(name=df.array_int_field, element_type=DataType.INT64), cf.gen_array_field(name=df.array_float_field, element_type=DataType.FLOAT), cf.gen_array_field(name=df.array_string_field, element_type=DataType.VARCHAR, max_length=100), cf.gen_array_field(name=df.array_bool_field, element_type=DataType.BOOL), cf.gen_float_vec_field(name=df.float_vec_field, dim=dim), cf.gen_sparse_vec_field(name=df.sparse_vec_field), ] data_fields = [f.name for f in fields if not f.to_dict().get("auto_id", False)] self._connect() c_name = cf.gen_unique_str("bulk_insert") schema = cf.gen_collection_schema(fields=fields, auto_id=auto_id, enable_dynamic_field=enable_dynamic_field) files = prepare_bulk_insert_new_json_files( minio_endpoint=self.minio_endpoint, bucket_name=self.bucket_name, rows=entities, dim=dim, data_fields=data_fields, enable_dynamic_field=enable_dynamic_field, force=True, include_meta=include_meta, sparse_format=sparse_format, schema=schema, ) self.collection_wrap.init_collection(c_name, schema=schema) # import data t0 = time.time() task_id, _ = self.utility_wrap.do_bulk_insert(collection_name=c_name, files=files) logging.info(f"bulk insert task ids:{task_id}") success, states = self.utility_wrap.wait_for_bulk_insert_tasks_completed(task_ids=[task_id], timeout=300) tt = time.time() - t0 log.info(f"bulk insert state:{success} in {tt} with states:{states}") assert success num_entities = self.collection_wrap.num_entities log.info(f" collection entities: {num_entities}") assert num_entities == entities # verify imported data is available for search index_params = ct.default_index float_vec_fields = [f.name for f in fields if "vec" in f.name and "float" in f.name] sparse_vec_fields = [f.name for f in fields if "vec" in f.name and "sparse" in f.name] for f in float_vec_fields: self.collection_wrap.create_index(field_name=f, index_params=index_params) for f in sparse_vec_fields: self.collection_wrap.create_index(field_name=f, index_params=ct.default_sparse_inverted_index) self.collection_wrap.load() log.info("wait for load finished and be ready for search") time.sleep(2) # log.info(f"query seg info: {self.utility_wrap.get_query_segment_info(c_name)[0]}") search_data = cf.gen_vectors(1, dim) search_params = ct.default_search_params for field_name in float_vec_fields: res, _ = self.collection_wrap.search( search_data, field_name, param=search_params, limit=1, output_fields=["*"], check_task=CheckTasks.check_search_results, check_items={"nq": 1, "limit": 1}, ) for hit in res: for r in hit: fields_from_search = r.fields.keys() for f in fields: assert f.name in fields_from_search if enable_dynamic_field and include_meta: assert "name" in fields_from_search assert "address" in fields_from_search search_data = cf.gen_sparse_vectors(1, dim) search_params = ct.default_sparse_search_params for field_name in sparse_vec_fields: res, _ = self.collection_wrap.search( search_data, field_name, param=search_params, limit=1, output_fields=["*"], check_task=CheckTasks.check_search_results, check_items={"nq": 1, "limit": 1}, ) for hit in res: for r in hit: fields_from_search = r.fields.keys() for f in fields: assert f.name in fields_from_search if enable_dynamic_field and include_meta: assert "name" in fields_from_search assert "address" in fields_from_search @pytest.mark.tags(CaseLabel.L2) @pytest.mark.parametrize("auto_id", [True, False]) @pytest.mark.parametrize("dim", [128]) # 128 @pytest.mark.parametrize("entities", [1000]) # 1000 @pytest.mark.parametrize("enable_dynamic_field", [True, False]) @pytest.mark.parametrize("sparse_format", ["doc", "coo"]) @pytest.mark.parametrize("nullable", [True, False]) def test_with_all_field_json_with_bulk_writer( self, auto_id, dim, entities, enable_dynamic_field, sparse_format, nullable ): """ collection schema 1: [pk, int64, float64, string float_vector] data file: vectors.npy and uid.npy, Steps: 1. create collection 2. import data 3. verify """ self._connect() fields = [ cf.gen_int64_field(name=df.pk_field, is_primary=True, auto_id=auto_id), cf.gen_int64_field(name=df.int_field, nullable=nullable), cf.gen_float_field(name=df.float_field, nullable=nullable), cf.gen_string_field(name=df.string_field, nullable=nullable), cf.gen_json_field(name=df.json_field, nullable=nullable), cf.gen_timestamptz_field(name=df.timestamp_field, nullable=nullable), cf.gen_array_field(name=df.array_int_field, element_type=DataType.INT64, nullable=nullable), cf.gen_array_field(name=df.array_float_field, element_type=DataType.FLOAT, nullable=nullable), cf.gen_array_field( name=df.array_string_field, element_type=DataType.VARCHAR, max_length=100, nullable=nullable ), cf.gen_array_field(name=df.array_bool_field, element_type=DataType.BOOL, nullable=nullable), cf.gen_float_vec_field(name=df.float_vec_field, dim=dim), cf.gen_float16_vec_field(name=df.fp16_vec_field, dim=dim), cf.gen_bfloat16_vec_field(name=df.bf16_vec_field, dim=dim), cf.gen_sparse_vec_field(name=df.sparse_vec_field), ] c_name = cf.gen_unique_str("bulk_insert") schema = cf.gen_collection_schema(fields=fields, auto_id=auto_id, enable_dynamic_field=enable_dynamic_field) self.collection_wrap.init_collection(c_name, schema=schema) with RemoteBulkWriter( schema=schema, remote_path="bulk_data", connect_param=RemoteBulkWriter.ConnectParam( bucket_name=self.bucket_name, endpoint=self.minio_endpoint, access_key="minioadmin", secret_key="minioadmin", ), file_type=BulkFileType.JSON, ) as remote_writer: json_value = [ # 1, # 1.0, # "1", # [1, 2, 3], # ["1", "2", "3"], # [1, 2, "3"], {"key": "value"}, {"number": 1}, {"name": fake.name()}, {"address": fake.address()}, ] for i in range(entities): row = { df.pk_field: i, df.int_field: 1 if not (nullable and random.random() < 0.5) else None, df.float_field: 1.0 if not (nullable and random.random() < 0.5) else None, df.string_field: "string" if not (nullable and random.random() < 0.5) else None, df.json_field: json_value[i % len(json_value)] if not (nullable and random.random() < 0.5) else None, df.timestamp_field: cf.gen_timestamptz_str() if not (nullable and random.random() < 0.5) else None, df.array_int_field: [1, 2] if not (nullable and random.random() < 0.5) else None, df.array_float_field: [1.0, 2.0] if not (nullable and random.random() < 0.5) else None, df.array_string_field: ["string1", "string2"] if not (nullable and random.random() < 0.5) else None, df.array_bool_field: [True, False] if not (nullable and random.random() < 0.5) else None, df.float_vec_field: cf.gen_vectors(1, dim)[0], df.fp16_vec_field: cf.gen_vectors(1, dim, vector_data_type=DataType.FLOAT_VECTOR)[0], df.bf16_vec_field: cf.gen_vectors(1, dim, vector_data_type=DataType.BFLOAT16_VECTOR)[0], df.sparse_vec_field: cf.gen_sparse_vectors(1, dim, sparse_format=sparse_format)[0], } if auto_id: row.pop(df.pk_field) if enable_dynamic_field: row["name"] = fake.name() row["address"] = fake.address() remote_writer.append_row(row) remote_writer.commit() files = remote_writer.batch_files # import data for f in files: t0 = time.time() task_id, _ = self.utility_wrap.do_bulk_insert(collection_name=c_name, files=f) logging.info(f"bulk insert task ids:{task_id}") success, states = self.utility_wrap.wait_for_bulk_insert_tasks_completed(task_ids=[task_id], timeout=300) tt = time.time() - t0 log.info(f"bulk insert state:{success} in {tt} with states:{states}") assert success num_entities = self.collection_wrap.num_entities log.info(f" collection entities: {num_entities}") assert num_entities == entities # verify imported data is available for search index_params = ct.default_index float_vec_fields = [f.name for f in fields if "vec" in f.name and "float" in f.name] sparse_vec_fields = [f.name for f in fields if "vec" in f.name and "sparse" in f.name] for f in float_vec_fields: self.collection_wrap.create_index(field_name=f, index_params=index_params) for f in sparse_vec_fields: self.collection_wrap.create_index(field_name=f, index_params=ct.default_sparse_inverted_index) # add json path index for json field json_path_index_params_double = { "index_type": "INVERTED", "params": {"json_cast_type": "double", "json_path": f"{df.json_field}['number']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_double) json_path_index_params_varchar = { "index_type": "INVERTED", "params": {"json_cast_type": "VARCHAR", "json_path": f"{df.json_field}['address']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_varchar) json_path_index_params_bool = { "index_type": "INVERTED", "params": {"json_cast_type": "Bool", "json_path": f"{df.json_field}['name']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_bool) json_path_index_params_not_exist = { "index_type": "INVERTED", "params": {"json_cast_type": "Double", "json_path": f"{df.json_field}['not_exist']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_not_exist) self.collection_wrap.load() log.info("wait for load finished and be ready for search") time.sleep(2) # log.info(f"query seg info: {self.utility_wrap.get_query_segment_info(c_name)[0]}") search_data = cf.gen_vectors(1, dim) search_params = ct.default_search_params res, _ = self.collection_wrap.search( search_data, df.float_vec_field, param=search_params, limit=1, output_fields=["*"], check_task=CheckTasks.check_search_results, check_items={"nq": 1, "limit": 1}, ) for hit in res: for r in hit: fields_from_search = r.fields.keys() for f in fields: assert f.name in fields_from_search if enable_dynamic_field: assert "name" in fields_from_search assert "address" in fields_from_search res, _ = self.collection_wrap.query(expr=f"{df.json_field}['number'] == 1", output_fields=[df.json_field]) if not nullable: assert len(res) == int(entities / len(json_value)) else: assert 0 < len(res) < int(entities / len(json_value)) @pytest.mark.tags(CaseLabel.L1) @pytest.mark.parametrize("auto_id", [True]) @pytest.mark.parametrize("dim", [128]) @pytest.mark.parametrize("entities", [1000]) @pytest.mark.parametrize("enable_dynamic_field", [True]) @pytest.mark.parametrize("sparse_format", ["doc"]) @pytest.mark.parametrize("file_format", ["parquet", "json"]) def test_with_all_field_and_bm25_function_with_bulk_writer( self, auto_id, dim, entities, enable_dynamic_field, sparse_format, file_format ): """ target: test bulk insert with all field and bm25 function method: create collection with all field and bm25 function, then import data with bulk writer expected: verify data imported correctly """ self._connect() fields = [ cf.gen_int64_field(name=df.pk_field, is_primary=True, auto_id=auto_id), cf.gen_int64_field(name=df.int_field), cf.gen_float_field(name=df.float_field), cf.gen_string_field(name=df.string_field), cf.gen_string_field(name=df.text_field, enable_analyzer=True, enable_match=True), cf.gen_json_field(name=df.json_field), cf.gen_array_field(name=df.array_int_field, element_type=DataType.INT64), cf.gen_array_field(name=df.array_float_field, element_type=DataType.FLOAT), cf.gen_array_field(name=df.array_string_field, element_type=DataType.VARCHAR, max_length=100), cf.gen_array_field(name=df.array_bool_field, element_type=DataType.BOOL), cf.gen_float_vec_field(name=df.float_vec_field, dim=dim), cf.gen_sparse_vec_field(name=df.sparse_vec_field), cf.gen_sparse_vec_field(name=df.bm25_sparse_vec_field), ] c_name = cf.gen_unique_str("bulk_insert") schema = cf.gen_collection_schema(fields=fields, auto_id=auto_id, enable_dynamic_field=enable_dynamic_field) bm25_function = Function( name="text_bm25_emb", function_type=FunctionType.BM25, input_field_names=[df.text_field], output_field_names=[df.bm25_sparse_vec_field], params={}, ) schema.add_function(bm25_function) self.collection_wrap.init_collection(c_name, schema=schema) documents = [] if file_format == "parquet": ff = BulkFileType.PARQUET elif file_format == "json": ff = BulkFileType.JSON else: raise Exception(f"not support file format:{file_format}") with RemoteBulkWriter( schema=schema, remote_path="bulk_data", connect_param=RemoteBulkWriter.ConnectParam( bucket_name=self.bucket_name, endpoint=self.minio_endpoint, access_key="minioadmin", secret_key="minioadmin", ), file_type=ff, ) as remote_writer: json_value = [ # 1, # 1.0, # "1", # [1, 2, 3], # ["1", "2", "3"], # [1, 2, "3"], {"key": "value"} ] for i in range(entities): text_value = f"{gen_utf8_string(i)} milvus BM25 text" row = { df.pk_field: i, df.int_field: 1, df.float_field: 1.0, df.string_field: "string", df.text_field: text_value, df.json_field: json_value[i % len(json_value)], df.array_int_field: [1, 2], df.array_float_field: [1.0, 2.0], df.array_string_field: ["string1", "string2"], df.array_bool_field: [True, False], df.float_vec_field: cf.gen_vectors(1, dim)[0], df.sparse_vec_field: cf.gen_sparse_vectors(1, dim, sparse_format=sparse_format)[0], } if auto_id: row.pop(df.pk_field) if enable_dynamic_field: row["name"] = fake.name() row["address"] = fake.address() documents.append(row[df.text_field]) remote_writer.append_row(row) remote_writer.commit() files = remote_writer.batch_files assert all(any(ord(ch) > 127 for ch in document) for document in documents) # import data for f in files: t0 = time.time() task_id, _ = self.utility_wrap.do_bulk_insert(collection_name=c_name, files=f) logging.info(f"bulk insert task ids:{task_id}") success, states = self.utility_wrap.wait_for_bulk_insert_tasks_completed(task_ids=[task_id], timeout=300) tt = time.time() - t0 log.info(f"bulk insert state:{success} in {tt} with states:{states}") assert success num_entities = self.collection_wrap.num_entities log.info(f" collection entities: {num_entities}") assert num_entities == entities # verify imported data is available for search index_params = ct.default_index float_vec_fields = [f.name for f in fields if "vec" in f.name and "float" in f.name] sparse_vec_fields = [f.name for f in fields if "vec" in f.name and "sparse" in f.name and "bm25" not in f.name] bm25_sparse_vec_fields = [f.name for f in fields if "vec" in f.name and "sparse" in f.name and "bm25" in f.name] for f in float_vec_fields: self.collection_wrap.create_index(field_name=f, index_params=index_params) for f in sparse_vec_fields: self.collection_wrap.create_index(field_name=f, index_params=ct.default_sparse_inverted_index) for f in bm25_sparse_vec_fields: self.collection_wrap.create_index(field_name=f, index_params=ct.default_text_sparse_inverted_index) self.collection_wrap.load() log.info("wait for load finished and be ready for search") time.sleep(2) # log.info(f"query seg info: {self.utility_wrap.get_query_segment_info(c_name)[0]}") search_data = cf.gen_vectors(1, dim) search_params = ct.default_search_params res, _ = self.collection_wrap.search( search_data, df.float_vec_field, param=search_params, limit=1, output_fields=["*"], check_task=CheckTasks.check_search_results, check_items={"nq": 1, "limit": 1}, ) for hit in res: for r in hit: fields_from_search = r.fields.keys() for f in fields: if f.name == df.bm25_sparse_vec_field: continue assert f.name in fields_from_search if enable_dynamic_field: assert "name" in fields_from_search assert "address" in fields_from_search # verify full text search word_freq = cf.analyze_documents(documents) token = word_freq.most_common(1)[0][0] search_data = [f" {token} " + fake.text()] search_params = ct.default_text_sparse_search_params res, _ = self.collection_wrap.search( search_data, df.bm25_sparse_vec_field, param=search_params, limit=1, output_fields=["*"], check_task=CheckTasks.check_search_results, check_items={"nq": 1, "limit": 1}, ) for hit in res: for r in hit: assert any(ord(ch) > 127 for ch in r.fields[df.text_field]) @pytest.mark.tags(CaseLabel.L2) @pytest.mark.parametrize("auto_id", [True, False]) @pytest.mark.parametrize("dim", [128]) # 128 @pytest.mark.parametrize("entities", [1000]) # 1000 @pytest.mark.parametrize("enable_dynamic_field", [True, False]) @pytest.mark.parametrize("nullable", [True, False]) def test_with_all_field_numpy_with_bulk_writer(self, auto_id, dim, entities, enable_dynamic_field, nullable): """ """ if nullable is True: pytest.skip("not support bulk writer numpy files in field(int_scalar) which has 'None' data") self._connect() fields = [ cf.gen_int64_field(name=df.pk_field, is_primary=True, auto_id=auto_id), cf.gen_int64_field(name=df.int_field, nullable=nullable), cf.gen_float_field(name=df.float_field), cf.gen_string_field(name=df.string_field), cf.gen_json_field(name=df.json_field), cf.gen_float_vec_field(name=df.float_vec_field, dim=dim), cf.gen_float16_vec_field(name=df.fp16_vec_field, dim=dim), cf.gen_bfloat16_vec_field(name=df.bf16_vec_field, dim=dim), ] c_name = cf.gen_unique_str("bulk_insert") schema = cf.gen_collection_schema(fields=fields, auto_id=auto_id, enable_dynamic_field=enable_dynamic_field) self.collection_wrap.init_collection(c_name, schema=schema) with RemoteBulkWriter( schema=schema, remote_path="bulk_data", connect_param=RemoteBulkWriter.ConnectParam( bucket_name=self.bucket_name, endpoint=self.minio_endpoint, access_key="minioadmin", secret_key="minioadmin", ), file_type=BulkFileType.NUMPY, ) as remote_writer: json_value = [ # 1, # 1.0, # "1", # [1, 2, 3], # ["1", "2", "3"], # [1, 2, "3"], {"key": "value"}, {"number": 1}, {"name": fake.name()}, {"address": fake.address()}, ] for i in range(entities): row = { df.pk_field: i, df.int_field: 1 if not (nullable and random.random() < 0.5) else None, df.float_field: 1.0, df.string_field: "string", df.json_field: json_value[i % len(json_value)], df.float_vec_field: cf.gen_vectors(1, dim)[0], df.fp16_vec_field: cf.gen_vectors(1, dim, vector_data_type=DataType.FLOAT16_VECTOR)[0], df.bf16_vec_field: cf.gen_vectors(1, dim, vector_data_type=DataType.BFLOAT16_VECTOR)[0], } if auto_id: row.pop(df.pk_field) if enable_dynamic_field: row["name"] = fake.name() row["address"] = fake.address() remote_writer.append_row(row) remote_writer.commit() files = remote_writer.batch_files # import data for f in files: t0 = time.time() task_id, _ = self.utility_wrap.do_bulk_insert(collection_name=c_name, files=f) logging.info(f"bulk insert task ids:{task_id}") success, states = self.utility_wrap.wait_for_bulk_insert_tasks_completed(task_ids=[task_id], timeout=300) tt = time.time() - t0 log.info(f"bulk insert state:{success} in {tt} with states:{states}") assert success num_entities = self.collection_wrap.num_entities log.info(f" collection entities: {num_entities}") assert num_entities == entities # verify imported data is available for search index_params = ct.default_index float_vec_fields = [f.name for f in fields if "vec" in f.name and "float" in f.name] sparse_vec_fields = [f.name for f in fields if "vec" in f.name and "sparse" in f.name] for f in float_vec_fields: self.collection_wrap.create_index(field_name=f, index_params=index_params) for f in sparse_vec_fields: self.collection_wrap.create_index(field_name=f, index_params=ct.default_sparse_inverted_index) # add json path index for json field json_path_index_params_double = { "index_type": "INVERTED", "params": {"json_cast_type": "double", "json_path": f"{df.json_field}['number']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_double) json_path_index_params_varchar = { "index_type": "INVERTED", "params": {"json_cast_type": "VARCHAR", "json_path": f"{df.json_field}['address']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_varchar) json_path_index_params_bool = { "index_type": "INVERTED", "params": {"json_cast_type": "Bool", "json_path": f"{df.json_field}['name']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_bool) json_path_index_params_not_exist = { "index_type": "INVERTED", "params": {"json_cast_type": "Double", "json_path": f"{df.json_field}['not_exist']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_not_exist) self.collection_wrap.load() log.info("wait for load finished and be ready for search") time.sleep(2) # log.info(f"query seg info: {self.utility_wrap.get_query_segment_info(c_name)[0]}") search_data = cf.gen_vectors(1, dim) search_params = ct.default_search_params res, _ = self.collection_wrap.search( search_data, df.float_vec_field, param=search_params, limit=1, output_fields=["*"], check_task=CheckTasks.check_search_results, check_items={"nq": 1, "limit": 1}, ) for hit in res: for r in hit: fields_from_search = r.fields.keys() for f in fields: assert f.name in fields_from_search if enable_dynamic_field: assert "name" in fields_from_search assert "address" in fields_from_search res, _ = self.collection_wrap.query(expr=f"{df.json_field}['number'] == 1", output_fields=[df.json_field]) assert len(res) == int(entities / len(json_value)) @pytest.mark.tags(CaseLabel.L2) @pytest.mark.parametrize("auto_id", [True, False]) @pytest.mark.parametrize("dim", [128]) # 128 @pytest.mark.parametrize("entities", [1000]) # 1000 @pytest.mark.parametrize("enable_dynamic_field", [True, False]) @pytest.mark.parametrize("sparse_format", ["doc", "coo"]) @pytest.mark.parametrize("nullable", [True, False]) def test_with_all_field_parquet_with_bulk_writer( self, auto_id, dim, entities, enable_dynamic_field, sparse_format, nullable ): """ """ self._connect() fields = [ cf.gen_int64_field(name=df.pk_field, is_primary=True, auto_id=auto_id), cf.gen_int64_field(name=df.int_field, nullable=nullable), cf.gen_float_field(name=df.float_field, nullable=nullable), cf.gen_string_field(name=df.string_field, nullable=nullable), cf.gen_json_field(name=df.json_field, nullable=nullable), cf.gen_timestamptz_field(name=df.timestamp_field, nullable=nullable), cf.gen_array_field(name=df.array_int_field, element_type=DataType.INT64, nullable=nullable), cf.gen_array_field(name=df.array_float_field, element_type=DataType.FLOAT, nullable=nullable), cf.gen_array_field( name=df.array_string_field, element_type=DataType.VARCHAR, max_length=100, nullable=nullable ), cf.gen_array_field(name=df.array_bool_field, element_type=DataType.BOOL, nullable=nullable), cf.gen_float_vec_field(name=df.float_vec_field, dim=dim), cf.gen_float16_vec_field(name=df.fp16_vec_field, dim=dim), cf.gen_bfloat16_vec_field(name=df.bf16_vec_field, dim=dim), cf.gen_sparse_vec_field(name=df.sparse_vec_field), ] c_name = cf.gen_unique_str("bulk_insert") schema = cf.gen_collection_schema(fields=fields, auto_id=auto_id, enable_dynamic_field=enable_dynamic_field) self.collection_wrap.init_collection(c_name, schema=schema) with RemoteBulkWriter( schema=schema, remote_path="bulk_data", connect_param=RemoteBulkWriter.ConnectParam( bucket_name=self.bucket_name, endpoint=self.minio_endpoint, access_key="minioadmin", secret_key="minioadmin", ), file_type=BulkFileType.JSON, ) as remote_writer: json_value = [ # 1, # 1.0, # "1", # [1, 2, 3], # ["1", "2", "3"], # [1, 2, "3"], {"key": "value"}, {"number": 1}, {"name": fake.name()}, {"address": fake.address()}, ] for i in range(entities): row = { df.pk_field: i, df.int_field: 1 if not (nullable and random.random() < 0.5) else None, df.float_field: 1.0 if not (nullable and random.random() < 0.5) else None, df.string_field: "string" if not (nullable and random.random() < 0.5) else None, df.json_field: json_value[i % len(json_value)] if not (nullable and random.random() < 0.5) else None, df.timestamp_field: cf.gen_timestamptz_str() if not (nullable and random.random() < 0.5) else None, df.array_int_field: [1, 2] if not (nullable and random.random() < 0.5) else None, df.array_float_field: [1.0, 2.0] if not (nullable and random.random() < 0.5) else None, df.array_string_field: ["string1", "string2"] if not (nullable and random.random() < 0.5) else None, df.array_bool_field: [True, False] if not (nullable and random.random() < 0.5) else None, df.float_vec_field: cf.gen_vectors(1, dim)[0], df.fp16_vec_field: cf.gen_vectors(1, dim, vector_data_type=DataType.FLOAT16_VECTOR)[0], df.bf16_vec_field: cf.gen_vectors(1, dim, vector_data_type=DataType.BFLOAT16_VECTOR)[0], df.sparse_vec_field: cf.gen_sparse_vectors(1, dim, sparse_format=sparse_format)[0], } if auto_id: row.pop(df.pk_field) if enable_dynamic_field: row["name"] = fake.name() row["address"] = fake.address() remote_writer.append_row(row) remote_writer.commit() files = remote_writer.batch_files # import data for f in files: t0 = time.time() task_id, _ = self.utility_wrap.do_bulk_insert(collection_name=c_name, files=f) logging.info(f"bulk insert task ids:{task_id}") success, states = self.utility_wrap.wait_for_bulk_insert_tasks_completed(task_ids=[task_id], timeout=300) tt = time.time() - t0 log.info(f"bulk insert state:{success} in {tt} with states:{states}") assert success num_entities = self.collection_wrap.num_entities log.info(f" collection entities: {num_entities}") assert num_entities == entities # verify imported data is available for search index_params = ct.default_index float_vec_fields = [f.name for f in fields if "vec" in f.name and "float" in f.name] sparse_vec_fields = [f.name for f in fields if "vec" in f.name and "sparse" in f.name] for f in float_vec_fields: self.collection_wrap.create_index(field_name=f, index_params=index_params) for f in sparse_vec_fields: self.collection_wrap.create_index(field_name=f, index_params=ct.default_sparse_inverted_index) # add json path index for json field json_path_index_params_double = { "index_type": "INVERTED", "params": {"json_cast_type": "double", "json_path": f"{df.json_field}['number']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_double) json_path_index_params_varchar = { "index_type": "INVERTED", "params": {"json_cast_type": "VARCHAR", "json_path": f"{df.json_field}['address']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_varchar) json_path_index_params_bool = { "index_type": "INVERTED", "params": {"json_cast_type": "Bool", "json_path": f"{df.json_field}['name']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_bool) json_path_index_params_not_exist = { "index_type": "INVERTED", "params": {"json_cast_type": "Double", "json_path": f"{df.json_field}['not_exist']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_not_exist) self.collection_wrap.load() log.info("wait for load finished and be ready for search") time.sleep(2) # log.info(f"query seg info: {self.utility_wrap.get_query_segment_info(c_name)[0]}") search_data = cf.gen_vectors(1, dim) search_params = ct.default_search_params res, _ = self.collection_wrap.search( search_data, df.float_vec_field, param=search_params, limit=1, output_fields=["*"], check_task=CheckTasks.check_search_results, check_items={"nq": 1, "limit": 1}, ) for hit in res: for r in hit: fields_from_search = r.fields.keys() for f in fields: assert f.name in fields_from_search if enable_dynamic_field: assert "name" in fields_from_search assert "address" in fields_from_search res, _ = self.collection_wrap.query(expr=f"{df.json_field}['number'] == 1", output_fields=[df.json_field]) if not nullable: assert len(res) == int(entities / len(json_value)) else: assert 0 < len(res) < int(entities / len(json_value)) @pytest.mark.tags(CaseLabel.L2) @pytest.mark.parametrize("auto_id", [True, False]) @pytest.mark.parametrize("dim", [128]) # 128 @pytest.mark.parametrize("entities", [1000]) # 1000 @pytest.mark.parametrize("enable_dynamic_field", [True, False]) @pytest.mark.parametrize("sparse_format", ["doc", "coo"]) @pytest.mark.parametrize("nullable", [True, False]) @pytest.mark.parametrize("add_field", [True, False]) def test_with_all_field_csv_with_bulk_writer( self, auto_id, dim, entities, enable_dynamic_field, sparse_format, nullable, add_field ): """ """ self._connect() fields = [ cf.gen_int64_field(name=df.pk_field, is_primary=True, auto_id=auto_id), cf.gen_int64_field(name=df.int_field, nullable=nullable), cf.gen_float_field(name=df.float_field, nullable=nullable), cf.gen_string_field(name=df.string_field, nullable=nullable), cf.gen_json_field(name=df.json_field, nullable=nullable), cf.gen_timestamptz_field(name=df.timestamp_field, nullable=nullable), cf.gen_array_field(name=df.array_int_field, element_type=DataType.INT64, nullable=nullable), cf.gen_array_field(name=df.array_float_field, element_type=DataType.FLOAT, nullable=nullable), cf.gen_array_field( name=df.array_string_field, element_type=DataType.VARCHAR, max_length=100, nullable=nullable ), cf.gen_array_field(name=df.array_bool_field, element_type=DataType.BOOL, nullable=nullable), cf.gen_float_vec_field(name=df.float_vec_field, dim=dim), cf.gen_float16_vec_field(name=df.fp16_vec_field, dim=dim), cf.gen_bfloat16_vec_field(name=df.bf16_vec_field, dim=dim), cf.gen_sparse_vec_field(name=df.sparse_vec_field), ] c_name = cf.gen_unique_str("bulk_insert") schema = cf.gen_collection_schema(fields=fields, auto_id=auto_id, enable_dynamic_field=enable_dynamic_field) self.collection_wrap.init_collection(c_name, schema=schema) with RemoteBulkWriter( schema=schema, remote_path="bulk_data", connect_param=RemoteBulkWriter.ConnectParam( bucket_name=self.bucket_name, endpoint=self.minio_endpoint, access_key="minioadmin", secret_key="minioadmin", ), file_type=BulkFileType.CSV, ) as remote_writer: json_value = [{"key": "value"}, {"number": 1}, {"name": fake.name()}, {"address": fake.address()}] for i in range(entities): row = { df.pk_field: i, df.int_field: 1 if not (nullable and random.random() < 0.5) else None, df.float_field: 1.0 if not (nullable and random.random() < 0.5) else None, df.string_field: "string" if not (nullable and random.random() < 0.5) else None, df.json_field: json_value[i % len(json_value)] if not (nullable and random.random() < 0.5) else None, df.timestamp_field: cf.gen_timestamptz_str() if not (nullable and random.random() < 0.5) else None, df.array_int_field: [1, 2] if not (nullable and random.random() < 0.5) else None, df.array_float_field: [1.0, 2.0] if not (nullable and random.random() < 0.5) else None, df.array_string_field: ["string1", "string2"] if not (nullable and random.random() < 0.5) else None, df.array_bool_field: [True, False] if not (nullable and random.random() < 0.5) else None, df.float_vec_field: cf.gen_vectors(1, dim)[0], df.fp16_vec_field: cf.gen_vectors(1, dim, vector_data_type=DataType.FLOAT16_VECTOR)[0], df.bf16_vec_field: cf.gen_vectors(1, dim, vector_data_type=DataType.BFLOAT16_VECTOR)[0], df.sparse_vec_field: cf.gen_sparse_vectors(1, dim, sparse_format=sparse_format)[0], } if auto_id: row.pop(df.pk_field) if enable_dynamic_field: row["name"] = fake.name() row["address"] = fake.address() remote_writer.append_row(row) remote_writer.commit() files = remote_writer.batch_files if add_field: self._connect(enable_milvus_client_api=True) self.client.add_collection_field( collection_name=c_name, field_name=df.new_field, data_type=DataType.INT64, nullable=True ) # import data for f in files: t0 = time.time() task_id, _ = self.utility_wrap.do_bulk_insert(collection_name=c_name, files=f) logging.info(f"bulk insert task ids:{task_id}") success, states = self.utility_wrap.wait_for_bulk_insert_tasks_completed(task_ids=[task_id], timeout=300) tt = time.time() - t0 log.info(f"bulk insert state:{success} in {tt} with states:{states}") assert success num_entities = self.collection_wrap.num_entities log.info(f" collection entities: {num_entities}") assert num_entities == entities # verify imported data is available for search index_params = ct.default_index float_vec_fields = [f.name for f in fields if "vec" in f.name and "float" in f.name] sparse_vec_fields = [f.name for f in fields if "vec" in f.name and "sparse" in f.name] for f in float_vec_fields: self.collection_wrap.create_index(field_name=f, index_params=index_params) for f in sparse_vec_fields: self.collection_wrap.create_index(field_name=f, index_params=ct.default_sparse_inverted_index) # add json path index for json field json_path_index_params_double = { "index_type": "INVERTED", "params": {"json_cast_type": "double", "json_path": f"{df.json_field}['number']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_double) json_path_index_params_varchar = { "index_type": "INVERTED", "params": {"json_cast_type": "VARCHAR", "json_path": f"{df.json_field}['address']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_varchar) json_path_index_params_bool = { "index_type": "INVERTED", "params": {"json_cast_type": "Bool", "json_path": f"{df.json_field}['name']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_bool) json_path_index_params_not_exist = { "index_type": "INVERTED", "params": {"json_cast_type": "Double", "json_path": f"{df.json_field}['not_exist']"}, } self.collection_wrap.create_index(field_name=df.json_field, index_params=json_path_index_params_not_exist) self.collection_wrap.load() log.info("wait for load finished and be ready for search") time.sleep(2) # log.info(f"query seg info: {self.utility_wrap.get_query_segment_info(c_name)[0]}") search_data = cf.gen_vectors(1, dim) search_params = ct.default_search_params res, _ = self.collection_wrap.search( search_data, df.float_vec_field, param=search_params, limit=1, output_fields=["*"], check_task=CheckTasks.check_search_results, check_items={"nq": 1, "limit": 1}, ) for hit in res: for r in hit: fields_from_search = r.fields.keys() for f in fields: assert f.name in fields_from_search if enable_dynamic_field: assert "name" in fields_from_search assert "address" in fields_from_search res, _ = self.collection_wrap.query(expr=f"{df.json_field}['number'] == 1", output_fields=[df.json_field]) if not nullable: assert len(res) == int(entities / len(json_value)) else: assert 0 < len(res) < int(entities / len(json_value)) if add_field: res, _ = self.collection_wrap.query( expr=f"{df.new_field} is not null", output_fields=[df.string_field, df.int_field, df.new_field] ) assert len(res) == 0 @pytest.mark.tags(CaseLabel.L2) @pytest.mark.parametrize("auto_id", [True]) @pytest.mark.parametrize("dim", [128]) # 128 @pytest.mark.parametrize("entities", [1000]) # 1000 @pytest.mark.parametrize("file_nums", [0, 10]) @pytest.mark.parametrize("array_len", [1]) def test_with_wrong_parquet_file_num(self, auto_id, dim, entities, file_nums, array_len): """ collection schema 1: [pk, int64, float64, string float_vector] data file: vectors.parquet and uid.parquet, Steps: 1. create collection 2. import data 3. verify failure, because only one file is allowed """ fields = [ cf.gen_int64_field(name=df.pk_field, is_primary=True, auto_id=auto_id), cf.gen_int64_field(name=df.int_field), cf.gen_float_field(name=df.float_field), cf.gen_double_field(name=df.double_field), cf.gen_array_field(name=df.array_int_field, element_type=DataType.INT64), cf.gen_array_field(name=df.array_float_field, element_type=DataType.FLOAT), cf.gen_array_field(name=df.array_string_field, element_type=DataType.VARCHAR, max_length=100), cf.gen_array_field(name=df.array_bool_field, element_type=DataType.BOOL), cf.gen_float_vec_field(name=df.vec_field, dim=dim), ] data_fields = [f.name for f in fields if not f.to_dict().get("auto_id", False)] files = prepare_bulk_insert_parquet_files( minio_endpoint=self.minio_endpoint, bucket_name=self.bucket_name, rows=entities, dim=dim, data_fields=data_fields, file_nums=file_nums, array_length=array_len, force=True, ) self._connect() c_name = cf.gen_unique_str("bulk_insert") schema = cf.gen_collection_schema(fields=fields, auto_id=auto_id) self.collection_wrap.init_collection(c_name, schema=schema) # import data error = {} if file_nums == 0: error = {ct.err_code: 1100, ct.err_msg: "import request is empty"} if file_nums > 1: error = {ct.err_code: 65535, ct.err_msg: "for Parquet import, accepts only one file"} self.utility_wrap.do_bulk_insert( collection_name=c_name, files=files, check_task=CheckTasks.err_res, check_items=error ) @pytest.mark.tags(CaseLabel.L2) @pytest.mark.parametrize("auto_id", [True, False]) @pytest.mark.parametrize("dim", [128]) @pytest.mark.parametrize("entities", [2000]) @pytest.mark.parametrize("file_nums", [5]) def test_multi_numpy_files_from_diff_folders(self, auto_id, dim, entities, file_nums): """ collection schema 1: [pk, float_vector] data file: .npy files in different folders Steps: 1. create collection, create index and load 2. import data 3. verify that import numpy files in a loop """ self._connect() c_name = cf.gen_unique_str("bulk_insert") fields = [ cf.gen_int64_field(name=df.pk_field, is_primary=True, auto_id=auto_id), cf.gen_int64_field(name=df.int_field), cf.gen_float_field(name=df.float_field), cf.gen_double_field(name=df.double_field), cf.gen_float_vec_field(name=df.vec_field, dim=dim), ] schema = cf.gen_collection_schema(fields=fields, auto_id=auto_id) self.collection_wrap.init_collection(c_name, schema=schema) # build index index_params = ct.default_index self.collection_wrap.create_index(field_name=df.vec_field, index_params=index_params) # load collection self.collection_wrap.load() data_fields = [f.name for f in fields if not f.to_dict().get("auto_id", False)] task_ids = [] for i in range(file_nums): files = prepare_bulk_insert_numpy_files( minio_endpoint=self.minio_endpoint, bucket_name=self.bucket_name, rows=entities, dim=dim, data_fields=data_fields, file_nums=1, force=True, ) task_id, _ = self.utility_wrap.do_bulk_insert(collection_name=c_name, files=files) task_ids.append(task_id) success, states = self.utility_wrap.wait_for_bulk_insert_tasks_completed(task_ids=task_ids, timeout=300) log.info(f"bulk insert state:{success}") assert success log.info(f" collection entities: {self.collection_wrap.num_entities}") assert self.collection_wrap.num_entities == entities * file_nums # verify search and query log.info("wait for load finished and be ready for search") self.collection_wrap.load(_refresh=True) time.sleep(2) search_data = cf.gen_vectors(1, dim) search_params = ct.default_search_params res, _ = self.collection_wrap.search( search_data, df.vec_field, param=search_params, limit=1, check_task=CheckTasks.check_search_results, check_items={"nq": 1, "limit": 1}, ) @pytest.mark.tags(CaseLabel.L2) @pytest.mark.parametrize("is_row_based", [True]) @pytest.mark.parametrize("auto_id", [True, False]) @pytest.mark.parametrize("par_key_field", [df.int_field, df.string_field]) def test_partition_key_on_json_file(self, is_row_based, auto_id, par_key_field): """ collection: auto_id, customized_id collection schema: [pk, int64, varchar, float_vector] Steps: 1. create collection with partition key enabled 2. import data 3. verify the data entities equal the import data and distributed by values of partition key field 4. load the collection 5. verify search successfully 6. verify query successfully """ dim = 12 entities = 200 self._connect() c_name = cf.gen_unique_str("bulk_partition_key") fields = [ cf.gen_int64_field(name=df.pk_field, is_primary=True, auto_id=auto_id), cf.gen_float_vec_field(name=df.float_vec_field, dim=dim), cf.gen_int64_field(name=df.int_field, is_partition_key=(par_key_field == df.int_field)), cf.gen_string_field(name=df.string_field, is_partition_key=(par_key_field == df.string_field)), cf.gen_bool_field(name=df.bool_field), cf.gen_float_field(name=df.float_field), cf.gen_array_field(name=df.array_int_field, element_type=DataType.INT64), ] data_fields = [f.name for f in fields if not f.to_dict().get("auto_id", False)] schema = cf.gen_collection_schema(fields=fields, auto_id=auto_id) files = prepare_bulk_insert_new_json_files( minio_endpoint=self.minio_endpoint, bucket_name=self.bucket_name, is_row_based=is_row_based, rows=entities, dim=dim, auto_id=auto_id, data_fields=data_fields, force=True, schema=schema, ) self.collection_wrap.init_collection(c_name, schema=schema, num_partitions=10) assert len(self.collection_wrap.partitions) == 10 # import data t0 = time.time() task_id, _ = self.utility_wrap.do_bulk_insert( collection_name=c_name, partition_name=None, files=files, ) logging.info(f"bulk insert task id:{task_id}") success, _ = self.utility_wrap.wait_for_bulk_insert_tasks_completed(task_ids=[task_id], timeout=300) tt = time.time() - t0 log.info(f"bulk insert state:{success} in {tt}") assert success num_entities = self.collection_wrap.num_entities log.info(f" collection entities: {num_entities}") assert num_entities == entities # verify imported data is available for search index_params = ct.default_index self.collection_wrap.create_index(field_name=df.float_vec_field, index_params=index_params) self.collection_wrap.load() log.info("wait for load finished and be ready for search") time.sleep(10) log.info(f"query seg info: {self.utility_wrap.get_query_segment_info(c_name)[0]}") nq = 2 topk = 2 search_data = cf.gen_vectors(nq, dim) search_params = ct.default_search_params res, _ = self.collection_wrap.search( search_data, df.float_vec_field, param=search_params, limit=topk, check_task=CheckTasks.check_search_results, check_items={"nq": nq, "limit": topk}, ) for hits in res: ids = hits.ids results, _ = self.collection_wrap.query(expr=f"{df.pk_field} in {ids}") assert len(results) == len(ids) # verify data was bulk inserted into different partitions num_entities = 0 empty_partition_num = 0 for p in self.collection_wrap.partitions: if p.num_entities == 0: empty_partition_num += 1 num_entities += p.num_entities assert num_entities == entities # verify error when trying to bulk insert into a specific partition err_msg = "not allow to set partition name for collection with partition key" task_id, _ = self.utility_wrap.do_bulk_insert( collection_name=c_name, partition_name=self.collection_wrap.partitions[0].name, files=files, check_task=CheckTasks.err_res, check_items={"err_code": 2100, "err_msg": err_msg}, ) @pytest.mark.tags(CaseLabel.L2) @pytest.mark.parametrize("auto_id", [True, False]) @pytest.mark.parametrize("dim", [13]) @pytest.mark.parametrize("entities", [150]) @pytest.mark.parametrize("file_nums", [10]) def test_partition_key_on_multi_numpy_files(self, auto_id, dim, entities, file_nums): """ collection schema 1: [pk, int64, float_vector, double] data file: .npy files in different folders Steps: 1. create collection with partition key enabled, create index and load 2. import data 3. verify that import numpy files in a loop """ self._connect() c_name = cf.gen_unique_str("bulk_ins_parkey") fields = [ cf.gen_int64_field(name=df.pk_field, is_primary=True), cf.gen_int64_field(name=df.int_field, is_partition_key=True), cf.gen_float_field(name=df.float_field), cf.gen_double_field(name=df.double_field), cf.gen_float_vec_field(name=df.vec_field, dim=dim), ] schema = cf.gen_collection_schema(fields=fields) self.collection_wrap.init_collection(c_name, schema=schema, num_partitions=10) # build index index_params = ct.default_index self.collection_wrap.create_index(field_name=df.vec_field, index_params=index_params) # load collection self.collection_wrap.load() data_fields = [f.name for f in fields if not f.to_dict().get("auto_id", False)] task_ids = [] for i in range(file_nums): files = prepare_bulk_insert_numpy_files( minio_endpoint=self.minio_endpoint, bucket_name=self.bucket_name, rows=entities, dim=dim, data_fields=data_fields, file_nums=1, force=True, ) task_id, _ = self.utility_wrap.do_bulk_insert(collection_name=c_name, files=files) task_ids.append(task_id) success, states = self.utility_wrap.wait_for_bulk_insert_tasks_completed(task_ids=task_ids, timeout=300) log.info(f"bulk insert state:{success}") assert success log.info(f" collection entities: {self.collection_wrap.num_entities}") assert self.collection_wrap.num_entities == entities * file_nums # verify imported data is indexed success = self.utility_wrap.wait_index_build_completed(c_name) assert success # verify search and query log.info("wait for load finished and be ready for search") self.collection_wrap.load(_refresh=True) time.sleep(2) search_data = cf.gen_vectors(1, dim) search_params = ct.default_search_params res, _ = self.collection_wrap.search( search_data, df.vec_field, param=search_params, limit=1, check_task=CheckTasks.check_search_results, check_items={"nq": 1, "limit": 1}, ) # verify data was bulk inserted into different partitions num_entities = 0 empty_partition_num = 0 for p in self.collection_wrap.partitions: if p.num_entities == 0: empty_partition_num += 1 num_entities += p.num_entities assert num_entities == entities * file_nums @pytest.mark.parametrize("pk_field", [df.pk_field, df.string_field]) @pytest.mark.tags(CaseLabel.L2) def test_bulk_import_random_pk_stats_task(self, pk_field): # connect -> prepare json data self._connect() collection_name = cf.gen_unique_str("stats_task") nb = 3000 fields = [] files = "" # prepare data: int64_pk -> json data; varchar_pk -> numpy data if pk_field == df.pk_field: fields = [ cf.gen_int64_field(name=df.pk_field, is_primary=True, auto_id=False), cf.gen_float_vec_field(name=df.float_vec_field, dim=ct.default_dim), ] data_fields = [f.name for f in fields if not f.to_dict().get("auto_id", False)] files = prepare_bulk_insert_new_json_files( minio_endpoint=self.minio_endpoint, bucket_name=self.bucket_name, is_row_based=True, rows=nb, dim=ct.default_dim, auto_id=False, data_fields=data_fields, force=True, shuffle=True, ) elif pk_field == df.string_field: fields = [ cf.gen_string_field(name=df.string_field, is_primary=True, auto_id=False), cf.gen_float_vec_field(name=df.float_vec_field, dim=ct.default_dim), ] data_fields = [f.name for f in fields if not f.to_dict().get("auto_id", False)] files = prepare_bulk_insert_numpy_files( minio_endpoint=self.minio_endpoint, bucket_name=self.bucket_name, rows=nb, dim=ct.default_dim, data_fields=data_fields, enable_dynamic_field=False, force=True, shuffle_pk=True, ) else: log.error(f"pk_field name {pk_field} not supported now, [{df.pk_field}, {df.string_field}] expected~") # create collection -> create vector index schema = cf.gen_collection_schema(fields=fields) self.collection_wrap.init_collection(collection_name, schema=schema) self.build_multi_index(index_params=DefaultVectorIndexParams.IVF_SQ8(df.float_vec_field)) # bulk_insert data t0 = time.time() task_id, _ = self.utility_wrap.do_bulk_insert(collection_name=collection_name, files=files) logging.info(f"bulk insert task ids:{task_id}") completed, _ = self.utility_wrap.wait_for_bulk_insert_tasks_completed(task_ids=[task_id], timeout=300) tt = time.time() - t0 log.info(f"bulk insert state:{completed} with latency {tt}") assert completed # load -> get_segment_info -> verify stats task self.collection_wrap.load() res_segment_info, _ = self.utility_wrap.get_query_segment_info(collection_name) assert len(res_segment_info) > 0 # maybe mix compaction to 1 segment cnt = 0 for r in res_segment_info: log.info(f"segmentID {r.segmentID}: state: {r.state}; num_rows: {r.num_rows}; is_sorted: {r.is_sorted} ") cnt += r.num_rows assert r.is_sorted is True assert cnt == nb # verify search self.collection_wrap.search( data=cf.gen_vectors(ct.default_nq, ct.default_dim, vector_data_type=DataType.FLOAT_VECTOR), anns_field=df.float_vec_field, param=DefaultVectorSearchParams.IVF_SQ8(), limit=ct.default_limit, check_task=CheckTasks.check_search_results, check_items={"nq": ct.default_nq, "limit": ct.default_limit}, ) class TestImportWithTextEmbeddingFunction(TestcaseBase): """ ****************************************************************** The following cases are used to test import with text embedding ****************************************************************** """ @pytest.mark.parametrize("file_format", ["json", "parquet", "numpy"]) @pytest.mark.parametrize("add_field", [True, False]) @pytest.mark.tags(CaseLabel.L1) def test_import_without_embedding(self, tei_endpoint, minio_host, minio_bucket, file_format, add_field): """ target: test import data without embedding method: 1. create collection 2. import data without embedding field 3. verify embeddings are generated expected: embeddings should be generated after import """ dim = 768 fields = [ FieldSchema(name="id", dtype=DataType.INT64, is_primary=True), FieldSchema(name="document", dtype=DataType.VARCHAR, max_length=65535), FieldSchema(name="dense", dtype=DataType.FLOAT_VECTOR, dim=dim), ] schema = CollectionSchema(fields=fields, description="test collection") text_embedding_function = Function( name="text_embedding", function_type=FunctionType.TEXTEMBEDDING, input_field_names=["document"], output_field_names="dense", params={ "provider": "TEI", "endpoint": tei_endpoint, }, ) schema.add_function(text_embedding_function) c_name = cf.gen_unique_str("import_without_embedding") collection_w = self.init_collection_wrap(name=c_name, schema=schema) # prepare import data without embedding nb = 1000 if file_format == "json": file_type = BulkFileType.JSON elif file_format == "numpy": file_type = BulkFileType.NUMPY else: file_type = BulkFileType.PARQUET with RemoteBulkWriter( schema=schema, remote_path="bulk_data", connect_param=RemoteBulkWriter.ConnectParam( bucket_name=minio_bucket, endpoint=f"{minio_host}:9000", access_key="minioadmin", secret_key="minioadmin", ), file_type=file_type, ) as remote_writer: for i in range(nb): row = {"id": i, "document": f"This is test document {i}"} remote_writer.append_row(row) remote_writer.commit() files = remote_writer.batch_files if add_field: self._connect(enable_milvus_client_api=True) self.client.add_collection_field( collection_name=c_name, field_name=df.new_field, data_type=DataType.INT64, nullable=True ) # import data for f in files: t0 = time.time() task_id, _ = self.utility_wrap.do_bulk_insert(collection_name=c_name, files=f) log.info(f"bulk insert task ids:{task_id}") success, states = self.utility_wrap.wait_for_bulk_insert_tasks_completed(task_ids=[task_id], timeout=300) tt = time.time() - t0 log.info(f"bulk insert state:{success} in {tt} with states:{states}") assert success num_entities = collection_w.num_entities log.info(f" collection entities: {num_entities}") assert num_entities == nb # create index and load index_params = { "index_type": "AUTOINDEX", "metric_type": "COSINE", "params": {}, } collection_w.create_index("dense", index_params) collection_w.load() # verify embeddings are generated res, _ = collection_w.query(expr="id >= 0", output_fields=["dense"]) assert len(res) == nb for r in res: assert "dense" in r assert len(r["dense"]) == dim if add_field: res, _ = collection_w.query(expr=f"{df.new_field} is not null", output_fields=[df.new_field]) assert len(res) == 0 class TestImportWithFunctionNegative(TestcaseBase): """ ****************************************************************** The following cases are used to test import with text embedding ****************************************************************** """ @pytest.mark.parametrize("file_format", ["json", "parquet"]) @pytest.mark.tags(CaseLabel.L2) def test_import_for_bm25_function_with_output_field(self, tei_endpoint, minio_host, minio_bucket, file_format): """ target: test import data for bm25 with output field method: 1. create collection 2. import data with output field 3. verify import failed expected: import failed """ fields = [ FieldSchema(name="id", dtype=DataType.INT64, is_primary=True), FieldSchema(name="document", dtype=DataType.VARCHAR, max_length=65535, enable_analyzer=True), FieldSchema(name="bm25", dtype=DataType.SPARSE_FLOAT_VECTOR), ] schema = CollectionSchema(fields=fields, description="test collection") bm25_function = Function( name="text_embedding", function_type=FunctionType.BM25, input_field_names=["document"], output_field_names=["bm25"], params={}, ) schema.add_function(bm25_function) c_name = cf.gen_unique_str("import_without_embedding") self.init_collection_wrap(name=c_name, schema=schema) # prepare import data with function output invalid_schema = CollectionSchema(fields=fields, description="test collection") nb = 1000 rng = np.random.default_rng() if file_format == "json": file_type = BulkFileType.JSON elif file_format == "numpy": file_type = BulkFileType.NUMPY else: file_type = BulkFileType.PARQUET with RemoteBulkWriter( schema=invalid_schema, remote_path="bulk_data", connect_param=RemoteBulkWriter.ConnectParam( bucket_name=minio_bucket, endpoint=f"{minio_host}:9000", access_key="minioadmin", secret_key="minioadmin", ), file_type=file_type, ) as remote_writer: for i in range(nb): row = { "id": i, "document": f"This is test document {i}", "bm25": {d: rng.random() for d in random.sample(range(1000), random.randint(20, 30))}, } remote_writer.append_row(row) remote_writer.commit() files = remote_writer.batch_files # import data for f in files: t0 = time.time() task_id, _ = self.utility_wrap.do_bulk_insert(collection_name=c_name, files=f) log.info(f"bulk insert task ids:{task_id}") success, states = self.utility_wrap.wait_for_bulk_insert_tasks_completed(task_ids=[task_id], timeout=300) tt = time.time() - t0 log.info(f"bulk insert state:{success} in {tt} with states:{states}") assert not success @pytest.mark.parametrize("file_format", ["json", "parquet"]) @pytest.mark.tags(CaseLabel.L1) def test_import_for_text_embedding_function_with_output_field( self, tei_endpoint, minio_host, minio_bucket, file_format ): """ target: test import data for text embedding function with output field method: 1. create collection 2. import data for text embedding function with output field 3. import failed expected: import failed """ dim = 768 fields = [ FieldSchema(name="id", dtype=DataType.INT64, is_primary=True), FieldSchema(name="document", dtype=DataType.VARCHAR, max_length=65535), FieldSchema(name="dense", dtype=DataType.FLOAT_VECTOR, dim=dim), ] schema = CollectionSchema(fields=fields, description="test collection") text_embedding_function = Function( name="text_embedding", function_type=FunctionType.TEXTEMBEDDING, input_field_names=["document"], output_field_names="dense", params={ "provider": "TEI", "endpoint": tei_endpoint, }, ) schema.add_function(text_embedding_function) c_name = cf.gen_unique_str("import_without_embedding") self.init_collection_wrap(name=c_name, schema=schema) # prepare import data embedding invalid_schema = CollectionSchema(fields=fields, description="test collection") nb = 1000 if file_format == "json": file_type = BulkFileType.JSON elif file_format == "numpy": file_type = BulkFileType.NUMPY else: file_type = BulkFileType.PARQUET with RemoteBulkWriter( schema=invalid_schema, remote_path="bulk_data", connect_param=RemoteBulkWriter.ConnectParam( bucket_name=minio_bucket, endpoint=f"{minio_host}:9000", access_key="minioadmin", secret_key="minioadmin", ), file_type=file_type, ) as remote_writer: for i in range(nb): row = { "id": i, "document": f"This is test document {i}", "dense": [random.random() for _ in range(dim)], } remote_writer.append_row(row) remote_writer.commit() files = remote_writer.batch_files # import data for f in files: t0 = time.time() task_id, _ = self.utility_wrap.do_bulk_insert(collection_name=c_name, files=f) log.info(f"bulk insert task ids:{task_id}") success, states = self.utility_wrap.wait_for_bulk_insert_tasks_completed(task_ids=[task_id], timeout=300) tt = time.time() - t0 log.info(f"bulk insert state:{success} in {tt} with states:{states}") assert not success