mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-18 12:08:35 +00:00
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:
@@ -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`<[`RefreshColumnResult`](../interfaces/RefreshColumnResult.md)>
|
||||||
|
|
||||||
|
A promise that resolves to the
|
||||||
|
number of rows filled and the new version number of the table.
|
||||||
|
|
||||||
|
***
|
||||||
|
|
||||||
### restore()
|
### restore()
|
||||||
|
|
||||||
```ts
|
```ts
|
||||||
|
|||||||
@@ -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;
|
||||||
|
```
|
||||||
@@ -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]);
|
||||||
});
|
});
|
||||||
});
|
});
|
||||||
|
|||||||
@@ -50,6 +50,7 @@ export {
|
|||||||
MergeResult,
|
MergeResult,
|
||||||
AddResult,
|
AddResult,
|
||||||
AddColumnsResult,
|
AddColumnsResult,
|
||||||
|
RefreshColumnResult,
|
||||||
AlterColumnsResult,
|
AlterColumnsResult,
|
||||||
UpdateFieldMetadataResult,
|
UpdateFieldMetadataResult,
|
||||||
DeleteResult,
|
DeleteResult,
|
||||||
|
|||||||
+21
-3
@@ -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> {
|
||||||
|
|||||||
@@ -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 {
|
||||||
|
|||||||
@@ -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
|
||||||
|
|
||||||
|
|||||||
@@ -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:
|
||||||
|
|||||||
@@ -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:
|
||||||
|
|||||||
@@ -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
@@ -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>()?;
|
||||||
|
|||||||
@@ -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>,
|
||||||
|
|||||||
@@ -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]
|
||||||
|
|||||||
@@ -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();
|
||||||
|
|||||||
@@ -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(())
|
||||||
/// # }
|
/// # }
|
||||||
/// ```
|
/// ```
|
||||||
|
|||||||
@@ -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;
|
||||||
|
|||||||
@@ -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")));
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user