From 4be1720c51d8f966d24172e06ae1d228bb6b1508 Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Wed, 14 Oct 2026 12:23:10 -0700 Subject: [PATCH] fix(python): support default blob conversion in pandas query results (#4308) --- python/python/lancedb/query.py | 116 +++++++++----- python/python/tests/test_blob.py | 6 +- python/python/tests/test_query.py | 209 +++++++++++++++++++++++++- python/python/tests/test_remote_db.py | 54 ++++++- 4 files changed, 341 insertions(+), 44 deletions(-) diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index 15a52c5e3..9734c9b43 100644 --- a/python/python/lancedb/query.py +++ b/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( diff --git a/python/python/tests/test_blob.py b/python/python/tests/test_blob.py index 71e78643c..8be90115c 100644 --- a/python/python/tests/test_blob.py +++ b/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") diff --git a/python/python/tests/test_query.py b/python/python/tests/test_query.py index a05247c5e..c8d0ad159 100644 --- a/python/python/tests/test_query.py +++ b/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", diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index f43972a22..717490dcd 100644 --- a/python/python/tests/test_remote_db.py +++ b/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():