wehub-resource-sync
498b235461
Build and test / Build and test AMD64 Ubuntu 22.04 (push) Failing after 0s
Publish Builder / amazonlinux2023 (push) Failing after 1s
Build and test / UT for Go (push) Has been skipped
Publish KRTE Images / KRTE (push) Failing after 1s
Build and test / Integration Test (push) Has been skipped
Build and test / Upload Code Coverage (push) Has been skipped
Publish Builder / rockylinux9 (push) Failing after 1s
Publish Builder / ubuntu22.04 (push) Failing after 0s
Publish Builder / ubuntu24.04 (push) Failing after 0s
Publish Gpu Builder / publish-gpu-builder (push) Failing after 1s
Publish Test Images / PyTest (push) Failing after 0s
Build and test / UT for Cpp (push) Has been cancelled
1207 lines
52 KiB
Python
1207 lines
52 KiB
Python
"""
|
|
Base class for CDC sync tests with common utilities.
|
|
"""
|
|
|
|
import logging
|
|
import random
|
|
import string
|
|
import time
|
|
from collections.abc import Callable
|
|
from datetime import datetime
|
|
from typing import Any
|
|
|
|
from pymilvus import DataType, MilvusClient
|
|
|
|
# Configure logging
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="%(asctime)s - %(name)s - %(levelname)s - [%(funcName)s:%(lineno)d] - %(message)s",
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class TestCDCSyncBase:
|
|
"""Base class for CDC sync tests with common utilities."""
|
|
|
|
SEARCH_OVERLAP_THRESHOLD = 0.8
|
|
|
|
FTS_SENTENCES = [
|
|
"Vector databases enable efficient similarity search at scale.",
|
|
"Machine learning models produce high-dimensional embeddings for semantic search.",
|
|
"BM25 is a classic full-text search ranking function used in information retrieval.",
|
|
"Approximate nearest neighbor algorithms trade accuracy for speed.",
|
|
"Milvus supports hybrid search combining dense and sparse vectors.",
|
|
"Transformer models like BERT generate contextual word embeddings.",
|
|
"Inverted indexes power traditional keyword-based search engines.",
|
|
"HNSW graph structure allows fast approximate nearest neighbor queries.",
|
|
"Data replication ensures high availability in distributed databases.",
|
|
"CDC stands for Change Data Capture, tracking database mutations in real time.",
|
|
"Sparse vectors represent text documents in a high-dimensional term space.",
|
|
"Index partitioning improves query throughput in large vector collections.",
|
|
"Cosine similarity measures the angle between two vectors in embedding space.",
|
|
"Product quantization compresses vectors to reduce memory footprint.",
|
|
"Segment-level operations allow incremental updates without full re-indexing.",
|
|
"Full-text search ranks documents by term frequency and inverse document frequency.",
|
|
"Recall measures the fraction of true neighbors found by an approximate search.",
|
|
"Distributed query execution parallelizes workloads across multiple nodes.",
|
|
"Schema evolution allows adding new fields without dropping existing collections.",
|
|
"Hybrid ranking merges keyword and semantic scores for improved relevance.",
|
|
]
|
|
|
|
@staticmethod
|
|
def gen_unique_name(prefix: str = "", length: int = 8, max_length: int = None) -> str:
|
|
"""Generate a unique string with prefix and timestamp."""
|
|
chars = string.ascii_letters + string.digits
|
|
random_str = "".join(random.choice(chars) for _ in range(length))
|
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3] # milliseconds
|
|
|
|
name = f"{prefix}_{timestamp}_{random_str}"
|
|
|
|
# If max_length is specified and name exceeds it, truncate intelligently
|
|
if max_length and len(name) > max_length:
|
|
# Keep the random suffix for uniqueness, truncate prefix and timestamp
|
|
suffix_len = length + 1 # +1 for underscore
|
|
available_len = max_length - suffix_len
|
|
|
|
if available_len > 0:
|
|
# Use shorter timestamp format for space
|
|
short_timestamp = datetime.now().strftime("%m%d_%H%M%S") # 11 chars
|
|
truncated_prefix = prefix[: available_len - len(short_timestamp) - 1] if len(prefix) > 0 else ""
|
|
name = (
|
|
f"{truncated_prefix}_{short_timestamp}_{random_str}"
|
|
if truncated_prefix
|
|
else f"{short_timestamp}_{random_str}"
|
|
)
|
|
else:
|
|
# Fallback: just use random string
|
|
name = random_str
|
|
|
|
return name
|
|
|
|
@staticmethod
|
|
def wait_for_sync(
|
|
check_func: Callable[[], bool],
|
|
timeout: int = 120,
|
|
operation_name: str = "operation",
|
|
) -> bool:
|
|
"""
|
|
Wait for sync operation to complete with progress logging.
|
|
|
|
Args:
|
|
check_func: Function that returns True when sync is complete
|
|
timeout: Timeout in seconds
|
|
operation_name: Name of operation for logging
|
|
|
|
Returns:
|
|
True if sync completed, False if timeout
|
|
"""
|
|
start_time = time.time()
|
|
check_interval = 2
|
|
last_log_time = start_time
|
|
|
|
logger.info(f"Starting sync wait for: {operation_name}")
|
|
|
|
while time.time() - start_time < timeout:
|
|
try:
|
|
if check_func():
|
|
elapsed = time.time() - start_time
|
|
logger.info(f"[SUCCESS] {operation_name} synced successfully in {elapsed:.2f}s")
|
|
return True
|
|
except Exception as e:
|
|
elapsed = time.time() - start_time
|
|
logger.warning(f"Sync check failed for {operation_name} at {elapsed:.1f}s: {e}")
|
|
|
|
elapsed = time.time() - start_time
|
|
# Log every 10 seconds or on first check
|
|
if elapsed - (last_log_time - start_time) >= 10 or elapsed <= check_interval:
|
|
progress = (elapsed / timeout) * 100
|
|
logger.info(f"[WAITING] {operation_name} sync... {elapsed:.1f}s elapsed ({progress:.1f}% of timeout)")
|
|
last_log_time = time.time()
|
|
|
|
time.sleep(check_interval)
|
|
|
|
elapsed = time.time() - start_time
|
|
logger.error(f"[FAILED] {operation_name} sync failed after {elapsed:.2f}s timeout")
|
|
return False
|
|
|
|
@staticmethod
|
|
def create_default_schema(client):
|
|
"""Create default collection schema for testing using MilvusClient API."""
|
|
|
|
# Create schema using MilvusClient API like in the example
|
|
schema = client.create_schema(enable_dynamic_field=True)
|
|
schema.add_field("id", DataType.INT64, is_primary=True)
|
|
schema.add_field("vector", DataType.FLOAT_VECTOR, dim=128)
|
|
|
|
return schema
|
|
|
|
@staticmethod
|
|
def create_comprehensive_schema(client):
|
|
"""Create schema with comprehensive data types for testing (max 4 vector fields)."""
|
|
|
|
schema = client.create_schema(enable_dynamic_field=True)
|
|
|
|
# Primary key
|
|
schema.add_field("id", DataType.INT64, is_primary=True, auto_id=True)
|
|
|
|
# Vector fields (limited to 4 total)
|
|
schema.add_field("float_vector", DataType.FLOAT_VECTOR, dim=128)
|
|
schema.add_field("float16_vector", DataType.FLOAT16_VECTOR, dim=64)
|
|
schema.add_field("binary_vector", DataType.BINARY_VECTOR, dim=128)
|
|
schema.add_field("sparse_vector", DataType.SPARSE_FLOAT_VECTOR)
|
|
|
|
# Scalar fields
|
|
schema.add_field("bool_field", DataType.BOOL)
|
|
schema.add_field("int8_field", DataType.INT8)
|
|
schema.add_field("int16_field", DataType.INT16)
|
|
schema.add_field("int32_field", DataType.INT32)
|
|
schema.add_field("int64_field", DataType.INT64)
|
|
schema.add_field("float_field", DataType.FLOAT)
|
|
schema.add_field("double_field", DataType.DOUBLE)
|
|
schema.add_field("varchar_field", DataType.VARCHAR, max_length=1000)
|
|
|
|
# Array fields
|
|
schema.add_field("bool_array", DataType.ARRAY, element_type=DataType.BOOL, max_capacity=100)
|
|
schema.add_field("int32_array", DataType.ARRAY, element_type=DataType.INT32, max_capacity=100)
|
|
schema.add_field("int64_array", DataType.ARRAY, element_type=DataType.INT64, max_capacity=100)
|
|
schema.add_field("float_array", DataType.ARRAY, element_type=DataType.FLOAT, max_capacity=100)
|
|
schema.add_field(
|
|
"double_array",
|
|
DataType.ARRAY,
|
|
element_type=DataType.DOUBLE,
|
|
max_capacity=100,
|
|
)
|
|
schema.add_field(
|
|
"varchar_array",
|
|
DataType.ARRAY,
|
|
element_type=DataType.VARCHAR,
|
|
max_capacity=100,
|
|
max_length=100,
|
|
)
|
|
|
|
# JSON field
|
|
schema.add_field("json_field", DataType.JSON)
|
|
|
|
return schema
|
|
|
|
@staticmethod
|
|
def create_comprehensive_manual_id_schema(client):
|
|
"""Create comprehensive schema with manual ID for upsert operations (max 4 vector fields)."""
|
|
|
|
schema = client.create_schema(enable_dynamic_field=True)
|
|
|
|
# Primary key - manual ID
|
|
schema.add_field("id", DataType.INT64, is_primary=True, auto_id=False)
|
|
|
|
# Vector fields (limited to 4 total)
|
|
schema.add_field("float_vector", DataType.FLOAT_VECTOR, dim=128)
|
|
schema.add_field("float16_vector", DataType.FLOAT16_VECTOR, dim=64)
|
|
schema.add_field("binary_vector", DataType.BINARY_VECTOR, dim=128)
|
|
schema.add_field("sparse_vector", DataType.SPARSE_FLOAT_VECTOR)
|
|
|
|
# Scalar fields
|
|
schema.add_field("bool_field", DataType.BOOL)
|
|
schema.add_field("int8_field", DataType.INT8)
|
|
schema.add_field("int16_field", DataType.INT16)
|
|
schema.add_field("int32_field", DataType.INT32)
|
|
schema.add_field("int64_field", DataType.INT64)
|
|
schema.add_field("float_field", DataType.FLOAT)
|
|
schema.add_field("double_field", DataType.DOUBLE)
|
|
schema.add_field("varchar_field", DataType.VARCHAR, max_length=1000)
|
|
|
|
# Array fields
|
|
schema.add_field("bool_array", DataType.ARRAY, element_type=DataType.BOOL, max_capacity=100)
|
|
schema.add_field("int32_array", DataType.ARRAY, element_type=DataType.INT32, max_capacity=100)
|
|
schema.add_field("int64_array", DataType.ARRAY, element_type=DataType.INT64, max_capacity=100)
|
|
schema.add_field("float_array", DataType.ARRAY, element_type=DataType.FLOAT, max_capacity=100)
|
|
schema.add_field(
|
|
"double_array",
|
|
DataType.ARRAY,
|
|
element_type=DataType.DOUBLE,
|
|
max_capacity=100,
|
|
)
|
|
schema.add_field(
|
|
"varchar_array",
|
|
DataType.ARRAY,
|
|
element_type=DataType.VARCHAR,
|
|
max_capacity=100,
|
|
max_length=100,
|
|
)
|
|
|
|
# JSON field
|
|
schema.add_field("json_field", DataType.JSON)
|
|
|
|
return schema
|
|
|
|
@staticmethod
|
|
def create_comprehensive_schema_alt(client):
|
|
"""Create alternative comprehensive schema with BFLOAT16_VECTOR and INT8_VECTOR (max 4 vector fields)."""
|
|
|
|
schema = client.create_schema(enable_dynamic_field=True)
|
|
|
|
# Primary key
|
|
schema.add_field("id", DataType.INT64, is_primary=True, auto_id=True)
|
|
|
|
# Vector fields (alternative set - limited to 4 total)
|
|
schema.add_field("float_vector", DataType.FLOAT_VECTOR, dim=128)
|
|
schema.add_field("bfloat16_vector", DataType.BFLOAT16_VECTOR, dim=64)
|
|
schema.add_field("int8_vector", DataType.INT8_VECTOR, dim=128)
|
|
schema.add_field("sparse_vector", DataType.SPARSE_FLOAT_VECTOR)
|
|
|
|
# Scalar fields
|
|
schema.add_field("bool_field", DataType.BOOL)
|
|
schema.add_field("int8_field", DataType.INT8)
|
|
schema.add_field("int16_field", DataType.INT16)
|
|
schema.add_field("int32_field", DataType.INT32)
|
|
schema.add_field("int64_field", DataType.INT64)
|
|
schema.add_field("float_field", DataType.FLOAT)
|
|
schema.add_field("double_field", DataType.DOUBLE)
|
|
schema.add_field("varchar_field", DataType.VARCHAR, max_length=1000)
|
|
|
|
# Array fields
|
|
schema.add_field("bool_array", DataType.ARRAY, element_type=DataType.BOOL, max_capacity=100)
|
|
schema.add_field("int32_array", DataType.ARRAY, element_type=DataType.INT32, max_capacity=100)
|
|
schema.add_field("int64_array", DataType.ARRAY, element_type=DataType.INT64, max_capacity=100)
|
|
schema.add_field("float_array", DataType.ARRAY, element_type=DataType.FLOAT, max_capacity=100)
|
|
schema.add_field(
|
|
"double_array",
|
|
DataType.ARRAY,
|
|
element_type=DataType.DOUBLE,
|
|
max_capacity=100,
|
|
)
|
|
schema.add_field(
|
|
"varchar_array",
|
|
DataType.ARRAY,
|
|
element_type=DataType.VARCHAR,
|
|
max_capacity=100,
|
|
max_length=100,
|
|
)
|
|
|
|
# JSON field
|
|
schema.add_field("json_field", DataType.JSON)
|
|
|
|
return schema
|
|
|
|
@staticmethod
|
|
def create_comprehensive_manual_id_schema_alt(client):
|
|
"""Create alternative comprehensive schema with manual ID, BFLOAT16_VECTOR and INT8_VECTOR (max 4 vector fields)."""
|
|
|
|
schema = client.create_schema(enable_dynamic_field=True)
|
|
|
|
# Primary key - manual ID
|
|
schema.add_field("id", DataType.INT64, is_primary=True, auto_id=False)
|
|
|
|
# Vector fields (alternative set - limited to 4 total)
|
|
schema.add_field("float_vector", DataType.FLOAT_VECTOR, dim=128)
|
|
schema.add_field("bfloat16_vector", DataType.BFLOAT16_VECTOR, dim=64)
|
|
schema.add_field("int8_vector", DataType.INT8_VECTOR, dim=128)
|
|
schema.add_field("sparse_vector", DataType.SPARSE_FLOAT_VECTOR)
|
|
|
|
# Scalar fields
|
|
schema.add_field("bool_field", DataType.BOOL)
|
|
schema.add_field("int8_field", DataType.INT8)
|
|
schema.add_field("int16_field", DataType.INT16)
|
|
schema.add_field("int32_field", DataType.INT32)
|
|
schema.add_field("int64_field", DataType.INT64)
|
|
schema.add_field("float_field", DataType.FLOAT)
|
|
schema.add_field("double_field", DataType.DOUBLE)
|
|
schema.add_field("varchar_field", DataType.VARCHAR, max_length=1000)
|
|
|
|
# Array fields
|
|
schema.add_field("bool_array", DataType.ARRAY, element_type=DataType.BOOL, max_capacity=100)
|
|
schema.add_field("int32_array", DataType.ARRAY, element_type=DataType.INT32, max_capacity=100)
|
|
schema.add_field("int64_array", DataType.ARRAY, element_type=DataType.INT64, max_capacity=100)
|
|
schema.add_field("float_array", DataType.ARRAY, element_type=DataType.FLOAT, max_capacity=100)
|
|
schema.add_field(
|
|
"double_array",
|
|
DataType.ARRAY,
|
|
element_type=DataType.DOUBLE,
|
|
max_capacity=100,
|
|
)
|
|
schema.add_field(
|
|
"varchar_array",
|
|
DataType.ARRAY,
|
|
element_type=DataType.VARCHAR,
|
|
max_capacity=100,
|
|
max_length=100,
|
|
)
|
|
|
|
# JSON field
|
|
schema.add_field("json_field", DataType.JSON)
|
|
|
|
return schema
|
|
|
|
@staticmethod
|
|
def create_manual_id_schema(client):
|
|
"""Create collection schema with manual ID for upsert operations."""
|
|
|
|
# Create schema using MilvusClient API with manual ID
|
|
schema = client.create_schema(enable_dynamic_field=True)
|
|
schema.add_field("id", DataType.INT64, is_primary=True, auto_id=False)
|
|
schema.add_field("vector", DataType.FLOAT_VECTOR, dim=128)
|
|
|
|
return schema
|
|
|
|
@staticmethod
|
|
def generate_test_data(count: int = 100) -> list[dict[str, Any]]:
|
|
"""Generate test data for insert operations."""
|
|
return [
|
|
{
|
|
"id": i,
|
|
"vector": [random.random() for _ in range(128)],
|
|
"text": f"test_text_{i}",
|
|
"number": i,
|
|
"metadata": {"type": "test", "value": i},
|
|
}
|
|
for i in range(count)
|
|
]
|
|
|
|
@staticmethod
|
|
def generate_test_data_with_id(count: int = 100, start_id: int = 0) -> list[dict[str, Any]]:
|
|
"""Generate test data with manual IDs for upsert operations."""
|
|
return [
|
|
{
|
|
"id": start_id + i,
|
|
"vector": [random.random() for _ in range(128)],
|
|
"text": f"test_text_{i}",
|
|
"number": i,
|
|
"metadata": {"type": "test", "value": i},
|
|
}
|
|
for i in range(count)
|
|
]
|
|
|
|
@staticmethod
|
|
def generate_comprehensive_test_data(count: int = 100) -> list[dict[str, Any]]:
|
|
"""Generate comprehensive test data with all data types using standard vector generation."""
|
|
|
|
# Generate vectors using standard method
|
|
float_vectors = TestCDCSyncBase._gen_vectors(count, 128, DataType.FLOAT_VECTOR)
|
|
float16_vectors = TestCDCSyncBase._gen_vectors(count, 64, DataType.FLOAT16_VECTOR)
|
|
binary_vectors = TestCDCSyncBase._gen_vectors(count, 128, DataType.BINARY_VECTOR)
|
|
sparse_vectors = TestCDCSyncBase._gen_vectors(count, 1000, DataType.SPARSE_FLOAT_VECTOR)
|
|
|
|
data = []
|
|
for i in range(count):
|
|
record = {
|
|
# Vector fields (limited to 4 total)
|
|
"float_vector": float_vectors[i],
|
|
"float16_vector": float16_vectors[i],
|
|
"binary_vector": binary_vectors[i],
|
|
"sparse_vector": sparse_vectors[i],
|
|
# Scalar fields
|
|
"bool_field": random.choice([True, False]),
|
|
"int8_field": random.randint(-128, 127),
|
|
"int16_field": random.randint(-32768, 32767),
|
|
"int32_field": random.randint(-2147483648, 2147483647),
|
|
"int64_field": random.randint(-9223372036854775808, 9223372036854775807),
|
|
"float_field": random.uniform(-1000.0, 1000.0),
|
|
"double_field": random.uniform(-1000.0, 1000.0),
|
|
"varchar_field": f"test_varchar_{i}_{random.randint(1000, 9999)}",
|
|
# Array fields
|
|
"bool_array": [random.choice([True, False]) for _ in range(random.randint(1, 10))],
|
|
"int32_array": [random.randint(-100, 100) for _ in range(random.randint(1, 10))],
|
|
"int64_array": [random.randint(-1000, 1000) for _ in range(random.randint(1, 10))],
|
|
"float_array": [random.uniform(-100.0, 100.0) for _ in range(random.randint(1, 10))],
|
|
"double_array": [random.uniform(-100.0, 100.0) for _ in range(random.randint(1, 10))],
|
|
"varchar_array": [f"array_str_{j}" for j in range(random.randint(1, 5))],
|
|
# JSON field
|
|
"json_field": {
|
|
"name": f"item_{i}",
|
|
"value": random.randint(1, 1000),
|
|
"tags": [f"tag_{j}" for j in range(random.randint(1, 3))],
|
|
"metadata": {
|
|
"created": f"2024-01-{random.randint(1, 28):02d}",
|
|
"score": random.uniform(0.0, 100.0),
|
|
},
|
|
},
|
|
}
|
|
data.append(record)
|
|
|
|
return data
|
|
|
|
@staticmethod
|
|
def generate_comprehensive_test_data_alt(count: int = 100) -> list[dict[str, Any]]:
|
|
"""Generate comprehensive test data with alternative vector types (BFLOAT16 + INT8)."""
|
|
|
|
# Generate vectors using standard method - alternative set
|
|
float_vectors = TestCDCSyncBase._gen_vectors(count, 128, DataType.FLOAT_VECTOR)
|
|
bfloat16_vectors = TestCDCSyncBase._gen_vectors(count, 64, DataType.BFLOAT16_VECTOR)
|
|
int8_vectors = TestCDCSyncBase._gen_vectors(count, 128, DataType.INT8_VECTOR)
|
|
sparse_vectors = TestCDCSyncBase._gen_vectors(count, 1000, DataType.SPARSE_FLOAT_VECTOR)
|
|
|
|
data = []
|
|
for i in range(count):
|
|
record = {
|
|
# Vector fields (alternative set - limited to 4 total)
|
|
"float_vector": float_vectors[i],
|
|
"bfloat16_vector": bfloat16_vectors[i],
|
|
"int8_vector": int8_vectors[i],
|
|
"sparse_vector": sparse_vectors[i],
|
|
# Scalar fields
|
|
"bool_field": random.choice([True, False]),
|
|
"int8_field": random.randint(-128, 127),
|
|
"int16_field": random.randint(-32768, 32767),
|
|
"int32_field": random.randint(-2147483648, 2147483647),
|
|
"int64_field": random.randint(-9223372036854775808, 9223372036854775807),
|
|
"float_field": random.uniform(-1000.0, 1000.0),
|
|
"double_field": random.uniform(-1000.0, 1000.0),
|
|
"varchar_field": f"test_varchar_{i}_{random.randint(1000, 9999)}",
|
|
# Array fields
|
|
"bool_array": [random.choice([True, False]) for _ in range(random.randint(1, 10))],
|
|
"int32_array": [random.randint(-100, 100) for _ in range(random.randint(1, 10))],
|
|
"int64_array": [random.randint(-1000, 1000) for _ in range(random.randint(1, 10))],
|
|
"float_array": [random.uniform(-100.0, 100.0) for _ in range(random.randint(1, 10))],
|
|
"double_array": [random.uniform(-100.0, 100.0) for _ in range(random.randint(1, 10))],
|
|
"varchar_array": [f"array_str_{j}" for j in range(random.randint(1, 5))],
|
|
# JSON field
|
|
"json_field": {
|
|
"name": f"item_{i}",
|
|
"value": random.randint(1, 1000),
|
|
"tags": [f"tag_{j}" for j in range(random.randint(1, 3))],
|
|
"metadata": {
|
|
"created": f"2024-01-{random.randint(1, 28):02d}",
|
|
"score": random.uniform(0.0, 100.0),
|
|
},
|
|
},
|
|
}
|
|
data.append(record)
|
|
|
|
return data
|
|
|
|
@staticmethod
|
|
def generate_comprehensive_test_data_alt_with_id(count: int = 100, start_id: int = 0) -> list[dict[str, Any]]:
|
|
"""Generate comprehensive test data with manual IDs and alternative vector types (BFLOAT16 + INT8)."""
|
|
|
|
# Generate vectors using standard method - alternative set
|
|
float_vectors = TestCDCSyncBase._gen_vectors(count, 128, DataType.FLOAT_VECTOR)
|
|
bfloat16_vectors = TestCDCSyncBase._gen_vectors(count, 64, DataType.BFLOAT16_VECTOR)
|
|
int8_vectors = TestCDCSyncBase._gen_vectors(count, 128, DataType.INT8_VECTOR)
|
|
sparse_vectors = TestCDCSyncBase._gen_vectors(count, 1000, DataType.SPARSE_FLOAT_VECTOR)
|
|
|
|
data = []
|
|
for i in range(count):
|
|
record = {
|
|
"id": start_id + i,
|
|
# Vector fields (alternative set - limited to 4 total)
|
|
"float_vector": float_vectors[i],
|
|
"bfloat16_vector": bfloat16_vectors[i],
|
|
"int8_vector": int8_vectors[i],
|
|
"sparse_vector": sparse_vectors[i],
|
|
# Scalar fields
|
|
"bool_field": random.choice([True, False]),
|
|
"int8_field": random.randint(-128, 127),
|
|
"int16_field": random.randint(-32768, 32767),
|
|
"int32_field": random.randint(-2147483648, 2147483647),
|
|
"int64_field": random.randint(-9223372036854775808, 9223372036854775807),
|
|
"float_field": random.uniform(-1000.0, 1000.0),
|
|
"double_field": random.uniform(-1000.0, 1000.0),
|
|
"varchar_field": f"test_varchar_{i}_{random.randint(1000, 9999)}",
|
|
# Array fields
|
|
"bool_array": [random.choice([True, False]) for _ in range(random.randint(1, 10))],
|
|
"int32_array": [random.randint(-100, 100) for _ in range(random.randint(1, 10))],
|
|
"int64_array": [random.randint(-1000, 1000) for _ in range(random.randint(1, 10))],
|
|
"float_array": [random.uniform(-100.0, 100.0) for _ in range(random.randint(1, 10))],
|
|
"double_array": [random.uniform(-100.0, 100.0) for _ in range(random.randint(1, 10))],
|
|
"varchar_array": [f"array_str_{j}" for j in range(random.randint(1, 5))],
|
|
# JSON field
|
|
"json_field": {
|
|
"name": f"item_{i}",
|
|
"value": random.randint(1, 1000),
|
|
"tags": [f"tag_{j}" for j in range(random.randint(1, 3))],
|
|
"metadata": {
|
|
"created": f"2024-01-{random.randint(1, 28):02d}",
|
|
"score": random.uniform(0.0, 100.0),
|
|
},
|
|
},
|
|
}
|
|
data.append(record)
|
|
|
|
return data
|
|
|
|
@staticmethod
|
|
def _gen_vectors(nb, dim, vector_data_type):
|
|
"""Generate vectors using standard Milvus testing methods."""
|
|
import numpy as np
|
|
|
|
vectors = []
|
|
if vector_data_type == DataType.FLOAT_VECTOR:
|
|
vectors = [[random.uniform(-1, 1) for _ in range(dim)] for _ in range(nb)]
|
|
elif vector_data_type == DataType.FLOAT16_VECTOR:
|
|
# Generate float16 vectors like common_func.py
|
|
for _ in range(nb):
|
|
raw_vector = [random.random() for _ in range(dim)]
|
|
fp16_vector = np.array(raw_vector, dtype=np.float16)
|
|
vectors.append(fp16_vector)
|
|
elif vector_data_type == DataType.BFLOAT16_VECTOR:
|
|
# Generate bfloat16 vectors like common_func.py
|
|
from ml_dtypes import bfloat16
|
|
|
|
for _ in range(nb):
|
|
raw_vector = [random.random() for _ in range(dim)]
|
|
bf16_vector = np.array(raw_vector, dtype=bfloat16)
|
|
vectors.append(bf16_vector)
|
|
elif vector_data_type == DataType.INT8_VECTOR:
|
|
# Generate int8 vectors like common_func.py
|
|
for _ in range(nb):
|
|
raw_vector = [random.randint(-128, 127) for _ in range(dim)]
|
|
int8_vector = np.array(raw_vector, dtype=np.int8)
|
|
vectors.append(int8_vector)
|
|
elif vector_data_type == DataType.BINARY_VECTOR:
|
|
# Generate binary vectors (dim bits = dim/8 bytes)
|
|
bytes_per_vector = dim // 8
|
|
for _ in range(nb):
|
|
binary_vec = np.random.randint(0, 256, size=bytes_per_vector, dtype=np.uint8).tobytes()
|
|
vectors.append(binary_vec)
|
|
elif vector_data_type == DataType.SPARSE_FLOAT_VECTOR:
|
|
# Sparse vectors have no fixed dim in schema; dim=0 is passed through
|
|
# from parametrized callers. Use a fixed pool so we always have a
|
|
# valid range for randint(5, ...).
|
|
pool = dim if dim >= 20 else 1000
|
|
for _ in range(nb):
|
|
sparse_indices = random.sample(range(pool), random.randint(5, 20))
|
|
sparse_values = [random.random() for _ in sparse_indices]
|
|
sparse_vector = {idx: val for idx, val in zip(sparse_indices, sparse_values)}
|
|
vectors.append(sparse_vector)
|
|
|
|
return vectors
|
|
|
|
@staticmethod
|
|
def generate_comprehensive_test_data_with_id(count: int = 100, start_id: int = 0) -> list[dict[str, Any]]:
|
|
"""Generate comprehensive test data with manual IDs for upsert operations using standard vector generation."""
|
|
|
|
# Generate vectors using standard method
|
|
float_vectors = TestCDCSyncBase._gen_vectors(count, 128, DataType.FLOAT_VECTOR)
|
|
float16_vectors = TestCDCSyncBase._gen_vectors(count, 64, DataType.FLOAT16_VECTOR)
|
|
binary_vectors = TestCDCSyncBase._gen_vectors(count, 128, DataType.BINARY_VECTOR)
|
|
sparse_vectors = TestCDCSyncBase._gen_vectors(count, 1000, DataType.SPARSE_FLOAT_VECTOR)
|
|
|
|
data = []
|
|
for i in range(count):
|
|
record = {
|
|
"id": start_id + i,
|
|
# Vector fields (limited to 4 total)
|
|
"float_vector": float_vectors[i],
|
|
"float16_vector": float16_vectors[i],
|
|
"binary_vector": binary_vectors[i],
|
|
"sparse_vector": sparse_vectors[i],
|
|
# Scalar fields
|
|
"bool_field": random.choice([True, False]),
|
|
"int8_field": random.randint(-128, 127),
|
|
"int16_field": random.randint(-32768, 32767),
|
|
"int32_field": random.randint(-2147483648, 2147483647),
|
|
"int64_field": random.randint(-9223372036854775808, 9223372036854775807),
|
|
"float_field": random.uniform(-1000.0, 1000.0),
|
|
"double_field": random.uniform(-1000.0, 1000.0),
|
|
"varchar_field": f"test_varchar_{i}_{random.randint(1000, 9999)}",
|
|
# Array fields
|
|
"bool_array": [random.choice([True, False]) for _ in range(random.randint(1, 10))],
|
|
"int32_array": [random.randint(-100, 100) for _ in range(random.randint(1, 10))],
|
|
"int64_array": [random.randint(-1000, 1000) for _ in range(random.randint(1, 10))],
|
|
"float_array": [random.uniform(-100.0, 100.0) for _ in range(random.randint(1, 10))],
|
|
"double_array": [random.uniform(-100.0, 100.0) for _ in range(random.randint(1, 10))],
|
|
"varchar_array": [f"array_str_{j}" for j in range(random.randint(1, 5))],
|
|
# JSON field
|
|
"json_field": {
|
|
"name": f"item_{i}",
|
|
"value": random.randint(1, 1000),
|
|
"tags": [f"tag_{j}" for j in range(random.randint(1, 3))],
|
|
"metadata": {
|
|
"created": f"2024-01-{random.randint(1, 28):02d}",
|
|
"score": random.uniform(0.0, 100.0),
|
|
},
|
|
},
|
|
}
|
|
data.append(record)
|
|
|
|
return data
|
|
|
|
@staticmethod
|
|
def generate_bfloat16_test_data(count: int = 100) -> list[dict[str, Any]]:
|
|
"""Generate test data with BFLOAT16_VECTOR for index testing using standard method."""
|
|
|
|
# Generate bfloat16 vectors using standard method
|
|
bfloat16_vectors = TestCDCSyncBase._gen_vectors(count, 64, DataType.BFLOAT16_VECTOR)
|
|
|
|
data = []
|
|
for i in range(count):
|
|
record = {
|
|
"bfloat16_vector": bfloat16_vectors[i],
|
|
}
|
|
data.append(record)
|
|
|
|
return data
|
|
|
|
@staticmethod
|
|
def generate_int8_test_data(count: int = 100) -> list[dict[str, Any]]:
|
|
"""Generate test data with INT8_VECTOR for index testing using standard method."""
|
|
|
|
# Generate int8 vectors using standard method
|
|
int8_vectors = TestCDCSyncBase._gen_vectors(count, 128, DataType.INT8_VECTOR)
|
|
|
|
data = []
|
|
for i in range(count):
|
|
record = {
|
|
"int8_vector": int8_vectors[i],
|
|
}
|
|
data.append(record)
|
|
|
|
return data
|
|
|
|
def cleanup_database(self, client: MilvusClient, db_name: str):
|
|
"""Clean up database if exists."""
|
|
try:
|
|
if db_name in client.list_databases():
|
|
logger.info(f"[CLEANUP] Cleaning up database: {db_name}")
|
|
client.drop_database(db_name)
|
|
logger.info(f"[SUCCESS] Database {db_name} cleaned up successfully")
|
|
else:
|
|
logger.debug(f"Database {db_name} does not exist, skipping cleanup")
|
|
except Exception as e:
|
|
logger.warning(f"[FAILED] Failed to cleanup database {db_name}: {e}")
|
|
|
|
def cleanup_collection(self, client: MilvusClient, collection_name: str):
|
|
"""Clean up collection if exists."""
|
|
try:
|
|
if client.has_collection(collection_name):
|
|
logger.info(f"[CLEANUP] Cleaning up collection: {collection_name}")
|
|
client.drop_collection(collection_name)
|
|
logger.info(f"[SUCCESS] Collection {collection_name} cleaned up successfully")
|
|
else:
|
|
logger.debug(f"Collection {collection_name} does not exist, skipping cleanup")
|
|
except Exception as e:
|
|
logger.warning(f"[FAILED] Failed to cleanup collection {collection_name}: {e}")
|
|
|
|
def cleanup_user(self, client: MilvusClient, username: str):
|
|
"""Clean up user if exists."""
|
|
try:
|
|
users = client.list_users()
|
|
if username in users:
|
|
logger.info(f"[CLEANUP] Cleaning up user: {username}")
|
|
client.drop_user(username)
|
|
logger.info(f"[SUCCESS] User {username} cleaned up successfully")
|
|
else:
|
|
logger.debug(f"User {username} does not exist, skipping cleanup")
|
|
except Exception as e:
|
|
logger.warning(f"[FAILED] Failed to cleanup user {username}: {e}")
|
|
|
|
def cleanup_role(self, client: MilvusClient, role_name: str):
|
|
"""Clean up role if exists."""
|
|
try:
|
|
roles = client.list_roles()
|
|
if role_name in roles:
|
|
logger.info(f"[CLEANUP] Cleaning up role: {role_name}")
|
|
|
|
# First, revoke all privileges from the role
|
|
try:
|
|
role_privileges = client.describe_role(role_name)
|
|
if isinstance(role_privileges, dict) and "privileges" in role_privileges:
|
|
privileges_list = role_privileges["privileges"]
|
|
elif isinstance(role_privileges, list):
|
|
privileges_list = role_privileges
|
|
else:
|
|
privileges_list = []
|
|
|
|
for privilege_info in privileges_list:
|
|
try:
|
|
client.revoke_privilege(
|
|
role_name=role_name,
|
|
object_type=privilege_info.get("object_type", "Collection"),
|
|
privilege=privilege_info.get("privilege"),
|
|
object_name=privilege_info.get("object_name", "*"),
|
|
)
|
|
logger.debug(
|
|
f"[CLEANUP] Revoked privilege {privilege_info.get('privilege')} from role {role_name}"
|
|
)
|
|
except Exception as revoke_e:
|
|
logger.debug(
|
|
f"[CLEANUP] Failed to revoke privilege {privilege_info.get('privilege')}: {revoke_e}"
|
|
)
|
|
except Exception as describe_e:
|
|
logger.debug(f"[CLEANUP] Failed to describe role privileges: {describe_e}")
|
|
|
|
# Then drop the role
|
|
client.drop_role(role_name)
|
|
logger.info(f"[SUCCESS] Role {role_name} cleaned up successfully")
|
|
else:
|
|
logger.debug(f"Role {role_name} does not exist, skipping cleanup")
|
|
except Exception as e:
|
|
logger.warning(f"[FAILED] Failed to cleanup role {role_name}: {e}")
|
|
|
|
def log_test_start(self, test_name: str, operation_type: str, resource_name: str = ""):
|
|
"""Log test case start with detailed information."""
|
|
separator = "=" * 80
|
|
logger.info(f"\n{separator}")
|
|
logger.info(f"[TEST_START] Starting test: {test_name}")
|
|
logger.info(f"[OPERATION] Operation type: {operation_type}")
|
|
if resource_name:
|
|
logger.info(f"[RESOURCE] Resource: {resource_name}")
|
|
logger.info(f"[TIME] Start time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
|
|
logger.info(f"{separator}")
|
|
|
|
def log_test_end(self, test_name: str, success: bool, duration: float = 0):
|
|
"""Log test case completion with result."""
|
|
separator = "=" * 80
|
|
status = "PASSED" if success else "FAILED"
|
|
logger.info(f"\n{separator}")
|
|
logger.info(f"[TEST_END] Test completed: {test_name}")
|
|
logger.info(f"[RESULT] Result: {status}")
|
|
if duration > 0:
|
|
logger.info(f"[DURATION] Duration: {duration:.2f}s")
|
|
logger.info(f"[TIME] End time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
|
|
logger.info(f"{separator}\n")
|
|
|
|
def log_operation(
|
|
self,
|
|
operation: str,
|
|
resource_type: str,
|
|
resource_name: str,
|
|
client_type: str = "upstream",
|
|
):
|
|
"""Log CDC operation execution."""
|
|
logger.info(f"[EXECUTE] Executing {operation} on {client_type}: {resource_type} '{resource_name}'")
|
|
|
|
def log_sync_verification(self, operation: str, resource_name: str, expected_state: str):
|
|
"""Log sync verification attempt."""
|
|
logger.info(f"[VERIFY] Verifying sync for {operation}: {resource_name} should be {expected_state}")
|
|
|
|
def log_data_operation(self, operation: str, collection_name: str, count: int = 0, details: str = ""):
|
|
"""Log data manipulation operations."""
|
|
if count > 0:
|
|
logger.info(f"[DATA] {operation} operation: {collection_name} - {count} records {details}")
|
|
else:
|
|
logger.info(f"[DATA] {operation} operation: {collection_name} {details}")
|
|
|
|
def log_resource_state(
|
|
self,
|
|
resource_type: str,
|
|
resource_name: str,
|
|
state: str,
|
|
client_type: str,
|
|
details: str = "",
|
|
):
|
|
"""Log current resource state."""
|
|
state_prefix = "[EXISTS]" if state == "exists" else "[MISSING]" if state == "missing" else "[UNKNOWN]"
|
|
logger.info(f"{state_prefix} {client_type.capitalize()} {resource_type} '{resource_name}': {state} {details}")
|
|
|
|
# -------------------------------------------------------------------------
|
|
# Part 1: Additional Schema Factories
|
|
# -------------------------------------------------------------------------
|
|
|
|
@staticmethod
|
|
def create_fts_schema(client, analyzer_type="standard"):
|
|
"""Create schema with BM25 full-text search function."""
|
|
from pymilvus import Function, FunctionType
|
|
|
|
schema = client.create_schema()
|
|
schema.add_field("id", DataType.INT64, is_primary=True, auto_id=True)
|
|
schema.add_field(
|
|
"text_field",
|
|
DataType.VARCHAR,
|
|
max_length=2048,
|
|
enable_analyzer=True,
|
|
analyzer_params={"type": analyzer_type},
|
|
)
|
|
schema.add_field("sparse_output", DataType.SPARSE_FLOAT_VECTOR, is_function_output=True)
|
|
schema.add_field("dense_vector", DataType.FLOAT_VECTOR, dim=128)
|
|
schema.add_field("category", DataType.VARCHAR, max_length=64)
|
|
|
|
bm25_fn = Function(
|
|
name="bm25_fn",
|
|
function_type=FunctionType.BM25,
|
|
input_field_names=["text_field"],
|
|
output_field_names=["sparse_output"],
|
|
)
|
|
schema.add_function(bm25_fn)
|
|
|
|
return schema
|
|
|
|
@staticmethod
|
|
def create_dynamic_schema(client):
|
|
"""Create schema with dynamic field enabled."""
|
|
schema = client.create_schema(enable_dynamic_field=True)
|
|
schema.add_field("id", DataType.INT64, is_primary=True, auto_id=True)
|
|
schema.add_field("float_vector", DataType.FLOAT_VECTOR, dim=128)
|
|
schema.add_field("varchar_field", DataType.VARCHAR, max_length=256)
|
|
return schema
|
|
|
|
@staticmethod
|
|
def create_nullable_schema(client):
|
|
"""Create schema with nullable fields."""
|
|
schema = client.create_schema()
|
|
schema.add_field("id", DataType.INT64, is_primary=True, auto_id=True)
|
|
schema.add_field("float_vector", DataType.FLOAT_VECTOR, dim=128)
|
|
schema.add_field("nullable_int64", DataType.INT64, nullable=True)
|
|
schema.add_field("nullable_varchar", DataType.VARCHAR, max_length=256, nullable=True)
|
|
schema.add_field("nullable_float", DataType.FLOAT, nullable=True)
|
|
schema.add_field("nullable_json", DataType.JSON, nullable=True)
|
|
schema.add_field(
|
|
"nullable_array",
|
|
DataType.ARRAY,
|
|
element_type=DataType.INT32,
|
|
max_capacity=100,
|
|
nullable=True,
|
|
)
|
|
return schema
|
|
|
|
@staticmethod
|
|
def create_default_values_schema(client):
|
|
"""Create schema with default values on scalar fields."""
|
|
schema = client.create_schema()
|
|
schema.add_field("id", DataType.INT64, is_primary=True, auto_id=True)
|
|
schema.add_field("float_vector", DataType.FLOAT_VECTOR, dim=128)
|
|
schema.add_field("default_int64", DataType.INT64, default_value=0)
|
|
schema.add_field("default_varchar", DataType.VARCHAR, max_length=256, default_value="default")
|
|
schema.add_field("default_float", DataType.FLOAT, default_value=1.0)
|
|
schema.add_field("default_bool", DataType.BOOL, default_value=False)
|
|
return schema
|
|
|
|
@staticmethod
|
|
def create_partition_key_schema(client, key_type="VarChar"):
|
|
"""Create schema with a partition key field."""
|
|
schema = client.create_schema()
|
|
schema.add_field("id", DataType.INT64, is_primary=True, auto_id=True)
|
|
schema.add_field("float_vector", DataType.FLOAT_VECTOR, dim=128)
|
|
if key_type == "VarChar":
|
|
schema.add_field(
|
|
"partition_key_field",
|
|
DataType.VARCHAR,
|
|
max_length=256,
|
|
is_partition_key=True,
|
|
)
|
|
else:
|
|
schema.add_field(
|
|
"partition_key_field",
|
|
DataType.INT64,
|
|
is_partition_key=True,
|
|
)
|
|
schema.add_field("data_field", DataType.VARCHAR, max_length=256)
|
|
return schema
|
|
|
|
@staticmethod
|
|
def create_clustering_key_schema(client):
|
|
"""Create schema with a clustering key field."""
|
|
schema = client.create_schema()
|
|
schema.add_field("id", DataType.INT64, is_primary=True, auto_id=True)
|
|
schema.add_field("float_vector", DataType.FLOAT_VECTOR, dim=128)
|
|
schema.add_field("clustering_key_field", DataType.INT64, is_clustering_key=True)
|
|
schema.add_field("data_field", DataType.VARCHAR, max_length=256)
|
|
return schema
|
|
|
|
@staticmethod
|
|
def create_single_vector_schema(client, vector_type="FLOAT_VECTOR", dim=128):
|
|
"""Create a simple single-vector schema for a given vector type."""
|
|
schema = client.create_schema()
|
|
schema.add_field("id", DataType.INT64, is_primary=True, auto_id=True)
|
|
dtype = getattr(DataType, vector_type)
|
|
if vector_type == "SPARSE_FLOAT_VECTOR":
|
|
schema.add_field("vector", dtype)
|
|
else:
|
|
schema.add_field("vector", dtype, dim=dim)
|
|
schema.add_field("int_field", DataType.INT64)
|
|
schema.add_field("varchar_field", DataType.VARCHAR, max_length=256)
|
|
schema.add_field("float_field", DataType.FLOAT)
|
|
return schema
|
|
|
|
# -------------------------------------------------------------------------
|
|
# Part 2: Additional Data Generators
|
|
# -------------------------------------------------------------------------
|
|
|
|
@classmethod
|
|
def generate_fts_data(cls, count=100) -> list[dict[str, Any]]:
|
|
"""Generate data for full-text search (FTS) collections."""
|
|
categories = ["catA", "catB", "catC", "catD"]
|
|
data = []
|
|
for i in range(count):
|
|
base_sentence = cls.FTS_SENTENCES[i % len(cls.FTS_SENTENCES)]
|
|
record = {
|
|
"text_field": f"{base_sentence} unique_{i}_{random.randint(10000, 99999)}",
|
|
"dense_vector": [random.random() for _ in range(128)],
|
|
"category": random.choice(categories),
|
|
}
|
|
data.append(record)
|
|
return data
|
|
|
|
@staticmethod
|
|
def generate_nullable_data(count=100, null_ratio=0.3) -> list[dict[str, Any]]:
|
|
"""Generate data with randomly null-ified nullable fields."""
|
|
data = []
|
|
for _ in range(count):
|
|
record: dict[str, Any] = {
|
|
"float_vector": [random.random() for _ in range(128)],
|
|
"nullable_int64": None if random.random() < null_ratio else random.randint(-1000, 1000),
|
|
"nullable_varchar": None if random.random() < null_ratio else f"varchar_{random.randint(1000, 9999)}",
|
|
"nullable_float": None if random.random() < null_ratio else random.uniform(-100.0, 100.0),
|
|
"nullable_json": None if random.random() < null_ratio else {"key": random.randint(1, 100)},
|
|
"nullable_array": None
|
|
if random.random() < null_ratio
|
|
else [random.randint(0, 100) for _ in range(random.randint(1, 10))],
|
|
}
|
|
data.append(record)
|
|
return data
|
|
|
|
@staticmethod
|
|
def generate_dynamic_data(count=100, extra_fields=None) -> list[dict[str, Any]]:
|
|
"""Generate data for collections with dynamic fields enabled."""
|
|
if extra_fields is None:
|
|
extra_fields = {"extra_int": int, "extra_str": str, "extra_float": float}
|
|
|
|
data = []
|
|
for i in range(count):
|
|
record: dict[str, Any] = {
|
|
"float_vector": [random.random() for _ in range(128)],
|
|
"varchar_field": f"dynamic_varchar_{i}_{random.randint(1000, 9999)}",
|
|
}
|
|
for field_name, field_type in extra_fields.items():
|
|
if field_type is int:
|
|
record[field_name] = random.randint(0, 10000)
|
|
elif field_type is str:
|
|
record[field_name] = f"{field_name}_val_{random.randint(1000, 9999)}"
|
|
elif field_type is float:
|
|
record[field_name] = random.uniform(0.0, 1000.0)
|
|
else:
|
|
record[field_name] = None
|
|
data.append(record)
|
|
return data
|
|
|
|
@classmethod
|
|
def generate_single_vector_data(cls, count=100, vector_type="FLOAT_VECTOR", dim=128) -> list[dict[str, Any]]:
|
|
"""Generate data for a single-vector collection."""
|
|
dtype = getattr(DataType, vector_type)
|
|
vectors = cls._gen_vectors(count, dim, dtype)
|
|
data = []
|
|
for i, vec in enumerate(vectors):
|
|
record = {
|
|
"vector": vec,
|
|
"int_field": random.randint(0, 1000),
|
|
"varchar_field": f"single_vec_{i}_{random.randint(1000, 9999)}",
|
|
"float_field": random.uniform(0.0, 1000.0),
|
|
}
|
|
data.append(record)
|
|
return data
|
|
|
|
# -------------------------------------------------------------------------
|
|
# Part 3: Verification Helpers
|
|
# -------------------------------------------------------------------------
|
|
|
|
@classmethod
|
|
def verify_data_sampling(
|
|
cls,
|
|
upstream_client,
|
|
downstream_client,
|
|
collection_name: str,
|
|
sample_ratio: float = 0.2,
|
|
output_fields: list[str] = None,
|
|
):
|
|
"""
|
|
Sample a fraction of records and compare them field-by-field between upstream
|
|
and downstream.
|
|
|
|
Returns:
|
|
(match_count, mismatch_count, mismatch_details)
|
|
"""
|
|
# Retrieve PKs from upstream. Milvus requires a limit when filter is empty.
|
|
all_up_results = upstream_client.query(
|
|
collection_name=collection_name,
|
|
filter="",
|
|
output_fields=["id"],
|
|
limit=16384,
|
|
)
|
|
all_pks = [r["id"] for r in all_up_results]
|
|
if not all_pks:
|
|
logger.warning(f"[VERIFY] No records found in upstream collection '{collection_name}'")
|
|
return 0, 0, []
|
|
|
|
sample_size = max(1, int(len(all_pks) * sample_ratio))
|
|
sampled_pks = random.sample(all_pks, min(sample_size, len(all_pks)))
|
|
logger.info(f"[VERIFY] Sampling {len(sampled_pks)} / {len(all_pks)} records from '{collection_name}'")
|
|
|
|
match_count = 0
|
|
mismatch_count = 0
|
|
mismatch_details: list[dict[str, Any]] = []
|
|
|
|
for pk in sampled_pks:
|
|
up_rows = upstream_client.query(
|
|
collection_name=collection_name,
|
|
filter=f"id == {pk}",
|
|
output_fields=output_fields or ["*"],
|
|
)
|
|
down_rows = downstream_client.query(
|
|
collection_name=collection_name,
|
|
filter=f"id == {pk}",
|
|
output_fields=output_fields or ["*"],
|
|
)
|
|
|
|
if not up_rows or not down_rows:
|
|
mismatch_count += 1
|
|
mismatch_details.append({"pk": pk, "reason": "missing on one side"})
|
|
continue
|
|
|
|
up_row = up_rows[0]
|
|
down_row = down_rows[0]
|
|
fields_match = True
|
|
field_diffs: dict[str, Any] = {}
|
|
|
|
for field, up_val in up_row.items():
|
|
down_val = down_row.get(field)
|
|
if isinstance(up_val, float) or isinstance(down_val, float):
|
|
try:
|
|
if abs(float(up_val) - float(down_val)) > 1e-6:
|
|
fields_match = False
|
|
field_diffs[field] = {"upstream": up_val, "downstream": down_val}
|
|
except (TypeError, ValueError):
|
|
if up_val != down_val:
|
|
fields_match = False
|
|
field_diffs[field] = {"upstream": up_val, "downstream": down_val}
|
|
else:
|
|
if up_val != down_val:
|
|
fields_match = False
|
|
field_diffs[field] = {"upstream": up_val, "downstream": down_val}
|
|
|
|
if fields_match:
|
|
match_count += 1
|
|
else:
|
|
mismatch_count += 1
|
|
mismatch_details.append({"pk": pk, "field_diffs": field_diffs})
|
|
|
|
logger.info(f"[VERIFY] Sampling result — match: {match_count}, mismatch: {mismatch_count}")
|
|
return match_count, mismatch_count, mismatch_details
|
|
|
|
@classmethod
|
|
def verify_search_consistency(
|
|
cls,
|
|
upstream_client,
|
|
downstream_client,
|
|
collection_name: str,
|
|
vectors: list[Any],
|
|
anns_field: str,
|
|
limit: int = 10,
|
|
metric_type: str = "COSINE",
|
|
):
|
|
"""
|
|
Run search queries on both sides and compute PK overlap ratios.
|
|
|
|
Returns:
|
|
(avg_overlap_ratio, all_upstream_pks, all_downstream_pks)
|
|
"""
|
|
all_upstream_pks: list[list[Any]] = []
|
|
all_downstream_pks: list[list[Any]] = []
|
|
overlap_ratios: list[float] = []
|
|
|
|
search_params = {"metric_type": metric_type}
|
|
|
|
for vec in vectors:
|
|
up_results = upstream_client.search(
|
|
collection_name=collection_name,
|
|
data=[vec],
|
|
anns_field=anns_field,
|
|
search_params=search_params,
|
|
limit=limit,
|
|
output_fields=["id"],
|
|
)
|
|
down_results = downstream_client.search(
|
|
collection_name=collection_name,
|
|
data=[vec],
|
|
anns_field=anns_field,
|
|
search_params=search_params,
|
|
limit=limit,
|
|
output_fields=["id"],
|
|
)
|
|
|
|
up_pks = set(hit["id"] for hit in up_results[0]) if up_results else set()
|
|
down_pks = set(hit["id"] for hit in down_results[0]) if down_results else set()
|
|
|
|
all_upstream_pks.append(list(up_pks))
|
|
all_downstream_pks.append(list(down_pks))
|
|
|
|
union_size = len(up_pks | down_pks)
|
|
if union_size > 0:
|
|
overlap = len(up_pks & down_pks) / union_size
|
|
else:
|
|
overlap = 1.0
|
|
overlap_ratios.append(overlap)
|
|
|
|
avg_overlap = sum(overlap_ratios) / len(overlap_ratios) if overlap_ratios else 0.0
|
|
logger.info(
|
|
f"[VERIFY] Search consistency for '{collection_name}': avg_overlap={avg_overlap:.4f} "
|
|
f"over {len(vectors)} queries"
|
|
)
|
|
return avg_overlap, all_upstream_pks, all_downstream_pks
|
|
|
|
@staticmethod
|
|
def verify_query_consistency(
|
|
upstream_client,
|
|
downstream_client,
|
|
collection_name: str,
|
|
filter_expr: str,
|
|
output_fields: list[str] = None,
|
|
):
|
|
"""
|
|
Query both upstream and downstream with the same filter and compare PK sets.
|
|
|
|
Returns:
|
|
(up_count, down_count, overlap_count, extra_up, extra_down)
|
|
"""
|
|
up_results = upstream_client.query(
|
|
collection_name=collection_name,
|
|
filter=filter_expr,
|
|
output_fields=output_fields or ["id"],
|
|
)
|
|
down_results = downstream_client.query(
|
|
collection_name=collection_name,
|
|
filter=filter_expr,
|
|
output_fields=output_fields or ["id"],
|
|
)
|
|
|
|
up_pks = set(r["id"] for r in up_results)
|
|
down_pks = set(r["id"] for r in down_results)
|
|
overlap = up_pks & down_pks
|
|
extra_up = up_pks - down_pks
|
|
extra_down = down_pks - up_pks
|
|
|
|
logger.info(
|
|
f"[VERIFY] Query consistency for '{collection_name}' filter='{filter_expr}': "
|
|
f"upstream={len(up_pks)}, downstream={len(down_pks)}, "
|
|
f"overlap={len(overlap)}, extra_up={len(extra_up)}, extra_down={len(extra_down)}"
|
|
)
|
|
return len(up_pks), len(down_pks), len(overlap), extra_up, extra_down
|
|
|
|
@staticmethod
|
|
def verify_iterator_consistency(
|
|
upstream_client,
|
|
downstream_client,
|
|
collection_name: str,
|
|
batch_size: int = 100,
|
|
):
|
|
"""
|
|
Iterate through all records on both sides and compare PK sets.
|
|
|
|
Returns:
|
|
(up_count, down_count, match_bool)
|
|
"""
|
|
|
|
def _collect_pks(client):
|
|
pks = set()
|
|
iterator = client.query_iterator(
|
|
collection_name=collection_name,
|
|
batch_size=batch_size,
|
|
output_fields=["id"],
|
|
)
|
|
while True:
|
|
batch = iterator.next()
|
|
if not batch:
|
|
iterator.close()
|
|
break
|
|
for row in batch:
|
|
pks.add(row["id"])
|
|
return pks
|
|
|
|
up_pks = _collect_pks(upstream_client)
|
|
down_pks = _collect_pks(downstream_client)
|
|
match = up_pks == down_pks
|
|
|
|
logger.info(
|
|
f"[VERIFY] Iterator consistency for '{collection_name}': "
|
|
f"upstream={len(up_pks)}, downstream={len(down_pks)}, match={match}"
|
|
)
|
|
return len(up_pks), len(down_pks), match
|