diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index af964d459..a8f02e458 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -423,6 +423,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( diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index d9755afa1..6f68692a2 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -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 _validate_query_vector, flatten_columns from . import _wal_hybrid # WAL-PK-FUSION: delete. from ._blob import ( @@ -51,6 +51,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 @@ -3073,14 +3074,18 @@ 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. Remote queries support + "descriptions"; "bytes" and "lazy" are not yet supported. **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): @@ -3108,6 +3113,31 @@ 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": + # 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'" + ) + async def _plain_scan_to_pandas( self, blob_mode: BlobMode, diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 3c741cf74..b21010681 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -259,10 +259,9 @@ class RemoteTable(Table): ) def to_lance(self, **kwargs) -> "lance.LanceDataset": - """Return the Lance dataset at this table handle's version and branch. + """Direct Lance dataset access is not supported for remote tables. - Requires the server to provide a dataset location accessible to the client. - Keyword arguments are forwarded to ``lance.dataset``. + Use ``table.search().to_arrow()`` to read query results through the server. """ return LOOP.run(self._table.to_lance(**kwargs)) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 3c6e6893b..a1821c8cf 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -1332,10 +1332,10 @@ class Table(ABC): raise NotImplementedError("to_arrow is not supported for this table type") def to_lance(self, **kwargs) -> lance.LanceDataset: - """Return the table as a lance.LanceDataset. + """Return a local table as a lance.LanceDataset. - For remote tables, the server must provide a dataset location that is - accessible to the client. + Remote tables cannot be opened as Lance datasets by the client. Use + ``table.search().to_arrow()`` to read them through the server. Returns ------- @@ -5614,7 +5614,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 ---------- @@ -5631,6 +5635,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: @@ -5660,7 +5670,8 @@ class AsyncTable: Parameters ---------- blob_mode: str, default "lazy" - Controls how Lance blob columns are returned. + Controls how Lance blob columns are returned. Remote tables support + "descriptions"; "bytes" and "lazy" are not yet supported. **kwargs Forwarded to PyArrow / Lance pandas conversion. @@ -5678,6 +5689,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): diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index ceeac224e..f8c78c492 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -807,67 +807,55 @@ def test_table_unimplemented_functions(): @pytest.mark.asyncio @pytest.mark.parametrize("version", [None, 1]) -async def test_remote_to_lance(tmp_path, version): +@pytest.mark.parametrize("expose_location", [False, True]) +async def test_remote_to_lance_rejection( + tmp_path, version, expose_location, monkeypatch +): import lance location = tmp_path / "test.lance" lance.write_dataset(pa.table({"id": [1]}), location) lance.write_dataset(pa.table({"id": [2]}), location, mode="append") + requests = [] + + def fail_dataset(*args, **kwargs): + raise AssertionError("remote storage must not be opened") + + monkeypatch.setattr(lance, "dataset", fail_dataset) def handler(request): + requests.append(request.path) assert request.path == "/v1/table/test/describe/" content_len = int(request.headers.get("Content-Length", 0)) body = json.loads(request.rfile.read(content_len)) if content_len else {} request.send_response(200) request.send_header("Content-Type", "application/json") request.end_headers() - request.wfile.write( - json.dumps( - { - "version": body.get("version") or 2, - "schema": {"fields": []}, - "location": str(location), - } - ).encode() - ) + response = { + "version": body.get("version") or 2, + "schema": {"fields": []}, + } + if expose_location: + response["location"] = str(location) + request.wfile.write(json.dumps(response).encode()) options = {"default_scan_options": {"with_row_id": True}} with mock_lancedb_connection(handler) as db: table = db.open_table("test", version=version) - sync_dataset = table.to_lance(**options) + opened_requests = list(requests) + with pytest.raises(NotImplementedError) as sync_error: + table.to_lance(**options) + assert requests == opened_requests async with mock_lancedb_connection_async(handler) as db: table = await db.open_table("test", version=version) - async_dataset = await table.to_lance(**options) - - for dataset in [sync_dataset, async_dataset]: - assert isinstance(dataset, lance.LanceDataset) - assert dataset.version == (version or 2) - result = dataset.to_table() - assert result["id"].to_pylist() == ([1] if version == 1 else [1, 2]) - assert "_rowid" in result.column_names - - -@pytest.mark.asyncio -async def test_remote_to_lance_without_server_location(): - def handler(request): - assert request.path == "/v1/table/test/describe/" - request.send_response(200) - request.send_header("Content-Type", "application/json") - request.end_headers() - request.wfile.write(b'{"version": 1, "schema": {"fields": []}}') - - with mock_lancedb_connection(handler) as db: - table = db.open_table("test") - with pytest.raises(NotImplementedError) as sync_error: - table.to_lance() - - async with mock_lancedb_connection_async(handler) as db: - table = await db.open_table("test") + opened_requests = list(requests) with pytest.raises(NotImplementedError) as async_error: - await table.to_lance() + await table.to_lance(**options) + assert requests == opened_requests - assert "Table URI not supported by the server" in str(sync_error.value) + assert "to_lance() is not supported for remote tables" in str(sync_error.value) + assert "query the server" in str(sync_error.value) assert str(sync_error.value) == str(async_error.value) @@ -3284,14 +3272,17 @@ def blob_descriptor_take_table(): @contextlib.contextmanager -def blob_remote_table(*, server_version=Version("0.5.0")): +def blob_remote_table(*, server_version=Version("0.5.0"), location=None): 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", str(server_version)) request.end_headers() - request.wfile.write(json.dumps(BLOB_DESCRIBE_RESPONSE).encode()) + response = dict(BLOB_DESCRIBE_RESPONSE) + if location is not None: + response["location"] = str(location) + request.wfile.write(json.dumps(response).encode()) elif request.path.startswith("/v1/table/test/blob/image/"): path = request.path.partition("?")[0] row_id = int(path.split("/")[-2]) @@ -3321,14 +3312,19 @@ 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["with_row_id"] is True - if body["columns"] == ["image"]: - # fetch_blob_files sizes its handles from a descriptor take. - assert body["filter"].startswith("_rowid IN") - response_table = blob_descriptor_take_table() + columns = body.get("columns") + assert columns in (None, ["id", "image"], ["id"], ["image"]) + if columns == ["id"]: + response_table = blob_query_response_table().select(["id"]) else: - assert body["columns"] == ["id", "image"] - response_table = blob_query_response_table() + if columns == ["image"]: + # fetch_blob_files sizes its handles from a descriptor take. + assert body["filter"].startswith("_rowid IN") + response_table = blob_descriptor_take_table() + else: + response_table = blob_query_response_table() + if not body.get("with_row_id", False): + response_table = response_table.drop(["_rowid"]) request.send_response(200) request.send_header("Content-Type", "application/vnd.apache.arrow.file") request.end_headers() @@ -3362,6 +3358,106 @@ def test_remote_blob_columns_and_fetch(): assert blobs.to_pylist() == [b"alpha", None, b"gamma"] +@pytest.mark.parametrize("blob_mode", ["descriptions", "bytes", "lazy"]) +def test_sync_remote_blob_pandas_rejects_storage_access( + tmp_path, monkeypatch, blob_mode +): + import lance + + location = tmp_path / "remote.lance" + lance.write_dataset(pa.table({"id": [7], "image": [b"server-owned"]}), location) + + def fail_dataset(*args, **kwargs): + raise AssertionError("remote storage must not be opened") + + monkeypatch.setattr(lance, "dataset", fail_dataset) + with blob_remote_table(location=location) as table: + with pytest.raises(RuntimeError, match="to_lance.*remote tables") as error: + table.search().select(["id", "image"]).to_pandas(blob_mode=blob_mode) + + assert isinstance(error.value.__cause__, NotImplementedError) + + +@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 +async def test_async_remote_blob_descriptions_to_pandas_uses_server(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") + + 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) + + query_df = await ( + table.query().select(["id", "image"]).to_pandas(blob_mode="descriptions") + ) + table_df = await table.to_pandas(blob_mode="descriptions") + + 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 + assert images[0]["size"] == 5 + assert "_lance_row_id" not in images[0] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("blob_mode", ["bytes", "lazy"]) +async def test_async_remote_blob_pandas_rejects_unpinned_fetch(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") + + async def fail_fetch(*args, **kwargs): + raise AssertionError("blob fetch must not use an unpinned remote table") + + monkeypatch.setattr(table, "uri", fail_uri) + monkeypatch.setattr(table, "to_lance", fail_to_lance) + monkeypatch.setattr(table, "fetch_blobs", fail_fetch) + monkeypatch.setattr(table, "fetch_blob_files", fail_fetch) + + pandas_kwargs = {} if blob_mode == "lazy" else {"blob_mode": blob_mode} + with pytest.raises(NotImplementedError, match="stable table snapshot"): + await table.query().select(["id", "image"]).to_pandas(**pandas_kwargs) + with pytest.raises(NotImplementedError, match="stable table snapshot"): + await table.to_pandas(**pandas_kwargs) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("blob_mode", ["bytes", "lazy"]) +async def test_async_remote_nonblob_projection_to_pandas(blob_mode): + with blob_remote_table() as remote_table: + df = await ( + remote_table._table.query().select(["id"]).to_pandas(blob_mode=blob_mode) + ) + + assert df["id"].tolist() == [1, 2, 3] + assert list(df.columns) == ["id"] + + def test_remote_blob_files_are_lazy_seekable_handles(): with blob_remote_table() as table: files = table.fetch_blob_files("image", [10, 20, 30, 40]) diff --git a/python/src/table.rs b/python/src/table.rs index 7102a8aa7..d82ebbaaf 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -917,6 +917,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 { + 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();