From 6c8e5e9d9816bba7fda02f10fb53595fd292a6a8 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 28 Aug 2026 16:44:17 +0800 Subject: [PATCH] refactor: inherit blob metadata for computed projections --- python/python/lancedb/_lancedb.pyi | 3 +- python/python/lancedb/remote/table.py | 3 +- python/python/lancedb/table.py | 56 +--- python/python/tests/test_table.py | 42 +-- python/src/table.rs | 27 +- rust/lancedb/src/remote/table.rs | 100 ++------ rust/lancedb/src/table.rs | 12 +- rust/lancedb/src/table/add_columns.rs | 37 +-- rust/lancedb/src/table/computed_columns.rs | 284 ++++++--------------- rust/lancedb/src/table/refresh.rs | 104 ++------ rust/lancedb/src/table/schema_evolution.rs | 2 +- 11 files changed, 158 insertions(+), 512 deletions(-) diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 15cfea29d..7d7ca7f2a 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -352,8 +352,7 @@ class Table: async def delete(self, filter: Union[str, PyExpr]) -> DeleteResult: ... async def add_columns(self, columns: list[tuple[str, str]]) -> AddColumnsResult: ... async def add_computed_columns( - self, - columns: list[tuple[str | pa.Field, str]], + self, columns: list[tuple[str, 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 683d5f3b0..55014a423 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -13,7 +13,6 @@ from typing import ( Iterable, List, Optional, - Sequence, Union, Literal, overload, @@ -984,7 +983,7 @@ class RemoteTable(Table): | FunctionApplication | None = None, *, - computed: Dict[str, str] | Sequence[tuple[str | pa.Field, str]] | None = None, + computed: Dict[str, 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 4a09c35e7..36c5b727d 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -913,26 +913,6 @@ def _normalize_progress(progress): return progress, False -def _normalize_computed_columns( - 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 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 - - class Table(ABC): """ A Table is a collection of Records in a LanceDB Database. @@ -2162,7 +2142,7 @@ class Table(ABC): | pa.Schema | None = None, *, - computed: Dict[str, str] | Sequence[tuple[str | pa.Field, str]] | None = None, + computed: Dict[str, str] | None = None, ): """ Add new columns with defined values. @@ -2184,15 +2164,12 @@ class Table(ABC): atomic binding; aliases come from ``rename(columns=...)``. Function columns are supported only on LanceDB Cloud and Enterprise. - computed: Dict[str, str] or Sequence[Tuple[str | pa.Field, str]], optional + computed: Dict[str, 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. + output field from its expression. A direct projection of a Blob v2 + field inherits Blob v2 semantics; other expressions derive their + ordinary Arrow type. Mapping 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 @@ -4324,7 +4301,7 @@ class LanceTable(Table): | pa.Schema | None = None, *, - computed: Dict[str, str] | Sequence[tuple[str | pa.Field, str]] | None = None, + computed: Dict[str, str] | None = None, ) -> AddColumnsResult: return LOOP.run(self._table.add_columns(transforms, computed=computed)) @@ -6272,7 +6249,7 @@ class AsyncTable: | pa.Schema | None = None, *, - computed: dict[str, str] | Sequence[tuple[str | pa.Field, str]] | None = None, + computed: dict[str, str] | None = None, ) -> AddColumnsResult: """ Add new columns with defined values. @@ -6292,15 +6269,12 @@ class AsyncTable: atomic binding; aliases come from ``rename(columns=...)``. Function columns are supported only on LanceDB Cloud and Enterprise. - computed: Dict[str, str] or Sequence[Tuple[str | pa.Field, str]], optional + computed: Dict[str, 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. + output field from its expression. A direct projection of a Blob v2 + field inherits Blob v2 semantics; other expressions derive their + ordinary Arrow type. Mapping 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 @@ -6358,9 +6332,7 @@ class AsyncTable: raise ValueError( "add_columns cannot take both transforms and computed columns" ) - return await self._inner.add_computed_columns( - _normalize_computed_columns(computed) - ) + return await self._inner.add_computed_columns(list(computed.items())) if transforms is None: raise ValueError("add_columns requires transforms or computed columns") if isinstance(transforms, pa.Schema): diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index d9e8030f9..fbdfac5d8 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -4087,7 +4087,7 @@ def test_computed_column_rejects_transforms_and_computed_together(tmp_path): table.add_columns({"a": "x + 1"}, computed={"b": "x * 2"}) -def test_computed_column_blob_input_and_explicit_output(tmp_path): +def test_computed_column_blob_projection_inherits_semantics(tmp_path): schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) db = lancedb.connect(tmp_path) table = db.create_table("computed_column_blob", schema=schema) @@ -4099,49 +4099,17 @@ def test_computed_column_blob_input_and_explicit_output(tmp_path): ] ) - table.add_columns( - computed=[ - (lancedb.blob("image_copy"), "image"), - ("payload_copy", "image_copy"), - ] - ) + table.add_columns(computed={"image_copy": "image", "second_copy": "image_copy"}) 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] - assert table.blob_columns() == ["image", "image_copy"] + assert table.refresh_column("second_copy").rows_filled == 2 + assert table.blob_columns() == ["image", "image_copy", "second_copy"] hits = table.search().with_row_id(True).limit(10).to_arrow() rows = sorted(zip(hits["id"].to_pylist(), hits["_rowid"].to_pylist())) - copied = table.fetch_blobs("image_copy", [row_id for _, row_id in rows]) + copied = table.fetch_blobs("second_copy", [row_id for _, row_id in rows]) assert copied.to_pylist() == [b"hello", b"", None] -def test_blob_output_declaration_rejects_eager_transforms(tmp_path): - db = lancedb.connect(tmp_path) - table = db.create_table("computed_column_blob_mixed", [{"x": 1}]) - with pytest.raises(ValueError): - table.add_columns( - {"a": "x + 1"}, - computed=[(lancedb.blob("b"), "x")], - ) - - -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="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 async def test_computed_column_async(tmp_path): db = await lancedb.connect_async(tmp_path) diff --git a/python/src/table.rs b/python/src/table.rs index cbefc4e5b..784d29136 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -13,16 +13,15 @@ use crate::{ }; use arrow::{ array::{Array, LargeBinaryArray}, - datatypes::{DataType, Field, Schema}, + datatypes::{DataType, Schema}, ffi_stream::ArrowArrayStreamReader, pyarrow::{FromPyArrow, PyArrowType, ToPyArrow}, }; use lancedb::blob::{BlobFile, BlobRangeRequest}; use lancedb::index::scalar::FtsIndexBuilder; use lancedb::table::{ - AddDataMode, ColumnAlteration, ComputedColumnDeclaration, Duration, FieldMetadataUpdate, - FtsToken as LanceDbFtsToken, NewColumnTransform, OptimizeAction, OptimizeOptions, Ref, - Table as LanceDbTable, + AddDataMode, ColumnAlteration, Duration, FieldMetadataUpdate, FtsToken as LanceDbFtsToken, + NewColumnTransform, OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable, }; use lancedb::tokenize as lancedb_tokenize; use pyo3::{ @@ -101,12 +100,6 @@ 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)] @@ -1584,21 +1577,13 @@ impl Table { pub fn add_computed_columns( self_: PyRef<'_, Self>, - columns: Vec<(ComputedColumnFieldArg, String)>, + columns: Vec<(String, String)>, ) -> PyResult> { let inner = self_.inner_ref()?.clone(); future_into_py(self_.py(), async move { let mut builder = inner.add_columns(); - 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); + for (name, expression) in columns { + builder = builder.computed(name, expression); } 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 cd45f7153..57d2dc47d 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, Schema as ArrowSchema, SchemaRef}; +use arrow_schema::{ArrowError, DataType, SchemaRef}; use async_trait::async_trait; use chrono::{DateTime, Utc}; use datafusion_common::DataFusionError; @@ -3174,40 +3174,24 @@ impl BaseTable for RemoteTable { } } - async fn add_computed_columns( - &self, - columns: &[crate::table::computed_columns::ComputedColumnDeclaration], - ) -> Result { + async fn add_computed_columns(&self, columns: &[(String, String)]) -> Result { self.check_mutable().await?; crate::table::computed_columns::ensure_no_function_bindings_for_mutation( self.schema().await?.as_ref(), "schema evolution", )?; - // 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::>>()?; + // The server plans the declaration against its table schema, including + // Blob v2 semantics inherited by a direct field projection. + let entries = columns + .iter() + .map( + |(name, expression)| lance_namespace::models::AddColumnsEntry { + name: name.clone(), + computed: Some(Some(expression.clone())), + ..Default::default() + }, + ) + .collect::>(); let mut body = serde_json::json!({ "new_columns": entries }); self.apply_branch_body(&mut body); let request = self @@ -7404,8 +7388,8 @@ mod tests { assert_eq!(result.version, if old_server { 0 } else { 43 }); } - /// An inferred declaration is sent as `{name, computed}` for the server to - /// plan; the client never types the expression itself. + /// A 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() { let table = Table::new_with_handler("my_table", |request| match request.url().path() { @@ -7435,58 +7419,6 @@ mod tests { assert_eq!(result.version, 7); } - #[tokio::test] - 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/" => { - let body = request.body().unwrap().as_bytes().unwrap(); - let value: serde_json::Value = serde_json::from_slice(body).unwrap(); - assert_eq!( - value["new_columns"], - serde_json::json!([{ - "computed": "image", - "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() - .status(200) - .body(r#"{"version": 8}"#.to_string()) - .unwrap() - } - path => panic!("Unexpected path: {path}"), - }); - - let result = table - .add_columns() - .computed_field(crate::blob("image_copy", true), "image") - .execute() - .await - .unwrap(); - assert_eq!(result.version, 8); - } - #[tokio::test] async fn test_add_scalar_function_column_sends_atomic_null_declaration() { let table = Table::new_with_handler("my_table", |request| { diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 57fe9dfa1..70b86f3cf 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -95,8 +95,7 @@ pub use cherry_pick::{ }; pub use chrono::Duration; pub use computed_columns::{ - ComputedColumn, ComputedColumnDeclaration, ComputedColumnKind, computed_column_from_field, - computed_columns, + ComputedColumn, ComputedColumnKind, computed_column_from_field, computed_columns, }; pub use delete::DeleteResult; use futures::future::join_all; @@ -752,10 +751,10 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { /// /// Where the declaration is planned depends on the backend: a local table /// validates and types the expression itself, while a remote one sends the - /// expression and any explicit output field for the server to plan. + /// expression for the server to plan. async fn add_computed_columns( &self, - _columns: &[computed_columns::ComputedColumnDeclaration], + _columns: &[(String, String)], ) -> Result { Err(Error::NotSupported { message: "computed columns are not supported on this table type".into(), @@ -3524,10 +3523,7 @@ impl BaseTable for NativeTable { Ok(result) } - async fn add_computed_columns( - &self, - columns: &[computed_columns::ComputedColumnDeclaration], - ) -> Result { + async fn add_computed_columns(&self, columns: &[(String, String)]) -> Result { let result = schema_evolution::execute_declare(self, columns).await?; self.bump_freshness(); Ok(result) diff --git a/rust/lancedb/src/table/add_columns.rs b/rust/lancedb/src/table/add_columns.rs index ae3ec7ec3..1ac0c6b4f 100644 --- a/rust/lancedb/src/table/add_columns.rs +++ b/rust/lancedb/src/table/add_columns.rs @@ -5,11 +5,9 @@ use std::sync::Arc; -use arrow_schema::Field as ArrowField; use lance::dataset::NewColumnTransform; use super::BaseTable; -use super::computed_columns::ComputedColumnDeclaration; use super::schema_evolution::AddColumnsResult; use crate::function::FunctionApplication; use crate::{Error, Result}; @@ -18,7 +16,7 @@ use crate::{Error, Result}; pub struct AddColumnsBuilder { parent: Arc, transform: Option, - computed: Vec, + computed: Vec<(String, String)>, function: Option<(FunctionApplication, Option)>, read_columns: Option>, } @@ -84,37 +82,8 @@ impl AddColumnsBuilder { /// # Ok(()) /// # } /// ``` - pub fn computed(self, name: impl Into, expression: impl Into) -> Self { - self.computed_column(ComputedColumnDeclaration::inferred(name, expression)) - } - - /// Add a computed column whose output schema is the supplied Arrow field. - /// - /// 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; - /// # async fn declare(table: &Table) -> Result<(), Box> { - /// table - /// .add_columns() - /// .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); + pub fn computed(mut self, name: impl Into, expression: impl Into) -> Self { + self.computed.push((name.into(), expression.into())); self } diff --git a/rust/lancedb/src/table/computed_columns.rs b/rust/lancedb/src/table/computed_columns.rs index d108b8bc2..4e8d8211e 100644 --- a/rust/lancedb/src/table/computed_columns.rs +++ b/rust/lancedb/src/table/computed_columns.rs @@ -10,9 +10,10 @@ //! //! The rule is tagged by kind ([`ComputedColumnKind`]) because kinds differ in //! 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. +//! its inputs and physical result type. A direct projection of a Blob v2 field +//! also inherits that field's semantic type while execution continues to use +//! `LargeBinary`. 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. @@ -24,19 +25,19 @@ use std::collections::{BTreeSet, HashMap, HashSet}; use std::sync::Arc; use arrow_schema::{DataType, Field as ArrowField, Fields, Schema as ArrowSchema, SchemaRef}; -use datafusion_common::tree_node::TreeNode; +use datafusion_common::{ScalarValue, tree_node::TreeNode}; +use datafusion_expr::Expr; use datafusion_physical_plan::PhysicalExpr; use lance::dataset::NewColumnTransform; use lance_arrow::FieldExt; -use lance_core::datatypes::{ - BLOB_V2_DESC_FIELD, BlobV2Layout, format_field_path_minimal, parse_field_path, -}; +use lance_core::datatypes::{BLOB_V2_DESC_FIELD, format_field_path_minimal, parse_field_path}; use lance_datafusion::planner::Planner; use lance_namespace::models::{JsonArrowDataType, JsonArrowField, JsonArrowSchema}; use serde::{Deserialize, Serialize}; use serde_json::Value; use crate::function::{FunctionApplication, FunctionBinding}; +use crate::utils::resolve_arrow_field_path; use crate::{Error, Result}; /// Field metadata key marking a column as computed. The value is `"true"`. @@ -69,63 +70,6 @@ pub const SQL_KIND: &str = "sql"; /// Value of [`KIND_META_KEY`] for a registered Function binding. pub const FUNCTION_KIND: &str = "function"; -#[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 { - 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 { - target: ComputedColumnTarget::Inferred(name.into()), - expression: expression.into(), - } - } - - /// 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 { - target: ComputedColumnTarget::Explicit(field), - expression: expression.into(), - } - } - - /// 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), - } - } -} - /// Synthetic result identity used when the entire Function result maps to one /// table column (scalar or struct-as-one-column). pub const WHOLE_RESULT_FIELD: &str = "$value"; @@ -1238,6 +1182,36 @@ pub(crate) struct BoundExpression { pub data_type: DataType, /// Blob v2 leaves the scan must materialize as `LargeBinary`. pub blob_paths: Vec, + /// A directly projected Blob v2 field whose semantics the output inherits. + projected_blob_field: Option, +} + +fn is_direct_field_projection(expr: &Expr) -> bool { + match expr { + Expr::Column(_) => true, + Expr::ScalarFunction(function) + if function.name() == "get_field" && function.args.len() == 2 => + { + is_direct_field_projection(&function.args[0]) + && matches!( + &function.args[1], + Expr::Literal(ScalarValue::Utf8(Some(_)), _) + ) + } + _ => false, + } +} + +fn projected_blob_field(schema: &ArrowSchema, expr: &Expr) -> Result> { + if !is_direct_field_projection(expr) { + return Ok(None); + } + let paths = Planner::column_names_in_expr(expr); + let [path] = paths.as_slice() else { + return Ok(None); + }; + let (_, field) = resolve_arrow_field_path(schema, path)?; + Ok(field.is_blob_v2().then_some(field)) } fn collect_blob_paths(field: &ArrowField, parent: &[String], paths: &mut Vec>) { @@ -1372,6 +1346,7 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result< let parsed = planner .parse_expr(expression) .map_err(|e| invalid(e.to_string()))?; + let projected_blob_field = projected_blob_field(schema.as_ref(), &parsed)?; // A declaration is evaluated more than once -- staging and writing are // separate passes, and a refresh years later replays the same text -- so @@ -1459,6 +1434,7 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result< format_field_path_minimal(&segments) }) .collect(), + projected_blob_field, }) } @@ -1474,10 +1450,7 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result< /// batch may declare `a` and then `b = a + 1` in one commit. Refresh order /// then matters, and refresh enforces it: `b` is refused while `a` still has /// unfilled rows. -fn plan_declarations( - schema: SchemaRef, - columns: &[ComputedColumnDeclaration], -) -> Result> { +fn plan_declarations(schema: SchemaRef, columns: &[(String, String)]) -> Result> { if columns.is_empty() { return Err(Error::InvalidInput { message: "at least one computed column is required".into(), @@ -1487,67 +1460,29 @@ fn plan_declarations( let mut schema = schema; let mut fields = Vec::with_capacity(columns.len()); - for declaration in columns { - let name = declaration.name(); + for (name, expression) in columns { if schema.field_with_name(name).is_ok() { return Err(Error::ColumnAlreadyExists { name: name.to_string(), }); } - let bound = bind(schema.clone(), name, declaration.expression())?; + let bound = bind(schema.clone(), name, expression)?; // Declared columns start entirely null, so nullability is a property // of the declaration rather than of what the expression yields. - 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!( - "explicit computed output field '{}' must be nullable", - field.name() - ), - }); - } - 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(); + let computed_metadata = computed_column_metadata(expression, &bound.inputs); + let field = match bound.projected_blob_field { + Some(source) => { + let mut metadata = source.metadata().clone(); + metadata.retain(|key, _| !is_declaration_key(key)); metadata.extend(computed_metadata); - field.clone().with_metadata(metadata) + source + .with_name(name) + .with_nullable(true) + .with_metadata(metadata) } + None => ArrowField::new(name, bound.data_type, true).with_metadata(computed_metadata), }; schema = Arc::new(ArrowSchema::new_with_metadata( schema @@ -1565,13 +1500,7 @@ fn plan_declarations( } pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result> { - let declarations = columns - .iter() - .map(|(name, expression)| { - ComputedColumnDeclaration::inferred(name.clone(), expression.clone()) - }) - .collect::>(); - plan_declarations(schema, &declarations) + plan_declarations(schema, columns) } /// Run the schema-level checks of @@ -1610,7 +1539,7 @@ pub fn validate_declarations(schema: SchemaRef, columns: &[(String, String)]) -> /// public way in. pub(crate) fn declare( schema: SchemaRef, - columns: &[ComputedColumnDeclaration], + columns: &[(String, String)], ) -> Result { let fields = plan_declarations(schema, columns)?; Ok(NewColumnTransform::AllNulls(Arc::new(ArrowSchema::new( @@ -1738,89 +1667,42 @@ 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())])); + #[test] + fn test_direct_blob_projection_inherits_semantics() { + let schema = Arc::new(ArrowSchema::new(vec![crate::blob("image", false)])); + let fields = plan( + schema, + &[ + ("first".to_string(), "image".to_string()), + ("second".to_string(), "first".to_string()), + ], + ) + .unwrap(); - 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); + for field in &fields { + assert!(field.is_blob_v2()); + assert!(field.is_nullable()); + } assert_eq!( - field.metadata().get("semantic").map(String::as_str), - Some("custom") - ); - assert_eq!( - field + fields[1] .metadata() - .get(COMPUTED_COLUMN_META_KEY) + .get(EXPRESSION_META_KEY) .map(String::as_str), - Some("true") + Some("first") ); } - #[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(); + #[test] + fn test_blob_expression_transformation_does_not_inherit_semantics() { + let schema = Arc::new(ArrowSchema::new(vec![crate::blob("image", true)])); + let fields = plan( + schema, + &[("payload".to_string(), "coalesce(image, image)".to_string())], + ) + .unwrap(); - 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") - )); + assert!(!fields[0].is_blob_v2()); + assert_eq!(fields[0].data_type(), &DataType::LargeBinary); } /// The binding reaches the schema only if `AllNulls` carries per-field diff --git a/rust/lancedb/src/table/refresh.rs b/rust/lancedb/src/table/refresh.rs index a376f0001..511fce8ff 100644 --- a/rust/lancedb/src/table/refresh.rs +++ b/rust/lancedb/src/table/refresh.rs @@ -1338,51 +1338,7 @@ mod tests { } #[tokio::test] - async fn test_refresh_materializes_top_level_blob_input() { - let tmp = tempfile::tempdir().unwrap(); - let table = create_blob_table( - tmp.path(), - blob_batch(vec![1, 2, 3], vec![Some(b"hello"), Some(b""), None]), - ) - .await; - table - .add_columns() - .computed("payload_copy", "image") - .execute() - .await - .unwrap(); - - let result = table.refresh_column("payload_copy").await.unwrap(); - assert_eq!(result.rows_filled, 2); - let batches = table - .query() - .select(Select::columns(&["payload_copy"])) - .execute() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - let payloads = batches[0] - .column_by_name("payload_copy") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(payloads.value(0), b"hello"); - assert_eq!(payloads.value(1), b""); - assert!(payloads.is_null(2)); - assert_eq!( - table - .count_rows(Some("payload_copy IS NULL".to_string())) - .await - .unwrap(), - 1 - ); - } - - #[tokio::test] - async fn test_refresh_publishes_explicit_blob_output() { + async fn test_refresh_inherits_and_publishes_blob_output() { use arrow_array::UInt64Array; use lance_arrow::{ BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY, BLOB_INLINE_SIZE_THRESHOLD_META_KEY, @@ -1402,7 +1358,7 @@ mod tests { .await; table .add_columns() - .computed_field(crate::blob("image_copy", true), "image") + .computed("image_copy", "image") .execute() .await .unwrap(); @@ -1515,28 +1471,8 @@ mod tests { } #[tokio::test] - async fn test_explicit_blob_output_requires_large_binary_expression() { - let tmp = tempfile::tempdir().unwrap(); - let table = create_blob_table(tmp.path(), blob_batch(vec![1], vec![Some(b"hello")])).await; - - let error = table - .add_columns() - .computed_field(crate::blob("invalid", true), "id + 1") - .execute() - .await - .unwrap_err(); - assert!(matches!( - error, - Error::InvalidExpression { column, message } - if column == "invalid" - && message.contains("explicit output field accepts LargeBinary") - && message.contains("expression yields Int32") - )); - } - - #[tokio::test] - async fn test_refresh_materializes_nested_struct_blob_input() { - use arrow_array::{Int32Array, StructArray}; + async fn test_refresh_inherits_nested_struct_blob_input() { + use arrow_array::{Int32Array, StructArray, UInt64Array}; use arrow_schema::{DataType, Field, Fields, Schema}; let tmp = tempfile::tempdir().unwrap(); @@ -1574,21 +1510,27 @@ mod tests { .rows_filled, 1 ); + assert_eq!( + table.blob_columns().await.unwrap(), + vec!["metadata.image".to_string(), "payload_copy".to_string()] + ); let batches = table .query() - .select(Select::columns(&["payload_copy"])) + .with_row_id() .execute() .await .unwrap() .try_collect::>() .await .unwrap(); - let payloads = batches[0] - .column_by_name("payload_copy") + let row_ids = batches[0] + .column_by_name(ROW_ID) .unwrap() .as_any() - .downcast_ref::() - .unwrap(); + .downcast_ref::() + .unwrap() + .values(); + let payloads = table.fetch_blobs("payload_copy", row_ids).await.unwrap(); assert_eq!(payloads.value(0), b"nested"); assert!(payloads.is_null(1)); } @@ -1655,8 +1597,8 @@ mod tests { } #[tokio::test] - async fn test_refresh_materializes_external_blob_input() { - use arrow_array::{Int32Array, StringArray}; + async fn test_refresh_inherits_external_blob_input() { + use arrow_array::{Int32Array, StringArray, UInt64Array}; use arrow_schema::{DataType, Field, Schema}; let tmp = tempfile::tempdir().unwrap(); @@ -1713,19 +1655,21 @@ mod tests { ); let batches = table .query() - .select(Select::columns(&["payload_copy"])) + .with_row_id() .execute() .await .unwrap() .try_collect::>() .await .unwrap(); - let payloads = batches[0] - .column_by_name("payload_copy") + let row_ids = batches[0] + .column_by_name(ROW_ID) .unwrap() .as_any() - .downcast_ref::() - .unwrap(); + .downcast_ref::() + .unwrap() + .values(); + let payloads = table.fetch_blobs("payload_copy", row_ids).await.unwrap(); assert_eq!(payloads.value(0), payload); } } diff --git a/rust/lancedb/src/table/schema_evolution.rs b/rust/lancedb/src/table/schema_evolution.rs index ae7583a31..4f8dc811a 100644 --- a/rust/lancedb/src/table/schema_evolution.rs +++ b/rust/lancedb/src/table/schema_evolution.rs @@ -124,7 +124,7 @@ pub(crate) async fn execute_add_columns( /// declaration metadata. pub(crate) async fn execute_declare( table: &NativeTable, - columns: &[computed_columns::ComputedColumnDeclaration], + columns: &[(String, String)], ) -> Result { use lance::dataset::mem_wal::DatasetMemWalExt;