From fcd26c425a8657501b9a438fd8ab35f086bc5943 Mon Sep 17 00:00:00 2001 From: LeonSGP <154585401+LeonSGP43@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:44:05 +0800 Subject: [PATCH] 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**. --- python/python/lancedb/pydantic.py | 12 +++++- python/python/tests/test_pydantic.py | 61 ++++++++++++++++++++++++++++ 2 files changed, 71 insertions(+), 2 deletions(-) diff --git a/python/python/lancedb/pydantic.py b/python/python/lancedb/pydantic.py index 8fe3c1c93..bb70c9a67 100644 --- a/python/python/lancedb/pydantic.py +++ b/python/python/lancedb/pydantic.py @@ -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() + ), ), ), ) diff --git a/python/python/tests/test_pydantic.py b/python/python/tests/test_pydantic.py index 24fb7ee89..6717bbc32 100644 --- a/python/python/tests/test_pydantic.py +++ b/python/python/tests/test_pydantic.py @@ -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",