fix(python): preserve JSON encoding in merge insert (#3976)

<!-- lance-gatekeeper-fix:v1 agent=40e5cf476a59265c71653574eda834d2
generation=1 -->

## Summary

- preserve incoming PyArrow `arrow.json` fields while schema
sanitization aligns input to a stored `lance.json` schema
- let Lance perform the required JSONB encoding instead of relabeling
raw JSON bytes as encoded storage
- cover both merge insert and the conditional add sanitization path with
end-to-end regression tests

## Root cause

Python schema sanitization aligns incoming data to the table schema
before passing it to Lance. Merge insert always takes this path, while
add takes it conditionally for preprocessing such as non-default
bad-vector handling or embedding functions. For JSON columns, the cast
changed logical `arrow.json` strings into the table's JSONB-backed
`lance.json` storage type without encoding the bytes, so Lance treated
raw JSON text as JSONB.

## Validation

- `cd python && uv run --extra tests pytest python/tests/test_table.py
-k 'merge_insert or add_sanitization_encodes_json' -q`
- targeted schema-cast and JSON encoding tests
- `ruff check .`
- `ruff format --check python/python/lancedb/table.py
python/python/tests/test_table.py`

Fixes #3923

---------

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>

python/python/lancedb/table.py

@@ -433,6 +433,20 @@ def _cast_to_target_schema(
     return pa.RecordBatchReader.from_batches(reordered_schema, gen())
 
 
+def _field_extension_name(field: pa.Field) -> Optional[str]:
+    extension_name = getattr(field.type, "extension_name", None)
+    if extension_name is not None:
+        return extension_name
+
+    metadata = field.metadata or {}
+    extension_name = metadata.get(b"ARROW:extension:name") or metadata.get(
+        "ARROW:extension:name"
+    )
+    if isinstance(extension_name, bytes):
+        return extension_name.decode()
+    return extension_name
+
+
 def _align_field_types(
     fields: List[pa.Field],
     target_fields: List[pa.Field],
@@ -445,6 +459,16 @@ def _align_field_types(
         target_field = next((f for f in target_fields if f.name == field.name), None)
         if target_field is None:
             raise ValueError(f"Field '{field.name}' not found in target schema")
+        # Preserve arrow.json input until it reaches Lance. 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.
+        if (
+            _field_extension_name(field) == "arrow.json"
+            and _field_extension_name(target_field) == "lance.json"
+        ):
+            new_fields.append(field)
+            continue
         if pa.types.is_struct(target_field.type):
             if pa.types.is_struct(field.type):
                 new_type = pa.struct(

python/python/tests/test_table.py

@@ -2772,6 +2772,56 @@ async def test_merge_insert_async(mem_db_async: AsyncConnection):
     assert (await table.to_arrow()).sort_by("a") == expected
 
 
+@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
+@pytest.mark.asyncio
+async def test_merge_insert_encodes_json(mem_db_async: AsyncConnection):
+    json_type = pa.json_()
+    schema = pa.schema([pa.field("id", pa.string()), pa.field("j", json_type)])
+
+    def json_table(rows):
+        json_values = pa.ExtensionArray.from_storage(
+            json_type,
+            pa.array([value for _, value in rows], type=json_type.storage_type),
+        )
+        return pa.Table.from_arrays(
+            [pa.array([row_id for row_id, _ in rows]), json_values], schema=schema
+        )
+
+    table = await mem_db_async.create_table("json_merge", schema=schema)
+    await table.add(json_table([("a", '{"k": 1}'), ("b", '{"k": 9}')]))
+
+    await (
+        table.merge_insert("id")
+        .when_matched_update_all()
+        .execute(json_table([("a", '{"k": 2}')]))
+    )
+
+    rows = sorted(await table.query().to_list(), key=lambda row: row["id"])
+    assert rows == [
+        {"id": "a", "j": '{"k":2}'},
+        {"id": "b", "j": '{"k":9}'},
+    ]
+    filtered = await table.query().where("json_extract(j, '$.k') = '2'").to_list()
+    assert filtered == [{"id": "a", "j": '{"k":2}'}]
+
+
+@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
+@pytest.mark.asyncio
+async def test_add_sanitization_encodes_json(mem_db_async: AsyncConnection):
+    json_type = pa.json_()
+    schema = pa.schema([pa.field("id", pa.string()), pa.field("j", json_type)])
+    json_values = pa.ExtensionArray.from_storage(
+        json_type, pa.array(['{"k": 3}'], type=json_type.storage_type)
+    )
+    data = pa.Table.from_arrays([pa.array(["c"]), json_values], schema=schema)
+
+    table = await mem_db_async.create_table("json_add", schema=schema)
+    await table.add(data, on_bad_vectors="fill")
+
+    rows = await table.query().where("json_extract(j, '$.k') = '3'").to_list()
+    assert rows == [{"id": "c", "j": '{"k":3}'}]
+
+
 def test_create_with_embedding_function(mem_db: DBConnection):
     class MyTable(LanceModel):
         text: str