fix(python): route async remote blob reads through server

This commit is contained in:
Gatefixer
2026-09-23 19:08:26 +00:00
parent df7bfcdb65
commit b1d9c4cb84
5 changed files with 148 additions and 7 deletions
+1
View File
@@ -393,6 +393,7 @@ class Table:
def name(self) -> str: ...
def __repr__(self) -> str: ...
def is_open(self) -> bool: ...
def _is_native(self) -> bool: ...
def close(self) -> None: ...
async def schema(self) -> pa.Schema: ...
async def add(
+76 -4
View File
@@ -39,7 +39,7 @@ from .expr import Expr
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 .schema import blob_column_paths, is_blob_like_field, schema_has_blob_field
from .util import flatten_columns
from ._blob import (
BLOB_MODE_TO_HANDLING,
@@ -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
@@ -281,6 +282,24 @@ def _scanner_to_pandas(
return tbl.to_pandas(**kwargs)
def _set_blob_frame_values(df: "pd.DataFrame", path: str, values: list) -> None:
"""Replace a blob output column, including a blob inside a struct column."""
if path in df.columns:
df[path] = values
return
top, *nested = path.split(".")
if top not in df.columns or not nested:
raise ValueError(f"blob output column {path!r} is missing from query results")
for record, value in zip(df[top], values):
for part in nested[:-1]:
if record is None:
break
record = record[part]
if record is not None:
record[nested[-1]] = value
def _finish_plain_scan_pandas(
scanner: _LanceScanner,
*,
@@ -3018,14 +3037,19 @@ 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. Local plain scan queries
use Lance native conversion; remote queries use server blob fetch
APIs for blob v2 "bytes" and "lazy" modes.
**kwargs
Forwarded to pyarrow.Table.to_pandas after query execution and
optional flattening.
"""
validate_blob_mode(blob_mode)
if self._table is not None and not self._table._inner._is_native():
return await self._remote_to_pandas(
flatten=flatten, timeout=timeout, blob_mode=blob_mode, **kwargs
)
if hasattr(self._inner, "output_schema"):
schema = await self.output_schema()
if _blob_mode_requires_native_pandas(blob_mode, schema):
@@ -3053,6 +3077,54 @@ class AsyncQueryBase(object):
)
return tbl.to_pandas(**kwargs)
async def _remote_to_pandas(
self,
*,
flatten: Optional[Union[int, bool]],
timeout: Optional[timedelta],
blob_mode: BlobMode,
**kwargs,
) -> "pd.DataFrame":
# to_arrow() runs the query on the server and keeps the row ids needed
# by the server-side blob fetch APIs in the v2 descriptors.
table = self._table
assert table is not None
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)
schema = await table.schema()
blob_sources = blob_v2_projection_sources(
schema, self.to_query_object().columns
)
if not projected_blob_paths.issubset(blob_sources):
raise NotImplementedError(
f"remote to_pandas(blob_mode={blob_mode!r}) requires row-addressable "
"blob v2 columns"
)
values = {}
for output_name, source_name in blob_sources.items():
if output_name not in projected_blob_paths:
continue
if blob_mode == "bytes":
values[output_name] = (
await table.fetch_blobs(source_name, tbl)
).to_pylist()
else:
values[output_name] = await table.fetch_blob_files(source_name, tbl)
tbl = strip_auto_row_ids(tbl, self._blob_paths)
df = flatten_columns(tbl, flatten).to_pandas(**kwargs)
for output_name, column_values in values.items():
_set_blob_frame_values(df, output_name, column_values)
return df
async def _plain_scan_to_pandas(
self,
blob_mode: BlobMode,
+16 -2
View File
@@ -5423,7 +5423,11 @@ class AsyncTable:
return AsyncQuery(self._inner.query(), self)
async def to_lance(self, **kwargs) -> lance.LanceDataset:
"""Return the Lance dataset backing this table.
"""Return the Lance dataset backing a local table.
Remote tables cannot be opened as Lance datasets by the client. Use
[to_arrow][lancedb.table.AsyncTable.to_arrow] or
[query][lancedb.table.AsyncTable.query] to read them through the server.
Parameters
----------
@@ -5440,6 +5444,12 @@ class AsyncTable:
>>> async def get_lance_dataset(table):
... return await table.to_lance()
"""
if not self._inner._is_native():
raise NotImplementedError(
"to_lance() is not supported for remote tables; "
"query the server instead"
)
try:
import lance
except ImportError:
@@ -5469,7 +5479,8 @@ class AsyncTable:
Parameters
----------
blob_mode: str, default "lazy"
Controls how Lance blob columns are returned.
Controls how Lance blob columns are returned. On remote tables,
blob v2 data is fetched through the server for "bytes" and "lazy".
**kwargs
Forwarded to PyArrow / Lance pandas conversion.
@@ -5487,6 +5498,9 @@ class AsyncTable:
)
return arrow_tbl.to_pandas(**kwargs)
if not self._inner._is_native():
return await self.query().to_pandas(blob_mode=blob_mode, **kwargs)
if blob_mode == "lazy" and get_uri_scheme(await self.uri()) == "memory":
return (await self.to_arrow()).to_pandas(**kwargs)
if blob_mode == "bytes" and blob_v2_column_paths(schema):
+50 -1
View File
@@ -2421,7 +2421,7 @@ def blob_remote_table(*, server_version=Version("0.5.0")):
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["columns"] == ["id", "image"]
assert body.get("columns") in (None, ["id", "image"])
assert body["with_row_id"] is True
response_table = blob_query_response_table()
request.send_response(200)
@@ -2457,6 +2457,55 @@ def test_remote_blob_columns_and_fetch():
assert blobs.to_pylist() == [b"alpha", None, b"gamma"]
@pytest.mark.asyncio
async def test_async_remote_to_lance_rejects_storage_access(monkeypatch):
with blob_remote_table() as remote_table:
table = remote_table._table
async def fail_uri():
raise AssertionError("remote table URI must not be requested")
monkeypatch.setattr(table, "uri", fail_uri)
with pytest.raises(NotImplementedError, match="remote tables"):
await table.to_lance()
@pytest.mark.asyncio
@pytest.mark.parametrize("blob_mode", ["descriptions", "bytes", "lazy"])
async def test_async_remote_blob_to_pandas_uses_server(monkeypatch, blob_mode):
with blob_remote_table() as remote_table:
table = remote_table._table
async def fail_uri():
raise AssertionError("remote table URI must not be requested")
async def fail_to_lance():
raise AssertionError("remote table storage must not be opened")
monkeypatch.setattr(table, "uri", fail_uri)
monkeypatch.setattr(table, "to_lance", fail_to_lance)
pandas_kwargs = {} if blob_mode == "lazy" else {"blob_mode": blob_mode}
query_df = (
await table.query().select(["id", "image"]).to_pandas(**pandas_kwargs)
)
table_df = await table.to_pandas(**pandas_kwargs)
for df in (query_df, table_df):
assert df["id"].tolist() == [1, 2, 3]
assert "_rowid" not in df.columns
images = df["image"].tolist()
assert images[1] is None
if blob_mode == "bytes":
assert images == [b"alpha", None, b"gamma"]
elif blob_mode == "lazy":
assert await images[0].aread() == b"alpha"
assert await images[2].aread() == b"gamma"
else:
assert images[0]["size"] == 5
assert "_lance_row_id" not in images[0]
def test_remote_blob_files_are_lazy_seekable_handles():
with blob_remote_table() as table:
files = table.fetch_blob_files("image", [10, 20, 30])
+5
View File
@@ -882,6 +882,11 @@ impl Table {
self.inner.is_some()
}
/// Whether this table has a local Lance dataset that the client may open.
pub fn _is_native(&self) -> PyResult<bool> {
Ok(self.inner_ref()?.as_native().is_some())
}
/// Closes the table, releasing any resources associated with it.
pub fn close(&mut self) {
self.inner.take();