Merge main into blob concurrency repair

python/python/tests/test_blob.py

@@ -7,10 +7,11 @@ import sys
 import textwrap
 
 import lance
+import pandas as pd
 import pyarrow as pa
 import pyarrow.compute as pc
 import pytest
-from lance.blob import BlobType as LanceBlobType
+from lance.blob import Blob, BlobType as LanceBlobType
 
 import lancedb
 from lancedb._blob import (
@@ -921,6 +922,70 @@ def test_fetch_blob_files_lazy_read():
     handles = table.fetch_blob_files("image", [by_id[1]])
     assert len(handles) == 1
     assert handles[0].read() == payload
+    handles[0].close()
+    assert handles[0].closed
+    with pytest.raises(RuntimeError, match="already closed"):
+        handles[0].read()
+
+
+def test_blob_file_finalizer_and_close_on_runtime_worker():
+    # A tokio panic in IOBase.__del__ is printed but does not fail the add call.
+    # Run in a child process so the assertion catches that panic on stderr.
+    script = textwrap.dedent(
+        """\
+        import gc
+        import threading
+
+        import lancedb
+        import pyarrow as pa
+
+        db = lancedb.connect("memory:///")
+        schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
+        table = db.create_table("finalize", schema=schema)
+        table.add([{"id": 1, "image": b"first"}])
+        hits = table.search().with_row_id(True).limit(1).to_arrow()
+        row_id = hits["_rowid"][0].as_py()
+        gc_handle, explicit_handle = table.fetch_blob_files("image", [row_id, row_id])
+
+        gc.disable()
+        cycle = [gc_handle]
+        cycle.append(cycle)
+        del gc_handle, cycle
+
+        observed = []
+        def on_progress(_):
+            if not observed:
+                observed.append((
+                    threading.current_thread() is not threading.main_thread(),
+                    explicit_handle.closed,
+                    gc.collect(),
+                ))
+                explicit_handle.close()
+
+        table.add([{"id": 2, "image": b"second"}], progress=on_progress)
+        gc.enable()
+
+        assert len(observed) == 1, observed
+        assert observed[0][0] is True, observed
+        assert observed[0][1] is False, observed
+        assert observed[0][2] > 0, observed
+        assert explicit_handle.closed
+        explicit_handle.close()  # Close remains idempotent after the worker call.
+        try:
+            explicit_handle.read()
+        except RuntimeError as error:
+            assert "already closed" in str(error)
+        else:
+            raise AssertionError("a closed blob file was readable")
+        print("worker close completed")
+        """
+    )
+    result = subprocess.run(
+        [sys.executable, "-c", script], capture_output=True, text=True, check=False
+    )
+    assert result.returncode == 0, result.stderr
+    assert "Cannot start a runtime from within a runtime" not in result.stderr
+    assert "worker close completed" in result.stdout
 
 
 def test_fetch_blob_files_null_alignment():
@@ -1368,3 +1433,129 @@ def test_add_external_uri_string_round_trips_with_flag(tmp_path):
     hits = table.search().to_arrow()
     blobs = table.fetch_blobs("image", hits)
     assert blobs[0].as_py() == payload
+
+
+@pytest.mark.parametrize("as_pandas", [False, True])
+def test_add_bytes_and_uri_in_one_batch_preserves_external_uri(tmp_path, as_pandas):
+    payload = b"external-payload"
+    blob_path = tmp_path / "payload.bin"
+    blob_path.write_bytes(payload)
+    rows = [
+        {"id": 1, "image": b"inline"},
+        {"id": 2, "image": blob_path.as_uri()},
+    ]
+    db = lancedb.connect(tmp_path / "db")
+    schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
+    table = db.create_table("mixed", schema=schema)
+
+    table.add(
+        pd.DataFrame(rows) if as_pandas else rows,
+        allow_external_blob_outside_bases=True,
+    )
+
+    by_id = _row_ids_by_id(table)
+    assert table.fetch_blobs("image", [by_id[1], by_id[2]]).to_pylist() == [
+        b"inline",
+        payload,
+    ]
+
+
+@pytest.mark.parametrize("as_pandas", [False, True])
+def test_add_mixed_python_blob_values_preserves_external_uris(tmp_path, as_pandas):
+    payload = b"external-payload"
+    blob_path = tmp_path / "payload.bin"
+    blob_path.write_bytes(payload)
+    uri = blob_path.as_uri()
+    rows = [
+        {"id": 1, "image": b"inline"},
+        {"id": 2, "image": uri},
+        {"id": 3, "image": {"data": b"from-dict"}},
+        {"id": 4, "image": Blob.from_bytes(b"from-blob")},
+        {"id": 5, "image": Blob.from_uri(uri)},
+        {"id": 6, "image": {"uri": uri}},
+        {"id": 7, "image": None},
+    ]
+    db = lancedb.connect(tmp_path / "db")
+    schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
+    table = db.create_table("mixed", schema=schema)
+
+    table.add(
+        pd.DataFrame(rows) if as_pandas else rows,
+        allow_external_blob_outside_bases=True,
+    )
+
+    by_id = _row_ids_by_id(table)
+    assert table.fetch_blobs("image", [by_id[i] for i in range(1, 8)]).to_pylist() == [
+        b"inline",
+        payload,
+        b"from-dict",
+        b"from-blob",
+        payload,
+        payload,
+        None,
+    ]
+
+
+@pytest.mark.parametrize("as_pandas", [False, True])
+def test_merge_insert_mixed_python_blobs_validates_external_uri(tmp_path, as_pandas):
+    blob_path = tmp_path / "payload.bin"
+    blob_path.write_bytes(b"external-payload")
+    table = _blob_table("merge_mixed_uri", [{"id": 1, "image": b"before"}])
+    rows = [
+        {"id": 1, "image": b"updated"},
+        {"id": 2, "image": blob_path.as_uri()},
+    ]
+
+    with pytest.raises(ValueError, match="allow_external_blob_outside_bases"):
+        (
+            table.merge_insert("id")
+            .when_matched_update_all()
+            .when_not_matched_insert_all()
+            .execute(pd.DataFrame(rows) if as_pandas else rows)
+        )
+    assert table.count_rows() == 1
+
+
+@pytest.mark.parametrize("as_pandas", [False, True])
+def test_merge_insert_python_blob_dicts_and_objects(as_pandas):
+    table = _blob_table("merge_python_blobs", [{"id": 1, "image": b"before"}])
+    rows = [
+        {"id": 1, "image": b"updated"},
+        {"id": 2, "image": {"data": b"from-dict"}},
+        {"id": 3, "image": Blob.from_bytes(b"from-blob")},
+    ]
+
+    result = (
+        table.merge_insert("id")
+        .when_matched_update_all()
+        .when_not_matched_insert_all()
+        .execute(pd.DataFrame(rows) if as_pandas else rows)
+    )
+
+    assert result.num_updated_rows == 1
+    assert result.num_inserted_rows == 2
+    by_id = _row_ids_by_id(table)
+    assert table.fetch_blobs("image", [by_id[i] for i in range(1, 4)]).to_pylist() == [
+        b"updated",
+        b"from-dict",
+        b"from-blob",
+    ]
+
+
+def test_create_table_with_python_blob_dict_and_object():
+    db = lancedb.connect("memory:///")
+    schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
+    table = db.create_table(
+        "create_python_blobs",
+        data=[
+            {"id": 1, "image": {"data": b"from-dict"}},
+            {"id": 2, "image": Blob.from_bytes(b"from-blob")},
+        ],
+        schema=schema,
+    )
+
+    by_id = _row_ids_by_id(table)
+    assert table.fetch_blobs("image", [by_id[1], by_id[2]]).to_pylist() == [
+        b"from-dict",
+        b"from-blob",
+    ]

python/python/tests/test_function_packaging.py

@@ -11,6 +11,7 @@ import json
 import subprocess
 import sys
 import textwrap
+import warnings
 from pathlib import Path
 from typing import Annotated, Optional
 
@@ -299,9 +300,11 @@ def test_code_rejects_submodules_duplicates_and_non_modules(helper_package):
                 return value
 
 
-def test_local_module_outside_code_is_rejected_at_registration(
-    tmp_path, helper_package
-):
+def test_local_module_outside_code_warns_at_registration(tmp_path, helper_package):
+    """The worker most likely cannot import a local module that is neither
+    shipped nor declared, but registration still succeeds: a declared
+    distribution may install it under another name, and callers may rewrite
+    the packaged source before registering it."""
     helper_package(
         "fx_local", {"__init__.py": "def normalize(value):\n    return value\n"}
     )
@@ -322,11 +325,14 @@ def test_local_module_outside_code_is_rejected_at_registration(
         )
         return getattr(importlib.import_module(name), name)
 
-    with pytest.raises(ValueError, match=r"local module.*code=\[fx_local\]"):
-        define("fx_uses_local", "")
+    with pytest.warns(UserWarning, match=r"local module.*code=\[fx_local\]"):
+        local = define("fx_uses_local", "")
+    assert "from fx_local import normalize" in artifact(local).decode()
 
     # A declared package that provides the module is trusted to install it.
-    installed = define("fx_uses_installed", "pip=['fx-local==1.0']")
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        installed = define("fx_uses_installed", "pip=['fx-local==1.0']")
     assert "from fx_local import normalize" in artifact(installed).decode()
 
 

python/python/tests/test_hybrid_query.py

@@ -315,6 +315,33 @@ def test_hybrid_query_offset(sync_table: Table):
     assert offset_result["_rowid"].to_pylist() == full["_rowid"].to_pylist()[2:]
 
 
+def test_hybrid_query_default_limit(sync_table: Table):
+    # The vector and FTS legs match disjoint rows, so fusing their results
+    # yields more rows than the default limit.
+    new_rows = []
+    for i in range(20):
+        new_rows.append({"text": "close_vec", "vector": [0.1, 0.1]})
+        new_rows.append({"text": "dog", "vector": [50.0 + i, 50.0 + i]})
+    sync_table.add(new_rows)
+
+    def query():
+        return (
+            sync_table.search(query_type="hybrid")
+            .vector([0.1, 0.1])
+            .text("dog")
+            .with_row_id(True)
+        )
+
+    # Like the async hybrid query, the default limit is 10.
+    result = query().to_arrow()
+    assert len(result) == 10
+
+    # The offset window is taken after the default limit is applied.
+    full = query().limit(15).to_arrow()
+    offset_result = query().offset(5).to_arrow()
+    assert offset_result["_rowid"].to_pylist() == full["_rowid"].to_pylist()[5:]
+
+
 def test_hybrid_query_minimum_nprobes_zero_raises(sync_table: Table):
     # minimum_nprobes(0) must raise the same validation error a plain vector
     # query raises, not silently no-op because 0 is falsy.

python/python/tests/test_materialized_views.py

@@ -520,10 +520,166 @@ def test_stored_queries_and_legacy_layouts_are_read():
         "SELECT * FROM people"
     )
 
+    for fmt, function in [(3, "vector_duplicate_pairs"), (4, "vector_dedup")]:
+        native = f"SELECT * FROM {function}('images', 3, 'phash', 4)"
+        assert read({"format": fmt, "query": native}).query == native
+
     # A newer writer's layout is reported, never guessed at.
     for newer in (
-        {"format": 3, "query": query},
+        {"format": 5, "query": query},
         {"kind": "select_v3", "source_table": "people"},
     ):
         with pytest.raises(NotImplementedError, match="cannot refresh"):
             read(newer)
+
+
+def indexed_dedup_source(db, name="images"):
+    import pyarrow as pa
+    from lancedb.index import IvfFlat
+
+    vectors = pa.FixedSizeListArray.from_arrays(
+        pa.array(
+            [v for i in range(256) for v in (float(i if i < 3 else i * 10), 0.0)],
+            type=pa.float32(),
+        ),
+        2,
+    )
+    table = db.create_table(name, pa.table({"id": range(256), "vector": vectors}))
+    table.create_index("vector", config=IvfFlat(num_partitions=1))
+    return table
+
+
+@pytest.mark.parametrize("stable", [False, True])
+def test_vector_dedup_source_captures_snapshot_and_preserves_source(tmp_path, stable):
+    db = lancedb.connect(
+        tmp_path,
+        storage_options={"new_table_enable_stable_row_ids": str(stable).lower()},
+    )
+    source = indexed_dedup_source(db, "images.with.dot")
+    version = source.version
+    expression = lancedb.vector_dedup(
+        source.name, column="vector", distance_threshold=1
+    )
+    view = db.create_materialized_view("clean", expression)
+    assert "vector_dedup(" in view.definition.query
+    assert f", {version}, 'vector', 1.0)" in view.definition.query
+    assert view.table.count_rows() == 255
+    assert [r["id"] for r in view.table.search().where("id < 3").to_list()] == [0, 2]
+    assert source.version == version
+    assert source.count_rows() == 256
+
+    source.delete("id = 0")
+    view.refresh_async().wait()
+    assert view.table.search().where("id = 0").to_list() == [
+        {"id": 0, "vector": [0.0, 0.0]}
+    ]
+    assert expression.dataset_version is None  # reusable immutable declaration
+
+
+def test_vector_dedup_create_job_and_explicit_snapshot(tmp_path):
+    db = lancedb.connect(tmp_path)
+    source = indexed_dedup_source(db)
+    version = source.version
+    source.delete("id = 0")
+    job = db.create_materialized_view_async(
+        "clean",
+        lancedb.vector_dedup(
+            "images", column="vector", distance_threshold=1, dataset_version=version
+        ),
+    )
+    assert job.wait() is None
+    view = db.open_materialized_view("clean")
+    assert view.table.count_rows() == 255
+    assert len(view.table.search().where("id = 0").to_list()) == 1
+
+
+@pytest.mark.asyncio
+async def test_async_vector_dedup_declaration_and_creation(tmp_path):
+    sync_db = lancedb.connect(tmp_path)
+    source = indexed_dedup_source(sync_db)
+    db = await lancedb.connect_async(tmp_path)
+    expression = lancedb.vector_dedup("images", column="vector", distance_threshold=1)
+    view = await db.create_materialized_view("clean", expression, with_no_data=True)
+    assert await view.table.count_rows() == 0
+    # The snapshot is captured at declaration even if refresh comes later.
+    source.delete("id = 0")
+    job = await view.refresh_async()
+    await job.wait()
+    assert await view.table.count_rows() == 255
+    another = await db.create_materialized_view_async(
+        "later", expression, with_no_data=True
+    )
+    await another.wait()
+    later = await db.open_materialized_view("later")
+    assert (await later.definition()).query != (await view.definition()).query
+    assert await later.table.count_rows() == 0
+
+
+@pytest.mark.parametrize(
+    "options", [{"select": ["id"]}, {"select": []}, {"where": "id > 0"}, {"limit": 1}]
+)
+def test_vector_dedup_rejects_extra_clauses_before_opening_source(tmp_path, options):
+    db = lancedb.connect(tmp_path)
+    expression = lancedb.vector_dedup("missing", column="vector", distance_threshold=1)
+    with pytest.raises(ValueError, match="cannot be combined"):
+        db.create_materialized_view("bad", expression, **options)
+    assert db.list_tables().tables == []
+
+
+@pytest.mark.parametrize("version", [0, -1, True, 1.5, "2", 2**64])
+def test_vector_dedup_rejects_invalid_snapshot(version):
+    with pytest.raises(ValueError, match="dataset_version"):
+        lancedb.vector_dedup(
+            "images", column="vector", distance_threshold=1, dataset_version=version
+        )
+
+
+@pytest.mark.parametrize(
+    "threshold", [float("nan"), float("inf"), -float("inf"), 1e100, 10**1000, True, "4"]
+)
+def test_vector_dedup_rejects_invalid_threshold(threshold):
+    with pytest.raises((ValueError, TypeError), match="distance_threshold"):
+        lancedb.vector_dedup("images", column="vector", distance_threshold=threshold)
+
+
+def test_vector_dedup_descriptor_is_frozen():
+    from dataclasses import FrozenInstanceError
+
+    expression = lancedb.vector_dedup("images", column="vector", distance_threshold=1)
+    with pytest.raises(FrozenInstanceError):
+        expression.column = "other"
+
+
+@pytest.mark.parametrize("with_no_data", [False, True])
+def test_remote_vector_dedup_uses_mv_endpoint_without_sql_client(with_no_data):
+    with mock_remote_materialized_view_create() as (host, requests):
+        db = lancedb.connect(
+            "db://dev",
+            api_key="fake",
+            host_override=host,
+            sql_host_override="grpc://127.0.0.1:1",
+            client_config={"retry_config": {"retries": 0}},
+        )
+        job = db.create_materialized_view_async(
+            "clean",
+            lancedb.vector_dedup(
+                "images.with'dot",
+                column="vector'field",
+                distance_threshold=4,
+                dataset_version=7,
+            ),
+            with_no_data=with_no_data,
+        )
+        assert job.id == "mv-create-123"
+    assert requests == [
+        (
+            "/v1/materialized_view/clean/create",
+            {
+                "query": (
+                    "SELECT * FROM vector_dedup('`images.with''dot`', "
+                    "7, 'vector''field', 4.0)"
+                ),
+                "with_no_data": with_no_data,
+            },
+        )
+    ]

python/python/tests/test_remote_db.py

@@ -2408,11 +2408,16 @@ def blob_remote_table(*, server_version=Version("0.5.0")):
         elif request.path.startswith("/v1/table/test/blob/image/"):
             path = request.path.partition("?")[0]
             row_id = int(path.split("/")[-2])
-            payload = {10: b"alpha", 20: None, 30: b"gamma"}[row_id]
+            payload = {10: b"alpha", 20: None, 30: b"gamma", 40: b""}[row_id]
             if payload is None:
                 request.send_response(204)
                 request.end_headers()
                 return
+            if not payload:
+                request.send_response(416)
+                request.send_header("Content-Range", "bytes */0")
+                request.end_headers()
+                return
             byte_range = request.headers["Range"].removeprefix("bytes=")
             start_text, end_text = byte_range.split("-", maxsplit=1)
             start = int(start_text)
@@ -2421,6 +2426,9 @@ def blob_remote_table(*, server_version=Version("0.5.0")):
             request.send_response(206)
             request.send_header("Content-Range", f"bytes {start}-{end}/{len(payload)}")
             request.send_header("Content-Length", str(len(chunk)))
+            request.send_header(
+                "x-lancedb-version", str(BLOB_DESCRIBE_RESPONSE["version"])
+            )
             request.end_headers()
             request.wfile.write(chunk)
         elif request.path == "/v1/table/test/query/":
@@ -2464,17 +2472,25 @@ def test_remote_blob_columns_and_fetch():
 
 def test_remote_blob_files_are_lazy_seekable_handles():
     with blob_remote_table() as table:
-        files = table.fetch_blob_files("image", [10, 20, 30])
+        files = table.fetch_blob_files("image", [10, 20, 30, 40])
 
-        assert len(files) == 3
-        alpha, null_row, gamma = files
+        assert len(files) == 4
+        alpha, null_row, gamma, empty = files
         assert null_row is None
         assert alpha is not None
         assert gamma is not None
+        assert empty is not None
         assert alpha.size() == 5
         assert alpha.read_range(1, 3) == b"lph"
         gamma.seek(2)
         assert gamma.read() == b"mma"
+        assert empty.size() == 0
+        assert empty.read() == b""
+        assert empty.read_range(0, 0) == b""
+        alpha.close()
+        assert alpha.closed
+        with pytest.raises(RuntimeError, match="already closed"):
+            alpha.read_range(0, 1)
 
 
 def test_remote_blob_fetch_accepts_query_table():
@@ -2629,6 +2645,26 @@ def test_remote_connection_jobs_surface():
             request.send_header("Content-Type", "application/json")
             request.end_headers()
             request.wfile.write(b'{"job_id": "job-1"}')
+        elif request.path == "/v1/jobs/pause":
+            if payload["job_id"] != "job-1":
+                request.send_response(404)
+                request.end_headers()
+                return
+            request.send_response(200)
+            request.send_header("Content-Type", "application/json")
+            request.end_headers()
+            request.wfile.write(b'{"job_id": "job-1", "paused": true}')
+        elif request.path == "/v1/jobs/resume":
+            if payload["job_id"] != "job-1":
+                request.send_response(404)
+                request.end_headers()
+                return
+            request.send_response(200)
+            request.send_header("Content-Type", "application/json")
+            request.end_headers()
+            request.wfile.write(
+                b'{"job_id": "job-1", "resumed": false, "still_pausing": true}'
+            )
         elif request.path == "/v1/jobs/query_events":
             query_events_payloads.append(payload)
             request.send_response(200)
@@ -2649,6 +2685,9 @@ def test_remote_connection_jobs_surface():
         assert db.cancel_job("job-1") is True
         assert db.cancel_job("missing") is False
 
+        assert db.pause_job("job-1") == "pausing"
+        assert db.resume_job("job-1") == "still_pausing"
+
         # Opening a job hands back a populated handle; a missing one fails.
         with pytest.raises(JobNotFoundError, match="missing"):
             db.open_job("missing")

python/python/tests/test_rerankers.py

@@ -395,6 +395,38 @@ def test_mrr_multivector_rewards_consensus():
     assert result["_rowid"].to_pylist()[0] == 2
 
 
+def test_rrf_multivector_ranks_each_result_list():
+    # RRF scores a document by its rank *within each* result list. The top hit
+    # of the second list must score the same as the top hit of the first list,
+    # not as if it were ranked after every row of the first list.
+    reranker = RRFReranker(K=60)
+
+    def ranking(row_ids):
+        return pa.table({"_rowid": pa.array(row_ids, type=pa.int64())})
+
+    rs1 = ranking([1, 2, 3])
+    rs2 = ranking([4, 5, 6])
+
+    result = reranker.rerank_multivector([rs1, rs2])
+    scores = dict(
+        zip(result["_rowid"].to_pylist(), result["_relevance_score"].to_pylist())
+    )
+
+    assert scores[1] == pytest.approx(1 / 61)
+    assert scores[4] == pytest.approx(1 / 61)
+    assert scores[5] == pytest.approx(1 / 62)
+    # the second-ranked hit of the second list beats the third hit of the first
+    assert scores[5] > scores[3]
+
+    # A document found by both lists sums its reciprocal ranks from each list.
+    result = reranker.rerank_multivector([ranking([1, 2]), ranking([2, 3])])
+    scores = dict(
+        zip(result["_rowid"].to_pylist(), result["_relevance_score"].to_pylist())
+    )
+    assert scores[2] == pytest.approx(1 / 62 + 1 / 61)
+    assert result["_rowid"].to_pylist() == [2, 1, 3]
+
+
 def test_rrf_reranker_distance():
     data = pa.table(
         {
@@ -481,6 +513,14 @@ def test_cross_encoder_reranker(tmp_path):
     _run_test_reranker(reranker, table, "single player experience", None, schema)
 
 
+def test_cross_encoder_reranker_kwargs(tmp_path):
+    pytest.importorskip("sentence_transformers")
+    reranker = CrossEncoderReranker(batch_size=4, max_length=64)
+    assert reranker.model.max_length == 64
+    table, schema = get_test_table(tmp_path)
+    _run_test_reranker(reranker, table, "single player experience", None, schema)
+
+
 def test_colbert_reranker(tmp_path):
     pytest.importorskip("rerankers")
     reranker = ColbertReranker()
@@ -536,28 +576,43 @@ class _FakeTypeSafeClient:
     def system_one(self, state, questions, model):
         self.requests.append((state, questions, model))
         query_words = set(state["query"].lower().split())
-        doc_words = set(state["document"].lower().split())
-        noul = len(query_words & doc_words) / len(query_words)
-        answers = {key: type("NoulAnswer", (), {"noul": noul})() for key in questions}
+        answers = {}
+        for key, question in questions.items():
+            document = (
+                state["document"]
+                if "document" in state
+                else question["instructions"]["document"]
+            )
+            doc_words = set(document.lower().split())
+            noul = len(query_words & doc_words) / len(query_words)
+            answers[key] = type("NoulAnswer", (), {"noul": noul})()
         return type("SystemOneResponse", (), {"answers": answers})()
 
 
-def test_typesafe_reranker_with_fake_client(tmp_path):
-    reranker = TypeSafeReranker(max_concurrency=4)
+@pytest.mark.parametrize("batch_size", [1, 40])
+@pytest.mark.parametrize("return_score", ["relevance", "all"])
+def test_typesafe_reranker_with_fake_client(tmp_path, batch_size, return_score):
+    reranker = TypeSafeReranker(
+        max_concurrency=4, batch_size=batch_size, return_score=return_score
+    )
     reranker._client = _FakeTypeSafeClient()
     table, schema = get_test_table(tmp_path)
     _run_test_reranker(reranker, table, "single player experience", None, schema)
 
     state, questions, model = reranker._client.requests[0]
     assert model == "jev-latest"
-    assert set(state) == {"query", "document"}
-    assert questions == {
-        "relevance": {
-            "type": "noul",
-            "instructions": reranker.instructions,
-            "criteria": reranker.criteria,
+    if batch_size == 1:
+        assert set(state) == {"query", "document"}
+        assert questions == {
+            "relevance": {
+                "type": "noul",
+                "instructions": reranker.instructions,
+                "criteria": reranker.criteria,
+            }
         }
-    }
+    else:
+        assert set(state) == {"query"}
+        assert 1 <= len(questions) <= batch_size
 
 
 def test_typesafe_reranker_scores_each_row():

python/python/tests/test_table.py

@@ -4307,7 +4307,8 @@ async def test_optimize(mem_db_async: AsyncConnection):
     assert stats.prune.bytes_removed == 0
     assert stats.prune.old_versions_removed == 0
 
-    stats = await table.optimize(cleanup_older_than=timedelta(seconds=0))
+    with pytest.warns(UserWarning, match="concurrent"):
+        stats = await table.optimize(cleanup_older_than=timedelta(seconds=0))
     assert stats.prune.bytes_removed > 0
     assert stats.prune.old_versions_removed == 3
 
@@ -4335,12 +4336,35 @@ async def test_optimize_delete_unverified(tmp_db_async: AsyncConnection, tmp_pat
 
     stats = await table.optimize(delete_unverified=False)
     assert stats.prune.old_versions_removed == 0
-    stats = await table.optimize(
-        cleanup_older_than=timedelta(seconds=0), delete_unverified=True
-    )
+    with pytest.warns(UserWarning, match="concurrent"):
+        stats = await table.optimize(
+            cleanup_older_than=timedelta(seconds=0), delete_unverified=True
+        )
     assert stats.prune.old_versions_removed == 2
 
 
+@pytest.mark.asyncio
+async def test_optimize_warns_on_zero_cleanup(mem_db_async: AsyncConnection):
+    table = await mem_db_async.create_table("test", data=[{"x": [1]}])
+    with pytest.warns(UserWarning, match="concurrent"):
+        await table.optimize(cleanup_older_than=timedelta(0))
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        await table.optimize(cleanup_older_than=timedelta(days=1))
+        await table.optimize()
+
+
+def test_optimize_warns_on_zero_cleanup_sync(mem_db: DBConnection):
+    table = mem_db.create_table("test", data=[{"x": [1]}])
+    with warnings.catch_warnings(record=True) as seen:
+        warnings.simplefilter("default")
+        table.optimize(cleanup_older_than=timedelta(0))
+        table.optimize(cleanup_older_than=timedelta(0))
+    assert [w.filename for w in seen] == [__file__, __file__]
+    assert all("concurrent" in str(w.message) for w in seen)
+
+
 def test_replace_field_metadata(tmp_path):
     db = lancedb.connect(tmp_path)
     table = db.create_table("my_table", data=[{"x": 0}])

python/python/tests/test_typesafe_reranker.py

@@ -0,0 +1,338 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright The LanceDB Authors
+
+import json
+from threading import Barrier, Event, Lock
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import lancedb
+import pyarrow as pa
+import pytest
+from lancedb.index import FTS
+from lancedb.rerankers import TypeSafeReranker
+
+
+def response(questions, probability=0.5):
+    return SimpleNamespace(
+        answers={key: SimpleNamespace(noul=probability) for key in reversed(questions)}
+    )
+
+
+def vector_results(documents):
+    return pa.table(
+        {
+            "text": pa.array(documents, type=pa.string()),
+            "_rowid": pa.array(range(len(documents)), type=pa.uint64()),
+            "_distance": pa.array(range(len(documents)), type=pa.float32()),
+        }
+    )
+
+
+def fake_reranker(**kwargs):
+    reranker = TypeSafeReranker(**kwargs)
+    reranker._client = Mock()
+    reranker._client.system_one.side_effect = lambda **kw: response(kw["questions"])
+    return reranker
+
+
+@pytest.mark.parametrize("batch_size", [0, -1, 1.5, 40.0, True, False, "40", None])
+def test_invalid_batch_size(batch_size):
+    with pytest.raises(ValueError, match="batch_size must be a positive integer"):
+        TypeSafeReranker(batch_size=batch_size)
+
+
+@pytest.mark.parametrize("batch_size", [1, 7, 40])
+@pytest.mark.parametrize("count", [0, 1, 39, 40, 41, 80, 81])
+def test_request_counts_and_boundaries(batch_size, count):
+    reranker = fake_reranker(batch_size=batch_size)
+    docs = [f"candidate {i}" for i in range(count)]
+    result = reranker.rerank_vector("query", vector_results(docs))
+    calls = reranker._client.system_one.call_args_list
+    assert len(calls) == (count + batch_size - 1) // batch_size
+    assert sorted(len(call.kwargs["questions"]) for call in calls) == sorted(
+        min(batch_size, count - start) for start in range(0, count, batch_size)
+    )
+    assert result["_rowid"].to_pylist() == list(range(count))
+
+
+@pytest.mark.parametrize("batch_size", [1, 40])
+@pytest.mark.parametrize(
+    "criteria", [{}, {"true": "  state.document fits\nexactly", "false": "No!"}]
+)
+def test_payload_isolation_and_custom_prompts(batch_size, criteria):
+    instructions = "Read state.document and state.query.\n Preserve this text. "
+    reranker = fake_reranker(
+        batch_size=batch_size,
+        instructions=instructions,
+        criteria=criteria,
+        model_name="jev-1.13.0",
+        column="body",
+    )
+    docs = [f"UNIQUE_CANDIDATE_{i}" for i in range(41)]
+    table = vector_results(docs).rename_columns(["body", "_rowid", "_distance"])
+    reranker.rerank_vector("shared query", table)
+    seen = []
+    for call in reranker._client.system_one.call_args_list:
+        state, questions = call.kwargs["state"], call.kwargs["questions"]
+        assert call.kwargs["model"] == "jev-1.13.0"
+        if batch_size == 1:
+            doc = state["document"]
+            assert state == {"query": "shared query", "document": doc}
+            assert list(questions) == ["relevance"]
+            expected_instructions = instructions
+        else:
+            assert state == {"query": "shared query"}
+        for question in questions.values():
+            if batch_size > 1:
+                doc = question["instructions"]["document"]
+                expected_instructions = {"question": instructions, "document": doc}
+            expected = {"type": "noul", "instructions": expected_instructions}
+            if criteria:
+                expected["criteria"] = criteria
+            # Exact equality rules out other candidates anywhere in this question.
+            assert question == expected
+            seen.append(doc)
+    assert sorted(seen) == sorted(docs)
+    assert reranker.instructions == instructions
+    assert reranker.criteria == criteria
+
+
+def test_default_payload():
+    reranker = fake_reranker()
+    reranker.rerank_vector("query", vector_results(["document"]))
+    reranker._client.system_one.assert_called_once_with(
+        state={"query": "query", "document": "document"},
+        questions={
+            "relevance": {
+                "type": "noul",
+                "instructions": reranker.instructions,
+                "criteria": reranker.criteria,
+            }
+        },
+        model="jev-latest",
+    )
+
+
+def test_response_mapping_with_duplicates_and_nulls():
+    reranker = fake_reranker(batch_size=2, max_concurrency=2, return_score="all")
+    second_batch_scored = Event()
+
+    def score(**kw):
+        questions = kw["questions"]
+        documents = [
+            question["instructions"]["document"] for question in questions.values()
+        ]
+        if documents == ["same", "same"]:
+            # Prepare the later batch first; result order must still match input.
+            assert second_batch_scored.wait(timeout=10)
+            scores = [0.0, 0.5]
+        else:
+            assert documents == ["different", "same"]
+            scores = [0.75, 1.0]
+        result = SimpleNamespace(
+            answers={
+                key: SimpleNamespace(noul=value)
+                for key, value in reversed(list(zip(questions, scores)))
+            }
+        )
+        second_batch_scored.set()
+        return result
+
+    reranker._client.system_one.side_effect = score
+    result = reranker.rerank_vector(
+        "query", vector_results(["same", None, "same", "different", "same"])
+    )
+    assert result["_rowid"].to_pylist() == [4, 3, 2, 0, 1]
+    assert result["_distance"].to_pylist() == [4, 3, 2, 0, 1]
+    assert result["_relevance_score"].to_pylist() == [1, 0.75, 0.5, 0, 0]
+    assert reranker._client.system_one.call_count == 2
+
+
+@pytest.mark.parametrize("batch_size", [1, 40])
+@pytest.mark.parametrize("problem", ["missing", "unexpected", "replaced", "no_noul"])
+def test_invalid_answers(batch_size, problem):
+    reranker = fake_reranker(batch_size=batch_size)
+
+    def invalid(**kw):
+        result = response(kw["questions"])
+        key = next(iter(result.answers))
+        if problem in ("missing", "replaced"):
+            del result.answers[key]
+        if problem in ("unexpected", "replaced"):
+            result.answers["unknown"] = SimpleNamespace(noul=0.5)
+        if problem == "no_noul":
+            result.answers[key] = SimpleNamespace()
+        return result
+
+    reranker._client.system_one.side_effect = invalid
+    with pytest.raises(ValueError, match="invalid relevance|answer IDs"):
+        reranker.rerank_vector("query", vector_results(["one", "two"]))
+
+
+@pytest.mark.parametrize("batch_size", [1, 40])
+@pytest.mark.parametrize(
+    "probability",
+    [None, True, False, "0.5", -0.01, 1.01, float("nan"), float("inf"), -float("inf")],
+)
+def test_invalid_probabilities(batch_size, probability):
+    reranker = fake_reranker(batch_size=batch_size)
+    reranker._client.system_one.side_effect = lambda **kw: response(
+        kw["questions"], probability
+    )
+    with pytest.raises(ValueError, match="invalid relevance probability"):
+        reranker.rerank_vector("query", vector_results(["one", "two"]))
+
+
+@pytest.mark.parametrize("batch_size", [1, 40])
+def test_api_failure_propagates(batch_size):
+    reranker = fake_reranker(batch_size=batch_size)
+    error = RuntimeError("SDK exhausted retries")
+    reranker._client.system_one.side_effect = error
+    with pytest.raises(RuntimeError) as caught:
+        reranker.rerank_vector("query", vector_results(["one"]))
+    assert caught.value is error
+    assert reranker._client.system_one.call_count == 1
+
+
+@pytest.mark.parametrize("batch_size", [1, 40])
+def test_concurrency_limits_requests(batch_size):
+    reranker = fake_reranker(batch_size=batch_size, max_concurrency=2)
+    barrier, lock = Barrier(2), Lock()
+    active = peak = 0
+
+    def blocking(**kw):
+        nonlocal active, peak
+        with lock:
+            active += 1
+            peak = max(peak, active)
+        barrier.wait(timeout=10)
+        with lock:
+            active -= 1
+        return response(kw["questions"])
+
+    reranker._client.system_one.side_effect = blocking
+    reranker.rerank_vector("query", vector_results(["doc"] * (6 * batch_size)))
+    assert peak == 2
+    assert active == 0
+    assert reranker._client.system_one.call_count == 6
+
+
+@pytest.mark.parametrize("batch_size", [1, 40])
+@pytest.mark.parametrize("return_score", ["relevance", "all"])
+@pytest.mark.parametrize("method", ["vector", "fts", "hybrid"])
+@pytest.mark.parametrize("docs", [[], [None, None], [None, "", None]])
+def test_empty_and_null_inputs(batch_size, return_score, method, docs):
+    reranker = fake_reranker(batch_size=batch_size, return_score=return_score)
+    vector = vector_results(docs)
+    fts = vector.rename_columns(["text", "_rowid", "_score"])
+    args = (
+        (vector, fts)
+        if method == "hybrid"
+        else (vector if method == "vector" else fts,)
+    )
+    result = getattr(reranker, f"rerank_{method}")("query", *args)
+    assert len(result) == len(docs)
+    assert result["_relevance_score"].type == pa.float32()
+    assert sorted(result["_relevance_score"].to_pylist()) == sorted(
+        0 if doc is None else 0.5 for doc in docs
+    )
+    assert reranker._client.system_one.call_count == docs.count("")
+    if return_score == "relevance":
+        assert "_distance" not in result.column_names
+        assert "_score" not in result.column_names
+
+
+@pytest.mark.parametrize("batch_size", [1, 40])
+@pytest.mark.parametrize("return_score", ["relevance", "all"])
+def test_hybrid_deduplicates_by_row_id_and_keeps_scores(batch_size, return_score):
+    reranker = fake_reranker(batch_size=batch_size, return_score=return_score)
+    vector = vector_results(["same", "same", "third"])
+    fts = vector.slice(1).rename_columns(["text", "_rowid", "_score"])
+    result = reranker.rerank_hybrid("query", vector.slice(0, 2), fts)
+    assert result["_rowid"].to_pylist() == [0, 1, 2]
+    assert (
+        sum(
+            len(call.kwargs["questions"])
+            for call in reranker._client.system_one.call_args_list
+        )
+        == 3
+    )
+    assert reranker._client.system_one.call_count == (3 if batch_size == 1 else 1)
+    if return_score == "all":
+        assert result["_distance"].to_pylist() == [0, 1, None]
+        assert result["_score"].to_pylist() == [None, 1, 2]
+    else:
+        assert result.column_names == ["text", "_rowid", "_relevance_score"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("batch_size", [1, 40])
+async def test_async_search_integrations(tmp_path, batch_size):
+    db = await lancedb.connect_async(tmp_path)
+    table = await db.create_table(
+        "docs",
+        [
+            {"text": "cat naps", "vector": [1.0, 0.0]},
+            {"text": "cat plays", "vector": [0.0, 1.0]},
+        ],
+    )
+    await table.create_index("text", config=FTS())
+    reranker = fake_reranker(batch_size=batch_size)
+    queries = [
+        table.query().nearest_to([1.0, 0.0]).rerank(reranker, query_string="cat"),
+        table.query().nearest_to_text("cat").rerank(reranker),
+        table.query().nearest_to([1.0, 0.0]).nearest_to_text("cat").rerank(reranker),
+    ]
+    for query in queries:
+        reranker._client.system_one.reset_mock()
+        result = await query.with_row_id().to_arrow()
+        assert sorted(result["_rowid"].to_pylist()) == [0, 1]
+        assert result["_relevance_score"].to_pylist() == [0.5, 0.5]
+        assert "_distance" not in result.column_names
+        assert "_score" not in result.column_names
+        assert reranker._client.system_one.call_count == (2 if batch_size == 1 else 1)
+
+
+@pytest.mark.parametrize("batch_size", [1, 40])
+@pytest.mark.parametrize("retry_once", [False, True])
+def test_sdk_transport_and_retries(batch_size, retry_once):
+    sdk = pytest.importorskip("typesafe_sdk")
+    httpx = pytest.importorskip("httpx2")
+    payloads = []
+
+    def handle(request):
+        payload = json.loads(request.content)
+        payloads.append(payload)
+        if retry_once and len(payloads) == 1:
+            return httpx.Response(503, headers={"retry-after-ms": "0"})
+        return httpx.Response(
+            200,
+            json={
+                "model": "jev-1.13.0",
+                "usage": {},
+                "answers": {
+                    key: {"type": "noul", "noul": 0.5}
+                    for key in reversed(payload["questions"])
+                },
+            },
+        )
+
+    with sdk.TypeSafeClient(
+        api_key="test-key", transport=httpx.MockTransport(handle)
+    ) as client:
+        reranker = TypeSafeReranker(batch_size=batch_size, max_concurrency=1)
+        reranker._client = client
+        result = reranker.rerank_vector("query", vector_results(["document"] * 80))
+    assert len(payloads) == (80 if batch_size == 1 else 2) + int(retry_once)
+    assert result["_relevance_score"].to_pylist() == [0.5] * 80
+    if retry_once:
+        assert payloads[0] == payloads[1]
+    for payload in payloads:
+        assert len(payload["questions"]) == batch_size
+        if batch_size == 40:
+            assert payload["state"] == {"query": "query"}
+            assert all(
+                question["instructions"]["document"] == "document"
+                for question in payload["questions"].values()
+            )

python/python/tests/test_util.py

@@ -213,6 +213,35 @@ def test_value_to_sql_numpy_scalars():
     assert value_to_sql(np.bool_(False)) == "FALSE"
 
 
+def test_value_to_sql_nan_and_infinity(tmp_path):
+    # str(float("nan")) is "nan", which SQL reads as a column named nan, so
+    # table.update(values={"x": float("nan")}) failed with "No field named nan".
+    import math
+
+    import numpy as np
+
+    db = lancedb.connect(tmp_path)
+    table = db.create_table(
+        "test",
+        pa.table(
+            {
+                "id": [1, 2, 3, 4],
+                "f64": [1.0, 2.0, 3.0, 4.0],
+                "f32": pa.array([1.0, 2.0, 3.0, 4.0], pa.float32()),
+            }
+        ),
+    )
+    table.update(where="id = 1", values={"f64": float("nan"), "f32": np.nan})
+    table.update(where="id = 2", values={"f64": math.inf, "f32": np.float32("inf")})
+    table.update(where="id = 3", values={"f64": -math.inf, "f32": -math.inf})
+
+    rows = {row["id"]: row for row in table.to_arrow().to_pylist()}
+    assert math.isnan(rows[1]["f64"]) and math.isnan(rows[1]["f32"])
+    assert rows[2]["f64"] == math.inf and rows[2]["f32"] == math.inf
+    assert rows[3]["f64"] == -math.inf and rows[3]["f32"] == -math.inf
+    assert rows[4]["f64"] == 4.0 and rows[4]["f32"] == 4.0
+
+
 def test_append_vector_columns():
     registry = EmbeddingFunctionRegistry.get_instance()
     registry.register("test")(MockTextEmbeddingFunction)
@@ -948,6 +977,55 @@ def test_cast_to_target_schema_coerces_binary_to_metadata_blob_struct():
     ]
 
 
+def test_cast_to_target_schema_labels_json_text_with_table_extension_metadata():
+    # A table created from pa.json_() stores the empty ARROW:extension:metadata that
+    # pyarrow exports next to the name. Lance keeps that key when it turns the
+    # arrow.json label back into lance.json on write and then requires the field to
+    # match the stored one exactly, so the label has to mirror the table's metadata.
+    target = pa.schema(
+        [
+            pa.field(
+                "payload",
+                pa.large_binary(),
+                metadata={
+                    b"ARROW:extension:name": b"lance.json",
+                    b"ARROW:extension:metadata": b"",
+                },
+            ),
+            pa.field(
+                "docs",
+                pa.list_(
+                    pa.field(
+                        "item",
+                        pa.large_binary(),
+                        metadata={b"ARROW:extension:name": b"lance.json"},
+                    )
+                ),
+            ),
+        ]
+    )
+    data = pa.table(
+        {
+            "payload": pa.array(['{"a": 1}', None], type=pa.string()),
+            "docs": pa.array([['{"b": 2}'], []], type=pa.list_(pa.string())),
+        }
+    )
+
+    output = _cast_to_target_schema(data.to_reader(), target).read_all()
+
+    payload = output.schema.field("payload")
+    assert payload.type == pa.string()
+    assert payload.metadata == {
+        b"ARROW:extension:name": b"arrow.json",
+        b"ARROW:extension:metadata": b"",
+    }
+    docs_item = output.schema.field("docs").type.value_field
+    assert docs_item.type == pa.string()
+    assert docs_item.metadata == {b"ARROW:extension:name": b"arrow.json"}
+    assert output["payload"].to_pylist() == ['{"a": 1}', None]
+    assert output["docs"].to_pylist() == [['{"b": 2}'], []]
+
+
 def test_cast_to_target_schema_coerces_nested_binary_blob():
     data = pa.table(
         {

python/python/tests/test_wal_hybrid.py

@@ -0,0 +1,219 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright The LanceDB Authors
+
+"""WAL-PK-FUSION: delete this file with `lancedb/_wal_hybrid.py`."""
+
+from unittest import mock
+
+import lancedb
+import pyarrow as pa
+import pytest
+from lancedb._wal_hybrid import WAL_ROW_ID_UNSUPPORTED, with_pk_fallback
+from lancedb.query import LanceHybridQueryBuilder
+from lancedb.rerankers import RRFReranker
+
+
+def _wal_table(schema: pa.Schema) -> mock.Mock:
+    """A MemWAL-backed table a previous query has already been refused
+    `_rowid` on, so the fusion goes straight to the primary key."""
+    table = mock.Mock()
+    table._hybrid_pk_fusion_learned.return_value = True
+    table.schema = schema
+    return table
+
+
+_PK_SCHEMA = pa.schema(
+    [
+        pa.field(
+            "id",
+            pa.string(),
+            metadata={b"lance-schema:unenforced-primary-key:position": b"1"},
+        ),
+        pa.field("text", pa.string()),
+    ]
+)
+
+
+def _wal_hybrid(**select):
+    builder = (
+        LanceHybridQueryBuilder(_wal_table(_PK_SCHEMA)).vector([0.1, 0.2]).text("puppy")
+    )
+    if "columns" in select:
+        builder = builder.select(select["columns"])
+    return builder
+
+
+def test_wal_hybrid_projects_the_key_instead_of_asking_for_row_ids():
+    builder = _wal_hybrid(columns=["text"])
+    builder._create_query_builders()
+
+    assert builder._pk_fusion.pk_columns == ["id"]
+    assert builder._pk_fusion.injected == ["id"]
+    for leg in (builder._vector_query, builder._fts_query):
+        assert leg._columns == ["text", "id"]
+        # A MemWAL table rejects with_row_id before it plans.
+        assert leg._with_row_id is not True
+
+
+def test_wal_hybrid_selecting_the_key_injects_nothing():
+    builder = _wal_hybrid(columns=["text", "id"])
+    builder._create_query_builders()
+    assert builder._pk_fusion.injected == []
+
+
+def test_wal_hybrid_without_a_projection_leaves_it_alone():
+    builder = _wal_hybrid()
+    builder._create_query_builders()
+    # Select-all already returns the key.
+    assert builder._pk_fusion.injected == []
+    assert builder._vector_query._columns is None
+
+
+def test_base_table_hybrid_still_joins_on_row_ids():
+    table = mock.Mock()
+    table._hybrid_pk_fusion_learned.return_value = False
+    builder = LanceHybridQueryBuilder(table).vector([0.1, 0.2]).text("puppy")
+    builder._create_query_builders()
+
+    # `to_arrow` is what turns row ids on for this path; all that matters here
+    # is that no key was resolved and no projection was touched.
+    assert builder._pk_fusion is None
+    assert builder._vector_query._columns is None
+
+
+def test_wal_hybrid_with_use_lsm_false_joins_on_row_ids():
+    builder = _wal_hybrid().use_lsm(False)
+    builder._create_query_builders()
+    assert builder._pk_fusion is None
+
+
+def test_wal_hybrid_refuses_an_explicit_with_row_id():
+    builder = _wal_hybrid().with_row_id(True)
+    with pytest.raises(NotImplementedError, match="use_lsm"):
+        builder._create_query_builders()
+
+
+def test_wal_hybrid_return_score_all_is_not_refused():
+    """`return_score="all"` turns row ids on for the reranker, not the caller."""
+    builder = _wal_hybrid().rerank(RRFReranker(return_score="all"))
+    builder._create_query_builders()
+    assert builder._pk_fusion.pk_columns == ["id"]
+
+
+def test_wal_hybrid_fuses_on_the_key_and_drops_it_again():
+    builder = _wal_hybrid(columns=["text"])
+    builder._create_query_builders()
+
+    vector_results = pa.table(
+        {"text": ["a", "b"], "id": ["a", "b"], "_distance": [0.1, 0.2]}
+    )
+    fts_results = pa.table({"text": ["b", "c"], "id": ["b", "c"], "_score": [1.0, 2.0]})
+    vector_results, fts_results = builder._pk_fusion.stamp(vector_results, fts_results)
+
+    combined = LanceHybridQueryBuilder._combine_hybrid_results(
+        fts_results=fts_results,
+        vector_results=vector_results,
+        norm="score",
+        fts_query="puppy",
+        reranker=RRFReranker(),
+        limit=10,
+        with_row_ids=True,
+    )
+    results = builder._finish_hybrid_results(combined)
+
+    assert "_rowid" not in results.column_names, "the surrogate is internal"
+    assert "id" not in results.column_names, "the caller selected only `text`"
+    assert sorted(results.column("text").to_pylist()) == ["a", "b", "c"]
+
+
+@pytest.mark.parametrize(
+    "row_id_first", [True, False], ids=["row_id_first", "rerank_first"]
+)
+def test_wal_hybrid_refuses_row_id_even_with_return_score_all(row_id_first: bool):
+    """`return_score="all"` sets the same flag the caller's request does, so
+    asking for both must still be refused — in either call order."""
+    builder = _wal_hybrid()
+    reranker = RRFReranker(return_score="all")
+    if row_id_first:
+        builder = builder.with_row_id(True).rerank(reranker)
+    else:
+        builder = builder.rerank(reranker).with_row_id(True)
+
+    with pytest.raises(NotImplementedError, match="use_lsm"):
+        builder._create_query_builders()
+
+
+def test_wal_hybrid_refuses_a_blob_projection_with_return_score_all():
+    """A projected blob needs a real row id to fetch through, which the
+    reranker turning `_with_row_id` on does not provide."""
+    schema = pa.schema(
+        [
+            pa.field(
+                "id",
+                pa.string(),
+                metadata={b"lance-schema:unenforced-primary-key:position": b"1"},
+            ),
+            pa.field("text", pa.string()),
+            lancedb.blob("blob"),
+        ]
+    )
+    builder = (
+        LanceHybridQueryBuilder(_wal_table(schema))
+        .vector([0.1, 0.2])
+        .text("puppy")
+        .select(["blob"])
+        .rerank(RRFReranker(return_score="all"))
+    )
+
+    with pytest.raises(NotImplementedError, match="blob"):
+        builder._create_query_builders()
+
+
+_REFUSAL = RuntimeError(
+    "Bad request: the MemWAL LSM scanner does not support with_row_id"
+)
+
+
+def test_fallback_retries_on_the_key_and_remembers_it():
+    table = mock.Mock()
+    attempts = []
+
+    def run():
+        attempts.append(len(attempts))
+        if len(attempts) == 1:
+            raise _REFUSAL
+        return "fused"
+
+    result = with_pk_fallback(
+        run, table, fused=lambda: False, caller_requested_row_id=False
+    )
+
+    assert result == "fused"
+    assert len(attempts) == 2
+    table._note_hybrid_pk_fusion.assert_called_once()
+
+
+def test_fallback_explains_the_refusal_to_a_caller_who_asked_for_row_ids():
+    def run():
+        raise _REFUSAL
+
+    with pytest.raises(NotImplementedError, match="use_lsm") as info:
+        with_pk_fallback(
+            run, mock.Mock(), fused=lambda: False, caller_requested_row_id=True
+        )
+    assert str(info.value) == WAL_ROW_ID_UNSUPPORTED
+
+
+@pytest.mark.parametrize(
+    "error, fused",
+    [(ValueError("unrelated"), False), (_REFUSAL, True)],
+    ids=["other_error", "already_fused"],
+)
+def test_fallback_reraises_what_it_cannot_retry(error, fused):
+    def run():
+        raise error
+
+    table = mock.Mock()
+    with pytest.raises(type(error)):
+        with_pk_fallback(run, table, fused=lambda: fused, caller_requested_row_id=False)
+    table._note_hybrid_pk_fusion.assert_not_called()