Skip to content

Commit

Permalink
chore(python): add docstring for limit behavior (#800)
Browse files Browse the repository at this point in the history
Closes #796
  • Loading branch information
changhiskhan authored and westonpace committed Apr 5, 2024
1 parent 881dfa0 commit 7581cbb
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 2 deletions.
14 changes: 12 additions & 2 deletions python/lancedb/query.py
Original file line number Diff line number Diff line change
Expand Up @@ -260,20 +260,30 @@ def to_pydantic(self, model: Type[LanceModel]) -> List[LanceModel]:
for row in self.to_arrow().to_pylist()
]

def limit(self, limit: int) -> LanceQueryBuilder:
def limit(self, limit: Union[int, None]) -> LanceQueryBuilder:
"""Set the maximum number of results to return.
Parameters
----------
limit: int
The maximum number of results to return.
By default the query is limited to the first 10.
Call this method and pass 0, a negative value,
or None to remove the limit.
*WARNING* if you have a large dataset, removing
the limit can potentially result in reading a
large amount of data into memory and cause
out of memory issues.
Returns
-------
LanceQueryBuilder
The LanceQueryBuilder object.
"""
self._limit = limit
if limit is None or limit <= 0:
self._limit = None
else:
self._limit = limit
return self

def select(self, columns: list) -> LanceQueryBuilder:
Expand Down
8 changes: 8 additions & 0 deletions python/tests/test_table.py
Original file line number Diff line number Diff line change
Expand Up @@ -569,6 +569,14 @@ def test_empty_query(db):
val = df.id.iloc[0]
assert val == 1

table = LanceTable.create(db, "my_table2", data=[{"id": i} for i in range(100)])
df = table.search().select(["id"]).to_pandas()
assert len(df) == 10
df = table.search().select(["id"]).limit(None).to_pandas()
assert len(df) == 100
df = table.search().select(["id"]).limit(-1).to_pandas()
assert len(df) == 100


def test_compact_cleanup(db):
table = LanceTable.create(
Expand Down

0 comments on commit 7581cbb

Please sign in to comment.