Merge remote-tracking branch 'origin/main' into gatekeeper/fix-3153-1

docs/src/python/python.md

@@ -17,25 +17,6 @@ The general flow of using the API is:
 pip install lancedb
 ```
 
-The core package does not require PyLance. When you need access to the underlying
-Lance dataset or GPU-accelerated indexing, add the `pylance` extra to the
-distribution you already installed.
-
-For the standard distribution:
-
-```shell
-pip install "lancedb[pylance]"
-```
-
-For the pre-Haswell compatibility distribution:
-
-```shell
-pip install "lancedb-compat[pylance]"
-```
-
-Use only the extra matching your installed distribution. Do not install both
-distributions because they share the `lancedb` namespace.
-
 The following methods describe the synchronous API client. There
 is also an [asynchronous API client](#connections-asynchronous).
 

python/README.md

@@ -8,25 +8,6 @@ A Python library for [LanceDB](https://github.com/lancedb/lancedb).
 pip install lancedb
 ```
 
-The core package does not require PyLance. When you need access to the underlying
-Lance dataset or GPU-accelerated indexing, add the `pylance` extra to the
-distribution you already installed.
-
-For the standard distribution:
-
-```bash
-pip install "lancedb[pylance]"
-```
-
-For the pre-Haswell compatibility distribution:
-
-```bash
-pip install "lancedb-compat[pylance]"
-```
-
-Use only the extra matching your installed distribution. Do not install both
-distributions because they share the `lancedb` namespace.
-
 ### Pre-Haswell x86_64 hosts: `lancedb-compat`
 
 The default `lancedb` wheel targets `x86-64-haswell` (AVX2 + FMA + F16C) for full performance on modern hardware. Pre-Haswell hosts — Intel Sandy Bridge / Ivy Bridge / Westmere; AMD Bulldozer / Piledriver / Steamroller — don't have AVX2 and crash with `Illegal instruction` at `import lancedb`.

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]

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
@@ -1595,7 +1597,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

@@ -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

@@ -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

@@ -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

@@ -117,14 +117,6 @@ _MODEL_BACKED_TOKENIZER_ERRORS = (
     "Failed to initialize default tokenizer",
 )
 
-_PYLANCE_INSTALL_ERROR = (
-    "The lance library is required to use this function. Install the PyLance "
-    "extra for the distribution already installed: "
-    '`pip install "lancedb[pylance]"` for `lancedb`, or '
-    '`pip install "lancedb-compat[pylance]"` for `lancedb-compat`. '
-    "Do not install both distributions because they share the `lancedb` namespace."
-)
-
 
 def _add_unique_note(exception: BaseException, note: str) -> None:
     existing_notes = getattr(exception, "__notes__", ()) or ()
@@ -358,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.
@@ -1255,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
@@ -2257,7 +2253,10 @@ class LanceTable(Table):
         try:
             import lance
         except ImportError:
-            raise ImportError(_PYLANCE_INSTALL_ERROR)
+            raise ImportError(
+                "The lance library is required to use this function. "
+                "Please install with `pip install pylance`."
+            )
 
         branch = self.current_branch()
         version = None if branch is not None else self.version
@@ -3274,7 +3273,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
@@ -3587,7 +3588,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
@@ -4023,7 +4026,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,
@@ -4197,7 +4200,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".
     """
@@ -4262,7 +4267,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(
@@ -4279,6 +4285,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}")
 
@@ -4762,7 +4778,10 @@ class AsyncTable:
         try:
             import lance
         except ImportError:
-            raise ImportError(_PYLANCE_INSTALL_ERROR)
+            raise ImportError(
+                "The lance library is required to use this function. "
+                "Please install with `pip install pylance`."
+            )
 
         # lance.dataset() can't open a branch directly, so open the base table
         # and check out the branch ref (a None branch resolves to main).
@@ -5120,7 +5139,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
@@ -5168,6 +5189,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

@@ -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_import.py

@@ -1,30 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright The LanceDB Authors
-
-import subprocess
-import sys
-
-
-def test_import_lancedb_without_pylance():
-    script = """
-import sys
-
-
-class BlockLanceImports:
-    def find_spec(self, fullname, path=None, target=None):
-        if fullname == "lance" or fullname.startswith("lance."):
-            raise ModuleNotFoundError(f"blocked optional dependency: {fullname}")
-        return None
-
-
-sys.meta_path.insert(0, BlockLanceImports())
-import lancedb
-"""
-
-    result = subprocess.run(
-        [sys.executable, "-c", script],
-        capture_output=True,
-        text=True,
-    )
-
-    assert result.returncode == 0, result.stderr

python/python/tests/test_table.py

@@ -1258,24 +1258,6 @@ def test_branch_to_lance_targets_branch(tmp_path):
     assert table.to_lance().count_rows() == 1
 
 
-def _assert_pylance_install_error(error: ImportError):
-    message = str(error)
-    assert 'pip install "lancedb[pylance]"' in message
-    assert 'pip install "lancedb-compat[pylance]"' in message
-    assert "distribution already installed" in message
-    assert "Do not install both distributions" in message
-
-
-def test_to_lance_recommends_pylance_extra(tmp_db):
-    table = tmp_db.create_table("t", [{"i": 1}])
-
-    with patch("builtins.__import__", side_effect=ImportError):
-        with pytest.raises(ImportError) as exc_info:
-            table.to_lance()
-
-    _assert_pylance_install_error(exc_info.value)
-
-
 @pytest.mark.asyncio
 async def test_async_to_lance(tmp_path):
     pytest.importorskip("lance")
@@ -1287,18 +1269,6 @@ async def test_async_to_lance(tmp_path):
     assert dataset.count_rows() == 1
 
 
-@pytest.mark.asyncio
-async def test_async_to_lance_recommends_pylance_extra(tmp_path):
-    db = await lancedb.connect_async(tmp_path)
-    table = await db.create_table("t", [{"i": 1}])
-
-    with patch("builtins.__import__", side_effect=ImportError):
-        with pytest.raises(ImportError) as exc_info:
-            await table.to_lance()
-
-    _assert_pylance_install_error(exc_info.value)
-
-
 @pytest.mark.asyncio
 async def test_async_branch_to_lance_targets_branch(tmp_path):
     pytest.importorskip("lance")
@@ -1797,6 +1767,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(

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<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);
         }