diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 90dc1cccf..74acc28ba 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -548,7 +548,9 @@ class DBConnection(EnforceOverrides): LanceTable(name='table4', ...) """ - raise NotImplementedError + raise NotImplementedError( + "create_table is not supported for this connection type" + ) def __getitem__(self, name: str) -> LanceTable: return self.open_table(name) @@ -603,7 +605,9 @@ class DBConnection(EnforceOverrides): ------- A LanceTable object representing the table. """ - raise NotImplementedError + raise NotImplementedError( + "open_table is not supported for this connection type" + ) def create_materialized_view( self, @@ -743,13 +747,17 @@ class DBConnection(EnforceOverrides): """ if namespace_path is None: namespace_path = [] - raise NotImplementedError + raise NotImplementedError( + "drop_table is not supported for this connection type" + ) def drop_table_async( self, name: str, namespace_path: Optional[List[str]] = None ) -> Job: """Start dropping a table and return its cleanup job.""" - raise NotImplementedError + raise NotImplementedError( + "drop_table_async is not supported for this connection type" + ) def rename_table( self, @@ -777,19 +785,26 @@ class DBConnection(EnforceOverrides): cur_namespace_path = [] if new_namespace_path is None: new_namespace_path = [] - raise NotImplementedError + raise NotImplementedError( + "rename_table is not supported for this connection type" + ) def drop_database(self): """ Drop database This is the same thing as dropping all the tables """ - raise NotImplementedError + raise NotImplementedError( + "drop_database is not supported for this connection type" + ) def drop_all_tables(self, namespace_path: Optional[List[str]] = None): """ Drop all tables from the database + Remote connections do not currently support this operation. Use + ``list_tables()`` and ``drop_table()`` to drop tables individually. + Parameters ---------- namespace_path: List[str], optional @@ -798,7 +813,9 @@ class DBConnection(EnforceOverrides): """ if namespace_path is None: namespace_path = [] - raise NotImplementedError + raise NotImplementedError( + "drop_all_tables is not supported for this connection type" + ) @property def uri(self) -> str: diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 08125cc42..5e9288471 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -841,6 +841,20 @@ class RemoteDBConnection(DBConnection): job = LOOP.run(self._conn.drop_table_async(name, namespace_path=namespace_path)) return Job(job if isinstance(job, AsyncJob) else AsyncJob(job)) + @override + def drop_all_tables(self, namespace_path: Optional[List[str]] = None): + """Drop all tables from the database. + + The remote API does not currently support this operation. Use + ``list_tables()`` and ``drop_table()`` to drop tables individually. + """ + LOOP.run(self._conn.drop_all_tables(namespace_path=namespace_path)) + + @override + def drop_database(self): + """Deprecated alias for ``drop_all_tables()``.""" + LOOP.run(self._conn.drop_database()) + @override def rename_table( self, diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 0744410b3..53591182e 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -7,6 +7,7 @@ import logging from functools import cached_property import os from typing import ( + TYPE_CHECKING, Any, Callable, Dict, @@ -82,6 +83,10 @@ from ..table import ( from ..types import BaseTokenizerType from ..util import infer_vector_column_name +if TYPE_CHECKING: + import lance + import polars as pl + class RemoteTable(Table): def __init__( @@ -243,12 +248,32 @@ class RemoteTable(Table): def to_arrow(self, *, blob_mode: BlobMode = "lazy") -> pa.Table: """to_arrow() is not supported for remote tables.""" validate_blob_mode(blob_mode) - raise NotImplementedError("to_arrow() is not supported for remote tables.") + raise NotImplementedError( + "to_arrow() is not supported for remote tables. " + "Use table.search().to_arrow() instead." + ) def to_pandas(self, blob_mode: BlobMode = "lazy", **kwargs): """to_pandas() is not supported for remote tables.""" validate_blob_mode(blob_mode) - raise NotImplementedError("to_pandas() is not supported for remote tables.") + raise NotImplementedError( + "to_pandas() is not supported for remote tables. " + "Use table.search().to_pandas() instead." + ) + + def to_lance(self, **kwargs) -> "lance.LanceDataset": + """Direct Lance dataset access is not supported for remote tables. + + Use ``table.search().to_arrow()`` to read query results through the server. + """ + return LOOP.run(self._table.to_lance(**kwargs)) + + def to_polars(self, **kwargs) -> "pl.LazyFrame": + """Use ``table.search().to_polars()`` to read remote query results.""" + raise NotImplementedError( + "to_polars() is not supported for remote tables. " + "Use table.search().to_polars() instead." + ) def checkout(self, version: Union[int, str]): result = LOOP.run(self._table.checkout(version)) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index b9557f219..c08c6390e 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -1356,13 +1356,13 @@ class Table(ABC): @abstractmethod def name(self) -> str: """The name of this Table""" - raise NotImplementedError + raise NotImplementedError("name is not supported for this table type") @property @abstractmethod def version(self) -> int: """The version of this Table""" - raise NotImplementedError + raise NotImplementedError("version is not supported for this table type") @property @abstractmethod @@ -1371,7 +1371,7 @@ class Table(ABC): of this Table """ - raise NotImplementedError + raise NotImplementedError("schema is not supported for this table type") @property @abstractmethod @@ -1400,7 +1400,7 @@ class Table(ABC): tags = table.tags.list() """ - raise NotImplementedError + raise NotImplementedError("tags is not supported for this table type") @property def branches(self) -> "Branches": @@ -1409,11 +1409,11 @@ class Table(ABC): Branches are isolated, writable lines of history forked from another branch (or version). Writes on a branch do not affect ``main``. """ - raise NotImplementedError + raise NotImplementedError("branches is not supported for this table type") def current_branch(self) -> Optional[str]: """The branch this table handle is scoped to, or ``None`` for ``main``.""" - raise NotImplementedError + raise NotImplementedError("current_branch is not supported for this table type") def __len__(self) -> int: """The number of rows in this Table""" @@ -1436,7 +1436,7 @@ class Table(ABC): filter: str, optional A SQL where clause to filter the rows to count. """ - raise NotImplementedError + raise NotImplementedError("count_rows is not supported for this table type") def to_pandas(self, blob_mode: BlobMode = "lazy", **kwargs) -> "pandas.DataFrame": """Return the table as a pandas DataFrame. @@ -1471,20 +1471,26 @@ class Table(ABC): ------- pa.Table """ - raise NotImplementedError + 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. + + Remote tables cannot be opened as Lance datasets by the client. Use + ``table.search().to_arrow()`` to read them through the server. Returns ------- lance.LanceDataset """ - raise NotImplementedError + raise NotImplementedError("to_lance is not supported for this table type") def to_polars(self, **kwargs) -> "pl.LazyFrame": """Return the table as a Polars LazyFrame. + For remote tables, use ``table.search().to_polars()`` to convert query + results to a Polars DataFrame. + Note ---- The Polars streaming engine is not supported because it does not currently @@ -1495,7 +1501,7 @@ class Table(ABC): ------- polars.LazyFrame """ - raise NotImplementedError + raise NotImplementedError("to_polars is not supported for this table type") # New unified API overload @overload @@ -1600,7 +1606,7 @@ class Table(ABC): ... "l2", vector_column_name="vector" ... ) """ - raise NotImplementedError + raise NotImplementedError("create_index is not supported for this table type") def create_index_async( self, @@ -1618,7 +1624,9 @@ class Table(ABC): be complete when returned; callers must not assume the index exists until :meth:`Job.wait` returns. """ - raise NotImplementedError + raise NotImplementedError( + "create_index_async is not supported for this table type" + ) def drop_index(self, name: str) -> None: """ @@ -1638,7 +1646,7 @@ class Table(ABC): Use [list_indices][lancedb.table.Table.list_indices] to find the names of the indices. """ - raise NotImplementedError + raise NotImplementedError("drop_index is not supported for this table type") def wait_for_index( self, index_names: Iterable[str], timeout: timedelta = timedelta(seconds=300) @@ -1655,14 +1663,14 @@ class Table(ABC): timeout: timedelta Timeout to wait for asynchronous indexing. The default is 5 minutes. """ - raise NotImplementedError + raise NotImplementedError("wait_for_index is not supported for this table type") @abstractmethod def stats(self) -> TableStatistics: """ Retrieve table and fragment statistics. """ - raise NotImplementedError + raise NotImplementedError("stats is not supported for this table type") @abstractmethod def create_scalar_index( @@ -1727,7 +1735,9 @@ class Table(ABC): ``my_col = 0 OR not_indexed = 1`` will not be able to use any scalar index on ``my_col``. """ - raise NotImplementedError + raise NotImplementedError( + "create_scalar_index is not supported for this table type" + ) def create_fts_index( self, @@ -1853,7 +1863,9 @@ class Table(ABC): ``LANCE_LANGUAGE_MODEL_HOME`` to override the default platform data directory under ``lance/language_models``. """ - raise NotImplementedError + raise NotImplementedError( + "create_fts_index is not supported for this table type" + ) @abstractmethod def add( @@ -1927,7 +1939,7 @@ class Table(ABC): AddResult An object containing the new version number of the table after adding data. """ - raise NotImplementedError + raise NotImplementedError("add is not supported for this table type") def merge_insert(self, on: Union[str, Iterable[str]]) -> LanceMergeInsertBuilder: """ @@ -2079,7 +2091,7 @@ class Table(ABC): - and also the "_distance" column which is the distance between the query vector and the returned vector. """ - raise NotImplementedError + raise NotImplementedError("search is not supported for this table type") @abstractmethod def take_offsets( @@ -2365,7 +2377,7 @@ class Table(ABC): x vector 0 3 [5.0, 6.0] """ - raise NotImplementedError + raise NotImplementedError("delete is not supported for this table type") @abstractmethod def update( @@ -2432,7 +2444,7 @@ class Table(ABC): 1 4 [5.0, 6.0] 2 3 [10.0, 10.0] """ - raise NotImplementedError + raise NotImplementedError("update is not supported for this table type") @abstractmethod def cleanup_old_versions( diff --git a/python/python/tests/test_errors.py b/python/python/tests/test_errors.py index 338b6a0d7..0e362e5d8 100644 --- a/python/python/tests/test_errors.py +++ b/python/python/tests/test_errors.py @@ -3,7 +3,28 @@ import pickle +import pytest + +from lancedb.db import DBConnection from lancedb.remote.errors import HttpError, LanceDBClientError, RetryError +from lancedb.table import Table + + +@pytest.mark.parametrize( + "base, method", + [ + (Table, "to_lance"), + (Table, "to_polars"), + (DBConnection, "drop_all_tables"), + (DBConnection, "drop_database"), + ], +) +def test_base_unsupported_methods_have_messages(base, method): + with pytest.raises(NotImplementedError) as exc_info: + getattr(base, method)(None) + message = str(exc_info.value) + assert method in message + assert "not supported" in message def test_pickle_lancedb_client_error(): diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 9655f5807..c1db92948 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -857,10 +857,116 @@ def test_table_unimplemented_functions(): with mock_lancedb_connection(handler) as db: table = db.create_table("test", [{"id": 1}]) - with pytest.raises(NotImplementedError): - table.to_arrow() - with pytest.raises(NotImplementedError): - table.to_pandas() + for method in ["to_arrow", "to_pandas", "to_polars"]: + with pytest.raises(NotImplementedError) as exc_info: + getattr(table, method)() + message = str(exc_info.value) + assert method in message + assert "not" in message and "supported" in message + assert "remote tables" in message + assert "cloud" not in message.lower() + assert f"search().{method}()" in message + + +@pytest.mark.parametrize("method", ["to_arrow", "to_pandas"]) +@pytest.mark.parametrize("blob_mode", ["lazy", "descriptions", "bytes", "invalid"]) +def test_remote_conversion_blob_mode_errors(method, blob_mode): + with blob_remote_table() as table: + if blob_mode == "invalid": + with pytest.raises(ValueError, match="blob_mode must be one of"): + getattr(table, method)(blob_mode=blob_mode) + else: + with pytest.raises(NotImplementedError) as error: + getattr(table, method)(blob_mode=blob_mode) + assert f"search().{method}()" in str(error.value) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("version", [None, 1]) +@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() + 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) + 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) + opened_requests = list(requests) + with pytest.raises(NotImplementedError) as async_error: + await table.to_lance(**options) + assert requests == opened_requests + + 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) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "method, kwargs", + [ + ("drop_all_tables", {}), + ("drop_all_tables", {"namespace_path": []}), + ("drop_all_tables", {"namespace_path": ["team"]}), + ("drop_database", {}), + ], +) +async def test_remote_drop_all_tables_not_supported(method, kwargs): + requests = [] + + def handler(request): + requests.append(request.path) + request.send_response(500) + request.end_headers() + + with mock_lancedb_connection(handler) as db: + with pytest.raises(NotImplementedError) as sync_error: + getattr(db, method)(**kwargs) + + async with mock_lancedb_connection_async(handler) as db: + with pytest.raises(NotImplementedError) as async_error: + await getattr(db, method)(**kwargs) + + assert "Dropping all tables is not currently supported in the remote API" in str( + sync_error.value + ) + assert "drop_table" in str(sync_error.value) + assert str(sync_error.value) == str(async_error.value) + assert requests == [] def test_table_to_pandas_not_supported(): @@ -1681,10 +1787,13 @@ def test_remote_python_only_messages(caplog): table = db.open_table("test", storage_options={}, index_cache_size=1) assert "storage_options is ignored for remote tables" in caplog.text assert "index_cache_size is ignored for remote tables" in caplog.text - for method in ["to_arrow", "to_pandas"]: + for method in ["to_arrow", "to_pandas", "to_polars"]: with pytest.raises(NotImplementedError) as error: getattr(table, method)() - assert str(error.value) == f"{method}() is not supported for remote tables." + assert str(error.value) == ( + f"{method}() is not supported for remote tables. " + f"Use table.search().{method}() instead." + ) for method in ["compact_files", "cleanup_old_versions"]: with pytest.warns(UserWarning) as warning: getattr(table, method)() @@ -3490,14 +3599,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]) @@ -3532,7 +3644,6 @@ def blob_remote_table(*, server_version=Version("0.5.0")): if columns == ["id"]: response_table = blob_query_response_table().select(["id"]) else: - assert body["with_row_id"] is True if columns == ["image"]: # Blob handles are sized from a descriptor take. assert body["filter"].startswith("_rowid IN") @@ -3545,6 +3656,8 @@ def blob_remote_table(*, server_version=Version("0.5.0")): ) 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.send_header("x-lancedb-version", "1") @@ -3579,6 +3692,32 @@ 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_avoids_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: + monkeypatch.setattr(table, "to_lance", fail_dataset) + query = table.search().select(["id", "image"]) + if blob_mode == "descriptions": + df = query.to_pandas(blob_mode=blob_mode) + assert df["id"].tolist() == [1, 2, 3] + assert "_rowid" not in df.columns + assert df["image"][0]["size"] == 5 + else: + with pytest.raises(NotImplementedError, match="stable table snapshot"): + query.to_pandas(blob_mode=blob_mode) + + @pytest.mark.asyncio @pytest.mark.parametrize("use_async", [False, True]) async def test_remote_fetch_blob_ranges(use_async): diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 8a456ab3b..cad60ab7d 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -1707,7 +1707,9 @@ impl Database for RemoteDatabase { // TODO: Implement namespace-aware drop_all_tables let _namespace_path = namespace_path; // Suppress unused warning for now Err(crate::Error::NotSupported { - message: "Dropping all tables is not currently supported in the remote API".to_string(), + message: "Dropping all tables is not currently supported in the remote API. \ + List tables and call drop_table for each table instead." + .to_string(), }) }