From 9c1adff4268fc739de36ca8d0f0dfebd2114b726 Mon Sep 17 00:00:00 2001 From: Gagan Bhullar Date: Thu, 8 Aug 2024 12:45:20 -0600 Subject: [PATCH] feat(python): add to_list to async api (#1520) PR fixes #1517 --- python/python/lancedb/query.py | 10 ++++++++++ python/python/tests/test_query.py | 8 ++++++++ 2 files changed, 18 insertions(+) diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index 7f76c78c..8564575f 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -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. diff --git a/python/python/tests/test_query.py b/python/python/tests/test_query.py index c569cd49..30eba26e 100644 --- a/python/python/tests/test_query.py +++ b/python/python/tests/test_query.py @@ -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]