mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-13 07:12:27 +00:00
fix(python): align sync remote methods with async behavior
This commit is contained in:
@@ -497,7 +497,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)
|
||||
@@ -552,7 +554,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,
|
||||
@@ -685,13 +689,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,
|
||||
@@ -719,19 +727,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
|
||||
@@ -740,7 +755,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:
|
||||
|
||||
@@ -819,6 +819,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,
|
||||
|
||||
@@ -7,6 +7,7 @@ import logging
|
||||
from functools import cached_property
|
||||
import os
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Callable,
|
||||
Dict,
|
||||
@@ -78,6 +79,10 @@ from ..table import (
|
||||
)
|
||||
from ..types import BaseTokenizerType
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import lance
|
||||
import polars as pl
|
||||
|
||||
|
||||
class RemoteTable(Table):
|
||||
def __init__(
|
||||
@@ -238,11 +243,32 @@ class RemoteTable(Table):
|
||||
|
||||
def to_arrow(self) -> pa.Table:
|
||||
"""to_arrow() is not yet supported on LanceDB cloud."""
|
||||
raise NotImplementedError("to_arrow() is not yet supported on LanceDB cloud.")
|
||||
raise NotImplementedError(
|
||||
"to_arrow() is not yet supported on LanceDB cloud. "
|
||||
"Use table.search().to_arrow() instead."
|
||||
)
|
||||
|
||||
def to_pandas(self, blob_mode: BlobMode = "lazy", **kwargs):
|
||||
"""to_pandas() is not yet supported on LanceDB cloud."""
|
||||
raise NotImplementedError("to_pandas() is not yet supported on LanceDB cloud.")
|
||||
raise NotImplementedError(
|
||||
"to_pandas() is not yet supported on LanceDB cloud. "
|
||||
"Use table.search().to_pandas() instead."
|
||||
)
|
||||
|
||||
def to_lance(self, **kwargs) -> "lance.LanceDataset":
|
||||
"""Return the Lance dataset at this table handle's version and branch.
|
||||
|
||||
Requires the server to provide a dataset location accessible to the client.
|
||||
Keyword arguments are forwarded to ``lance.dataset``.
|
||||
"""
|
||||
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 on LanceDB cloud. "
|
||||
"Use table.search().to_polars() instead."
|
||||
)
|
||||
|
||||
def checkout(self, version: Union[int, str]):
|
||||
result = LOOP.run(self._table.checkout(version))
|
||||
|
||||
@@ -1222,13 +1222,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
|
||||
@@ -1237,7 +1237,7 @@ class Table(ABC):
|
||||
of this Table
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
raise NotImplementedError("schema is not supported for this table type")
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
@@ -1266,7 +1266,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":
|
||||
@@ -1275,11 +1275,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"""
|
||||
@@ -1302,7 +1302,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.
|
||||
@@ -1329,20 +1329,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.
|
||||
|
||||
For remote tables, the server must provide a dataset location that is
|
||||
accessible to the client.
|
||||
|
||||
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
|
||||
@@ -1353,7 +1359,7 @@ class Table(ABC):
|
||||
-------
|
||||
polars.LazyFrame
|
||||
"""
|
||||
raise NotImplementedError
|
||||
raise NotImplementedError("to_polars is not supported for this table type")
|
||||
|
||||
# New unified API overload
|
||||
@overload
|
||||
@@ -1458,7 +1464,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,
|
||||
@@ -1476,7 +1482,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:
|
||||
"""
|
||||
@@ -1496,7 +1504,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)
|
||||
@@ -1513,14 +1521,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(
|
||||
@@ -1585,7 +1593,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,
|
||||
@@ -1706,7 +1716,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(
|
||||
@@ -1779,7 +1791,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:
|
||||
"""
|
||||
@@ -1931,7 +1943,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(
|
||||
@@ -2193,7 +2205,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(
|
||||
@@ -2260,7 +2272,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(
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -616,10 +616,113 @@ 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 f"search().{method}()" in message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("version", [None, 1])
|
||||
async def test_remote_to_lance(tmp_path, version):
|
||||
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")
|
||||
|
||||
def handler(request):
|
||||
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()
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
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")
|
||||
with pytest.raises(NotImplementedError) as async_error:
|
||||
await table.to_lance()
|
||||
|
||||
assert "Table URI not supported by 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():
|
||||
|
||||
Reference in New Issue
Block a user