fix(python): allow keeping NaN vectors

python/python/tests/test_table.py

@@ -1766,6 +1766,40 @@ def test_add_with_nans(mem_db: DBConnection):
     assert np.allclose(filled_vectors[22.0], np.array([5.0, 0.0]))
 
 
+def test_add_with_non_finite_values_keep(mem_db: DBConnection):
+    schema = pa.schema([pa.field("data", pa.list_(pa.float32(), 4))])
+    table = mem_db.create_table("test", schema=schema)
+    batch = pa.table(
+        {
+            "data": pa.array(
+                [[np.nan, np.inf, -np.inf, -0.0]],
+                type=schema.field("data").type,
+            )
+        },
+        schema=schema,
+    )
+
+    with pytest.raises(ValueError, match="NaN"):
+        table.add(batch)
+
+    table.add(batch, on_bad_vectors="keep")
+
+    values = table.to_arrow()["data"][0].as_py()
+    assert np.isnan(values[0])
+    assert np.isposinf(values[1])
+    assert np.isneginf(values[2])
+    assert values[3] == 0.0
+    assert np.signbit(values[3])
+
+
+def test_add_keep_rejects_wrong_dimension(mem_db: DBConnection):
+    schema = pa.schema([pa.field("vector", pa.list_(pa.float32(), 2))])
+    table = mem_db.create_table("test", schema=schema)
+
+    with pytest.raises((ValueError, RuntimeError), match="variable length"):
+        table.add([{"vector": [1.0]}], on_bad_vectors="keep")
+
+
 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

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright The LanceDB Authors
 
 
+import math
 import os
 import pathlib
 from typing import Optional
@@ -364,7 +365,7 @@ def test_fill_bad_vector_values_arrow_types(vector_type, vectors, expected):
     assert actual.to_pylist() == expected
 
 
-@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"])
+@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null", "keep"])
 def test_handle_bad_vectors_nan(on_bad_vectors):
     vector = pa.array([[1.0, float("nan")], [3.0, 4.0]])
     data = pa.table({"vector": vector})
@@ -379,8 +380,9 @@ def test_handle_bad_vectors_nan(on_bad_vectors):
         assert output == (
             "ValueError: Vector column 'vector' has NaNs. 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."
+            "and fill_value=<value> to replace them, set on_bad_vectors='null' "
+            "to replace them with null, or set on_bad_vectors='keep' to preserve "
+            "them."
         )
         return
     else:
@@ -396,10 +398,26 @@ def test_handle_bad_vectors_nan(on_bad_vectors):
         expected = pa.array([[1.0, 42.0], [3.0, 4.0]])
     elif on_bad_vectors == "null":
         expected = pa.array([None, [3.0, 4.0]])
+    elif on_bad_vectors == "keep":
+        actual = output["vector"].to_pylist()
+        assert actual[0][0] == 1.0
+        assert math.isnan(actual[0][1])
+        assert actual[1] == [3.0, 4.0]
+        return
 
     assert output["vector"].combine_chunks() == expected
 
 
+def test_handle_bad_vectors_keep_rejects_wrong_dimension():
+    data = pa.table({"vector": [[1.0, 2.0], [3.0]]})
+
+    with pytest.raises(ValueError, match="only preserves vectors containing NaNs"):
+        _handle_bad_vectors(
+            data.to_reader(),
+            on_bad_vectors="keep",
+        ).read_all()
+
+
 def test_handle_bad_vectors_noop():
     # ChunkedArray should be preserved as-is
     vector = pa.chunked_array(