feat(refresh): priority as a per-refresh knob; fix batch_size on RemoteTable

Thread priority (Kueue tier) through refresh_column at every layer (Python sync+async
+ RemoteTable -> pyo3 -> Rust client trait/public/remote -> REST body), mirroring
num_workers/batch_size. The function keeps its priority as a default; the per-refresh
value overrides. Also adds the previously-missed batch_size to RemoteTable.refresh_column
(the REST sync path). cargo check (lancedb --features remote --tests, lancedb-python) +
ruff clean.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Wyatt Alt
2026-06-14 09:18:23 -07:00
parent 78884d2755
commit cef840463a
5 changed files with 58 additions and 17 deletions
+9 -1
View File
@@ -934,6 +934,8 @@ class RemoteTable(Table):
where: Optional[str] = None,
num_workers: Optional[int] = None,
max_workers: Optional[int] = None,
batch_size: Optional[int] = None,
priority: Optional[str] = None,
) -> str:
"""Trigger recompute of computed columns (REFRESH COLUMN).
@@ -941,6 +943,11 @@ class RemoteTable(Table):
binding; columns bound to the same struct-returning function
refresh together. Returns the refresh job id. Server-backed
feature (LanceDB Enterprise / Cloud).
num_workers / max_workers / batch_size / priority are per-refresh
scheduling knobs (how to run THIS refresh) and override any default
the function carries. `priority` is a Kueue tier
(training | interactive | backfill).
"""
if isinstance(columns, str):
columns = [columns]
@@ -950,10 +957,11 @@ class RemoteTable(Table):
where=where,
num_workers=num_workers,
max_workers=max_workers,
batch_size=batch_size,
priority=priority,
)
)
def alter_columns(
self, *alterations: Iterable[Dict[str, str]]
) -> AlterColumnsResult:
+24 -12
View File
@@ -702,7 +702,6 @@ def _normalize_progress(progress):
return progress, False
def _computed_groups(computed):
"""Group computed columns by expression, preserving declaration order
(struct-returning functions need their columns adjacent so schema order
@@ -846,7 +845,7 @@ class Table(ABC):
import warnings
warnings.warn(
'add_computed_column is deprecated; use add_columns(computed='
"add_computed_column is deprecated; use add_columns(computed="
'{"vec": embed("data")}).',
DeprecationWarning,
stacklevel=2,
@@ -3888,7 +3887,11 @@ class LanceTable(Table):
def add_columns(
self,
transforms: Dict[str, str] | pa.field | List[pa.field] | pa.Schema | None = None,
transforms: Dict[str, str]
| pa.field
| List[pa.field]
| pa.Schema
| None = None,
*,
computed: Optional[Dict] = None,
) -> Optional[AddColumnsResult]:
@@ -3914,6 +3917,7 @@ class LanceTable(Table):
num_workers: Optional[int] = None,
max_workers: Optional[int] = None,
batch_size: Optional[int] = None,
priority: Optional[str] = None,
) -> str:
"""Trigger recompute of computed columns (REFRESH COLUMN).
@@ -3922,9 +3926,10 @@ class LanceTable(Table):
refresh together. Returns the refresh job id. Server-backed
feature (LanceDB Enterprise / Cloud).
num_workers / max_workers / batch_size are per-refresh scheduling
knobs (how to run THIS refresh) and override any default the
function carries.
num_workers / max_workers / batch_size / priority are per-refresh
scheduling knobs (how to run THIS refresh) and override any default
the function carries. `priority` is a Kueue tier
(training | interactive | backfill).
"""
if isinstance(columns, str):
columns = [columns]
@@ -3935,10 +3940,10 @@ class LanceTable(Table):
num_workers=num_workers,
max_workers=max_workers,
batch_size=batch_size,
priority=priority,
)
)
def alter_columns(
self, *alterations: Iterable[Dict[str, str]]
) -> AlterColumnsResult:
@@ -5698,13 +5703,15 @@ class AsyncTable:
num_workers: Optional[int] = None,
max_workers: Optional[int] = None,
batch_size: Optional[int] = None,
priority: Optional[str] = None,
) -> str:
"""Trigger recompute of computed columns (REFRESH COLUMN).
Returns the refresh job id. Server-backed feature.
num_workers / max_workers / batch_size are per-refresh scheduling
knobs (how to run THIS refresh); they override any default the
function carries."""
num_workers / max_workers / batch_size / priority are per-refresh
scheduling knobs (how to run THIS refresh); they override any default
the function carries. `priority` is a Kueue tier
(training | interactive | backfill)."""
if isinstance(columns, str):
columns = [columns]
return await self._inner.refresh_column(
@@ -5713,11 +5720,16 @@ class AsyncTable:
num_workers=num_workers,
max_workers=max_workers,
batch_size=batch_size,
priority=priority,
)
async def add_columns(
self,
transforms: dict[str, str] | pa.field | List[pa.field] | pa.Schema | None = None,
transforms: dict[str, str]
| pa.field
| List[pa.field]
| pa.Schema
| None = None,
*,
computed: Optional[Dict] = None,
) -> Optional[AddColumnsResult]:
@@ -5778,7 +5790,7 @@ class AsyncTable:
import warnings
warnings.warn(
'add_computed_column is deprecated; use add_columns(computed='
"add_computed_column is deprecated; use add_columns(computed="
'{"col": fn("input_col")}).',
DeprecationWarning,
stacklevel=2,
+10 -2
View File
@@ -1308,7 +1308,7 @@ impl Table {
})
}
#[pyo3(signature = (columns, where_clause=None, num_workers=None, max_workers=None, batch_size=None))]
#[pyo3(signature = (columns, where_clause=None, num_workers=None, max_workers=None, batch_size=None, priority=None))]
pub fn refresh_column(
self_: PyRef<'_, Self>,
columns: Vec<String>,
@@ -1316,11 +1316,19 @@ impl Table {
num_workers: Option<u32>,
max_workers: Option<u32>,
batch_size: Option<u32>,
priority: Option<String>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
inner
.refresh_column(&columns, where_clause, num_workers, max_workers, batch_size)
.refresh_column(
&columns,
where_clause,
num_workers,
max_workers,
batch_size,
priority,
)
.await
.infer_error()
})
+5 -1
View File
@@ -2452,6 +2452,7 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
num_workers: Option<u32>,
max_workers: Option<u32>,
batch_size: Option<u32>,
priority: Option<String>,
) -> Result<String> {
let mut body = serde_json::json!({ "columns": columns });
if let Some(w) = where_clause {
@@ -2466,6 +2467,9 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
if let Some(n) = batch_size {
body["batch_size"] = n.into();
}
if let Some(p) = priority {
body["priority"] = serde_json::Value::String(p);
}
let request = self
.client
.post(&format!("/v1/table/{}/refresh_column", self.identifier))
@@ -2992,7 +2996,7 @@ mod tests {
});
let job_id = table
.refresh_column(&["vec".to_string()], None, Some(2), None, None)
.refresh_column(&["vec".to_string()], None, Some(2), None, None, None)
.await
.unwrap();
assert_eq!(job_id, "j-9");
+10 -1
View File
@@ -688,6 +688,7 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync {
_num_workers: Option<u32>,
_max_workers: Option<u32>,
_batch_size: Option<u32>,
_priority: Option<String>,
) -> Result<String> {
Err(Error::NotSupported {
message: "refresh_column is not supported by this table".into(),
@@ -1556,9 +1557,17 @@ impl Table {
num_workers: Option<u32>,
max_workers: Option<u32>,
batch_size: Option<u32>,
priority: Option<String>,
) -> Result<String> {
self.inner
.refresh_column(columns, where_clause, num_workers, max_workers, batch_size)
.refresh_column(
columns,
where_clause,
num_workers,
max_workers,
batch_size,
priority,
)
.await
}