diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index 21ffc8860..8d6993ca6 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -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 ( diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 6c2800ed1..edbcc5ccc 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -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] diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 7f30e4d7c..75cd71bdc 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -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)) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 09166519a..2571fbda7 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -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") diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 4cac132d0..f05dd148f 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -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 diff --git a/python/src/table.rs b/python/src/table.rs index d50d537d2..33aba3792 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -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>, + columns: Vec<(String, String, String)>, ) -> PyResult> { 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)) diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index cc53e331d..619018dab 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -7450,7 +7450,10 @@ mod tests { let result = table .add_columns() - .computed_blob("image_copy", "image") + .computed_column(crate::table::ComputedColumnDeclaration::blob( + "image_copy", + "image", + )) .execute() .await .unwrap(); diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 96e041320..8ba9d5e05 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -95,7 +95,8 @@ pub use cherry_pick::{ }; pub use chrono::Duration; pub use computed_columns::{ - ComputedColumn, ComputedColumnKind, computed_column_from_field, computed_columns, + ComputedColumn, ComputedColumnDeclaration, ComputedColumnKind, ComputedColumnOutput, + computed_column_from_field, computed_columns, }; pub use delete::DeleteResult; use futures::future::join_all; diff --git a/rust/lancedb/src/table/add_columns.rs b/rust/lancedb/src/table/add_columns.rs index 0d8aab73f..6265e0197 100644 --- a/rust/lancedb/src/table/add_columns.rs +++ b/rust/lancedb/src/table/add_columns.rs @@ -83,35 +83,32 @@ impl AddColumnsBuilder { /// # Ok(()) /// # } /// ``` - pub fn computed(mut self, name: impl Into, expression: impl Into) -> Self { - self.computed - .push(ComputedColumnDeclaration::inferred(name, expression)); - self + pub fn computed(self, name: impl Into, expression: impl Into) -> Self { + self.computed_column(ComputedColumnDeclaration::inferred(name, expression)) } - /// Add a Blob v2 column defined by a `LargeBinary` expression and filled - /// by a later refresh. + /// Add one computed-column declaration with explicit output semantics. /// - /// Blob inputs in the expression are materialized as their payload bytes. - /// The expression result is wrapped back into the Blob v2 logical type - /// before publication, so queries and [`Table::blob_columns`](super::Table::blob_columns) - /// continue to recognize the output as a Blob column. + /// Use [`ComputedColumnDeclaration::blob`] when a `LargeBinary` expression + /// should be published as Blob v2. Declarations share one ordered stream, + /// so a later expression may reference a column declared earlier in the + /// same builder. /// /// ``` /// # use lancedb::Table; + /// # use lancedb::table::ComputedColumnDeclaration; /// # async fn declare(table: &Table) -> Result<(), Box> { /// table /// .add_columns() - /// .computed_blob("image_copy", "image") + /// .computed_column(ComputedColumnDeclaration::blob("image_copy", "image")) /// .execute() /// .await?; /// table.refresh_column("image_copy").await?; /// # Ok(()) /// # } /// ``` - pub fn computed_blob(mut self, name: impl Into, expression: impl Into) -> Self { - self.computed - .push(ComputedColumnDeclaration::blob(name, expression)); + pub fn computed_column(mut self, declaration: ComputedColumnDeclaration) -> Self { + self.computed.push(declaration); self } diff --git a/rust/lancedb/src/table/refresh.rs b/rust/lancedb/src/table/refresh.rs index a8837eb5a..8272b08d7 100644 --- a/rust/lancedb/src/table/refresh.rs +++ b/rust/lancedb/src/table/refresh.rs @@ -513,6 +513,7 @@ mod tests { use crate::connect; use crate::query::{ExecutableQuery, QueryBase, Select}; + use crate::table::ComputedColumnDeclaration; use crate::{Error, Result, Table}; async fn table_with(name: &str, values: Vec) -> Table { @@ -1344,7 +1345,7 @@ mod tests { .await; table .add_columns() - .computed_blob("image_copy", "image") + .computed_column(ComputedColumnDeclaration::blob("image_copy", "image")) .execute() .await .unwrap(); @@ -1463,7 +1464,7 @@ mod tests { let error = table .add_columns() - .computed_blob("invalid", "id + 1") + .computed_column(ComputedColumnDeclaration::blob("invalid", "id + 1")) .execute() .await .unwrap_err();