feat(python): add to_list to async api (#1520)

PR fixes #1517
This commit is contained in:
Gagan Bhullar
2024-08-08 12:45:20 -06:00
committed by GitHub
parent f9d5fa88a1
commit 9c1adff426
2 changed files with 18 additions and 0 deletions

View File

@@ -1217,6 +1217,16 @@ class AsyncQueryBase(object):
await batch_iter.read_all(), schema=batch_iter.schema
)
async def to_list(self) -> List[dict]:
"""
Execute the query and return the results as a list of dictionaries.
Each list entry is a dictionary with the selected column names as keys,
or all table columns if `select` is not called. The vector and the "_distance"
fields are returned whether or not they're explicitly selected.
"""
return (await self.to_arrow()).to_pylist()
async def to_pandas(self) -> "pd.DataFrame":
"""
Execute the query and collect the results into a pandas DataFrame.

View File

@@ -354,3 +354,11 @@ async def test_query_camelcase_async(tmp_path):
result = await table.query().select(["camelCase"]).to_arrow()
assert result == pa.table({"camelCase": pa.array([1, 2])})
@pytest.mark.asyncio
async def test_query_to_list_async(table_async: AsyncTable):
list = await table_async.query().to_list()
assert len(list) == 2
assert list[0]["vector"] == [1, 2]
assert list[1]["vector"] == [3, 4]