From 566517be05ff39d3fc6445c5031abccceecf11eb Mon Sep 17 00:00:00 2001 From: Raphael Fakhri Date: Thu, 22 Oct 2026 22:24:23 +0300 Subject: [PATCH] fix(python): apply pyarrow expression filters in PyarrowDatasetAdapter (#4592) `PyarrowDatasetAdapter` lets DuckDB query a LanceDB table, but DuckDB pushes `WHERE` clauses into `Dataset.scanner(filter=...)` as `pyarrow.compute.Expression` objects. The adapter forwarded them to `Table.count_rows` and `LanceQueryBuilder.where`, which accept only SQL strings, so every filtered DuckDB query failed with `TypeError: argument 'filter': 'Expression' object is not an instance of 'str'`. This change: - Applies a `pyarrow.compute.Expression` filter locally with `pyarrow.dataset.Scanner.from_batches`. The scan reads all columns and pyarrow applies the filter and projection, because DuckDB can filter on a column it does not project. - Applies `head` limits after the filter, so the limit counts matching rows. - Routes `PyarrowDatasetAdapter.count_rows(filter)` through the scanner, so it accepts both filter types. SQL string filters keep the existing path and still run inside LanceDB. Fixes #4591 ### Testing - New `test_dataset_adapter_expression_filter` in `test_pyarrow.py` covers `count_rows`, `to_table`, `head` and a projection that excludes the filtered column. - New `test_query_with_filter` in `test_duckdb.py` runs DuckDB queries with a `WHERE` clause. - Both tests fail without the fix and pass with it. `ruff check` and `ruff format --check` pass. --- python/python/lancedb/integrations/pyarrow.py | 24 ++++++++++++++++++- python/python/tests/test_duckdb.py | 16 +++++++++++++ python/python/tests/test_pyarrow.py | 17 +++++++++++++ 3 files changed, 56 insertions(+), 1 deletion(-) 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