diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index 8d6993ca6..21ffc8860 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, ComputedColumn as ComputedColumn, Table +from .table import AsyncTable, 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 edbcc5ccc..15cfea29d 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/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] diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 75cd71bdc..683d5f3b0 100644 --- a/python/python/lancedb/remote/table.py +++ b/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)) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 2571fbda7..4a09c35e7 100644 --- a/python/python/lancedb/table.py +++ b/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 diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 38a950999..d9e8030f9 100644 --- a/python/python/tests/test_table.py +++ b/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 diff --git a/python/src/table.rs b/python/src/table.rs index 33aba3792..cbefc4e5b 100644 --- a/python/src/table.rs +++ b/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), +} + /// 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> { 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); diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 619018dab..cd45f7153 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -55,7 +55,7 @@ use crate::{ }; use arrow_array::{LargeBinaryArray, RecordBatch, RecordBatchReader}; use arrow_ipc::reader::{FileReader, StreamReader}; -use arrow_schema::{ArrowError, DataType, SchemaRef}; +use arrow_schema::{ArrowError, DataType, Schema as ArrowSchema, SchemaRef}; use async_trait::async_trait; use chrono::{DateTime, Utc}; use datafusion_common::DataFusionError; @@ -3183,21 +3183,31 @@ impl BaseTable for RemoteTable { self.schema().await?.as_ref(), "schema evolution", )?; - // The server plans the declaration: expression validation, type - // inference and the persisted binding all happen there. - let entries = columns - .iter() - .map(|column| { - let mut entry = serde_json::json!({ - "name": column.name, - "computed": column.expression, - }); - if column.output == crate::table::computed_columns::ComputedColumnOutput::BlobV2 { - entry["computed_output"] = serde_json::json!("blob_v2"); - } - entry - }) - .collect::>(); + // The server plans the declaration: expression validation and the + // persisted binding happen there. An explicit field carries its own + // output schema and extension metadata. + let entries = + columns + .iter() + .map(|column| -> Result { + let mut entry = serde_json::json!({ "computed": column.expression() }); + if let Some(field) = column.field() { + let mut schema = lance_namespace::schema::arrow_schema_to_json( + &ArrowSchema::new(vec![field.clone()]), + )?; + entry["field"] = serde_json::to_value(schema.fields.remove(0)).map_err( + |error| Error::Runtime { + message: format!( + "failed to serialize explicit computed output field: {error}" + ), + }, + )?; + } else { + entry["name"] = serde_json::json!(column.name()); + } + Ok(entry) + }) + .collect::>>()?; let mut body = serde_json::json!({ "new_columns": entries }); self.apply_branch_body(&mut body); let request = self @@ -7394,7 +7404,7 @@ mod tests { assert_eq!(result.version, if old_server { 0 } else { 43 }); } - /// A declaration is sent as `{name, computed}` entries for the server to + /// An inferred declaration is sent as `{name, computed}` for the server to /// plan; the client never types the expression itself. #[tokio::test] async fn test_add_computed_columns_sends_the_expression() { @@ -7426,7 +7436,7 @@ mod tests { } #[tokio::test] - async fn test_add_blob_computed_column_sends_explicit_output_semantics() { + async fn test_add_computed_column_sends_explicit_output_field() { let table = Table::new_with_handler("my_table", |request| match request.url().path() { "/v1/table/my_table/describe/" => simple_describe_response(), "/v1/table/my_table/add_columns/" => { @@ -7435,9 +7445,29 @@ mod tests { assert_eq!( value["new_columns"], serde_json::json!([{ - "name": "image_copy", "computed": "image", - "computed_output": "blob_v2" + "field": { + "metadata": { + "ARROW:extension:name": "lance.blob.v2" + }, + "name": "image_copy", + "nullable": true, + "type": { + "type": "struct", + "fields": [ + { + "name": "data", + "nullable": true, + "type": { "type": "large_binary" } + }, + { + "name": "uri", + "nullable": true, + "type": { "type": "utf8" } + } + ] + } + } }]) ); http::Response::builder() @@ -7450,10 +7480,7 @@ mod tests { let result = table .add_columns() - .computed_column(crate::table::ComputedColumnDeclaration::blob( - "image_copy", - "image", - )) + .computed_field(crate::blob("image_copy", true), "image") .execute() .await .unwrap(); diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 8ba9d5e05..57fe9dfa1 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -95,8 +95,8 @@ pub use cherry_pick::{ }; pub use chrono::Duration; pub use computed_columns::{ - ComputedColumn, ComputedColumnDeclaration, ComputedColumnKind, ComputedColumnOutput, - computed_column_from_field, computed_columns, + ComputedColumn, ComputedColumnDeclaration, ComputedColumnKind, computed_column_from_field, + computed_columns, }; pub use delete::DeleteResult; use futures::future::join_all; @@ -751,8 +751,8 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { /// Declare computed columns, each defined by a SQL expression. /// /// Where the declaration is planned depends on the backend: a local table - /// validates and types the expression itself, a remote one sends the text - /// for the server to plan. + /// validates and types the expression itself, while a remote one sends the + /// expression and any explicit output field for the server to plan. async fn add_computed_columns( &self, _columns: &[computed_columns::ComputedColumnDeclaration], diff --git a/rust/lancedb/src/table/add_columns.rs b/rust/lancedb/src/table/add_columns.rs index 6265e0197..ae3ec7ec3 100644 --- a/rust/lancedb/src/table/add_columns.rs +++ b/rust/lancedb/src/table/add_columns.rs @@ -5,6 +5,7 @@ use std::sync::Arc; +use arrow_schema::Field as ArrowField; use lance::dataset::NewColumnTransform; use super::BaseTable; @@ -87,26 +88,31 @@ impl AddColumnsBuilder { self.computed_column(ComputedColumnDeclaration::inferred(name, expression)) } - /// Add one computed-column declaration with explicit output semantics. + /// Add a computed column whose output schema is the supplied Arrow field. /// - /// 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. + /// Extension semantics such as Blob v2 come from the field metadata. The + /// expression's inferred type must be compatible with the field semantics. /// /// ``` /// # use lancedb::Table; - /// # use lancedb::table::ComputedColumnDeclaration; /// # async fn declare(table: &Table) -> Result<(), Box> { /// table /// .add_columns() - /// .computed_column(ComputedColumnDeclaration::blob("image_copy", "image")) + /// .computed_field(lancedb::blob("image_copy", true), "image") /// .execute() /// .await?; /// table.refresh_column("image_copy").await?; /// # Ok(()) /// # } /// ``` + pub fn computed_field(self, field: ArrowField, expression: impl Into) -> Self { + self.computed_column(ComputedColumnDeclaration::with_field(field, expression)) + } + + /// Add one pre-built computed-column declaration. + /// + /// Declarations share one ordered stream, so a later expression may + /// reference a column declared earlier in the same builder. pub fn computed_column(mut self, declaration: ComputedColumnDeclaration) -> Self { self.computed.push(declaration); self diff --git a/rust/lancedb/src/table/computed_columns.rs b/rust/lancedb/src/table/computed_columns.rs index e6f2d0ae9..d108b8bc2 100644 --- a/rust/lancedb/src/table/computed_columns.rs +++ b/rust/lancedb/src/table/computed_columns.rs @@ -9,12 +9,13 @@ //! refresh fills the rows. //! //! The rule is tagged by kind ([`ComputedColumnKind`]) because kinds differ in -//! where the column's type and inputs come from. A SQL expression is -//! self-describing -- both are derived from the expression, so a caller writes -//! neither -- while a kind resolved through a registry cannot be typed without -//! consulting it. Registered Functions use an exact remote version plus a -//! schema-level Function binding; unknown newer kinds remain readable and fail -//! closed before mutation. +//! where the column's type and inputs come from. A SQL expression determines +//! its inputs and physical result type; a caller may also supply an Arrow field +//! when the output needs extension metadata that expressions do not carry. +//! A kind resolved through a registry cannot be typed without consulting it. +//! Registered Functions use an exact remote version plus a schema-level +//! Function binding; unknown newer kinds remain readable and fail closed +//! before mutation. //! //! [`computed_columns`] and [`computed_column_from_field`] read declarations //! back off a schema. @@ -27,7 +28,9 @@ use datafusion_common::tree_node::TreeNode; use datafusion_physical_plan::PhysicalExpr; use lance::dataset::NewColumnTransform; use lance_arrow::FieldExt; -use lance_core::datatypes::{BLOB_V2_DESC_FIELD, format_field_path_minimal, parse_field_path}; +use lance_core::datatypes::{ + BLOB_V2_DESC_FIELD, BlobV2Layout, format_field_path_minimal, parse_field_path, +}; use lance_datafusion::planner::Planner; use lance_namespace::models::{JsonArrowDataType, JsonArrowField, JsonArrowSchema}; use serde::{Deserialize, Serialize}; @@ -66,43 +69,59 @@ pub const SQL_KIND: &str = "sql"; /// Value of [`KIND_META_KEY`] for a registered Function binding. pub const FUNCTION_KIND: &str = "function"; -/// How a SQL computed declaration chooses its stored output type. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -#[non_exhaustive] -pub enum ComputedColumnOutput { - /// Use the expression's inferred Arrow type. - Inferred, - /// Store a materialized `LargeBinary` result as a Blob v2 column. - BlobV2, +#[derive(Debug, Clone, PartialEq, Eq)] +enum ComputedColumnTarget { + Inferred(String), + Explicit(ArrowField), } /// One SQL computed-column declaration before it is planned. #[derive(Debug, Clone, PartialEq, Eq)] pub struct ComputedColumnDeclaration { - /// Name of the column to declare. - pub name: String, - /// Immutable SQL expression evaluated by refresh. - pub expression: String, - /// Stored semantic type of the result. - pub output: ComputedColumnOutput, + target: ComputedColumnTarget, + expression: String, } impl ComputedColumnDeclaration { /// Declare a computed column whose type is inferred from its expression. pub fn inferred(name: impl Into, expression: impl Into) -> Self { Self { - name: name.into(), + target: ComputedColumnTarget::Inferred(name.into()), expression: expression.into(), - output: ComputedColumnOutput::Inferred, } } - /// Declare a Blob v2 output backed by a `LargeBinary` expression. - pub fn blob(name: impl Into, expression: impl Into) -> Self { + /// Declare a computed column with an explicit Arrow field. + /// + /// The field is the source of truth for the output's name, type, + /// nullability, and extension metadata. Its semantics must accept the + /// expression's inferred type, and it must be nullable because declaration + /// commits the column before refresh fills its values. + pub fn with_field(field: ArrowField, expression: impl Into) -> Self { Self { - name: name.into(), + target: ComputedColumnTarget::Explicit(field), expression: expression.into(), - output: ComputedColumnOutput::BlobV2, + } + } + + /// Name of the column to declare. + pub fn name(&self) -> &str { + match &self.target { + ComputedColumnTarget::Inferred(name) => name, + ComputedColumnTarget::Explicit(field) => field.name(), + } + } + + /// Immutable SQL expression evaluated by refresh. + pub fn expression(&self) -> &str { + &self.expression + } + + /// Explicit output field, or `None` when the expression determines it. + pub fn field(&self) -> Option<&ArrowField> { + match &self.target { + ComputedColumnTarget::Inferred(_) => None, + ComputedColumnTarget::Explicit(field) => Some(field), } } } @@ -1149,15 +1168,20 @@ pub(crate) fn ensure_no_foreign_declarations<'a>( fields: impl IntoIterator>, ) -> Result<()> { for field in fields { - if field.metadata().keys().any(|k| is_declaration_key(k)) { - return Err(Error::InvalidInput { - message: format!( - "field '{}' carries computed-column metadata; declare computed columns \ - with add_columns().computed()", - field.name() - ), - }); - } + ensure_no_foreign_declaration(field)?; + } + Ok(()) +} + +fn ensure_no_foreign_declaration(field: &ArrowField) -> Result<()> { + if field.metadata().keys().any(|k| is_declaration_key(k)) { + return Err(Error::InvalidInput { + message: format!( + "field '{}' carries computed-column metadata; declare computed columns \ + with add_columns().computed()", + field.name() + ), + }); } Ok(()) } @@ -1464,35 +1488,65 @@ fn plan_declarations( let mut fields = Vec::with_capacity(columns.len()); for declaration in columns { - if schema.field_with_name(&declaration.name).is_ok() { + let name = declaration.name(); + if schema.field_with_name(name).is_ok() { return Err(Error::ColumnAlreadyExists { - name: declaration.name.clone(), + name: name.to_string(), }); } - let bound = bind(schema.clone(), &declaration.name, &declaration.expression)?; + let bound = bind(schema.clone(), name, declaration.expression())?; // Declared columns start entirely null, so nullability is a property // of the declaration rather than of what the expression yields. - let metadata = computed_column_metadata(&declaration.expression, &bound.inputs); - let field = match declaration.output { - ComputedColumnOutput::Inferred => { - ArrowField::new(&declaration.name, bound.data_type, true).with_metadata(metadata) - } - ComputedColumnOutput::BlobV2 => { - if bound.data_type != DataType::LargeBinary { - return Err(Error::InvalidExpression { - column: declaration.name.clone(), + let computed_metadata = computed_column_metadata(declaration.expression(), &bound.inputs); + let field = match declaration.field() { + None => ArrowField::new(name, bound.data_type, true).with_metadata(computed_metadata), + Some(field) => { + ensure_no_foreign_declaration(field)?; + if !field.is_nullable() { + return Err(Error::InvalidInput { message: format!( - "a Blob v2 computed output requires a LargeBinary expression, got {}", - bound.data_type + "explicit computed output field '{}' must be nullable", + field.name() ), }); } - let field = crate::blob::blob(&declaration.name, true); - let mut blob_metadata = field.metadata().clone(); - blob_metadata.extend(metadata); - field.with_metadata(blob_metadata) + let expression_type = if field.is_blob_v2() { + match field.data_type() { + DataType::Struct(fields) + if BlobV2Layout::classify(fields) == Some(BlobV2Layout::Logical) => + { + // A Blob field stores descriptors, but refresh + // publishes materialized bytes through Lance's + // Blob conversion path. + DataType::LargeBinary + } + data_type => { + return Err(Error::InvalidInput { + message: format!( + "explicit Blob v2 output field '{}' has unsupported logical type {}", + field.name(), + data_type + ), + }); + } + } + } else { + field.data_type().clone() + }; + if expression_type != bound.data_type { + return Err(Error::InvalidExpression { + column: name.to_string(), + message: format!( + "explicit output field accepts {}, but the expression yields {}", + expression_type, bound.data_type, + ), + }); + } + let mut metadata = field.metadata().clone(); + metadata.extend(computed_metadata); + field.clone().with_metadata(metadata) } }; schema = Arc::new(ArrowSchema::new_with_metadata( @@ -1684,6 +1738,91 @@ mod tests { ); } + #[tokio::test] + async fn test_explicit_field_is_the_output_schema() { + let table = table_with_ints("explicit_field").await; + let field = ArrowField::new("copy", DataType::Int32, true) + .with_metadata(HashMap::from([("semantic".into(), "custom".into())])); + + table + .add_columns() + .computed_field(field, "x") + .execute() + .await + .unwrap(); + + let schema = table.schema().await.unwrap(); + let field = schema.field_with_name("copy").unwrap(); + assert_eq!(field.data_type(), &DataType::Int32); + assert_eq!( + field.metadata().get("semantic").map(String::as_str), + Some("custom") + ); + assert_eq!( + field + .metadata() + .get(COMPUTED_COLUMN_META_KEY) + .map(String::as_str), + Some("true") + ); + } + + #[tokio::test] + async fn test_explicit_field_must_be_nullable() { + let table = table_with_ints("explicit_field_nullable").await; + let error = table + .add_columns() + .computed_field(ArrowField::new("copy", DataType::Int32, false), "x") + .execute() + .await + .unwrap_err(); + + assert!( + matches!(error, Error::InvalidInput { message } if message.contains("must be nullable")) + ); + } + + #[tokio::test] + async fn test_explicit_field_cannot_supply_binding_metadata() { + let table = table_with_ints("explicit_field_binding").await; + let field = + ArrowField::new("copy", DataType::Int32, true).with_metadata(HashMap::from([( + EXPRESSION_META_KEY.to_string(), + "other".to_string(), + )])); + let error = table + .add_columns() + .computed_field(field, "x") + .execute() + .await + .unwrap_err(); + + assert!(matches!( + error, + Error::InvalidInput { message } + if message.contains("carries computed-column metadata") + )); + } + + #[tokio::test] + async fn test_explicit_blob_field_requires_a_logical_blob_layout() { + let table = table_with_ints("explicit_blob_layout").await; + let metadata = crate::blob("copy", true).metadata().clone(); + let field = ArrowField::new("copy", DataType::LargeBinary, true).with_metadata(metadata); + let error = table + .add_columns() + .computed_field(field, "x") + .execute() + .await + .unwrap_err(); + + assert!(matches!( + error, + Error::InvalidInput { message } + if message.contains("unsupported logical type LargeBinary") + )); + } + /// The binding reaches the schema only if `AllNulls` carries per-field /// metadata through the commit. The whole representation rests on it. #[tokio::test] diff --git a/rust/lancedb/src/table/refresh.rs b/rust/lancedb/src/table/refresh.rs index 8272b08d7..a376f0001 100644 --- a/rust/lancedb/src/table/refresh.rs +++ b/rust/lancedb/src/table/refresh.rs @@ -33,9 +33,10 @@ use std::collections::HashSet; use std::sync::Arc; use arrow_array::{ - Array, ArrayRef, BooleanArray, LargeBinaryArray, RecordBatch, RecordBatchOptions, + Array, ArrayRef, BooleanArray, LargeBinaryArray, RecordBatch, RecordBatchOptions, StructArray, + new_null_array, }; -use arrow_schema::Schema as ArrowSchema; +use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; use datafusion_expr::ColumnarValue; use futures::{Stream, StreamExt, TryStreamExt}; use lance::Dataset; @@ -374,7 +375,10 @@ fn configure_blob_inputs( Ok(()) } -fn blob_array_from_binary(array: &ArrayRef) -> lance_core::Result { +fn blob_array_from_binary( + array: &ArrayRef, + target_field: &ArrowField, +) -> lance_core::Result { let values = array .as_any() .downcast_ref::() @@ -392,7 +396,39 @@ fn blob_array_from_binary(array: &ArrayRef) -> lance_core::Result { builder.push_bytes(values.value(index))?; } } - builder.finish() + let minimal = builder.finish()?; + let minimal = minimal + .as_any() + .downcast_ref::() + .ok_or_else(|| lance_core::Error::internal("Blob builder returned a non-struct array"))?; + let DataType::Struct(target_fields) = target_field.data_type() else { + return Err(lance_core::Error::invalid_input(format!( + "Blob v2 output field '{}' has non-struct type {}", + target_field.name(), + target_field.data_type() + ))); + }; + let columns = target_fields + .iter() + .map(|field| match field.name().as_str() { + "data" | "uri" => minimal + .column_by_name(field.name()) + .cloned() + .ok_or_else(|| { + lance_core::Error::internal(format!("Blob builder omitted '{}'", field.name())) + }), + "position" | "size" => Ok(new_null_array(field.data_type(), minimal.len())), + name => Err(lance_core::Error::invalid_input(format!( + "Blob v2 output field '{}' has unsupported logical child '{name}'", + target_field.name() + ))), + }) + .collect::>>()?; + Ok(Arc::new(StructArray::try_new( + target_fields.clone(), + columns, + minimal.nulls().cloned(), + )?)) } /// How many rows of one fragment would gain a value. @@ -495,7 +531,7 @@ async fn fill_stream( let computed = evaluate(&bound, &evaluation_batch(&batch, &bound, Some(&keep))?)?; let merged = arrow_select::zip::zip(&fill, &computed, existing)?; let merged = if output_is_blob { - blob_array_from_binary(&merged)? + blob_array_from_binary(&merged, projected.field(0))? } else { merged }; @@ -507,13 +543,15 @@ async fn fill_stream( mod tests { use std::sync::Arc; - use arrow_array::{Array, Int32Array, LargeBinaryArray, RecordBatch, record_batch}; + use arrow_array::{ + Array, ArrayRef, Int32Array, LargeBinaryArray, RecordBatch, StructArray, record_batch, + }; + use arrow_schema::Field as ArrowField; use futures::TryStreamExt; use lance_core::ROW_ID; 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 { @@ -566,6 +604,25 @@ mod tests { table.add(batch).execute().await.unwrap(); } + #[test] + fn test_blob_output_matches_complete_logical_field() { + let values: ArrayRef = Arc::new(LargeBinaryArray::from(vec![ + Some(b"hello".as_slice()), + None, + ])); + let field = ArrowField::new( + "image", + lance_core::datatypes::BLOB_V2_LOGICAL_TYPE.clone(), + true, + ); + + let output = super::blob_array_from_binary(&values, &field).unwrap(); + assert_eq!(output.data_type(), field.data_type()); + let output = output.as_any().downcast_ref::().unwrap(); + assert_eq!(output.column_by_name("position").unwrap().null_count(), 2); + assert_eq!(output.column_by_name("size").unwrap().null_count(), 2); + } + /// The gate's reproducer: `b = coalesce(a, 0)` refreshed before `a` /// must not bake zeros from `a`'s placeholder null. It is refused, and /// names the input, until `a` is filled -- after every append too. @@ -1345,7 +1402,7 @@ mod tests { .await; table .add_columns() - .computed_column(ComputedColumnDeclaration::blob("image_copy", "image")) + .computed_field(crate::blob("image_copy", true), "image") .execute() .await .unwrap(); @@ -1464,14 +1521,16 @@ mod tests { let error = table .add_columns() - .computed_column(ComputedColumnDeclaration::blob("invalid", "id + 1")) + .computed_field(crate::blob("invalid", true), "id + 1") .execute() .await .unwrap_err(); assert!(matches!( error, Error::InvalidExpression { column, message } - if column == "invalid" && message.contains("requires a LargeBinary expression") + if column == "invalid" + && message.contains("explicit output field accepts LargeBinary") + && message.contains("expression yields Int32") )); }