mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 14:16:54 +00:00
fix(python): honor bad-vector policies for pandas NaNs (#4503)
Pandas converts NaN vector components to Arrow nulls. Both Python's NaN mask and Rust's default append validator ignored null children, so `add()` silently stored incomplete vectors despite `on_bad_vectors="error"`. Reject null components in the existing Rust ingestion validator, and include null components in Python's bad-vector detection and filling. `drop` removes affected rows, `fill` repairs missing components while preserving valid ones, and `null` replaces affected vectors. Explicit null components in Arrow and list inputs follow the same policies. Rust allows whole null vectors and retains the explicit `NaNVectorBehavior::Keep` opt-out. The ingestion guide and synchronous/asynchronous `create_table()` option docs describe null components, all four policies, and filling that preserves valid components. Validation: - Reproduced the default pandas append failure with the regression test before changing the Rust validator. - Rebuilt the native Python extension from the merged sources. The Python table and utility suites passed: 481 passed, 1 skipped. Coverage includes create/append/overwrite, pandas/rows/Arrow inputs, every bad-vector policy, float16/32/64, sliced and chunked arrays, all-null components, and hidden children of null lists. - All 43 focused Rust tests passed: 7 validator tests and 36 ingestion tests. - Rust Clippy with `remote`, tests, and examples passed. Ruff lint passes for the Python package, formatting checks pass for modified Python files, and `cargo fmt --all` is clean. - Documentation follow-up: all 56 focused vector-policy tests passed; Python package lint and creation-docstring formatting checks passed. Repository-wide Ruff reports 49 existing errors in untouched documentation examples. Create/overwrite retain their existing `RuntimeError` wrapping for errors raised while consuming lazy Arrow readers. Fixes #4435 <!-- lance-gatekeeper-fix:v1 agent=b5108e27d28ef613a78c1e4354c03c48 generation=1 --> --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Co-authored-by: Xuanwo <github@xuanwo.io>
This commit is contained in:
co-authored by
Xuanwo
parent
256db55a51
commit
7ed03e1c40
@@ -2170,6 +2170,71 @@ def test_add_with_nans(mem_db: DBConnection):
|
||||
assert np.allclose(filled_vectors[22.0], np.array([5.0, 0.0]))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("operation", ["create", "append", "overwrite"])
|
||||
@pytest.mark.parametrize("input_type", ["pandas", "rows", "arrow"])
|
||||
@pytest.mark.parametrize("on_bad_vectors", [None, "error", "drop", "fill", "null"])
|
||||
def test_write_vectors_with_null_elements(
|
||||
mem_db: DBConnection, operation, input_type, on_bad_vectors
|
||||
):
|
||||
schema = pa.schema({"vector": pa.list_(pa.float32(), 2), "id": pa.int64()})
|
||||
vectors = [[3.0, 4.0], [np.nan, 1.0], [np.nan, np.nan]]
|
||||
if input_type == "pandas":
|
||||
pd = pytest.importorskip("pandas")
|
||||
|
||||
data = pd.DataFrame(
|
||||
{
|
||||
"vector": [np.array(v, dtype=np.float32) for v in vectors],
|
||||
"id": [0, 1, 2],
|
||||
}
|
||||
)
|
||||
else:
|
||||
# Pandas turns the NaNs above into null child elements. Explicit nulls
|
||||
# in other input formats must follow the same bad-vector policy.
|
||||
data = [
|
||||
{"vector": [3.0, 4.0], "id": 0},
|
||||
{"vector": [None, 1.0], "id": 1},
|
||||
{"vector": [None, None], "id": 2},
|
||||
]
|
||||
if input_type == "arrow":
|
||||
data = pa.Table.from_pylist(data, schema=schema)
|
||||
|
||||
options = {} if on_bad_vectors is None else {"on_bad_vectors": on_bad_vectors}
|
||||
if on_bad_vectors == "fill":
|
||||
options["fill_value"] = 42.0
|
||||
|
||||
table = None
|
||||
if operation != "create":
|
||||
table = mem_db.create_table("test", schema=schema)
|
||||
|
||||
def write():
|
||||
if operation == "create":
|
||||
return mem_db.create_table("test", data=data, schema=schema, **options)
|
||||
table.add(data, mode=operation, **options)
|
||||
return table
|
||||
|
||||
if on_bad_vectors in (None, "error"):
|
||||
# Create/overwrite can wrap errors from lazy Arrow readers in RuntimeError.
|
||||
expected_error = (
|
||||
ValueError if operation == "append" else (ValueError, RuntimeError)
|
||||
)
|
||||
with pytest.raises(expected_error, match="NaN|null"):
|
||||
write()
|
||||
if table is not None:
|
||||
assert table.count_rows() == 0
|
||||
return
|
||||
|
||||
table = write()
|
||||
if on_bad_vectors == "drop":
|
||||
expected = [{"vector": [3.0, 4.0], "id": 0}]
|
||||
else:
|
||||
expected = [
|
||||
{"vector": [3.0, 4.0], "id": 0},
|
||||
{"vector": [42.0, 1.0] if on_bad_vectors == "fill" else None, "id": 1},
|
||||
{"vector": [42.0, 42.0] if on_bad_vectors == "fill" else None, "id": 2},
|
||||
]
|
||||
assert table.to_arrow().to_pylist() == expected
|
||||
|
||||
|
||||
def test_add_with_empty_fixed_size_list_drops_bad_rows(mem_db: DBConnection):
|
||||
class Schema(LanceModel):
|
||||
text: str
|
||||
|
||||
Reference in New Issue
Block a user