fix(python): support default blob conversion in pandas query results (#4308)

python/python/lancedb/query.py

@@ -128,6 +128,11 @@ def _blob_bytes_without_snapshot_error() -> NotImplementedError:
     )
 
 
+def _is_remote_query_table(table: Any) -> bool:
+    async_table = getattr(table, "_table", table)
+    return not async_table._inner._is_native()
+
+
 def _query_needs_lance_scanner(builder: Any) -> bool:
     """Whether the query sets options that to_arrow ignores."""
     to_query_object = getattr(builder, "to_query_object", None)
@@ -347,6 +352,28 @@ def _finish_plain_scan_pandas(
     return df
 
 
+def _finish_remote_pandas(
+    tbl: pa.Table,
+    *,
+    blob_mode: BlobMode,
+    blob_paths: Iterable[str],
+    flatten: Optional[Union[int, bool]],
+    **kwargs,
+) -> pd.DataFrame:
+    # A live remote table can be replaced between query and blob fetch.
+    # Row ids and version numbers do not identify the producing table.
+    if not blob_column_paths(tbl.schema):
+        return flatten_columns(tbl, flatten).to_pandas(**kwargs)
+    if blob_mode == "descriptions":
+        tbl = strip_auto_row_ids(tbl, blob_paths)
+        return flatten_columns(tbl, flatten).to_pandas(**kwargs)
+    raise NotImplementedError(
+        f"remote to_pandas(blob_mode={blob_mode!r}) cannot safely materialize "
+        "blob columns without a stable table snapshot; "
+        "use blob_mode='descriptions'"
+    )
+
+
 async def _finish_plain_scan_pandas_async(
     scanner: _LanceScanner,
     *,
@@ -1107,34 +1134,50 @@ class LanceQueryBuilder(ABC):
             The maximum time to wait for the query to complete.
             If None, wait indefinitely.
         blob_mode: str, default "lazy"
-            Controls how blob columns are returned. ``"lazy"`` returns
-            :class:`~lancedb._blob.BlobFile` handles and needs a plain scan.
-            ``"descriptions"`` and ``"bytes"`` work for every query shape. A
-            plain scan without ``timeout`` reads the local dataset through
-            Lance's scanner, which also returns legacy blob columns as bytes.
-            Other shapes use :meth:`to_arrow` with the same ``blob_mode``.
+            Controls how blob columns are returned. Eligible local plain scans
+            use Lance's scanner, which returns lazy blob handles in ``"lazy"``
+            mode and also supports legacy blob columns. Other local blob v2
+            queries return descriptors in ``"lazy"`` and ``"descriptions"``
+            modes, or payloads in ``"bytes"`` mode, through :meth:`to_arrow`.
+            Remote blob queries support only ``"descriptions"`` until a stable
+            table snapshot can bind the query and blob fetches.
         **kwargs
             Forwarded to pyarrow.Table.to_pandas after query execution and
             optional flattening.
         """
         validate_blob_mode(blob_mode)
+        if _is_remote_query_table(self._table):
+            blob_sources = self._validated_blob_sources(blob_mode)
+            return _finish_remote_pandas(
+                self.to_arrow(timeout=timeout),
+                blob_mode=blob_mode,
+                blob_paths=blob_sources,
+                flatten=flatten,
+                **kwargs,
+            )
+
         output_schema = getattr(self, "output_schema", None)
         if output_schema is not None and schema_has_blob_field(output_schema()):
             if (flatten is None or blob_mode != "lazy") and timeout is None:
                 df = self._plain_scan_to_pandas(blob_mode, flatten=flatten, **kwargs)
                 if df is not None:
                     return df
-            if blob_mode == "lazy" or _query_needs_lance_scanner(self):
+            if _query_needs_lance_scanner(self):
                 raise _unsupported_blob_pandas_error(
                     "this query shape cannot use Lance native pandas conversion"
                 )
 
-        tbl = flatten_columns(self._to_arrow_in_mode(timeout, blob_mode), flatten)
-        if blob_mode == "lazy" and schema_has_blob_field(tbl.schema):
+        arrow_mode = "descriptions" if blob_mode == "lazy" else blob_mode
+        tbl = self._to_arrow_in_mode(timeout, arrow_mode)
+        if (
+            blob_mode == "lazy"
+            and schema_has_blob_field(tbl.schema)
+            and not blob_v2_projection_sources(self._table.schema, self._columns)
+        ):
             raise _unsupported_blob_pandas_error(
-                "this query shape cannot use Lance native pandas conversion"
+                "the query has no fetchable blob v2 source"
             )
-        return tbl.to_pandas(**kwargs)
+        return flatten_columns(tbl, flatten).to_pandas(**kwargs)
 
     @abstractmethod
     def to_arrow(
@@ -3364,13 +3407,13 @@ class AsyncQueryBase(object):
             If not specified, no timeout is applied. If the query does not
             complete within the specified time, an error will be raised.
         blob_mode: str, default "lazy"
-            Controls how blob columns are returned. ``"lazy"`` returns
-            :class:`~lancedb._blob.BlobFile` handles and needs a plain scan.
-            ``"descriptions"`` and ``"bytes"`` work for every query shape. A
-            plain scan without ``timeout`` reads the local dataset through
-            Lance's scanner, which also returns legacy blob columns as bytes.
-            Other shapes use :meth:`to_arrow` with the same ``blob_mode``.
-            Remote queries support ``"descriptions"`` only.
+            Controls how blob columns are returned. Eligible local plain scans
+            use Lance's scanner, which returns lazy blob handles in ``"lazy"``
+            mode and also supports legacy blob columns. Other local blob v2
+            queries return descriptors in ``"lazy"`` and ``"descriptions"``
+            modes, or payloads in ``"bytes"`` mode, through :meth:`to_arrow`.
+            Remote blob queries support only ``"descriptions"`` until a stable
+            table snapshot can bind the query and blob fetches.
         **kwargs
             Forwarded to pyarrow.Table.to_pandas after query execution and
             optional flattening.
@@ -3391,19 +3434,22 @@ class AsyncQueryBase(object):
                 )
                 if df is not None:
                     return df
-            if blob_mode == "lazy" or _query_needs_lance_scanner(self):
+            if _query_needs_lance_scanner(self):
                 raise _unsupported_blob_pandas_error(
                     "this query shape cannot use Lance native pandas conversion"
                 )
 
-        tbl = flatten_columns(
-            await self.to_arrow(timeout=timeout, blob_mode=blob_mode), flatten
-        )
-        if blob_mode == "lazy" and schema_has_blob_field(tbl.schema):
+        arrow_mode = "descriptions" if blob_mode == "lazy" else blob_mode
+        tbl = await self.to_arrow(timeout=timeout, blob_mode=arrow_mode)
+        if (
+            blob_mode == "lazy"
+            and schema_has_blob_field(tbl.schema)
+            and not await self._validated_blob_sources("descriptions")
+        ):
             raise _unsupported_blob_pandas_error(
-                "this query shape cannot use Lance native pandas conversion"
+                "the query has no fetchable blob v2 source"
             )
-        return tbl.to_pandas(**kwargs)
+        return flatten_columns(tbl, flatten).to_pandas(**kwargs)
 
     async def _remote_to_pandas(
         self,
@@ -3413,21 +3459,13 @@ class AsyncQueryBase(object):
         blob_mode: BlobMode,
         **kwargs,
     ) -> "pd.DataFrame":
-        # A live remote table can be replaced between query and blob fetch.
-        # Row ids and version numbers do not identify the producing table.
         tbl = await self.to_arrow(timeout=timeout)
-        projected_blob_paths = set(blob_column_paths(tbl.schema))
-        if not projected_blob_paths:
-            return flatten_columns(tbl, flatten).to_pandas(**kwargs)
-
-        if blob_mode == "descriptions":
-            tbl = strip_auto_row_ids(tbl, self._blob_paths)
-            return flatten_columns(tbl, flatten).to_pandas(**kwargs)
-
-        raise NotImplementedError(
-            f"remote to_pandas(blob_mode={blob_mode!r}) cannot safely materialize "
-            "blob columns without a stable table snapshot; "
-            "use blob_mode='descriptions'"
+        return _finish_remote_pandas(
+            tbl,
+            blob_mode=blob_mode,
+            blob_paths=self._blob_paths,
+            flatten=flatten,
+            **kwargs,
         )
 
     async def _plain_scan_to_pandas(

python/python/tests/test_blob.py

@@ -1632,7 +1632,7 @@ def test_fetch_blob_files_lazy_read():
     assert handles[0].read() == payload
     handles[0].close()
     assert handles[0].closed
-    with pytest.raises(RuntimeError, match="already closed"):
+    with pytest.raises(ValueError, match="read of closed file"):
         handles[0].read()
 
 
@@ -1681,8 +1681,8 @@ def test_blob_file_finalizer_and_close_on_runtime_worker():
         explicit_handle.close()  # Close remains idempotent after the worker call.
         try:
             explicit_handle.read()
-        except RuntimeError as error:
-            assert "already closed" in str(error)
+        except ValueError as error:
+            assert "read of closed file" in str(error)
         else:
             raise AssertionError("a closed blob file was readable")
         print("worker close completed")

python/python/tests/test_query.py

@@ -111,6 +111,26 @@ async def _create_blob_v2_query_table_async(db, name):
     return table
 
 
+def _nested_blob_v2_query_schema():
+    return pa.schema(
+        [
+            pa.field("id", pa.int64()),
+            pa.field("vector", pa.list_(pa.float32(), list_size=2)),
+            pa.field(
+                "info",
+                pa.struct([pa.field("label", pa.string()), lancedb.blob("blob")]),
+            ),
+        ]
+    )
+
+
+def _nested_blob_v2_query_rows():
+    return [
+        {"id": 1, "vector": [1.0, 0.0], "info": {"label": "one", "blob": b"one"}},
+        {"id": 2, "vector": [2.0, 0.0], "info": {"label": "two", "blob": b"two"}},
+    ]
+
+
 def _assert_lazy_blob(value, expected: bytes):
     assert hasattr(value, "readall")
     assert value.readall() == expected
@@ -503,7 +523,7 @@ def test_vector_query_to_pandas_blob_mode_requires_native_path(tmp_db):
         storage_options=LEGACY_BLOB_STORAGE_OPTIONS,
     )
 
-    with pytest.raises(RuntimeError, match="Lance native pandas conversion"):
+    with pytest.raises(RuntimeError, match="no fetchable blob v2 source"):
         table.search([1.0, 0.0]).select(["blob", "vector"]).limit(1).to_pandas(
             blob_mode="lazy"
         )
@@ -535,6 +555,193 @@ def test_vector_query_to_pandas_legacy_blob_descriptions_and_bytes(tmp_db):
     ]
 
 
+@pytest.mark.parametrize("blob_mode", ["default", "lazy", "bytes", "descriptions"])
+@pytest.mark.parametrize("with_row_id", [None, True, False])
+def test_vector_query_to_pandas_blob_v2_modes(tmp_db, blob_mode, with_row_id):
+    pytest.importorskip("lance")
+    table = _create_blob_v2_query_table(tmp_db, f"test_vector_blob_v2_{blob_mode}")
+
+    query = table.search([1.0, 0.0]).select(["id", "blob", "vector"]).limit(1)
+    if with_row_id is not None:
+        query = query.with_row_id(with_row_id)
+    df = (
+        query.to_pandas()
+        if blob_mode == "default"
+        else query.to_pandas(blob_mode=blob_mode)
+    )
+
+    assert df["id"].tolist() == [1]
+    assert ("_rowid" in df.columns) == (with_row_id is True)
+    if blob_mode == "bytes":
+        assert df["blob"].tolist() == [b"one"]
+    else:
+        descriptor = df["blob"].iloc[0]
+        assert descriptor != b"one"
+        assert not hasattr(descriptor, "readall")
+        assert "_lance_row_id" not in descriptor
+
+
+@pytest.mark.parametrize("blob_mode", ["lazy", "bytes", "descriptions"])
+def test_blob_v2_non_scan_query_shapes_to_pandas(tmp_db, blob_mode):
+    pytest.importorskip("lance")
+    table = _create_blob_v2_query_table(tmp_db, "test_blob_v2_query_shapes")
+    table.create_index("tag", config=FTS())
+
+    queries = [
+        table.search("keep", query_type="fts").limit(2),
+        table.search(query_type="hybrid").vector([1.0, 0.0]).text("keep").limit(2),
+        table.take_offsets([1, 2]).select(["id", "blob"]),
+    ]
+    for query in queries:
+        df = query.to_pandas(blob_mode=blob_mode)
+        assert len(df) == 2
+        if blob_mode == "bytes":
+            assert df["blob"].tolist() == [
+                {1: b"one", 2: b"two", 3: b"three", 4: b"four"}[row_id]
+                for row_id in df["id"]
+            ]
+        else:
+            assert all("_lance_row_id" not in blob for blob in df["blob"])
+        assert "_rowid" not in df.columns
+
+
+@pytest.mark.parametrize("rename_parent", [False, True])
+def test_vector_query_to_pandas_nested_blob_bytes(tmp_db, rename_parent):
+    pytest.importorskip("lance")
+    table = tmp_db.create_table(
+        f"test_vector_nested_blob_{rename_parent}",
+        schema=_nested_blob_v2_query_schema(),
+    )
+    table.add(_nested_blob_v2_query_rows())
+
+    projection = {"id": "id", "renamed": "info"} if rename_parent else ["id", "info"]
+    df = (
+        table.search([1.0, 0.0])
+        .select(projection)
+        .limit(2)
+        .to_pandas(blob_mode="bytes")
+    )
+
+    assert df["id"].tolist() == [1, 2]
+    parent = "renamed" if rename_parent else "info"
+    assert df[parent].tolist() == [
+        {"label": "one", "blob": b"one"},
+        {"label": "two", "blob": b"two"},
+    ]
+    for blob_mode in ("lazy", "descriptions"):
+        descriptors = (
+            table.search([1.0, 0.0])
+            .select(projection)
+            .limit(1)
+            .to_pandas(blob_mode=blob_mode)
+        )
+        assert "_lance_row_id" not in descriptors[parent].iloc[0]["blob"]
+
+
+@pytest.mark.asyncio
+async def test_async_vector_query_to_pandas_nested_blob_bytes(tmp_db_async):
+    pytest.importorskip("lance")
+    table = await tmp_db_async.create_table(
+        "test_async_vector_nested_blob",
+        schema=_nested_blob_v2_query_schema(),
+    )
+    await table.add(_nested_blob_v2_query_rows())
+
+    df = await (
+        table.vector_search([1.0, 0.0])
+        .select(["id", "info"])
+        .limit(2)
+        .to_pandas(blob_mode="bytes")
+    )
+
+    assert df["id"].tolist() == [1, 2]
+    assert df["info"].tolist() == [
+        {"label": "one", "blob": b"one"},
+        {"label": "two", "blob": b"two"},
+    ]
+
+
+def test_blob_v2_vector_projection_and_plain_scan_options(tmp_db, monkeypatch):
+    pytest.importorskip("lance")
+    table = _create_blob_v2_query_table(tmp_db, "test_blob_v2_pandas_options")
+
+    df = (
+        table.search([1.0, 0.0])
+        .select({"payload": "blob", "id": "id"})
+        .limit(1)
+        .to_pandas(blob_mode="bytes")
+    )
+    assert df["payload"].tolist() == [b"one"]
+
+    query = table.search().where("id = 1")
+
+    def fail_to_arrow(*args, **kwargs):
+        raise AssertionError("the native scanner should handle bytes with flatten")
+
+    monkeypatch.setattr(query, "to_arrow", fail_to_arrow)
+    assert query.to_pandas(blob_mode="bytes", flatten=True)["blob"].tolist() == [b"one"]
+
+    flattened_lazy = table.search().where("id = 1").to_pandas(flatten=True)
+    assert flattened_lazy["id"].tolist() == [1]
+    assert not any("_lance_row_id" in col for col in flattened_lazy.columns)
+
+    timed = (
+        table.search()
+        .where("id = 1")
+        .to_pandas(blob_mode="bytes", timeout=timedelta(seconds=10))
+    )
+    assert timed["blob"].tolist() == [b"one"]
+
+    timed_lazy = table.search().where("id = 1").to_pandas(timeout=timedelta(seconds=10))
+    assert timed_lazy["id"].tolist() == [1]
+
+
+def test_blob_v2_pandas_preserves_fetch_error(tmp_db, monkeypatch):
+    pytest.importorskip("lance")
+    table = _create_blob_v2_query_table(tmp_db, "test_blob_v2_fetch_error")
+
+    def missing_blob(*args, **kwargs):
+        raise FileNotFoundError("missing obj.bin")
+
+    monkeypatch.setattr(table, "fetch_blobs", missing_blob)
+    with pytest.raises(FileNotFoundError, match="missing obj.bin"):
+        table.search().to_pandas(blob_mode="bytes")
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("blob_mode", ["lazy", "bytes", "descriptions"])
+@pytest.mark.parametrize("with_row_id", [False, True])
+async def test_async_vector_query_to_pandas_blob_v2_modes(
+    tmp_db_async, blob_mode, with_row_id
+):
+    pytest.importorskip("lance")
+    table = await _create_blob_v2_query_table_async(
+        tmp_db_async, f"test_async_vector_blob_v2_{blob_mode}"
+    )
+
+    query = table.vector_search([1.0, 0.0]).limit(1)
+    if with_row_id:
+        query = query.with_row_id()
+    df = await query.to_pandas(blob_mode=blob_mode)
+    assert df["id"].tolist() == [1]
+    assert ("_rowid" in df.columns) == with_row_id
+    if blob_mode == "bytes":
+        assert df["blob"].tolist() == [b"one"]
+    else:
+        assert "_lance_row_id" not in df["blob"].iloc[0]
+
+
+@pytest.mark.asyncio
+async def test_async_take_query_to_pandas_blob_v2_bytes(tmp_db_async):
+    pytest.importorskip("lance")
+    table = await _create_blob_v2_query_table_async(
+        tmp_db_async, "test_async_take_blob_v2"
+    )
+
+    df = await table.take_offsets([1, 2]).to_pandas(blob_mode="bytes")
+    assert set(df["blob"]) == {b"two", b"three"}
+
+
 def test_order_by_plain_query(mem_db):
     table = mem_db.create_table(
         "test_order_by",

python/python/tests/test_remote_db.py

@@ -3539,7 +3539,7 @@ def test_remote_blob_files_are_lazy_seekable_handles():
         assert empty.read_range(0, 0) == b""
         alpha.close()
         assert alpha.closed
-        with pytest.raises(RuntimeError, match="already closed"):
+        with pytest.raises(ValueError, match="read of closed file"):
             alpha.read_range(0, 1)
 
 
@@ -3584,6 +3584,56 @@ def test_remote_blob_query_stashes_row_ids_for_fetch():
     assert blobs.to_pylist() == [b"alpha", None, b"gamma"]
 
 
+def test_remote_blob_query_to_pandas():
+    with blob_remote_table() as table:
+        query = table.search().select(["id", "image"]).limit(3)
+        df = query.to_pandas(blob_mode="descriptions", split_blocks=True)
+
+    assert df["id"].tolist() == [1, 2, 3]
+    assert "_rowid" not in df.columns
+    assert df["image"].iloc[1] is None
+    assert "_lance_row_id" not in df["image"].iloc[0]
+
+
+@pytest.mark.parametrize("blob_mode", ["default", "lazy", "bytes"])
+def test_remote_blob_query_to_pandas_rejects_unpinned_fetch(monkeypatch, blob_mode):
+    with blob_remote_table() as table:
+
+        def fail_fetch(*args, **kwargs):
+            raise AssertionError("blob fetch must not use an unpinned remote table")
+
+        monkeypatch.setattr(table, "fetch_blobs", fail_fetch)
+        query = table.search().select(["id", "image"]).limit(3)
+        pandas_kwargs = {} if blob_mode == "default" else {"blob_mode": blob_mode}
+        with pytest.raises(NotImplementedError, match="stable table snapshot"):
+            query.to_pandas(**pandas_kwargs)
+
+
+@pytest.mark.parametrize("blob_mode", ["lazy", "bytes"])
+def test_remote_nonblob_projection_to_pandas(blob_mode):
+    with blob_remote_table() as table:
+        df = table.search().select(["id"]).to_pandas(blob_mode=blob_mode)
+
+    assert df["id"].tolist() == [1, 2, 3]
+    assert list(df.columns) == ["id"]
+
+
+@pytest.mark.asyncio
+async def test_async_remote_blob_query_to_pandas_flatten():
+    with blob_remote_table() as table:
+        df = (
+            await table._table.query()
+            .select(["id", "image"])
+            .limit(3)
+            .to_pandas(blob_mode="descriptions", flatten=True, split_blocks=True)
+        )
+
+    assert df["id"].tolist() == [1, 2, 3]
+    assert df["image.size"].iloc[0] == 5
+    assert "image._lance_row_id" not in df.columns
+    assert "_rowid" not in df.columns
+
+
 def test_remote_blob_query_survives_a_server_that_ignores_the_row_id_request():
     def handler(request):
         if request.path == "/v1/table/test/describe/":
@@ -3613,6 +3663,8 @@ def test_remote_blob_query_survives_a_server_that_ignores_the_row_id_request():
         assert "_lance_row_id" not in hits.schema.field("image").type.names
         with pytest.raises(ValueError, match="pass a list of row ids"):
             table.fetch_blobs("image", hits)
+        with pytest.raises(NotImplementedError, match="stable table snapshot"):
+            table.search().select(["id", "image"]).to_pandas(blob_mode="bytes")
 
 
 def test_remote_blob_byte_apis_not_supported_on_old_server():