diff --git a/python/lancedb/query.py b/python/lancedb/query.py index 532566da..886ec966 100644 --- a/python/lancedb/query.py +++ b/python/lancedb/query.py @@ -260,20 +260,30 @@ class LanceQueryBuilder(ABC): 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: diff --git a/python/tests/test_table.py b/python/tests/test_table.py index 4c769009..3f096a53 100644 --- a/python/tests/test_table.py +++ b/python/tests/test_table.py @@ -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(