diff --git a/python/python/lancedb/integrations/pyarrow.py b/python/python/lancedb/integrations/pyarrow.py index 085964690..1626a6b57 100644 --- a/python/python/lancedb/integrations/pyarrow.py +++ b/python/python/lancedb/integrations/pyarrow.py @@ -52,6 +52,9 @@ class PyarrowScannerAdapter(pa.dataset.Scanner): raise NotImplementedError("memory_pool not supported") def count_rows(self): + if isinstance(self.filter, pa.compute.Expression): + # stream the batches so memory does not grow with the matching rows + return sum(batch.num_rows for batch in self.to_reader()) return self.table.count_rows(self.filter) def from_batches(self, **kwargs): @@ -83,6 +86,8 @@ class PyarrowScannerAdapter(pa.dataset.Scanner): return self.to_reader().read_all() def to_reader(self, *, limit: Optional[int] = None): + if isinstance(self.filter, pa.compute.Expression): + return self._to_reader_with_expression(limit=limit) query = self.table.search() # Disable the builtin limit if limit is None: @@ -98,13 +103,30 @@ class PyarrowScannerAdapter(pa.dataset.Scanner): query = query.where(self.filter, prefilter=True) return query.to_batches(batch_size=self.batch_size) + def _to_reader_with_expression(self, *, limit: Optional[int] = None): + # LanceDB filters are SQL strings, but engines such as DuckDB push down + # pyarrow compute expressions. Scan the table and apply the filter and + # projection with pyarrow, since the filter may reference columns that + # are not part of the projection. + if limit is not None and limit <= 0: + raise ValueError("limit must be positive") + query = self.table.search().limit(self.table.count_rows()) + scanner = pa.dataset.Scanner.from_batches( + query.to_batches(batch_size=self.batch_size), + columns=self.columns, + filter=self.filter, + ) + if limit is None: + return scanner.to_reader() + return scanner.head(limit).to_reader() + class PyarrowDatasetAdapter(pa.dataset.Dataset): def __init__(self, table: Table): self.table = table def count_rows(self, filter: Optional[Filter] = None): - return self.table.count_rows(filter) + return self.scanner(filter=filter).count_rows() def get_fragments(self, filter: Optional[Filter] = None): raise NotImplementedError diff --git a/python/python/tests/test_duckdb.py b/python/python/tests/test_duckdb.py index a307149ea..a24933210 100644 --- a/python/python/tests/test_duckdb.py +++ b/python/python/tests/test_duckdb.py @@ -22,3 +22,19 @@ def test_basic_query(tmp_path): results = duck_conn.sql("SELECT SUM(y) FROM adapter").fetchall() assert results[0][0] == 26 + + +def test_query_with_filter(tmp_path): + data = pa.table({"x": [1, 2, 3, 4], "y": [5, 6, 7, 8]}) + conn = lancedb.connect(tmp_path) + tbl = conn.create_table("test", data) + + adapter = PyarrowDatasetAdapter(tbl) # noqa: F841 + + duck_conn = duckdb.connect() + + results = duck_conn.sql("SELECT SUM(y) FROM adapter WHERE x > 2").fetchall() + assert results[0][0] == 15 + + results = duck_conn.sql("SELECT y FROM adapter WHERE x > 2 ORDER BY y").fetchall() + assert results == [(7,), (8,)] diff --git a/python/python/tests/test_pyarrow.py b/python/python/tests/test_pyarrow.py index 7d99d19c6..272eee2ea 100644 --- a/python/python/tests/test_pyarrow.py +++ b/python/python/tests/test_pyarrow.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright The LanceDB Authors import pyarrow as pa +import pyarrow.compute as pc import lancedb from lancedb.integrations.pyarrow import PyarrowDatasetAdapter @@ -48,3 +49,19 @@ def test_dataset_adapter(tmp_path): assert adapter.head(10).num_rows == 0 assert adapter.scanner().projected_schema == pa.schema({"x": pa.int64()}) + + +def test_dataset_adapter_expression_filter(tmp_path): + data = pa.table({"x": [1, 2, 3, 4], "y": [5, 6, 7, 8]}) + conn = lancedb.connect(tmp_path) + tbl = conn.create_table("test", data) + + adapter = PyarrowDatasetAdapter(tbl) + filt = pc.field("x") > 2 + + assert adapter.count_rows(filt) == 2 + assert adapter.to_table(filter=filt) == data.slice(2, 2) + assert adapter.head(1, filter=filt) == data.slice(2, 1) + # The filter can reference a column that is not projected + assert adapter.to_table(columns=["y"], filter=filt) == pa.table({"y": [7, 8]}) + assert adapter.scanner(filter=filt).count_rows() == 2