rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-08-18 12:08:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(python): allow keeping NaN vectors

Browse Source
This commit is contained in:
Gatefixer
2026-08-06 03:37:37 +00:00
parent c7ea91f3ea
commit 626e1a001e
10 changed files with 120 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/python/lancedb/_lancedb.pyi
+1
View File
@@ -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]
python/python/lancedb/db.py
+6 -2
View File
@@ -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
python/python/lancedb/merge.py
+3 -1
View File
@@ -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
python/python/lancedb/remote/db.py
+3 -1
View File
@@ -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".
python/python/lancedb/remote/table.py
+3 -1
View File
@@ -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
python/python/lancedb/table.py
+32 -8
View File
@@ -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=<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=<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):
python/python/lancedb/types.py
+1 -1
View File
@@ -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"]
python/python/tests/test_table.py
+34
View File
@@ -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
+21 -3
View File
@@ -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(
python/src/table.rs
+16 -2
View File
@@ -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<Py<PyAny>>,
write_parallelism: Option<usize>,
on_nan_vectors: Option<String>,
) -> PyResult<Bound<'a, PyAny>> {
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);
}