fix(python): fill bad vector values element-wise (#3613)

## Summary

Fix `on_bad_vectors="fill"` so it replaces only invalid or missing
vector values instead of replacing the entire vector row.

Fixes #3026.

## Reasoning

The old Python sanitizer detected whether a vector row was bad at row
granularity. For `fill`, it then used that row-level flag to replace the
whole vector with `[fill_value] * dim`. That meant an input like `[1.0,
NaN, 3.0]` became `[0.0, 0.0, 0.0]`, even though the documented and more
useful behavior is to preserve valid values and fill only the bad
element.

I checked whether this should be a Rust-side fix so TypeScript users
would benefit too. Today, Rust core exposes `NaNVectorBehavior::{Error,
Keep}` for rejecting or keeping NaN vectors, while the Python
`on_bad_vectors` API (`error`, `drop`, `fill`, `null`) is implemented in
the Python ingestion sanitizer before data reaches Rust. TypeScript does
not expose the Python `on_bad_vectors="fill"` behavior today. Moving
this exact behavior to Rust would be a broader cross-language API
change, so this PR keeps the fix scoped to the currently affected Python
API.

## What changed

- Added a small helper that fills bad vector rows by preserving valid
elements, replacing NaN elements with `fill_value`, truncating vectors
longer than the expected dimension, and padding short vectors with
`fill_value`.
- Kept the existing fast path unchanged: the helper only runs after bad
vectors are detected and `on_bad_vectors="fill"` is selected.
- Updated sanitizer and table tests to assert element-wise NaN
replacement and short-vector padding for both `create_table` and `add`.

## Validation

- `uv run ruff format .`
- `uv run ruff check .`
- `cd python && uv run --no-sync pytest
python/tests/test_util.py::test_handle_bad_vectors_jagged
python/tests/test_util.py::test_handle_bad_vectors_nan
python/tests/test_table.py::test_create_with_nans
python/tests/test_table.py::test_add_with_nans -vv`

Targeted pytest result: `10 passed`.

## Why this fix is Python-side (and not Rust)

The problematic behavior lives in Python’s `on_bad_vectors` sanitizer,
before data is handed off to Rust. Rust currently only exposes
`NaNVectorBehavior::{Error, Keep}` for add operations, while Python has
the richer `on_bad_vectors={"error","drop","fill","null"}` API.
TypeScript does not currently expose the Python-style fill behavior, so
moving this exact fix into Rust would require designing a broader
cross-language bad-vector handling API.

This PR keeps the change scoped to the existing affected surface:
Python’s `on_bad_vectors="fill"` path. This way, Python users
immediately benefit.
This commit is contained in:
Prashanth Rao
2026-07-14 16:43:17 -04:00
committed by GitHub
parent 137eac9b50
commit 3b626efa47
3 changed files with 117 additions and 19 deletions
+24 -9
View File
@@ -1611,16 +1611,23 @@ def test_create_with_nans(mem_db: DBConnection):
"fill_test",
data=[
{"vector": [3.1, 4.1], "item": "foo", "price": 10.0},
{"vector": [2.1, 4.1], "item": "foo", "price": 9.0},
{"vector": [np.nan], "item": "bar", "price": 20.0},
{"vector": [np.nan, np.nan], "item": "bar", "price": 20.0},
{"vector": [np.nan, 5.0], "item": "bar", "price": 21.0},
{"vector": [5], "item": "bar", "price": 22.0},
],
on_bad_vectors="fill",
fill_value=0.0,
)
assert len(table) == 3
assert len(table) == 5
arrow_tbl = table.search().where("item == 'bar'").to_arrow()
v = arrow_tbl["vector"].to_pylist()[0]
assert np.allclose(v, np.array([0.0, 0.0]))
filled_vectors = {
row["price"]: row["vector"]
for row in arrow_tbl.select(["price", "vector"]).to_pylist()
}
assert np.allclose(filled_vectors[20.0], np.array([0.0, 0.0]))
assert np.allclose(filled_vectors[21.0], np.array([0.0, 5.0]))
assert np.allclose(filled_vectors[22.0], np.array([5.0, 0.0]))
def test_add_with_nans(mem_db: DBConnection):
@@ -1663,15 +1670,21 @@ def test_add_with_nans(mem_db: DBConnection):
data=[
{"vector": [3.1, 4.1], "item": "foo", "price": 10.0},
{"vector": [np.nan], "item": "bar", "price": 20.0},
{"vector": [np.nan, np.nan], "item": "bar", "price": 20.0},
{"vector": [np.nan, 5.0], "item": "bar", "price": 21.0},
{"vector": [5], "item": "bar", "price": 22.0},
],
on_bad_vectors="fill",
fill_value=0.0,
)
assert len(table) == 3
assert len(table) == 4
arrow_tbl = table.search().where("item == 'bar'").to_arrow()
v = arrow_tbl["vector"].to_pylist()[0]
assert np.allclose(v, np.array([0.0, 0.0]))
filled_vectors = {
row["price"]: row["vector"]
for row in arrow_tbl.select(["price", "vector"]).to_pylist()
}
assert np.allclose(filled_vectors[20.0], np.array([0.0, 0.0]))
assert np.allclose(filled_vectors[21.0], np.array([0.0, 5.0]))
assert np.allclose(filled_vectors[22.0], np.array([5.0, 0.0]))
def test_add_with_empty_fixed_size_list_drops_bad_rows(mem_db: DBConnection):
@@ -1832,7 +1845,9 @@ def test_on_bad_vectors_fill_preserves_arrow_nested_vector_type(mem_db: DBConnec
fill_value=0.0,
)
assert table.to_arrow()["vector"].to_pylist() == [[1.0, 2.0], [0.0, 0.0]]
vector = table.to_arrow()["vector"]
assert vector.type == pa.list_(pa.float32())
assert vector.to_pylist() == [[1.0, 2.0], [0.0, 3.0]]
@pytest.mark.parametrize(