diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 47e727f99..de9c7aafb 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -269,6 +269,7 @@ class Table: mode: Literal["append", "overwrite"], progress: Optional[Any] = None, write_parallelism: Optional[int] = None, + on_nan_vectors: Optional[Literal["error", "keep"]] = None, ) -> AddResult: ... async def update( self, updates: Dict[str, str], where: Optional[str] diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index b48e84cc5..3c32c974a 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -333,7 +333,9 @@ class DBConnection(EnforceOverrides): 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". + One of "error", "drop", "fill", "null", or "keep". With "keep", + vectors containing NaNs are preserved, but vectors with the wrong + dimension still raise an error. fill_value: float The value to use when filling vectors. Only used if on_bad_vectors="fill". storage_options: dict, optional @@ -1581,7 +1583,9 @@ class AsyncConnection(object): 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". + One of "error", "drop", "fill", "null", or "keep". With "keep", + vectors containing NaNs are preserved, but vectors with the wrong + dimension still raise an error. fill_value: float The value to use when filling vectors. Only used if on_bad_vectors="fill". storage_options: dict, optional diff --git a/python/python/lancedb/merge.py b/python/python/lancedb/merge.py index 115a47e2d..b387d7c69 100644 --- a/python/python/lancedb/merge.py +++ b/python/python/lancedb/merge.py @@ -175,7 +175,9 @@ class LanceMergeInsertBuilder(object): can be anything you use for [`add`][lancedb.table.Table.add] 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". + One of "error", "drop", "fill", "null", or "keep". With "keep", + vectors containing NaNs are preserved, but vectors with the wrong + dimension still raise an error. fill_value: float, default 0. The value to use when filling vectors. Only used if on_bad_vectors="fill". timeout: Optional[timedelta], default None diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 332886590..a8198a3b0 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -543,7 +543,9 @@ class RemoteDBConnection(DBConnection): to "exist_ok". 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". + One of "error", "drop", "fill", "null", or "keep". With "keep", + vectors containing NaNs are preserved, but vectors with the wrong + dimension still raise an error. fill_value: float The value to use when filling vectors. Only used if on_bad_vectors="fill". diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index acc2f4c9d..4c32c9889 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -629,7 +629,9 @@ class RemoteTable(Table): "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". + One of "error", "drop", "fill", "null", or "keep". With "keep", + vectors containing NaNs are preserved, but vectors with the wrong + dimension still raise an error. 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 diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 31a70c298..8b1a83bd4 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -350,8 +350,10 @@ def _sanitize_data( in the input table before casting. metadata : Optional[dict], default None The embedding metadata to add to the schema. - on_bad_vectors : Literal["error", "drop", "fill", "null"], default "error" + on_bad_vectors : Literal["error", "drop", "fill", "null", "keep"], default "error" What to do if any of the vectors are not the same size or contains NaNs. + With "keep", vectors containing NaNs are preserved, but vectors with the + wrong dimension still raise an error. 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. @@ -1247,7 +1249,9 @@ class Table(ABC): "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". + One of "error", "drop", "fill", "null", or "keep". With "keep", + vectors containing NaNs are preserved but are not indexed for vector + search; vectors with the wrong dimension still raise an error. 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 @@ -3265,7 +3269,9 @@ class LanceTable(Table): "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", "null". + One of "error", "drop", "fill", "null", or "keep". With "keep", + vectors containing NaNs are preserved, but vectors with the wrong + dimension still raise an error. 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 @@ -3578,7 +3584,9 @@ class LanceTable(Table): 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", "null". + One of "error", "drop", "fill", "null", or "keep". With "keep", + vectors containing NaNs are preserved, but vectors with the wrong + dimension still raise an error. fill_value: float, default 0. The value to use when filling vectors. Only used if on_bad_vectors="fill". embedding_functions: list of EmbeddingFunctionModel, default None @@ -4014,7 +4022,7 @@ class LanceTable(Table): def _handle_bad_vectors( reader: pa.RecordBatchReader, - on_bad_vectors: Literal["error", "drop", "fill", "null"] = "error", + on_bad_vectors: OnBadVectorsType = "error", fill_value: float = 0.0, target_schema: Optional[pa.Schema] = None, metadata: Optional[dict] = None, @@ -4188,7 +4196,9 @@ def _handle_bad_vector_column( 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. - One of "error", "drop", "fill", "null". + One of "error", "drop", "fill", "null", or "keep". With "keep", + vectors containing NaNs are preserved, but vectors with the wrong dimension + still raise an error. fill_value: float, default 0.0 The value to use when filling vectors. Only used if on_bad_vectors="fill". """ @@ -4253,7 +4263,8 @@ def _handle_bad_vector_column( f"Vector column '{vector_column_name}' has NaNs. " "Set on_bad_vectors='drop' to remove them, " "set on_bad_vectors='fill' and fill_value= to replace them, " - "or set on_bad_vectors='null' to replace them with null." + "set on_bad_vectors='null' to replace them with null, " + "or set on_bad_vectors='keep' to preserve them." ) elif on_bad_vectors == "null": vec_arr = pc.if_else( @@ -4270,6 +4281,16 @@ def _handle_bad_vector_column( "`fill_value` must not be None if `on_bad_vectors` is 'fill'" ) vec_arr = _fill_bad_vector_values(vec_arr, dim, fill_value) + elif on_bad_vectors == "keep": + if pc.any(has_wrong_dim).as_py(): + raise ValueError( + f"Vector column '{vector_column_name}' has variable length " + "vectors. on_bad_vectors='keep' only preserves vectors " + "containing NaNs. Set on_bad_vectors='drop' to remove " + "wrong-size vectors, set on_bad_vectors='fill' and " + "fill_value= to replace them, or set " + "on_bad_vectors='null' to replace them with null." + ) else: raise ValueError(f"Invalid value for on_bad_vectors: {on_bad_vectors}") @@ -5114,7 +5135,9 @@ class AsyncTable: "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", "null". + One of "error", "drop", "fill", "null", or "keep". With "keep", + vectors containing NaNs are preserved but are not indexed for vector + search; vectors with the wrong dimension still raise an error. fill_value: float, default 0. The value to use when filling vectors. Only used if on_bad_vectors="fill". progress: callable or tqdm-like, optional @@ -5162,6 +5185,7 @@ class AsyncTable: mode or "append", progress=progress, write_parallelism=write_parallelism, + on_nan_vectors="keep" if on_bad_vectors == "keep" else None, ) except RuntimeError as e: if "Cast error" in str(e): diff --git a/python/python/lancedb/types.py b/python/python/lancedb/types.py index 567eecb4e..b2cc79fef 100644 --- a/python/python/lancedb/types.py +++ b/python/python/lancedb/types.py @@ -24,7 +24,7 @@ DistanceType = Literal["l2", "cosine", "dot"] DistanceTypeWithHamming = Literal["l2", "cosine", "dot", "hamming"] # Vector handling literals -OnBadVectorsType = Literal["error", "drop", "fill", "null"] +OnBadVectorsType = Literal["error", "drop", "fill", "null", "keep"] # Mode literals AddMode = Literal["append", "overwrite"] diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 542c944c8..c4d85ea4f 100644 --- a/python/python/tests/test_table.py +++ b/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 diff --git a/python/python/tests/test_util.py b/python/python/tests/test_util.py index a9b66b2dd..525eda135 100644 --- a/python/python/tests/test_util.py +++ b/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= to replace them, or set on_bad_vectors='null' " - "to replace them with null." + "and fill_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( diff --git a/python/src/table.rs b/python/src/table.rs index f20bcf2bf..3835afe53 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -21,7 +21,8 @@ use lancedb::blob::{BlobFile, BlobRangeRequest}; use lancedb::index::scalar::FtsIndexBuilder; use lancedb::table::{ AddDataMode, ColumnAlteration, Duration, FieldMetadataUpdate, FtsToken as LanceDbFtsToken, - NewColumnTransform, OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable, + NaNVectorBehavior, NewColumnTransform, OptimizeAction, OptimizeOptions, Ref, + Table as LanceDbTable, }; use lancedb::tokenize as lancedb_tokenize; use pyo3::{ @@ -642,13 +643,14 @@ impl Table { }) } - #[pyo3(signature = (data, mode, progress=None, write_parallelism=None))] + #[pyo3(signature = (data, mode, progress=None, write_parallelism=None, on_nan_vectors=None))] pub fn add<'a>( self_: PyRef<'a, Self>, data: PyScannable, mode: String, progress: Option>, write_parallelism: Option, + on_nan_vectors: Option, ) -> PyResult> { let mut op = self_.inner_ref()?.add(data); if mode == "append" { @@ -658,6 +660,18 @@ impl Table { } else { return Err(PyValueError::new_err(format!("Invalid mode: {}", mode))); } + match on_nan_vectors.as_deref() { + None | Some("error") => {} + Some("keep") => { + op = op.on_nan_vectors(NaNVectorBehavior::Keep); + } + Some(other) => { + return Err(PyValueError::new_err(format!( + "Invalid on_nan_vectors: {}", + other + ))); + } + } if let Some(write_parallelism) = write_parallelism { op = op.write_parallelism(write_parallelism); }