Files
2026-07-13 12:47:42 +08:00

127 lines
3.3 KiB
Python

from zvec import (
CollectionOption,
IndexOption,
OptimizeOption,
InvertIndexParam,
HnswIndexParam,
IVFIndexParam,
FlatIndexParam,
DataType,
IndexType,
QuantizeType,
)
SUPPORT_SCALAR_DATA_TYPES = [
DataType.BOOL,
DataType.FLOAT,
DataType.DOUBLE,
DataType.INT32,
DataType.INT64,
DataType.UINT32,
DataType.UINT64,
DataType.STRING,
DataType.ARRAY_BOOL,
DataType.ARRAY_FLOAT,
DataType.ARRAY_DOUBLE,
DataType.ARRAY_INT32,
DataType.ARRAY_INT64,
DataType.ARRAY_UINT32,
DataType.ARRAY_UINT64,
DataType.ARRAY_STRING,
]
DEFAULT_SCALAR_FIELD_NAME = {
DataType.BOOL: "bool_field",
DataType.FLOAT: "float_field",
DataType.DOUBLE: "double_field",
DataType.INT32: "int32_field",
DataType.INT64: "int64_field",
DataType.UINT32: "uint32_field",
DataType.UINT64: "uint64_field",
DataType.STRING: "string_field",
DataType.ARRAY_BOOL: "array_bool_field",
DataType.ARRAY_FLOAT: "array_float_field",
DataType.ARRAY_DOUBLE: "array_double_field",
DataType.ARRAY_INT32: "array_int32_field",
DataType.ARRAY_INT64: "array_int64_field",
DataType.ARRAY_UINT32: "array_uint32_field",
DataType.ARRAY_UINT64: "array_uint64_field",
DataType.ARRAY_STRING: "array_string_field",
}
SUPPORT_SCALAR_INDEX_TYPES = [
IndexType.INVERT,
]
SUPPORT_VECTOR_DATA_TYPES = [
DataType.VECTOR_FP16,
DataType.VECTOR_FP32,
DataType.VECTOR_INT8,
DataType.SPARSE_VECTOR_FP32,
DataType.SPARSE_VECTOR_FP16,
]
SUPPORT_VECTOR_INDEX_TYPES = [
IndexType.FLAT,
IndexType.HNSW,
IndexType.IVF,
]
DEFAULT_VECTOR_FIELD_NAME = {
DataType.VECTOR_FP16: "vector_fp16_field",
DataType.VECTOR_FP32: "vector_fp32_field",
DataType.VECTOR_INT8: "vector_int8_field",
DataType.SPARSE_VECTOR_FP32: "sparse_vector_fp32_field",
DataType.SPARSE_VECTOR_FP16: "sparse_vector_fp16_field",
}
DEFAULT_VECTOR_DIMENSION = 128
VECTOR_DIMENSION_1024 = 4
SUPPORT_VECTOR_DATA_TYPE_INDEX_MAP = {
DataType.VECTOR_FP16: [IndexType.FLAT, IndexType.HNSW, IndexType.IVF],
DataType.VECTOR_FP32: [IndexType.FLAT, IndexType.HNSW, IndexType.IVF],
DataType.VECTOR_INT8: [IndexType.FLAT, IndexType.HNSW],
DataType.SPARSE_VECTOR_FP32: [IndexType.FLAT, IndexType.HNSW],
DataType.SPARSE_VECTOR_FP16: [IndexType.FLAT, IndexType.HNSW],
}
SUPPORT_VECTOR_DATA_TYPE_INDEX_MAP_PARAMS = [
(data_type, index_type)
for data_type, index_types in SUPPORT_VECTOR_DATA_TYPE_INDEX_MAP.items()
for index_type in index_types
]
DEFAULT_INDEX_PARAMS = {
IndexType.FLAT: FlatIndexParam(),
IndexType.HNSW: HnswIndexParam(),
IndexType.IVF: IVFIndexParam(),
IndexType.INVERT: InvertIndexParam(),
}
SUPPORT_VECTOR_DATA_TYPE_QUANT_MAP = {
DataType.VECTOR_FP32: [QuantizeType.FP16, QuantizeType.INT8, QuantizeType.INT4],
DataType.SPARSE_VECTOR_FP32: [QuantizeType.FP16],
}
SUPPORT_ADD_COLUMN_DATA_TYPE = [
DataType.INT32,
DataType.UINT32,
DataType.INT64,
DataType.UINT64,
DataType.FLOAT,
DataType.DOUBLE,
]
NOT_SUPPORT_ADD_COLUMN_DATA_TYPE = [
DataType.BOOL,
DataType.STRING,
DataType.ARRAY_BOOL,
DataType.ARRAY_INT32,
DataType.ARRAY_INT64,
DataType.ARRAY_UINT32,
DataType.ARRAY_UINT64,
DataType.ARRAY_FLOAT,
DataType.ARRAY_DOUBLE,
DataType.ARRAY_STRING,
]