rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-11-13 14:16:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: support batch size options for column refreshes (#4777)

Browse Source 
Python column refresh methods rejected `batch_size`, and the Rust remote
backend never included it in the backfill request. Callers could not
reduce the rows passed to a slow Function to stay within the executor's
per-call deadline.

Add optional, keyword-only `batch_size` arguments throughout the
synchronous and asynchronous Python APIs and bindings. Add compatible
Rust `RefreshColumnOptions` APIs while preserving existing Rust method
signatures. Reject non-positive sizes before submitting a job, send
positive sizes in `backfill_column`, and omit the field when unset while
retaining branch scoping.

Batch size controls remote Function calls. Embedded expression refreshes
accept and validate the option without changing execution, and blocking
remote refreshes retain their existing unsupported behavior; remote
callers use `refresh_column_async` and wait on the returned job.

Validation:

- Reproduced the original Python `TypeError` from the committed method.
- Five focused Rust tests passed, covering request contents, omission,
branch scoping, and rejection before any request.
- `cargo check --quiet --features remote --tests --examples` passed.
- `cargo clippy --quiet --features remote --tests --examples -- -D
warnings` passed.
- Nineteen targeted Python tests passed against the rebuilt extension,
including synchronous and asynchronous remote requests, invalid sizes,
and default-build behavior. Eight embedded refresh cases require the
experimental-functions build and were skipped.
- Rust formatting and Python formatting/lint checks passed, including
CI's pinned Ruff version. Repository-wide Ruff reports 48 existing
findings in untouched documentation files; none are in changed files.

Fixes #4776

<!-- lance-gatekeeper-fix:v1 agent=4906ff3f5f7797a1d49a13803302021c
generation=1 -->

---------

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
This commit is contained in:
lancedb-gatefixer[bot]
2026-10-31 10:00:10 -07:00
committed by GitHub
parent 569952d150
commit 9804f26bc3
9 changed files with 338 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/python/lancedb/_lancedb.pyi
+6 -2
View File
@@ -557,8 +557,12 @@ class Table:
auto_refresh_enabled: Optional[bool] = None,
auto_refresh_interval_seconds: Optional[int] = None,
) -> AddColumnsResult: ...
async def refresh_column(self, column: str) -> RefreshColumnResult: ...
async def refresh_column_async(self, column: str) -> Job: ...
async def refresh_column(
self, column: str, *, batch_size: Optional[int] = None
) -> RefreshColumnResult: ...
async def refresh_column_async(
self, column: str, *, batch_size: Optional[int] = None
) -> Job: ...
async def function_errors(
self,
job_id: Optional[str] = None,
python/python/lancedb/remote/table.py
+8 -4
View File
@@ -1108,11 +1108,15 @@ class RemoteTable(Table):
)
)
def refresh_column(self, column: str):
return LOOP.run(self._table.refresh_column(column))
def refresh_column(self, column: str, *, batch_size: Optional[int] = None):
return LOOP.run(self._table.refresh_column(column, batch_size=batch_size))
def refresh_column_async(self, column: str) -> Job[RefreshColumnResult]:
return Job(LOOP.run(self._table.refresh_column_async(column)))
def refresh_column_async(
self, column: str, *, batch_size: Optional[int] = None
) -> Job[RefreshColumnResult]:
return Job(
LOOP.run(self._table.refresh_column_async(column, batch_size=batch_size))
)
def function_errors(
self,
python/python/lancedb/table.py
+62 -25
View File
@@ -3075,7 +3075,9 @@ class Table(ABC):
"""
@abstractmethod
def refresh_column(self, column: str) -> "RefreshColumnResult":
def refresh_column(
self, column: str, *, batch_size: Optional[int] = None
) -> "RefreshColumnResult":
"""
Fill the rows of a computed column that hold no value yet.
@@ -3087,14 +3089,18 @@ class Table(ABC):
by the next one, and rows whose inputs changed since they were computed
are recomputed; everything else is left as it is.
LanceDB Embedded only. In LanceDB Distributed, a refresh runs as a
server job, through
Supported for local tables only. For remote tables, a refresh runs as
a server job, through
[`refresh_column_async`][lancedb.table.Table.refresh_column_async].
Parameters
----------
column: str
The name of the computed column to fill.
batch_size: int, optional
Maximum rows per Function call on the remote server. Must be
greater than zero. None uses the server default. Has no effect on
local expression-backed columns.
Returns
-------
@@ -3104,20 +3110,31 @@ class Table(ABC):
"""
@abstractmethod
def refresh_column_async(self, column: str) -> Job[RefreshColumnJobResult]:
def refresh_column_async(
self, column: str, *, batch_size: Optional[int] = None
) -> Job[RefreshColumnJobResult]:
"""
Like :meth:`refresh_column`, but returns a handle to the refresh job
instead of blocking until it completes.
LanceDB Embedded refresh requires a source build with
``experimental-functions``.
LanceDB Distributed jobs do not require this feature.
Local refresh requires a source build with ``experimental-functions``.
Remote jobs do not require this feature.
The job may already be complete when returned; callers must not assume
the column is filled until :meth:`Job.wait` returns. Invalid input --
an unknown column, or one that is not computed -- raises here rather
than failing the job. In LanceDB Embedded the job runs in-process; in
LanceDB Distributed it is the server's backfill job.
than failing the job. For local tables, the job runs in-process; for
remote tables, it is the server's backfill job.
Parameters
----------
column: str
The name of the computed column to fill.
batch_size: int, optional
Maximum rows per Function call on the remote server. Must be
greater than zero. None uses the server default. Lower this for
slow Functions to stay within the executor's per-call deadline.
Has no effect on local expression-backed columns.
Returns
-------
@@ -5156,18 +5173,24 @@ class LanceTable(Table):
)
)
def refresh_column(self, column: str) -> "RefreshColumnResult":
def refresh_column(
self, column: str, *, batch_size: Optional[int] = None
) -> "RefreshColumnResult":
"""Fill a computed column's unfilled rows and recompute those whose
inputs changed. See
[`AsyncTable.refresh_column`][lancedb.AsyncTable.refresh_column]."""
return LOOP.run(self._table.refresh_column(column))
return LOOP.run(self._table.refresh_column(column, batch_size=batch_size))
def refresh_column_async(self, column: str) -> Job[RefreshColumnJobResult]:
def refresh_column_async(
self, column: str, *, batch_size: Optional[int] = None
) -> Job[RefreshColumnJobResult]:
"""Fill a computed column's unfilled rows and recompute those whose
inputs changed, returning a handle to the refresh job. See
[`Table.refresh_column_async`][lancedb.table.Table.refresh_column_async].
"""
return Job(LOOP.run(self._table.refresh_column_async(column)))
return Job(
LOOP.run(self._table.refresh_column_async(column, batch_size=batch_size))
)
def function_errors(
self,
@@ -7422,7 +7445,9 @@ class AsyncTable:
else:
return await self._inner.add_columns(list(transforms.items()))
async def refresh_column(self, column: str) -> RefreshColumnResult:
async def refresh_column(
self, column: str, *, batch_size: Optional[int] = None
) -> RefreshColumnResult:
"""
Fill the rows of a computed column that hold no value yet.
@@ -7434,21 +7459,25 @@ class AsyncTable:
by the next one, and rows whose inputs changed since they were computed
are recomputed; everything else is left as it is.
LanceDB Embedded only. In LanceDB Distributed, a refresh runs as a
server job, through
Supported for local tables only. For remote tables, a refresh runs as
a server job, through
[`refresh_column_async`][lancedb.table.Table.refresh_column_async].
Parameters
----------
column: str
The name of the computed column to fill.
batch_size: int, optional
Maximum rows per Function call on the remote server. Must be
greater than zero. None uses the server default. Has no effect on
local expression-backed columns.
Returns
-------
RefreshColumnResult
The number of rows filled and the new version of the table.
"""
return await self._inner.refresh_column(column)
return await self._inner.refresh_column(column, batch_size=batch_size)
async def function_errors(
self,
@@ -7494,22 +7523,30 @@ class AsyncTable:
)
async def refresh_column_async(
self, column: str
self, column: str, *, batch_size: Optional[int] = None
) -> AsyncJob[RefreshColumnJobResult]:
"""
Like :meth:`refresh_column`, but returns a handle to the refresh job
instead of blocking until it completes.
LanceDB Embedded refresh requires a source build with
``experimental-functions``.
LanceDB Distributed jobs do not require this feature.
Local refresh requires a source build with ``experimental-functions``.
Remote jobs do not require this feature.
The job may already be complete when returned; callers must not assume
the column is filled until :meth:`AsyncJob.wait` resolves. Invalid
input -- an unknown column, or one that is not computed -- raises here
rather than failing the job. In LanceDB Embedded the job runs
in-process; in LanceDB Distributed it is the server's
backfill job.
rather than failing the job. For local tables, the job runs in-process;
for remote tables, it is the server's backfill job.
Parameters
----------
column: str
The name of the computed column to fill.
batch_size: int, optional
Maximum rows per Function call on the remote server. Must be
greater than zero. None uses the server default. Lower this for
slow Functions to stay within the executor's per-call deadline.
Has no effect on local expression-backed columns.
Returns
-------
@@ -7533,7 +7570,7 @@ class AsyncTable:
'finished'
"""
return _typed_job(
await self._inner.refresh_column_async(column),
await self._inner.refresh_column_async(column, batch_size=batch_size),
RefreshColumnJobResult.from_json,
)
python/python/tests/test_remote_db.py
+24 -7
View File
@@ -1520,7 +1520,14 @@ def test_remote_create_index_async_returns_job():
job.cancel()
def test_remote_refresh_async_returns_typed_terminal_result():
@pytest.mark.asyncio
@pytest.mark.parametrize("async_table", [False, True])
@pytest.mark.parametrize(
"refresh_kwargs", [{}, {"batch_size": None}, {"batch_size": 1}, {"batch_size": 8}]
)
async def test_remote_refresh_async_returns_typed_terminal_result(
async_table, refresh_kwargs
):
terminal_result = {
"rows_assigned": 12,
"rows_failed": 0,
@@ -1533,7 +1540,10 @@ def test_remote_refresh_async_returns_typed_terminal_result():
content_len = int(request.headers.get("Content-Length", 0))
body = request.rfile.read(content_len) if content_len > 0 else b""
if request.path == "/v1/table/test/backfill_column":
assert json.loads(body)["column"] == "derived"
expected = {"column": "derived"}
if refresh_kwargs.get("batch_size") is not None:
expected["batch_size"] = refresh_kwargs["batch_size"]
assert json.loads(body) == expected
request.send_response(202)
request.send_header("Content-Type", "application/json")
request.end_headers()
@@ -1582,11 +1592,18 @@ def test_remote_refresh_async_returns_typed_terminal_result():
request.send_response(404)
request.end_headers()
with mock_lancedb_connection(handler) as db:
table = db.create_table("test", [{"id": 1}])
job = table.refresh_column_async("derived")
assert job.id == "refresh-1"
result = job.wait(timeout=timedelta(seconds=30))
if async_table:
async with mock_lancedb_connection_async(handler) as db:
table = await db.create_table("test", [{"id": 1}])
job = await table.refresh_column_async("derived", **refresh_kwargs)
assert job.id == "refresh-1"
result = await job.wait(timeout=timedelta(seconds=30))
else:
with mock_lancedb_connection(handler) as db:
table = db.create_table("test", [{"id": 1}])
job = table.refresh_column_async("derived", **refresh_kwargs)
assert job.id == "refresh-1"
result = job.wait(timeout=timedelta(seconds=30))
assert isinstance(result, lancedb.RefreshColumnResult)
assert result.model_dump() == terminal_result
python/python/tests/test_table.py
+33 -8
View File
@@ -5551,14 +5551,15 @@ async def test_async_search_runs_embedding_on_dedicated_executor(
@pytest.mark.experimental_functions
def test_computed_column_declare_and_refresh(tmp_path):
@pytest.mark.parametrize("batch_size", [None, 8])
def test_computed_column_declare_and_refresh(tmp_path, batch_size):
db = lancedb.connect(tmp_path)
table = db.create_table("computed", [{"x": 1}, {"x": 2}])
table.add_columns(computed={"doubled": "x * 2"})
assert table.to_arrow()["doubled"].to_pylist() == [None, None]
result = table.refresh_column("doubled")
result = table.refresh_column("doubled", batch_size=batch_size)
assert result.rows_filled == 2
assert sorted(table.to_arrow()["doubled"].to_pylist()) == [2, 4]
@@ -5600,16 +5601,38 @@ def test_computed_column_blob_projection_inherits_semantics(tmp_path):
@pytest.mark.experimental_functions
@pytest.mark.asyncio
async def test_computed_column_async(tmp_path):
@pytest.mark.parametrize("batch_size", [None, 8])
async def test_computed_column_async(tmp_path, batch_size):
db = await lancedb.connect_async(tmp_path)
table = await db.create_table("computed_async", [{"x": 3}])
await table.add_columns(computed={"tripled": "x * 3"})
await table.refresh_column("tripled")
await table.refresh_column("tripled", batch_size=batch_size)
assert (await table.to_arrow())["tripled"].to_pylist() == [9]
@pytest.mark.parametrize("method", ["refresh_column", "refresh_column_async"])
@pytest.mark.parametrize("batch_size", [0, -1])
def test_refresh_column_rejects_invalid_batch_size(tmp_path, method, batch_size):
db = lancedb.connect(tmp_path)
table = db.create_table("invalid_refresh_batch", [{"x": 1}])
with pytest.raises(ValueError, match="batch_size must be greater than zero"):
getattr(table, method)("x", batch_size=batch_size)
@pytest.mark.asyncio
@pytest.mark.parametrize("method", ["refresh_column", "refresh_column_async"])
@pytest.mark.parametrize("batch_size", [0, -1])
async def test_async_refresh_column_rejects_invalid_batch_size(
tmp_path, method, batch_size
):
db = await lancedb.connect_async(tmp_path)
table = await db.create_table("invalid_refresh_batch", [{"x": 1}])
with pytest.raises(ValueError, match="batch_size must be greater than zero"):
await getattr(table, method)("x", batch_size=batch_size)
def test_function_errors_are_remote_only(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("t", [{"x": 1}])
@@ -5620,12 +5643,13 @@ def test_function_errors_are_remote_only(tmp_path):
@pytest.mark.experimental_functions
def test_refresh_column_async_returns_job(tmp_path):
@pytest.mark.parametrize("batch_size", [None, 8])
def test_refresh_column_async_returns_job(tmp_path, batch_size):
db = lancedb.connect(tmp_path)
table = db.create_table("computed_job", [{"x": 1}, {"x": 2}])
table.add_columns(computed={"doubled": "x * 2"})
job = table.refresh_column_async("doubled")
job = table.refresh_column_async("doubled", batch_size=batch_size)
assert job.id is None # in-process jobs have no server id
result = job.wait()
assert isinstance(result, lancedb.RefreshColumnResult)
@@ -5650,12 +5674,13 @@ def test_refresh_column_async_returns_job(tmp_path):
@pytest.mark.experimental_functions
@pytest.mark.asyncio
async def test_refresh_column_async_job_async_table(tmp_path):
@pytest.mark.parametrize("batch_size", [None, 8])
async def test_refresh_column_async_job_async_table(tmp_path, batch_size):
db = await lancedb.connect_async(tmp_path)
table = await db.create_table("computed_job_async", [{"x": 3}])
await table.add_columns(computed={"tripled": "x * 3"})
job = await table.refresh_column_async("tripled")
job = await table.refresh_column_async("tripled", batch_size=batch_size)
result = await job.wait()
assert isinstance(result, lancedb.RefreshColumnResult)
assert result.rows_assigned == 1