mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-18 20:18:37 +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:
Prashanth Rao
GitHub
@@ -4154,17 +4154,58 @@ def _handle_bad_vector_column(
|
||||
raise ValueError(
|
||||
"`fill_value` must not be None if `on_bad_vectors` is 'fill'"
|
||||
)
|
||||
vec_arr = pc.if_else(
|
||||
is_bad,
|
||||
pa.scalar([fill_value] * dim, type=vec_arr.type),
|
||||
vec_arr,
|
||||
)
|
||||
vec_arr = _fill_bad_vector_values(vec_arr, dim, fill_value)
|
||||
else:
|
||||
raise ValueError(f"Invalid value for on_bad_vectors: {on_bad_vectors}")
|
||||
|
||||
return data.set_column(position, vector_column_name, vec_arr)
|
||||
|
||||
|
||||
def _fill_bad_vector_values(
|
||||
arr: Union[pa.Array, pa.ChunkedArray],
|
||||
dim: int,
|
||||
fill_value: float,
|
||||
) -> pa.Array:
|
||||
if not isinstance(arr, pa.ChunkedArray):
|
||||
arr = pa.chunked_array([arr])
|
||||
arr = arr.combine_chunks()
|
||||
|
||||
# A fixed-size slice truncates long vectors and pads short vectors with nulls.
|
||||
# Slice an array marking the original child nulls in parallel so padding nulls
|
||||
# can be distinguished from null values that were already present.
|
||||
sliced = pc.list_slice(arr, 0, dim, return_fixed_size_list=True)
|
||||
child_nulls = pc.is_null(arr.values)
|
||||
parent_nulls = pc.is_null(arr)
|
||||
if pa.types.is_list(arr.type):
|
||||
original_child_nulls = pa.ListArray.from_arrays(
|
||||
arr.offsets, child_nulls, mask=parent_nulls
|
||||
)
|
||||
elif pa.types.is_large_list(arr.type):
|
||||
original_child_nulls = pa.LargeListArray.from_arrays(
|
||||
arr.offsets, child_nulls, mask=parent_nulls
|
||||
)
|
||||
else:
|
||||
original_child_nulls = pa.FixedSizeListArray.from_arrays(
|
||||
child_nulls, arr.type.list_size, mask=parent_nulls
|
||||
)
|
||||
sliced_child_nulls = pc.list_slice(
|
||||
original_child_nulls, 0, dim, return_fixed_size_list=True
|
||||
)
|
||||
needs_fill = pc.is_null(sliced_child_nulls.values)
|
||||
|
||||
values = sliced.values
|
||||
if pa.types.is_floating(values.type):
|
||||
values_for_nan_check = (
|
||||
values.cast(pa.float32()) if pa.types.is_float16(values.type) else values
|
||||
)
|
||||
needs_fill = pc.or_kleene(needs_fill, pc.is_nan(values_for_nan_check))
|
||||
|
||||
fill_scalar = pa.scalar(fill_value).cast(values.type)
|
||||
filled_values = pc.if_else(needs_fill, fill_scalar, values)
|
||||
filled = pa.FixedSizeListArray.from_arrays(filled_values, dim)
|
||||
return filled.cast(arr.type)
|
||||
|
||||
|
||||
def has_nan_values(arr: Union[pa.ListArray, pa.ChunkedArray]) -> pa.BooleanArray:
|
||||
if isinstance(arr, pa.ChunkedArray):
|
||||
values = pa.chunked_array([chunk.flatten() for chunk in arr.chunks])
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -13,6 +13,7 @@ from lancedb.embeddings.registry import EmbeddingFunctionRegistry
|
||||
from lancedb.table import (
|
||||
_append_vector_columns,
|
||||
_cast_to_target_schema,
|
||||
_fill_bad_vector_values,
|
||||
_handle_bad_vectors,
|
||||
_into_pyarrow_reader,
|
||||
_infer_target_schema,
|
||||
@@ -287,7 +288,9 @@ def test_append_vector_columns():
|
||||
|
||||
@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"])
|
||||
def test_handle_bad_vectors_jagged(on_bad_vectors):
|
||||
vector = pa.array([[1.0, 2.0], [3.0], [4.0, 5.0]])
|
||||
vector = pa.array(
|
||||
[[1.0, 2.0], [3.0], [4.0, 5.0], [6.0, 7.0, 8.0], [None, 9.0], None]
|
||||
)
|
||||
schema = pa.schema({"vector": pa.list_(pa.float64())})
|
||||
data = pa.table({"vector": vector}, schema=schema)
|
||||
|
||||
@@ -313,15 +316,54 @@ def test_handle_bad_vectors_jagged(on_bad_vectors):
|
||||
).read_all()
|
||||
|
||||
if on_bad_vectors == "drop":
|
||||
expected = pa.array([[1.0, 2.0], [4.0, 5.0]])
|
||||
expected = pa.array([[1.0, 2.0], [4.0, 5.0], [None, 9.0]])
|
||||
elif on_bad_vectors == "fill":
|
||||
expected = pa.array([[1.0, 2.0], [42.0, 42.0], [4.0, 5.0]])
|
||||
expected = pa.array(
|
||||
[
|
||||
[1.0, 2.0],
|
||||
[3.0, 42.0],
|
||||
[4.0, 5.0],
|
||||
[6.0, 7.0],
|
||||
[None, 9.0],
|
||||
[42.0, 42.0],
|
||||
]
|
||||
)
|
||||
elif on_bad_vectors == "null":
|
||||
expected = pa.array([[1.0, 2.0], None, [4.0, 5.0]])
|
||||
expected = pa.array([[1.0, 2.0], None, [4.0, 5.0], None, [None, 9.0], None])
|
||||
|
||||
assert output["vector"].combine_chunks() == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("vector_type", "vectors", "expected"),
|
||||
[
|
||||
(
|
||||
pa.list_(pa.float64()),
|
||||
[[1.0, float("nan")], [2.0], None, [None, 3.0], [4.0, 5.0, 6.0]],
|
||||
[[1.0, 42.0], [2.0, 42.0], [42.0, 42.0], [None, 3.0], [4.0, 5.0]],
|
||||
),
|
||||
(
|
||||
pa.large_list(pa.float64()),
|
||||
[[1.0, float("nan")], [2.0], None, [None, 3.0], [4.0, 5.0, 6.0]],
|
||||
[[1.0, 42.0], [2.0, 42.0], [42.0, 42.0], [None, 3.0], [4.0, 5.0]],
|
||||
),
|
||||
(
|
||||
pa.list_(pa.float64(), 2),
|
||||
[[1.0, float("nan")], None, [None, 3.0]],
|
||||
[[1.0, 42.0], [42.0, 42.0], [None, 3.0]],
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_fill_bad_vector_values_arrow_types(vector_type, vectors, expected):
|
||||
arr = pa.array([[0.0, 0.0], *vectors, [9.0, 9.0]], type=vector_type)
|
||||
arr = arr.slice(1, len(vectors))
|
||||
|
||||
actual = _fill_bad_vector_values(arr, dim=2, fill_value=42.0)
|
||||
|
||||
assert actual.type == vector_type
|
||||
assert actual.to_pylist() == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"])
|
||||
def test_handle_bad_vectors_nan(on_bad_vectors):
|
||||
vector = pa.array([[1.0, float("nan")], [3.0, 4.0]])
|
||||
@@ -351,7 +393,7 @@ def test_handle_bad_vectors_nan(on_bad_vectors):
|
||||
if on_bad_vectors == "drop":
|
||||
expected = pa.array([[3.0, 4.0]])
|
||||
elif on_bad_vectors == "fill":
|
||||
expected = pa.array([[42.0, 42.0], [3.0, 4.0]])
|
||||
expected = pa.array([[1.0, 42.0], [3.0, 4.0]])
|
||||
elif on_bad_vectors == "null":
|
||||
expected = pa.array([None, [3.0, 4.0]])
|
||||
|
||||
|
||||
Reference in New Issue
Block a user