mirror of
https://github.com/lancedb/lancedb.git
synced 2026-05-24 07:20:40 +00:00
feat(python): align to_pandas pandas kwargs (#3397)
## Feature This PR aligns LanceDB Python `to_pandas()` APIs with Lance pandas conversion capabilities while keeping LanceDB query-specific semantics intact. - Adds `blob_mode` and pandas `**kwargs` support to local table `to_pandas()`. - Delegates local `LanceTable.to_pandas()` to Lance dataset `to_pandas(blob_mode=..., **kwargs)`. - Keeps remote table `to_pandas()` unsupported with `NotImplementedError`. - Allows sync and async query `to_pandas()` to forward pandas kwargs after LanceDB `flatten` and `timeout` handling. Why we need this feature: Users can access Lance blob-aware pandas conversion from LanceDB local tables and can pass PyArrow pandas conversion options through table/query APIs without losing existing `flatten` or `timeout` behavior. How it works: The table API exposes a `BlobMode` literal type for `lazy`, `bytes`, and `descriptions`. Local tables call through to the backing Lance dataset. Query APIs do not add `blob_mode`; they materialize Arrow results, apply LanceDB flattening when requested, and then call `to_pandas(**kwargs)`. ## Validation - `uv run --frozen --extra tests pytest python/tests/test_table.py::test_table_to_pandas_default_matches_arrow python/tests/test_table.py::test_table_to_pandas_blob_bytes python/tests/test_table.py::test_table_to_pandas_kwargs python/tests/test_query.py::test_query_to_pandas_kwargs python/tests/test_query.py::test_query_timeout python/tests/test_remote_db.py::test_table_to_pandas_not_supported` - `uv run --frozen --extra dev ruff check python/lancedb/table.py python/lancedb/query.py python/lancedb/remote/table.py python/tests/test_table.py python/tests/test_query.py python/tests/test_remote_db.py` - `uv run --frozen --extra tests pytest python/tests/test_table.py python/tests/test_query.py python/tests/test_remote_db.py` Note: `python/uv.lock` was intentionally not committed in this branch.
This commit is contained in:
@@ -47,6 +47,85 @@ def test_basic(mem_db: DBConnection):
|
||||
assert table.to_arrow() == expected_data
|
||||
|
||||
|
||||
def test_table_to_pandas_default_matches_arrow(tmp_db: DBConnection):
|
||||
pd = pytest.importorskip("pandas")
|
||||
data = pa.table({"id": [1, 2], "text": ["one", "two"]})
|
||||
table = tmp_db.create_table("test_to_pandas_old_call", data=data)
|
||||
|
||||
expected = data.to_pandas()
|
||||
pd.testing.assert_frame_equal(table.to_pandas(), expected)
|
||||
|
||||
|
||||
def test_table_to_pandas_blob_bytes(tmp_db: DBConnection):
|
||||
pytest.importorskip("lance")
|
||||
data = pa.table(
|
||||
{
|
||||
"id": pa.array([1, 2], pa.int64()),
|
||||
"blob": pa.array([b"hello", b"world"], pa.large_binary()),
|
||||
},
|
||||
schema=pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field(
|
||||
"blob", pa.large_binary(), metadata={"lance-encoding:blob": "true"}
|
||||
),
|
||||
]
|
||||
),
|
||||
)
|
||||
table = tmp_db.create_table("test_to_pandas_blob_bytes", data=data)
|
||||
|
||||
df = table.to_pandas(blob_mode="bytes")
|
||||
|
||||
assert df["blob"].tolist() == [b"hello", b"world"]
|
||||
|
||||
|
||||
def test_table_to_pandas_kwargs(tmp_db: DBConnection):
|
||||
pd = pytest.importorskip("pandas")
|
||||
data = pa.table({"id": pa.array([1, 2], pa.int64())})
|
||||
table = tmp_db.create_table("test_to_pandas_kwargs", data=data)
|
||||
|
||||
df = table.to_pandas(types_mapper=pd.ArrowDtype)
|
||||
|
||||
assert str(df["id"].dtype) == "int64[pyarrow]"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_table_to_pandas_blob_bytes(tmp_db_async: AsyncConnection):
|
||||
pytest.importorskip("lance")
|
||||
data = pa.table(
|
||||
{
|
||||
"id": pa.array([1, 2], pa.int64()),
|
||||
"blob": pa.array([b"hello", b"world"], pa.large_binary()),
|
||||
},
|
||||
schema=pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field(
|
||||
"blob", pa.large_binary(), metadata={"lance-encoding:blob": "true"}
|
||||
),
|
||||
]
|
||||
),
|
||||
)
|
||||
table = await tmp_db_async.create_table(
|
||||
"test_async_to_pandas_blob_bytes", data=data
|
||||
)
|
||||
|
||||
df = await table.to_pandas(blob_mode="bytes")
|
||||
|
||||
assert df["blob"].tolist() == [b"hello", b"world"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_table_to_pandas_kwargs(tmp_db_async: AsyncConnection):
|
||||
pd = pytest.importorskip("pandas")
|
||||
data = pa.table({"id": pa.array([1, 2], pa.int64())})
|
||||
table = await tmp_db_async.create_table("test_async_to_pandas_kwargs", data=data)
|
||||
|
||||
df = await table.to_pandas(types_mapper=pd.ArrowDtype)
|
||||
|
||||
assert str(df["id"].dtype) == "int64[pyarrow]"
|
||||
|
||||
|
||||
def test_create_table_infers_large_int_vectors(mem_db: DBConnection):
|
||||
data = [{"vector": [0, 300]}]
|
||||
|
||||
|
||||
Reference in New Issue
Block a user