mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-12 14:52:26 +00:00
fix(python): coerce blob URI inputs before Arrow cast
This commit is contained in:
@@ -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.
|
||||
|
||||
@@ -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")])
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user