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(

rust/lancedb/src/table/add_data.rs

@@ -43,10 +43,12 @@ pub struct AddResult {
 
 #[derive(Debug, Default, Clone, Copy)]
 pub enum NaNVectorBehavior {
-    /// Reject any vectors containing NaN values (the default)
+    /// Reject any vectors containing NaN values or null elements (the default).
+    /// Whole null vectors are allowed.
     #[default]
     Error,
-    /// Allow NaN values to be added, but they will not be indexed for search
+    /// Allow NaN values and null elements to be added, but they will not be indexed
+    /// for search.
     Keep,
 }
 
@@ -102,11 +104,12 @@ impl AddDataBuilder {
         self
     }
 
-    /// Configure how to handle NaN values in vector columns.
+    /// Configure how to handle NaN values and null elements in vector columns.
     ///
-    /// By default, any vectors containing NaN values will be rejected with an
-    /// error, since NaNs cannot be indexed for search. Setting this to `Keep`
-    /// will allow NaN values to be added to the table, but they will not be
+    /// By default, any vectors containing NaN values or null elements will be
+    /// rejected with an error, since these vectors cannot be indexed for search.
+    /// Whole null vectors are allowed. Setting this to `Keep` will allow NaN values
+    /// and null elements to be added to the table, but these vectors will not be
     /// indexed and will not be searchable.
     pub fn on_nan_vectors(mut self, behavior: NaNVectorBehavior) -> Self {
         self.on_nan_vectors = behavior;
@@ -853,8 +856,11 @@ mod tests {
         );
     }
 
+    #[rstest::rstest]
+    #[case::nan(Some(f32::NAN))]
+    #[case::null(None)]
     #[tokio::test]
-    async fn test_add_rejects_nan_vectors() {
+    async fn test_add_rejects_nan_or_null_vectors(#[case] bad_value: Option<f32>) {
         let schema = Arc::new(Schema::new(vec![Field::new(
             "embedding",
             DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4),
@@ -874,7 +880,12 @@ mod tests {
                 FixedSizeListArray::try_new(
                     Arc::new(Field::new("item", DataType::Float32, true)),
                     4,
-                    Arc::new(Float32Array::from(vec![0.1, 0.2, f32::NAN, 0.4])),
+                    Arc::new(Float32Array::from(vec![
+                        Some(0.1),
+                        Some(0.2),
+                        bad_value,
+                        Some(0.4),
+                    ])),
                     None,
                 )
                 .unwrap(),
@@ -884,8 +895,8 @@ mod tests {
         let res = table.add(batch.clone()).execute().await;
         let err = res.unwrap_err();
         assert!(
-            err.to_string().contains("NaN"),
-            "Expected error mentioning NaN values, but got: {err:?}"
+            err.to_string().contains("NaN or null"),
+            "Expected error mentioning NaN or null values, but got: {err:?}"
         );
 
         table

rust/lancedb/src/table/datafusion/reject_nan.rs

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright The LanceDB Authors
 
-//! A DataFusion projection that rejects vectors containing NaN values.
+//! A DataFusion projection that rejects vectors containing NaN or null elements.
 
 use std::sync::{Arc, LazyLock};
 
@@ -31,10 +31,12 @@ fn is_vector_field(field: &Field) -> bool {
     }
 }
 
-/// Wraps the input plan with a projection that checks vector columns for NaN values.
+/// Wraps the input plan with a projection that checks vector columns for NaN or
+/// null elements.
 ///
 /// Non-vector columns pass through unchanged. Vector columns are wrapped with a
-/// UDF that returns the column as-is if no NaNs are present, or errors otherwise.
+/// UDF that returns the column as-is if no NaN or null elements are present, or
+/// errors otherwise. Whole null vectors pass through unchanged.
 pub fn reject_nan_vectors(input: Arc<dyn ExecutionPlan>) -> Result<Arc<dyn ExecutionPlan>> {
     let schema = input.schema();
     let config = Arc::new(ConfigOptions::default());
@@ -70,7 +72,7 @@ pub fn reject_nan_vectors(input: Arc<dyn ExecutionPlan>) -> Result<Arc<dyn Execu
 }
 
 /// A scalar UDF that passes through FixedSizeList arrays unchanged, but errors
-/// if any float values in the list are NaN.
+/// if any float values in a valid list are NaN or null.
 #[derive(Debug, Hash, PartialEq, Eq)]
 struct RejectNanUdf {
     signature: Signature,
@@ -122,11 +124,14 @@ fn check_no_nans(array: &dyn Array) -> datafusion_common::Result<()> {
             )
         })?;
 
-    // Only inspect elements that are both in a valid parent row and non-null
-    // themselves. Values backing null parent rows or null child elements may
-    // contain garbage (including NaN) per the Arrow spec.
-    let has_nan = (0..fsl.len()).filter(|i| fsl.is_valid(*i)).any(|i| {
+    // Values backing null parent rows may contain garbage per the Arrow spec.
+    // A null child in a valid row is an invalid vector component, however:
+    // pandas can represent NaN components as Arrow nulls.
+    let has_bad_values = (0..fsl.len()).filter(|i| fsl.is_valid(*i)).any(|i| {
         let row = fsl.value(i);
+        if row.null_count() > 0 {
+            return true;
+        }
         match row.data_type() {
             DataType::Float16 => row
                 .as_any()
@@ -150,10 +155,10 @@ fn check_no_nans(array: &dyn Array) -> datafusion_common::Result<()> {
         }
     });
 
-    if has_nan {
+    if has_bad_values {
         return Err(datafusion_common::DataFusionError::ArrowError(
             Box::new(arrow_schema::ArrowError::ComputeError(
-                "Vector column contains NaN values".to_string(),
+                "Vector column contains NaN or null values".to_string(),
             )),
             None,
         ));
@@ -165,7 +170,7 @@ fn check_no_nans(array: &dyn Array) -> datafusion_common::Result<()> {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use arrow_array::Float32Array;
+    use arrow_array::{Float16Array, Float32Array, Float64Array};
 
     #[test]
     fn test_passes_clean_vectors() {
@@ -209,23 +214,43 @@ mod tests {
     }
 
     #[test]
-    fn test_skips_null_elements_within_valid_row() {
-        // A valid row with null child elements: the underlying buffer may hold
-        // NaN but the null bitmap says they're absent — should not reject.
-        let values = Float32Array::from(vec![
-            Some(1.0),
-            None, // null element — buffer may contain NaN
-            Some(3.0),
-            None, // null element
-        ]);
+    fn test_rejects_null_elements_within_valid_row() {
+        let arrays: Vec<Arc<dyn Array>> = vec![
+            Arc::new(Float16Array::from(vec![Some(half::f16::ONE), None])),
+            Arc::new(Float32Array::from(vec![Some(1.0), None])),
+            Arc::new(Float64Array::from(vec![Some(1.0), None])),
+        ];
+        for values in arrays {
+            let fsl = FixedSizeListArray::try_new(
+                Arc::new(Field::new("item", values.data_type().clone(), true)),
+                2,
+                values,
+                None,
+            )
+            .unwrap();
+            let err = check_no_nans(&fsl).unwrap_err();
+            assert!(err.to_string().contains("null"));
+        }
+    }
+
+    #[test]
+    fn test_skips_null_elements_in_null_rows_and_outside_slice() {
         let fsl = FixedSizeListArray::try_new(
             Arc::new(Field::new("item", DataType::Float32, true)),
             2,
-            Arc::new(values),
-            None, // both rows are valid
+            Arc::new(Float32Array::from(vec![
+                None,
+                Some(1.0),
+                Some(2.0),
+                Some(3.0),
+                None,
+                None,
+            ])),
+            Some(vec![true, true, false].into()),
         )
         .unwrap();
-        assert!(check_no_nans(&fsl).is_ok());
+        assert!(check_no_nans(&fsl.slice(1, 2)).is_ok());
+        assert!(check_no_nans(&fsl.slice(0, 2)).is_err());
     }
 
     #[test]