fix!: normalize blob layouts and support external byte ranges (#4589)

python/python/lancedb/table.py

@@ -588,7 +588,7 @@ def _cast_to_target_schema(
 def _preserve_blob_v2_fields(
     inferred_schema: pa.Schema, existing_schema: pa.Schema
 ) -> pa.Schema:
-    """Keep existing blob fields when an overwrite infers a replacement schema."""
+    """Keep inferred blob semantics without overriding explicit replacement types."""
     existing_fields = {field.name: field for field in existing_schema}
     fields = [
         _preserve_blob_v2_field(field, existing_fields[field.name])
@@ -600,6 +600,9 @@ def _preserve_blob_v2_fields(
 
 
 def _preserve_blob_v2_field(inferred: pa.Field, existing: pa.Field) -> pa.Field:
+    if is_blob_v2_field(inferred):
+        # An explicitly declared blob layout is part of the replacement schema.
+        return inferred
     if is_blob_v2_field(existing):
         return existing
 
@@ -637,12 +640,10 @@ def _preserve_blob_v2_field(inferred: pa.Field, existing: pa.Field) -> pa.Field:
 def _coerce_blob_write_columns(
     batch: pa.RecordBatch, target_schema: pa.Schema
 ) -> pa.RecordBatch:
-    """Materialize blob storage structs before the stream leaves Python.
+    """Materialize Python bytes and URIs as blob structs before Arrow casting.
 
-    merge_insert requires its source reader to already match the table's
-    physical schema. Unlike add and insert, it does not pass through
-    LanceDB's Rust blob coercion, so preserving binary input here would
-    reach Lance as binary and fail the schema check.
+    Existing blob structs retain their input layout. Rust normalizes them to
+    the stored schema and rejects populated ranges on two-field targets.
     """
     columns = []
     fields = []
@@ -973,6 +974,17 @@ def _align_field(field: pa.Field, target_field: pa.Field) -> pa.Field:
                 "(Utf8, LargeUtf8, Utf8View), or a Struct with a "
                 "'data' or 'uri' child"
             )
+        if pa.types.is_struct(source_type) or (
+            isinstance(source_type, pa.ExtensionType)
+            and (
+                not isinstance(target_field.type, pa.ExtensionType)
+                or type(source_type) is type(target_field.type)
+            )
+        ):
+            # Keep the source layout until Rust can fill missing range fields or
+            # validate every batch before dropping all-null range fields. Arrow
+            # casting here would reject differing layouts before that validation.
+            return field.with_metadata(target_field.metadata)
     # 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.
@@ -1832,8 +1844,8 @@ class Table(ABC):
             - pyarrow.Table or pyarrow.RecordBatch
         mode: str
             The mode to use when writing the data. Valid values are
-            "append" and "overwrite". Overwrite infers a replacement schema
-            but preserves existing blob v2 fields present in the new data.
+            "append" and "overwrite". Overwrite adopts explicitly declared blob
+            layouts, while inferred bytes retain the existing blob v2 fields.
         on_bad_vectors: str, default "error"
             What to do if any of the vectors are not the same size or contains NaNs.
             One of "error", "drop", "fill".
@@ -4111,8 +4123,8 @@ class LanceTable(Table):
             The data to insert into the table.
         mode: str
             The mode to use when writing the data. Valid values are
-            "append" and "overwrite". Overwrite infers a replacement schema
-            but preserves existing blob v2 fields present in the new data.
+            "append" and "overwrite". Overwrite adopts explicitly declared blob
+            layouts, while inferred bytes retain the existing blob v2 fields.
         on_bad_vectors: str, default "error"
             What to do if any of the vectors are not the same size or contains NaNs.
             One of "error", "drop", "fill", "null".
@@ -6145,8 +6157,8 @@ class AsyncTable:
             - pyarrow.Table or pyarrow.RecordBatch
         mode: str
             The mode to use when writing the data. Valid values are
-            "append" and "overwrite". Overwrite infers a replacement schema
-            but preserves existing blob v2 fields present in the new data.
+            "append" and "overwrite". Overwrite adopts explicitly declared blob
+            layouts, while inferred bytes retain the existing blob v2 fields.
         on_bad_vectors: str, default "error"
             What to do if any of the vectors are not the same size or contains NaNs.
             One of "error", "drop", "fill", "null".

python/python/tests/test_blob.py

@@ -63,6 +63,25 @@ def _blob_array(name, values):
     return pa.ExtensionArray.from_storage(blob_type, storage)
 
 
+def _blob_layout_data(rows, with_ranges, *, as_extension=True):
+    field = lancedb.blob("image")
+    storage_type = field.type.storage_type
+    if not with_ranges or not as_extension:
+        storage_type = pa.struct(list(storage_type)[: (4 if with_ranges else 2)])
+        field = pa.field(
+            "image",
+            storage_type,
+            metadata={"ARROW:extension:name": "lance.blob.v2"},
+        )
+    images = pa.array([row["image"] for row in rows], type=storage_type)
+    if isinstance(field.type, pa.ExtensionType):
+        images = pa.ExtensionArray.from_storage(field.type, images)
+    return pa.Table.from_arrays(
+        [pa.array([row["id"] for row in rows], type=pa.int64()), images],
+        schema=pa.schema([pa.field("id", pa.int64()), field]),
+    )
+
+
 def _row_ids_by_id(table):
     hits = table.search().with_row_id(True).limit(1000).to_arrow()
     assert "_rowid" in hits.column_names
@@ -998,6 +1017,133 @@ def test_merge_insert_writes_python_bytes():
     assert blobs.to_pylist() == [b"updated", b"inserted"]
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("use_async", [False, True])
+@pytest.mark.parametrize("input_type", ["extension", "struct", "reader"])
+@pytest.mark.parametrize("target_with_ranges", [False, True])
+async def test_merge_insert_normalizes_blob_layouts(
+    use_async, input_type, target_with_ranges
+):
+    data = _blob_layout_data(
+        [{"id": 1, "image": {"data": b"before"}}], target_with_ranges
+    )
+    table = lancedb.connect("memory:///").create_table("merge_layouts", data=data)
+    original_schema = table.schema
+    updates = _blob_layout_data(
+        [
+            {"id": 1, "image": {"data": b"updated"}},
+            {"id": 2, "image": {"data": b"inserted"}},
+            {"id": 3, "image": None},
+        ],
+        not target_with_ranges,
+        as_extension=input_type != "struct",
+    )
+    if input_type == "reader":
+        updates = updates.to_reader(max_chunksize=1)
+    writer = table._table if use_async else table
+    result = (
+        writer.merge_insert("id")
+        .when_matched_update_all()
+        .when_not_matched_insert_all()
+        .execute(updates)
+    )
+    if use_async:
+        result = await result
+
+    assert result.num_updated_rows == 1
+    assert result.num_inserted_rows == 2
+    assert table.schema == original_schema
+    by_id = _row_ids_by_id(table)
+    assert table.fetch_blobs("image", [by_id[i] for i in (1, 2, 3)]).to_pylist() == [
+        b"updated",
+        b"inserted",
+        None,
+    ]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("use_async", [False, True])
+@pytest.mark.parametrize(
+    "byte_range", [{"position": 0}, {"size": 6}, {"position": 4, "size": 6}]
+)
+async def test_merge_insert_legacy_blob_rejects_ranges_atomically(
+    use_async, byte_range
+):
+    data = _blob_layout_data([{"id": 1, "image": {"data": b"before"}}], False)
+    table = lancedb.connect("memory:///").create_table("reject_ranges", data=data)
+    original_schema, original_version = table.schema, table.version
+    updates = _blob_layout_data(
+        [
+            {"id": 1, "image": {"data": b"updated"}},
+            {"id": 2, "image": {"uri": "file:///payload", **byte_range}},
+        ],
+        True,
+    ).to_reader(max_chunksize=1)
+    writer = table._table if use_async else table
+    with pytest.raises(RuntimeError, match="cannot store byte ranges"):
+        result = (
+            writer.merge_insert("id")
+            .when_matched_update_all()
+            .when_not_matched_insert_all()
+            .execute(updates)
+        )
+        if use_async:
+            await result
+
+    assert table.version == original_version
+    assert table.schema == original_schema
+    assert table.count_rows() == 1
+    assert table.fetch_blobs("image", [_row_ids_by_id(table)[1]]).to_pylist() == [
+        b"before"
+    ]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("use_async", [False, True])
+@pytest.mark.parametrize("input_type", ["extension", "struct", "reader"])
+@pytest.mark.parametrize("source_with_ranges", [False, True])
+async def test_add_overwrite_adopts_explicit_blob_layout(
+    tmp_path, use_async, input_type, source_with_ranges
+):
+    data = _blob_layout_data(
+        [{"id": 1, "image": {"data": b"before"}}], not source_with_ranges
+    )
+    table = lancedb.connect("memory:///").create_table("overwrite_layouts", data=data)
+    payload = b"0123456789abcdef"
+    path = tmp_path / "payload.bin"
+    path.write_bytes(payload)
+    value = (
+        {"uri": path.as_uri(), "position": 4, "size": 6}
+        if source_with_ranges
+        else {"data": b"updated"}
+    )
+    updates = _blob_layout_data(
+        [{"id": 2, "image": value}, {"id": 3, "image": None}],
+        source_with_ranges,
+        as_extension=input_type != "struct",
+    )
+    if input_type == "reader":
+        updates = updates.to_reader(max_chunksize=1)
+    writer = table._table if use_async else table
+    result = writer.add(
+        updates, mode="overwrite", allow_external_blob_outside_bases=True
+    )
+    if use_async:
+        await result
+
+    blob_type = table.schema.field("image").type
+    if isinstance(blob_type, pa.ExtensionType):
+        blob_type = blob_type.storage_type
+    assert blob_type.num_fields == (4 if source_with_ranges else 2)
+    assert table.blob_columns() == ["image"]
+    assert table.count_rows() == 2
+    by_id = _row_ids_by_id(table)
+    assert table.fetch_blobs("image", [by_id[2], by_id[3]]).to_pylist() == [
+        payload[4:10] if source_with_ranges else b"updated",
+        None,
+    ]
+
+
 @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"

python/python/tests/test_remote_db.py

@@ -3254,6 +3254,89 @@ BLOB_DESCRIBE_RESPONSE = {
 }
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("use_async", [False, True])
+@pytest.mark.parametrize("mode", ["merge", "overwrite"])
+@pytest.mark.parametrize("target_with_ranges", [False, True])
+async def test_remote_writes_normalize_blob_layouts(
+    use_async, mode, target_with_ranges
+):
+    description = json.loads(json.dumps(BLOB_DESCRIBE_RESPONSE))
+    if target_with_ranges:
+        description["schema"]["fields"][1]["type"]["fields"].extend(
+            {"name": name, "type": {"type": "uint64"}, "nullable": True}
+            for name in ("position", "size")
+        )
+    received = []
+
+    def handler(request):
+        if request.path == "/v1/table/test/describe/":
+            response = description
+        else:
+            expected = "merge_insert" if mode == "merge" else "insert"
+            assert urlparse(request.path).path == f"/v1/table/test/{expected}/"
+            if request.headers.get("Transfer-Encoding") == "chunked":
+                body = bytearray()
+                while size := int(request.rfile.readline().split(b";")[0], 16):
+                    body.extend(request.rfile.read(size))
+                    assert request.rfile.read(2) == b"\r\n"
+                assert request.rfile.read(2) == b"\r\n"
+            else:
+                body = request.rfile.read(int(request.headers["Content-Length"]))
+            received.append(pa.ipc.open_stream(body).read_all())
+            response = {
+                "version": 2,
+                "num_updated_rows": 0,
+                "num_inserted_rows": 2,
+                "num_deleted_rows": 0,
+            }
+        request.send_response(200)
+        request.send_header("Content-Type", "application/json")
+        request.send_header("phalanx-version", "0.5.0")
+        request.end_headers()
+        request.wfile.write(json.dumps(response).encode())
+
+    field = lancedb.blob("image")
+    storage_type = field.type.storage_type
+    if target_with_ranges:
+        storage_type = pa.struct(list(storage_type)[:2])
+        field = pa.field(
+            "image",
+            storage_type,
+            metadata={"ARROW:extension:name": "lance.blob.v2"},
+        )
+    images = pa.array([{"data": b"payload"}, None], type=storage_type)
+    if isinstance(field.type, pa.ExtensionType):
+        images = pa.ExtensionArray.from_storage(field.type, images)
+    data = pa.Table.from_arrays(
+        [pa.array([1, 2], type=pa.int64()), images],
+        schema=pa.schema([pa.field("id", pa.int64()), field]),
+    )
+
+    with mock_lancedb_connection(handler) as db:
+        remote_table = db.open_table("test")
+        table = remote_table._table if use_async else remote_table
+        if mode == "merge":
+            result = (
+                table.merge_insert("id").when_not_matched_insert_all().execute(data)
+            )
+        else:
+            result = table.add(data, mode="overwrite")
+        if use_async:
+            await result
+
+    assert len(received) == 1
+    images = received[0]["image"].combine_chunks()
+    if isinstance(images.type, pa.ExtensionType):
+        images = images.storage
+    with_ranges = target_with_ranges if mode == "merge" else not target_with_ranges
+    assert images.type.num_fields == (4 if with_ranges else 2)
+    expected = {"data": b"payload", "uri": None}
+    if with_ranges:
+        expected.update(position=None, size=None)
+    assert images.to_pylist() == [expected, None]
+
+
 def blob_query_response_table():
     image_field = pa.field(
         "image",