mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-30 23:36:52 +00:00
fix(python): restore blob extensions in table schemas (#4728)
Reading a Blob v2 table in a fresh process could return its storage struct if neither `blob()` nor pylance had registered the extension. Both local and remote schema reads were affected, so `field.type` could depend on import order. Restore metadata-only blob fields on every table schema read through the existing BlobType resolver and Arrow's schema deserializer. Preserve nested fields, nullability, custom metadata, and both supported storage layouts. The fallback deserializer now honors the stored shape too. Fresh-process regressions cover local and remote tables, sync/async reads, and operation with and without pylance. The four native cases failed before removing the native-table guard. Existing merge-insert tests now require the restored extension type while continuing to verify updated and inserted blob contents. Validation: - All 12 fresh-process local/remote schema regression cases pass. - `uv run --extra tests pytest python/tests/test_table.py python/tests/test_blob.py python/tests/test_namespace.py python/tests/test_util.py -q -m 'not slow and not s3_test'`: 835 passed, 11 skipped. - Python package Ruff lint and formatting checks pass. - Repository-wide `ruff check .` reports 48 pre-existing diagnostics in unchanged `docs/web-tests` files. Fixes #4719 <!-- lance-gatekeeper-fix:v1 agent=67c656bb5889e698efce5f2fd9d58ea7 generation=1 --> --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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]:
|
||||
"""
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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"])
|
||||
|
||||
Reference in New Issue
Block a user