fix(python): honor bad-vector policies for pandas NaNs

python/python/lancedb/table.py

@@ -499,10 +499,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
@@ -1740,8 +1741,8 @@ class Table(ABC):
             The mode to use when writing the data. Valid values are
             "append" and "overwrite".
         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
@@ -3982,7 +3983,7 @@ class LanceTable(Table):
             The mode to use when writing the data. Valid values are
             "append" and "overwrite".
         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".
@@ -4299,7 +4300,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".
@@ -4974,7 +4975,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".
@@ -4993,7 +4994,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))
@@ -5037,7 +5039,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."
@@ -5073,29 +5075,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
@@ -5109,6 +5093,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:
@@ -5119,6 +5104,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)
     has_nan_indices = pc.unique(pc.filter(values_indices, values_has_nan))
     indices = pa.array(range(len(arr)), type=pa.uint32())
@@ -5999,7 +5987,7 @@ class AsyncTable:
             The mode to use when writing the data. Valid values are
             "append" and "overwrite".
         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".

python/python/tests/test_table.py

@@ -1940,6 +1940,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":
+        import pandas as pd
+
+        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

python/python/tests/test_util.py

@@ -20,6 +20,7 @@ from lancedb.table import (
     _infer_target_schema,
     _merge_metadata,
     _sanitize_data,
+    has_nan_values,
     sanitize_create_table,
 )
 import pyarrow as pa
@@ -346,7 +347,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(
             [
@@ -354,12 +355,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
 
@@ -370,17 +371,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]],
         ),
     ],
 )
@@ -407,7 +408,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."
@@ -430,6 +431,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]
+
+
 def test_handle_bad_vectors_noop():
     # ChunkedArray should be preserved as-is
     vector = pa.chunked_array(