mirror of
https://github.com/lancedb/lancedb.git
synced 2026-09-01 11:08:55 +00:00
refactor: unify computed column declarations
This commit is contained in:
@@ -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 (
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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
@@ -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))
|
||||
|
||||
Reference in New Issue
Block a user