mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 14:16:54 +00:00
Raphael Fakhri
566517be05
`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.
41 lines
1.2 KiB
Python
41 lines
1.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
|
|
|
import duckdb
|
|
import pyarrow as pa
|
|
|
|
import lancedb
|
|
from lancedb.integrations.pyarrow import PyarrowDatasetAdapter
|
|
|
|
|
|
def test_basic_query(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(x) FROM adapter").fetchall()
|
|
assert results[0][0] == 10
|
|
|
|
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,)]
|