diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 7d7ca7f2a..6c2800ed1 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -352,7 +352,9 @@ 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, str]] + self, + columns: list[tuple[str, str]], + blob_columns: Optional[list[tuple[str, str]]] = None, ) -> 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 02748b9bc..7f30e4d7c 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -976,8 +976,13 @@ class RemoteTable(Table): | None = None, *, computed: Dict[str, str] | None = None, + computed_blobs: Dict[str, str] | None = None, ) -> AddColumnsResult: - return LOOP.run(self._table.add_columns(transforms, computed=computed)) + return LOOP.run( + self._table.add_columns( + transforms, computed=computed, computed_blobs=computed_blobs + ) + ) def refresh_column(self, column: str): return LOOP.run(self._table.refresh_column(column)) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 535ed7c0d..09166519a 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -2143,6 +2143,7 @@ class Table(ABC): | None = None, *, computed: Dict[str, str] | None = None, + computed_blobs: Dict[str, str] | None = None, ): """ Add new columns with defined values. @@ -2184,6 +2185,12 @@ 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 ------- @@ -4300,8 +4307,13 @@ class LanceTable(Table): | None = None, *, computed: Dict[str, str] | None = None, + computed_blobs: Dict[str, str] | None = None, ) -> AddColumnsResult: - return LOOP.run(self._table.add_columns(transforms, computed=computed)) + return LOOP.run( + self._table.add_columns( + transforms, computed=computed, computed_blobs=computed_blobs + ) + ) def refresh_column(self, column: str) -> "RefreshColumnResult": """Fill a computed column's unfilled rows. See @@ -6248,6 +6260,7 @@ class AsyncTable: | None = None, *, computed: dict[str, str] | None = None, + computed_blobs: dict[str, str] | None = None, ) -> AddColumnsResult: """ Add new columns with defined values. @@ -6283,6 +6296,11 @@ 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 ------- @@ -6306,7 +6324,7 @@ class AsyncTable: function_output_name, function_application = next(iter(transforms.items())) if function_application is not None: - if computed: + if computed or computed_blobs: raise ValueError( "add_columns cannot mix a Function application with SQL " "computed columns" @@ -6322,12 +6340,15 @@ class AsyncTable: {isinstance(f, pa.Field) for f in transforms} ): transforms = pa.schema(transforms) - if computed: + if computed or computed_blobs: if transforms: raise ValueError( "add_columns cannot take both transforms and computed columns" ) - return await self._inner.add_computed_columns(list(computed.items())) + return await self._inner.add_computed_columns( + list((computed or {}).items()), + list((computed_blobs or {}).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 0be9e139d..4cac132d0 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -4087,6 +4087,42 @@ def test_computed_column_rejects_transforms_and_computed_together(tmp_path): table.add_columns({"a": "x + 1"}, computed={"b": "x * 2"}) +def test_computed_blob_input_and_explicit_output(tmp_path): + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) + db = lancedb.connect(tmp_path) + table = db.create_table("computed_blob", schema=schema) + table.add( + [ + {"id": 1, "image": b"hello"}, + {"id": 2, "image": b""}, + {"id": 3, "image": None}, + ] + ) + + table.add_columns( + computed={"payload_copy": "image"}, + computed_blobs={"image_copy": "image"}, + ) + assert table.refresh_column("payload_copy").rows_filled == 2 + assert table.refresh_column("image_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"] + + 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]) + assert copied.to_pylist() == [b"hello", b"", None] + + +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"}) + + @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 784d29136..d50d537d2 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -1575,9 +1575,11 @@ impl Table { }) } + #[pyo3(signature = (columns, blob_columns=None))] pub fn add_computed_columns( self_: PyRef<'_, Self>, columns: Vec<(String, String)>, + blob_columns: Option>, ) -> PyResult> { let inner = self_.inner_ref()?.clone(); future_into_py(self_.py(), async move { @@ -1585,6 +1587,9 @@ impl Table { 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); + } 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 1afc2615a..cc53e331d 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -3174,7 +3174,10 @@ impl BaseTable for RemoteTable { } } - async fn add_computed_columns(&self, columns: &[(String, String)]) -> Result { + async fn add_computed_columns( + &self, + columns: &[crate::table::computed_columns::ComputedColumnDeclaration], + ) -> Result { self.check_mutable().await?; crate::table::computed_columns::ensure_no_function_bindings_for_mutation( self.schema().await?.as_ref(), @@ -3184,13 +3187,16 @@ impl BaseTable for RemoteTable { // inference and the persisted binding all happen there. let entries = columns .iter() - .map( - |(name, expression)| lance_namespace::models::AddColumnsEntry { - name: name.clone(), - computed: Some(Some(expression.clone())), - ..Default::default() - }, - ) + .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::>(); let mut body = serde_json::json!({ "new_columns": entries }); self.apply_branch_body(&mut body); @@ -7419,6 +7425,38 @@ mod tests { assert_eq!(result.version, 7); } + #[tokio::test] + async fn test_add_blob_computed_column_sends_explicit_output_semantics() { + 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!([{ + "name": "image_copy", + "computed": "image", + "computed_output": "blob_v2" + }]) + ); + http::Response::builder() + .status(200) + .body(r#"{"version": 8}"#.to_string()) + .unwrap() + } + path => panic!("Unexpected path: {path}"), + }); + + let result = table + .add_columns() + .computed_blob("image_copy", "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 8436657ca..96e041320 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -754,7 +754,7 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { /// for the server to plan. async fn add_computed_columns( &self, - _columns: &[(String, String)], + _columns: &[computed_columns::ComputedColumnDeclaration], ) -> Result { Err(Error::NotSupported { message: "computed columns are not supported on this table type".into(), @@ -3523,7 +3523,10 @@ impl BaseTable for NativeTable { Ok(result) } - async fn add_computed_columns(&self, columns: &[(String, String)]) -> Result { + async fn add_computed_columns( + &self, + columns: &[computed_columns::ComputedColumnDeclaration], + ) -> 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 1ac0c6b4f..0d8aab73f 100644 --- a/rust/lancedb/src/table/add_columns.rs +++ b/rust/lancedb/src/table/add_columns.rs @@ -8,6 +8,7 @@ use std::sync::Arc; 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}; @@ -16,7 +17,7 @@ use crate::{Error, Result}; pub struct AddColumnsBuilder { parent: Arc, transform: Option, - computed: Vec<(String, String)>, + computed: Vec, function: Option<(FunctionApplication, Option)>, read_columns: Option>, } @@ -83,7 +84,34 @@ impl AddColumnsBuilder { /// # } /// ``` pub fn computed(mut self, name: impl Into, expression: impl Into) -> Self { - self.computed.push((name.into(), expression.into())); + self.computed + .push(ComputedColumnDeclaration::inferred(name, expression)); + self + } + + /// Add a Blob v2 column defined by a `LargeBinary` expression and filled + /// by a later refresh. + /// + /// Blob inputs in the expression are materialized as their payload bytes. + /// The expression result is wrapped back into the Blob v2 logical type + /// before publication, so queries and [`Table::blob_columns`](super::Table::blob_columns) + /// continue to recognize the output as a Blob column. + /// + /// ``` + /// # use lancedb::Table; + /// # async fn declare(table: &Table) -> Result<(), Box> { + /// table + /// .add_columns() + /// .computed_blob("image_copy", "image") + /// .execute() + /// .await?; + /// table.refresh_column("image_copy").await?; + /// # Ok(()) + /// # } + /// ``` + pub fn computed_blob(mut self, name: impl Into, expression: impl Into) -> Self { + self.computed + .push(ComputedColumnDeclaration::blob(name, expression)); self } diff --git a/rust/lancedb/src/table/computed_columns.rs b/rust/lancedb/src/table/computed_columns.rs index 0f89ca612..e6f2d0ae9 100644 --- a/rust/lancedb/src/table/computed_columns.rs +++ b/rust/lancedb/src/table/computed_columns.rs @@ -19,13 +19,15 @@ //! [`computed_columns`] and [`computed_column_from_field`] read declarations //! back off a schema. -use std::collections::{BTreeSet, HashMap}; +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_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_datafusion::planner::Planner; use lance_namespace::models::{JsonArrowDataType, JsonArrowField, JsonArrowSchema}; use serde::{Deserialize, Serialize}; @@ -64,6 +66,47 @@ 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, +} + +/// 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, +} + +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(), + 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 { + Self { + name: name.into(), + expression: expression.into(), + output: ComputedColumnOutput::BlobV2, + } + } +} + /// 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"; @@ -1162,15 +1205,124 @@ pub(crate) struct BoundExpression { /// The columns the expression names, as written; nested inputs keep /// their dotted path. pub inputs: Vec, - /// The top-level columns evaluation reads, in [`Self::read_schema`] - /// order. A nested input appears through its root. + /// The top-level columns evaluation reads, in physical-expression order. + /// A nested input appears through its root. pub roots: Vec, - /// The projected schema evaluation runs against. - pub read_schema: SchemaRef, /// The compiled expression. pub physical: Arc, /// The type the expression yields. pub data_type: DataType, + /// Blob v2 leaves the scan must materialize as `LargeBinary`. + pub blob_paths: Vec, +} + +fn collect_blob_paths(field: &ArrowField, parent: &[String], paths: &mut Vec>) { + let mut path = parent.to_vec(); + path.push(field.name().clone()); + if field.is_blob_v2() { + paths.push(path); + return; + } + match field.data_type() { + DataType::Struct(children) => { + for child in children { + collect_blob_paths(child, &path, paths); + } + } + DataType::List(child) + | DataType::LargeList(child) + | DataType::FixedSizeList(child, _) + | DataType::Map(child, _) => collect_blob_paths(child, &path, paths), + _ => {} + } +} + +fn schema_blob_paths(schema: &ArrowSchema) -> Vec> { + let mut paths = Vec::new(); + for field in schema.fields() { + collect_blob_paths(field, &[], &mut paths); + } + paths +} + +fn transform_blob_field( + field: &ArrowField, + parent: &[String], + materialized: &HashSet>, +) -> ArrowField { + let mut path = parent.to_vec(); + path.push(field.name().clone()); + if field.is_blob_v2() { + if materialized.contains(&path) { + return ArrowField::new(field.name(), DataType::LargeBinary, field.is_nullable()); + } + return ArrowField::new( + field.name(), + BLOB_V2_DESC_FIELD.data_type().clone(), + field.is_nullable(), + ) + .with_metadata(BLOB_V2_DESC_FIELD.metadata().clone()); + } + + let data_type = match field.data_type() { + DataType::Struct(children) => DataType::Struct( + children + .iter() + .map(|child| Arc::new(transform_blob_field(child, &path, materialized))) + .collect(), + ), + DataType::List(child) => { + DataType::List(Arc::new(transform_blob_field(child, &path, materialized))) + } + DataType::LargeList(child) => { + DataType::LargeList(Arc::new(transform_blob_field(child, &path, materialized))) + } + DataType::FixedSizeList(child, size) => DataType::FixedSizeList( + Arc::new(transform_blob_field(child, &path, materialized)), + *size, + ), + DataType::Map(child, sorted) => DataType::Map( + Arc::new(transform_blob_field(child, &path, materialized)), + *sorted, + ), + _ => return field.clone(), + }; + ArrowField::new(field.name(), data_type, field.is_nullable()) + .with_metadata(field.metadata().clone()) +} + +fn blob_runtime_schema(schema: &ArrowSchema, materialized: &HashSet>) -> SchemaRef { + Arc::new(ArrowSchema::new_with_metadata( + schema + .fields() + .iter() + .map(|field| Arc::new(transform_blob_field(field, &[], materialized))) + .collect::(), + schema.metadata().clone(), + )) +} + +fn referenced_blob_paths(schema: &ArrowSchema, inputs: &[String]) -> Result>> { + let input_paths = inputs + .iter() + .map(|input| { + parse_field_path(input).map_err(|error| Error::InvalidInput { + message: format!("invalid computed-column input path '{input}': {error}"), + }) + }) + .collect::>>()?; + Ok(schema_blob_paths(schema) + .into_iter() + .filter(|blob_path| { + input_paths.iter().any(|input_path| { + input_path.len() <= blob_path.len() + && input_path + .iter() + .zip(blob_path) + .all(|(input, blob)| input == blob) + }) + }) + .collect()) } /// Parse, resolve and compile `expression` against `schema`. @@ -1185,7 +1337,14 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result< message, }; - let planner = Planner::new(schema.clone()); + // Blob v2 is a semantic type whose runtime expression ABI is + // `LargeBinary`. Parse against that ABI first so a direct Blob reference + // is not mistaken for its storage descriptor struct. + let all_blob_paths = schema_blob_paths(schema.as_ref()) + .into_iter() + .collect::>(); + let parsing_schema = blob_runtime_schema(schema.as_ref(), &all_blob_paths); + let planner = Planner::new(parsing_schema); let parsed = planner .parse_expr(expression) .map_err(|e| invalid(e.to_string()))?; @@ -1218,13 +1377,19 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result< inputs.sort(); inputs.dedup(); + let blob_paths = referenced_blob_paths(schema.as_ref(), &inputs)?; + let runtime_schema = blob_runtime_schema( + schema.as_ref(), + &blob_paths.iter().cloned().collect::>(), + ); + // A nested input is recorded by its path but read through its root // column; Schema::index_of resolves top-level names only. Resolved here // rather than left to the planner so an unknown column names itself in // the error instead of surfacing as a plan failure. let mut indices = Vec::with_capacity(inputs.len()); for input in &inputs { - let index = schema + let index = runtime_schema .index_of(root(input)) .map_err(|_| invalid(format!("unknown column '{input}'")))?; if !indices.contains(&index) { @@ -1237,7 +1402,7 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result< // compiles the expression has to be built on the projected schema // evaluation will actually read. let read_schema = Arc::new( - schema + runtime_schema .project(&indices) .map_err(|e| invalid(e.to_string()))?, ); @@ -1247,7 +1412,8 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result< .map(|field| field.name().clone()) .collect(); - let optimized = planner + let runtime_planner = Planner::new(runtime_schema); + let optimized = runtime_planner .optimize_expr(parsed) .map_err(|e| invalid(e.to_string()))?; let physical = Planner::new(read_schema.clone()) @@ -1260,9 +1426,15 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result< Ok(BoundExpression { inputs, roots, - read_schema, physical, data_type, + blob_paths: blob_paths + .iter() + .map(|path| { + let segments = path.iter().map(String::as_str).collect::>(); + format_field_path_minimal(&segments) + }) + .collect(), }) } @@ -1278,7 +1450,10 @@ 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. -pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result> { +fn plan_declarations( + schema: SchemaRef, + columns: &[ComputedColumnDeclaration], +) -> Result> { if columns.is_empty() { return Err(Error::InvalidInput { message: "at least one computed column is required".into(), @@ -1288,17 +1463,38 @@ pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result { + 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(), + message: format!( + "a Blob v2 computed output requires a LargeBinary expression, got {}", + bound.data_type + ), + }); + } + let field = crate::blob::blob(&declaration.name, true); + let mut blob_metadata = field.metadata().clone(); + blob_metadata.extend(metadata); + field.with_metadata(blob_metadata) + } + }; schema = Arc::new(ArrowSchema::new_with_metadata( schema .fields() @@ -1314,6 +1510,16 @@ pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result Result> { + let declarations = columns + .iter() + .map(|(name, expression)| { + ComputedColumnDeclaration::inferred(name.clone(), expression.clone()) + }) + .collect::>(); + plan_declarations(schema, &declarations) +} + /// Run the schema-level checks of /// [`AddColumnsBuilder::computed`](super::AddColumnsBuilder::computed) against /// `schema` without committing: the Function-binding guard and the planning of @@ -1350,9 +1556,9 @@ pub fn validate_declarations(schema: SchemaRef, columns: &[(String, String)]) -> /// public way in. pub(crate) fn declare( schema: SchemaRef, - columns: &[(String, String)], + columns: &[ComputedColumnDeclaration], ) -> Result { - let fields = plan(schema, columns)?; + let fields = plan_declarations(schema, columns)?; Ok(NewColumnTransform::AllNulls(Arc::new(ArrowSchema::new( fields, )))) diff --git a/rust/lancedb/src/table/refresh.rs b/rust/lancedb/src/table/refresh.rs index bc2cc38d1..a8837eb5a 100644 --- a/rust/lancedb/src/table/refresh.rs +++ b/rust/lancedb/src/table/refresh.rs @@ -29,9 +29,12 @@ //! inputs masked to null first, so a poison value in a row nobody is filling //! cannot fail the refresh. +use std::collections::HashSet; use std::sync::Arc; -use arrow_array::{ArrayRef, BooleanArray, RecordBatch, RecordBatchOptions}; +use arrow_array::{ + Array, ArrayRef, BooleanArray, LargeBinaryArray, RecordBatch, RecordBatchOptions, +}; use arrow_schema::Schema as ArrowSchema; use datafusion_expr::ColumnarValue; use futures::{Stream, StreamExt, TryStreamExt}; @@ -40,7 +43,7 @@ use lance::dataset::WriteDestination; use lance::dataset::fragment::FileFragment; use lance::dataset::transaction::Operation; use lance_core::ROW_ID; -use lance_core::datatypes::Schema as LanceSchema; +use lance_core::datatypes::{BlobHandling, Schema as LanceSchema}; use serde::{Deserialize, Serialize}; use super::computed_columns::{BoundExpression, ComputedColumnKind, computed_column_from_field}; @@ -104,6 +107,7 @@ async fn execute_refresh_column_with_source( fields: vec![field.clone()], metadata: Default::default(), }; + let output_is_blob = field.is_blob_v2(); let mut rows_filled = 0u64; let mut replacements = Vec::new(); @@ -113,7 +117,8 @@ async fn execute_refresh_column_with_source( continue; } rows_filled += gained; - let values = fill_stream(&dataset, &fragment, bound.clone(), column).await?; + let values = + fill_stream(&dataset, &fragment, bound.clone(), column, output_is_blob).await?; replacements.push(fragment.write_columns(values, &column_schema).await?); } @@ -294,12 +299,15 @@ fn evaluation_batch( mask_out: Option<&BooleanArray>, ) -> lance_core::Result { let mut columns = Vec::with_capacity(bound.roots.len()); + let mut fields = Vec::with_capacity(bound.roots.len()); for name in &bound.roots { - let column = batch.column_by_name(name).ok_or_else(|| { + let index = batch.schema_ref().index_of(name).map_err(|_| { lance_core::Error::invalid_input(format!( "refreshing a computed column read no {name} column" )) })?; + let column = batch.column(index); + fields.push(batch.schema_ref().field(index).clone()); // Rows outside the mask must not reach the expression: a value in a // deleted or already-filled row can be one it would choke on. columns.push(match mask_out { @@ -308,7 +316,7 @@ fn evaluation_batch( }); } Ok(RecordBatch::try_new_with_options( - bound.read_schema.clone(), + Arc::new(ArrowSchema::new(fields)), columns, &RecordBatchOptions::new().with_row_count(Some(batch.num_rows())), )?) @@ -329,6 +337,64 @@ fn evaluate(bound: &BoundExpression, batch: &RecordBatch) -> lance_core::Result< } } +fn materialized_blob_ids(schema: &LanceSchema, paths: &[String]) -> Result> { + paths + .iter() + .map(|path| { + let field = schema + .resolve(path) + .and_then(|fields| fields.last().copied()) + .ok_or_else(|| Error::InvalidInput { + message: format!("computed Blob input '{path}' no longer exists"), + })?; + if !field.is_blob_v2() { + return Err(Error::InvalidInput { + message: format!("computed Blob input '{path}' is no longer Blob v2"), + }); + } + u32::try_from(field.id).map_err(|_| Error::InvalidInput { + message: format!( + "computed Blob input '{path}' has invalid field id {}", + field.id + ), + }) + }) + .collect() +} + +fn configure_blob_inputs( + scanner: &mut lance::dataset::scanner::Scanner, + schema: &LanceSchema, + bound: &BoundExpression, + extra_blob_id: Option, +) -> Result<()> { + let mut ids = materialized_blob_ids(schema, &bound.blob_paths)?; + ids.extend(extra_blob_id); + scanner.blob_handling(BlobHandling::SomeBlobsBinary(ids)); + Ok(()) +} + +fn blob_array_from_binary(array: &ArrayRef) -> lance_core::Result { + let values = array + .as_any() + .downcast_ref::() + .ok_or_else(|| { + lance_core::Error::invalid_input(format!( + "a Blob v2 computed output produced {}, expected LargeBinary", + array.data_type() + )) + })?; + let mut builder = lance::blob::BlobArrayBuilder::new(values.len()); + for index in 0..values.len() { + if values.is_null(index) { + builder.push_null()?; + } else { + builder.push_bytes(values.value(index))?; + } + } + builder.finish() +} + /// How many rows of one fragment would gain a value. /// /// Scans only the unfilled live rows -- deleted rows never reach the @@ -347,6 +413,7 @@ async fn count_fragment_gains( .with_row_id() .filter(&format!("{} IS NULL", quote_identifier(column)))? .project(&bound.roots)?; + configure_blob_inputs(&mut scanner, dataset.schema(), bound, None)?; let mut gained = 0u64; let mut batches = scanner.try_into_stream().await?; @@ -368,6 +435,7 @@ async fn fill_stream( fragment: &FileFragment, bound: Arc, column: &str, + output_is_blob: bool, ) -> Result> + Send + use<>> { let mut projection: Vec = bound.roots.clone(); projection.push(column.to_string()); @@ -377,6 +445,20 @@ async fn fill_stream( .with_row_id() .include_deleted_rows() .project(&projection)?; + let output_blob_id = output_is_blob + .then(|| { + dataset + .schema() + .field(column) + .and_then(|field| u32::try_from(field.id).ok()) + }) + .flatten(); + configure_blob_inputs( + &mut scanner, + dataset.schema(), + bound.as_ref(), + output_blob_id, + )?; let projected = Arc::new(ArrowSchema::new(vec![ ArrowSchema::from(dataset.schema()) @@ -412,6 +494,11 @@ 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)? + } else { + merged + }; Ok(RecordBatch::try_new(projected.clone(), vec![merged])?) })) } @@ -420,8 +507,9 @@ async fn fill_stream( mod tests { use std::sync::Arc; - use arrow_array::{Int32Array, record_batch}; + use arrow_array::{Array, Int32Array, LargeBinaryArray, RecordBatch, record_batch}; use futures::TryStreamExt; + use lance_core::ROW_ID; use crate::connect; use crate::query::{ExecutableQuery, QueryBase, Select}; @@ -1164,4 +1252,420 @@ mod tests { let err = table.refresh_column("embedding").await.unwrap_err(); assert!(matches!(err, Error::NotSupported { message } if message.contains("udf"))); } + + fn blob_batch(ids: Vec, payloads: Vec>) -> RecordBatch { + use arrow_array::Int32Array; + use arrow_schema::{Field, Schema}; + + let mut builder = lance::blob::BlobArrayBuilder::new(payloads.len()); + for payload in payloads { + match payload { + Some(payload) => builder.push_bytes(payload).unwrap(), + None => builder.push_null().unwrap(), + } + } + RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("id", arrow_schema::DataType::Int32, false), + crate::blob("image", true), + ])), + vec![Arc::new(Int32Array::from(ids)), builder.finish().unwrap()], + ) + .unwrap() + } + + async fn create_blob_table(path: &std::path::Path, batch: RecordBatch) -> Table { + let conn = connect(path.to_str().unwrap()).execute().await.unwrap(); + conn.create_table("blobs", batch).execute().await.unwrap() + } + + #[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() { + use arrow_array::UInt64Array; + use lance_arrow::{ + BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY, BLOB_INLINE_SIZE_THRESHOLD_META_KEY, + }; + use lance_core::datatypes::BlobKind; + + use crate::table::schema_evolution::FieldMetadataUpdate; + + let tmp = tempfile::tempdir().unwrap(); + let table = create_blob_table( + tmp.path(), + blob_batch( + vec![1, 2, 3, 4], + vec![Some(b"hello"), Some(b"ab"), Some(b""), None], + ), + ) + .await; + table + .add_columns() + .computed_blob("image_copy", "image") + .execute() + .await + .unwrap(); + table + .update_field_metadata(&[FieldMetadataUpdate::new("image_copy") + .set(BLOB_INLINE_SIZE_THRESHOLD_META_KEY, "1") + .set(BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY, "4")]) + .await + .unwrap(); + + let first_refresh = table.refresh_column("image_copy").await.unwrap(); + assert_eq!(first_refresh.rows_filled, 3); + assert_eq!( + table.blob_columns().await.unwrap(), + vec!["image".to_string(), "image_copy".to_string()] + ); + + let batches = table + .query() + .with_row_id() + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let batch = arrow_select::concat::concat_batches(&batches[0].schema(), &batches).unwrap(); + assert!( + batch + .column_by_name("image_copy") + .unwrap() + .as_any() + .is::() + ); + let row_ids = batch + .column_by_name(ROW_ID) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + let original = table.fetch_blobs("image", &row_ids).await.unwrap(); + let copied = table.fetch_blobs("image_copy", &row_ids).await.unwrap(); + assert_eq!(original, copied); + let ids = batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let files = table + .fetch_blob_files("image_copy", &row_ids) + .await + .unwrap(); + let mut layouts = ids + .values() + .iter() + .copied() + .zip(files) + .map(|(id, file)| (id, file.and_then(|file| file.kind()))) + .collect::>(); + layouts.sort_by_key(|(id, _)| *id); + assert_eq!( + layouts, + vec![ + (1, Some(BlobKind::Dedicated)), + (2, Some(BlobKind::Packed)), + (3, Some(BlobKind::Inline)), + (4, None), + ] + ); + + table + .add(blob_batch(vec![5], vec![Some(b"appended")])) + .execute() + .await + .unwrap(); + table + .optimize(crate::table::OptimizeAction::Compact { + options: crate::table::CompactionOptions::default(), + remap_options: None, + }) + .await + .unwrap(); + assert_eq!( + table + .refresh_column("image_copy") + .await + .unwrap() + .rows_filled, + 1 + ); + assert_eq!( + table + .refresh_column("image_copy") + .await + .unwrap() + .rows_filled, + 0 + ); + + table.checkout(first_refresh.version).await.unwrap(); + assert_eq!(table.count_rows(None).await.unwrap(), 4); + assert_eq!( + table.blob_columns().await.unwrap(), + vec!["image".to_string(), "image_copy".to_string()] + ); + table.checkout_latest().await.unwrap(); + } + + #[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_blob("invalid", "id + 1") + .execute() + .await + .unwrap_err(); + assert!(matches!( + error, + Error::InvalidExpression { column, message } + if column == "invalid" && message.contains("requires a LargeBinary expression") + )); + } + + #[tokio::test] + async fn test_refresh_materializes_nested_struct_blob_input() { + use arrow_array::{Int32Array, StructArray}; + use arrow_schema::{DataType, Field, Fields, Schema}; + + let tmp = tempfile::tempdir().unwrap(); + let mut blob_builder = lance::blob::BlobArrayBuilder::new(2); + blob_builder.push_bytes(b"nested").unwrap(); + blob_builder.push_null().unwrap(); + let blob_field = crate::blob("image", true); + let metadata_fields = Fields::from(vec![blob_field.clone()]); + let metadata = StructArray::new( + metadata_fields.clone(), + vec![blob_builder.finish().unwrap()], + None, + ); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("metadata", DataType::Struct(metadata_fields), true), + ])), + vec![Arc::new(Int32Array::from(vec![1, 2])), Arc::new(metadata)], + ) + .unwrap(); + let table = create_blob_table(tmp.path(), batch).await; + table + .add_columns() + .computed("payload_copy", "metadata.image") + .execute() + .await + .unwrap(); + + assert_eq!( + table + .refresh_column("payload_copy") + .await + .unwrap() + .rows_filled, + 1 + ); + 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"nested"); + assert!(payloads.is_null(1)); + } + + #[tokio::test] + async fn test_refresh_preserves_list_shape_when_materializing_blob_input() { + use arrow_array::{Int32Array, ListArray}; + use arrow_buffer::{OffsetBuffer, ScalarBuffer}; + use arrow_schema::{DataType, Field, Schema}; + + let tmp = tempfile::tempdir().unwrap(); + let mut blob_builder = lance::blob::BlobArrayBuilder::new(3); + blob_builder.push_bytes(b"a").unwrap(); + blob_builder.push_bytes(b"bb").unwrap(); + blob_builder.push_null().unwrap(); + let item = Arc::new(crate::blob("item", true)); + let images = ListArray::new( + item.clone(), + OffsetBuffer::new(ScalarBuffer::from(vec![0, 2, 3])), + blob_builder.finish().unwrap(), + None, + ); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("images", DataType::List(item), true), + ])), + vec![Arc::new(Int32Array::from(vec![1, 2])), Arc::new(images)], + ) + .unwrap(); + let table = create_blob_table(tmp.path(), batch).await; + table + .add_columns() + .computed("image_payloads", "images") + .execute() + .await + .unwrap(); + + assert_eq!( + table + .refresh_column("image_payloads") + .await + .unwrap() + .rows_filled, + 2 + ); + let batches = table + .query() + .select(Select::columns(&["image_payloads"])) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let output = batches[0] + .column_by_name("image_payloads") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(output.value_offsets(), &[0, 2, 3]); + assert!(output.values().as_any().is::()); + } + + #[tokio::test] + async fn test_refresh_materializes_external_blob_input() { + use arrow_array::{Int32Array, StringArray}; + use arrow_schema::{DataType, Field, Schema}; + + let tmp = tempfile::tempdir().unwrap(); + let payload = b"external-payload"; + let path = tmp.path().join("payload.bin"); + std::fs::write(&path, payload).unwrap(); + let uri = url::Url::from_file_path(path).unwrap().to_string(); + let conn = connect(tmp.path().join("db").to_str().unwrap()) + .execute() + .await + .unwrap(); + let table = conn + .create_empty_table( + "external", + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + crate::blob("image", true), + ])), + ) + .execute() + .await + .unwrap(); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("image", DataType::Utf8, true), + ])), + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(StringArray::from(vec![Some(uri)])), + ], + ) + .unwrap(); + table + .add(batch) + .allow_external_blob_outside_bases(true) + .execute() + .await + .unwrap(); + table + .add_columns() + .computed("payload_copy", "image") + .execute() + .await + .unwrap(); + + assert_eq!( + table + .refresh_column("payload_copy") + .await + .unwrap() + .rows_filled, + 1 + ); + 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), payload); + } } diff --git a/rust/lancedb/src/table/schema_evolution.rs b/rust/lancedb/src/table/schema_evolution.rs index 4f8dc811a..ae7583a31 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: &[(String, String)], + columns: &[computed_columns::ComputedColumnDeclaration], ) -> Result { use lance::dataset::mem_wal::DatasetMemWalExt;