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  --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Co-authored-by: Xuanwo <github@xuanwo.io>
2026-10-27 14:06:14 +00:00 · 2026-10-26 17:46:35 +08:00
parent 256db55a51
commit 7ed03e1c40
7 changed files with 214 additions and 77 deletions
@@ -428,10 +428,13 @@ class DBConnection(EnforceOverrides):
            it will not add the provided data but will validate against any
            schema that's specified.
        on_bad_vectors: str, default "error"
-            What to do if any of the vectors are not the same size or contains NaNs.
-            One of "error", "drop", "fill".
+            How to handle vectors with incorrect sizes, NaNs, or null elements.
+            One of "error", "drop", "fill", "null".
        fill_value: float
            The value to use when filling vectors. Only used if on_bad_vectors="fill".
+            NaN or null elements and missing trailing elements are replaced with this
+            value. Valid elements are preserved, and vectors longer than the expected
+            dimension are truncated.
        storage_options: dict, optional
            Additional options for the storage backend. Options already set on the
            connection will be inherited by the table, but can be overridden here.
@@ -2550,10 +2553,13 @@ class AsyncConnection(object):
            it will not add the provided data but will validate against any
            schema that's specified.
        on_bad_vectors: str, default "error"
-            What to do if any of the vectors are not the same size or contains NaNs.
-            One of "error", "drop", "fill".
+            How to handle vectors with incorrect sizes, NaNs, or null elements.
+            One of "error", "drop", "fill", "null".
        fill_value: float
            The value to use when filling vectors. Only used if on_bad_vectors="fill".
+            NaN or null elements and missing trailing elements are replaced with this
+            value. Valid elements are preserved, and vectors longer than the expected
+            dimension are truncated.
        storage_options: dict, optional
            Additional options for the storage backend. Options already set on the
            connection will be inherited by the table, but can be overridden here.
@@ -547,10 +547,11 @@ def _sanitize_data(
    metadata : Optional[dict], default None
        The embedding metadata to add to the schema.
    on_bad_vectors : Literal["error", "drop", "fill", "null"], default "error"
-        What to do if any of the vectors are not the same size or contains NaNs.
+        How to handle vectors with incorrect sizes, NaNs, or null elements.
    fill_value : float, default 0.0
        The value to use when filling vectors. Only used if on_bad_vectors="fill".
-        All entries in the vector will be set to this value.
+        NaN or null elements and missing trailing elements are replaced with this
+        value. Valid elements are preserved.
    """
    # At this point, the table might not match the schema we are targeting:
    # 1. There might be embedding columns missing that will be added
@@ -1898,8 +1899,8 @@ class Table(ABC):
            "append" and "overwrite". Overwrite adopts explicitly declared blob
            layouts, while inferred bytes retain the existing blob v2 fields.
        on_bad_vectors: str, default "error"
-            What to do if any of the vectors are not the same size or contains NaNs.
-            One of "error", "drop", "fill".
+            How to handle vectors with incorrect sizes, NaNs, or null elements.
+            One of "error", "drop", "fill", "null".
        fill_value: float, default 0.
            The value to use when filling vectors. Only used if on_bad_vectors="fill".
        progress: bool, callable, or tqdm-like, optional
@@ -4254,7 +4255,7 @@ class LanceTable(Table):
            "append" and "overwrite". Overwrite adopts explicitly declared blob
            layouts, while inferred bytes retain the existing blob v2 fields.
        on_bad_vectors: str, default "error"
-            What to do if any of the vectors are not the same size or contains NaNs.
+            How to handle vectors with incorrect sizes, NaNs, or null elements.
            One of "error", "drop", "fill", "null".
        fill_value: float, default 0.
            The value to use when filling vectors. Only used if on_bad_vectors="fill".
@@ -4571,7 +4572,7 @@ class LanceTable(Table):
            otherwise just open the table, it will not add the provided
            data but will validate against any schema that's specified.
        on_bad_vectors: str, default "error"
-            What to do if any of the vectors are not the same size or contains NaNs.
+            How to handle vectors with incorrect sizes, NaNs, or null elements.
            One of "error", "drop", "fill", "null".
        fill_value: float, default 0.
            The value to use when filling vectors. Only used if on_bad_vectors="fill".
@@ -5260,7 +5261,7 @@ def _handle_bad_vector_column(
    vector_column_name: str
        The name of the vector column.
    on_bad_vectors: str, default "error"
-        What to do if any of the vectors are not the same size or contains NaNs.
+        How to handle vectors with incorrect sizes, NaNs, or null elements.
        One of "error", "drop", "fill", "null".
    fill_value: float, default 0.0
        The value to use when filling vectors. Only used if on_bad_vectors="fill".
@@ -5279,7 +5280,8 @@ def _handle_bad_vector_column(
        data = data.set_column(position, vector_column_name, vec_arr)

    if expected_value_type is not None and (
-        pa.types.is_integer(vec_arr.type.value_type)
+        pa.types.is_null(vec_arr.type.value_type)
+        or pa.types.is_integer(vec_arr.type.value_type)
        or pa.types.is_unsigned_integer(vec_arr.type.value_type)
    ):
        vec_arr = pa.array(vec_arr.to_pylist(), type=pa.list_(expected_value_type))
@@ -5323,7 +5325,7 @@ def _handle_bad_vector_column(
                )
            else:
                raise ValueError(
-                    f"Vector column '{vector_column_name}' has NaNs. "
+                    f"Vector column '{vector_column_name}' has NaNs or null values. "
                    "Set on_bad_vectors='drop' to remove them, "
                    "set on_bad_vectors='fill' and fill_value=<value> to replace them, "
                    "or set on_bad_vectors='null' to replace them with null."
@@ -5359,29 +5361,11 @@ def _fill_bad_vector_values(
    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.
+    # Both padding and existing null elements need to be filled.
    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
+    needs_fill = pc.is_null(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
@@ -5395,6 +5379,7 @@ def _fill_bad_vector_values(


 def has_nan_values(arr: Union[pa.ListArray, pa.ChunkedArray]) -> pa.BooleanArray:
+    """Find vectors with NaN or null components, including pandas-converted NaNs."""
    if isinstance(arr, pa.ChunkedArray):
        values = pa.chunked_array([chunk.flatten() for chunk in arr.chunks])
    else:
@@ -5405,6 +5390,9 @@ def has_nan_values(arr: Union[pa.ListArray, pa.ChunkedArray]) -> pa.BooleanArray
        values_has_nan = pc.is_nan(values.cast(pa.float32()))
    else:
        values_has_nan = pc.is_nan(values)
+    # Pandas represents NaN components as Arrow nulls. is_nan returns null
+    # for these values, so count them as invalid before filtering parent indices.
+    values_has_nan = pc.fill_null(values_has_nan, True)
    values_indices = pc.list_parent_indices(arr)
    if arr.null_count:
        # Match flatten(), which excludes child spans under null parent lists.
@@ -6318,7 +6306,7 @@ class AsyncTable:
            "append" and "overwrite". Overwrite adopts explicitly declared blob
            layouts, while inferred bytes retain the existing blob v2 fields.
        on_bad_vectors: str, default "error"
-            What to do if any of the vectors are not the same size or contains NaNs.
+            How to handle vectors with incorrect sizes, NaNs, or null elements.
            One of "error", "drop", "fill", "null".
        fill_value: float, default 0.
            The value to use when filling vectors. Only used if on_bad_vectors="fill".
@@ -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
@@ -406,7 +406,7 @@ 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], [None, 9.0]])
+        expected = pa.array([[1.0, 2.0], [4.0, 5.0]])
    elif on_bad_vectors == "fill":
        expected = pa.array(
            [
@@ -414,12 +414,12 @@ def test_handle_bad_vectors_jagged(on_bad_vectors):
                [3.0, 42.0],
                [4.0, 5.0],
                [6.0, 7.0],
-                [None, 9.0],
+                [42.0, 9.0],
                [42.0, 42.0],
            ]
        )
    elif on_bad_vectors == "null":
-        expected = pa.array([[1.0, 2.0], None, [4.0, 5.0], None, [None, 9.0], None])
+        expected = pa.array([[1.0, 2.0], None, [4.0, 5.0], None, None, None])

    assert output["vector"].combine_chunks() == expected

@@ -430,17 +430,17 @@ def test_handle_bad_vectors_jagged(on_bad_vectors):
        (
            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]],
+            [[1.0, 42.0], [2.0, 42.0], [42.0, 42.0], [42.0, 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]],
+            [[1.0, 42.0], [2.0, 42.0], [42.0, 42.0], [42.0, 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]],
+            [[1.0, 42.0], [42.0, 42.0], [42.0, 3.0]],
        ),
    ],
 )
@@ -467,7 +467,7 @@ def test_handle_bad_vectors_nan(on_bad_vectors):
            ).read_all()
        output = exception_output(e)
        assert output == (
-            "ValueError: Vector column 'vector' has NaNs. Set "
+            "ValueError: Vector column 'vector' has NaNs or null values. Set "
            "on_bad_vectors='drop' to remove them, set on_bad_vectors='fill' "
            "and fill_value=<value> to replace them, or set on_bad_vectors='null' "
            "to replace them with null."
@@ -490,6 +490,46 @@ def test_handle_bad_vectors_nan(on_bad_vectors):
    assert output["vector"].combine_chunks() == expected


+@pytest.mark.parametrize("value_type", [pa.float16(), pa.float32(), pa.float64()])
+@pytest.mark.parametrize(
+    "list_type", [pa.list_, pa.large_list, lambda t: pa.list_(t, 2)]
+)
+@pytest.mark.parametrize("chunked", [False, True])
+def test_has_nan_values_with_null_elements(value_type, list_type, chunked):
+    arr = pa.array(
+        [[None, None], [1.0, 2.0], [None, 3.0], None, [4.0, float("nan")], [5.0, 6.0]],
+        type=list_type(value_type),
+    ).slice(1, 5)
+    if chunked:
+        arr = pa.chunked_array([arr.slice(0, 2), arr.slice(2)])
+    # Whole null vectors have no components to inspect. Their handling is
+    # separate from missing or NaN components in a valid vector.
+    assert has_nan_values(arr).to_pylist() == [False, True, False, True, False]
+
+
+@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"])
+def test_sanitize_vectors_with_only_null_elements(on_bad_vectors):
+    data = pa.table({"vector": [[None, None]]})
+    schema = pa.schema({"vector": pa.list_(pa.float32(), 2)})
+
+    def sanitize():
+        return _sanitize_data(
+            data, schema, on_bad_vectors=on_bad_vectors, fill_value=42.0
+        ).read_all()
+
+    if on_bad_vectors == "error":
+        with pytest.raises(ValueError, match="null"):
+            sanitize()
+        return
+
+    expected = {
+        "drop": [],
+        "fill": [{"vector": [42.0, 42.0]}],
+        "null": [{"vector": None}],
+    }
+    assert sanitize().to_pylist() == expected[on_bad_vectors]
+
+
@pytest.mark.parametrize("dtype", [pa.float16(), pa.float32(), pa.float64()])
@pytest.mark.parametrize("large", [False, True])
@pytest.mark.parametrize("chunked", [False, True])