diff --git a/rust/lancedb/src/table/computed_columns.rs b/rust/lancedb/src/table/computed_columns.rs index e86b350ee..fe4c17b1f 100644 --- a/rust/lancedb/src/table/computed_columns.rs +++ b/rust/lancedb/src/table/computed_columns.rs @@ -491,24 +491,6 @@ pub(crate) fn ensure_supported_function_metadata(schema: &ArrowSchema) -> Result ensure_computed_inputs_acyclic(schema, &bindings) } -/// Refuse `operation` outright on a table with a Function binding. For -/// operations that cannot say which columns they touch; the others use -/// [`ensure_not_function_bound`]. -pub(crate) fn ensure_no_function_bindings_for_mutation( - schema: &ArrowSchema, - operation: &str, -) -> Result<()> { - ensure_supported_function_metadata(schema)?; - if !function_bindings(schema)?.is_empty() { - return Err(Error::NotSupported { - message: format!( - "{operation} is not supported on a table with registered Function bindings" - ), - }); - } - Ok(()) -} - /// Refuse `operation` only when a path in `touched` names a column a Function /// binding depends on: an input's root, an output, or the assignment column. /// A binding stores those columns' exact Arrow fields, so editing one strands it. @@ -2406,6 +2388,111 @@ impl lance_datafusion::utils:: } } +/// Test fixtures for a table carrying a Function binding, shared by the +/// write-path tests. No local path declares a binding, so these stamp one +/// the way the server does. +#[cfg(test)] +pub(crate) mod test_support { + use arrow_array::record_batch; + + use super::{ + FUNCTION_ASSIGNMENT_OUTPUT_ORDINAL, FUNCTION_BINDINGS_META_KEY, + ensure_supported_function_metadata, function_bindings_metadata, + function_computed_column_metadata, + }; + use crate::Table; + use crate::connect; + use crate::function::FunctionBinding; + + /// A table carrying the fixture binding: `title` and `body` are its + /// inputs, `search_text` and `search_token_count` its outputs, `spare` + /// nobody's. Stamped the way the server does it, since no local path + /// declares a binding. + pub async fn bound_table() -> Table { + let batch = record_batch!( + ("title", Utf8, ["a"]), + ("body", Utf8, ["b"]), + ("search_text", Utf8, ["a b"]), + ("search_token_count", Int64, [2]), + ("spare", Int32, [1]) + ) + .unwrap(); + bound_table_from(batch).await + } + + /// [`bound_table`] over a caller's batch, which must carry `title`, `body`, + /// `search_text` and `search_token_count`. + pub async fn bound_table_from(batch: arrow_array::RecordBatch) -> Table { + let conn = connect("memory://").execute().await.unwrap(); + let table = conn.create_table("bound", batch).execute().await.unwrap(); + stamp_bindings(&table, &[text_features_binding("")]).await; + table + } + + /// The fixture binding, its id and output names suffixed so a table can + /// carry more than one. + pub fn text_features_binding(suffix: &str) -> FunctionBinding { + let raw = include_str!( + "../../tests/fixtures/first_class_functions/v1/remote_function_binding.json" + ); + if suffix.is_empty() { + return FunctionBinding::from_json(raw).unwrap(); + } + FunctionBinding::from_json( + &raw.replace("fb_01K3TEXT", &format!("fb_01K3TEXT{suffix}")) + .replace("search_text", &format!("search_text{suffix}")) + .replace("search_token_count", &format!("search_token_count{suffix}")), + ) + .unwrap() + } + + /// Stamp bindings the way the server does it, since no local path + /// declares one: the envelope on the schema, a declaration marker on + /// every output field. + pub async fn stamp_bindings(table: &Table, bindings: &[FunctionBinding]) { + let native = table.as_native().unwrap(); + let mut dataset = native.dataset.get().await.unwrap().as_ref().clone(); + dataset + .update_schema_metadata(vec![( + FUNCTION_BINDINGS_META_KEY.to_string(), + Some(function_bindings_metadata(bindings).unwrap()), + )]) + .await + .unwrap(); + let inputs = ["title".to_string(), "body".to_string()]; + let outputs = bindings + .iter() + .flat_map(|binding| { + let assignment = binding.assignment().map(|assignment| { + ( + assignment.output_name.clone(), + FUNCTION_ASSIGNMENT_OUTPUT_ORDINAL, + ) + }); + binding + .outputs() + .iter() + .map(|output| (output.output_name.clone(), output.output_ordinal)) + .chain(assignment) + .map(|(name, ordinal)| { + ( + dataset.schema().field(&name).unwrap().id as u32, + function_computed_column_metadata( + binding.binding_id(), + ordinal, + &inputs, + ), + ) + }) + .collect::>() + }) + .collect::>(); + dataset.replace_field_metadata(outputs).await.unwrap(); + native.dataset.update(dataset); + ensure_supported_function_metadata(&table.schema().await.unwrap()).unwrap(); + } +} + #[cfg(test)] mod tests { /// The gate's reproducer: the validator applies the same schema-level diff --git a/rust/lancedb/src/table/merge.rs b/rust/lancedb/src/table/merge.rs index 24c5f55b3..91a0ed6ea 100644 --- a/rust/lancedb/src/table/merge.rs +++ b/rust/lancedb/src/table/merge.rs @@ -238,6 +238,21 @@ impl MergeInsertBuilder { canonicalize_merge_filter(self.when_not_matched_by_source_delete_filt.take())?; Ok(()) } + + /// The existing columns a matched row takes from `source`: every source + /// column but the join keys, which a match leaves equal. An insert-only + /// merge rewrites none; its new rows get Function outputs from the next + /// refresh, as appended rows do. + pub(crate) fn function_guarded_columns(&self, source: &Fields) -> Vec { + if !self.when_matched_update_all { + return Vec::new(); + } + source + .iter() + .filter(|field| !self.on.iter().any(|key| key == field.name())) + .map(|field| field.name().clone()) + .collect() + } } fn canonicalize_merge_filter(filter: Option) -> Result> { @@ -326,9 +341,10 @@ pub(crate) async fn execute_merge_insert( ) -> Result { params.canonicalize_filters()?; let schema = table.schema().await?; - super::computed_columns::ensure_no_function_bindings_for_mutation( + super::computed_columns::ensure_not_function_bound( schema.as_ref(), "merge_insert", + params.function_guarded_columns(new_data.schema().fields()), )?; // Normalize blob layouts and JSON extension metadata just as on append. // Keep arrow.json text labelled until Lance encodes it as JSONB. @@ -889,6 +905,87 @@ mod tests { assert_eq!(result.num_deleted_rows, (ROW_COUNT - 1) as u64); assert_eq!(table.count_rows(None).await.unwrap(), 1); } + + fn reader(batch: RecordBatch) -> Box { + let schema = batch.schema(); + Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema)) + } + + /// Inserting rows through a merge is an append: the new rows' Function + /// outputs come from the next refresh, so a binding does not refuse it. + #[tokio::test] + async fn merge_insert_only_is_allowed_on_a_function_table() { + let table = crate::table::computed_columns::test_support::bound_table_from( + arrow_array::record_batch!( + ("id", Int32, [1]), + ("note", Utf8, ["n"]), + ("title", Utf8, ["a"]), + ("body", Utf8, ["b"]), + ("search_text", Utf8, ["a b"]), + ("search_token_count", Int64, [2]) + ) + .unwrap(), + ) + .await; + let mut merge = table.merge_insert(&["id"]); + merge.when_not_matched_insert_all(); + let result = merge + .execute(reader( + arrow_array::record_batch!( + ("id", Int32, [2]), + ("note", Utf8, ["m"]), + ("title", Utf8, ["c"]), + ("body", Utf8, ["d"]) + ) + .unwrap(), + )) + .await + .unwrap(); + assert_eq!(result.num_inserted_rows, 1); + assert_eq!(table.count_rows(None).await.unwrap(), 2); + } + + /// A matched row takes every non-key source column, so the merge is + /// refused only when one of those is a binding's input or output. + #[tokio::test] + async fn merge_update_on_a_function_table_refuses_only_bound_columns() { + let table = crate::table::computed_columns::test_support::bound_table_from( + arrow_array::record_batch!( + ("id", Int32, [1]), + ("note", Utf8, ["n"]), + ("title", Utf8, ["a"]), + ("body", Utf8, ["b"]), + ("search_text", Utf8, ["a b"]), + ("search_token_count", Int64, [2]) + ) + .unwrap(), + ) + .await; + let mut merge = table.merge_insert(&["id"]); + merge.when_matched_update_all(None); + let result = merge + .execute(reader( + arrow_array::record_batch!(("id", Int32, [1]), ("note", Utf8, ["m"])).unwrap(), + )) + .await + .unwrap(); + assert_eq!(result.num_updated_rows, 1); + + let mut merge = table.merge_insert(&["id"]); + merge.when_matched_update_all(None); + let err = merge + .execute(reader( + arrow_array::record_batch!(("id", Int32, [1]), ("title", Utf8, ["z"])).unwrap(), + )) + .await + .unwrap_err(); + assert!( + matches!(&err, crate::Error::InvalidInput { message } + if message.contains("merge_insert of 'title'") + && message.contains("a Function binding reads or writes it")), + "{err:?}" + ); + } } #[cfg(test)] diff --git a/rust/lancedb/src/table/schema_evolution.rs b/rust/lancedb/src/table/schema_evolution.rs index 1ad532d0a..f80e89ee7 100644 --- a/rust/lancedb/src/table/schema_evolution.rs +++ b/rust/lancedb/src/table/schema_evolution.rs @@ -459,96 +459,16 @@ mod tests { use crate::function::FunctionBinding; use crate::query::{ExecutableQuery, QueryBase, Select}; use crate::table::NewColumnTransform; + use crate::table::computed_columns::test_support::{ + bound_table, stamp_bindings, text_features_binding, + }; use crate::table::computed_columns::{ - FUNCTION_ASSIGNMENT_OUTPUT_ORDINAL, FUNCTION_BINDINGS_META_KEY, - ensure_supported_function_metadata, function_bindings, function_bindings_metadata, - function_computed_column_metadata, is_declaration_key, + FUNCTION_BINDINGS_META_KEY, ensure_supported_function_metadata, function_bindings, + is_declaration_key, }; use crate::{Error, Table}; use std::collections::HashMap; - /// A table carrying the fixture binding: `title` and `body` are its - /// inputs, `search_text` and `search_token_count` its outputs, `spare` - /// nobody's. Stamped the way the server does it, since no local path - /// declares a binding. - async fn bound_table() -> Table { - let conn = connect("memory://").execute().await.unwrap(); - let batch = record_batch!( - ("title", Utf8, ["a"]), - ("body", Utf8, ["b"]), - ("search_text", Utf8, ["a b"]), - ("search_token_count", Int64, [2]), - ("spare", Int32, [1]) - ) - .unwrap(); - let table = conn.create_table("bound", batch).execute().await.unwrap(); - stamp_bindings(&table, &[text_features_binding("")]).await; - table - } - - /// The fixture binding, its id and output names suffixed so a table can - /// carry more than one. - fn text_features_binding(suffix: &str) -> FunctionBinding { - let raw = include_str!( - "../../tests/fixtures/first_class_functions/v1/remote_function_binding.json" - ); - if suffix.is_empty() { - return FunctionBinding::from_json(raw).unwrap(); - } - FunctionBinding::from_json( - &raw.replace("fb_01K3TEXT", &format!("fb_01K3TEXT{suffix}")) - .replace("search_text", &format!("search_text{suffix}")) - .replace("search_token_count", &format!("search_token_count{suffix}")), - ) - .unwrap() - } - - /// Stamp bindings the way the server does it, since no local path - /// declares one: the envelope on the schema, a declaration marker on - /// every output field. - async fn stamp_bindings(table: &Table, bindings: &[FunctionBinding]) { - let native = table.as_native().unwrap(); - let mut dataset = native.dataset.get().await.unwrap().as_ref().clone(); - dataset - .update_schema_metadata(vec![( - FUNCTION_BINDINGS_META_KEY.to_string(), - Some(function_bindings_metadata(bindings).unwrap()), - )]) - .await - .unwrap(); - let inputs = ["title".to_string(), "body".to_string()]; - let outputs = bindings - .iter() - .flat_map(|binding| { - let assignment = binding.assignment().map(|assignment| { - ( - assignment.output_name.clone(), - FUNCTION_ASSIGNMENT_OUTPUT_ORDINAL, - ) - }); - binding - .outputs() - .iter() - .map(|output| (output.output_name.clone(), output.output_ordinal)) - .chain(assignment) - .map(|(name, ordinal)| { - ( - dataset.schema().field(&name).unwrap().id as u32, - function_computed_column_metadata( - binding.binding_id(), - ordinal, - &inputs, - ), - ) - }) - .collect::>() - }) - .collect::>(); - dataset.replace_field_metadata(outputs).await.unwrap(); - native.dataset.update(dataset); - ensure_supported_function_metadata(&table.schema().await.unwrap()).unwrap(); - } - fn metadata_update(path: &str) -> FieldMetadataUpdate { FieldMetadataUpdate { path: path.into(), diff --git a/rust/lancedb/src/table/update.rs b/rust/lancedb/src/table/update.rs index cec787ff8..aa325cd7c 100644 --- a/rust/lancedb/src/table/update.rs +++ b/rust/lancedb/src/table/update.rs @@ -94,9 +94,10 @@ pub(crate) async fn execute_update( // 1. Snapshot the current dataset let dataset = table.dataset.get().await?; crate::materialized_view::ensure_not_a_view(&table.name, &dataset.schema().metadata, "update")?; - super::computed_columns::ensure_no_function_bindings_for_mutation( + super::computed_columns::ensure_not_function_bound( &arrow_schema::Schema::from(dataset.schema()), "update", + update.columns.iter().map(|(name, _)| name.as_str()), )?; super::computed_columns::ensure_not_written( &arrow_schema::Schema::from(dataset.schema()), @@ -445,4 +446,27 @@ mod tests { tbl.update().column("i", "i+1").execute().await.unwrap(); assert_eq!(0, tbl.count_rows(Some("i == 0".to_string())).await.unwrap()); } + + /// A Function binding protects only the columns it reads or writes; any + /// other column updates as on a plain table. + #[tokio::test] + async fn update_on_a_function_table_refuses_only_bound_columns() { + let table = crate::table::computed_columns::test_support::bound_table().await; + let result = table.update().column("spare", "7").execute().await.unwrap(); + assert_eq!(result.rows_updated, 1); + for column in ["title", "search_text"] { + let err = table + .update() + .column(column, "'x'") + .execute() + .await + .unwrap_err(); + assert!( + matches!(&err, crate::Error::InvalidInput { message } + if message.contains(&format!("update of '{column}'")) + && message.contains("a Function binding reads or writes it")), + "{column}: {err:?}" + ); + } + } }