mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-19 04:28:35 +00:00
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:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user