fix(python): coerce blob URI inputs before Arrow cast

This commit is contained in:
Gatefixer
2026-09-23 18:38:55 +00:00
parent df7bfcdb65
commit 87c573fc92
3 changed files with 114 additions and 9 deletions
+37 -9
View File
@@ -563,12 +563,7 @@ def _coerce_blob_list_values(
def _coerce_value_to_blob(values: pa.Array, target_field: pa.Field) -> pa.Array:
if pa.types.is_null(values.type):
data = pa.nulls(len(values), type=pa.large_binary())
elif pa.types.is_large_binary(values.type):
data = values
else:
data = values.cast(pa.large_binary())
value_field = "uri" if _is_string_like(values.type) else "data"
length = len(values)
storage_type = target_field.type
if isinstance(storage_type, pa.ExtensionType):
@@ -576,8 +571,11 @@ def _coerce_value_to_blob(values: pa.Array, target_field: pa.Field) -> pa.Array:
storage_fields = list(storage_type)
children = []
for storage_field in storage_fields:
if storage_field.name == "data":
children.append(data)
if storage_field.name == value_field and not pa.types.is_null(values.type):
if values.type == storage_field.type:
children.append(values)
else:
children.append(values.cast(storage_field.type))
else:
children.append(pa.nulls(length, type=storage_field.type))
storage = pa.StructArray.from_arrays(
@@ -597,7 +595,19 @@ def _physical_array_and_type(array: pa.Array) -> tuple[pa.Array, pa.DataType]:
def _can_coerce_to_blob(data_type: pa.DataType) -> bool:
return _is_binary_like(data_type) or pa.types.is_null(data_type)
return (
_is_binary_like(data_type)
or _is_string_like(data_type)
or pa.types.is_null(data_type)
)
def _is_string_like(data_type: pa.DataType) -> bool:
return (
pa.types.is_string(data_type)
or pa.types.is_large_string(data_type)
or pa.types.is_string_view(data_type)
)
def _is_binary_like(data_type: pa.DataType) -> bool:
@@ -759,6 +769,24 @@ def _arrow_json_storage_type(input_type: pa.DataType) -> Optional[pa.DataType]:
def _align_field(field: pa.Field, target_field: pa.Field) -> pa.Field:
if is_blob_v2_field(target_field):
source_type = field.type
if pa.types.is_struct(source_type):
if not any(child.name in ("data", "uri") for child in source_type):
raise ValueError(
f"blob struct input for column '{field.name}' must contain "
"a 'data' or 'uri' child"
)
elif not _can_coerce_to_blob(source_type) and not isinstance(
source_type, pa.ExtensionType
):
raise ValueError(
f"cannot coerce column '{field.name}' with type {source_type} "
"into a blob v2 struct. expected binary bytes "
"(Binary, LargeBinary, BinaryView), strings "
"(Utf8, LargeUtf8, Utf8View), or a Struct with a "
"'data' or 'uri' child"
)
# LanceDB exposes stored JSON columns as lance.json (JSONB-backed LargeBinary), but
# casting the input to that storage type here merely relabels the raw JSON bytes as
# JSONB. Lance must see arrow.json so it can perform the JSONB encoding.
+44
View File
@@ -434,6 +434,50 @@ def test_merge_insert_writes_python_bytes():
assert blobs.to_pylist() == [b"updated", b"inserted"]
@pytest.mark.parametrize("write", ["merge_insert", "create_table"])
def test_uri_string_blob_write_reaches_external_base_validation(tmp_path, write):
blob_path = tmp_path / "payload.bin"
blob_path.write_bytes(b"payload")
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
rows = pa.table(
{"id": pa.array([1], pa.int64()), "image": pa.array([blob_path.as_uri()])}
)
db = lancedb.connect(tmp_path)
# These APIs do not expose add's outside-base opt-in. The URI should reach
# the blob policy check instead of failing at PyArrow's string-to-struct cast.
with pytest.raises(ValueError, match="outside registered external bases"):
if write == "merge_insert":
table = db.create_table("merge_uri", schema=schema)
table.merge_insert("id").when_not_matched_insert_all().execute(rows)
else:
db.create_table("create_uri", data=rows, schema=schema)
@pytest.mark.parametrize("write", ["merge_insert", "create_table"])
@pytest.mark.parametrize(
"source_type, value",
[
(pa.binary(4), b"abcd"),
(pa.list_(pa.uint8()), [1, 2]),
(pa.int64(), 42),
],
)
def test_blob_write_rejects_unsupported_input(write, source_type, value):
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
rows = pa.table(
{"id": pa.array([1], pa.int64()), "image": pa.array([value], source_type)}
)
db = lancedb.connect("memory:///")
with pytest.raises(ValueError, match="cannot coerce column 'image' with type"):
if write == "merge_insert":
table = db.create_table("merge_invalid", schema=schema)
table.merge_insert("id").when_not_matched_insert_all().execute(rows)
else:
db.create_table("create_invalid", data=rows, schema=schema)
def test_merge_insert_bytes_after_reopen_without_touching_blob_type(tmp_path):
db = lancedb.connect(tmp_path)
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
+33
View File
@@ -922,6 +922,39 @@ def test_cast_to_target_schema_coerces_binary_to_blob_v2():
]
@pytest.mark.parametrize(
"string_type", [pa.string(), pa.large_string(), pa.string_view()]
)
def test_cast_to_target_schema_coerces_string_to_blob_v2_uri(string_type):
data = pa.table({"image": pa.array(["file:///payload", None], type=string_type)})
target = pa.schema([lancedb.blob("image")])
output = _cast_to_target_schema(data.to_reader(), target).read_all()
image = output["image"].chunk(0)
assert type(image.type) is lancedb.BlobType
assert image.storage.to_pylist() == [
{"data": None, "uri": "file:///payload", "position": None, "size": None},
None,
]
@pytest.mark.parametrize(
"source_type, value",
[
(pa.binary(4), b"abcd"),
(pa.list_(pa.uint8()), [1, 2]),
(pa.int64(), 42),
],
)
def test_cast_to_target_schema_rejects_unsupported_blob_input(source_type, value):
data = pa.table({"image": pa.array([value], type=source_type)})
target = pa.schema([lancedb.blob("image")])
with pytest.raises(ValueError, match="cannot coerce column 'image' with type"):
_cast_to_target_schema(data.to_reader(), target)
def test_cast_to_target_schema_coerces_binary_to_metadata_blob_struct():
storage = lancedb.blob("image").type.storage_type
target = pa.schema(