wehub/alibaba--zvec
1
0
Fork 0
Code Issues Pull Requests Actions 2 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
be3ef883e104e6e259d681b996bf142f5921c12f
alibaba--zvec/python/tests/detail/doc_helper.py
T
wehub-resource-sync be3ef883e1 chore: import upstream snapshot with attribution
2026-07-13 12:47:42 +08:00

465 lines
20 KiB
Python
Raw Blame History

from zvec import CollectionSchema, Doc
from support_helper import *
import numpy as np
from typing import Literal, Optional, Union, Tuple
import random
import string
import math
def generate_constant_vector(
i: int, dimension: int, dtype: Literal["int8", "float16", "float32"] = "float32"
):
if dtype == "int8":
vec = [(i % 127)] * dimension
vec[i % dimension] = (i + 1) % 127
else:
base_val = (i % 1000) / 256.0
special_val = ((i + 1) % 1000) / 256.0
vec = [base_val] * dimension
vec[i % dimension] = special_val
return vec
def generate_constant_vector_recall(
i: int, dimension: int, dtype: Literal["int8", "float16", "float32"] = "float32"
):
if dtype == "int8":
vec = [(i % 127)] * dimension
vec[i % dimension] = (i + 1) % 127
else:
base_val = math.sin((i) * 1000) / 256.0
special_val = math.sin((i + 1) * 1000) / 256.0
vec = [base_val] * dimension
vec[i % dimension] = special_val
return vec
def generate_sparse_vector(i: int):
return {i: i + 0.1}
def generate_vectordict(i: int, schema: CollectionSchema) -> Doc:
doc_fields = {}
doc_vectors = {}
doc_fields = {}
doc_vectors = {}
for field in schema.fields:
if field.data_type == DataType.BOOL:
doc_fields[field.name] = i % 2 == 0
elif field.data_type == DataType.INT32:
doc_fields[field.name] = i
elif field.data_type == DataType.UINT32:
doc_fields[field.name] = i
elif field.data_type == DataType.INT64:
doc_fields[field.name] = i
elif field.data_type == DataType.UINT64:
doc_fields[field.name] = i
elif field.data_type == DataType.FLOAT:
doc_fields[field.name] = float(i) + 0.1
elif field.data_type == DataType.DOUBLE:
doc_fields[field.name] = float(i) + 0.11
elif field.data_type == DataType.STRING:
doc_fields[field.name] = f"test_{i}"
elif field.data_type == DataType.ARRAY_BOOL:
doc_fields[field.name] = [i % 2 == 0, i % 3 == 0]
elif field.data_type == DataType.ARRAY_INT32:
doc_fields[field.name] = [i, i + 1, i + 2]
elif field.data_type == DataType.ARRAY_UINT32:
doc_fields[field.name] = [i, i + 1, i + 2]
elif field.data_type == DataType.ARRAY_INT64:
doc_fields[field.name] = [i, i + 1, i + 2]
elif field.data_type == DataType.ARRAY_UINT64:
doc_fields[field.name] = [i, i + 1, i + 2]
elif field.data_type == DataType.ARRAY_FLOAT:
doc_fields[field.name] = [float(i + 0.1), float(i + 1.1), float(i + 2.1)]
elif field.data_type == DataType.ARRAY_DOUBLE:
doc_fields[field.name] = [float(i + 0.11), float(i + 1.11), float(i + 2.11)]
elif field.data_type == DataType.ARRAY_STRING:
doc_fields[field.name] = [f"test_{i}", f"test_{i + 1}", f"test_{i + 2}"]
else:
raise ValueError(f"Unsupported field type: {field.data_type}")
for vector in schema.vectors:
if vector.data_type == DataType.VECTOR_FP16:
doc_vectors[vector.name] = generate_constant_vector(
i, vector.dimension, "float16"
)
elif vector.data_type == DataType.VECTOR_FP32:
doc_vectors[vector.name] = generate_constant_vector(
i, vector.dimension, "float32"
)
elif vector.data_type == DataType.VECTOR_INT8:
doc_vectors[vector.name] = generate_constant_vector(
i,
vector.dimension,
"int8",
)
elif vector.data_type == DataType.SPARSE_VECTOR_FP32:
doc_vectors[vector.name] = generate_sparse_vector(i)
elif vector.data_type == DataType.SPARSE_VECTOR_FP16:
doc_vectors[vector.name] = generate_sparse_vector(i)
else:
raise ValueError(f"Unsupported vector type: {vector.data_type}")
return doc_fields, doc_vectors
def generate_vectordict_recall(i: int, schema: CollectionSchema) -> Doc:
doc_fields = {}
doc_vectors = {}
doc_fields = {}
doc_vectors = {}
for field in schema.fields:
if field.data_type == DataType.BOOL:
doc_fields[field.name] = i % 2 == 0
elif field.data_type == DataType.INT32:
doc_fields[field.name] = i
elif field.data_type == DataType.UINT32:
doc_fields[field.name] = i
elif field.data_type == DataType.INT64:
doc_fields[field.name] = i
elif field.data_type == DataType.UINT64:
doc_fields[field.name] = i
elif field.data_type == DataType.FLOAT:
doc_fields[field.name] = float(i) + 0.1
elif field.data_type == DataType.DOUBLE:
doc_fields[field.name] = float(i) + 0.11
elif field.data_type == DataType.STRING:
doc_fields[field.name] = f"test_{i}"
elif field.data_type == DataType.ARRAY_BOOL:
doc_fields[field.name] = [i % 2 == 0, i % 3 == 0]
elif field.data_type == DataType.ARRAY_INT32:
doc_fields[field.name] = [i, i + 1, i + 2]
elif field.data_type == DataType.ARRAY_UINT32:
doc_fields[field.name] = [i, i + 1, i + 2]
elif field.data_type == DataType.ARRAY_INT64:
doc_fields[field.name] = [i, i + 1, i + 2]
elif field.data_type == DataType.ARRAY_UINT64:
doc_fields[field.name] = [i, i + 1, i + 2]
elif field.data_type == DataType.ARRAY_FLOAT:
doc_fields[field.name] = [float(i + 0.1), float(i + 1.1), float(i + 2.1)]
elif field.data_type == DataType.ARRAY_DOUBLE:
doc_fields[field.name] = [float(i + 0.11), float(i + 1.11), float(i + 2.11)]
elif field.data_type == DataType.ARRAY_STRING:
doc_fields[field.name] = [f"test_{i}", f"test_{i + 1}", f"test_{i + 2}"]
else:
raise ValueError(f"Unsupported field type: {field.data_type}")
for vector in schema.vectors:
if vector.data_type == DataType.VECTOR_FP16:
doc_vectors[vector.name] = generate_constant_vector_recall(
i, vector.dimension, "float16"
)
elif vector.data_type == DataType.VECTOR_FP32:
doc_vectors[vector.name] = generate_constant_vector_recall(
i, vector.dimension, "float32"
)
elif vector.data_type == DataType.VECTOR_INT8:
doc_vectors[vector.name] = generate_constant_vector_recall(
i,
vector.dimension,
"int8",
)
elif vector.data_type == DataType.SPARSE_VECTOR_FP32:
doc_vectors[vector.name] = generate_sparse_vector(i)
elif vector.data_type == DataType.SPARSE_VECTOR_FP16:
doc_vectors[vector.name] = generate_sparse_vector(i)
else:
raise ValueError(f"Unsupported vector type: {vector.data_type}")
return doc_fields, doc_vectors
def generate_vectordict_update(i: int, schema: CollectionSchema) -> Doc:
doc_fields = {}
doc_vectors = {}
doc_fields = {}
doc_vectors = {}
for field in schema.fields:
if field.data_type == DataType.BOOL:
doc_fields[field.name] = (i + 1) % 2 == 0
elif field.data_type == DataType.INT32:
doc_fields[field.name] = i + 1
elif field.data_type == DataType.UINT32:
doc_fields[field.name] = i + 1
elif field.data_type == DataType.INT64:
doc_fields[field.name] = i + 1
elif field.data_type == DataType.UINT64:
doc_fields[field.name] = i + 1
elif field.data_type == DataType.FLOAT:
doc_fields[field.name] = float(i + 1) + 0.1
elif field.data_type == DataType.DOUBLE:
doc_fields[field.name] = float(i + 1) + 0.11
elif field.data_type == DataType.STRING:
doc_fields[field.name] = f"test_{i + 1}"
elif field.data_type == DataType.ARRAY_BOOL:
doc_fields[field.name] = [(i + 1) % 2 == 0, (i + 1) % 3 == 0]
elif field.data_type == DataType.ARRAY_INT32:
doc_fields[field.name] = [i + 1, i + 1, i + 2]
elif field.data_type == DataType.ARRAY_UINT32:
doc_fields[field.name] = [i + 1, i + 1, i + 2]
elif field.data_type == DataType.ARRAY_INT64:
doc_fields[field.name] = [i + 1, i + 1, i + 2]
elif field.data_type == DataType.ARRAY_UINT64:
doc_fields[field.name] = [i + 1, i + 1, i + 2]
elif field.data_type == DataType.ARRAY_FLOAT:
doc_fields[field.name] = [float(i + 1.1), float(i + 2.1), float(i + 3.1)]
elif field.data_type == DataType.ARRAY_DOUBLE:
doc_fields[field.name] = [float(i + 1.11), float(i + 2.11), float(i + 3.11)]
elif field.data_type == DataType.ARRAY_STRING:
doc_fields[field.name] = [f"test_{i + 1}", f"test_{i + 2}", f"test_{i + 3}"]
else:
raise ValueError(f"Unsupported field type: {field.data_type}")
for vector in schema.vectors:
if vector.data_type == DataType.VECTOR_FP16:
doc_vectors[vector.name] = generate_constant_vector(
i + 1, vector.dimension, "float16"
)
elif vector.data_type == DataType.VECTOR_FP32:
doc_vectors[vector.name] = generate_constant_vector(
i + 1, vector.dimension, "float32"
)
elif vector.data_type == DataType.VECTOR_INT8:
doc_vectors[vector.name] = generate_constant_vector(
i + 1,
vector.dimension,
"int8",
)
elif vector.data_type == DataType.SPARSE_VECTOR_FP32:
doc_vectors[vector.name] = generate_sparse_vector(i + 1)
elif vector.data_type == DataType.SPARSE_VECTOR_FP16:
doc_vectors[vector.name] = generate_sparse_vector(i + 1)
else:
raise ValueError(f"Unsupported vector type: {vector.data_type}")
return doc_fields, doc_vectors
def generate_doc(i: int, schema: CollectionSchema) -> Doc:
doc_fields = {}
doc_vectors = {}
doc_fields, doc_vectors = generate_vectordict(i, schema)
doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors)
return doc
def generate_doc_recall(i: int, schema: CollectionSchema) -> Doc:
doc_fields = {}
doc_vectors = {}
doc_fields, doc_vectors = generate_vectordict_recall(i, schema)
doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors)
return doc
def generate_update_doc(i: int, schema: CollectionSchema) -> Doc:
doc_fields = {}
doc_vectors = {}
doc_fields, doc_vectors = generate_vectordict_update(i, schema)
doc = Doc(id=str(i), fields=doc_fields, vectors=doc_vectors)
return doc
def generate_doc_random(i, schema: CollectionSchema) -> Doc:
doc_fields = {}
doc_vectors = {}
random.seed(i)
for field in schema.fields:
if field.data_type == DataType.BOOL:
doc_fields[field.name] = random.choice([True, False])
elif field.data_type == DataType.INT32:
doc_fields[field.name] = random.randint(-2147483648, 2147483647)
elif field.data_type == DataType.UINT32:
doc_fields[field.name] = random.randint(0, 4294967295)
elif field.data_type == DataType.INT64:
doc_fields[field.name] = random.randint(
-9223372036854775808, 9223372036854775807
)
elif field.data_type == DataType.UINT64:
doc_fields[field.name] = random.randint(0, 18446744073709551615)
elif field.data_type == DataType.FLOAT:
doc_fields[field.name] = random.uniform(-3.4028235e38, 3.4028235e38)
elif field.data_type == DataType.DOUBLE:
doc_fields[field.name] = random.uniform(
-1.7976931348623157e308, 1.7976931348623157e308
)
elif field.data_type == DataType.STRING:
length = random.randint(1, 999)
doc_fields[field.name] = "".join(
random.choices(string.ascii_letters + string.digits, k=length)
)
elif field.data_type == DataType.ARRAY_BOOL:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.choice([True, False]) for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_INT32:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.randint(-2147483648, 2147483647) for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_UINT32:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.randint(0, 4294967295) for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_INT64:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.randint(-9223372036854775808, 9223372036854775807)
for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_UINT64:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.randint(0, 18446744073709551615) for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_FLOAT:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.uniform(-3.4028235e38, 3.4028235e38) for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_DOUBLE:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.uniform(-1.7976931348623157e308, 1.7976931348623157e308)
for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_STRING:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
"".join(
random.choices(
string.ascii_letters + string.digits, k=random.randint(1, 100)
)
)
for _ in range(array_length)
]
else:
raise ValueError(f"Unsupported field type: {field.data_type}")
for vector in schema.vectors:
if vector.data_type == DataType.VECTOR_FP16:
doc_vectors[vector.name] = generate_constant_vector(
random.randint(1, 100), DEFAULT_VECTOR_DIMENSION, "float16"
)
elif vector.data_type == DataType.VECTOR_FP32:
doc_vectors[vector.name] = generate_constant_vector(
random.randint(1, 100), DEFAULT_VECTOR_DIMENSION, "float32"
)
elif vector.data_type == DataType.VECTOR_INT8:
doc_vectors[vector.name] = generate_constant_vector(
random.randint(1, 100), DEFAULT_VECTOR_DIMENSION, "int8"
)
elif vector.data_type == DataType.SPARSE_VECTOR_FP32:
doc_vectors[vector.name] = generate_sparse_vector(random.randint(1, 100))
elif vector.data_type == DataType.SPARSE_VECTOR_FP16:
doc_vectors[vector.name] = generate_sparse_vector(random.randint(1, 100))
else:
raise ValueError(f"Unsupported vector type: {vector.data_type}")
doc = Doc(id=i, fields=doc_fields, vectors=doc_vectors)
return doc
def generate_vectordict_random(schema: CollectionSchema):
doc_fields = {}
doc_vectors = {}
for field in schema.fields:
if field.data_type == DataType.BOOL:
doc_fields[field.name] = random.choice([True, False])
elif field.data_type == DataType.INT32:
doc_fields[field.name] = random.randint(-2147483648, 2147483647)
elif field.data_type == DataType.UINT32:
doc_fields[field.name] = random.randint(0, 4294967295)
elif field.data_type == DataType.INT64:
doc_fields[field.name] = random.randint(
-9223372036854775808, 9223372036854775807
)
elif field.data_type == DataType.UINT64:
doc_fields[field.name] = random.randint(0, 18446744073709551615)
elif field.data_type == DataType.FLOAT:
doc_fields[field.name] = random.uniform(-3.4028235e38, 3.4028235e38)
elif field.data_type == DataType.DOUBLE:
doc_fields[field.name] = random.uniform(
-1.7976931348623157e308, 1.7976931348623157e308
)
elif field.data_type == DataType.STRING:
length = random.randint(1, 999)
doc_fields[field.name] = "".join(
random.choices(string.ascii_letters + string.digits, k=length)
)
elif field.data_type == DataType.ARRAY_BOOL:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.choice([True, False]) for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_INT32:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.randint(-2147483648, 2147483647) for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_UINT32:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.randint(0, 4294967295) for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_INT64:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.randint(-9223372036854775808, 9223372036854775807)
for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_UINT64:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.randint(0, 18446744073709551615) for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_FLOAT:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.uniform(-3.4028235e38, 3.4028235e38) for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_DOUBLE:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
random.uniform(-1.7976931348623157e308, 1.7976931348623157e308)
for _ in range(array_length)
]
elif field.data_type == DataType.ARRAY_STRING:
array_length = random.randint(0, 10)
doc_fields[field.name] = [
"".join(
random.choices(
string.ascii_letters + string.digits, k=random.randint(1, 100)
)
)
for _ in range(array_length)
]
else:
raise ValueError(f"Unsupported field type: {field.data_type}")
for vector in schema.vectors:
if vector.data_type == DataType.VECTOR_FP16:
doc_vectors[vector.name] = generate_constant_vector(
random.randint(1, 100), vector.dimension, "float16"
)
elif vector.data_type == DataType.VECTOR_FP32:
doc_vectors[vector.name] = generate_constant_vector(
random.randint(1, 100), vector.dimension, "float32"
)
elif vector.data_type == DataType.VECTOR_INT8:
doc_vectors[vector.name] = generate_constant_vector(
random.randint(1, 100), vector.dimension, "int8"
)
elif vector.data_type == DataType.SPARSE_VECTOR_FP32:
doc_vectors[vector.name] = generate_sparse_vector(random.randint(1, 100))
elif vector.data_type == DataType.SPARSE_VECTOR_FP16:
doc_vectors[vector.name] = generate_sparse_vector(random.randint(1, 100))
else:
raise ValueError(f"Unsupported vector type: {vector.data_type}")
return doc_fields, doc_vectors
Reference in New Issue View Git Blame Copy Permalink