diff --git a/rust/lancedb/src/table/computed_columns.rs b/rust/lancedb/src/table/computed_columns.rs index 4e8d8211e..ed568aa9b 100644 --- a/rust/lancedb/src/table/computed_columns.rs +++ b/rust/lancedb/src/table/computed_columns.rs @@ -574,7 +574,14 @@ fn resolve_field_path<'a>(schema: &'a ArrowSchema, path: &str) -> Result<&'a Arr Ok(field) } -fn canonical_input_arrow_type(field: &JsonArrowField) -> Result { +fn canonical_input_arrow_type(field: &ArrowField) -> Result { + if field.is_blob_v2() { + return Ok("blob_v2".to_string()); + } + let schema = + lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![field.clone()])) + .map_err(|e| invalid_function(format!("invalid Function input schema: {e}")))?; + let field = schema.fields.into_iter().next().unwrap(); if field.r#type.fields.is_none() && field.r#type.length.is_none() { Ok(field.r#type.r#type.clone()) } else { @@ -663,6 +670,90 @@ fn parse_output_arrow_type(raw: &str) -> Result { Ok(data_type) } +fn parse_output_arrow_field(name: &str, nullable: bool, raw: &str) -> Result { + if raw.trim() == "blob_v2" { + let schema = + lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![crate::blob( + name, nullable, + )])) + .map_err(|e| invalid_function(format!("invalid Function Blob type: {e}")))?; + return Ok(schema.fields.into_iter().next().unwrap()); + } + + fn parse_nested(name: &str, nullable: bool, raw: &str) -> Result { + let raw = raw.trim(); + if raw == "blob_v2" { + return parse_output_arrow_field(name, nullable, raw); + } + if let Some(inner) = raw + .strip_prefix("list<") + .and_then(|value| value.strip_suffix('>')) + { + let mut data_type = JsonArrowDataType::new("list".to_string()); + data_type.fields = Some(vec![parse_nested("item", false, inner)?]); + return Ok(JsonArrowField::new(name.to_string(), nullable, data_type)); + } + if let Some(inner) = raw + .strip_prefix("large_list<") + .and_then(|value| value.strip_suffix('>')) + { + let mut data_type = JsonArrowDataType::new("large_list".to_string()); + data_type.fields = Some(vec![parse_nested("item", false, inner)?]); + return Ok(JsonArrowField::new(name.to_string(), nullable, data_type)); + } + if let Some((inner, size)) = split_fixed_size_list(raw) { + let mut data_type = JsonArrowDataType::new("fixed_size_list".to_string()); + data_type.fields = Some(vec![parse_nested("item", false, inner)?]); + data_type.length = Some(i64::from(size)); + return Ok(JsonArrowField::new(name.to_string(), nullable, data_type)); + } + Ok(JsonArrowField::new( + name.to_string(), + nullable, + parse_output_arrow_type(raw)?, + )) + } + + let field = parse_nested(name, nullable, raw)?; + lance_namespace::schema::convert_json_arrow_field(&field) + .map_err(|e| invalid_function(format!("unsupported Function Arrow type '{raw}': {e}")))?; + Ok(field) +} + +fn function_fields_equivalent(actual: &ArrowField, expected: &ArrowField) -> bool { + if actual.is_blob_v2() || expected.is_blob_v2() { + return actual.is_blob_v2() + && expected.is_blob_v2() + && actual.data_type() == expected.data_type(); + } + match (actual.data_type(), expected.data_type()) { + (DataType::Struct(actual), DataType::Struct(expected)) => { + actual.len() == expected.len() + && actual.iter().zip(expected).all(|(actual, expected)| { + actual.name() == expected.name() + && actual.is_nullable() == expected.is_nullable() + && function_fields_equivalent(actual, expected) + }) + } + (DataType::List(actual), DataType::List(expected)) + | (DataType::LargeList(actual), DataType::LargeList(expected)) => { + actual.name() == expected.name() + && actual.is_nullable() == expected.is_nullable() + && function_fields_equivalent(actual, expected) + } + ( + DataType::FixedSizeList(actual, actual_size), + DataType::FixedSizeList(expected, expected_size), + ) => { + actual_size == expected_size + && actual.name() == expected.name() + && actual.is_nullable() == expected.is_nullable() + && function_fields_equivalent(actual, expected) + } + _ => actual.data_type() == expected.data_type(), + } +} + fn ensure_binding_matches_schema(schema: &ArrowSchema, binding: &FunctionBinding) -> Result<()> { let mut input_fields = Vec::with_capacity(binding.inputs().len()); for input in binding.inputs() { @@ -694,12 +785,7 @@ fn ensure_binding_matches_schema(schema: &ArrowSchema, binding: &FunctionBinding input.nullable, ) .with_metadata(field.metadata().clone()); - let json = lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![ - parameter_field.clone(), - ])) - .map_err(|e| invalid_function(format!("invalid Function input schema: {e}")))?; - let json_field = json.fields.into_iter().next().unwrap(); - if canonical_input_arrow_type(&json_field)? != input.arrow_type { + if canonical_input_arrow_type(¶meter_field)? != input.arrow_type { return Err(invalid_function(format!( "Function input '{}' type no longer matches binding '{}'", input.field_path, @@ -737,21 +823,17 @@ fn ensure_binding_matches_schema(schema: &ArrowSchema, binding: &FunctionBinding binding.binding_id() ))); } - let expected_type = parse_output_arrow_type(&output.arrow_type)?; - let expected_type = lance_namespace::schema::convert_json_arrow_type(&expected_type) + let expected_field = parse_output_arrow_field(field.name(), true, &output.arrow_type)?; + let expected_field = lance_namespace::schema::convert_json_arrow_field(&expected_field) .map_err(|e| invalid_function(format!("invalid Function output type: {e}")))?; - if field.data_type() != &expected_type { + if !function_fields_equivalent(field, &expected_field) { return Err(invalid_function(format!( "Function output '{}' type no longer matches binding '{}'", output.output_name, binding.binding_id() ))); } - output_fields.push(ArrowField::new( - field.name().clone(), - field.data_type().clone(), - true, - )); + output_fields.push(expected_field); } let output_schema = lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(output_fields)) @@ -845,15 +927,10 @@ pub(crate) fn plan_function_application( field.is_nullable(), ) .with_metadata(field.metadata().clone()); - let input_schema = lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![ - parameter_field.clone(), - ])) - .map_err(|e| invalid_function(format!("invalid Function input schema: {e}")))?; - let json_field = input_schema.fields.into_iter().next().unwrap(); input_bindings.push(FunctionInputTarget { parameter: input.parameter.clone(), field_path: path.to_string(), - arrow_type: canonical_input_arrow_type(&json_field)?, + arrow_type: canonical_input_arrow_type(¶meter_field)?, nullable: field.is_nullable(), }); input_fields.push(parameter_field); @@ -882,16 +959,19 @@ pub(crate) fn plan_function_application( "Function logical outputs must be non-nullable during NULL assignment", )); } - let data_type = - parse_output_arrow_type(output.arrow_type.as_deref().ok_or_else(|| { + let output_field = parse_output_arrow_field( + name, + true, + output.arrow_type.as_deref().ok_or_else(|| { invalid_function("scalar Function output is missing its Arrow type") - })?)?; + })?, + )?; outputs.push(FunctionOutputTarget { result_field: WHOLE_RESULT_FIELD.to_string(), output_name: name.to_string(), output_ordinal: 0, }); - output_fields.push(JsonArrowField::new(name.to_string(), true, data_type)); + output_fields.push(output_field); } "named_struct" => { if output.fields.is_empty() { @@ -935,13 +1015,7 @@ pub(crate) fn plan_function_application( let fields = output .fields .iter() - .map(|field| { - Ok(JsonArrowField::new( - field.name.clone(), - false, - parse_output_arrow_type(&field.arrow_type)?, - )) - }) + .map(|field| parse_output_arrow_field(&field.name, false, &field.arrow_type)) .collect::>>()?; let mut data_type = JsonArrowDataType::new("struct".to_string()); data_type.fields = Some(fields); @@ -968,11 +1042,7 @@ pub(crate) fn plan_function_application( output_name: name.clone(), output_ordinal: ordinal as u32, }); - output_fields.push(JsonArrowField::new( - name.clone(), - true, - parse_output_arrow_type(&field.arrow_type)?, - )); + output_fields.push(parse_output_arrow_field(name, true, &field.arrow_type)?); } } } @@ -2572,6 +2642,44 @@ mod tests { .unwrap() } + fn blob_function_application() -> FunctionApplication { + FunctionApplication::from_json( + r#"{ + "function":{"name":"copy_blob","version":"fv_exact"}, + "inputs":[ + {"parameter":"image","kind":"column","value":{"path":"image"}} + ], + "output":{"kind":"scalar","arrow_type":"blob_v2","nullable":false} + }"#, + ) + .unwrap() + } + + #[test] + fn test_blob_function_input_uses_the_semantic_binding_type() { + let schema = ArrowSchema::new(vec![crate::blob("image", true)]); + let plan = + plan_function_application(&schema, &blob_function_application(), Some("copy")).unwrap(); + + assert_eq!(plan.input_bindings[0].arrow_type, "blob_v2"); + let input = lance_namespace::schema::convert_json_arrow_field(&plan.input_schema.fields[0]) + .unwrap(); + assert!(input.is_blob_v2()); + } + + #[test] + fn test_blob_function_output_is_plannable() { + let schema = ArrowSchema::new(vec![crate::blob("image", true)]); + let plan = + plan_function_application(&schema, &blob_function_application(), Some("copy")).unwrap(); + + let output = + lance_namespace::schema::convert_json_arrow_field(&plan.output_schema.fields[0]) + .unwrap(); + assert_eq!(output.name(), "copy"); + assert!(output.is_blob_v2()); + } + fn function_binding_schema(title_nullable: bool, body_nullable: bool) -> ArrowSchema { ArrowSchema::new(vec![ ArrowField::new("title", DataType::Utf8, title_nullable),