rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-11-13 14:16:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: allow hidden function assignment columns (#4779)

Browse Source 
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

<!-- lance-gatekeeper-fix:v1 agent=c55eff66e22d736a6b954a9eeec88898
generation=1 -->

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
This commit is contained in:
lancedb-gatefixer[bot]
2026-10-31 09:07:09 -07:00
committed by GitHub
parent 784d649380
commit 569952d150
2 changed files with 255 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/lancedb/src/remote/table.rs
+195 -7
View File
@@ -763,8 +763,8 @@ impl<S: HttpSend> RemoteTable<S> {
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<S: HttpSend> RemoteTable<S> {
}
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<String>,
}
/// Keep server-owned Function assignment columns out of client schemas.
fn schema_from_description(schema: JsonSchema) -> Result<SchemaRef> {
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::<HashSet<_>>();
if assignments.is_empty() {
return Ok(Arc::new(schema));
}
let fields = schema
.fields()
.iter()
.filter(|field| !assignments.contains(field.name().as_str()))
.cloned()
.collect::<Vec<_>>();
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<S: HttpSend>(
});
}
let arrow_schema: arrow_schema::Schema = description.schema.try_into()?;
Ok(Arc::new(arrow_schema))
schema_from_description(description.schema)
}
impl<S: HttpSend> std::fmt::Display for RemoteTable<S> {
@@ -2318,7 +2345,7 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
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<S: HttpSend> BaseTable for RemoteTable<S> {
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| {
rust/lancedb/src/table/computed_columns.rs
+60 -55
View File
@@ -407,12 +407,23 @@ pub(crate) fn ensure_supported_function_metadata(schema: &ArrowSchema) -> Result
.iter()
.map(|binding| (binding.binding_id(), binding))
.collect::<HashMap<_, _>>();
let assignment_columns = bindings
.iter()
.filter_map(|binding| {
binding
.assignment()
.map(|assignment| assignment.output_name.as_str())
})
.collect::<HashSet<_>>();
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::<u32>().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::<Vec<_>>();
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();