mirror of https://github.com/milvus-io/milvus.git
test: add partition key isolation test case (#39403)
/kind improvement --------- Signed-off-by: zhuwenxing <wenxing.zhu@zilliz.com>pull/39503/head
zhuwenxing
GitHub
parent
eab102e976
commit
40e6fcd868
|
|
@ -0,0 +1,357 @@
|
||||||
|
from base.client_base import TestcaseBase
|
||||||
|
from common import common_func as cf
|
||||||
|
from common.common_type import CaseLabel
|
||||||
|
from utils.util_log import test_log as log
|
||||||
|
import time
|
||||||
|
import pytest
|
||||||
|
import random
|
||||||
|
from pymilvus import (
|
||||||
|
list_collections,
|
||||||
|
FieldSchema, CollectionSchema, DataType,
|
||||||
|
Collection, utility
|
||||||
|
)
|
||||||
|
import pandas as pd
|
||||||
|
import faker
|
||||||
|
fake = faker.Faker()
|
||||||
|
|
||||||
|
|
||||||
|
prefix = "par_key_isolation_"
|
||||||
|
|
||||||
|
|
||||||
|
class TestPartitionKeyIsolation(TestcaseBase):
|
||||||
|
""" Test case of partition key isolation"""
|
||||||
|
@pytest.mark.tags(CaseLabel.L3)
|
||||||
|
def test_par_key_isolation_with_valid_expr(self):
|
||||||
|
# create
|
||||||
|
self._connect()
|
||||||
|
collection_name = cf.gen_unique_str(prefix)
|
||||||
|
partition_key = "scalar_6"
|
||||||
|
enable_isolation = "true"
|
||||||
|
if collection_name in list_collections():
|
||||||
|
log.info(f"collection {collection_name} exists, drop it")
|
||||||
|
Collection(name=collection_name).drop()
|
||||||
|
fields = [
|
||||||
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
|
||||||
|
FieldSchema(name="scalar_3", dtype=DataType.VARCHAR, max_length=1000,
|
||||||
|
is_partition_key=bool(partition_key == "scalar_3")),
|
||||||
|
FieldSchema(name="scalar_6", dtype=DataType.VARCHAR, max_length=1000,
|
||||||
|
is_partition_key=bool(partition_key == "scalar_6")),
|
||||||
|
FieldSchema(name="scalar_9", dtype=DataType.VARCHAR, max_length=1000,
|
||||||
|
is_partition_key=bool(partition_key == "scalar_9")),
|
||||||
|
FieldSchema(name="scalar_12", dtype=DataType.VARCHAR, max_length=1000,
|
||||||
|
is_partition_key=bool(partition_key == "scalar_12")),
|
||||||
|
FieldSchema(name="scalar_5_linear", dtype=DataType.VARCHAR, max_length=1000,
|
||||||
|
is_partition_key=bool(partition_key == "scalar_5_linear")),
|
||||||
|
FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=768)
|
||||||
|
]
|
||||||
|
schema = CollectionSchema(fields=fields, description="test collection", enable_dynamic_field=True,
|
||||||
|
num_partitions=1)
|
||||||
|
collection = Collection(name=collection_name, schema=schema, num_partitions=1)
|
||||||
|
|
||||||
|
collection.set_properties({"partitionkey.isolation": enable_isolation})
|
||||||
|
log.info(f"collection {collection_name} created: {collection.describe()}")
|
||||||
|
index_params = {"metric_type": "L2", "index_type": "HNSW", "params": {"M": 30, "efConstruction": 360}}
|
||||||
|
log.info(f"collection {collection_name} created")
|
||||||
|
batch_size = 1000
|
||||||
|
data_size = 10000
|
||||||
|
epoch = data_size // batch_size
|
||||||
|
remainder = data_size % batch_size
|
||||||
|
all_data = []
|
||||||
|
for i in range(epoch + 1):
|
||||||
|
if i == epoch:
|
||||||
|
if remainder == 0:
|
||||||
|
break
|
||||||
|
batch_size = remainder
|
||||||
|
start_idx = i * batch_size
|
||||||
|
end_idx = (i + 1) * batch_size
|
||||||
|
t0 = time.time()
|
||||||
|
data = {
|
||||||
|
"id": [i for i in range(start_idx, end_idx)],
|
||||||
|
"scalar_3": [str(i % 3) for i in range(start_idx, end_idx)],
|
||||||
|
"scalar_6": [str(i % 6) for i in range(start_idx, end_idx)],
|
||||||
|
"scalar_9": [str(i % 9) for i in range(start_idx, end_idx)],
|
||||||
|
"scalar_12": [str(i % 12) for i in range(start_idx, end_idx)],
|
||||||
|
"scalar_5_linear": [str(i % 5) for i in range(start_idx, end_idx)],
|
||||||
|
"emb": [[random.random() for _ in range(768)] for _ in range(batch_size)]
|
||||||
|
}
|
||||||
|
df = pd.DataFrame(data)
|
||||||
|
all_data.append(df)
|
||||||
|
log.info(f"generate test data {batch_size} cost time {time.time() - t0}")
|
||||||
|
collection.insert(df)
|
||||||
|
all_df = pd.concat(all_data)
|
||||||
|
collection.compact()
|
||||||
|
collection.wait_for_compaction_completed()
|
||||||
|
t0 = time.time()
|
||||||
|
collection.create_index("emb", index_params=index_params)
|
||||||
|
index_list = utility.list_indexes(collection_name=collection_name)
|
||||||
|
for index_name in index_list:
|
||||||
|
progress = utility.index_building_progress(collection_name=collection_name, index_name=index_name)
|
||||||
|
while progress["pending_index_rows"] > 0:
|
||||||
|
time.sleep(30)
|
||||||
|
progress = utility.index_building_progress(collection_name=collection_name, index_name=index_name)
|
||||||
|
log.info(f"collection {collection_name} index {index_name} progress: {progress}")
|
||||||
|
log.info(f"collection {collection_name} index {index_name} progress: {progress}")
|
||||||
|
tt = time.time() - t0
|
||||||
|
log.info(f"create index cost time {tt}")
|
||||||
|
collection.compact()
|
||||||
|
collection.wait_for_compaction_completed()
|
||||||
|
t0 = time.time()
|
||||||
|
collection.load()
|
||||||
|
log.info(f"load collection cost time {time.time() - t0}")
|
||||||
|
num = collection.num_entities
|
||||||
|
log.info(f"collection {collection_name} loaded, num_entities: {num}")
|
||||||
|
|
||||||
|
valid_expressions = [
|
||||||
|
"scalar_6 == '1' and scalar_12 == '1'",
|
||||||
|
"scalar_6 == '1' and scalar_12 > '1'",
|
||||||
|
"scalar_6 == '3' and (scalar_12 == '1' or scalar_3 != '1')",
|
||||||
|
"scalar_6 == '2' and ('4' < scalar_12 < '6' or scalar_3 == '1')",
|
||||||
|
"scalar_6 == '5' and scalar_12 in ['1', '3', '5']",
|
||||||
|
"scalar_6 == '1'"
|
||||||
|
]
|
||||||
|
for expr in valid_expressions:
|
||||||
|
res = collection.search(
|
||||||
|
data=[[random.random() for _ in range(768)]],
|
||||||
|
anns_field="emb",
|
||||||
|
expr=expr,
|
||||||
|
param={"metric_type": "L2", "params": {"nprobe": 16}},
|
||||||
|
limit=10000,
|
||||||
|
output_fields=["scalar_3", "scalar_6", "scalar_12"]
|
||||||
|
)
|
||||||
|
log.info(f"search res {res}")
|
||||||
|
true_res = all_df.query(expr)
|
||||||
|
log.info(f"true res {true_res}")
|
||||||
|
assert len(res[0]) == len(true_res)
|
||||||
|
|
||||||
|
@pytest.mark.tags(CaseLabel.L3)
|
||||||
|
def test_par_key_isolation_with_unsupported_expr(self):
|
||||||
|
# create
|
||||||
|
self._connect()
|
||||||
|
collection_name = cf.gen_unique_str(prefix)
|
||||||
|
partition_key = "scalar_6"
|
||||||
|
enable_isolation = "true"
|
||||||
|
if collection_name in list_collections():
|
||||||
|
log.info(f"collection {collection_name} exists, drop it")
|
||||||
|
Collection(name=collection_name).drop()
|
||||||
|
fields = [
|
||||||
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
|
||||||
|
FieldSchema(name="scalar_3", dtype=DataType.VARCHAR, max_length=1000,
|
||||||
|
is_partition_key=bool(partition_key == "scalar_3")),
|
||||||
|
FieldSchema(name="scalar_6", dtype=DataType.VARCHAR, max_length=1000,
|
||||||
|
is_partition_key=bool(partition_key == "scalar_6")),
|
||||||
|
FieldSchema(name="scalar_9", dtype=DataType.VARCHAR, max_length=1000,
|
||||||
|
is_partition_key=bool(partition_key == "scalar_9")),
|
||||||
|
FieldSchema(name="scalar_12", dtype=DataType.VARCHAR, max_length=1000,
|
||||||
|
is_partition_key=bool(partition_key == "scalar_12")),
|
||||||
|
FieldSchema(name="scalar_5_linear", dtype=DataType.VARCHAR, max_length=1000,
|
||||||
|
is_partition_key=bool(partition_key == "scalar_5_linear")),
|
||||||
|
FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=768)
|
||||||
|
]
|
||||||
|
schema = CollectionSchema(fields=fields, description="test collection", enable_dynamic_field=True,
|
||||||
|
num_partitions=1)
|
||||||
|
collection = Collection(name=collection_name, schema=schema, num_partitions=1)
|
||||||
|
|
||||||
|
collection.set_properties({"partitionkey.isolation": enable_isolation})
|
||||||
|
log.info(f"collection {collection_name} created: {collection.describe()}")
|
||||||
|
index_params = {"metric_type": "L2", "index_type": "HNSW", "params": {"M": 30, "efConstruction": 360}}
|
||||||
|
log.info(f"collection {collection_name} created")
|
||||||
|
batch_size = 1000
|
||||||
|
data_size = 10000
|
||||||
|
epoch = data_size // batch_size
|
||||||
|
remainder = data_size % batch_size
|
||||||
|
for i in range(epoch + 1):
|
||||||
|
if i == epoch:
|
||||||
|
if remainder == 0:
|
||||||
|
break
|
||||||
|
batch_size = remainder
|
||||||
|
start_idx = i * batch_size
|
||||||
|
end_idx = (i + 1) * batch_size
|
||||||
|
t0 = time.time()
|
||||||
|
data = {
|
||||||
|
"id": [i for i in range(start_idx, end_idx)],
|
||||||
|
"scalar_3": [str(i % 3) for i in range(start_idx, end_idx)],
|
||||||
|
"scalar_6": [str(i % 6) for i in range(start_idx, end_idx)],
|
||||||
|
"scalar_9": [str(i % 9) for i in range(start_idx, end_idx)],
|
||||||
|
"scalar_12": [str(i % 12) for i in range(start_idx, end_idx)],
|
||||||
|
"scalar_5_linear": [str(i % 5) for i in range(start_idx, end_idx)],
|
||||||
|
"emb": [[random.random() for _ in range(768)] for _ in range(batch_size)]
|
||||||
|
}
|
||||||
|
df = pd.DataFrame(data)
|
||||||
|
log.info(f"generate test data {batch_size} cost time {time.time() - t0}")
|
||||||
|
collection.insert(df)
|
||||||
|
collection.compact()
|
||||||
|
collection.wait_for_compaction_completed()
|
||||||
|
t0 = time.time()
|
||||||
|
collection.create_index("emb", index_params=index_params)
|
||||||
|
index_list = utility.list_indexes(collection_name=collection_name)
|
||||||
|
for index_name in index_list:
|
||||||
|
progress = utility.index_building_progress(collection_name=collection_name, index_name=index_name)
|
||||||
|
while progress["pending_index_rows"] > 0:
|
||||||
|
time.sleep(30)
|
||||||
|
progress = utility.index_building_progress(collection_name=collection_name, index_name=index_name)
|
||||||
|
log.info(f"collection {collection_name} index {index_name} progress: {progress}")
|
||||||
|
log.info(f"collection {collection_name} index {index_name} progress: {progress}")
|
||||||
|
tt = time.time() - t0
|
||||||
|
log.info(f"create index cost time {tt}")
|
||||||
|
collection.compact()
|
||||||
|
collection.wait_for_compaction_completed()
|
||||||
|
t0 = time.time()
|
||||||
|
collection.load()
|
||||||
|
log.info(f"load collection cost time {time.time() - t0}")
|
||||||
|
num = collection.num_entities
|
||||||
|
log.info(f"collection {collection_name} loaded, num_entities: {num}")
|
||||||
|
|
||||||
|
invalid_expressions = [
|
||||||
|
"scalar_6 in ['1', '2']",
|
||||||
|
"scalar_6 not in ['1', '2']",
|
||||||
|
"scalar_6 == '1' or scalar_3 == '1'",
|
||||||
|
"scalar_6 != '1'",
|
||||||
|
"scalar_6 > '1'",
|
||||||
|
"'1' < scalar_6 < '3'",
|
||||||
|
"scalar_3 == '1'" # scalar_3 is not partition key
|
||||||
|
]
|
||||||
|
false_result = []
|
||||||
|
for expr in invalid_expressions:
|
||||||
|
try:
|
||||||
|
res = collection.search(
|
||||||
|
data=[[random.random() for _ in range(768)]],
|
||||||
|
anns_field="emb",
|
||||||
|
expr=expr,
|
||||||
|
param={"metric_type": "L2", "params": {"nprobe": 16}},
|
||||||
|
limit=10,
|
||||||
|
output_fields=["scalar_6"]
|
||||||
|
)
|
||||||
|
log.info(f"search with {expr} get res {res}")
|
||||||
|
false_result.append(expr)
|
||||||
|
except Exception as e:
|
||||||
|
log.info(f"search with unsupported expr {expr} get {e}")
|
||||||
|
if len(false_result) > 0:
|
||||||
|
log.info(f"search with unsupported expr {false_result}, but not raise error\n")
|
||||||
|
assert False
|
||||||
|
|
||||||
|
@pytest.mark.tags(CaseLabel.L3)
|
||||||
|
def test_par_key_isolation_without_partition_key(self):
|
||||||
|
# create
|
||||||
|
self._connect()
|
||||||
|
collection_name = cf.gen_unique_str(prefix)
|
||||||
|
partition_key = "None"
|
||||||
|
enable_isolation = "true"
|
||||||
|
if collection_name in list_collections():
|
||||||
|
log.info(f"collection {collection_name} exists, drop it")
|
||||||
|
Collection(name=collection_name).drop()
|
||||||
|
fields = [
|
||||||
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
|
||||||
|
FieldSchema(name="scalar_3", dtype=DataType.VARCHAR, max_length=1000,
|
||||||
|
is_partition_key=bool(partition_key == "scalar_3")),
|
||||||
|
FieldSchema(name="scalar_6", dtype=DataType.VARCHAR, max_length=1000,
|
||||||
|
is_partition_key=bool(partition_key == "scalar_6")),
|
||||||
|
FieldSchema(name="scalar_9", dtype=DataType.VARCHAR, max_length=1000,
|
||||||
|
is_partition_key=bool(partition_key == "scalar_9")),
|
||||||
|
FieldSchema(name="scalar_12", dtype=DataType.VARCHAR, max_length=1000,
|
||||||
|
is_partition_key=bool(partition_key == "scalar_12")),
|
||||||
|
FieldSchema(name="scalar_5_linear", dtype=DataType.VARCHAR, max_length=1000,
|
||||||
|
is_partition_key=bool(partition_key == "scalar_5_linear")),
|
||||||
|
FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=768)
|
||||||
|
]
|
||||||
|
schema = CollectionSchema(fields=fields, description="test collection", enable_dynamic_field=True,
|
||||||
|
num_partitions=1)
|
||||||
|
collection = Collection(name=collection_name, schema=schema)
|
||||||
|
try:
|
||||||
|
collection.set_properties({"partitionkey.isolation": enable_isolation})
|
||||||
|
assert False
|
||||||
|
except Exception as e:
|
||||||
|
log.info(f"set_properties failed without partition key {e}")
|
||||||
|
assert "partition key isolation mode is enabled but no partition key field is set" in str(e)
|
||||||
|
|
||||||
|
@pytest.mark.tags(CaseLabel.L3)
|
||||||
|
def test_set_par_key_isolation_after_vector_indexed(self):
|
||||||
|
# create
|
||||||
|
self._connect()
|
||||||
|
collection_name = cf.gen_unique_str(prefix)
|
||||||
|
partition_key = "scalar_6"
|
||||||
|
enable_isolation = "false"
|
||||||
|
if collection_name in list_collections():
|
||||||
|
log.info(f"collection {collection_name} exists, drop it")
|
||||||
|
Collection(name=collection_name).drop()
|
||||||
|
fields = [
|
||||||
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
|
||||||
|
FieldSchema(name="scalar_3", dtype=DataType.VARCHAR, max_length=1000,
|
||||||
|
is_partition_key=bool(partition_key == "scalar_3")),
|
||||||
|
FieldSchema(name="scalar_6", dtype=DataType.VARCHAR, max_length=1000,
|
||||||
|
is_partition_key=bool(partition_key == "scalar_6")),
|
||||||
|
FieldSchema(name="scalar_9", dtype=DataType.VARCHAR, max_length=1000,
|
||||||
|
is_partition_key=bool(partition_key == "scalar_9")),
|
||||||
|
FieldSchema(name="scalar_12", dtype=DataType.VARCHAR, max_length=1000,
|
||||||
|
is_partition_key=bool(partition_key == "scalar_12")),
|
||||||
|
FieldSchema(name="scalar_5_linear", dtype=DataType.VARCHAR, max_length=1000,
|
||||||
|
is_partition_key=bool(partition_key == "scalar_5_linear")),
|
||||||
|
FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=768)
|
||||||
|
]
|
||||||
|
schema = CollectionSchema(fields=fields, description="test collection", enable_dynamic_field=True,
|
||||||
|
num_partitions=1)
|
||||||
|
collection = Collection(name=collection_name, schema=schema, num_partitions=1)
|
||||||
|
|
||||||
|
collection.set_properties({"partitionkey.isolation": enable_isolation})
|
||||||
|
log.info(f"collection {collection_name} created: {collection.describe()}")
|
||||||
|
index_params = {"metric_type": "L2", "index_type": "HNSW", "params": {"M": 30, "efConstruction": 360}}
|
||||||
|
log.info(f"collection {collection_name} created")
|
||||||
|
batch_size = 1000
|
||||||
|
data_size = 10000
|
||||||
|
epoch = data_size // batch_size
|
||||||
|
remainder = data_size % batch_size
|
||||||
|
for i in range(epoch + 1):
|
||||||
|
if i == epoch:
|
||||||
|
if remainder == 0:
|
||||||
|
break
|
||||||
|
batch_size = remainder
|
||||||
|
start_idx = i * batch_size
|
||||||
|
end_idx = (i + 1) * batch_size
|
||||||
|
t0 = time.time()
|
||||||
|
data = {
|
||||||
|
"id": [i for i in range(start_idx, end_idx)],
|
||||||
|
"scalar_3": [str(i % 3) for i in range(start_idx, end_idx)],
|
||||||
|
"scalar_6": [str(i % 6) for i in range(start_idx, end_idx)],
|
||||||
|
"scalar_9": [str(i % 9) for i in range(start_idx, end_idx)],
|
||||||
|
"scalar_12": [str(i % 12) for i in range(start_idx, end_idx)],
|
||||||
|
"scalar_5_linear": [str(i % 5) for i in range(start_idx, end_idx)],
|
||||||
|
"emb": [[random.random() for _ in range(768)] for _ in range(batch_size)]
|
||||||
|
}
|
||||||
|
df = pd.DataFrame(data)
|
||||||
|
log.info(f"generate test data {batch_size} cost time {time.time() - t0}")
|
||||||
|
collection.insert(df)
|
||||||
|
collection.compact()
|
||||||
|
collection.wait_for_compaction_completed()
|
||||||
|
t0 = time.time()
|
||||||
|
collection.create_index("emb", index_params=index_params)
|
||||||
|
index_list = utility.list_indexes(collection_name=collection_name)
|
||||||
|
for index_name in index_list:
|
||||||
|
progress = utility.index_building_progress(collection_name=collection_name, index_name=index_name)
|
||||||
|
while progress["pending_index_rows"] > 0:
|
||||||
|
time.sleep(30)
|
||||||
|
progress = utility.index_building_progress(collection_name=collection_name, index_name=index_name)
|
||||||
|
log.info(f"collection {collection_name} index {index_name} progress: {progress}")
|
||||||
|
log.info(f"collection {collection_name} index {index_name} progress: {progress}")
|
||||||
|
tt = time.time() - t0
|
||||||
|
log.info(f"create index cost time {tt}")
|
||||||
|
result = True
|
||||||
|
try:
|
||||||
|
collection.set_properties({"partitionkey.isolation": "true"})
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
result = False
|
||||||
|
log.info(f"set_properties after vector indexed {e}")
|
||||||
|
assert result is False
|
||||||
|
collection.drop_index()
|
||||||
|
collection.set_properties({"partitionkey.isolation": "true"})
|
||||||
|
collection.create_index("emb", index_params=index_params)
|
||||||
|
collection.load()
|
||||||
|
res = collection.search(
|
||||||
|
data=[[random.random() for _ in range(768)]],
|
||||||
|
anns_field="emb",
|
||||||
|
expr="scalar_6 == '1' and scalar_3 == '1'",
|
||||||
|
param={"metric_type": "L2", "params": {"nprobe": 16}},
|
||||||
|
limit=10,
|
||||||
|
output_fields=["scalar_6", "scalar_3"]
|
||||||
|
)
|
||||||
|
log.info(f"search res {res}")
|
||||||
|
assert len(res[0]) > 0
|
||||||
Loading…
Reference in New Issue