wehub/milvus-io--milvus
1
0
Fork 0
Code Issues Pull Requests Actions 4 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
efb6aee7926e5b239bcef3ae9b52cb542ead9403
milvus-io--milvus/tests/python_client/common/external_table_common.py
T
wehub-resource-sync 498b235461
Build and test / Build and test AMD64 Ubuntu 22.04 (push) Failing after 0s
Details
Publish Builder / amazonlinux2023 (push) Failing after 1s
Details
Build and test / UT for Go (push) Has been skipped
Details
Publish KRTE Images / KRTE (push) Failing after 1s
Details
Build and test / Integration Test (push) Has been skipped
Details
Build and test / Upload Code Coverage (push) Has been skipped
Details
Publish Builder / rockylinux9 (push) Failing after 1s
Details
Publish Builder / ubuntu22.04 (push) Failing after 0s
Details
Publish Builder / ubuntu24.04 (push) Failing after 0s
Details
Publish Gpu Builder / publish-gpu-builder (push) Failing after 1s
Details
Publish Test Images / PyTest (push) Failing after 0s
Details
Build and test / UT for Cpp (push) Has been cancelled
Details
chore: import upstream snapshot with attribution
2026-07-13 12:31:17 +08:00

1267 lines
50 KiB
Python
Raw Blame History

"""Shared helpers for MilvusClient external table tests."""
import io
import json
import os
import tempfile
import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq
from common import common_func as cf
from common import common_type as ct
from minio import Minio
from pymilvus import DataType
# ============================================================
# Constants & Configuration
# ============================================================
REFRESH_TIMEOUT = 180
# Multi-file format coverage: each format E2E writes FORMAT_NUM_FILES sibling
# data files (or fragments / appended snapshots), each carrying
# FORMAT_ROWS_PER_FILE rows, exercising the multi-file refresh path.
FORMAT_NUM_FILES = 3
FORMAT_ROWS_PER_FILE = 3000
FORMAT_TOTAL_ROWS = FORMAT_NUM_FILES * FORMAT_ROWS_PER_FILE
BASIC_FORMATS = ("parquet", "lance-table", "iceberg-table", "vortex")
BASIC_FORMAT_IDS = ("parquet", "lance", "iceberg", "vortex")
def _minio_address(minio_host):
address = minio_host or "localhost"
if address.startswith("http://"):
address = address.removeprefix("http://")
elif address.startswith("https://"):
address = address.removeprefix("https://")
if ":" not in address:
address = f"{address}:9000"
return address
def get_minio_config(minio_host=None, minio_bucket=None):
return {
"address": _minio_address(minio_host),
"access_key": "minioadmin",
"secret_key": "minioadmin",
"bucket": minio_bucket or cf.param_info.param_bucket_name or "milvus-bucket",
"secure": str(minio_host or "").startswith("https://"),
}
def build_external_source(cfg, key_prefix):
"""Build a full s3:// URL that Milvus's extfs resolver can use directly."""
# Trailing slash matters: refresh lists objects under the prefix.
return f"s3://{cfg['address']}/{cfg['bucket']}/{key_prefix}/"
def _minio_endpoint_url(cfg):
scheme = "https" if cfg["secure"] else "http"
return f"{scheme}://{cfg['address']}"
def build_external_spec(cfg=None, fmt="parquet", cloud_provider="minio", **extra):
"""Build an external_spec accepted by current master validation.
Current server-side validation requires external sources to carry a
self-contained extfs block. The tests use Milvus-form S3 URLs
(s3://<endpoint>/<bucket>/<prefix>/), so cloud_provider=minio keeps the
endpoint from the URI instead of deriving a cloud endpoint.
"""
cfg = cfg or get_minio_config()
spec = {
"format": fmt,
"extfs": {
"cloud_provider": cloud_provider,
"region": "us-east-1",
"access_key_id": cfg["access_key"],
"access_key_value": cfg["secret_key"],
"use_ssl": "true" if cfg["secure"] else "false",
},
}
spec.update(extra)
return json.dumps(spec, separators=(",", ":"))
def new_minio_client(cfg):
return Minio(
cfg["address"],
access_key=cfg["access_key"],
secret_key=cfg["secret_key"],
secure=cfg["secure"],
)
# ============================================================
# Parquet Helpers
# ============================================================
_PARQUET_COMPRESSION = "snappy"
def _float_vectors(ids, dim):
return np.array(
[[float(i) * 0.1 + d for d in range(dim)] for i in ids],
dtype=np.float32,
)
def _float16_vectors(ids, dim):
return np.array(
[[float(i) * 0.1 + d for d in range(dim)] for i in ids],
dtype=np.float16,
)
def _bfloat16_vectors(ids, dim):
"""Build bfloat16 via numpy uint16 view; pyarrow lacks native bfloat16 bitcast."""
# bfloat16 shares the top 16 bits of float32's IEEE 754 representation.
f = np.array(
[[float(i) * 0.1 + d for d in range(dim)] for i in ids],
dtype=np.float32,
)
return (f.view(np.uint32) >> 16).astype(np.uint16)
def _int8_vectors(ids, dim):
return np.array(
[[(i + d) % 127 - 63 for d in range(dim)] for i in ids],
dtype=np.int8,
)
def _binary_vectors_bytes(ids, dim):
"""Return (num_rows, dim/8) uint8 array; bit set patterns depend on id."""
nbytes = dim // 8
rng = np.random.default_rng(seed=42)
arr = rng.integers(low=0, high=256, size=(len(ids), nbytes), dtype=np.uint8)
# make it deterministic per id: XOR with id-derived bytes
for row_idx, i in enumerate(ids):
arr[row_idx][0] = i & 0xFF
return arr
def _vector_byte_width(vec_dtype, dim):
if vec_dtype == DataType.FLOAT_VECTOR:
return dim * 4
if vec_dtype in (DataType.FLOAT16_VECTOR, DataType.BFLOAT16_VECTOR):
return dim * 2
if vec_dtype == DataType.INT8_VECTOR:
return dim
if vec_dtype == DataType.BINARY_VECTOR:
return dim // 8
raise ValueError(f"unsupported vec dtype {vec_dtype}")
def _fixed_size_binary_vector_array(raw_rows, byte_width):
return pa.array(
[np.ascontiguousarray(row).tobytes() for row in raw_rows],
type=pa.binary(byte_width),
)
def _fixed_size_uint8_vector_array(raw_rows, byte_width):
raw = b"".join(np.ascontiguousarray(row).tobytes() for row in raw_rows)
return pa.FixedSizeListArray.from_arrays(pa.array(raw, type=pa.uint8()), list_size=byte_width)
def gen_parquet_bytes(num_rows, start_id, scalar_name, arrow_type, value_fn, dim=ct.default_dim):
"""Generate a Parquet file with columns: id, <scalar_name>, embedding (float_vec)."""
ids = list(range(start_id, start_id + num_rows))
vectors = _float_vectors(ids, dim)
table = pa.table(
{
"id": pa.array(ids, type=pa.int64()),
scalar_name: pa.array([value_fn(i) for i in ids], type=arrow_type),
"embedding": pa.FixedSizeListArray.from_arrays(vectors.flatten(), list_size=dim),
}
)
buf = io.BytesIO()
pq.write_table(table, buf, compression=_PARQUET_COMPRESSION)
return buf.getvalue()
def gen_basic_parquet_bytes(num_rows, start_id, dim=ct.default_dim):
"""Generate a Parquet file with columns: id, value (float), embedding."""
return gen_parquet_bytes(
num_rows,
start_id,
"value",
pa.float32(),
lambda i: float(i) * 1.5,
dim=dim,
)
def gen_parquet_bytes_with_codec(num_rows, start_id, codec, dim=ct.default_dim):
"""Like gen_basic_parquet_bytes but with an explicit Parquet compression codec."""
ids = list(range(start_id, start_id + num_rows))
vectors = _float_vectors(ids, dim)
table = pa.table(
{
"id": pa.array(ids, type=pa.int64()),
"value": pa.array([float(i) * 1.5 for i in ids], type=pa.float32()),
"embedding": pa.FixedSizeListArray.from_arrays(vectors.flatten(), list_size=dim),
}
)
buf = io.BytesIO()
pq.write_table(table, buf, compression=codec)
return buf.getvalue()
def gen_array_parquet_bytes(num_rows, start_id, arr_name, arr_element_arrow_type, arr_value_fn, dim=ct.default_dim):
"""Parquet with an Array-typed column: id, <arr_name> (list<element>), embedding."""
ids = list(range(start_id, start_id + num_rows))
vectors = _float_vectors(ids, dim)
table = pa.table(
{
"id": pa.array(ids, type=pa.int64()),
arr_name: pa.array([arr_value_fn(i) for i in ids], type=pa.list_(arr_element_arrow_type)),
"embedding": pa.FixedSizeListArray.from_arrays(vectors.flatten(), list_size=dim),
}
)
buf = io.BytesIO()
pq.write_table(table, buf, compression=_PARQUET_COMPRESSION)
return buf.getvalue()
def gen_dtype_mismatch_parquet_bytes(
num_rows, start_id, mismatch_name, mismatch_array, dim=ct.default_dim, include_embedding=True
):
"""Parquet with a deliberately incompatible user column.
The id column and optional embedding are valid so failures can be
attributed to the requested Milvus type vs external Arrow type mapping.
"""
ids = list(range(start_id, start_id + num_rows))
columns = {
"id": pa.array(ids, type=pa.int64()),
mismatch_name: mismatch_array,
}
if include_embedding:
vectors = _float_vectors(ids, dim)
columns["embedding"] = pa.FixedSizeListArray.from_arrays(
vectors.flatten(),
list_size=dim,
)
table = pa.table(columns)
buf = io.BytesIO()
pq.write_table(table, buf, compression=_PARQUET_COMPRESSION)
return buf.getvalue()
def gen_vector_variant_parquet_bytes(num_rows, start_id, vec_field, vec_dtype, dim):
"""Generate parquet with id + <vec_field> of the given vector family + no scalar."""
ids = list(range(start_id, start_id + num_rows))
if vec_dtype == DataType.FLOAT_VECTOR:
arr = _float_vectors(ids, dim)
parquet_col = pa.FixedSizeListArray.from_arrays(arr.flatten(), list_size=dim)
elif vec_dtype == DataType.FLOAT16_VECTOR:
parquet_col = _fixed_size_binary_vector_array(
_float16_vectors(ids, dim),
_vector_byte_width(vec_dtype, dim),
)
elif vec_dtype == DataType.BFLOAT16_VECTOR:
parquet_col = _fixed_size_binary_vector_array(
_bfloat16_vectors(ids, dim),
_vector_byte_width(vec_dtype, dim),
)
elif vec_dtype == DataType.INT8_VECTOR:
arr = _int8_vectors(ids, dim)
parquet_col = pa.FixedSizeListArray.from_arrays(
pa.array(arr.flatten(), type=pa.int8()),
list_size=dim,
)
elif vec_dtype == DataType.BINARY_VECTOR:
parquet_col = _fixed_size_binary_vector_array(
_binary_vectors_bytes(ids, dim),
_vector_byte_width(vec_dtype, dim),
)
else:
raise ValueError(f"unsupported vec_dtype {vec_dtype}")
table = pa.table(
{
"id": pa.array(ids, type=pa.int64()),
vec_field: parquet_col,
}
)
buf = io.BytesIO()
pq.write_table(table, buf, compression=_PARQUET_COMPRESSION)
return buf.getvalue()
def gen_all_scalar_parquet_bytes(num_rows, start_id, dim=ct.default_dim):
"""Single parquet with every stable scalar type + FloatVector.
Columns: id, val_bool, val_int8, val_int16, val_int32, val_int64,
val_float, val_double, val_varchar, val_json, embedding.
"""
ids = list(range(start_id, start_id + num_rows))
vectors = _float_vectors(ids, dim)
table = pa.table(
{
"id": pa.array(ids, type=pa.int64()),
"val_bool": pa.array([i % 2 == 0 for i in ids], type=pa.bool_()),
"val_int8": pa.array([i % 100 for i in ids], type=pa.int8()),
"val_int16": pa.array([i * 3 for i in ids], type=pa.int16()),
"val_int32": pa.array([i * 7 for i in ids], type=pa.int32()),
"val_int64": pa.array([i * 1000 for i in ids], type=pa.int64()),
"val_float": pa.array([float(i) * 1.5 for i in ids], type=pa.float32()),
"val_double": pa.array([float(i) * 2.5 for i in ids], type=pa.float64()),
"val_varchar": pa.array([f"s_{i:05d}" for i in ids], type=pa.string()),
"val_json": pa.array([json.dumps({"k": i, "g": i % 3}) for i in ids], type=pa.string()),
"embedding": pa.FixedSizeListArray.from_arrays(vectors.flatten(), list_size=dim),
}
)
buf = io.BytesIO()
pq.write_table(table, buf, compression=_PARQUET_COMPRESSION)
return buf.getvalue()
def gen_multi_vector_parquet_bytes(num_rows, start_id, dim=ct.default_dim):
"""Parquet with two vector fields (float + binary) + id + scalar."""
ids = list(range(start_id, start_id + num_rows))
float_arr = _float_vectors(ids, dim)
bin_arr = _binary_vectors_bytes(ids, dim)
table = pa.table(
{
"id": pa.array(ids, type=pa.int64()),
"value": pa.array([float(i) for i in ids], type=pa.float32()),
"dense_vec": pa.FixedSizeListArray.from_arrays(float_arr.flatten(), list_size=dim),
"bin_vec": _fixed_size_binary_vector_array(
bin_arr,
_vector_byte_width(DataType.BINARY_VECTOR, dim),
),
}
)
buf = io.BytesIO()
pq.write_table(table, buf, compression=_PARQUET_COMPRESSION)
return buf.getvalue()
def gen_timestamptz_parquet_bytes(num_rows, start_id, dim=ct.default_dim, ts_unit="us"):
"""Parquet with id + ts (arrow timestamp) + embedding.
Milvus reads Timestamptz as int64 microseconds; the arrow timestamp column
is normalized to int64 by NormalizeExternalArrow on the load path.
"""
ids = list(range(start_id, start_id + num_rows))
vectors = _float_vectors(ids, dim)
base = 1_700_000_000
if ts_unit == "us":
ts_vals = [(base + i) * 1_000_000 for i in ids]
elif ts_unit == "ms":
ts_vals = [(base + i) * 1_000 for i in ids]
else:
ts_vals = [base + i for i in ids]
table = pa.table(
{
"id": pa.array(ids, type=pa.int64()),
"ts": pa.array(ts_vals, type=pa.timestamp(ts_unit)),
"embedding": pa.FixedSizeListArray.from_arrays(vectors.flatten(), list_size=dim),
}
)
buf = io.BytesIO()
pq.write_table(table, buf, compression=_PARQUET_COMPRESSION)
return buf.getvalue()
def gen_geometry_wkt_parquet_bytes(num_rows, start_id, dim=ct.default_dim):
"""Parquet with Geometry stored as WKT strings, matching user datasets."""
ids = list(range(start_id, start_id + num_rows))
vectors = _float_vectors(ids, dim)
table = pa.table(
{
"id": pa.array(ids, type=pa.int64()),
"geo": pa.array([f"POINT({i}.0 {i}.0)" for i in ids], type=pa.string()),
"embedding": pa.FixedSizeListArray.from_arrays(
vectors.flatten(),
list_size=dim,
),
}
)
buf = io.BytesIO()
pq.write_table(table, buf, compression=_PARQUET_COMPRESSION)
return buf.getvalue()
def gen_all_null_columns_parquet_bytes(num_rows, start_id, dim=ct.default_dim):
"""Parquet where every user scalar column is entirely null.
External collections force user fields to nullable=true (schema.go:2350).
A column with no non-null values exercises that path: arrow reads back
ListArray with null_count == num_rows, and Milvus must surface those
nulls without crashing the segment writer.
"""
ids = list(range(start_id, start_id + num_rows))
vectors = _float_vectors(ids, dim)
table = pa.table(
{
"id": pa.array(ids, type=pa.int64()),
"n_int": pa.array([None] * num_rows, type=pa.int32()),
"n_float": pa.array([None] * num_rows, type=pa.float32()),
"n_varchar": pa.array([None] * num_rows, type=pa.string()),
"n_json": pa.array([None] * num_rows, type=pa.string()),
"embedding": pa.FixedSizeListArray.from_arrays(vectors.flatten(), list_size=dim),
}
)
buf = io.BytesIO()
pq.write_table(table, buf, compression=_PARQUET_COMPRESSION)
return buf.getvalue()
def gen_large_parquet_with_row_groups(num_rows, start_id, row_group_size, dim=ct.default_dim):
"""Generate a single large parquet file with explicit row_group_size.
Forces parquet to split internally into multiple row groups so the
Milvus reader exercises the multi-row-group path. The same file may
also be split into multiple Milvus segments by the loon FFI layer
when num_rows exceeds the segment-size threshold.
"""
ids = list(range(start_id, start_id + num_rows))
vectors = _float_vectors(ids, dim)
table = pa.table(
{
"id": pa.array(ids, type=pa.int64()),
"value": pa.array([float(i) * 1.5 for i in ids], type=pa.float32()),
"embedding": pa.FixedSizeListArray.from_arrays(vectors.flatten(), list_size=dim),
}
)
buf = io.BytesIO()
pq.write_table(
table,
buf,
compression=_PARQUET_COMPRESSION,
row_group_size=row_group_size,
)
return buf.getvalue()
# ============================================================
# Full-matrix schema (DataType x Index)
# ============================================================
# A single collection schema covering every supported scalar + vector
# DataType, with same-dtype duplicates wired to different index types.
# Each entry is (milvus_field_name, milvus_DataType, index_type, index_metric,
# index_params, ext_field_name, arrow_writer_type, value_fn,
# add_field_extra). Vector entries treat dim/element_type via
# add_field_extra rather than a value_fn.
#
# The same field list is consumed by both the parquet/lance arrow-table
# generator and the iceberg multi-column writer (vector columns degrade to
# raw binary blobs in iceberg because Iceberg has no fixed-size-list).
FULL_MATRIX_SCALAR_FIELDS = [
# (field_name, DataType, pa type, milvus add_field extras, value_fn)
("b_inv", DataType.BOOL, pa.bool_(), {}, lambda i: i % 2 == 0),
("i8_bmp", DataType.INT8, pa.int8(), {}, lambda i: i % 100),
("i16_inv", DataType.INT16, pa.int16(), {}, lambda i: (i * 3) % 32000),
("i32_stl", DataType.INT32, pa.int32(), {}, lambda i: i * 7),
("i64_inv", DataType.INT64, pa.int64(), {}, lambda i: i * 1000),
("f_inv", DataType.FLOAT, pa.float32(), {}, lambda i: float(i) * 1.5),
("d_inv", DataType.DOUBLE, pa.float64(), {}, lambda i: float(i) * 2.5),
("vc_trie", DataType.VARCHAR, pa.string(), {"max_length": 64}, lambda i: f"s_{i:05d}"),
("txt", DataType.TEXT, pa.string(), {"max_length": 1024}, lambda i: f"text document {i}"),
("j", DataType.JSON, pa.string(), {}, lambda i: json.dumps({"k": i, "g": i % 3})),
("ts", DataType.TIMESTAMPTZ, pa.timestamp("us"), {}, lambda i: (1_700_000_000 + i) * 1_000_000),
]
# Scalar indexes wired per field. None means no scalar index built.
FULL_MATRIX_SCALAR_INDEXES = {
"b_inv": "INVERTED",
"i8_bmp": "BITMAP",
"i16_inv": "INVERTED",
"i32_stl": "STL_SORT",
"i64_inv": "INVERTED",
"f_inv": "INVERTED",
"d_inv": "INVERTED",
"vc_trie": "TRIE",
# txt (TEXT), j (JSON), and ts (Timestamptz) are stored without scalar index.
}
FULL_MATRIX_ARRAY_FIELD = ("arr_int", DataType.INT64, lambda i: [i, i + 1, i + 2])
# A fixed POINT(0 0) WKB blob, sufficient to verify Geometry round-trip.
# WKB: byte order (1=little-endian) | uint32 type(1=POINT) | float64 x | float64 y
_GEOMETRY_WKB = (
b"\x01\x01\x00\x00\x00"
+ b"\x00" * 8 # x = 0.0
+ b"\x00" * 8 # y = 0.0
)
# Full-matrix vectors keep their own fixed production-realistic dim so the
# format-index matrix stays stable even if the common default changes.
FULL_MATRIX_DIM = 128
FULL_MATRIX_BINARY_DIM = 128 # multiple of 8
# Vector fields. Each entry: (name, DataType, dim, index_type, metric, params).
FULL_MATRIX_VECTOR_FIELDS = [
("fv_auto", DataType.FLOAT_VECTOR, FULL_MATRIX_DIM, "AUTOINDEX", "L2", {}),
("fv_flat", DataType.FLOAT_VECTOR, FULL_MATRIX_DIM, "FLAT", "L2", {}),
("fv_hnsw", DataType.FLOAT_VECTOR, FULL_MATRIX_DIM, "HNSW", "L2", {"M": 8, "efConstruction": 64}),
("fv_ivf", DataType.FLOAT_VECTOR, FULL_MATRIX_DIM, "IVF_FLAT", "L2", {"nlist": 8}),
("f16v", DataType.FLOAT16_VECTOR, FULL_MATRIX_DIM, "AUTOINDEX", "L2", {}),
("bf16v", DataType.BFLOAT16_VECTOR, FULL_MATRIX_DIM, "AUTOINDEX", "L2", {}),
("binv_flat", DataType.BINARY_VECTOR, FULL_MATRIX_BINARY_DIM, "BIN_FLAT", "HAMMING", {}),
("binv_ivf", DataType.BINARY_VECTOR, FULL_MATRIX_BINARY_DIM, "BIN_IVF_FLAT", "HAMMING", {"nlist": 8}),
("i8v", DataType.INT8_VECTOR, FULL_MATRIX_DIM, "AUTOINDEX", "L2", {}),
]
def _full_matrix_arrow_columns(
num_rows, start_id, dim=FULL_MATRIX_DIM, bin_dim=FULL_MATRIX_BINARY_DIM, excluded_fields=(), vortex_compatible=False
):
"""Build a dict {column_name -> pyarrow.Array} that fits both parquet and
Lance. Column layout matches FULL_MATRIX_SCALAR_FIELDS / VECTOR_FIELDS;
each milvus field reads from a uniquely-named column.
Pass excluded_fields to omit specific scalar columns (e.g. for vortex,
where the server reader can't sample Arrow String buffers)."""
ids = list(range(start_id, start_id + num_rows))
excluded = set(excluded_fields)
columns = {"id": pa.array(ids, type=pa.int64())}
# Scalars
for name, _dtype, arrow_type, _extra, value_fn in FULL_MATRIX_SCALAR_FIELDS:
if name in excluded:
continue
columns[name] = pa.array([value_fn(i) for i in ids], type=arrow_type)
# Array<Int64>
arr_name, _arr_dtype, arr_value_fn = FULL_MATRIX_ARRAY_FIELD
if arr_name not in excluded:
columns[arr_name] = pa.array([arr_value_fn(i) for i in ids], type=pa.list_(pa.int64()))
# Geometry as binary (WKB).
if "geo" not in excluded:
columns["geo"] = pa.array([_GEOMETRY_WKB] * num_rows, type=pa.binary())
# Vectors. Layout per type:
# FloatVector / Float16Vector / Int8Vector -> FixedSizeList<element, dim>
# BFloat16/BinaryVector -> fixed_size_binary raw bytes
# Vortex 0.56 cannot write FixedSizeBinary; Vortex data uses
# FixedSizeList<UInt8> for BF16 and binary byte-vector payloads.
fv_arr = _float_vectors(ids, dim).flatten()
f16_arr = _float16_vectors(ids, dim)
bf16_arr = _bfloat16_vectors(ids, dim)
i8_arr = _int8_vectors(ids, dim).flatten()
bin_arr = _binary_vectors_bytes(ids, bin_dim)
for name, vtype, vdim, _idx, _metric, _params in FULL_MATRIX_VECTOR_FIELDS:
if name in excluded:
continue
if vtype == DataType.FLOAT_VECTOR:
columns[name] = pa.FixedSizeListArray.from_arrays(
pa.array(fv_arr, type=pa.float32()),
list_size=vdim,
)
elif vtype == DataType.FLOAT16_VECTOR:
byte_width = _vector_byte_width(vtype, vdim)
columns[name] = (
pa.FixedSizeListArray.from_arrays(pa.array(f16_arr.flatten(), type=pa.float16()), list_size=vdim)
if vortex_compatible
else _fixed_size_binary_vector_array(f16_arr, byte_width)
)
elif vtype == DataType.BFLOAT16_VECTOR:
byte_width = _vector_byte_width(vtype, vdim)
columns[name] = (
_fixed_size_uint8_vector_array(bf16_arr, byte_width)
if vortex_compatible
else _fixed_size_binary_vector_array(bf16_arr, byte_width)
)
elif vtype == DataType.INT8_VECTOR:
columns[name] = pa.FixedSizeListArray.from_arrays(
pa.array(i8_arr, type=pa.int8()),
list_size=vdim,
)
elif vtype == DataType.BINARY_VECTOR:
byte_width = _vector_byte_width(vtype, vdim)
columns[name] = (
_fixed_size_uint8_vector_array(bin_arr, byte_width)
if vortex_compatible
else _fixed_size_binary_vector_array(bin_arr, byte_width)
)
else:
raise ValueError(f"unsupported vector dtype {vtype}")
return columns
def gen_full_matrix_parquet_bytes(num_rows, start_id, dim=FULL_MATRIX_DIM, bin_dim=FULL_MATRIX_BINARY_DIM):
"""Single parquet file with every full-matrix field."""
columns = _full_matrix_arrow_columns(num_rows, start_id, dim=dim, bin_dim=bin_dim)
table = pa.table(columns)
buf = io.BytesIO()
pq.write_table(table, buf, compression=_PARQUET_COMPRESSION)
return buf.getvalue()
# ============================================================
# MinIO Helpers
# ============================================================
def upload_parquet(minio_client, bucket, key, data):
minio_client.put_object(
bucket,
key,
io.BytesIO(data),
length=len(data),
content_type="application/octet-stream",
)
def cleanup_minio_prefix(minio_client, bucket, prefix):
for obj in minio_client.list_objects(bucket, prefix=prefix, recursive=True):
minio_client.remove_object(bucket, obj.object_name)
# ============================================================
# Full-Matrix Schema Helpers
# ============================================================
def build_full_matrix_schema(
ops, client, ext_path, ext_spec=None, dim=FULL_MATRIX_DIM, bin_dim=FULL_MATRIX_BINARY_DIM, excluded_fields=()
):
"""Schema covering every supported DataType, with same-dtype duplicates
wired to different index types. See FULL_MATRIX_SCALAR_FIELDS /
FULL_MATRIX_VECTOR_FIELDS for the full layout.
"""
schema = ops.create_schema(client, external_source=ext_path, external_spec=ext_spec or build_external_spec())[0]
ops.add_field(schema, "id", DataType.INT64, external_field="id")
excluded = set(excluded_fields)
for name, dtype, _arrow, extra, _value_fn in FULL_MATRIX_SCALAR_FIELDS:
if name in excluded:
continue
ops.add_field(schema, name, dtype, external_field=name, **extra)
arr_name, arr_elem_dtype, _ = FULL_MATRIX_ARRAY_FIELD
if arr_name not in excluded:
ops.add_field(
schema, arr_name, DataType.ARRAY, element_type=arr_elem_dtype, max_capacity=8, external_field=arr_name
)
if "geo" not in excluded:
ops.add_field(schema, "geo", DataType.GEOMETRY, external_field="geo")
for name, vtype, vdim, _idx, _metric, _params in FULL_MATRIX_VECTOR_FIELDS:
if name in excluded:
continue
ops.add_field(schema, name, vtype, dim=vdim, external_field=name)
return schema
def create_full_matrix_indexes(ops, client, collection_name, excluded_fields=()):
"""Build every scalar + vector index per FULL_MATRIX configuration in
a single create_index call.
"""
index_params = ops.prepare_index_params(client)[0]
excluded = set(excluded_fields)
# Scalar indexes
for field, idx_type in FULL_MATRIX_SCALAR_INDEXES.items():
if field in excluded:
continue
index_params.add_index(
field_name=field,
index_type=idx_type,
)
# Vector indexes
for name, _vtype, _vdim, idx_type, metric, params in FULL_MATRIX_VECTOR_FIELDS:
if name in excluded:
continue
kwargs = {
"field_name": name,
"index_type": idx_type,
"metric_type": metric,
}
if params:
kwargs["params"] = params
index_params.add_index(**kwargs)
ops.create_index(client, collection_name, index_params)
def query_count(client, collection_name):
res = client.query(collection_name, filter="", output_fields=["count(*)"])
return res[0]["count(*)"]
def upload_basic_data(
minio_client, cfg, ext_key, num_rows=ct.default_nb, start_id=0, filename="data.parquet", dim=ct.default_dim
):
"""Upload a basic parquet file and return full minio key."""
key = f"{ext_key}/{filename}"
upload_parquet(
minio_client,
cfg["bucket"],
key,
gen_basic_parquet_bytes(num_rows, start_id, dim=dim),
)
return key
def require_format_dependencies(fmt):
if fmt == "lance-table":
import lance # noqa: F401
elif fmt == "iceberg-table":
from pyiceberg.catalog.sql import SqlCatalog # noqa: F401
elif fmt == "vortex":
import vortex.io as vortex_io # noqa: F401
def _basic_format_arrow_table(num_rows, start_id, dim=ct.default_dim, include_score=False):
ids = list(range(start_id, start_id + num_rows))
vectors = _float_vectors(ids, dim)
columns = {
"id": pa.array(ids, type=pa.int64()),
"value": pa.array([float(i) * 1.5 for i in ids], type=pa.float32()),
"embedding": pa.FixedSizeListArray.from_arrays(vectors.flatten(), list_size=dim),
}
if include_score:
columns["score"] = pa.array([float(i) * 0.01 for i in ids], type=pa.float64())
return pa.table(columns)
def _basic_format_table_to_parquet_bytes(table):
buf = io.BytesIO()
pq.write_table(table, buf, compression=_PARQUET_COMPRESSION)
return buf.getvalue()
def write_basic_format_dataset(
fmt,
minio_client,
cfg,
ext_url,
ext_key,
batches,
dim=ct.default_dim,
include_score=False,
):
"""Write the basic id/value/embedding dataset for one external format.
Returns (external_source, external_spec). Iceberg rewrites external_source
to the metadata.json URI.
"""
require_format_dependencies(fmt)
if fmt == "parquet":
for idx, (start_id, num_rows) in enumerate(batches):
upload_parquet(
minio_client,
cfg["bucket"],
f"{ext_key}/file_{idx:03d}.parquet",
gen_basic_parquet_bytes(num_rows, start_id, dim=dim)
if not include_score
else _basic_format_table_to_parquet_bytes(
_basic_format_arrow_table(num_rows, start_id, dim=dim, include_score=True)
),
)
return ext_url, build_external_spec(cfg, fmt=fmt)
if fmt == "lance-table":
_write_lance_to_minio_batches(
minio_client,
cfg["bucket"],
ext_key,
batches,
dim=dim,
include_score=include_score,
)
return ext_url, build_external_spec(cfg, fmt=fmt)
if fmt == "iceberg-table":
snapshot_id, iceberg_url = _build_iceberg_table_in_minio(
ext_key,
cfg,
batches=batches,
dim=dim,
include_score=include_score,
)
return iceberg_url, build_external_spec(cfg, fmt=fmt, snapshot_id=int(snapshot_id))
if fmt == "vortex":
for idx, (start_id, num_rows) in enumerate(batches):
table = _basic_format_arrow_table(num_rows, start_id, dim=dim, include_score=include_score)
write_vortex_table(minio_client, cfg["bucket"], f"{ext_key}/file_{idx:03d}.vortex", table)
return ext_url, build_external_spec(cfg, fmt=fmt)
raise AssertionError(f"unsupported format: {fmt}")
# ============================================================
# External file format helpers
# ============================================================
def _write_lance_to_minio_batches(minio_client, bucket, key_prefix, batches, dim=ct.default_dim, include_score=False):
"""Multi-fragment variant: each (start_id, count) tuple in `batches` is
appended into the same Lance dataset, producing one fragment file per
batch under data/. Used to exercise multi-data-file refresh.
"""
import shutil
import tempfile
import lance
tmpdir = tempfile.mkdtemp(prefix="ext_lance_")
local_path = os.path.join(tmpdir, "dataset.lance")
try:
for idx, (start_id, num_rows) in enumerate(batches):
table = _basic_format_arrow_table(num_rows, start_id, dim=dim, include_score=include_score)
mode = "create" if idx == 0 else "append"
lance.write_dataset(table, local_path, mode=mode)
for root, _dirs, files in os.walk(local_path):
for fname in files:
absolute = os.path.join(root, fname)
relative = os.path.relpath(absolute, local_path)
minio_client.fput_object(bucket, f"{key_prefix}/{relative}", absolute)
finally:
shutil.rmtree(tmpdir, ignore_errors=True)
def write_vortex_table(minio_client, bucket, key, table):
"""Write a pyarrow table as a Vortex file and upload it to MinIO."""
import vortex.io as vortex_io
with tempfile.NamedTemporaryFile(suffix=".vortex", delete=False) as tmp:
tmp_path = tmp.name
try:
vortex_io.write(table, tmp_path)
minio_client.fput_object(bucket, key, tmp_path, content_type="application/octet-stream")
finally:
try:
os.unlink(tmp_path)
except OSError:
pass
def _build_iceberg_table_in_minio(prefix, cfg, batches, dim=ct.default_dim, include_score=False):
"""Create an Iceberg table whose warehouse lives directly on MinIO.
Writing the warehouse to MinIO (not locally) is required because Iceberg
metadata.json records absolute data-file URIs at write time; if we wrote
to file:// then uploaded, those URIs would be unreachable from the
Milvus reader.
Vector encoding: each row stores its vector as a fixed-size binary blob of
dim * sizeof(float) bytes. Milvus maps external_field=embedding to a
FloatVector field and reinterprets the raw bytes as float32.
Note: list<float32> looks more natural but pyarrow >=15 promotes long
lists to large_list, which the vector normalizer rejects
(Util.cpp:1761). The binary encoding sidesteps that.
`batches` is a list of (start_id, count) tuples. Each batch is appended
to the same iceberg table as a separate parquet data file under data/.
The final snapshot includes every batch.
Returns (snapshot_id, ext_url) where ext_url is the s3:// URL of the
table's metadata.json in Milvus URI form.
"""
import shutil
import struct
import tempfile
from pyiceberg.catalog.sql import SqlCatalog
from pyiceberg.schema import Schema as IbgSchema
from pyiceberg.types import DoubleType, FixedType, FloatType, LongType, NestedField
tmp = tempfile.mkdtemp(prefix="ibg_cat_")
try:
cat = SqlCatalog(
"milvus_test",
**{
"uri": f"sqlite:///{tmp}/cat.db",
"warehouse": f"s3://{cfg['bucket']}/{prefix}",
"s3.endpoint": _minio_endpoint_url(cfg),
"s3.access-key-id": cfg["access_key"],
"s3.secret-access-key": cfg["secret_key"],
"s3.region": "us-east-1",
"s3.path-style-access": "true",
},
)
cat.create_namespace("ext")
ibg_fields = [
NestedField(1, "id", LongType(), required=False),
NestedField(2, "value", FloatType(), required=False),
NestedField(3, "embedding", FixedType(_vector_byte_width(DataType.FLOAT_VECTOR, dim)), required=False),
]
arrow_fields = [
pa.field("id", pa.int64(), nullable=True),
pa.field("value", pa.float32(), nullable=True),
pa.field("embedding", pa.binary(_vector_byte_width(DataType.FLOAT_VECTOR, dim)), nullable=True),
]
if include_score:
ibg_fields.append(NestedField(4, "score", DoubleType(), required=False))
arrow_fields.append(pa.field("score", pa.float64(), nullable=True))
ibg_schema = IbgSchema(*ibg_fields)
arrow_schema = pa.schema(arrow_fields)
tbl = cat.create_table("ext.t", schema=ibg_schema)
for start_id, num_rows in batches:
ids = list(range(start_id, start_id + num_rows))
vec_bytes = [struct.pack(f"<{dim}f", *[float(i) * 0.1 + j for j in range(dim)]) for i in ids]
columns = {
"id": pa.array(ids, type=pa.int64()),
"value": pa.array([float(i) * 1.5 for i in ids], type=pa.float32()),
"embedding": pa.array(vec_bytes, type=pa.binary(_vector_byte_width(DataType.FLOAT_VECTOR, dim))),
}
if include_score:
columns["score"] = pa.array([float(i) * 0.01 for i in ids], type=pa.float64())
arrow_table = pa.table(columns, schema=arrow_schema)
tbl.append(arrow_table)
snap = tbl.current_snapshot().snapshot_id
# metadata_location: s3://<bucket>/<prefix>/.../metadata/000NN-uuid.metadata.json
# -> Milvus form: s3://<address>/<bucket>/...
meta_loc = tbl.metadata_location
if not meta_loc.startswith("s3://"):
raise RuntimeError(f"unexpected metadata_loc scheme: {meta_loc}")
bucket_and_key = meta_loc[len("s3://") :]
ext_url = f"s3://{cfg['address']}/{bucket_and_key}"
return snap, ext_url
finally:
shutil.rmtree(tmp, ignore_errors=True)
# Iceberg full-matrix is a strict subset of FULL_MATRIX_* because Iceberg V2
# only supports int (32-bit) and long (64-bit) integers; no Int8/Int16. We
# reuse the same field names where possible to keep query expectations
# identical between formats; Int8 -> IntegerType promoted to Milvus INT32
# rather than INT8, otherwise NormalizeExternalArrow refuses the cast.
ICEBERG_FULL_MATRIX_SCALAR_FIELDS = [
# (field_name, milvus DataType, iceberg type ctor (lambda field_id: type),
# arrow type, milvus add_field extras, value_fn)
("b_inv", DataType.BOOL, lambda fid: ("BooleanType",), pa.bool_(), {}, lambda i: i % 2 == 0),
("i32_stl", DataType.INT32, lambda fid: ("IntegerType",), pa.int32(), {}, lambda i: i * 7),
("i64_inv", DataType.INT64, lambda fid: ("LongType",), pa.int64(), {}, lambda i: i * 1000),
("f_inv", DataType.FLOAT, lambda fid: ("FloatType",), pa.float32(), {}, lambda i: float(i) * 1.5),
("d_inv", DataType.DOUBLE, lambda fid: ("DoubleType",), pa.float64(), {}, lambda i: float(i) * 2.5),
("vc_trie", DataType.VARCHAR, lambda fid: ("StringType",), pa.string(), {"max_length": 64}, lambda i: f"s_{i:05d}"),
("j", DataType.JSON, lambda fid: ("StringType",), pa.string(), {}, lambda i: json.dumps({"k": i, "g": i % 3})),
(
"ts",
DataType.TIMESTAMPTZ,
lambda fid: ("TimestamptzType",),
pa.timestamp("us", tz="UTC"),
{},
lambda i: (1_700_000_000 + i) * 1_000_000,
),
]
ICEBERG_FULL_MATRIX_SCALAR_INDEXES = {
"b_inv": "INVERTED",
"i32_stl": "STL_SORT",
"i64_inv": "INVERTED",
"f_inv": "INVERTED",
"d_inv": "INVERTED",
"vc_trie": "TRIE",
}
def _build_iceberg_full_matrix_table(prefix, cfg, num_rows, dim=FULL_MATRIX_DIM, bin_dim=FULL_MATRIX_BINARY_DIM):
"""Iceberg variant of the full-matrix dataset.
Vectors are stored as raw binary blobs because Iceberg has no
fixed-size-list. JSON / Geometry use String / Binary respectively.
Returns (snapshot_id, ext_url).
"""
import shutil
import struct
import tempfile
from pyiceberg.catalog.sql import SqlCatalog
from pyiceberg.schema import Schema as IbgSchema
from pyiceberg.types import (
BinaryType,
BooleanType,
DoubleType,
FixedType,
FloatType,
IntegerType,
ListType,
LongType,
NestedField,
StringType,
TimestamptzType,
)
iceberg_type_lookup = {
"BooleanType": BooleanType(),
"IntegerType": IntegerType(),
"LongType": LongType(),
"FloatType": FloatType(),
"DoubleType": DoubleType(),
"StringType": StringType(),
"TimestamptzType": TimestamptzType(),
}
ids = list(range(num_rows))
next_field_id = [3] # field-id assignment, id=1, schema starts at 2 then bumps
def fid():
x = next_field_id[0]
next_field_id[0] += 1
return x
nested_fields = [NestedField(1, "id", LongType(), required=False)]
arrow_fields = [pa.field("id", pa.int64(), nullable=True)]
arrow_columns = {"id": pa.array(ids, type=pa.int64())}
# Scalars
for name, _dtype, ibg_factory, arrow_type, _extra, value_fn in ICEBERG_FULL_MATRIX_SCALAR_FIELDS:
ibg_type = iceberg_type_lookup[ibg_factory(0)[0]]
nested_fields.append(NestedField(fid(), name, ibg_type, required=False))
arrow_fields.append(pa.field(name, arrow_type, nullable=True))
arrow_columns[name] = pa.array([value_fn(i) for i in ids], type=arrow_type)
# Array<Int64>: ListType<LongType>
arr_name, _arr_dtype, arr_value_fn = FULL_MATRIX_ARRAY_FIELD
nested_fields.append(
NestedField(
fid(),
arr_name,
ListType(element_id=fid(), element_type=LongType(), element_required=False),
required=False,
),
)
arrow_fields.append(
pa.field(arr_name, pa.list_(pa.field("element", pa.int64(), nullable=True)), nullable=True),
)
arrow_columns[arr_name] = pa.array(
[arr_value_fn(i) for i in ids],
type=pa.list_(pa.field("element", pa.int64(), nullable=True)),
)
# Geometry -> BinaryType (WKB)
nested_fields.append(NestedField(fid(), "geo", BinaryType(), required=False))
arrow_fields.append(pa.field("geo", pa.binary(), nullable=True))
arrow_columns["geo"] = pa.array([_GEOMETRY_WKB] * num_rows, type=pa.binary())
# Vector fields -> FixedType blobs of dim*sizeof(elem) bytes.
fv_arr = _float_vectors(ids, dim)
f16_arr = _float16_vectors(ids, dim).view(np.uint16)
bf16_arr = _bfloat16_vectors(ids, dim) # already raw uint16-equivalent
i8_arr = _int8_vectors(ids, dim)
bin_arr = _binary_vectors_bytes(ids, bin_dim)
def vec_blobs(name, vtype, vdim):
if vtype == DataType.FLOAT_VECTOR:
return [struct.pack(f"<{vdim}f", *fv_arr[i].tolist()) for i in range(num_rows)]
if vtype == DataType.FLOAT16_VECTOR:
return [f16_arr[i].astype(np.uint16).tobytes() for i in range(num_rows)]
if vtype == DataType.BFLOAT16_VECTOR:
return [bf16_arr[i].tobytes() for i in range(num_rows)]
if vtype == DataType.INT8_VECTOR:
return [i8_arr[i].tobytes() for i in range(num_rows)]
if vtype == DataType.BINARY_VECTOR:
return [bin_arr[i].tobytes() for i in range(num_rows)]
raise ValueError(f"unsupported vec type {vtype}")
for name, vtype, vdim, _idx, _metric, _params in FULL_MATRIX_VECTOR_FIELDS:
byte_width = _vector_byte_width(vtype, vdim)
nested_fields.append(NestedField(fid(), name, FixedType(byte_width), required=False))
arrow_fields.append(pa.field(name, pa.binary(byte_width), nullable=True))
arrow_columns[name] = pa.array(vec_blobs(name, vtype, vdim), type=pa.binary(byte_width))
ibg_schema = IbgSchema(*nested_fields)
arrow_schema = pa.schema(arrow_fields)
arrow_table = pa.table(arrow_columns, schema=arrow_schema)
tmp = tempfile.mkdtemp(prefix="ibg_fm_")
try:
cat = SqlCatalog(
"milvus_test_fm",
**{
"uri": f"sqlite:///{tmp}/cat.db",
"warehouse": f"s3://{cfg['bucket']}/{prefix}",
"s3.endpoint": _minio_endpoint_url(cfg),
"s3.access-key-id": cfg["access_key"],
"s3.secret-access-key": cfg["secret_key"],
"s3.region": "us-east-1",
"s3.path-style-access": "true",
},
)
cat.create_namespace("ext")
tbl = cat.create_table("ext.t", schema=ibg_schema)
tbl.append(arrow_table)
snap = tbl.current_snapshot().snapshot_id
meta_loc = tbl.metadata_location
if not meta_loc.startswith("s3://"):
raise RuntimeError(f"unexpected metadata_loc scheme: {meta_loc}")
bucket_and_key = meta_loc[len("s3://") :]
ext_url = f"s3://{cfg['address']}/{bucket_and_key}"
return snap, ext_url
finally:
shutil.rmtree(tmp, ignore_errors=True)
def build_iceberg_full_matrix_schema(
ops, client, ext_path, ext_spec, dim=FULL_MATRIX_DIM, bin_dim=FULL_MATRIX_BINARY_DIM
):
"""Milvus schema for the iceberg full-matrix dataset. Matches
ICEBERG_FULL_MATRIX_SCALAR_FIELDS for scalars and FULL_MATRIX_VECTOR_FIELDS
for vectors (vectors are stored as iceberg binary blobs but exposed as
Milvus vector types; NormalizeExternalArrow reinterprets the bytes)."""
schema = ops.create_schema(client, external_source=ext_path, external_spec=ext_spec or build_external_spec())[0]
ops.add_field(schema, "id", DataType.INT64, external_field="id")
for name, dtype, _ibg, _arrow, extra, _value_fn in ICEBERG_FULL_MATRIX_SCALAR_FIELDS:
ops.add_field(schema, name, dtype, external_field=name, **extra)
arr_name, arr_elem_dtype, _ = FULL_MATRIX_ARRAY_FIELD
ops.add_field(
schema, arr_name, DataType.ARRAY, element_type=arr_elem_dtype, max_capacity=8, external_field=arr_name
)
ops.add_field(schema, "geo", DataType.GEOMETRY, external_field="geo")
for name, vtype, vdim, _idx, _metric, _params in FULL_MATRIX_VECTOR_FIELDS:
ops.add_field(schema, name, vtype, dim=vdim, external_field=name)
return schema
def create_iceberg_full_matrix_indexes(ops, client, collection_name):
"""Wire the iceberg-subset scalar indexes + the same 9 vector indexes
used by parquet/lance full-matrix."""
index_params = ops.prepare_index_params(client)[0]
for field, idx_type in ICEBERG_FULL_MATRIX_SCALAR_INDEXES.items():
index_params.add_index(field_name=field, index_type=idx_type)
for name, _vtype, _vdim, idx_type, metric, params in FULL_MATRIX_VECTOR_FIELDS:
kwargs = {"field_name": name, "index_type": idx_type, "metric_type": metric}
if params:
kwargs["params"] = params
index_params.add_index(**kwargs)
ops.create_index(client, collection_name, index_params)
def _iceberg_full_matrix_assert(ops, client, coll, expected_count):
"""Iceberg-specific subset of _full_matrix_assert_basic: scalar set is
smaller (no Int8/Int16/Bitmap)."""
assert ops.query_count(client, coll) == expected_count
output_fields = ["id"] + [f for f, _, _, _, _, _ in ICEBERG_FULL_MATRIX_SCALAR_FIELDS]
output_fields += [FULL_MATRIX_ARRAY_FIELD[0], "geo"]
rows = ops.query(client, coll, filter="id == 42", output_fields=output_fields)[0]
assert len(rows) == 1, f"id=42 missing: {rows}"
r = rows[0]
assert r["id"] == 42
for name, _dtype, _ibg, _arrow, _extra, value_fn in ICEBERG_FULL_MATRIX_SCALAR_FIELDS:
expected = value_fn(42)
if name == "j":
parsed = r["j"] if isinstance(r["j"], dict) else json.loads(r["j"])
assert parsed.get("k") == 42, f"json[k]: {parsed}"
elif name in ("f_inv", "d_inv"):
assert abs(r[name] - expected) < 1e-3, f"{name}: {r[name]}"
elif name == "ts":
assert isinstance(r[name], str) and r[name], "ts empty"
else:
assert r[name] == expected, f"{name}: {r[name]} != {expected}"
arr_name, _, arr_value_fn = FULL_MATRIX_ARRAY_FIELD
assert r[arr_name] == arr_value_fn(42), f"{arr_name}: {r[arr_name]}"
for name, _vtype, vdim, _idx, metric, _params in FULL_MATRIX_VECTOR_FIELDS:
query_vec = _full_matrix_query_vec(name, vdim)
hits = ops.search(
client,
coll,
data=query_vec,
limit=3,
anns_field=name,
output_fields=["id"],
search_params={"metric_type": metric},
)[0][0]
assert len(hits) == 3, f"[{name}] search returned {len(hits)} hits"
top_hit = hits[0]
assert top_hit["id"] == 0, f"[{name}] expected exact vector id=0 as top1, got {hits}"
assert abs(top_hit["distance"]) < 1e-5, f"[{name}] exact vector distance mismatch: {top_hit}"
# Smaller row count keeps full-matrix tests under ~2 min each: each test
# builds 21 fields with 10+ scalar indexes and 9 vector indexes.
FULL_MATRIX_NB = 200
def _full_matrix_query_vec(vec_field, dim):
"""Build a search query vector matching the dtype of the given vector
field and exactly matching row id=0. pymilvus is dtype-strict on the
wire: int8 needs np.int8, bf16 needs raw bytes, etc."""
for name, vtype, vdim, _idx, _metric, _params in FULL_MATRIX_VECTOR_FIELDS:
if name != vec_field:
continue
if vtype == DataType.BINARY_VECTOR:
return [_binary_vectors_bytes([0], vdim)[0].tobytes()]
if vtype == DataType.INT8_VECTOR:
return [_int8_vectors([0], vdim)[0]]
if vtype == DataType.FLOAT16_VECTOR:
return [_float16_vectors([0], vdim)[0]]
if vtype == DataType.BFLOAT16_VECTOR:
return [_bfloat16_vectors([0], vdim)[0].tobytes()]
return _float_vectors([0], vdim).tolist()
raise KeyError(vec_field)
def _full_matrix_assert_basic(ops, client, coll, expected_count, excluded_fields=()):
"""Shared assertions: count + per-type spot-checks + search hits per
vector field. Used by every full-matrix format test."""
assert ops.query_count(client, coll) == expected_count
excluded = set(excluded_fields)
# Spot-check id=42: every scalar field round-trips with the expected
# value derived from the same value_fn used to build the parquet.
output_fields = ["id"] + [f for f, _, _, _, _ in FULL_MATRIX_SCALAR_FIELDS if f not in excluded]
arr_name = FULL_MATRIX_ARRAY_FIELD[0]
if arr_name not in excluded:
output_fields.append(arr_name)
if "geo" not in excluded:
output_fields.append("geo")
rows = ops.query(client, coll, filter="id == 42", output_fields=output_fields)[0]
assert len(rows) == 1, f"id=42 missing: {rows}"
r = rows[0]
assert r["id"] == 42
for name, _dtype, _arrow, _extra, value_fn in FULL_MATRIX_SCALAR_FIELDS:
if name in excluded:
continue
expected = value_fn(42)
if name == "j":
parsed = r["j"] if isinstance(r["j"], dict) else json.loads(r["j"])
assert parsed.get("k") == 42, f"json[k] mismatch: {parsed}"
elif name == "f_inv":
assert abs(r[name] - expected) < 1e-3, f"{name}: {r[name]} != {expected}"
elif name == "d_inv":
assert abs(r[name] - expected) < 1e-3, f"{name}: {r[name]} != {expected}"
elif name == "ts":
# Server returns an ISO-8601 string; we just assert non-empty.
assert isinstance(r[name], str) and r[name], f"ts empty: {r[name]}"
else:
assert r[name] == expected, f"{name}: {r[name]} != {expected}"
arr_name, _, arr_value_fn = FULL_MATRIX_ARRAY_FIELD
if arr_name not in excluded:
assert r[arr_name] == arr_value_fn(42), f"{arr_name}: {r[arr_name]}"
# One search per vector field confirms every (type x index x metric)
# combination produced a usable index.
for name, _vtype, vdim, _idx, metric, _params in FULL_MATRIX_VECTOR_FIELDS:
if name in excluded:
continue
query_vec = _full_matrix_query_vec(name, vdim)
hits = ops.search(
client,
coll,
data=query_vec,
limit=3,
anns_field=name,
output_fields=["id"],
search_params={"metric_type": metric},
)[0][0]
assert len(hits) == 3, f"[{name}] search returned {len(hits)} hits"
top_hit = hits[0]
assert top_hit["id"] == 0, f"[{name}] expected exact vector id=0 as top1, got {hits}"
assert abs(top_hit["distance"]) < 1e-5, f"[{name}] exact vector distance mismatch: {top_hit}"
Reference in New Issue View Git Blame Copy Permalink