From d5dac65a21e4fb28ea909388bf49021ac1e4f265 Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Sun, 30 Aug 2026 23:33:30 -0700 Subject: [PATCH] feat: support Blob v2 UDF signatures (#4091) Make Function authoring and declaration planning treat Blob v2 as a scalar semantic type while preserving exact Blob metadata in binding schemas. Covers scalar Blob outputs, expanded named-struct outputs, and whole-result structs with Blob children. --- python/python/lancedb/functions.py | 93 ++++- .../tests/test_first_class_function_slice2.py | 118 ++++++ rust/lancedb/src/function.rs | 3 + rust/lancedb/src/table/computed_columns.rs | 378 ++++++++++++++++-- 4 files changed, 552 insertions(+), 40 deletions(-) diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index 9be63a558..bac6a762f 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -49,6 +49,8 @@ from pydantic import ( model_validator, ) +from .schema import is_blob_v2_field as _is_blob_v2_field + _Int32 = conint(strict=True, ge=-(2**31), le=2**31 - 1) _UInt32 = conint(strict=True, ge=0, le=2**32 - 1) _UInt64 = conint(strict=True, ge=0, le=2**64 - 1) @@ -518,6 +520,7 @@ class RefreshColumnResult(_RemoteValue): _FUNCTION_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_.-]*$") +_FUNCTION_BLOB_V2_TYPE = "blob_v2" _GRAMMAR_PRIMITIVES = ( @@ -581,20 +584,90 @@ def _validate_exact_arrow_field(field: pa.Field) -> None: "unsupported Arrow type for Function signature: field names " "must not be empty" ) - if field.metadata: + if _is_blob_v2_field(field): + if not _has_supported_blob_v2_layout(field): + raise TypeError( + "unsupported Arrow type for Function signature: lance.blob.v2 " + f"requires a supported Blob storage layout, got {field}" + ) + elif field.metadata: raise TypeError( "unsupported Arrow type for Function signature: field metadata " f"is not supported, got {field}" ) +def _has_supported_blob_v2_layout(field: pa.Field) -> bool: + data_type = field.type + if isinstance(data_type, pa.ExtensionType): + data_type = data_type.storage_type + if not pa.types.is_struct(data_type): + return False + + fields = tuple(data_type) + + def matches(spec, compare_nullable) -> bool: + return len(fields) == len(spec) and all( + actual.name == name + and actual.type == expected_type + and (not check_nullable or actual.nullable == nullable) + for actual, (name, expected_type, nullable), check_nullable in zip( + fields, spec, compare_nullable + ) + ) + + logical_minimal = ( + ("data", pa.large_binary(), True), + ("uri", pa.utf8(), True), + ) + logical_full = logical_minimal + ( + ("position", pa.uint64(), True), + ("size", pa.uint64(), True), + ) + prepared = ( + ("kind", pa.uint8(), True), + ("data", pa.large_binary(), True), + ("uri", pa.utf8(), True), + ("blob_id", pa.uint32(), True), + ("blob_size", pa.uint64(), True), + ("position", pa.uint64(), True), + ) + descriptor = ( + ("kind", pa.uint8(), False), + ("position", pa.uint64(), False), + ("size", pa.uint64(), False), + ("blob_id", pa.uint32(), False), + ("blob_uri", pa.utf8(), False), + ) + return ( + matches(logical_minimal, (True, True)) + or matches(logical_full, (True, True, False, False)) + or matches(prepared, (True,) * len(prepared)) + or matches(descriptor, (False,) * len(descriptor)) + ) + + +def _canonical_arrow_field(field: pa.Field) -> str: + _validate_exact_arrow_field(field) + if _is_blob_v2_field(field): + return _FUNCTION_BLOB_V2_TYPE + return _canonical_arrow_type(field.type) + + def _exact_arrow_field(field: pa.Field) -> dict[str, Any]: _validate_exact_arrow_field(field) - return { + 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" + ) + value = { "name": field.name, "nullable": field.nullable, "type": _exact_arrow_type(field.type), } + return value def _exact_arrow_type(data_type: pa.DataType) -> dict[str, Any]: @@ -718,7 +791,11 @@ def _function_output(output: pa.DataType | pa.Field | pa.Schema) -> FunctionOutp if output.metadata: raise TypeError("Function output schema metadata is not supported") fields = tuple(output) - elif isinstance(output, pa.Field) and pa.types.is_struct(output.type): + elif ( + isinstance(output, pa.Field) + and not _is_blob_v2_field(output) + and pa.types.is_struct(output.type) + ): _validate_exact_arrow_field(output) if output.nullable: raise ValueError("Function output must be non-nullable") @@ -740,7 +817,7 @@ def _function_output(output: pa.DataType | pa.Field | pa.Schema) -> FunctionOutp raise ValueError("Function output must be non-nullable") return FunctionOutput( kind="scalar", - arrow_type=_canonical_arrow_type(field.type), + arrow_type=_canonical_arrow_field(field), nullable=False, ) @@ -758,7 +835,7 @@ def _function_output(output: pa.DataType | pa.Field | pa.Schema) -> FunctionOutp fields=tuple( FunctionResultField( name=field.name, - arrow_type=_canonical_arrow_type(field.type), + arrow_type=_canonical_arrow_field(field), nullable=False, ) for field in fields @@ -792,7 +869,7 @@ def _infer_signature( inputs = tuple( FunctionParameter( name=field.name, - arrow_type=_canonical_arrow_type(field.type), + arrow_type=_canonical_arrow_field(field), nullable=field.nullable, ) for field in input_schema @@ -815,7 +892,9 @@ def _infer_signature( inputs.append( FunctionParameter( name=parameter.name, - arrow_type=_canonical_arrow_type(data_type), + arrow_type=_canonical_arrow_field( + pa.field(parameter.name, data_type, nullable=nullable) + ), nullable=nullable, ) ) diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index cf1542b55..518a62bae 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -578,6 +578,124 @@ def test_explicit_arrow_schema_is_deterministic(): assert signature.output.nullable is False +def test_blob_fields_use_the_scalar_function_semantic_type(): + @udf( + input_schema=pa.schema([lancedb.blob("image", nullable=False)]), + output_schema=lancedb.blob("result", nullable=False), + ) + def copy_blob(image): + return image + + signature = copy_blob.registration_request.signature + assert signature.inputs[0].arrow_type == "blob_v2" + assert signature.output.kind == "scalar" + assert signature.output.arrow_type == "blob_v2" + + +def test_named_struct_function_can_include_a_blob_result_field(): + @udf( + input_schema=pa.schema([lancedb.blob("image", nullable=False)]), + output_schema=pa.schema( + [ + lancedb.blob("thumbnail", nullable=False), + pa.field("width", pa.int32(), nullable=False), + ] + ), + ) + def inspect_blob(image): + return {"thumbnail": image, "width": 1} + + output = inspect_blob.registration_request.signature.output + assert output.kind == "named_struct" + assert [(field.name, field.arrow_type) for field in output.fields] == [ + ("thumbnail", "blob_v2"), + ("width", "int32"), + ] + + +def test_metadata_marked_blob_field_uses_the_semantic_type(): + extension = lancedb.blob("image", nullable=False).type + storage = ( + extension.storage_type if isinstance(extension, pa.ExtensionType) else extension + ) + metadata_blob = pa.field( + "image", + storage, + nullable=False, + metadata={"ARROW:extension:name": "lance.blob.v2"}, + ) + + @udf( + input_schema=pa.schema([metadata_blob]), + output_schema=pa.field("size", pa.int64(), nullable=False), + ) + def blob_size(image): + return len(image) + + assert blob_size.registration_request.signature.inputs[0].arrow_type == "blob_v2" + + +def test_blob_marker_rejects_invalid_storage_layout(): + malformed = pa.field( + "image", + pa.int64(), + nullable=False, + metadata={"ARROW:extension:name": "lance.blob.v2"}, + ) + + with pytest.raises(TypeError, match="requires a supported Blob storage layout"): + + @udf( + input_schema=pa.schema([malformed]), + output_schema=pa.field("size", pa.int64(), nullable=False), + ) + def blob_size(image): + 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): + super().__init__(pa.int64(), "test.function.extension") + + def __arrow_ext_serialize__(self): + return b"" + + @classmethod + def __arrow_ext_deserialize__(cls, storage_type, serialized): + return cls() + + nested = pa.field( + "value", + pa.struct([pa.field("extended", TestExtension(), nullable=False)]), + nullable=False, + ) + with pytest.raises(TypeError, match="unsupported Arrow type"): + + @udf( + input_schema=pa.schema([nested]), + output_schema=pa.field("result", pa.int64(), nullable=False), + ) + def extension_value(value): + return value["extended"] + + def test_nested_struct_output_uses_canonical_exact_json(): token = pa.struct( [ diff --git a/rust/lancedb/src/function.rs b/rust/lancedb/src/function.rs index 4b31a4376..79693c031 100644 --- a/rust/lancedb/src/function.rs +++ b/rust/lancedb/src/function.rs @@ -15,6 +15,9 @@ use serde_json::Value; use crate::{Error, Result}; +/// Semantic Function type for a Blob v2 value. +pub const FUNCTION_BLOB_V2_TYPE: &str = "blob_v2"; + fn invalid_json(error: impl std::fmt::Display) -> Error { Error::InvalidInput { message: format!("invalid remote Function JSON: {error}"), diff --git a/rust/lancedb/src/table/computed_columns.rs b/rust/lancedb/src/table/computed_columns.rs index 6dc3ffad5..77e3d0a4d 100644 --- a/rust/lancedb/src/table/computed_columns.rs +++ b/rust/lancedb/src/table/computed_columns.rs @@ -29,14 +29,16 @@ use datafusion_common::{ScalarValue, tree_node::TreeNode}; use datafusion_expr::Expr; use datafusion_physical_plan::PhysicalExpr; use lance::dataset::NewColumnTransform; -use lance_arrow::FieldExt; -use lance_core::datatypes::{BLOB_V2_DESC_FIELD, format_field_path_minimal, parse_field_path}; +use lance_arrow::{ARROW_EXT_NAME_KEY, BLOB_V2_EXT_NAME, FieldExt}; +use lance_core::datatypes::{ + BLOB_V2_DESC_FIELD, BlobV2Layout, format_field_path_minimal, parse_field_path, +}; use lance_datafusion::planner::Planner; use lance_namespace::models::{JsonArrowDataType, JsonArrowField, JsonArrowSchema}; use serde::{Deserialize, Serialize}; use serde_json::Value; -use crate::function::{FunctionApplication, FunctionBinding}; +use crate::function::{FUNCTION_BLOB_V2_TYPE, FunctionApplication, FunctionBinding}; use crate::utils::resolve_arrow_field_path; use crate::{Error, Result}; @@ -581,6 +583,23 @@ fn resolve_field_path<'a>(schema: &'a ArrowSchema, path: &str) -> Result Result { + let is_blob_v2 = field + .metadata + .as_ref() + .and_then(|metadata| metadata.get(ARROW_EXT_NAME_KEY)) + .map(String::as_str) + == Some(BLOB_V2_EXT_NAME); + if is_blob_v2 { + 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() + ))); + } + 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()) } else { @@ -590,6 +609,14 @@ fn canonical_input_arrow_type(field: &JsonArrowField) -> Result { } } +fn has_supported_blob_v2_layout(field: &ArrowField) -> bool { + field.is_blob_v2() + && matches!( + field.data_type(), + DataType::Struct(fields) if BlobV2Layout::classify(fields).is_some() + ) +} + /// `fixed_size_list` -> (`item`, `size`); the comma must sit outside /// any nested `<...>`. fn split_fixed_size_list(raw: &str) -> Option<(&str, i32)> { @@ -669,6 +696,76 @@ fn parse_output_arrow_type(raw: &str) -> Result { Ok(data_type) } +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), + ])) + .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)?, + )) +} + +fn function_output_field_matches(expected: &ArrowField, actual: &ArrowField) -> bool { + expected.name() == actual.name() + && expected.is_nullable() == actual.is_nullable() + && if expected.is_blob_v2() { + has_supported_blob_v2_layout(expected) && has_supported_blob_v2_layout(actual) + } else { + function_output_type_matches(expected.data_type(), actual.data_type()) + } +} + +fn function_output_type_matches(expected: &DataType, actual: &DataType) -> bool { + if expected == actual { + return true; + } + match (expected, actual) { + (DataType::Struct(expected), DataType::Struct(actual)) => { + expected.len() == actual.len() + && expected + .iter() + .zip(actual) + .all(|(expected, actual)| function_output_field_matches(expected, actual)) + } + (DataType::List(expected), DataType::List(actual)) + | (DataType::LargeList(expected), DataType::LargeList(actual)) => { + function_output_field_matches(expected, actual) + } + ( + DataType::FixedSizeList(expected, expected_size), + DataType::FixedSizeList(actual, actual_size), + ) => expected_size == actual_size && function_output_field_matches(expected, actual), + (DataType::Map(expected, expected_sorted), DataType::Map(actual, actual_sorted)) => { + expected_sorted == actual_sorted && function_output_field_matches(expected, actual) + } + _ => false, + } +} + +fn function_output_type_has_blob(data_type: &DataType) -> bool { + match data_type { + DataType::Struct(fields) => fields + .iter() + .any(|field| field.is_blob_v2() || function_output_type_has_blob(field.data_type())), + DataType::List(field) + | DataType::LargeList(field) + | DataType::FixedSizeList(field, _) + | DataType::Map(field, _) => { + field.is_blob_v2() || function_output_type_has_blob(field.data_type()) + } + _ => false, + } +} + 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() { @@ -749,10 +846,18 @@ 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) - .map_err(|e| invalid_function(format!("invalid Function output type: {e}")))?; - if field.data_type() != &expected_type { + let (type_matches, has_semantic_blob) = if output.arrow_type == FUNCTION_BLOB_V2_TYPE { + (has_supported_blob_v2_layout(field), true) + } else { + let expected_type = parse_output_arrow_type(&output.arrow_type)?; + let expected_type = lance_namespace::schema::convert_json_arrow_type(&expected_type) + .map_err(|e| invalid_function(format!("invalid Function output type: {e}")))?; + ( + function_output_type_matches(&expected_type, field.data_type()), + function_output_type_has_blob(&expected_type), + ) + }; + if !type_matches { return Err(invalid_function(format!( "Function output '{}' type no longer matches binding '{}'", output.output_name, @@ -781,15 +886,21 @@ fn ensure_binding_matches_schema(schema: &ArrowSchema, binding: &FunctionBinding binding.binding_id() ))); } - output_fields.push(ArrowField::new( - field.name().clone(), - field.data_type().clone(), - true, - )); - } - let output_schema = - lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(output_fields)) + if has_semantic_blob { + output_fields.push(function_output_field( + field.name(), + true, + &output.arrow_type, + )?); + } else { + let json = lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![ + ArrowField::new(field.name().clone(), field.data_type().clone(), true), + ])) .map_err(|e| invalid_function(format!("invalid Function output schema: {e}")))?; + output_fields.push(json.fields.into_iter().next().unwrap()); + } + } + let output_schema = JsonArrowSchema::new(output_fields); let output_schema = serde_json::to_value(output_schema).map_err(|e| { invalid_function(format!( "could not encode exact Function output schema: {e}" @@ -918,16 +1029,15 @@ 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(|| { - invalid_function("scalar Function output is missing its Arrow type") - })?)?; + let arrow_type = 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(function_output_field(name, true, arrow_type)?); } "named_struct" => { if output.fields.is_empty() { @@ -971,13 +1081,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| function_output_field(&field.name, false, &field.arrow_type)) .collect::>>()?; let mut data_type = JsonArrowDataType::new("struct".to_string()); data_type.fields = Some(fields); @@ -1004,11 +1108,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(function_output_field(name, true, &field.arrow_type)?); } } } @@ -1645,7 +1745,7 @@ mod tests { } use arrow_array::record_batch; - use arrow_schema::DataType; + use arrow_schema::{DataType, TimeUnit}; use futures::TryStreamExt; use lance::dataset::ColumnAlteration; @@ -2606,6 +2706,73 @@ mod tests { .unwrap() } + fn blob_application(output: &str) -> FunctionApplication { + FunctionApplication::from_json(&format!( + r#"{{ + "function":{{"name":"blob_features","version":"fv_blob"}}, + "inputs":[ + {{"parameter":"image","kind":"column","value":{{"path":"image"}}}} + ], + "output":{output} + }}"# + )) + .unwrap() + } + + fn binding_from_plan(plan: &FunctionDeclarationPlan) -> FunctionBinding { + let inputs = plan + .input_bindings + .iter() + .enumerate() + .map(|(index, input)| { + serde_json::json!({ + "parameter": input.parameter, + "field_id": index, + "field_path": input.field_path, + "arrow_type": input.arrow_type, + "nullable": input.nullable, + }) + }) + .collect::>(); + let outputs = plan + .outputs + .iter() + .zip(&plan.output_schema.fields) + .enumerate() + .map(|(index, (output, field))| { + serde_json::json!({ + "result_field": output.result_field, + "output_name": output.output_name, + "output_field_id": 100 + index, + "output_ordinal": output.output_ordinal, + "arrow_type": canonical_input_arrow_type(field).unwrap(), + "nullable": false, + }) + }) + .collect::>(); + FunctionBinding::from_json( + &serde_json::json!({ + "binding_id": "fb_blob", + "function": plan.application.function(), + "inputs": inputs, + "outputs": outputs, + "input_schema": plan.input_schema, + "output_schema": plan.output_schema, + }) + .to_string(), + ) + .unwrap() + } + + fn full_blob_field(name: &str, nullable: bool) -> ArrowField { + ArrowField::new( + name, + DataType::Struct(lance_core::datatypes::BLOB_V2_LOGICAL_FIELDS.clone()), + nullable, + ) + .with_metadata(crate::blob(name, nullable).metadata().clone()) + } + fn function_binding_schema(title_nullable: bool, body_nullable: bool) -> ArrowSchema { ArrowSchema::new(vec![ ArrowField::new("title", DataType::Utf8, title_nullable), @@ -2879,6 +3046,151 @@ mod tests { ); } + #[test] + fn test_blob_function_plans_semantic_input_and_scalar_output() { + let schema = ArrowSchema::new(vec![crate::blob("image", false)]); + let application = + blob_application(r#"{"kind":"scalar","arrow_type":"blob_v2","nullable":false}"#); + let plan = plan_function_application(&schema, &application, Some("thumbnail")).unwrap(); + + assert_eq!(plan.input_bindings[0].arrow_type, FUNCTION_BLOB_V2_TYPE); + let input_schema = + lance_namespace::schema::convert_json_arrow_schema(&plan.input_schema).unwrap(); + assert!(input_schema.field(0).is_blob_v2()); + let output_schema = + lance_namespace::schema::convert_json_arrow_schema(&plan.output_schema).unwrap(); + assert!(output_schema.field(0).is_blob_v2()); + } + + #[test] + fn test_blob_scalar_binding_accepts_full_logical_layout() { + let input = crate::blob("image", false); + let application = + blob_application(r#"{"kind":"scalar","arrow_type":"blob_v2","nullable":false}"#); + let plan = plan_function_application( + &ArrowSchema::new(vec![input.clone()]), + &application, + Some("thumbnail"), + ) + .unwrap(); + let binding = binding_from_plan(&plan); + let mut metadata = full_blob_field("thumbnail", true).metadata().clone(); + metadata.extend(function_computed_column_metadata( + binding.binding_id(), + 0, + &["image".into()], + )); + let output = full_blob_field("thumbnail", true).with_metadata(metadata); + + ensure_binding_matches_schema(&ArrowSchema::new(vec![input, output]), &binding).unwrap(); + } + + #[test] + fn test_blob_binding_rejects_marker_on_invalid_storage_layout() { + let input = crate::blob("image", false); + let application = + blob_application(r#"{"kind":"scalar","arrow_type":"blob_v2","nullable":false}"#); + let plan = plan_function_application( + &ArrowSchema::new(vec![input.clone()]), + &application, + Some("thumbnail"), + ) + .unwrap(); + let binding = binding_from_plan(&plan); + let malformed = ArrowField::new("thumbnail", DataType::Int64, true) + .with_metadata(crate::blob("thumbnail", true).metadata().clone()); + + ensure_binding_matches_schema(&ArrowSchema::new(vec![input, malformed]), &binding) + .unwrap_err(); + } + + #[test] + fn test_blob_input_rejects_marker_on_invalid_storage_layout() { + let malformed = ArrowField::new("image", DataType::Int64, false) + .with_metadata(crate::blob("image", false).metadata().clone()); + let application = + blob_application(r#"{"kind":"scalar","arrow_type":"blob_v2","nullable":false}"#); + + plan_function_application( + &ArrowSchema::new(vec![malformed]), + &application, + Some("thumbnail"), + ) + .unwrap_err(); + } + + #[test] + fn test_non_blob_input_does_not_require_json_round_trip() { + let json = lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![ + ArrowField::new("event_time", DataType::Time64(TimeUnit::Microsecond), false), + ])) + .unwrap(); + + assert_eq!( + canonical_input_arrow_type(&json.fields[0]).unwrap(), + "time64" + ); + } + + #[test] + fn test_blob_named_struct_plans_expanded_and_whole_outputs() { + let schema = ArrowSchema::new(vec![crate::blob("image", false)]); + let application = blob_application( + r#"{"kind":"named_struct","fields":[ + {"name":"thumbnail","arrow_type":"blob_v2","nullable":false}, + {"name":"width","arrow_type":"int32","nullable":false} + ]}"#, + ); + + let expanded = plan_function_application(&schema, &application, None).unwrap(); + let expanded_schema = + lance_namespace::schema::convert_json_arrow_schema(&expanded.output_schema).unwrap(); + assert!(expanded_schema.field(0).is_blob_v2()); + assert_eq!(expanded_schema.field(1).data_type(), &DataType::Int32); + + let whole = plan_function_application(&schema, &application, Some("analysis")).unwrap(); + let whole_schema = + lance_namespace::schema::convert_json_arrow_schema(&whole.output_schema).unwrap(); + let DataType::Struct(fields) = whole_schema.field(0).data_type() else { + panic!("whole Function output should be a struct"); + }; + assert!(fields[0].is_blob_v2()); + assert_eq!(fields[1].data_type(), &DataType::Int32); + } + + #[test] + fn test_blob_whole_struct_binding_accepts_full_logical_layout() { + let input = crate::blob("image", false); + let application = blob_application( + r#"{"kind":"named_struct","fields":[ + {"name":"thumbnail","arrow_type":"blob_v2","nullable":false}, + {"name":"width","arrow_type":"int32","nullable":false} + ]}"#, + ); + let plan = plan_function_application( + &ArrowSchema::new(vec![input.clone()]), + &application, + Some("analysis"), + ) + .unwrap(); + let binding = binding_from_plan(&plan); + + let output = ArrowField::new( + "analysis", + DataType::Struct(Fields::from(vec![ + full_blob_field("thumbnail", false), + ArrowField::new("width", DataType::Int32, false), + ])), + true, + ) + .with_metadata(function_computed_column_metadata( + binding.binding_id(), + 0, + &["image".into()], + )); + ensure_binding_matches_schema(&ArrowSchema::new(vec![input, output]), &binding).unwrap(); + } + #[test] fn test_function_mapping_and_sibling_collisions_fail_before_request() { let unknown = named_struct_application(r#"{"missing":"renamed"}"#);