Skip to content

Commit

Permalink
chore(python): add docstring for limit behavior
Browse files Browse the repository at this point in the history
Closes #796
  • Loading branch information
changhiskhan committed Jan 9, 2024
1 parent 99ba533 commit 72bdcde
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 in any 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 72bdcde

Please sign in to comment.