mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-15 07:06:52 +00:00
fix(python): accept with_row_id in take_offsets/take_row_ids on every Table implementation (#4374)
`Table.take_offsets` and `Table.take_row_ids` declare a keyword-only `with_row_id: bool = False`, but none of the three implementations accept it — `LanceTable`, `AsyncTable` and `RemoteTable` all raise `TypeError: ... got an unexpected keyword argument 'with_row_id'` on the documented call. The builders they return already support it (`LanceTakeQueryBuilder.with_row_id()` / `AsyncTakeQuery.with_row_id()`), so this is a signature that never reached the implementations. The base docstrings also don't mention the parameter and still say `take_offsets` returns a `pa.RecordBatch` (it returns a `LanceTakeQueryBuilder`), so those are corrected too. Each implementation now accepts the kwarg and applies `.with_row_id()` to the builder it returns. Tests: sync and async assert `_rowid` appears in the result columns; the remote test uses the suite's request-capturing fixture and asserts `with_row_id: true` reaches the wire (and is absent by default). All three fail on `main` with the `TypeError`. ruff format/check clean. A related mismatch I left alone: the deprecated `LanceTable.create_scalar_index`/`create_fts_index` also drop a `wait_timeout` the base declares — deprecated since 0.25, so not worth touching here. Co-authored-by: Xuanwo <github@xuanwo.io>
This commit is contained in:
co-authored by
Xuanwo
parent
3321860e9f
commit
9da4658d57
@@ -2324,6 +2324,33 @@ def test_take_queries(tmp_path):
|
||||
]
|
||||
|
||||
|
||||
def test_take_queries_with_row_id(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("test", pa.table({"idx": range(100)}))
|
||||
|
||||
by_row_id = table.take_row_ids([5, 2, 17], with_row_id=True).to_arrow()
|
||||
assert "_rowid" in by_row_id.column_names
|
||||
by_offset = table.take_offsets([5, 2, 17], with_row_id=True).to_arrow()
|
||||
assert "_rowid" in by_offset.column_names
|
||||
|
||||
# The column is opt-in and must stay out of the default projection.
|
||||
assert "_rowid" not in table.take_row_ids([5, 2, 17]).to_arrow().column_names
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_take_queries_with_row_id_async(tmp_path):
|
||||
db = await lancedb.connect_async(tmp_path)
|
||||
table = await db.create_table("test", pa.table({"idx": range(100)}))
|
||||
|
||||
by_row_id = await table.take_row_ids([5, 2, 17], with_row_id=True).to_arrow()
|
||||
assert "_rowid" in by_row_id.column_names
|
||||
by_offset = await table.take_offsets([5, 2, 17], with_row_id=True).to_arrow()
|
||||
assert "_rowid" in by_offset.column_names
|
||||
|
||||
default = await table.take_row_ids([5, 2, 17]).to_arrow()
|
||||
assert "_rowid" not in default.column_names
|
||||
|
||||
|
||||
def test_take_queries_to_batches(tmp_path):
|
||||
# Regression test for the sync take-query path: `to_batches` previously
|
||||
# raised ``AttributeError: 'AsyncTakeQuery' object has no attribute
|
||||
|
||||
Reference in New Issue
Block a user