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:
Joaquin Hui Gomez
2026-10-24 05:10:07 +08:00
committed by GitHub
co-authored by Xuanwo
parent 3321860e9f
commit 9da4658d57
4 changed files with 102 additions and 18 deletions
+27
View File
@@ -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