From 651379742e5337d723ebffd4f4c7fde521f65717 Mon Sep 17 00:00:00 2001 From: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 19:09:52 +0000 Subject: [PATCH] fix(python): support blob v2 pandas conversion for query results --- python/python/lancedb/query.py | 188 ++++++++++++++++++-------- python/python/tests/test_query.py | 121 +++++++++++++++-- python/python/tests/test_remote_db.py | 73 ++++++++++ 3 files changed, 313 insertions(+), 69 deletions(-) diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index cef89e63c..05f310a22 100644 --- a/python/python/lancedb/query.py +++ b/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. diff --git a/python/python/tests/test_query.py b/python/python/tests/test_query.py index e3d55917d..0f6fff18b 100644 --- a/python/python/tests/test_query.py +++ b/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): diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index c03862a49..3d53a9abc 100644 --- a/python/python/tests/test_remote_db.py +++ b/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():