refactor: carry explicit computed output fields

python/python/lancedb/__init__.py

@@ -38,7 +38,7 @@ from .materialized_view import (
     MaterializedView,
     MaterializedViewDefinition,
 )
-from .table import AsyncTable, ComputedColumn as ComputedColumn, Table
+from .table import AsyncTable, Table
 from .types import BaseTokenizerType
 from ._lancedb import Session
 from .namespace import (

python/python/lancedb/_lancedb.pyi

@@ -353,7 +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, Literal["inferred", "blob_v2"]]],
+        columns: list[tuple[str | pa.Field, str]],
     ) -> AddColumnsResult: ...
     async def add_function_columns(
         self, application_json: str, output_name: Optional[str]

python/python/lancedb/remote/table.py

@@ -13,6 +13,7 @@ from typing import (
     Iterable,
     List,
     Optional,
+    Sequence,
     Union,
     Literal,
     overload,
@@ -71,7 +72,6 @@ from ..table import (
     AsyncTable,
     BlobMode,
     Branches,
-    ComputedColumn,
     IndexStatistics,
     Query,
     Table,
@@ -984,7 +984,7 @@ class RemoteTable(Table):
         | FunctionApplication
         | None = None,
         *,
-        computed: Dict[str, str | ComputedColumn] | None = None,
+        computed: Dict[str, str] | Sequence[tuple[str | pa.Field, str]] | None = None,
     ) -> AddColumnsResult:
         return LOOP.run(self._table.add_columns(transforms, computed=computed))
 

python/python/lancedb/table.py

@@ -913,44 +913,23 @@ 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:
+    computed: Dict[str, str] | Sequence[tuple[str | pa.Field, str]],
+) -> list[tuple[str | pa.Field, str]]:
+    columns: list[tuple[str | pa.Field, str]] = []
+    declarations = computed.items() if isinstance(computed, dict) else computed
+    for declaration in declarations:
+        if not isinstance(declaration, (tuple, list)) or len(declaration) != 2:
             raise TypeError(
-                "computed values must be SQL expression strings or "
-                "ComputedColumn values"
+                "computed sequences must contain (column name or pyarrow Field, "
+                "SQL expression) pairs"
             )
+        field, expression = declaration
+        if not isinstance(field, (str, pa.Field)):
+            raise TypeError("computed targets must be column names or pyarrow Fields")
+        if not isinstance(expression, str):
+            raise TypeError("computed values must be SQL expression strings")
+        columns.append((field, expression))
     return columns
 
 
@@ -2183,7 +2162,7 @@ class Table(ABC):
         | pa.Schema
         | None = None,
         *,
-        computed: Dict[str, str | ComputedColumn] | None = None,
+        computed: Dict[str, str] | Sequence[tuple[str | pa.Field, str]] | None = None,
     ):
         """
         Add new columns with defined values.
@@ -2205,13 +2184,15 @@ class Table(ABC):
             atomic binding; aliases come from ``rename(columns=...)``.
             Function columns are supported only on LanceDB Cloud and
             Enterprise.
-        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. 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.
+        computed: Dict[str, str] or Sequence[Tuple[str | pa.Field, str]], optional
+            A mapping from output column names to SQL expressions derives each
+            output field from its expression. An ordered sequence may instead
+            use a pyarrow Field as a target, supplying its name, type,
+            nullability, and extension metadata; use
+            ``(lancedb.blob("name"), expression)`` for a Blob v2 output.
+            Explicit fields must be nullable and their expression result type
+            must be compatible. Mapping or sequence order is declaration and
+            dependency order.
 
             Unlike ``transforms``, the expression is stored rather than
             evaluated now: the column is committed with no values, and rows get
@@ -4343,7 +4324,7 @@ class LanceTable(Table):
         | pa.Schema
         | None = None,
         *,
-        computed: Dict[str, str | ComputedColumn] | None = None,
+        computed: Dict[str, str] | Sequence[tuple[str | pa.Field, str]] | None = None,
     ) -> AddColumnsResult:
         return LOOP.run(self._table.add_columns(transforms, computed=computed))
 
@@ -6291,7 +6272,7 @@ class AsyncTable:
         | pa.Schema
         | None = None,
         *,
-        computed: dict[str, str | ComputedColumn] | None = None,
+        computed: dict[str, str] | Sequence[tuple[str | pa.Field, str]] | None = None,
     ) -> AddColumnsResult:
         """
         Add new columns with defined values.
@@ -6311,11 +6292,14 @@ class AsyncTable:
             atomic binding; aliases come from ``rename(columns=...)``.
             Function columns are supported only on LanceDB Cloud and
             Enterprise.
-        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. Use
-            ``ComputedColumn.blob(expression)`` to publish a ``LargeBinary``
-            result as Blob v2. Mapping insertion order is the declaration and
+        computed: Dict[str, str] or Sequence[Tuple[str | pa.Field, str]], optional
+            A mapping from output column names to SQL expressions derives each
+            output field from its expression. An ordered sequence may instead
+            use a pyarrow Field as a target, supplying its name, type,
+            nullability, and extension metadata; use
+            ``(lancedb.blob("name"), expression)`` for a Blob v2 output.
+            Explicit fields must be nullable and their expression result type
+            must be compatible. Mapping or sequence order is declaration and
             dependency order.
 
             Unlike ``transforms``, the expression is stored rather than

python/python/tests/test_table.py

@@ -4100,10 +4100,10 @@ def test_computed_column_blob_input_and_explicit_output(tmp_path):
     )
 
     table.add_columns(
-        computed={
-            "image_copy": lancedb.ComputedColumn.blob("image"),
-            "payload_copy": "image_copy",
-        }
+        computed=[
+            (lancedb.blob("image_copy"), "image"),
+            ("payload_copy", "image_copy"),
+        ]
     )
     assert table.refresh_column("image_copy").rows_filled == 2
     assert table.refresh_column("payload_copy").rows_filled == 2
@@ -4124,18 +4124,22 @@ def test_blob_output_declaration_rejects_eager_transforms(tmp_path):
     with pytest.raises(ValueError):
         table.add_columns(
             {"a": "x + 1"},
-            computed={"b": lancedb.ComputedColumn.blob("x")},
+            computed=[(lancedb.blob("b"), "x")],
         )
 
 
-def test_computed_column_validates_explicit_output():
-    assert lancedb.ComputedColumn("x + 1").output == "inferred"
-    assert lancedb.ComputedColumn.blob("image").output == "blob_v2"
+def test_computed_column_validates_declaration_mapping(tmp_path):
+    db = lancedb.connect(tmp_path)
+    table = db.create_table("computed_column_mapping", [{"x": 1}])
 
-    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]
+    with pytest.raises(
+        TypeError, match="targets must be column names or pyarrow Fields"
+    ):
+        table.add_columns(computed={42: "x"})  # type: ignore[dict-item]
+    with pytest.raises(TypeError, match="values must be SQL expression strings"):
+        table.add_columns(computed={"copy": 42})  # type: ignore[dict-item]
+    with pytest.raises(TypeError, match="sequences must contain"):
+        table.add_columns(computed=[("copy", "x", "extra")])  # type: ignore[list-item]
 
 
 @pytest.mark.asyncio

python/src/table.rs

@@ -13,7 +13,7 @@ use crate::{
 };
 use arrow::{
     array::{Array, LargeBinaryArray},
-    datatypes::{DataType, Schema},
+    datatypes::{DataType, Field, Schema},
     ffi_stream::ArrowArrayStreamReader,
     pyarrow::{FromPyArrow, PyArrowType, ToPyArrow},
 };
@@ -101,6 +101,12 @@ enum PredicateArg {
     Sql(String),
 }
 
+#[derive(FromPyObject)]
+pub enum ComputedColumnFieldArg {
+    Name(String),
+    Field(PyArrowType<Field>),
+}
+
 /// Statistics about a compaction operation.
 #[pyclass(get_all, from_py_object)]
 #[derive(Clone, Debug)]
@@ -1578,19 +1584,18 @@ impl Table {
 
     pub fn add_computed_columns(
         self_: PyRef<'_, Self>,
-        columns: Vec<(String, String, String)>,
+        columns: Vec<(ComputedColumnFieldArg, 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, 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}'"
-                        )));
+            for (field, expression) in columns {
+                let declaration = match field {
+                    ComputedColumnFieldArg::Name(name) => {
+                        ComputedColumnDeclaration::inferred(name, expression)
+                    }
+                    ComputedColumnFieldArg::Field(PyArrowType(field)) => {
+                        ComputedColumnDeclaration::with_field(field, expression)
                     }
                 };
                 builder = builder.computed_column(declaration);