fix(python): preserve integer Vector and MultiVector values (#4450)

When `Vector` or `MultiVector` declares an integer Arrow value type,
Pydantic currently validates every element as a float. This silently
rounds integers above `2**53`, and converts `uint64`'s maximum value
into an out-of-range float before table insertion.

Use Pydantic's integer schema for integer Arrow value types in both
validators. Floating-point types retain the existing validation. Tests
exercise model serialization and real local table creation, insertion,
and reads for signed and unsigned integer boundaries and floating-point
controls, plus dimensions and optional/nullability behavior. Fractional
floats are rejected for integer declarations; integral floats remain
accepted. This does not extend ANN index support.

Validation on macOS arm64, with the current main native extension built
using Rust 1.97.0:

- Python 3.11.14 and 3.13.11, Pydantic 2.13.5: full `pytest -m "not slow
and not s3_test" -q python/tests` suite: **1,537 passed, 37 skipped, 68
deselected** on each runtime. Comparable unchanged-main baselines:
**1,521 passed, 37 skipped, 68 deselected** each.
- Added cases on unchanged main: **8 failed, 8 passed**; complete
`test_pydantic.py` after the fix: **40 passed** on each runtime.
- Documentation tests with the existing embeddings dependencies in an
isolated environment, Python 3.11.14: **77 passed, 6 skipped**, matching
the unchanged-main baseline. The workflow's Linux/Python 3.13 doctest
matrix was not reproduced locally.
- Repository Ruff lint/format, Pyright, all applicable pre-commit hooks,
compile checks, and `git diff --check` passed. Eight public SDK table
round trips passed on each final installed wheel runtime.

Slow, credential-dependent S3 integration tests, the remaining
OS/runtime matrix, and the minimum Pydantic/PyArrow dependency matrix
were not run locally. The native extension was built from unchanged
current main; the final local wheel repack contains the verified Python
change with that identical native extension. A separate candidate
maturin packaging attempt triggered Rust recompilation and was stopped,
so an official candidate distribution build was not completed. The
project requires Pydantic `>=2.7.4,<3`; Pydantic v1 support was removed
in #3990.

AI assistance: developed and verified with **gpt-6.1-sol**.
This commit is contained in:
LeonSGP
2026-10-23 02:44:05 +08:00
committed by GitHub
parent 1e313d3c0d
commit fcd26c425a
2 changed files with 71 additions and 2 deletions
+10 -2
View File
@@ -127,7 +127,11 @@ def Vector(
core_schema.list_schema(
min_length=dim,
max_length=dim,
items_schema=core_schema.float_schema(),
items_schema=(
core_schema.int_schema()
if pa.types.is_integer(value_type)
else core_schema.float_schema()
),
),
)
@@ -208,7 +212,11 @@ def MultiVector(
items_schema=core_schema.list_schema(
min_length=dim,
max_length=dim,
items_schema=core_schema.float_schema(),
items_schema=(
core_schema.int_schema()
if pa.types.is_integer(value_type)
else core_schema.float_schema()
),
),
),
)
+61
View File
@@ -415,6 +415,67 @@ def test_nullable_vector():
assert schema == pa.schema([pa.field("vec", pa.list_(pa.float32(), 16), True)])
@pytest.mark.parametrize("vector_type", [Vector, MultiVector])
@pytest.mark.parametrize(
"value_type, values",
[
(pa.int64(), [-(2**53 + 1), 2**53 + 3]),
(pa.uint64(), [2**53 + 1, 2**64 - 1]),
(pa.int8(), [-128, 127]),
(pa.uint8(), [0, 255]),
(pa.float32(), [1.25, 2.5]),
(pa.float64(), [1.25, 2.5]),
],
)
def test_typed_vector_round_trip(tmp_path, vector_type, value_type, values):
import lancedb
annotation = vector_type(2, value_type=value_type)
expected = [values, values] if vector_type is MultiVector else values
class Row(LanceModel):
id: int
payload: annotation
row = Row(id=1, payload=expected)
table = lancedb.connect(tmp_path).create_table("typed_vectors", [row], schema=Row)
table.add([Row(id=2, payload=expected)])
assert table.schema == Row.to_arrow_schema()
assert table.search().to_arrow().to_pylist() == [
{"id": 1, "payload": expected},
{"id": 2, "payload": expected},
]
assert json.loads(row.model_dump_json())["payload"] == expected
@pytest.mark.parametrize("vector_type", [Vector, MultiVector])
@pytest.mark.parametrize("nullable", [False, True])
def test_integer_vector_validation(vector_type, nullable):
annotation = vector_type(2, value_type=pa.int64(), nullable=nullable)
class Row(LanceModel):
payload: annotation
def shape(values):
return [values] if vector_type is MultiVector else values
assert Row.to_arrow_schema().field("payload").nullable is nullable
assert Row(payload=shape([1.0, -2.0])).model_dump()["payload"] == shape([1, -2])
with pytest.raises(pydantic.ValidationError, match="valid integer"):
Row(payload=shape([1.5, 2]))
with pytest.raises(pydantic.ValidationError):
Row(payload=shape([1]))
with pytest.raises(pydantic.ValidationError):
Row(payload=shape([1, 2, 3]))
with pytest.raises(pydantic.ValidationError):
Row(payload=None)
class OptionalRow(LanceModel):
payload: Optional[annotation] = None
assert OptionalRow(payload=None).payload is None
def test_bare_vector_raises_clear_error():
namespace = {
"__name__": "test_model_without_pyarrow",