From 9804f26bc3fc88788bfcf0249d0bc34e3108f677 Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Sat, 31 Oct 2026 10:00:10 -0700 Subject: [PATCH] fix: support batch size options for column refreshes (#4777) 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 --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> --- python/python/lancedb/_lancedb.pyi | 8 ++- python/python/lancedb/remote/table.py | 12 ++-- python/python/lancedb/table.py | 87 +++++++++++++++++++-------- python/python/tests/test_remote_db.py | 31 +++++++--- python/python/tests/test_table.py | 41 ++++++++++--- python/src/table.rs | 32 ++++++++-- rust/lancedb/src/remote/table.rs | 75 ++++++++++++++++++++++- rust/lancedb/src/table.rs | 76 ++++++++++++++++++++++- rust/lancedb/src/table/refresh.rs | 28 +++++++++ 9 files changed, 338 insertions(+), 52 deletions(-) diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 656a2be5c..9dfed229d 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -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, diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 318c798c6..67c58936f 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -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, diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index d60e89ad6..1a10f43b3 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -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, ) diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 77e76684c..dee1a7323 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -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 diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 88f5f834c..221b4eef2 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -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 diff --git a/python/src/table.rs b/python/src/table.rs index 440e494b7..83e18ef72 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -30,7 +30,7 @@ use lancedb::index::scalar::FtsIndexBuilder; use lancedb::table::{ AddDataMode, ColumnAlteration, CompactionMode, CompactionOptions, Duration, FieldMetadataUpdate, FtsToken as LanceDbFtsToken, IndexRemapMode, NewColumnTransform, - OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable, + OptimizeAction, OptimizeOptions, Ref, RefreshColumnOptions, Table as LanceDbTable, }; use lancedb::tokenize as lancedb_tokenize; use pyo3::{ @@ -42,6 +42,15 @@ use pyo3::{ mod scannable; +fn refresh_column_options(batch_size: Option) -> PyResult { + let mut options = RefreshColumnOptions::default(); + options.batch_size = batch_size + .map(usize::try_from) + .transpose() + .map_err(|_| PyValueError::new_err("batch_size must be greater than zero"))?; + Ok(options) +} + /// Convert `RTEStats` to a Python dict, preserving the per-table-shard list. /// /// Deliberately not flattened to a table-level summary: a table is N @@ -2062,21 +2071,36 @@ impl Table { }) } - pub fn refresh_column(self_: PyRef<'_, Self>, column: String) -> PyResult> { + #[pyo3(signature = (column, *, batch_size=None))] + pub fn refresh_column( + self_: PyRef<'_, Self>, + column: String, + batch_size: Option, + ) -> PyResult> { + let options = refresh_column_options(batch_size)?; let inner = self_.inner_ref()?.clone(); future_into_py(self_.py(), async move { - let result = inner.refresh_column(column).await.infer_error()?; + let result = inner + .refresh_column_with_options(column, options) + .await + .infer_error()?; Ok(RefreshColumnResult::from(result)) }) } + #[pyo3(signature = (column, *, batch_size=None))] pub fn refresh_column_async( self_: PyRef<'_, Self>, column: String, + batch_size: Option, ) -> PyResult> { + let options = refresh_column_options(batch_size)?; let inner = self_.inner_ref()?.clone(); future_into_py(self_.py(), async move { - let job = inner.refresh_column_async(column).await.infer_error()?; + let job = inner + .refresh_column_async_with_options(column, options) + .await + .infer_error()?; Ok(crate::job::Job::new_typed(job)) }) } diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 9ffb88b30..fad790391 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -3779,8 +3779,21 @@ impl BaseTable for RemoteTable { &self, column: &str, ) -> Result> { + self.refresh_column_async_with_options(column, Default::default()) + .await + } + + async fn refresh_column_async_with_options( + &self, + column: &str, + options: crate::table::RefreshColumnOptions, + ) -> Result> { + options.validate()?; self.check_mutable().await?; let mut body = serde_json::json!({ "column": column }); + if let Some(batch_size) = options.batch_size { + body["batch_size"] = serde_json::json!(batch_size); + } self.apply_branch_body(&mut body); let request = self .client @@ -10495,7 +10508,7 @@ mod tests { assert_eq!(request.url().path(), "/v1/table/my_table/backfill_column"); let body = request.body().unwrap().as_bytes().unwrap(); let value: serde_json::Value = serde_json::from_slice(body).unwrap(); - assert_eq!(value["column"], "doubled"); + assert_eq!(value, serde_json::json!({"column": "doubled"})); http::Response::builder() .status(202) .body(r#"{"job_id": "j-42"}"#) @@ -10513,6 +10526,66 @@ mod tests { ); } + #[rstest] + #[case(None)] + #[case(Some(1))] + #[case(Some(8))] + #[tokio::test] + async fn test_refresh_column_async_batch_size(#[case] batch_size: Option) { + use crate::table::{Ref, RefreshColumnOptions}; + + let table = + Table::new_with_handler("my_table", move |request| match request.url().path() { + "/v1/table/my_table/branches/create/" => http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap(), + "/v1/table/my_table/describe/" => simple_describe_response(), + "/v1/table/my_table/backfill_column" => { + let mut expected = serde_json::json!({"column": "quality", "branch": "exp"}); + if let Some(batch_size) = batch_size { + expected["batch_size"] = serde_json::json!(batch_size); + } + assert_eq!(request_body_json(&request), expected); + http::Response::builder() + .status(202) + .body(r#"{"job_id": "j-42"}"#.to_string()) + .unwrap() + } + path => panic!("unexpected request: {path}"), + }); + let branch = table + .create_branch("exp", Ref::Version(None, None)) + .await + .unwrap(); + let options = RefreshColumnOptions { batch_size }; + let job = branch + .refresh_column_async_with_options("quality", options) + .await + .unwrap(); + assert_eq!(job.id(), Some("j-42")); + } + + #[tokio::test] + async fn test_refresh_column_rejects_zero_batch_size_before_request() { + use crate::table::RefreshColumnOptions; + + let table = Table::new_with_handler("my_table", |_| -> http::Response { + panic!("invalid batch size must not send a request") + }); + let options = RefreshColumnOptions::default().batch_size(0); + let error = table + .refresh_column_with_options("quality", options.clone()) + .await + .unwrap_err(); + assert!(matches!(error, Error::InvalidInput { message } if message.contains("batch_size"))); + let error = table + .refresh_column_async_with_options("quality", options) + .await + .unwrap_err(); + assert!(matches!(error, Error::InvalidInput { message } if message.contains("batch_size"))); + } + #[rstest] #[case(false, "Function is unavailable")] #[case(true, "Function name was not found")] diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 180efb410..ca1ee32e4 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -111,7 +111,7 @@ pub use lance_index::optimize::OptimizeOptions; pub use optimize::{ CompactionMode, CompactionOptions, IndexRemapMode, OptimizeAction, OptimizeStats, }; -pub use refresh::RefreshColumnResult; +pub use refresh::{RefreshColumnOptions, RefreshColumnResult}; pub use rte_stats::{MemtableStats, RTEStats, SsTableStats, TableShardStats}; pub use schema_evolution::{ AddColumnsResult, AlterColumnsResult, DropColumnsResult, FieldMetadataUpdate, @@ -879,6 +879,16 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { message: "this operation is only supported on LanceDB Embedded".into(), }) } + /// Refresh with options; expression-backed local refreshes ignore the + /// Function batch size. + async fn refresh_column_with_options( + &self, + column: &str, + options: RefreshColumnOptions, + ) -> Result { + options.validate()?; + self.refresh_column(column).await + } /// Fill a computed column's unfilled rows and recompute those whose /// inputs changed, returning a [`Job`] tracking the operation. async fn refresh_column_async( @@ -889,6 +899,15 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { message: "this operation is only supported on LanceDB Embedded".into(), }) } + /// Start a refresh job with options. + async fn refresh_column_async_with_options( + &self, + column: &str, + options: RefreshColumnOptions, + ) -> Result> { + options.validate()?; + self.refresh_column_async(column).await + } /// The per-row errors Function refreshes recorded on this table; see /// [`Table::function_errors`]. The default returns `NotSupported`. async fn function_errors( @@ -1936,6 +1955,33 @@ impl Table { self.inner.refresh_column(column.as_ref()).await } + /// Like [`Table::refresh_column`], with explicit refresh options. + /// + /// [`RefreshColumnOptions::batch_size()`] caps rows per Function call in + /// LanceDB Distributed and must be greater than zero when set. It has no + /// effect on expression-backed columns in LanceDB Embedded. Remote refreshes + /// require [`Table::refresh_column_async_with_options`]. + /// + /// ``` + /// # use lancedb::{Table, table::RefreshColumnOptions}; + /// # async fn refresh(table: &Table) -> Result<(), Box> { + /// let result = table.refresh_column_with_options( + /// "doubled", RefreshColumnOptions::default(), + /// ).await?; + /// println!("filled {} rows", result.rows_filled); + /// # Ok(()) + /// # } + /// ``` + pub async fn refresh_column_with_options( + &self, + column: impl AsRef, + options: RefreshColumnOptions, + ) -> Result { + self.inner + .refresh_column_with_options(column.as_ref(), options) + .await + } + /// Like [`Table::refresh_column`], but returns a [`Job`] tracking the /// operation instead of blocking until it completes. /// @@ -1965,6 +2011,34 @@ impl Table { self.inner.refresh_column_async(column.as_ref()).await } + /// Like [`Table::refresh_column_async`], with explicit refresh options. + /// + /// Set [`RefreshColumnOptions::batch_size()`] to cap the rows each remote + /// Function call receives, for example to keep slow Functions within the + /// executor's per-call deadline. An unset batch size uses the server default; + /// zero is rejected before submitting a job. Embedded expression refreshes + /// ignore the batch size. + /// + /// ``` + /// # use lancedb::{Table, table::RefreshColumnOptions}; + /// # async fn refresh(table: &Table) -> Result<(), Box> { + /// let job = table.refresh_column_async_with_options( + /// "quality", RefreshColumnOptions::default().batch_size(8), + /// ).await?; + /// let result = job.wait().await?; + /// # Ok(()) + /// # } + /// ``` + pub async fn refresh_column_async_with_options( + &self, + column: impl AsRef, + options: RefreshColumnOptions, + ) -> Result> { + self.inner + .refresh_column_async_with_options(column.as_ref(), options) + .await + } + /// The per-row errors Function refreshes recorded on this table: the /// rows a refresh skipped under its skip policy, with the failing input /// and the error, plus a summary for any fragment whose detail was diff --git a/rust/lancedb/src/table/refresh.rs b/rust/lancedb/src/table/refresh.rs index 19accc2c2..7000be879 100644 --- a/rust/lancedb/src/table/refresh.rs +++ b/rust/lancedb/src/table/refresh.rs @@ -54,6 +54,34 @@ use super::{BaseTable, NativeTable}; use crate::job::Job; use crate::{Error, Result}; +/// Options for refreshing a computed column. +#[derive(Debug, Clone, Default)] +#[non_exhaustive] +pub struct RefreshColumnOptions { + /// Maximum rows per Function call in LanceDB Distributed. + /// + /// Must be greater than zero. `None` uses the server default. This has no + /// effect on expression-backed columns in LanceDB Embedded. + pub batch_size: Option, +} + +impl RefreshColumnOptions { + /// Limit each Function call to at most `batch_size` rows. + pub fn batch_size(mut self, batch_size: usize) -> Self { + self.batch_size = Some(batch_size); + self + } + + pub(crate) fn validate(&self) -> Result<()> { + if self.batch_size == Some(0) { + return Err(Error::InvalidInput { + message: "batch_size must be greater than zero".into(), + }); + } + Ok(()) + } +} + /// The result of refreshing a computed column. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)] pub struct RefreshColumnResult {