refactor: unify computed column declarations

This commit is contained in:
Xuanwo
2026-08-28 14:57:04 +08:00
parent c93911033f
commit 0de1d77f1b
10 changed files with 126 additions and 73 deletions
+1 -1
View File
@@ -38,7 +38,7 @@ from .materialized_view import (
MaterializedView,
MaterializedViewDefinition,
)
from .table import AsyncTable, Table
from .table import AsyncTable, ComputedColumn as ComputedColumn, Table
from .types import BaseTokenizerType
from ._lancedb import Session
from .namespace import (
+1 -2
View File
@@ -353,8 +353,7 @@ class Table:
async def add_columns(self, columns: list[tuple[str, str]]) -> AddColumnsResult: ...
async def add_computed_columns(
self,
columns: list[tuple[str, str]],
blob_columns: Optional[list[tuple[str, str]]] = None,
columns: list[tuple[str, str, Literal["inferred", "blob_v2"]]],
) -> AddColumnsResult: ...
async def add_function_columns(
self, application_json: str, output_name: Optional[str]
+12 -8
View File
@@ -67,7 +67,16 @@ from ..query import (
LanceTakeQueryBuilder,
LanceVectorQueryBuilder,
)
from ..table import AsyncTable, BlobMode, Branches, IndexStatistics, Query, Table, Tags
from ..table import (
AsyncTable,
BlobMode,
Branches,
ComputedColumn,
IndexStatistics,
Query,
Table,
Tags,
)
from ..types import BaseTokenizerType
@@ -975,14 +984,9 @@ class RemoteTable(Table):
| FunctionApplication
| None = None,
*,
computed: Dict[str, str] | None = None,
computed_blobs: Dict[str, str] | None = None,
computed: Dict[str, str | ComputedColumn] | None = None,
) -> AddColumnsResult:
return LOOP.run(
self._table.add_columns(
transforms, computed=computed, computed_blobs=computed_blobs
)
)
return LOOP.run(self._table.add_columns(transforms, computed=computed))
def refresh_column(self, column: str):
return LOOP.run(self._table.refresh_column(column))
+58 -30
View File
@@ -913,6 +913,47 @@ def _normalize_progress(progress):
return progress, False
@dataclass(frozen=True)
class ComputedColumn:
"""A computed-column expression with explicit output semantics.
Plain strings passed to ``Table.add_columns(computed=...)`` infer their
output type. Use :meth:`blob` when a ``LargeBinary`` expression should be
published as Blob v2.
"""
expression: str
output: Literal["inferred", "blob_v2"] = "inferred"
def __post_init__(self):
if not isinstance(self.expression, str):
raise TypeError("ComputedColumn.expression must be a string")
if self.output not in ("inferred", "blob_v2"):
raise ValueError("ComputedColumn.output must be 'inferred' or 'blob_v2'")
@classmethod
def blob(cls, expression: str) -> ComputedColumn:
"""Publish a ``LargeBinary`` expression result as Blob v2."""
return cls(expression=expression, output="blob_v2")
def _normalize_computed_columns(
computed: Dict[str, str | ComputedColumn],
) -> list[tuple[str, str, Literal["inferred", "blob_v2"]]]:
columns: list[tuple[str, str, Literal["inferred", "blob_v2"]]] = []
for name, declaration in computed.items():
if isinstance(declaration, str):
columns.append((name, declaration, "inferred"))
elif isinstance(declaration, ComputedColumn):
columns.append((name, declaration.expression, declaration.output))
else:
raise TypeError(
"computed values must be SQL expression strings or "
"ComputedColumn values"
)
return columns
class Table(ABC):
"""
A Table is a collection of Records in a LanceDB Database.
@@ -2142,8 +2183,7 @@ class Table(ABC):
| pa.Schema
| None = None,
*,
computed: Dict[str, str] | None = None,
computed_blobs: Dict[str, str] | None = None,
computed: Dict[str, str | ComputedColumn] | None = None,
):
"""
Add new columns with defined values.
@@ -2165,10 +2205,13 @@ class Table(ABC):
atomic binding; aliases come from ``rename(columns=...)``.
Function columns are supported only on LanceDB Cloud and
Enterprise.
computed: Dict[str, str], optional
computed: Dict[str, str | ComputedColumn], optional
A map of column name to a SQL expression defining the column. The
column's type and inputs are derived from the expression, so no
data type is supplied.
data type is supplied. Use ``ComputedColumn.blob(expression)``
when a ``LargeBinary`` result should be stored as Blob v2. All
entries share mapping insertion order, including dependencies
between inferred and Blob outputs.
Unlike ``transforms``, the expression is stored rather than
evaluated now: the column is committed with no values, and rows get
@@ -2185,12 +2228,6 @@ class Table(ABC):
server, and the refresh runs as a server job -- see
[`refresh_column_async`][lancedb.table.Table.refresh_column_async].
Cannot be combined with ``transforms``.
computed_blobs: Dict[str, str], optional
A map of Blob v2 output column names to SQL expressions returning
``LargeBinary`` payload bytes. Blob inputs named by an expression
are materialized as bytes, and refresh stores the result as Blob
v2 so ``blob_columns()`` and Blob read APIs continue to recognize
it. Cannot be combined with ``transforms``.
Returns
-------
@@ -4306,14 +4343,9 @@ class LanceTable(Table):
| pa.Schema
| None = None,
*,
computed: Dict[str, str] | None = None,
computed_blobs: Dict[str, str] | None = None,
computed: Dict[str, str | ComputedColumn] | None = None,
) -> AddColumnsResult:
return LOOP.run(
self._table.add_columns(
transforms, computed=computed, computed_blobs=computed_blobs
)
)
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
@@ -6259,8 +6291,7 @@ class AsyncTable:
| pa.Schema
| None = None,
*,
computed: dict[str, str] | None = None,
computed_blobs: dict[str, str] | None = None,
computed: dict[str, str | ComputedColumn] | None = None,
) -> AddColumnsResult:
"""
Add new columns with defined values.
@@ -6280,9 +6311,12 @@ class AsyncTable:
atomic binding; aliases come from ``rename(columns=...)``.
Function columns are supported only on LanceDB Cloud and
Enterprise.
computed: Dict[str, str], optional
computed: Dict[str, str | ComputedColumn], optional
A map of column name to a SQL expression defining the column. The
column's type and inputs are derived from the expression.
column's type and inputs are derived from the expression. Use
``ComputedColumn.blob(expression)`` to publish a ``LargeBinary``
result as Blob v2. Mapping insertion order is the declaration and
dependency order.
Unlike ``transforms``, the expression is stored rather than
evaluated now: the column is committed with no values, and rows get
@@ -6296,11 +6330,6 @@ class AsyncTable:
On LanceDB Cloud and Enterprise the expression is planned by
the server. Cannot be combined with ``transforms``.
computed_blobs: Dict[str, str], optional
A map of Blob v2 output column names to SQL expressions returning
``LargeBinary`` payload bytes. Blob inputs are materialized as
payload bytes during refresh. Cannot be combined with
``transforms``.
Returns
-------
@@ -6324,7 +6353,7 @@ class AsyncTable:
function_output_name, function_application = next(iter(transforms.items()))
if function_application is not None:
if computed or computed_blobs:
if computed:
raise ValueError(
"add_columns cannot mix a Function application with SQL "
"computed columns"
@@ -6340,14 +6369,13 @@ class AsyncTable:
{isinstance(f, pa.Field) for f in transforms}
):
transforms = pa.schema(transforms)
if computed or computed_blobs:
if computed:
if transforms:
raise ValueError(
"add_columns cannot take both transforms and computed columns"
)
return await self._inner.add_computed_columns(
list((computed or {}).items()),
list((computed_blobs or {}).items()),
_normalize_computed_columns(computed)
)
if transforms is None:
raise ValueError("add_columns requires transforms or computed columns")
+19 -4
View File
@@ -4100,11 +4100,13 @@ def test_computed_blob_input_and_explicit_output(tmp_path):
)
table.add_columns(
computed={"payload_copy": "image"},
computed_blobs={"image_copy": "image"},
computed={
"image_copy": lancedb.ComputedColumn.blob("image"),
"payload_copy": "image_copy",
}
)
assert table.refresh_column("payload_copy").rows_filled == 2
assert table.refresh_column("image_copy").rows_filled == 2
assert table.refresh_column("payload_copy").rows_filled == 2
values = table.to_arrow()["payload_copy"].combine_chunks().to_pylist()
assert values == [b"hello", b"", None]
@@ -4120,7 +4122,20 @@ def test_computed_blob_rejects_eager_transforms(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("computed_blob_mixed", [{"x": 1}])
with pytest.raises(ValueError):
table.add_columns({"a": "x + 1"}, computed_blobs={"b": "x"})
table.add_columns(
{"a": "x + 1"},
computed={"b": lancedb.ComputedColumn.blob("x")},
)
def test_computed_column_validates_explicit_output():
assert lancedb.ComputedColumn("x + 1").output == "inferred"
assert lancedb.ComputedColumn.blob("image").output == "blob_v2"
with pytest.raises(TypeError, match="expression must be a string"):
lancedb.ComputedColumn(42) # type: ignore[arg-type]
with pytest.raises(ValueError, match="output must be"):
lancedb.ComputedColumn("x", output="binary") # type: ignore[arg-type]
@pytest.mark.asyncio
+15 -10
View File
@@ -20,8 +20,9 @@ use arrow::{
use lancedb::blob::{BlobFile, BlobRangeRequest};
use lancedb::index::scalar::FtsIndexBuilder;
use lancedb::table::{
AddDataMode, ColumnAlteration, Duration, FieldMetadataUpdate, FtsToken as LanceDbFtsToken,
NewColumnTransform, OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable,
AddDataMode, ColumnAlteration, ComputedColumnDeclaration, Duration, FieldMetadataUpdate,
FtsToken as LanceDbFtsToken, NewColumnTransform, OptimizeAction, OptimizeOptions, Ref,
Table as LanceDbTable,
};
use lancedb::tokenize as lancedb_tokenize;
use pyo3::{
@@ -1575,20 +1576,24 @@ impl Table {
})
}
#[pyo3(signature = (columns, blob_columns=None))]
pub fn add_computed_columns(
self_: PyRef<'_, Self>,
columns: Vec<(String, String)>,
blob_columns: Option<Vec<(String, String)>>,
columns: Vec<(String, String, String)>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
let mut builder = inner.add_columns();
for (name, expression) in columns {
builder = builder.computed(name, expression);
}
for (name, expression) in blob_columns.unwrap_or_default() {
builder = builder.computed_blob(name, expression);
for (name, expression, output) in columns {
let declaration = match output.as_str() {
"inferred" => ComputedColumnDeclaration::inferred(name, expression),
"blob_v2" => ComputedColumnDeclaration::blob(name, expression),
output => {
return Err(PyValueError::new_err(format!(
"unsupported computed-column output '{output}'"
)));
}
};
builder = builder.computed_column(declaration);
}
let result = builder.execute().await.infer_error()?;
Ok(AddColumnsResult::from(result))