rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2025-12-26 06:39:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(python): flatten in AsyncQuery (#1967)

Browse Source 
PR fixes #1949

---------

Co-authored-by: Will Jones <willjones127@gmail.com>
This commit is contained in:
Gagan Bhullar
2025-01-06 11:52:03 -07:00
committed by GitHub
parent 2c05ffed52
commit b474f98049
3 changed files with 66 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
python/python/tests/test_query.py
View File

@@ -52,6 +52,17 @@ async def table_async(tmp_path) -> AsyncTable:
return await conn.create_table("test", data)
@pytest_asyncio.fixture
async def table_struct_async(tmp_path) -> AsyncTable:
conn = await lancedb.connect_async(
tmp_path, read_consistency_interval=timedelta(seconds=0)
)
struct = pa.array([{"n_legs": 2, "animals": "Parrot"}, {"year": 2022, "n_legs": 4}])
month = pa.array([4, 6])
table = pa.Table.from_arrays([struct, month], names=["a", "month"])
return await conn.create_table("test_struct", table)
def test_cast(table):
class TestModel(LanceModel):
vector: Vector(2)
@@ -400,6 +411,15 @@ async def test_query_to_pandas_async(table_async: AsyncTable):
assert df.shape == (0, 4)
@pytest.mark.asyncio
async def test_query_to_pandas_flatten_async(table_struct_async: AsyncTable):
df = await table_struct_async.query().to_pandas()
assert df.shape == (2, 2)
df = await table_struct_async.query().to_pandas(flatten=True)
assert df.shape == (2, 4)
@pytest.mark.asyncio
async def test_query_to_polars_async(table_async: AsyncTable):
df = await table_async.query().to_polars()