wehub-resource-sync
270 lines
10 KiB
Python
270 lines
10 KiB
Python
# Copyright 2025-present the zvec project
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
import pytest
|
|
|
|
|
|
from zvec._zvec import _Doc
|
|
from zvec import FieldSchema, VectorSchema, Doc, DataType
|
|
|
|
|
|
# ----------------------------
|
|
# PyDoc Test Case
|
|
# ----------------------------
|
|
class TestPyDoc:
|
|
def test_default(self):
|
|
Doc(id="1")
|
|
|
|
def test_with_single_vector(self):
|
|
doc = Doc(id="1", vectors={"dense": [1, 2, 3]})
|
|
assert doc is not None
|
|
assert doc.id == "1"
|
|
assert doc.vector("dense") == [1, 2, 3]
|
|
|
|
def test_with_hybrid_vectors(self):
|
|
doc = Doc(
|
|
id="1", vectors={"dense": [1, 2, 3], "sparse": {1: 1.0, 2: 2.0, 3: 3.0}}
|
|
)
|
|
assert doc is not None
|
|
assert doc.id == "1"
|
|
assert doc.vector("dense") == [1, 2, 3]
|
|
assert doc.vector("sparse") == {1: 1.0, 2: 2.0, 3: 3.0}
|
|
|
|
def test_with_multi_vectors(self):
|
|
doc = Doc(
|
|
id="1",
|
|
vectors={
|
|
"image": [1, 2, 3],
|
|
"description": [4, 5, 6],
|
|
"keys": {1: 1.0, 2: 2.0, 3: 3.0},
|
|
},
|
|
fields={"author": "Tom", "age": 19, "is_male": True, "weight": 60.5},
|
|
)
|
|
assert doc is not None
|
|
assert doc.id == "1"
|
|
assert doc.vector("image") == [1, 2, 3]
|
|
assert doc.vector("description") == [4, 5, 6]
|
|
assert doc.vector("keys") == {1: 1.0, 2: 2.0, 3: 3.0}
|
|
assert doc.field("author") == "Tom"
|
|
assert doc.field("age") == 19
|
|
assert doc.field("is_male") == True
|
|
assert doc.field("weight") == 60.5
|
|
|
|
def test_with_numpy_array(self):
|
|
import numpy as np
|
|
|
|
doc = Doc._from_tuple(
|
|
(
|
|
"1",
|
|
0.0,
|
|
None,
|
|
{
|
|
"image": np.array([1, 2, 3]),
|
|
"description": np.random.random(512),
|
|
"keys": {1: 1.0, 2: 2.0, 3: 3.0},
|
|
},
|
|
)
|
|
)
|
|
assert doc is not None
|
|
assert doc.id == "1"
|
|
assert doc.vector("image") == [1, 2, 3]
|
|
assert doc.vector("keys") == {1: 1.0, 2: 2.0, 3: 3.0}
|
|
|
|
|
|
# ----------------------------
|
|
# CppDoc Test Case
|
|
# ----------------------------
|
|
class TestCppDoc:
|
|
def test_default(self):
|
|
doc = _Doc()
|
|
assert doc is not None
|
|
|
|
def test_doc_set_pk(self):
|
|
doc = _Doc()
|
|
doc.set_pk("1")
|
|
assert doc.pk() == "1"
|
|
|
|
def test_doc_set_score(self):
|
|
doc = _Doc()
|
|
doc.set_score(0.9)
|
|
assert math.isclose(doc.score(), 0.9, rel_tol=1e-6)
|
|
|
|
def test_doc_get_null_field(self):
|
|
doc = _Doc()
|
|
schema = FieldSchema("author", DataType.STRING, nullable=True)
|
|
doc.set_any("author", schema._get_object(), None)
|
|
assert doc.has_field("author")
|
|
assert doc.get_any("author", schema.data_type) is None
|
|
|
|
def test_doc_get_set_has_null_field(self):
|
|
doc = _Doc()
|
|
schema = FieldSchema("author", DataType.STRING, nullable=False)
|
|
with pytest.raises(ValueError):
|
|
doc.set_any("author", schema._get_object(), None)
|
|
|
|
def test_doc_get_set_has_string_field(self):
|
|
doc = _Doc()
|
|
schema = FieldSchema("author", DataType.STRING)
|
|
doc.set_any("author", schema._get_object(), "Tom")
|
|
assert doc.has_field("author")
|
|
assert doc.get_any("author", DataType.STRING) == "Tom"
|
|
|
|
def test_doc_get_set_has_bool_field(self):
|
|
doc = _Doc()
|
|
schema = FieldSchema("is_male", DataType.BOOL)
|
|
doc.set_any("is_male", schema._get_object(), True)
|
|
assert doc.has_field("is_male")
|
|
assert doc.get_any("is_male", DataType.BOOL) == True
|
|
|
|
def test_doc_get_set_has_int32_field(self):
|
|
doc = _Doc()
|
|
schema = FieldSchema("age", DataType.INT32)
|
|
doc.set_any("age", schema._get_object(), 19)
|
|
assert doc.has_field("age")
|
|
assert doc.get_any("age", DataType.INT32) == 19
|
|
|
|
def test_doc_get_set_has_int64_field(self):
|
|
doc = _Doc()
|
|
schema = FieldSchema("id", DataType.INT64)
|
|
doc.set_any("id", schema._get_object(), 1111111111111111111)
|
|
assert doc.has_field("id")
|
|
assert doc.get_any("id", DataType.INT64) == 1111111111111111111
|
|
|
|
def test_doc_get_set_has_float_field(self):
|
|
doc = _Doc()
|
|
schema = FieldSchema("weight", DataType.FLOAT)
|
|
doc.set_any("weight", schema._get_object(), 60.5)
|
|
assert doc.has_field("weight")
|
|
assert math.isclose(doc.get_any("weight", DataType.FLOAT), 60.5, rel_tol=1e-6)
|
|
|
|
def test_doc_get_set_has_double_field(self):
|
|
doc = _Doc()
|
|
schema = FieldSchema("height", DataType.DOUBLE)
|
|
doc.set_any("height", schema._get_object(), 1.77777777777)
|
|
assert doc.has_field("height")
|
|
assert math.isclose(
|
|
doc.get_any("height", DataType.DOUBLE), 1.7777777777, rel_tol=1e-9
|
|
)
|
|
|
|
def test_doc_get_set_has_uint32_field(self):
|
|
doc = _Doc()
|
|
schema = FieldSchema("id", DataType.UINT32)
|
|
doc.set_any("id", schema._get_object(), 4294967295)
|
|
assert doc.has_field("id")
|
|
assert doc.get_any("id", DataType.UINT32) == 4294967295
|
|
|
|
def test_doc_get_set_has_uint64_field(self):
|
|
doc = _Doc()
|
|
schema = FieldSchema("id", DataType.UINT64)
|
|
doc.set_any("id", schema._get_object(), 18446744073709551615)
|
|
assert doc.has_field("id")
|
|
assert doc.get_any("id", DataType.UINT64) == 18446744073709551615
|
|
|
|
def test_doc_get_set_has_array_string_field(self):
|
|
doc = _Doc()
|
|
schema = FieldSchema("tags", DataType.ARRAY_STRING)
|
|
doc.set_any("tags", schema._get_object(), ["tag1", "tag2", "tag3"])
|
|
assert doc.has_field("tags")
|
|
assert doc.get_any("tags", DataType.ARRAY_STRING) == ["tag1", "tag2", "tag3"]
|
|
|
|
def test_doc_get_set_has_array_int32_field(self):
|
|
doc = _Doc()
|
|
schema = FieldSchema("ids", DataType.ARRAY_INT32)
|
|
doc.set_any("ids", schema._get_object(), [1, 2, 3])
|
|
assert doc.has_field("ids")
|
|
assert doc.get_any("ids", DataType.ARRAY_INT32) == [1, 2, 3]
|
|
|
|
def test_doc_get_set_has_array_int64_field(self):
|
|
doc = _Doc()
|
|
schema = FieldSchema("ids", DataType.ARRAY_INT64)
|
|
doc.set_any("ids", schema._get_object(), [1, 2, 3])
|
|
assert doc.has_field("ids")
|
|
assert doc.get_any("ids", DataType.ARRAY_INT64) == [1, 2, 3]
|
|
|
|
def test_doc_get_set_has_array_float_field(self):
|
|
doc = _Doc()
|
|
schema = FieldSchema("weights", DataType.ARRAY_FLOAT)
|
|
doc.set_any("weights", schema._get_object(), [1.0, 2.0, 3.0])
|
|
assert doc.has_field("weights")
|
|
assert doc.get_any("weights", DataType.ARRAY_FLOAT) == [1.0, 2.0, 3.0]
|
|
|
|
def test_doc_get_set_has_array_double_field(self):
|
|
doc = _Doc()
|
|
schema = FieldSchema("heights", DataType.ARRAY_DOUBLE)
|
|
doc.set_any("heights", schema._get_object(), [1.0, 2.0, 3.0])
|
|
assert doc.has_field("heights")
|
|
assert doc.get_any("heights", DataType.ARRAY_DOUBLE) == [1.0, 2.0, 3.0]
|
|
|
|
def test_doc_get_set_has_array_bool_field(self):
|
|
doc = _Doc()
|
|
schema = FieldSchema("bools", DataType.ARRAY_BOOL)
|
|
doc.set_any("bools", schema._get_object(), [True, False, True])
|
|
assert doc.has_field("bools")
|
|
assert doc.get_any("bools", DataType.ARRAY_BOOL) == [True, False, True]
|
|
|
|
def test_doc_get_set_has_vector_fp16(self):
|
|
doc = _Doc()
|
|
schema = VectorSchema("image", DataType.VECTOR_FP16)
|
|
doc.set_any("image", schema._get_object(), [1.0, 2.0, 3.0])
|
|
assert doc.has_field("image")
|
|
image_vector = doc.get_any("image", DataType.VECTOR_FP16)
|
|
assert image_vector is not None
|
|
for i in range(len(image_vector)):
|
|
assert math.isclose(image_vector[i], [1.0, 2.0, 3.0][i], rel_tol=1e-6)
|
|
|
|
def test_doc_get_set_has_vector_fp32(self):
|
|
doc = _Doc()
|
|
schema = VectorSchema("image", DataType.VECTOR_FP32)
|
|
doc.set_any("image", schema._get_object(), [1.111111, 2.222222, 3.333333])
|
|
assert doc.has_field("image")
|
|
vector = doc.get_any("image", DataType.VECTOR_FP32)
|
|
assert vector is not None
|
|
for i in range(len(vector)):
|
|
assert math.isclose(
|
|
vector[i], [1.111111, 2.222222, 3.333333][i], rel_tol=1e-6
|
|
)
|
|
|
|
def test_doc_get_set_has_vector_int8(self):
|
|
doc = _Doc()
|
|
schema = VectorSchema("image", DataType.VECTOR_INT8)
|
|
doc.set_any("image", schema._get_object(), [1, 2, 3])
|
|
assert doc.has_field("image")
|
|
assert doc.get_any("image", DataType.VECTOR_INT8) == [1, 2, 3]
|
|
|
|
def test_doc_get_set_has_sparse_vector_fp32(self):
|
|
doc = _Doc()
|
|
sparse = {1: 1.111111, 2: 2.222222, 3: 3.333333}
|
|
schema = VectorSchema("key", DataType.SPARSE_VECTOR_FP32)
|
|
doc.set_any("key", schema._get_object(), sparse)
|
|
assert doc.has_field("key")
|
|
vector = doc.get_any("key", DataType.SPARSE_VECTOR_FP32)
|
|
assert vector is not None
|
|
assert isinstance(vector, dict)
|
|
for key, value in sparse.items():
|
|
assert math.isclose(vector[key], value, rel_tol=1e-6)
|
|
|
|
def test_doc_get_set_has_sparse_vector_fp16(self):
|
|
doc = _Doc()
|
|
sparse = {1: 1.1, 2: 2.2, 3: 3.3}
|
|
schema = VectorSchema("key", DataType.SPARSE_VECTOR_FP16)
|
|
doc.set_any("key", schema._get_object(), sparse)
|
|
assert doc.has_field("key")
|
|
vector = doc.get_any("key", DataType.SPARSE_VECTOR_FP16)
|
|
assert vector is not None
|
|
assert isinstance(vector, dict)
|
|
for key, value in sparse.items():
|
|
assert math.isclose(vector[key], value, rel_tol=1e-1)
|