diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index 445d31027..7f5776b32 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -697,6 +697,7 @@ _GRAMMAR_PRIMITIVES = ( (pa.string(), "utf8"), (pa.large_string(), "large_utf8"), (pa.binary(), "binary"), + (pa.large_binary(), "large_binary"), (pa.date32(), "date32"), (pa.date64(), "date64"), ) @@ -846,11 +847,7 @@ def _exact_blob_storage_type(field: pa.Field) -> dict[str, Any]: { "name": child.name, "nullable": child.nullable, - "type": ( - {"type": "large_binary"} - if pa.types.is_large_binary(child.type) - else _exact_arrow_type(child.type) - ), + "type": _exact_arrow_type(child.type), } for child in storage ], diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index 975310ed4..9ee3e388a 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -622,7 +622,6 @@ def test_canonical_arrow_type_prefers_the_compact_grammar(): for outside in [ pa.timestamp("us"), pa.decimal128(10, 2), - pa.large_binary(), pa.binary(4), pa.duration("s"), pa.list_(pa.float32(), 0), @@ -1280,6 +1279,28 @@ def test_explicit_large_utf8_schemas_use_the_canonical_function_name(): } +def test_large_binary_schemas_use_the_canonical_function_name(): + payload = pa.field("data", pa.large_binary(), nullable=True) + + @udf( + input_schema=pa.schema([payload]), + output_schema=pa.field("result", pa.large_binary(), nullable=False), + ) + def preserve(data): + return data + + signature = preserve.registration_request.signature + assert signature.inputs[0].arrow_type == "large_binary" + assert signature.output.arrow_type == "large_binary" + cells = pa.large_list(pa.field("item", pa.struct([payload]), nullable=True)) + assert json.loads(_canonical_arrow_type(cells))["fields"][0]["type"] == { + "type": "struct", + "fields": [ + {"name": "data", "nullable": True, "type": {"type": "large_binary"}} + ], + } + + def test_nested_struct_output_uses_canonical_exact_json(): token = pa.struct( [ diff --git a/rust/lancedb/tests/fixtures/first_class_functions/v1/arrow_types.json b/rust/lancedb/tests/fixtures/first_class_functions/v1/arrow_types.json index 38d17821b..26e8db430 100644 --- a/rust/lancedb/tests/fixtures/first_class_functions/v1/arrow_types.json +++ b/rust/lancedb/tests/fixtures/first_class_functions/v1/arrow_types.json @@ -90,6 +90,12 @@ "type": "binary" } }, + { + "arrow_type": "large_binary", + "json": { + "type": "large_binary" + } + }, { "arrow_type": "date32", "json": { @@ -192,6 +198,21 @@ ] } }, + { + "arrow_type": "large_list", + "json": { + "type": "large_list", + "fields": [ + { + "name": "item", + "nullable": false, + "type": { + "type": "large_binary" + } + } + ] + } + }, { "arrow_type": "large_list", "json": {