From 569952d15049d6584090d183109e232c2b376e73 Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Sat, 31 Oct 2026 09:07:09 -0700 Subject: [PATCH] fix: allow hidden function assignment columns (#4779) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A second flattened Function declaration fails when the remote describe schema omits an existing binding's internal assignment field. The client incorrectly requires that server-owned field and compares it as a physical user output. Exclude assignment fields from binding/schema validation and the output-schema comparison. Filter declared assignment fields from remote schemas, including seeded caches and snapshot handles, while retaining schema metadata. The regression tests declare two all-nullable flattened bindings consecutively with no assignment declaration, a hidden assignment field, or a returned assignment field. They also cover cached and snapshot schemas and server-owned assignment types and metadata. All four regression cases failed before the fix. Validation: - `cargo fmt --all` - `cargo test --quiet --profile ci --features remote -p lancedb --lib -- table::computed_columns::tests remote::table::tests` — 379 passed, 30 ignored. - `cargo clippy --quiet --profile ci --features remote -p lancedb --tests --examples` Fixes #4778 Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> --- rust/lancedb/src/remote/table.rs | 202 ++++++++++++++++++++- rust/lancedb/src/table/computed_columns.rs | 115 ++++++------ 2 files changed, 255 insertions(+), 62 deletions(-) diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 681ea453e..9ffb88b30 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -763,8 +763,8 @@ impl RemoteTable { return; }; self.track_read_version(description.version); - if let Ok(schema) = arrow_schema::Schema::try_from(description.schema) { - self.schema_cache.seed(Arc::new(schema)); + if let Ok(schema) = schema_from_description(description.schema) { + self.schema_cache.seed(schema); } } @@ -836,7 +836,7 @@ impl RemoteTable { } let description = self.describe_read_snapshot(read_snapshot).await?; - Ok(Arc::new(description.schema.try_into()?)) + schema_from_description(description.schema) } async fn resolve_tag_version_with_request( @@ -1596,6 +1596,34 @@ struct TableDescription { location: Option, } +/// Keep server-owned Function assignment columns out of client schemas. +fn schema_from_description(schema: JsonSchema) -> Result { + let schema: arrow_schema::Schema = schema.try_into()?; + // Reading a schema must remain possible with newer binding metadata. + let bindings = crate::table::computed_columns::function_bindings(&schema).unwrap_or_default(); + let assignments = bindings + .iter() + .filter_map(|binding| { + binding + .assignment() + .map(|assignment| assignment.output_name.as_str()) + }) + .collect::>(); + if assignments.is_empty() { + return Ok(Arc::new(schema)); + } + let fields = schema + .fields() + .iter() + .filter(|field| !assignments.contains(field.name().as_str())) + .cloned() + .collect::>(); + Ok(Arc::new(arrow_schema::Schema::new_with_metadata( + fields, + schema.metadata, + ))) +} + /// How a response body frames its Arrow IPC payload. `/query` answers with file framing /// and `fetch_blobs` with stream framing, so the reader is picked per response. enum ArrowIpcFraming { @@ -1723,8 +1751,7 @@ async fn fetch_schema( }); } - let arrow_schema: arrow_schema::Schema = description.schema.try_into()?; - Ok(Arc::new(arrow_schema)) + schema_from_description(description.schema) } impl std::fmt::Display for RemoteTable { @@ -2318,7 +2345,7 @@ impl BaseTable for RemoteTable { schema, location, } = description; - let schema = Arc::new(arrow_schema::Schema::try_from(schema)?); + let schema = schema_from_description(schema)?; let snapshot = self.with_branch(self.branch.clone()); *snapshot.version.write().await = Some(version); *snapshot.location.write().await = location; @@ -2336,7 +2363,7 @@ impl BaseTable for RemoteTable { schema, location, } = description; - let schema = Arc::new(arrow_schema::Schema::try_from(schema)?); + let schema = schema_from_description(schema)?; let snapshot = self.with_branch(self.branch.clone()); *snapshot.version.write().await = Some(version); *snapshot.location.write().await = location; @@ -10175,6 +10202,167 @@ mod tests { assert_eq!(result.version, 10); } + #[rstest] + #[case(false, false)] + #[case(true, false)] + #[case(true, true)] + #[tokio::test] + async fn test_add_flattened_function_bindings_with_server_owned_assignment( + #[case] has_assignment: bool, + #[case] has_assignment_field: bool, + ) { + use crate::function::{FunctionApplication, FunctionBinding}; + use crate::table::computed_columns::{ + FUNCTION_BINDINGS_META_KEY, function_bindings_metadata, + }; + + let mut raw_binding: serde_json::Value = serde_json::from_str(include_str!( + "../../tests/fixtures/first_class_functions/v1/remote_function_binding.json" + )) + .unwrap(); + for output in raw_binding["outputs"].as_array_mut().unwrap() { + output["nullable"] = json!(true); + } + let assignment_name = "__function_assignment_fb_01K3TEXT"; + if has_assignment { + raw_binding["assignment"] = json!({ + "output_name": assignment_name, + "output_field_id": -1, + }); + raw_binding["output_schema"]["fields"] + .as_array_mut() + .unwrap() + .push(json!({ + "name": assignment_name, + "nullable": true, + "type": {"type": "bool"}, + })); + } + let function = raw_binding["function"].clone(); + let binding: FunctionBinding = serde_json::from_value(raw_binding).unwrap(); + let mut schema = fixture_bound_schema(); + schema.metadata.insert( + FUNCTION_BINDINGS_META_KEY.into(), + function_bindings_metadata(&[binding]).unwrap(), + ); + let visible_schema = Arc::new(schema.clone()); + if has_assignment_field { + let mut fields = schema.fields().to_vec(); + // Even the assignment field's physical type and metadata are server-owned. + fields.push(Arc::new(Field::new( + assignment_name, + DataType::Int32, + false, + ))); + schema = Schema::new_with_metadata(fields, schema.metadata().clone()); + } + let description = describe_response(&schema); + let seeded = RemoteTable::new_mock( + "my_table".into(), + move |_| { + http::Response::builder() + .status(200) + .body(description.clone()) + .unwrap() + }, + None, + ); + seeded.seed_schema(&describe_response(&schema)); + assert_eq!(seeded.schema().await.unwrap(), visible_schema); + seeded.invalidate_schema_cache(); + assert_eq!( + seeded + .schema_read_snapshot(seeded.snapshot_read_state().await) + .await + .unwrap(), + visible_schema, + ); + assert_eq!( + seeded + .query_snapshot() + .await + .unwrap() + .schema() + .await + .unwrap(), + visible_schema, + ); + assert_eq!( + seeded + .checkout_current() + .await + .unwrap() + .schema() + .await + .unwrap(), + visible_schema, + ); + + let initial_schema = Schema::new(schema.fields()[..2].to_vec()); + let writes = Arc::new(AtomicUsize::new(0)); + let handler_writes = writes.clone(); + let table = Table::new_with_handler("my_table", move |request| { + let body = match request.url().path() { + "/v1/table/my_table/describe/" => { + if handler_writes.load(Ordering::SeqCst) == 0 { + describe_response(&initial_schema) + } else { + describe_response(&schema) + } + } + "/v1/table/my_table/add_columns/" => { + let actual: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()) + .unwrap(); + let index = handler_writes.fetch_add(1, Ordering::SeqCst); + let prefix = if index == 0 { "search" } else { "secondary" }; + assert_eq!( + actual["new_columns"], + json!([ + {"name": format!("{prefix}_text"), "all_null": true}, + {"name": format!("{prefix}_token_count"), "all_null": true} + ]) + ); + json!({"version": index + 2}).to_string() + } + path => panic!("Unexpected path: {path}"), + }; + http::Response::builder().status(200).body(body).unwrap() + }); + + for prefix in ["search", "secondary"] { + let application = FunctionApplication::from_json( + &json!({ + "function": function, + "inputs": [ + {"parameter": "title", "kind": "column", "value": {"path": "title"}}, + {"parameter": "body", "kind": "column", "value": {"path": "body"}} + ], + "output": {"kind": "named_struct", "fields": [ + {"name": "normalized_text", "arrow_type": "utf8", "nullable": true}, + {"name": "token_count", "arrow_type": "int64", "nullable": true} + ]}, + "columns": { + "normalized_text": format!("{prefix}_text"), + "token_count": format!("{prefix}_token_count") + } + }) + .to_string(), + ) + .unwrap(); + table + .add_columns() + .function(application) + .execute() + .await + .unwrap(); + if prefix == "search" { + assert_eq!(table.schema().await.unwrap(), visible_schema); + } + } + assert_eq!(writes.load(Ordering::SeqCst), 2); + } + #[tokio::test] async fn test_add_fixed_size_list_function_column_declares_the_vector_type() { let table = Table::new_with_handler("my_table", |request| { diff --git a/rust/lancedb/src/table/computed_columns.rs b/rust/lancedb/src/table/computed_columns.rs index f2ec6dc4a..ecdaf6c7a 100644 --- a/rust/lancedb/src/table/computed_columns.rs +++ b/rust/lancedb/src/table/computed_columns.rs @@ -407,12 +407,23 @@ pub(crate) fn ensure_supported_function_metadata(schema: &ArrowSchema) -> Result .iter() .map(|binding| (binding.binding_id(), binding)) .collect::>(); + let assignment_columns = bindings + .iter() + .filter_map(|binding| { + binding + .assignment() + .map(|assignment| assignment.output_name.as_str()) + }) + .collect::>(); for field in schema.fields() { - if field - .metadata() - .get(COMPUTED_COLUMN_META_KEY) - .map(String::as_str) - != Some("true") + // Assignment columns are server-owned bookkeeping, even when describe + // includes them. Only user-visible outputs require reciprocal metadata. + if assignment_columns.contains(field.name().as_str()) + || field + .metadata() + .get(COMPUTED_COLUMN_META_KEY) + .map(String::as_str) + != Some("true") { continue; } @@ -1361,6 +1372,7 @@ fn ensure_binding_matches_schema(schema: &ArrowSchema, binding: &FunctionBinding &output.arrow_type, )?); } + let mut declared_output_schema = binding.output_schema().cloned(); if let Some(assignment) = binding.assignment() { if binding .outputs() @@ -1372,49 +1384,17 @@ fn ensure_binding_matches_schema(schema: &ArrowSchema, binding: &FunctionBinding binding.binding_id() ))); } - let field = schema - .field_with_name(&assignment.output_name) - .map_err(|_| { - invalid_function(format!( - "Function binding '{}' assignment column '{}' is missing", - binding.binding_id(), - assignment.output_name - )) - })?; - let metadata = field.metadata(); - if field.data_type() != &DataType::Boolean - || !field.is_nullable() - || metadata.get(COMPUTED_COLUMN_META_KEY).map(String::as_str) != Some("true") - || metadata.get(KIND_META_KEY).map(String::as_str) != Some(FUNCTION_KIND) - || metadata - .get(FUNCTION_BINDING_ID_META_KEY) - .map(String::as_str) - != Some(binding.binding_id()) - || metadata - .get(FUNCTION_OUTPUT_ORDINAL_META_KEY) - .and_then(|value| value.parse::().ok()) - != Some(FUNCTION_ASSIGNMENT_OUTPUT_ORDINAL) + // Servers may include the internal assignment in the persisted output + // schema while omitting it from describe. Compare only user outputs. + if let Some(fields) = declared_output_schema + .as_mut() + .and_then(|schema| schema.get_mut("fields")) + .and_then(Value::as_array_mut) { - return Err(invalid_function(format!( - "Function binding '{}' assignment column no longer matches its declaration", - binding.binding_id() - ))); + fields.retain(|field| { + field.get("name").and_then(Value::as_str) != Some(assignment.output_name.as_str()) + }); } - let json = lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![ - ArrowField::new(assignment.output_name.clone(), DataType::Boolean, true), - ])) - .map_err(|e| invalid_function(format!("invalid Function assignment schema: {e}")))?; - output_fields.push(json.fields.into_iter().next().unwrap()); - } else if binding.outputs().iter().all(|output| output.nullable) - && binding - .outputs() - .iter() - .all(|output| output.result_field != WHOLE_RESULT_FIELD) - { - return Err(invalid_function(format!( - "Function binding '{}' has no flattened-result assignment column", - binding.binding_id() - ))); } let output_schema = JsonArrowSchema::new(output_fields); let output_schema = serde_json::to_value(output_schema).map_err(|e| { @@ -1422,7 +1402,7 @@ fn ensure_binding_matches_schema(schema: &ArrowSchema, binding: &FunctionBinding "could not encode exact Function output schema: {e}" )) })?; - if binding.output_schema() != Some(&output_schema) { + if declared_output_schema.as_ref() != Some(&output_schema) { return Err(invalid_function(format!( "Function binding '{}' output schema does not match physical siblings", binding.binding_id() @@ -3895,7 +3875,7 @@ mod tests { } #[test] - fn test_binding_preserves_all_nullable_outputs_with_an_assignment_column() { + fn test_binding_accepts_server_owned_assignment_columns() { let mut raw_binding: Value = serde_json::from_str(include_str!( "../../tests/fixtures/first_class_functions/v1/remote_function_binding.json" )) @@ -3904,16 +3884,11 @@ mod tests { raw_binding["outputs"][1]["nullable"] = Value::Bool(true); let without_assignment: FunctionBinding = serde_json::from_value(raw_binding.clone()).unwrap(); - let error = ensure_binding_matches_schema( + ensure_binding_matches_schema( &valid_function_binding_schema(true, true, &without_assignment), &without_assignment, ) - .unwrap_err(); - assert!( - error - .to_string() - .contains("flattened-result assignment column") - ); + .unwrap(); raw_binding["assignment"] = serde_json::json!({ "output_name": "__function_assignment_fb_01K3TEXT", @@ -3945,6 +3920,36 @@ mod tests { ); ensure_supported_function_metadata(&schema).unwrap(); + let assignment_name = &binding.assignment().unwrap().output_name; + let fields = schema + .fields() + .iter() + .filter(|field| field.name() != assignment_name) + .cloned() + .collect::>(); + let hidden = ArrowSchema::new_with_metadata(fields.clone(), schema.metadata().clone()); + ensure_supported_function_metadata(&hidden).unwrap(); + + // The assignment field's type and metadata belong to the server. + let mut fields = fields; + fields.push(Arc::new( + ArrowField::new(assignment_name, DataType::Int32, false).with_metadata(HashMap::from( + [ + (COMPUTED_COLUMN_META_KEY.into(), "true".into()), + (KIND_META_KEY.into(), "server_owned".into()), + ], + )), + )); + let changed = ArrowSchema::new_with_metadata(fields, schema.metadata().clone()); + ensure_supported_function_metadata(&changed).unwrap(); + + // Describe metadata may also omit the internal output-schema field. + let mut raw = serde_json::to_value(&binding).unwrap(); + raw["output_schema"]["fields"].as_array_mut().unwrap().pop(); + let binding_without_assignment_field: FunctionBinding = + serde_json::from_value(raw).unwrap(); + ensure_binding_matches_schema(&hidden, &binding_without_assignment_field).unwrap(); + let mut metadata: Value = serde_json::from_str(schema.metadata().get(FUNCTION_BINDINGS_META_KEY).unwrap()) .unwrap();