feat: refresh computed columns (#3938)

table.refresh_column("doubled") fills the rows of a declared column that
hold no value, in two passes per fragment: the first scans only the
unfilled
live rows to count exact gains and decide staging, the second streams
the
fragment's physical rows into a standalone column file published in one
DataReplacement -- committed under the dataset's own session -- so peak
memory is bounded by a scan batch. A row that holds a value keeps it;
deleted and already-filled rows never reach the expression, so a poison
value in them cannot fail the refresh. Refresh refuses under an LSM
write
spec, including the mem-wal catch-up flag that outlives unset and marks
retained SSTable rows.

---

<sub>Stack created with <a
href="https://github.com/github/gh-stack">GitHub Stacks CLI</a> • <a
href="https://gh.io/stacks-feedback">Give Feedback 💬</a></sub>
This commit is contained in:
Wyatt Alt
2026-08-14 14:43:41 -07:00
committed by GitHub
parent def869bb78
commit fc0d917d32
18 changed files with 1042 additions and 29 deletions
+30 -3
View File
@@ -70,9 +70,9 @@ abstract addColumns(newColumnTransforms): Promise<AddColumnsResult>
Add new columns with defined values. Add new columns with defined values.
The `{ computed }` form stores the expression rather than evaluating it The `{ computed }` form stores the expression rather than evaluating it
now: the column is committed with no values, and a later refresh fills now: the column is committed with no values, and rows get them from
the rows. Declaring one therefore costs the same on a large table as on [Table#refreshColumn](Table.md#refreshcolumn). Declaring one therefore costs the same on a
an empty one. large table as on an empty one.
A refresh does not revisit rows it has already filled, so mutating an A refresh does not revisit rows it has already filled, so mutating an
input leaves the value computed at fill time; recomputing means dropping input leaves the value computed at fill time; recomputing means dropping
@@ -108,6 +108,7 @@ containing the new version number of the table after adding the columns.
```ts ```ts
await table.addColumns({ computed: [{ name: "doubled", valueSql: "x * 2" }] }); await table.addColumns({ computed: [{ name: "doubled", valueSql: "x * 2" }] });
const { rowsFilled } = await table.refreshColumn("doubled");
``` ```
*** ***
@@ -743,6 +744,32 @@ for await (const batch of table.query()) {
*** ***
### refreshColumn()
```ts
abstract refreshColumn(column): Promise<RefreshColumnResult>
```
Fill the rows of a computed column that hold no value yet.
Rows appended since the last refresh are filled by the next one; rows
already filled are left as they are, so the call is idempotent and does
not observe a mutated input. Local tables only.
#### Parameters
* **column**: `string`
The name of the computed column to fill.
#### Returns
`Promise`&lt;[`RefreshColumnResult`](../interfaces/RefreshColumnResult.md)&gt;
A promise that resolves to the
number of rows filled and the new version number of the table.
***
### restore() ### restore()
```ts ```ts
+1
View File
@@ -105,6 +105,7 @@
- [OptimizeOptions](interfaces/OptimizeOptions.md) - [OptimizeOptions](interfaces/OptimizeOptions.md)
- [OptimizeStats](interfaces/OptimizeStats.md) - [OptimizeStats](interfaces/OptimizeStats.md)
- [QueryExecutionOptions](interfaces/QueryExecutionOptions.md) - [QueryExecutionOptions](interfaces/QueryExecutionOptions.md)
- [RefreshColumnResult](interfaces/RefreshColumnResult.md)
- [RemovalStats](interfaces/RemovalStats.md) - [RemovalStats](interfaces/RemovalStats.md)
- [RenameTableOptions](interfaces/RenameTableOptions.md) - [RenameTableOptions](interfaces/RenameTableOptions.md)
- [RestNamespaceConfig](interfaces/RestNamespaceConfig.md) - [RestNamespaceConfig](interfaces/RestNamespaceConfig.md)
@@ -0,0 +1,23 @@
[**@lancedb/lancedb**](../README.md) • **Docs**
***
[@lancedb/lancedb](../globals.md) / RefreshColumnResult
# Interface: RefreshColumnResult
## Properties
### rowsFilled
```ts
rowsFilled: number;
```
***
### version
```ts
version: number;
```
+25 -2
View File
@@ -3348,14 +3348,37 @@ describe("computed columns", () => {
}); });
afterEach(() => tmpDir.removeCallback()); afterEach(() => tmpDir.removeCallback());
it("declares a column with no values", async () => { it("declares a column and fills it on refresh", async () => {
const db = await connect(tmpDir.name); const db = await connect(tmpDir.name);
const table = await db.createTable("computed", [{ x: 1 }, { x: 2 }]); const table = await db.createTable("computed", [{ x: 1 }, { x: 2 }]);
await table.addColumns({ await table.addColumns({
computed: [{ name: "doubled", valueSql: "x * 2" }], computed: [{ name: "doubled", valueSql: "x * 2" }],
}); });
const rows = await table.query().toArray(); let rows = await table.query().toArray();
expect(rows.map((r) => r.doubled)).toEqual([null, null]); expect(rows.map((r) => r.doubled)).toEqual([null, null]);
const result = await table.refreshColumn("doubled");
expect(result.rowsFilled).toBe(2);
rows = await table.query().toArray();
expect(rows.map((r) => r.doubled).sort()).toEqual([2, 4]);
});
it("fills rows added since the last refresh", async () => {
const db = await connect(tmpDir.name);
const table = await db.createTable("computed_append", [{ x: 1 }]);
await table.addColumns({
computed: [{ name: "doubled", valueSql: "x * 2" }],
});
await table.refreshColumn("doubled");
await table.add([{ x: 5 }]);
const result = await table.refreshColumn("doubled");
expect(result.rowsFilled).toBe(1);
const rows = await table.query().toArray();
expect(rows.map((r) => r.doubled).sort()).toEqual([10, 2]);
}); });
}); });
+1
View File
@@ -50,6 +50,7 @@ export {
MergeResult, MergeResult,
AddResult, AddResult,
AddColumnsResult, AddColumnsResult,
RefreshColumnResult,
AlterColumnsResult, AlterColumnsResult,
UpdateFieldMetadataResult, UpdateFieldMetadataResult,
DeleteResult, DeleteResult,
+21 -3
View File
@@ -33,6 +33,7 @@ import {
Job, Job,
Branches as NativeBranches, Branches as NativeBranches,
OptimizeStats, OptimizeStats,
RefreshColumnResult,
TableStatistics, TableStatistics,
Tags, Tags,
UpdateFieldMetadataResult, UpdateFieldMetadataResult,
@@ -527,9 +528,9 @@ export abstract class Table {
* Add new columns with defined values. * Add new columns with defined values.
* *
* The `{ computed }` form stores the expression rather than evaluating it * The `{ computed }` form stores the expression rather than evaluating it
* now: the column is committed with no values, and a later refresh fills * now: the column is committed with no values, and rows get them from
* the rows. Declaring one therefore costs the same on a large table as on * {@link Table#refreshColumn}. Declaring one therefore costs the same on a
* an empty one. * large table as on an empty one.
* *
* A refresh does not revisit rows it has already filled, so mutating an * A refresh does not revisit rows it has already filled, so mutating an
* input leaves the value computed at fill time; recomputing means dropping * input leaves the value computed at fill time; recomputing means dropping
@@ -549,6 +550,7 @@ export abstract class Table {
* @example * @example
* ```ts * ```ts
* await table.addColumns({ computed: [{ name: "doubled", valueSql: "x * 2" }] }); * await table.addColumns({ computed: [{ name: "doubled", valueSql: "x * 2" }] });
* const { rowsFilled } = await table.refreshColumn("doubled");
* ``` * ```
*/ */
abstract addColumns( abstract addColumns(
@@ -560,6 +562,18 @@ export abstract class Table {
| { computed: AddColumnsSql[] }, | { computed: AddColumnsSql[] },
): Promise<AddColumnsResult>; ): Promise<AddColumnsResult>;
/**
* Fill the rows of a computed column that hold no value yet.
*
* Rows appended since the last refresh are filled by the next one; rows
* already filled are left as they are, so the call is idempotent and does
* not observe a mutated input. Local tables only.
* @param {string} column The name of the computed column to fill.
* @returns {Promise<RefreshColumnResult>} A promise that resolves to the
* number of rows filled and the new version number of the table.
*/
abstract refreshColumn(column: string): Promise<RefreshColumnResult>;
/** /**
* Alter the name or nullability of columns. * Alter the name or nullability of columns.
* @param {ColumnAlteration[]} columnAlterations One or more alterations to * @param {ColumnAlteration[]} columnAlterations One or more alterations to
@@ -1161,6 +1175,10 @@ export class LocalTable extends Table {
throw new Error("Invalid input type for addColumns"); throw new Error("Invalid input type for addColumns");
} }
async refreshColumn(column: string): Promise<RefreshColumnResult> {
return await this.inner.refreshColumn(column);
}
async alterColumns( async alterColumns(
columnAlterations: ColumnAlteration[], columnAlterations: ColumnAlteration[],
): Promise<AlterColumnsResult> { ): Promise<AlterColumnsResult> {
+25
View File
@@ -361,6 +361,16 @@ impl Table {
Ok(res.into()) Ok(res.into())
} }
#[napi(catch_unwind)]
pub async fn refresh_column(&self, column: String) -> napi::Result<RefreshColumnResult> {
let res = self
.inner_ref()?
.refresh_column(column)
.await
.default_error()?;
Ok(res.into())
}
#[napi(catch_unwind)] #[napi(catch_unwind)]
pub async fn add_columns_with_schema( pub async fn add_columns_with_schema(
&self, &self,
@@ -1210,6 +1220,21 @@ pub struct AddColumnsResult {
pub version: i64, pub version: i64,
} }
#[napi(object)]
pub struct RefreshColumnResult {
pub rows_filled: i64,
pub version: i64,
}
impl From<lancedb::table::RefreshColumnResult> for RefreshColumnResult {
fn from(value: lancedb::table::RefreshColumnResult) -> Self {
Self {
rows_filled: value.rows_filled as i64,
version: value.version as i64,
}
}
}
impl From<lancedb::table::AddColumnsResult> for AddColumnsResult { impl From<lancedb::table::AddColumnsResult> for AddColumnsResult {
fn from(value: lancedb::table::AddColumnsResult) -> Self { fn from(value: lancedb::table::AddColumnsResult) -> Self {
Self { Self {
+5
View File
@@ -341,6 +341,7 @@ class Table:
async def add_computed_columns( async def add_computed_columns(
self, columns: list[tuple[str, str]] self, columns: list[tuple[str, str]]
) -> AddColumnsResult: ... ) -> AddColumnsResult: ...
async def refresh_column(self, column: str) -> RefreshColumnResult: ...
async def add_columns_with_schema(self, schema: pa.Schema) -> AddColumnsResult: ... async def add_columns_with_schema(self, schema: pa.Schema) -> AddColumnsResult: ...
async def alter_columns( async def alter_columns(
self, columns: list[dict[str, Any]] self, columns: list[dict[str, Any]]
@@ -686,6 +687,10 @@ class LsmWriteSpec:
class AddColumnsResult: class AddColumnsResult:
version: int version: int
class RefreshColumnResult:
rows_filled: int
version: int
class AlterColumnsResult: class AlterColumnsResult:
version: int version: int
+3
View File
@@ -970,6 +970,9 @@ class RemoteTable(Table):
) )
return LOOP.run(self._table.add_columns(transforms)) return LOOP.run(self._table.add_columns(transforms))
def refresh_column(self, column: str):
raise NotImplementedError("computed columns are supported only on local tables")
def alter_columns( def alter_columns(
self, *alterations: Iterable[Dict[str, str]] self, *alterations: Iterable[Dict[str, str]]
) -> AlterColumnsResult: ) -> AlterColumnsResult:
+68 -7
View File
@@ -176,6 +176,7 @@ if TYPE_CHECKING:
CompactionStats, CompactionStats,
Tag, Tag,
AddColumnsResult, AddColumnsResult,
RefreshColumnResult,
AddResult, AddResult,
AlterColumnsResult, AlterColumnsResult,
UpdateFieldMetadataResult, UpdateFieldMetadataResult,
@@ -1943,9 +1944,10 @@ class Table(ABC):
data type is supplied. data type is supplied.
Unlike ``transforms``, the expression is stored rather than Unlike ``transforms``, the expression is stored rather than
evaluated now: the column is committed with no values, and a evaluated now: the column is committed with no values, and rows get
later refresh fills the rows. Declaring one therefore costs the them from [`refresh_column`][lancedb.table.Table.refresh_column].
same on a large table as on an empty one. Declaring one therefore costs the same on a large table as on an
empty one.
A refresh does not revisit rows it has already filled, so mutating A refresh does not revisit rows it has already filled, so mutating
an input leaves the value computed at fill time; recomputing means an input leaves the value computed at fill time; recomputing means
@@ -1967,8 +1969,37 @@ class Table(ABC):
>>> table = db.create_table("computed_demo", [{"x": 1}, {"x": 2}]) >>> table = db.create_table("computed_demo", [{"x": 1}, {"x": 2}])
>>> table.add_columns(computed={"doubled": "x * 2"}) >>> table.add_columns(computed={"doubled": "x * 2"})
AddColumnsResult(version=2) AddColumnsResult(version=2)
>>> table.to_arrow()["doubled"].to_pylist() >>> table.refresh_column("doubled")
[None, None] RefreshColumnResult(rows_filled=2, version=3)
>>> table.to_arrow().sort_by("x").to_pandas()
x doubled
0 1 2
1 2 4
"""
@abstractmethod
def refresh_column(self, column: str) -> "RefreshColumnResult":
"""
Fill the rows of a computed column that hold no value yet.
Declared with ``add_columns(computed=...)``, a column starts empty and
gets its values here. Rows appended since the last refresh are filled
by the next one; rows already filled are left as they are, so the call
is idempotent and does not observe a mutated input.
Local tables only; LanceDB Cloud and Enterprise raise
``NotImplementedError``.
Parameters
----------
column: str
The name of the computed column to fill.
Returns
-------
RefreshColumnResult
rows_filled: the number of rows given a value.
version: the new version number of the table.
""" """
@abstractmethod @abstractmethod
@@ -3984,6 +4015,11 @@ class LanceTable(Table):
) -> AddColumnsResult: ) -> AddColumnsResult:
return LOOP.run(self._table.add_columns(transforms, computed=computed)) return LOOP.run(self._table.add_columns(transforms, computed=computed))
def refresh_column(self, column: str) -> "RefreshColumnResult":
"""Fill a computed column's unfilled rows. See
[`AsyncTable.refresh_column`][lancedb.AsyncTable.refresh_column]."""
return LOOP.run(self._table.refresh_column(column))
def alter_columns( def alter_columns(
self, *alterations: Iterable[Dict[str, str]] self, *alterations: Iterable[Dict[str, str]]
) -> AlterColumnsResult: ) -> AlterColumnsResult:
@@ -5922,8 +5958,9 @@ class AsyncTable:
column's type and inputs are derived from the expression. column's type and inputs are derived from the expression.
Unlike ``transforms``, the expression is stored rather than Unlike ``transforms``, the expression is stored rather than
evaluated now: the column is committed with no values, and a evaluated now: the column is committed with no values, and rows get
later refresh fills the rows. them from
[`refresh_column`][lancedb.table.AsyncTable.refresh_column].
A refresh does not revisit rows it has already filled, so mutating A refresh does not revisit rows it has already filled, so mutating
an input leaves the value computed at fill time. While a an input leaves the value computed at fill time. While a
@@ -5957,6 +5994,30 @@ class AsyncTable:
else: else:
return await self._inner.add_columns(list(transforms.items())) return await self._inner.add_columns(list(transforms.items()))
async def refresh_column(self, column: str) -> RefreshColumnResult:
"""
Fill the rows of a computed column that hold no value yet.
Declared with ``add_columns(computed=...)``, a column starts empty and
gets its values here. Rows appended since the last refresh are filled
by the next one; rows already filled are left as they are, so the call
is idempotent and does not observe a mutated input.
Local tables only; LanceDB Cloud and Enterprise raise
``NotImplementedError``.
Parameters
----------
column: str
The name of the computed column to fill.
Returns
-------
RefreshColumnResult
The number of rows filled and the new version of the table.
"""
return await self._inner.refresh_column(column)
async def alter_columns( async def alter_columns(
self, *alterations: Iterable[dict[str, Any]] self, *alterations: Iterable[dict[str, Any]]
) -> AlterColumnsResult: ) -> AlterColumnsResult:
+19 -4
View File
@@ -3856,16 +3856,20 @@ async def test_async_search_runs_embedding_on_dedicated_executor(
) )
def test_computed_column_declares_all_null(tmp_path): def test_computed_column_declare_and_refresh(tmp_path):
db = lancedb.connect(tmp_path) db = lancedb.connect(tmp_path)
table = db.create_table("computed", [{"x": 1}, {"x": 2}]) table = db.create_table("computed", [{"x": 1}, {"x": 2}])
table.add_columns(computed={"doubled": "x * 2"}) table.add_columns(computed={"doubled": "x * 2"})
assert table.to_arrow()["doubled"].to_pylist() == [None, None] assert table.to_arrow()["doubled"].to_pylist() == [None, None]
# The declaration is durable field metadata. result = table.refresh_column("doubled")
field = table.schema.field("doubled") assert result.rows_filled == 2
assert field.metadata[b"computed_column.expression"] == b"x * 2" assert sorted(table.to_arrow()["doubled"].to_pylist()) == [2, 4]
table.add([{"x": 5}])
assert table.refresh_column("doubled").rows_filled == 1
assert sorted(table.to_arrow()["doubled"].to_pylist()) == [2, 4, 10]
def test_computed_column_rejects_transforms_and_computed_together(tmp_path): def test_computed_column_rejects_transforms_and_computed_together(tmp_path):
@@ -3873,3 +3877,14 @@ def test_computed_column_rejects_transforms_and_computed_together(tmp_path):
table = db.create_table("computed_mixed", [{"x": 1}]) table = db.create_table("computed_mixed", [{"x": 1}])
with pytest.raises(ValueError): with pytest.raises(ValueError):
table.add_columns({"a": "x + 1"}, computed={"b": "x * 2"}) table.add_columns({"a": "x + 1"}, computed={"b": "x * 2"})
@pytest.mark.asyncio
async def test_computed_column_async(tmp_path):
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")
assert (await table.to_arrow())["tripled"].to_pylist() == [9]
+3 -1
View File
@@ -16,7 +16,8 @@ use query::{FTSQuery, HybridQuery, Query, VectorQuery};
use session::Session; use session::Session;
use table::{ use table::{
AddColumnsResult, AddResult, AlterColumnsResult, DeleteResult, DropColumnsResult, FtsToken, AddColumnsResult, AddResult, AlterColumnsResult, DeleteResult, DropColumnsResult, FtsToken,
LsmWriteSpec, MergeResult, PyBlobFile, Table, UpdateFieldMetadataResult, UpdateResult, LsmWriteSpec, MergeResult, PyBlobFile, RefreshColumnResult, Table, UpdateFieldMetadataResult,
UpdateResult,
}; };
pub mod arrow; pub mod arrow;
@@ -57,6 +58,7 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<VectorQuery>()?; m.add_class::<VectorQuery>()?;
m.add_class::<RecordBatchStream>()?; m.add_class::<RecordBatchStream>()?;
m.add_class::<AddColumnsResult>()?; m.add_class::<AddColumnsResult>()?;
m.add_class::<RefreshColumnResult>()?;
m.add_class::<AlterColumnsResult>()?; m.add_class::<AlterColumnsResult>()?;
m.add_class::<UpdateFieldMetadataResult>()?; m.add_class::<UpdateFieldMetadataResult>()?;
m.add_class::<AddResult>()?; m.add_class::<AddResult>()?;
+34
View File
@@ -415,6 +415,32 @@ pub struct AddColumnsResult {
pub version: u64, pub version: u64,
} }
#[pyclass(get_all, from_py_object)]
#[derive(Clone, Debug)]
pub struct RefreshColumnResult {
pub rows_filled: u64,
pub version: u64,
}
#[pymethods]
impl RefreshColumnResult {
pub fn __repr__(&self) -> String {
format!(
"RefreshColumnResult(rows_filled={}, version={})",
self.rows_filled, self.version
)
}
}
impl From<lancedb::table::RefreshColumnResult> for RefreshColumnResult {
fn from(result: lancedb::table::RefreshColumnResult) -> Self {
Self {
rows_filled: result.rows_filled,
version: result.version,
}
}
}
#[pymethods] #[pymethods]
impl AddColumnsResult { impl AddColumnsResult {
pub fn __repr__(&self) -> String { pub fn __repr__(&self) -> String {
@@ -1525,6 +1551,14 @@ impl Table {
}) })
} }
pub fn refresh_column(self_: PyRef<'_, Self>, column: String) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
let result = inner.refresh_column(column).await.infer_error()?;
Ok(RefreshColumnResult::from(result))
})
}
pub fn add_columns_with_schema( pub fn add_columns_with_schema(
self_: PyRef<'_, Self>, self_: PyRef<'_, Self>,
schema: PyArrowType<Schema>, schema: PyArrowType<Schema>,
+6
View File
@@ -6479,6 +6479,12 @@ mod tests {
matches!(&err, Error::NotSupported { message } if message.contains("local tables")), matches!(&err, Error::NotSupported { message } if message.contains("local tables")),
"{err:?}" "{err:?}"
); );
let err = table.refresh_column("doubled").await.unwrap_err();
assert!(
matches!(&err, Error::NotSupported { message } if message.contains("local tables")),
"{err:?}"
);
} }
#[tokio::test] #[tokio::test]
+39
View File
@@ -78,6 +78,7 @@ pub mod merge;
pub mod optimize; pub mod optimize;
mod primary_key; mod primary_key;
pub mod query; pub mod query;
pub mod refresh;
pub mod schema_evolution; pub mod schema_evolution;
pub mod update; pub mod update;
pub mod write_progress; pub mod write_progress;
@@ -101,6 +102,7 @@ pub use lance::dataset::scanner::DatasetRecordBatchStream;
pub use lance_index::optimize::OptimizeOptions; pub use lance_index::optimize::OptimizeOptions;
pub use lsm_stats::{BucketStats, GenerationStats, LsmStats, MemtableStats}; pub use lsm_stats::{BucketStats, GenerationStats, LsmStats, MemtableStats};
pub use optimize::{CompactionOptions, OptimizeAction, OptimizeStats}; pub use optimize::{CompactionOptions, OptimizeAction, OptimizeStats};
pub use refresh::RefreshColumnResult;
pub use schema_evolution::{ pub use schema_evolution::{
AddColumnsResult, AlterColumnsResult, DropColumnsResult, FieldMetadataUpdate, AddColumnsResult, AlterColumnsResult, DropColumnsResult, FieldMetadataUpdate,
UpdateFieldMetadataResult, UpdateFieldMetadataResult,
@@ -754,6 +756,14 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync {
message: "computed columns are not supported on this table type".into(), message: "computed columns are not supported on this table type".into(),
}) })
} }
/// Fill a computed column's unfilled rows.
///
/// The default returns `NotSupported`; Lance-backed tables override it.
async fn refresh_column(&self, _column: &str) -> Result<RefreshColumnResult> {
Err(Error::NotSupported {
message: "computed columns are supported only on local tables".into(),
})
}
/// Alter columns in the table. /// Alter columns in the table.
async fn alter_columns(&self, alterations: &[ColumnAlteration]) -> Result<AlterColumnsResult>; async fn alter_columns(&self, alterations: &[ColumnAlteration]) -> Result<AlterColumnsResult>;
/// Drop columns from the table. /// Drop columns from the table.
@@ -1646,6 +1656,29 @@ impl Table {
AddColumnsBuilder::new(self.inner.clone()) AddColumnsBuilder::new(self.inner.clone())
} }
/// Fill the fragments of a computed column that hold no values yet.
///
/// Declared with
/// [`AddColumnsBuilder::computed`](add_columns::AddColumnsBuilder::computed),
/// a column starts empty and gets its values here. Fragments appended
/// since the last refresh are filled by the next one; fragments already
/// filled are left as they are, so the call is idempotent and does not
/// observe a mutated input.
///
/// Local tables only.
///
/// ```
/// # use lancedb::Table;
/// # async fn refresh(table: &Table) -> Result<(), Box<dyn std::error::Error>> {
/// let result = table.refresh_column("doubled").await?;
/// println!("filled {} rows at version {}", result.rows_filled, result.version);
/// # Ok(())
/// # }
/// ```
pub async fn refresh_column(&self, column: impl AsRef<str>) -> Result<RefreshColumnResult> {
self.inner.refresh_column(column.as_ref()).await
}
/// Change a column's name or nullability. /// Change a column's name or nullability.
pub async fn alter_columns( pub async fn alter_columns(
&self, &self,
@@ -3353,6 +3386,12 @@ impl BaseTable for NativeTable {
Ok(result) Ok(result)
} }
async fn refresh_column(&self, column: &str) -> Result<RefreshColumnResult> {
let result = refresh::execute_refresh_column(self, column).await?;
self.bump_freshness();
Ok(result)
}
async fn alter_columns(&self, alterations: &[ColumnAlteration]) -> Result<AlterColumnsResult> { async fn alter_columns(&self, alterations: &[ColumnAlteration]) -> Result<AlterColumnsResult> {
let result = schema_evolution::execute_alter_columns(self, alterations).await?; let result = schema_evolution::execute_alter_columns(self, alterations).await?;
self.bump_freshness(); self.bump_freshness();
+6 -3
View File
@@ -51,9 +51,10 @@ impl AddColumnsBuilder {
/// expression. /// expression.
/// ///
/// The column is committed with no values, so declaring one costs the same /// The column is committed with no values, so declaring one costs the same
/// on an empty table as on a large one. Rows get values from a later /// on an empty table as on a large one. Rows get values from
/// refresh, which fills every fragment that has none -- including /// [`Table::refresh_column`](super::Table::refresh_column), which fills
/// fragments appended since the last refresh. /// every fragment that has none -- including fragments appended since the
/// last refresh.
/// ///
/// Refresh does not revisit a fragment it has filled, so mutating an input /// Refresh does not revisit a fragment it has filled, so mutating an input
/// leaves the value computed at fill time; recomputing means dropping the /// leaves the value computed at fill time; recomputing means dropping the
@@ -71,6 +72,8 @@ impl AddColumnsBuilder {
/// .computed("doubled", "x * 2") /// .computed("doubled", "x * 2")
/// .execute() /// .execute()
/// .await?; /// .await?;
/// let filled = table.refresh_column("doubled").await?;
/// println!("filled {} rows", filled.rows_filled);
/// # Ok(()) /// # Ok(())
/// # } /// # }
/// ``` /// ```
+25 -6
View File
@@ -23,6 +23,7 @@ use std::sync::Arc;
use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema, SchemaRef}; use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema, SchemaRef};
use datafusion_common::tree_node::TreeNode; use datafusion_common::tree_node::TreeNode;
use datafusion_physical_plan::PhysicalExpr;
use lance::dataset::NewColumnTransform; use lance::dataset::NewColumnTransform;
use lance_datafusion::planner::Planner; use lance_datafusion::planner::Planner;
@@ -296,11 +297,18 @@ pub(crate) fn root(path: &str) -> &str {
path.split('.').next().unwrap_or(path) path.split('.').next().unwrap_or(path)
} }
/// A declaration's expression bound to a schema. /// A declaration's expression bound to a schema, ready to evaluate.
pub(crate) struct BoundExpression { pub(crate) struct BoundExpression {
/// The columns the expression names, as written; nested inputs keep /// The columns the expression names, as written; nested inputs keep
/// their dotted path. /// their dotted path.
pub inputs: Vec<String>, pub inputs: Vec<String>,
/// The top-level columns evaluation reads, in [`Self::read_schema`]
/// order. A nested input appears through its root.
pub roots: Vec<String>,
/// The projected schema evaluation runs against.
pub read_schema: SchemaRef,
/// The compiled expression.
pub physical: Arc<dyn PhysicalExpr>,
/// The type the expression yields. /// The type the expression yields.
pub data_type: DataType, pub data_type: DataType,
} }
@@ -373,6 +381,12 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result<
.project(&indices) .project(&indices)
.map_err(|e| invalid(e.to_string()))?, .map_err(|e| invalid(e.to_string()))?,
); );
let roots = read_schema
.fields()
.iter()
.map(|field| field.name().clone())
.collect();
let optimized = planner let optimized = planner
.optimize_expr(parsed) .optimize_expr(parsed)
.map_err(|e| invalid(e.to_string()))?; .map_err(|e| invalid(e.to_string()))?;
@@ -383,7 +397,13 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result<
.data_type(read_schema.as_ref()) .data_type(read_schema.as_ref())
.map_err(|e| invalid(e.to_string()))?; .map_err(|e| invalid(e.to_string()))?;
Ok(BoundExpression { inputs, data_type }) Ok(BoundExpression {
inputs,
roots,
read_schema,
physical,
data_type,
})
} }
/// Resolve `(name, expression)` pairs against `schema` into fields carrying /// Resolve `(name, expression)` pairs against `schema` into fields carrying
@@ -905,7 +925,7 @@ mod tests {
); );
// The declaration survives the refused change. // The declaration survives the refused change.
assert_eq!(declared(&table).await.len(), 1); table.refresh_column("doubled").await.unwrap();
} }
/// A declaration cannot be edited, fabricated or erased through field /// A declaration cannot be edited, fabricated or erased through field
@@ -947,13 +967,12 @@ mod tests {
.unwrap_err(); .unwrap_err();
assert!(matches!(err, Error::InvalidInput { .. }), "{err:?}"); assert!(matches!(err, Error::InvalidInput { .. }), "{err:?}");
// Ordinary metadata on a computed column still merges, leaving the // Ordinary metadata on a computed column still merges.
// declaration intact.
table table
.update_field_metadata(&[FieldMetadataUpdate::new("doubled").set("note", "hi")]) .update_field_metadata(&[FieldMetadataUpdate::new("doubled").set("note", "hi")])
.await .await
.unwrap(); .unwrap();
assert_eq!(declared(&table).await.len(), 1); table.refresh_column("doubled").await.unwrap();
} }
/// The gate's reproducer: only refresh materializes a declared column; /// The gate's reproducer: only refresh materializes a declared column;
+708
View File
@@ -0,0 +1,708 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
//! Filling computed columns.
//!
//! A row without a value gets one; a row that has one keeps it. Refresh is
//! therefore idempotent and does not observe input mutation -- once a row is
//! filled, changing what the expression reads leaves the stored result alone.
//!
//! Two passes per fragment. The first scans only the unfilled live rows and
//! evaluates the expression over them, which yields the exact fill count and
//! decides whether the fragment is staged at all -- a fragment where nothing
//! would change stages nothing, which is what lets an expression yielding
//! null settle instead of restaging forever. The second streams the
//! fragment's physical rows into `write_column` a batch at a time, so peak
//! memory is bounded by a scan batch. The expression is evaluated by this
//! module, never through a projection alias, and only over rows being
//! filled: every other row -- deleted, or already holding a value -- has its
//! inputs masked to null first, so a poison value in a row nobody is filling
//! cannot fail the refresh.
use std::sync::Arc;
use arrow_array::{ArrayRef, BooleanArray, RecordBatch, RecordBatchOptions};
use arrow_schema::Schema as ArrowSchema;
use datafusion_expr::ColumnarValue;
use futures::{Stream, StreamExt, TryStreamExt};
use lance::Dataset;
use lance::dataset::WriteDestination;
use lance::dataset::fragment::FileFragment;
use lance::dataset::transaction::Operation;
use lance_core::ROW_ID;
use lance_core::datatypes::Schema as LanceSchema;
use serde::{Deserialize, Serialize};
use super::computed_columns::{BoundExpression, ComputedColumnKind, computed_column_from_field};
use super::{BaseTable, NativeTable};
use crate::{Error, Result};
/// The result of refreshing a computed column.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
pub struct RefreshColumnResult {
/// Rows that had a value computed.
#[serde(default)]
pub rows_filled: u64,
/// The commit version associated with the operation.
#[serde(default)]
pub version: u64,
}
/// Internal implementation of the refresh logic.
pub(crate) async fn execute_refresh_column(
table: &NativeTable,
column: &str,
) -> Result<RefreshColumnResult> {
table.dataset.ensure_mutable()?;
ensure_no_lsm_write_spec(table).await?;
let dataset = table.dataset.get().await?;
let expression = declared_expression(&dataset, column)?;
let schema = Arc::new(ArrowSchema::from(dataset.schema()));
let bound = Arc::new(super::computed_columns::bind(schema, column, &expression)?);
let field = dataset
.schema()
.field(column)
.ok_or_else(|| Error::ColumnNotFound {
name: column.to_string(),
})?;
// The dataset's own field, so the identity write_column checks against the
// manifest holds by construction.
let column_schema = LanceSchema {
fields: vec![field.clone()],
metadata: Default::default(),
};
let mut rows_filled = 0u64;
let mut replacements = Vec::new();
for fragment in dataset.get_fragments() {
let gained = count_fragment_gains(&dataset, &fragment, &bound, column).await?;
if gained == 0 {
continue;
}
rows_filled += gained;
let values = fill_stream(&dataset, &fragment, bound.clone(), column).await?;
replacements.push(fragment.write_column(values, &column_schema).await?);
}
if replacements.is_empty() {
return Ok(RefreshColumnResult {
rows_filled: 0,
version: dataset.version().version,
});
}
let read_version = dataset.version().version;
// The dataset's own session, so registrations and caches survive the
// commit being installed on the handle.
let session = dataset.session();
let new_dataset = Dataset::commit(
WriteDestination::Dataset(dataset.clone()),
Operation::DataReplacement { replacements },
Some(read_version),
None,
None,
session,
false,
)
.await?;
let version = new_dataset.version().version;
table.dataset.update(new_dataset);
Ok(RefreshColumnResult {
rows_filled,
version,
})
}
/// Refuse to refresh under an LSM write spec.
///
/// Refresh enumerates base fragments, and a write spec keeps visible rows in
/// un-compacted MemWAL tiers it cannot reach -- success would silently omit
/// readable rows.
async fn ensure_no_lsm_write_spec(table: &NativeTable) -> Result<()> {
// The catch-up flag outlives unset and marks retained SSTable rows.
let catchup = table.dataset.get().await?.manifest().reader_feature_flags
& lance_table::feature_flags::FLAG_MEM_WAL_INDEX_CATCHUP
!= 0;
if catchup || table.get_lsm_write_spec().await?.is_some() {
return Err(Error::NotSupported {
message: "refresh_column is not supported on a table with an LSM write \
spec: rows in un-compacted tiers are invisible to refresh"
.into(),
});
}
Ok(())
}
/// The SQL expression `column` is declared with.
fn declared_expression(dataset: &Dataset, column: &str) -> Result<String> {
let schema = ArrowSchema::from(dataset.schema());
let field = schema
.field_with_name(column)
.map_err(|_| Error::ColumnNotFound {
name: column.to_string(),
})?;
let declaration =
computed_column_from_field(field).ok_or_else(|| Error::NotAComputedColumn {
name: column.to_string(),
})?;
match declaration.kind {
ComputedColumnKind::Sql { expression } => Ok(expression),
ComputedColumnKind::Unrecognized { kind } => Err(Error::NotSupported {
message: format!(
"computed column '{column}' is defined by '{kind}', which this version of \
lancedb cannot evaluate"
),
}),
}
}
/// Quote `name` as a lance SQL identifier.
///
/// Lance's dialect delimits with backticks, so a double-quoted name would
/// parse as a string literal rather than a column.
fn quote_identifier(name: &str) -> String {
format!("`{}`", name.replace('`', "``"))
}
/// Assemble the batch evaluation runs against: the bound roots, in read-schema
/// order. Built by name so scan-side column order never matters.
fn evaluation_batch(
batch: &RecordBatch,
bound: &BoundExpression,
mask_out: Option<&BooleanArray>,
) -> lance_core::Result<RecordBatch> {
let mut columns = Vec::with_capacity(bound.roots.len());
for name in &bound.roots {
let column = batch.column_by_name(name).ok_or_else(|| {
lance_core::Error::invalid_input(format!(
"refreshing a computed column read no {name} column"
))
})?;
// Rows outside the mask must not reach the expression: a value in a
// deleted or already-filled row can be one it would choke on.
columns.push(match mask_out {
Some(mask) => arrow::compute::nullif(column, mask)?,
None => column.clone(),
});
}
Ok(RecordBatch::try_new_with_options(
bound.read_schema.clone(),
columns,
&RecordBatchOptions::new().with_row_count(Some(batch.num_rows())),
)?)
}
/// Evaluate the expression over `batch`, materializing a constant result to
/// the batch's length.
fn evaluate(bound: &BoundExpression, batch: &RecordBatch) -> lance_core::Result<ArrayRef> {
let value = bound
.physical
.evaluate(batch)
.map_err(lance_core::Error::from)?;
match value {
ColumnarValue::Array(array) => Ok(array),
scalar => scalar
.into_array(batch.num_rows())
.map_err(lance_core::Error::from),
}
}
/// How many rows of one fragment would gain a value.
///
/// Scans only the unfilled live rows -- deleted rows never reach the
/// expression here, the filter having already excluded them -- and counts the
/// non-null results. Exact, so it is both the staging decision and the
/// fragment's contribution to `rows_filled`.
async fn count_fragment_gains(
dataset: &Dataset,
fragment: &FileFragment,
bound: &BoundExpression,
column: &str,
) -> Result<u64> {
let mut scanner = dataset.scan();
scanner
.with_fragments(vec![fragment.metadata().clone()])
.with_row_id()
.filter(&format!("{} IS NULL", quote_identifier(column)))?
.project(&bound.roots)?;
let mut gained = 0u64;
let mut batches = scanner.try_into_stream().await?;
while let Some(batch) = batches.try_next().await? {
let evaluated = evaluate(bound, &evaluation_batch(&batch, bound, None)?)?;
gained += (batch.num_rows() - evaluated.null_count()) as u64;
}
Ok(gained)
}
/// Stream one fragment's column in physical order, filling the unfilled live
/// rows and keeping every other value.
///
/// Deleted rows are carried through so the values line up positionally with
/// the fragment's data files; they are never read back, but the column file
/// has to cover them.
async fn fill_stream(
dataset: &Dataset,
fragment: &FileFragment,
bound: Arc<BoundExpression>,
column: &str,
) -> Result<impl Stream<Item = lance_core::Result<RecordBatch>> + Send + use<>> {
let mut projection: Vec<String> = bound.roots.clone();
projection.push(column.to_string());
let mut scanner = dataset.scan();
scanner
.with_fragments(vec![fragment.metadata().clone()])
.with_row_id()
.include_deleted_rows()
.project(&projection)?;
let projected = Arc::new(ArrowSchema::new(vec![
ArrowSchema::from(dataset.schema())
.field_with_name(column)
.map_err(|_| Error::ColumnNotFound {
name: column.to_string(),
})?
.clone(),
]));
let column = column.to_string();
let batches = scanner.try_into_stream().await?;
Ok(batches.map(move |batch| {
let batch = batch?;
let missing = |name: &str| {
lance_core::Error::invalid_input(format!(
"refreshing a computed column read no {name} column"
))
};
let existing = batch
.column_by_name(&column)
.ok_or_else(|| missing(&column))?;
let row_ids = batch
.column_by_name(ROW_ID)
.ok_or_else(|| missing(ROW_ID))?;
// Only an unfilled live row gains a value; a deleted row has a null
// row id and keeps its (null) slot.
let unfilled = arrow::compute::is_null(existing.as_ref())?;
let live = arrow::compute::is_not_null(row_ids.as_ref())?;
let fill = arrow::compute::and(&unfilled, &live)?;
let keep = arrow::compute::not(&fill)?;
let computed = evaluate(&bound, &evaluation_batch(&batch, &bound, Some(&keep))?)?;
let merged = arrow_select::zip::zip(&fill, &computed, existing)?;
Ok(RecordBatch::try_new(projected.clone(), vec![merged])?)
}))
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use arrow_array::{Int32Array, record_batch};
use futures::TryStreamExt;
use crate::connect;
use crate::query::{ExecutableQuery, QueryBase, Select};
use crate::{Error, Result, Table};
async fn table_with(name: &str, values: Vec<i32>) -> Table {
let conn = connect("memory://").execute().await.unwrap();
let batch = record_batch!(("x", Int32, values)).unwrap();
conn.create_table(name, batch).execute().await.unwrap()
}
async fn declare_doubled(table: &Table) -> Result<u64> {
Ok(table
.add_columns()
.computed("doubled", "x * 2")
.execute()
.await?
.version)
}
async fn read(table: &Table, column: &str) -> Vec<Option<i32>> {
let batches = table
.query()
.select(Select::columns(&[column]))
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
let mut values: Vec<Option<i32>> = batches
.iter()
.flat_map(|batch| {
batch[column]
.as_any()
.downcast_ref::<Int32Array>()
.unwrap()
.iter()
.collect::<Vec<_>>()
})
.collect();
values.sort();
values
}
async fn append(table: &Table, values: Vec<i32>) {
let batch = record_batch!(("x", Int32, values)).unwrap();
table.add(batch).execute().await.unwrap();
}
#[tokio::test]
async fn test_refresh_fills_a_declared_column() {
let table = table_with("refresh_fills", vec![1, 2, 3]).await;
let declared = declare_doubled(&table).await.unwrap();
assert_eq!(read(&table, "doubled").await, vec![None, None, None]);
let result = table.refresh_column("doubled").await.unwrap();
assert!(result.version > declared);
assert_eq!(result.rows_filled, 3);
assert_eq!(
read(&table, "doubled").await,
vec![Some(2), Some(4), Some(6)]
);
}
/// Values written after the last refresh must be reachable by another one.
#[tokio::test]
async fn test_refresh_fills_rows_appended_since_the_last_refresh() {
let table = table_with("refresh_appended", vec![1, 2]).await;
declare_doubled(&table).await.unwrap();
table.refresh_column("doubled").await.unwrap();
append(&table, vec![5, 6]).await;
assert_eq!(
read(&table, "doubled").await,
vec![None, None, Some(2), Some(4)]
);
let result = table.refresh_column("doubled").await.unwrap();
assert_eq!(result.rows_filled, 2);
assert_eq!(
read(&table, "doubled").await,
vec![Some(2), Some(4), Some(10), Some(12)]
);
}
#[tokio::test]
async fn test_refresh_with_nothing_to_fill() {
let table = table_with("refresh_noop", vec![1, 2, 3]).await;
declare_doubled(&table).await.unwrap();
table.refresh_column("doubled").await.unwrap();
let again = table.refresh_column("doubled").await.unwrap();
assert_eq!(again.rows_filled, 0);
assert_eq!(
read(&table, "doubled").await,
vec![Some(2), Some(4), Some(6)]
);
}
/// A row is filled only by gaining a value, so an expression yielding null
/// settles at once instead of re-selecting the same rows forever. Nothing
/// is staged, so the version does not move either.
#[tokio::test]
async fn test_refresh_converges_on_a_null_result() {
let table = table_with("refresh_null_result", vec![1, 2, 3]).await;
let declared = table
.add_columns()
.computed("maybe", "nullif(x, x)")
.execute()
.await
.unwrap()
.version;
let first = table.refresh_column("maybe").await.unwrap();
assert_eq!(first.rows_filled, 0);
assert_eq!(first.version, declared);
assert_eq!(read(&table, "maybe").await, vec![None, None, None]);
let again = table.refresh_column("maybe").await.unwrap();
assert_eq!(again.rows_filled, 0);
assert_eq!(again.version, declared);
}
/// The contract's boundary: a filled fragment is not revisited, so
/// mutating an input leaves the value computed at fill time.
#[tokio::test]
async fn test_refresh_does_not_observe_input_mutation() {
let table = table_with("refresh_mutation", vec![1]).await;
declare_doubled(&table).await.unwrap();
table.refresh_column("doubled").await.unwrap();
assert_eq!(read(&table, "doubled").await, vec![Some(2)]);
table.update().column("x", "3").execute().await.unwrap();
let again = table.refresh_column("doubled").await.unwrap();
assert_eq!(again.rows_filled, 0);
assert_eq!(read(&table, "doubled").await, vec![Some(2)]);
}
/// A row rewrite before the first refresh materializes the declared
/// column as null behind a covering data file. Those rows are still
/// unfilled and a later refresh has to reach them.
#[tokio::test]
async fn test_update_before_the_first_refresh() {
let table = table_with("refresh_update_first", vec![1]).await;
declare_doubled(&table).await.unwrap();
table.update().column("x", "3").execute().await.unwrap();
let result = table.refresh_column("doubled").await.unwrap();
assert_eq!(result.rows_filled, 1);
assert_eq!(read(&table, "doubled").await, vec![Some(6)]);
}
/// The contract holds row by row, not fragment by fragment: revisiting a
/// fragment to fill one row must not recompute a filled row sitting beside
/// it, even where the input behind it has since changed.
#[tokio::test]
async fn test_refresh_does_not_recompute_a_filled_row_beside_an_unfilled_one() {
let table = table_with("refresh_mixed", vec![1, 2]).await;
declare_doubled(&table).await.unwrap();
table.refresh_column("doubled").await.unwrap();
append(&table, vec![5]).await;
table
.update()
.column("x", "100")
.only_if("x = 1")
.execute()
.await
.unwrap();
table
.optimize(crate::table::OptimizeAction::Compact {
options: crate::table::CompactionOptions::default(),
remap_options: None,
})
.await
.unwrap();
let result = table.refresh_column("doubled").await.unwrap();
assert_eq!(result.rows_filled, 1);
// 2 is the mutated row keeping the value it was filled with, not 200.
assert_eq!(
read(&table, "doubled").await,
vec![Some(2), Some(4), Some(10)]
);
}
/// Filling a fragment must not disturb the values it already holds, which
/// is what makes a compaction-mixed fragment safe to revisit.
#[tokio::test]
async fn test_refresh_preserves_already_filled_rows() {
let table = table_with("refresh_preserves", vec![1, 2]).await;
declare_doubled(&table).await.unwrap();
table.refresh_column("doubled").await.unwrap();
append(&table, vec![5]).await;
table
.optimize(crate::table::OptimizeAction::Compact {
options: crate::table::CompactionOptions::default(),
remap_options: None,
})
.await
.unwrap();
let result = table.refresh_column("doubled").await.unwrap();
assert_eq!(result.rows_filled, 1);
assert_eq!(
read(&table, "doubled").await,
vec![Some(2), Some(4), Some(10)]
);
}
#[tokio::test]
async fn test_refresh_leaves_deleted_rows_alone() {
let table = table_with("refresh_deleted", vec![1, 2, 3, 4]).await;
declare_doubled(&table).await.unwrap();
table.delete("x = 2").await.unwrap();
let result = table.refresh_column("doubled").await.unwrap();
assert_eq!(result.rows_filled, 3);
assert_eq!(
read(&table, "doubled").await,
vec![Some(2), Some(6), Some(8)]
);
}
#[tokio::test]
async fn test_refresh_a_constant_expression() {
let table = table_with("refresh_constant", vec![1, 2, 3]).await;
table
.add_columns()
.computed("answer", "42")
.execute()
.await
.unwrap();
let result = table.refresh_column("answer").await.unwrap();
assert_eq!(result.rows_filled, 3);
}
/// A name needing quotes reaches the evaluator intact: it is carried as a
/// projection alias, never spliced into SQL text.
#[tokio::test]
async fn test_refresh_a_column_whose_name_needs_quoting() {
let table = table_with("refresh_quoted", vec![1, 2, 3]).await;
table
.add_columns()
.computed("double value", "x * 2")
.execute()
.await
.unwrap();
let result = table.refresh_column("double value").await.unwrap();
assert_eq!(result.rows_filled, 3);
assert_eq!(
read(&table, "double value").await,
vec![Some(2), Some(4), Some(6)]
);
}
/// A fragment spanning several scan batches exercises the streamed fill:
/// the probe buffers only until the first gained value and the rest flows
/// through write_column a batch at a time.
#[tokio::test]
async fn test_refresh_streams_a_multi_batch_fragment() {
let values: Vec<i32> = (0..20_000).collect();
let table = table_with("refresh_multi_batch", values.clone()).await;
declare_doubled(&table).await.unwrap();
let result = table.refresh_column("doubled").await.unwrap();
assert_eq!(result.rows_filled, 20_000);
let read_back = read(&table, "doubled").await;
assert_eq!(read_back.len(), 20_000);
let mut expected: Vec<Option<i32>> = values.iter().map(|v| Some(v * 2)).collect();
expected.sort();
assert_eq!(read_back, expected);
}
/// The gate's reproducer: the commit must reuse the configured session,
/// or registrations and caches vanish from the handle after a refresh.
#[tokio::test]
async fn test_refresh_preserves_the_configured_session() {
let session = Arc::new(lance::session::Session::default());
let conn = crate::connect("memory://")
.session(session.clone())
.execute()
.await
.unwrap();
let batch = record_batch!(("x", Int32, [1, 2])).unwrap();
let table = conn
.create_table("session_kept", batch)
.execute()
.await
.unwrap();
declare_doubled(&table).await.unwrap();
table.refresh_column("doubled").await.unwrap();
let dataset = table.as_native().unwrap().dataset.get().await.unwrap();
assert!(Arc::ptr_eq(&dataset.session(), &session));
}
/// Both orders of declare+spec are refused at the source (see the
/// schema_evolution tests); refresh's own check covers a dataset another
/// writer left in that state.
#[tokio::test]
async fn test_refresh_refuses_a_foreign_lsm_state() {
use crate::table::LsmWriteSpec;
let tmp_dir = tempfile::tempdir().unwrap();
let conn = connect(tmp_dir.path().to_str().unwrap())
.execute()
.await
.unwrap();
let schema = Arc::new(arrow_schema::Schema::new(vec![arrow_schema::Field::new(
"x",
arrow_schema::DataType::Int32,
false,
)]));
let batch =
arrow_array::RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1]))])
.unwrap();
let table = conn.create_table("lsm", batch).execute().await.unwrap();
table.set_unenforced_primary_key(["x"]).await.unwrap();
table
.set_lsm_write_spec(LsmWriteSpec::unsharded())
.await
.unwrap();
super::super::computed_columns::add_foreign_kind(&table, "doubled", "sql").await;
let err = table.refresh_column("doubled").await.unwrap_err();
assert!(
matches!(&err, Error::NotSupported { message } if message.contains("LSM")),
"{err:?}"
);
}
/// After catch-up activation and unset, no spec remains but the catch-up
/// flag still marks retained SSTable rows; refresh refuses on the flag.
#[tokio::test]
async fn test_refresh_refuses_retained_catchup_state() {
use crate::table::LsmWriteSpec;
let tmp_dir = tempfile::tempdir().unwrap();
let conn = connect(tmp_dir.path().to_str().unwrap())
.execute()
.await
.unwrap();
let schema = Arc::new(arrow_schema::Schema::new(vec![arrow_schema::Field::new(
"x",
arrow_schema::DataType::Int32,
false,
)]));
let batch = arrow_array::RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Int32Array::from(vec![1]))],
)
.unwrap();
let table = conn
.create_table("catchup", batch.clone())
.execute()
.await
.unwrap();
table.set_unenforced_primary_key(["x"]).await.unwrap();
table
.set_lsm_write_spec(LsmWriteSpec::unsharded())
.await
.unwrap();
table.require_mem_wal_index_catchup().await.unwrap();
let mut merge = table.merge_insert(&["x"]);
merge
.when_matched_update_all(None)
.when_not_matched_insert_all()
.use_lsm(true);
merge
.execute(Box::new(arrow_array::RecordBatchIterator::new(
vec![Ok(batch)],
schema,
)))
.await
.unwrap();
table.unset_lsm_write_spec().await.unwrap();
super::super::computed_columns::add_foreign_kind(&table, "doubled", "sql").await;
let err = table.refresh_column("doubled").await.unwrap_err();
assert!(
matches!(&err, Error::NotSupported { message } if message.contains("LSM")),
"{err:?}"
);
}
/// A declaration of a kind this version cannot evaluate is refused by
/// name, rather than mistaken for a plain column or fed to the SQL path.
#[tokio::test]
async fn test_refresh_rejects_a_kind_it_cannot_evaluate() {
let table = table_with("refresh_foreign", vec![1, 2, 3]).await;
super::super::computed_columns::add_foreign_kind(&table, "embedding", "udf").await;
let err = table.refresh_column("embedding").await.unwrap_err();
assert!(matches!(err, Error::NotSupported { message } if message.contains("udf")));
}
}