diff --git a/python/python/lancedb/schema.py b/python/python/lancedb/schema.py index 8bec17058..342053d6a 100644 --- a/python/python/lancedb/schema.py +++ b/python/python/lancedb/schema.py @@ -32,8 +32,8 @@ _resolved_blob_type = None class _FallbackBlobType(pa.ExtensionType): """lance.blob.v2 extension type used when pylance is not installed.""" - def __init__(self) -> None: - pa.ExtensionType.__init__(self, _BLOB_V2_STORAGE_TYPE, _BLOB_EXTENSION_NAME) + def __init__(self, storage_type: pa.DataType = _BLOB_V2_STORAGE_TYPE) -> None: + pa.ExtensionType.__init__(self, storage_type, _BLOB_EXTENSION_NAME) def __arrow_ext_serialize__(self) -> bytes: return b"" @@ -42,7 +42,7 @@ class _FallbackBlobType(pa.ExtensionType): def __arrow_ext_deserialize__( cls, storage_type: pa.DataType, serialized: bytes ) -> "_FallbackBlobType": - return cls() + return cls(storage_type) def __reduce__(self): return type(self).__arrow_ext_deserialize__, ( @@ -183,6 +183,24 @@ def _resolve_blob_type(): return _resolved_blob_type +def _restore_blob_extensions(schema: pa.Schema) -> pa.Schema: + """Restore blob fields read before their extension type was registered.""" + unrestored = _collect_blob_paths( + schema, + lambda field: ( + not isinstance(field.type, pa.ExtensionType) + and _metadata_marks_blob_v2(field.metadata or {}) + ), + ) + if not unrestored: + return schema + + _resolve_blob_type() + # Arrow's deserializer uses the registered class and preserves nested fields, + # storage layouts, nullability, and unrelated metadata. Only the schema is read. + return pa.ipc.read_schema(schema.serialize()) + + def blob( name: str, nullable: bool = True, diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 59c88027d..98a841bd0 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -110,6 +110,7 @@ from .util import ( ) from .index import lang_mapping from .schema import ( + _restore_blob_extensions, blob_v2_column_paths, is_blob_v2_field, row_addressable_blob_v2_paths, @@ -6023,10 +6024,10 @@ class AsyncTable: async def schema(self) -> pa.Schema: """The [Arrow Schema](https://arrow.apache.org/docs/python/api/datatypes.html#) - of this Table + of this Table. Blob v2 fields retain their Arrow extension type. """ - return await self._inner.schema() + return _restore_blob_extensions(await self._inner.schema()) async def embedding_functions(self) -> Dict[str, EmbeddingFunctionConfig]: """ diff --git a/python/python/tests/test_blob.py b/python/python/tests/test_blob.py index 96bf88c24..b1da5472a 100644 --- a/python/python/tests/test_blob.py +++ b/python/python/tests/test_blob.py @@ -231,6 +231,78 @@ def test_blob_type_works_without_pylance(): assert result.returncode == 0, result.stderr +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("with_pylance", [False, True]) +def test_local_schema_restores_blob_extensions(tmp_path, use_async, with_pylance): + db = lancedb.connect(tmp_path) + db.create_table( + "images", + schema=pa.schema( + [ + pa.field("id", pa.int64()), + lancedb.blob("image", nullable=False, inline_size_threshold=8192), + ], + metadata={"schema-key": "schema-value"}, + ), + ) + + # Reopening must restore the type before blob() or pylance registers it in + # the reader process, including when pylance is not installed. + script = textwrap.dedent( + """\ + import asyncio + import importlib.abc + import sys + import pyarrow as pa + + if sys.argv[3] == "False": + class NoLance(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path, target=None): + if fullname == "lance" or fullname.startswith("lance."): + raise ModuleNotFoundError(fullname, name="lance") + sys.meta_path.insert(0, NoLance()) + import lancedb + + assert "lance.blob" not in sys.modules + + async def read_schema(): + db = await lancedb.connect_async(sys.argv[1]) + table = await db.open_table("images") + return await table.schema() + + if sys.argv[2] == "True": + schema = asyncio.run(read_schema()) + else: + schema = lancedb.connect(sys.argv[1]).open_table("images").schema + + image = schema.field("image") + assert isinstance(image.type, pa.ExtensionType), image + assert image.type == lancedb.blob("x").type + assert type(image.type) is lancedb.BlobType + assert not image.nullable + assert image.metadata == { + b"lance-encoding:blob-inline-size-threshold": b"8192" + } + assert schema.field("id") == pa.field("id", pa.int64()) + assert schema.metadata == {b"schema-key": b"schema-value"} + """ + ) + result = subprocess.run( + [ + sys.executable, + "-c", + script, + str(tmp_path), + str(use_async), + str(with_pylance), + ], + capture_output=True, + text=True, + timeout=30, + ) + assert result.returncode == 0, result.stderr + + def test_blob_resolves_pylance_type_without_eager_import(): script = textwrap.dedent( """\ @@ -1295,12 +1367,13 @@ def test_merge_insert_bytes_after_reopen_without_touching_blob_type(tmp_path): script = textwrap.dedent( f"""\ import lancedb + import pyarrow as pa db = lancedb.connect({str(tmp_path)!r}) table = db.open_table("images") image_type = table.schema.field("image").type - if type(image_type).__name__ != "StructType": - raise SystemExit(f"expected StructType, got {{type(image_type)}}") + if not isinstance(image_type, pa.ExtensionType): + raise SystemExit(f"expected ExtensionType, got {{type(image_type)}}") result = ( table.merge_insert("id") .when_matched_update_all() @@ -1339,12 +1412,13 @@ def test_merge_insert_bytes_after_reopen_without_pylance(tmp_path): script = _HIDE_LANCE_BLOB + textwrap.dedent( f"""\ import lancedb + import pyarrow as pa db = lancedb.connect({str(tmp_path)!r}) table = db.open_table("images") image_type = table.schema.field("image").type - if type(image_type).__name__ != "StructType": - raise SystemExit(f"expected StructType, got {{type(image_type)}}") + if not isinstance(image_type, pa.ExtensionType): + raise SystemExit(f"expected ExtensionType, got {{type(image_type)}}") result = ( table.merge_insert("id") .when_matched_update_all() @@ -1374,7 +1448,7 @@ def test_merge_insert_bytes_after_reopen_without_pylance(tmp_path): assert result.returncode == 0, result.stderr -def test_merge_insert_blob_array_into_reopened_unregistered_table(tmp_path): +def test_merge_insert_blob_array_into_reopened_table(tmp_path): db = lancedb.connect(tmp_path) schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) table = db.create_table("images", schema=schema) @@ -1388,10 +1462,8 @@ def test_merge_insert_blob_array_into_reopened_unregistered_table(tmp_path): db = lancedb.connect({str(tmp_path)!r}) table = db.open_table("images") image_type = table.schema.field("image").type - if type(image_type).__name__ != "StructType": - raise SystemExit( - f"expected StructType before lance import, got {{type(image_type)}}" - ) + if not isinstance(image_type, pa.ExtensionType): + raise SystemExit(f"expected ExtensionType, got {{type(image_type)}}") import lance diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index fe271b1fd..32fec4b57 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -10,7 +10,9 @@ import json import multiprocessing as mp import pickle import re +import subprocess import sys +import textwrap import threading import time from unittest.mock import MagicMock, patch @@ -3534,6 +3536,128 @@ BLOB_DESCRIBE_RESPONSE = { } +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("with_pylance", [False, True]) +@pytest.mark.parametrize("nested", [False, True]) +def test_remote_schema_restores_blob_extensions(use_async, with_pylance, nested): + description = json.loads(json.dumps(BLOB_DESCRIBE_RESPONSE)) + without_ranges = description["schema"]["fields"].pop() + without_ranges["name"] = "without_ranges" + image = json.loads(json.dumps(without_ranges)) + image["name"] = "image" + image["nullable"] = False + image["metadata"]["lance-encoding:blob-inline-size-threshold"] = "8192" + image["type"]["fields"].extend( + {"name": name, "type": {"type": "uint64"}, "nullable": True} + for name in ("position", "size") + ) + plain = { + "name": "plain", + "type": image["type"], + "nullable": True, + "metadata": {"app-key": "app-value"}, + } + fields = [image, without_ranges, plain] + if nested: + fields = [ + { + "name": "nested", + "type": {"type": "struct", "fields": fields}, + "nullable": False, + }, + { + "name": "images", + "type": {"type": "list", "fields": [image]}, + "nullable": True, + }, + ] + description["schema"]["fields"].extend(fields) + description["schema"]["metadata"] = {"schema-key": "schema-value"} + + def handler(request): + assert request.path == "/v1/table/test/describe/" + send_json(request, description) + + # A fresh interpreter must read the schema before blob() or a pylance import + # registers the extension type and hides the original failure. + script = textwrap.dedent( + """\ + import asyncio + import importlib.abc + import sys + import pyarrow as pa + + if sys.argv[3] == "False": + class NoLance(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path, target=None): + if fullname == "lance" or fullname.startswith("lance."): + raise ModuleNotFoundError(fullname, name="lance") + sys.meta_path.insert(0, NoLance()) + import lancedb + + assert "lance.blob" not in sys.modules + options = dict(api_key="fake", host_override=sys.argv[1]) + + async def read_schema(): + db = await lancedb.connect_async("db://dev", **options) + table = await db.open_table("test") + return await table.schema() + + if sys.argv[2] == "True": + schema = asyncio.run(read_schema()) + else: + db = lancedb.connect("db://dev", **options) + schema = db.open_table("test").schema + + assert schema.metadata == {b"schema-key": b"schema-value"} + assert schema.field("id") == pa.field("id", pa.int64(), nullable=False) + if sys.argv[4] == "True": + assert not schema.field("nested").nullable + fields = schema.field("nested").type + assert schema.field("images").type.value_field.equals( + fields.field("image"), check_metadata=True + ) + else: + fields = schema + + image = fields.field("image") + assert isinstance(image.type, pa.ExtensionType), image + assert image.type == lancedb.blob("x").type + assert type(image.type) is lancedb.BlobType + assert not image.nullable + assert image.metadata == { + b"lance-encoding:blob-inline-size-threshold": b"8192" + } + plain = fields.field("plain") + assert plain.type == image.type.storage_type + assert plain.metadata == {b"app-key": b"app-value"} + without_ranges = fields.field("without_ranges") + assert isinstance(without_ranges.type, pa.ExtensionType), without_ranges + assert type(without_ranges.type) is lancedb.BlobType + assert without_ranges.type.storage_type == pa.struct([ + pa.field("data", pa.large_binary()), pa.field("uri", pa.string()) + ]) + assert without_ranges.nullable + """ + ) + with mock_lancedb_connection(handler) as db: + result = subprocess.run( + [ + sys.executable, + "-c", + script, + db.host_override, + str(use_async), + str(with_pylance), + str(nested), + ], + capture_output=True, + text=True, + timeout=30, + ) + assert result.returncode == 0, result.stderr + + @pytest.mark.asyncio @pytest.mark.parametrize("use_async", [False, True]) @pytest.mark.parametrize("mode", ["merge", "overwrite"])