diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index 3bdd117cb..2815f7f17 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -521,6 +521,12 @@ class RefreshColumnResult(_RemoteValue): _FUNCTION_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_.-]*$") _FUNCTION_BLOB_V2_TYPE = "blob_v2" +_ARROW_EXTENSION_NAME_KEY = "ARROW:extension:name" +_BLOB_V2_EXTENSION_NAME = "lance.blob.v2" +_NESTED_BLOB_COLLECTION_ERROR = ( + "unsupported Arrow type for Function signature: Blob v2 fields nested under " + "collection types are not supported" +) _GRAMMAR_PRIMITIVES = ( @@ -591,6 +597,19 @@ def _validate_exact_arrow_field(field: pa.Field) -> None: "unsupported Arrow type for Function signature: lance.blob.v2 " f"requires a supported Blob storage layout, got {field}" ) + metadata = { + (key.decode() if isinstance(key, bytes) else key): ( + value.decode() if isinstance(value, bytes) else value + ) + for key, value in (field.metadata or {}).items() + } + if metadata and metadata != { + _ARROW_EXTENSION_NAME_KEY: _BLOB_V2_EXTENSION_NAME + }: + raise TypeError( + "unsupported Arrow type for Function signature: lance.blob.v2 " + "field metadata must contain only its canonical extension marker" + ) elif field.metadata: raise TypeError( "unsupported Arrow type for Function signature: field metadata " @@ -655,23 +674,84 @@ def _canonical_arrow_field(field: pa.Field) -> str: return _canonical_arrow_type(field.type) -def _exact_arrow_field(field: pa.Field) -> dict[str, Any]: +def _blob_storage_type(field: pa.Field) -> pa.DataType: + data_type = field.type + if isinstance(data_type, pa.ExtensionType): + return data_type.storage_type + return data_type + + +def _exact_blob_storage_type(field: pa.Field) -> dict[str, Any]: + storage = _blob_storage_type(field) + if not pa.types.is_struct(storage): + raise TypeError( + "unsupported Arrow type for Function signature: lance.blob.v2 " + "requires struct storage" + ) + return { + "type": "struct", + "fields": [ + { + "name": child.name, + "nullable": child.nullable, + "type": ( + {"type": "large_binary"} + if pa.types.is_large_binary(child.type) + else _exact_arrow_type(child.type) + ), + } + for child in storage + ], + } + + +def _data_type_has_blob_v2(data_type: pa.DataType) -> bool: + if pa.types.is_struct(data_type): + return any( + _is_blob_v2_field(field) or _data_type_has_blob_v2(field.type) + for field in data_type + ) + if ( + pa.types.is_list(data_type) + or pa.types.is_large_list(data_type) + or pa.types.is_fixed_size_list(data_type) + ): + field = data_type.value_field + return _is_blob_v2_field(field) or _data_type_has_blob_v2(field.type) + if pa.types.is_map(data_type): + return any( + _is_blob_v2_field(field) or _data_type_has_blob_v2(field.type) + for field in (data_type.key_field, data_type.item_field) + ) + return False + + +def _exact_arrow_field( + field: pa.Field, *, inside_collection: bool = False +) -> dict[str, Any]: _validate_exact_arrow_field(field) if _is_blob_v2_field(field): - raise TypeError( - "unsupported Arrow type for Function signature: nested Blob v2 " - "fields are not supported; declare Blob parameters or named result " - "fields directly" - ) + if inside_collection: + raise TypeError(_NESTED_BLOB_COLLECTION_ERROR) + return { + "name": field.name, + "nullable": field.nullable, + "type": _exact_blob_storage_type(field), + "metadata": { + _ARROW_EXTENSION_NAME_KEY: _BLOB_V2_EXTENSION_NAME, + }, + } value = { "name": field.name, "nullable": field.nullable, - "type": _exact_arrow_type(field.type), + "type": _exact_arrow_type(field.type, inside_collection=inside_collection), } return value -def _exact_arrow_type(data_type: pa.DataType) -> dict[str, Any]: +def _exact_arrow_type( + data_type: pa.DataType, *, inside_collection: bool = False +) -> dict[str, Any]: for candidate, name in _GRAMMAR_PRIMITIVES: if data_type == candidate: return {"type": name} @@ -685,7 +765,10 @@ def _exact_arrow_type(data_type: pa.DataType) -> dict[str, Any]: ) return { "type": "struct", - "fields": [_exact_arrow_field(field) for field in fields], + "fields": [ + _exact_arrow_field(field, inside_collection=inside_collection) + for field in fields + ], } if ( pa.types.is_list(data_type) @@ -710,11 +793,15 @@ def _exact_arrow_type(data_type: pa.DataType) -> dict[str, Any]: if pa.types.is_large_list(data_type) else "fixed_size_list" ), - "fields": [_exact_arrow_field(data_type.value_field)], + "fields": [ + _exact_arrow_field(data_type.value_field, inside_collection=True) + ], } if pa.types.is_fixed_size_list(data_type): value["length"] = data_type.list_size return value + if pa.types.is_map(data_type) and _data_type_has_blob_v2(data_type): + raise TypeError(_NESTED_BLOB_COLLECTION_ERROR) raise TypeError(f"unsupported Arrow type for Function signature: {data_type}") diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index b1308a2bc..d8a9aeebf 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -668,6 +668,167 @@ def test_blob_fields_use_the_scalar_function_semantic_type(): assert signature.output.arrow_type == "blob_v2" +def test_whole_named_struct_function_can_include_a_blob_result_field(): + @udf( + input_schema=pa.schema([lancedb.blob("image", nullable=False)]), + output_schema=pa.field( + "payload", + pa.struct( + [ + pa.field("mime_type", pa.string(), nullable=False), + lancedb.blob("image", nullable=False), + ] + ), + nullable=False, + ), + ) + def inspect_blob(image): + return {"mime_type": "image/png", "image": image} + + output = inspect_blob.registration_request.signature.output + assert output.kind == "named_struct" + assert [(field.name, field.arrow_type) for field in output.fields] == [ + ("mime_type", "utf8"), + ("image", "blob_v2"), + ] + + +def test_struct_blob_signature_fields_preserve_exact_metadata_and_nullability(): + nested_input = pa.field( + "payload", + pa.struct( + [ + pa.field("mime_type", pa.string(), nullable=False), + pa.field( + "nested", + pa.struct([lancedb.blob("image", nullable=True)]), + nullable=True, + ), + ] + ), + nullable=True, + ) + nested_output = pa.field( + "result", + pa.struct( + [ + pa.field("mime_type", pa.string(), nullable=False), + pa.field( + "nested", + pa.struct([lancedb.blob("image", nullable=True)]), + nullable=False, + ), + ] + ), + nullable=False, + ) + + @udf(input_schema=pa.schema([nested_input]), output_schema=nested_output) + def copy_payload(payload): + return payload + + signature = copy_payload.registration_request.signature + input_type = json.loads(signature.inputs[0].arrow_type) + assert input_type["fields"][1]["nullable"] is True + input_blob = input_type["fields"][1]["type"]["fields"][0] + assert input_blob["nullable"] is True + assert input_blob["metadata"] == {"ARROW:extension:name": "lance.blob.v2"} + + assert signature.output.kind == "named_struct" + nested_result = next( + field for field in signature.output.fields if field.name == "nested" + ) + output_type = json.loads(nested_result.arrow_type) + output_blob = output_type["fields"][0] + assert output_blob["nullable"] is True + assert output_blob["metadata"] == {"ARROW:extension:name": "lance.blob.v2"} + + +def test_struct_blob_signature_supports_multiple_struct_levels(): + recursive = pa.field( + "value", + pa.struct( + [ + pa.field( + "level_1", + pa.struct( + [ + pa.field( + "level_2", + pa.struct([lancedb.blob("image", nullable=False)]), + nullable=False, + ) + ] + ), + nullable=False, + ) + ] + ), + nullable=False, + ) + + @udf( + input_schema=pa.schema([recursive]), + output_schema=pa.field("size", pa.int64(), nullable=False), + ) + def blob_size(value): + return len(value["level_1"]["level_2"]["image"]) + + encoded = json.loads(blob_size.registration_request.signature.inputs[0].arrow_type) + blob = encoded["fields"][0]["type"]["fields"][0]["type"]["fields"][0] + assert blob["metadata"]["ARROW:extension:name"] == "lance.blob.v2" + + +@pytest.mark.parametrize( + "data_type", + [ + pa.list_(lancedb.blob("item", nullable=False)), + pa.large_list(lancedb.blob("item", nullable=False)), + pa.list_(lancedb.blob("item", nullable=False), 2), + pa.map_(pa.string(), lancedb.blob("value", nullable=False).type), + ], +) +def test_blob_signature_rejects_collection_ancestors(data_type): + with pytest.raises( + TypeError, + match="Blob v2 fields nested under collection types are not supported", + ): + + @udf( + input_schema=pa.schema([pa.field("value", data_type, nullable=False)]), + output_schema=pa.field("size", pa.int64(), nullable=False), + ) + def blob_size(value): + return len(value) + + +def test_blob_signature_rejects_collection_below_a_struct(): + nested = pa.field( + "value", + pa.struct( + [ + pa.field( + "images", + pa.list_(lancedb.blob("item", nullable=False)), + nullable=False, + ) + ] + ), + nullable=False, + ) + with pytest.raises( + TypeError, + match="Blob v2 fields nested under collection types are not supported", + ): + + @udf( + input_schema=pa.schema([nested]), + output_schema=pa.field("size", pa.int64(), nullable=False), + ) + def blob_size(value): + return len(value["images"]) + + def test_named_struct_function_can_include_a_blob_result_field(): @udf( input_schema=pa.schema([lancedb.blob("image", nullable=False)]), @@ -729,22 +890,6 @@ def test_blob_marker_rejects_invalid_storage_layout(): return len(image) -def test_nested_blob_signature_field_has_a_clear_error(): - nested = pa.field( - "value", - pa.struct([lancedb.blob("image", nullable=False)]), - nullable=False, - ) - with pytest.raises(TypeError, match="nested Blob v2 fields are not supported"): - - @udf( - input_schema=pa.schema([nested]), - output_schema=pa.field("size", pa.int64(), nullable=False), - ) - def blob_size(value): - return len(value["image"]) - - def test_nested_non_blob_extension_is_not_silently_unwrapped(): class TestExtension(pa.ExtensionType): def __init__(self): diff --git a/rust/lancedb/src/table/computed_columns.rs b/rust/lancedb/src/table/computed_columns.rs index 77e3d0a4d..085876a81 100644 --- a/rust/lancedb/src/table/computed_columns.rs +++ b/rust/lancedb/src/table/computed_columns.rs @@ -589,16 +589,13 @@ fn canonical_input_arrow_type(field: &JsonArrowField) -> Result { .and_then(|metadata| metadata.get(ARROW_EXT_NAME_KEY)) .map(String::as_str) == Some(BLOB_V2_EXT_NAME); - if is_blob_v2 { + if is_blob_v2 || field.r#type.fields.is_some() { let arrow_field = lance_namespace::schema::convert_json_arrow_field(field) .map_err(|e| invalid_function(format!("invalid Function input field: {e}")))?; - if !has_supported_blob_v2_layout(&arrow_field) { - return Err(invalid_function(format!( - "Function input '{}' has an invalid Blob v2 storage layout", - arrow_field.name() - ))); + validate_function_blob_nesting(&arrow_field, false)?; + if is_blob_v2 { + return Ok(FUNCTION_BLOB_V2_TYPE.to_string()); } - return Ok(FUNCTION_BLOB_V2_TYPE.to_string()); } if field.r#type.fields.is_none() && field.r#type.length.is_none() { Ok(field.r#type.r#type.clone()) @@ -617,6 +614,34 @@ fn has_supported_blob_v2_layout(field: &ArrowField) -> bool { ) } +fn validate_function_blob_nesting(field: &ArrowField, inside_collection: bool) -> Result<()> { + if field.is_blob_v2() { + if inside_collection { + return Err(invalid_function(format!( + "Function field '{}' nests Blob v2 under a collection, which Function signatures do not support", + field.name() + ))); + } + if !has_supported_blob_v2_layout(field) { + return Err(invalid_function(format!( + "Function field '{}' has an invalid Blob v2 storage layout", + field.name() + ))); + } + return Ok(()); + } + match field.data_type() { + DataType::Struct(fields) => fields + .iter() + .try_for_each(|field| validate_function_blob_nesting(field, inside_collection)), + DataType::List(field) + | DataType::LargeList(field) + | DataType::FixedSizeList(field, _) + | DataType::Map(field, _) => validate_function_blob_nesting(field, true), + _ => Ok(()), + } +} + /// `fixed_size_list` -> (`item`, `size`); the comma must sit outside /// any nested `<...>`. fn split_fixed_size_list(raw: &str) -> Option<(&str, i32)> { @@ -697,21 +722,22 @@ fn parse_output_arrow_type(raw: &str) -> Result { } fn function_output_field(name: &str, nullable: bool, raw: &str) -> Result { - if raw == FUNCTION_BLOB_V2_TYPE { - return lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![ - crate::blob(name, nullable), - ])) + let field = if raw == FUNCTION_BLOB_V2_TYPE { + lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![crate::blob( + name, nullable, + )])) .map_err(|e| invalid_function(format!("could not encode Blob v2 output field: {e}")))? .fields .into_iter() .next() - .ok_or_else(|| invalid_function("Blob v2 output field is missing")); - } - Ok(JsonArrowField::new( - name.to_string(), - nullable, - parse_output_arrow_type(raw)?, - )) + .ok_or_else(|| invalid_function("Blob v2 output field is missing"))? + } else { + JsonArrowField::new(name.to_string(), nullable, parse_output_arrow_type(raw)?) + }; + let arrow_field = lance_namespace::schema::convert_json_arrow_field(&field) + .map_err(|e| invalid_function(format!("invalid Function output field: {e}")))?; + validate_function_blob_nesting(&arrow_field, false)?; + Ok(field) } fn function_output_field_matches(expected: &ArrowField, actual: &ArrowField) -> bool { @@ -2719,6 +2745,28 @@ mod tests { .unwrap() } + fn exact_arrow_type(field: ArrowField) -> String { + let json = + lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![field])).unwrap(); + serde_json::to_string(json.fields[0].r#type.as_ref()).unwrap() + } + + fn single_input_application(path: &str) -> FunctionApplication { + FunctionApplication::from_json( + &serde_json::json!({ + "function": {"name": "inspect", "version": "fv_nested_blob"}, + "inputs": [{ + "parameter": "value", + "kind": "column", + "value": {"path": path} + }], + "output": {"kind": "scalar", "arrow_type": "int64", "nullable": false} + }) + .to_string(), + ) + .unwrap() + } + fn binding_from_plan(plan: &FunctionDeclarationPlan) -> FunctionBinding { let inputs = plan .input_bindings @@ -3158,6 +3206,128 @@ mod tests { assert_eq!(fields[1].data_type(), &DataType::Int32); } + #[test] + fn test_struct_blob_input_preserves_exact_schema_and_nullability() { + let payload = ArrowField::new( + "payload", + DataType::Struct(Fields::from(vec![ + ArrowField::new("mime_type", DataType::Utf8, false), + ArrowField::new( + "nested", + DataType::Struct(Fields::from(vec![crate::blob("image", true)])), + true, + ), + ])), + true, + ); + let plan = plan_function_application( + &ArrowSchema::new(vec![payload]), + &single_input_application("payload"), + Some("size"), + ) + .unwrap(); + + let declared: JsonArrowDataType = + serde_json::from_str(&plan.input_bindings[0].arrow_type).unwrap(); + let DataType::Struct(fields) = + lance_namespace::schema::convert_json_arrow_type(&declared).unwrap() + else { + panic!("expected a struct Function input") + }; + assert!(fields[1].is_nullable()); + let DataType::Struct(nested) = fields[1].data_type() else { + panic!("expected a recursive struct Function input") + }; + assert!(nested[0].is_blob_v2()); + assert!(nested[0].is_nullable()); + + let exact = lance_namespace::schema::convert_json_arrow_schema(&plan.input_schema).unwrap(); + let DataType::Struct(fields) = exact.field(0).data_type() else { + panic!("expected exact input schema to retain the struct") + }; + let DataType::Struct(nested) = fields[1].data_type() else { + panic!("expected exact input schema to retain the nested struct") + }; + assert!(nested[0].is_blob_v2()); + } + + #[test] + fn test_recursive_blob_result_plans_one_whole_named_struct_column() { + let details_type = exact_arrow_type(ArrowField::new( + "details", + DataType::Struct(Fields::from(vec![crate::blob("image", true)])), + false, + )); + let application = FunctionApplication::from_json( + &serde_json::json!({ + "function": {"name": "inspect", "version": "fv_nested_blob"}, + "inputs": [], + "output": { + "kind": "named_struct", + "fields": [ + {"name": "mime_type", "arrow_type": "utf8", "nullable": false}, + {"name": "details", "arrow_type": details_type, "nullable": false} + ] + } + }) + .to_string(), + ) + .unwrap(); + let plan = plan_function_application(&ArrowSchema::empty(), &application, Some("payload")) + .unwrap(); + + assert_eq!(plan.outputs.len(), 1); + assert_eq!(plan.outputs[0].result_field, WHOLE_RESULT_FIELD); + let schema = + lance_namespace::schema::convert_json_arrow_schema(&plan.output_schema).unwrap(); + assert_eq!(schema.field(0).name(), "payload"); + let DataType::Struct(fields) = schema.field(0).data_type() else { + panic!("whole named result must be one struct column") + }; + assert_eq!( + fields.iter().map(|field| field.name()).collect::>(), + ["mime_type", "details"] + ); + let DataType::Struct(details) = fields[1].data_type() else { + panic!("expected recursive result struct") + }; + assert!(details[0].is_blob_v2()); + assert!(!fields.iter().any(|field| field.name() == "payload")); + } + + #[test] + fn test_blob_children_under_collections_are_rejected() { + let collections = vec![ + DataType::List(Arc::new(crate::blob("item", false))), + DataType::LargeList(Arc::new(crate::blob("item", false))), + DataType::FixedSizeList(Arc::new(crate::blob("item", false)), 2), + DataType::Map( + Arc::new(ArrowField::new( + "entries", + DataType::Struct(Fields::from(vec![ + ArrowField::new("key", DataType::Utf8, false), + crate::blob("value", false), + ])), + false, + )), + false, + ), + ]; + for data_type in collections { + let schema = ArrowSchema::new(vec![ArrowField::new("value", data_type, false)]); + let error = plan_function_application( + &schema, + &single_input_application("value"), + Some("size"), + ) + .unwrap_err(); + assert!( + error.to_string().contains("under a collection"), + "got: {error}" + ); + } + } + #[test] fn test_blob_whole_struct_binding_accepts_full_logical_layout() { let input = crate::blob("image", false);