fix(python): support blob v2 pandas conversion for query results

python/python/lancedb/query.py

@@ -40,7 +40,7 @@ from .rerankers.base import Reranker
 from .rerankers.rrf import RRFReranker
 from .rerankers.util import check_reranker_result
 from .schema import is_blob_like_field, schema_has_blob_field
-from .util import flatten_columns
+from .util import flatten_columns, get_uri_scheme
 from ._blob import (
     BLOB_MODE_TO_HANDLING,
     FetchBlobsAsync,
@@ -50,6 +50,7 @@ from ._blob import (
     finalize_blob_query_table,
     replace_v2_blob_columns_with_bytes,
     replace_v2_blob_columns_with_bytes_sync,
+    strip_auto_row_ids,
     validate_blob_mode,
 )
 from .types import BlobMode, QueryProjection
@@ -97,10 +98,6 @@ class _LanceScanner(Protocol):
     def to_reader(self): ...
 
 
-def _blob_mode_requires_native_pandas(blob_mode: BlobMode, schema: pa.Schema) -> bool:
-    return blob_mode in BLOB_MODE_TO_HANDLING and schema_has_blob_field(schema)
-
-
 def _unsupported_blob_pandas_error(reason: str) -> RuntimeError:
     return RuntimeError(
         "blob columns require Lance native scanner conversion for query "
@@ -109,6 +106,11 @@ def _unsupported_blob_pandas_error(reason: str) -> RuntimeError:
     )
 
 
+def _is_remote_query_table(table: Any) -> bool:
+    async_table = getattr(table, "_table", table)
+    return get_uri_scheme(async_table._inner.database().uri) == "db"
+
+
 def _query_is_plain_scan(query: Query) -> bool:
     return (
         query.vector is None
@@ -310,6 +312,33 @@ def _finish_plain_scan_pandas(
     return df
 
 
+def _finish_arrow_pandas(
+    tbl: pa.Table,
+    *,
+    blob_mode: BlobMode,
+    blob_sources: dict[str, str],
+    fetch_blobs: FetchBlobsSync,
+    flatten: Optional[Union[int, bool]],
+    **kwargs,
+) -> pd.DataFrame:
+    if (
+        blob_mode != "descriptions"
+        and not blob_sources
+        and schema_has_blob_field(tbl.schema)
+    ):
+        raise _unsupported_blob_pandas_error(
+            "this query shape cannot use Lance native pandas conversion "
+            "for blob columns without a fetchable source"
+        )
+    if blob_mode == "bytes":
+        tbl = replace_v2_blob_columns_with_bytes_sync(tbl, blob_sources, fetch_blobs)
+    elif blob_sources:
+        # Non-scanner queries return descriptors in both descriptions and lazy
+        # modes. The row id is internal to fetching and not part of the view.
+        tbl = strip_auto_row_ids(tbl, blob_sources)
+    return flatten_columns(tbl, flatten).to_pandas(**kwargs)
+
+
 async def _finish_plain_scan_pandas_async(
     scanner: _LanceScanner,
     *,
@@ -339,6 +368,31 @@ async def _finish_plain_scan_pandas_async(
     return df
 
 
+async def _finish_arrow_pandas_async(
+    tbl: pa.Table,
+    *,
+    blob_mode: BlobMode,
+    blob_sources: dict[str, str],
+    fetch_blobs: FetchBlobsAsync,
+    flatten: Optional[Union[int, bool]],
+    **kwargs,
+) -> pd.DataFrame:
+    if (
+        blob_mode != "descriptions"
+        and not blob_sources
+        and schema_has_blob_field(tbl.schema)
+    ):
+        raise _unsupported_blob_pandas_error(
+            "this query shape cannot use Lance native pandas conversion "
+            "for blob columns without a fetchable source"
+        )
+    if blob_mode == "bytes":
+        tbl = await replace_v2_blob_columns_with_bytes(tbl, blob_sources, fetch_blobs)
+    elif blob_sources:
+        tbl = strip_auto_row_ids(tbl, blob_sources)
+    return flatten_columns(tbl, flatten).to_pandas(**kwargs)
+
+
 # Pydantic validation function for vector queries
 def ensure_vector_query(
     val: Any,
@@ -1058,41 +1112,37 @@ 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 for plain scan queries.
-            Vector, FTS, hybrid, and other non-native query shapes keep the
-            existing Arrow conversion path and only support blob descriptions.
+            Controls how blob columns are returned. For blob v2 columns, queries
+            that cannot use a native Lance scanner return descriptors in both
+            ``"lazy"`` and ``"descriptions"`` modes. Use ``"bytes"`` to fetch
+            full payloads for those queries.
         **kwargs
             Forwarded to pyarrow.Table.to_pandas after query execution and
             optional flattening.
         """
         validate_blob_mode(blob_mode)
         output_schema = getattr(self, "output_schema", None)
-        if output_schema is not None:
+        if output_schema is not None and not _is_remote_query_table(self._table):
             schema = output_schema()
-            if _blob_mode_requires_native_pandas(blob_mode, schema):
-                native_error = None
-                if (flatten is None or blob_mode == "descriptions") and timeout is None:
-                    try:
-                        df = self._plain_scan_to_pandas(
-                            blob_mode, flatten=flatten, **kwargs
-                        )
-                        if df is not None:
-                            return df
-                    except Exception as err:
-                        native_error = err
-                reason = (
-                    "this query shape cannot use Lance native pandas conversion"
-                    if native_error is None
-                    else str(native_error)
-                )
-                raise _unsupported_blob_pandas_error(reason) from native_error
+            if (
+                schema_has_blob_field(schema)
+                and timeout is None
+                and (flatten is None or blob_mode != "lazy")
+            ):
+                df = self._plain_scan_to_pandas(blob_mode, flatten=flatten, **kwargs)
+                if df is not None:
+                    return df
 
-        tbl = flatten_columns(self.to_arrow(timeout=timeout), flatten)
-        if _blob_mode_requires_native_pandas(blob_mode, tbl.schema):
-            raise _unsupported_blob_pandas_error(
-                "this query shape cannot use Lance native pandas conversion"
-            )
-        return tbl.to_pandas(**kwargs)
+        tbl = self.to_arrow(timeout=timeout)
+        blob_sources = blob_v2_projection_sources(self._table.schema, self._columns)
+        return _finish_arrow_pandas(
+            tbl,
+            blob_mode=blob_mode,
+            blob_sources=blob_sources,
+            fetch_blobs=self._table.fetch_blobs,
+            flatten=flatten,
+            **kwargs,
+        )
 
     @abstractmethod
     def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table:
@@ -1551,7 +1601,13 @@ class LanceQueryBuilder(ABC):
         if not _query_is_plain_scan(query):
             return None
 
-        dataset = self._table.to_lance()
+        if _is_remote_query_table(self._table):
+            return None
+
+        try:
+            dataset = self._table.to_lance()
+        except NotImplementedError:
+            return None
         blob_auto_row_id = self._blob_auto_row_id_enabled()
         blob_sources = (
             blob_v2_projection_sources(self._table.schema, query.columns)
@@ -3018,40 +3074,48 @@ 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 for plain scan queries.
-            Vector, FTS, hybrid, and other non-native query shapes keep the
-            existing Arrow conversion path and only support blob descriptions.
+            Controls how blob columns are returned. For blob v2 columns, queries
+            that cannot use a native Lance scanner return descriptors in both
+            ``"lazy"`` and ``"descriptions"`` modes. Use ``"bytes"`` to fetch
+            full payloads for those queries.
         **kwargs
             Forwarded to pyarrow.Table.to_pandas after query execution and
             optional flattening.
         """
         validate_blob_mode(blob_mode)
-        if hasattr(self._inner, "output_schema"):
+        if hasattr(self._inner, "output_schema") and (
+            self._table is None or not _is_remote_query_table(self._table)
+        ):
             schema = await self.output_schema()
-            if _blob_mode_requires_native_pandas(blob_mode, schema):
-                native_error = None
-                if (flatten is None or blob_mode == "descriptions") and timeout is None:
-                    try:
-                        df = await self._plain_scan_to_pandas(
-                            blob_mode, flatten=flatten, **kwargs
-                        )
-                        if df is not None:
-                            return df
-                    except Exception as err:
-                        native_error = err
-                reason = (
-                    "this query shape cannot use Lance native pandas conversion"
-                    if native_error is None
-                    else str(native_error)
+            if (
+                schema_has_blob_field(schema)
+                and timeout is None
+                and (flatten is None or blob_mode != "lazy")
+            ):
+                df = await self._plain_scan_to_pandas(
+                    blob_mode, flatten=flatten, **kwargs
                 )
-                raise _unsupported_blob_pandas_error(reason) from native_error
+                if df is not None:
+                    return df
 
-        tbl = flatten_columns(await self.to_arrow(timeout=timeout), flatten)
-        if _blob_mode_requires_native_pandas(blob_mode, tbl.schema):
-            raise _unsupported_blob_pandas_error(
-                "this query shape cannot use Lance native pandas conversion"
-            )
-        return tbl.to_pandas(**kwargs)
+        tbl = await self.to_arrow(timeout=timeout)
+        if self._table is None:
+            if blob_mode != "descriptions" and schema_has_blob_field(tbl.schema):
+                raise _unsupported_blob_pandas_error(
+                    "the query has no table from which to fetch blobs"
+                )
+            return flatten_columns(tbl, flatten).to_pandas(**kwargs)
+        schema = await self._table.schema()
+        projection = _query_request_projection(self._inner.to_query_request())
+        blob_sources = blob_v2_projection_sources(schema, projection)
+        return await _finish_arrow_pandas_async(
+            tbl,
+            blob_mode=blob_mode,
+            blob_sources=blob_sources,
+            fetch_blobs=self._table.fetch_blobs,
+            flatten=flatten,
+            **kwargs,
+        )
 
     async def _plain_scan_to_pandas(
         self,
@@ -3066,6 +3130,9 @@ class AsyncQueryBase(object):
         if not _query_is_plain_scan(query):
             return None
 
+        if _is_remote_query_table(self._table):
+            return None
+
         schema = await self._table.schema()
         blob_auto_row_id = blob_auto_row_id_for_scan(
             schema,
@@ -4290,7 +4357,8 @@ class BaseQueryBuilder(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 for plain scan queries.
+            Controls how blob columns are returned. For blob v2 take queries,
+            ``"lazy"`` returns descriptors and ``"bytes"`` fetches payloads.
         **kwargs
             Forwarded to pyarrow.Table.to_pandas after query execution and
             optional flattening.

python/python/tests/test_query.py

@@ -508,18 +508,121 @@ def test_vector_query_to_pandas_blob_mode_requires_native_path(tmp_db):
         )
 
 
-def test_vector_query_to_pandas_blob_descriptions_requires_plain_scan(tmp_db):
+@pytest.mark.parametrize("blob_mode", ["default", "lazy", "bytes", "descriptions"])
+def test_vector_query_to_pandas_blob_v2_modes(tmp_db, blob_mode):
     pytest.importorskip("lance")
-    table = tmp_db.create_table(
-        "test_vector_query_blob_descriptions",
-        _blob_query_data(),
-        storage_options=LEGACY_BLOB_STORAGE_OPTIONS,
+    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)
+    df = (
+        query.to_pandas()
+        if blob_mode == "default"
+        else query.to_pandas(blob_mode=blob_mode)
     )
 
-    with pytest.raises(RuntimeError, match="plain scan query"):
-        table.search([1.0, 0.0]).select(["blob", "vector"]).limit(1).to_pandas(
-            blob_mode="descriptions"
-        )
+    assert df["id"].tolist() == [1]
+    assert "_rowid" not in df.columns
+    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
+
+
+def test_blob_v2_non_scan_query_shapes_to_pandas(tmp_db):
+    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="bytes")
+        assert len(df) == 2
+        assert df["blob"].tolist() == [
+            {1: b"one", 2: b"two", 3: b"three", 4: b"four"}[row_id]
+            for row_id in df["id"]
+        ]
+        assert "_rowid" not in df.columns
+
+
+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"])
+async def test_async_vector_query_to_pandas_blob_v2_modes(tmp_db_async, blob_mode):
+    pytest.importorskip("lance")
+    table = await _create_blob_v2_query_table_async(
+        tmp_db_async, f"test_async_vector_blob_v2_{blob_mode}"
+    )
+
+    df = await table.vector_search([1.0, 0.0]).limit(1).to_pandas(blob_mode=blob_mode)
+    assert df["id"].tolist() == [1]
+    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):

python/python/tests/test_remote_db.py

@@ -2491,6 +2491,77 @@ def test_remote_blob_query_stashes_row_ids_for_fetch():
     assert blobs.to_pylist() == [b"alpha", None, b"gamma"]
 
 
+@pytest.mark.parametrize("blob_mode", ["default", "lazy", "bytes", "descriptions"])
+def test_remote_blob_query_to_pandas(blob_mode):
+    with blob_remote_table() as table:
+        query = table.search().select(["id", "image"]).limit(3)
+        df = (
+            query.to_pandas()
+            if blob_mode == "default"
+            else query.to_pandas(blob_mode=blob_mode)
+        )
+
+    assert df["id"].tolist() == [1, 2, 3]
+    assert "_rowid" not in df.columns
+    if blob_mode == "bytes":
+        assert df["image"].tolist() == [b"alpha", None, b"gamma"]
+    else:
+        assert df["image"].iloc[1] is None
+        assert "_lance_row_id" not in df["image"].iloc[0]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("blob_mode", ["lazy", "bytes", "descriptions"])
+async def test_async_remote_blob_query_to_pandas(blob_mode):
+    def handler(request):
+        if request.path == "/v1/table/test/describe/":
+            request.send_response(200)
+            request.send_header("Content-Type", "application/json")
+            request.send_header("phalanx-version", "0.5.0")
+            request.end_headers()
+            request.wfile.write(json.dumps(BLOB_DESCRIBE_RESPONSE).encode())
+        elif request.path == "/v1/table/test/query/":
+            content_len = int(request.headers.get("Content-Length", 0))
+            body = json.loads(request.rfile.read(content_len))
+            assert body["with_row_id"] is True
+            response_table = blob_query_response_table()
+            request.send_response(200)
+            request.send_header("Content-Type", "application/vnd.apache.arrow.file")
+            request.end_headers()
+            with pa.ipc.new_file(request.wfile, response_table.schema) as writer:
+                writer.write_table(response_table)
+        elif request.path == "/v1/table/test/fetch_blobs/":
+            content_len = int(request.headers.get("Content-Length", 0))
+            body = json.loads(request.rfile.read(content_len))
+            assert body["row_ids"] == [10, 20, 30]
+            response_table = pa.table(
+                {"image": pa.array([b"alpha", None, b"gamma"], pa.large_binary())}
+            )
+            request.send_response(200)
+            request.send_header("Content-Type", "application/vnd.apache.arrow.stream")
+            request.end_headers()
+            with pa.ipc.new_stream(request.wfile, response_table.schema) as writer:
+                writer.write_table(response_table)
+        else:
+            request.send_response(404)
+            request.end_headers()
+
+    async with mock_lancedb_connection_async(handler) as db:
+        table = await db.open_table("test")
+        df = (
+            await table.query()
+            .select(["id", "image"])
+            .limit(3)
+            .to_pandas(blob_mode=blob_mode)
+        )
+
+    assert df["id"].tolist() == [1, 2, 3]
+    if blob_mode == "bytes":
+        assert df["image"].tolist() == [b"alpha", None, b"gamma"]
+    else:
+        assert "_lance_row_id" not in df["image"].iloc[0]
+
+
 def test_remote_blob_query_survives_a_server_that_ignores_the_row_id_request():
     def handler(request):
         if request.path == "/v1/table/test/describe/":
@@ -2520,6 +2591,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(ValueError, match="pass a list of row ids"):
+            table.search().select(["id", "image"]).to_pandas(blob_mode="bytes")
 
 
 def test_remote_blob_byte_apis_not_supported_on_old_server():