fix(python): preserve nested extension carriers

python/python/lancedb/table.py

@@ -452,18 +452,96 @@ def _field_extension_name(field: pa.Field) -> Optional[str]:
     return extension_name
 
 
+_JSON_EXTENSION_NAMES = {"arrow.json", "lance.json"}
+_BLOB_EXTENSION_NAME = "lance.blob.v2"
+
+
+def _field_contains_write_extension(field: pa.Field) -> bool:
+    extension_name = _field_extension_name(field)
+    if (
+        extension_name in _JSON_EXTENSION_NAMES
+        or extension_name == _BLOB_EXTENSION_NAME
+    ):
+        return True
+    if pa.types.is_struct(field.type):
+        return any(_field_contains_write_extension(child) for child in field.type)
+    if (
+        pa.types.is_list(field.type)
+        or pa.types.is_large_list(field.type)
+        or pa.types.is_fixed_size_list(field.type)
+    ):
+        return _field_contains_write_extension(field.type.value_field)
+    return False
+
+
+def _with_field_type(
+    field: pa.Field,
+    data_type: pa.DataType,
+    *,
+    name: Optional[str] = None,
+    metadata: Optional[dict] = None,
+) -> pa.Field:
+    return pa.field(
+        name or field.name,
+        data_type,
+        nullable=field.nullable,
+        metadata=field.metadata if metadata is None else metadata,
+    )
+
+
+def _prepare_extension_field(
+    field: pa.Field, target_field: pa.Field
+) -> Tuple[pa.Field, bool]:
+    extension_name = _field_extension_name(target_field)
+    if extension_name in _JSON_EXTENSION_NAMES:
+        metadata = dict(field.metadata or {})
+        metadata[b"ARROW:extension:name"] = b"arrow.json"
+        return _with_field_type(field, pa.string(), metadata=metadata), True
+    if extension_name == _BLOB_EXTENSION_NAME and pa.types.is_null(field.type):
+        return _with_field_type(field, pa.large_binary()), True
+
+    if pa.types.is_struct(field.type) and pa.types.is_struct(target_field.type):
+        target_children = {child.name: child for child in target_field.type}
+        children = []
+        changed = False
+        for child in field.type:
+            target_child = target_children.get(child.name)
+            if target_child is None:
+                children.append(child)
+                continue
+            prepared, child_changed = _prepare_extension_field(child, target_child)
+            children.append(prepared)
+            changed = changed or child_changed
+        if changed:
+            return _with_field_type(field, pa.struct(children)), True
+
+    if _is_list_like(field.type) and _is_list_like(target_field.type):
+        prepared, changed = _prepare_extension_field(
+            field.type.value_field, target_field.type.value_field
+        )
+        if changed:
+            prepared = _with_field_type(
+                prepared, prepared.type, name=target_field.type.value_field.name
+            )
+            if pa.types.is_list(field.type):
+                data_type = pa.list_(prepared)
+            elif pa.types.is_large_list(field.type):
+                data_type = pa.large_list(prepared)
+            else:
+                data_type = pa.list_(prepared, field.type.list_size)
+            return _with_field_type(field, data_type), True
+
+    return field, False
+
+
 def _prepare_extension_list(data: DATA, target_schema: pa.Schema) -> DATA:
     """Give inferred list columns the logical type required by extensions."""
     if not isinstance(data, list) or not data or not isinstance(data[0], dict):
         return data
 
     target_fields = {field.name: field for field in target_schema}
-    extension_names = {
-        name: _field_extension_name(field) for name, field in target_fields.items()
-    }
     if not any(
-        name in {"arrow.json", "lance.json", "lance.blob.v2"}
-        for name in extension_names.values()
+        _field_contains_write_extension(field) for field in target_fields.values()
     ):
         return data
 
@@ -471,23 +549,13 @@ def _prepare_extension_list(data: DATA, target_schema: pa.Schema) -> DATA:
     fields = []
     changed = False
     for field in inferred.schema:
-        extension_name = extension_names.get(field.name)
-        if extension_name in {"arrow.json", "lance.json"}:
-            json_factory = getattr(pa, "json_", None)
-            if json_factory is not None:
-                field = pa.field(field.name, json_factory(), nullable=field.nullable)
-            else:
-                field = pa.field(
-                    field.name,
-                    pa.string(),
-                    nullable=field.nullable,
-                    metadata={b"ARROW:extension:name": b"arrow.json"},
-                )
-            changed = True
-        elif extension_name == "lance.blob.v2" and pa.types.is_null(field.type):
-            field = pa.field(field.name, pa.large_binary(), nullable=field.nullable)
-            changed = True
-        fields.append(field)
+        target_field = target_fields.get(field.name)
+        if target_field is None:
+            fields.append(field)
+            continue
+        prepared, field_changed = _prepare_extension_field(field, target_field)
+        fields.append(prepared)
+        changed = changed or field_changed
 
     if not changed:
         return inferred
@@ -496,6 +564,29 @@ def _prepare_extension_list(data: DATA, target_schema: pa.Schema) -> DATA:
     return pa.Table.from_pylist(data, schema=insert_schema)
 
 
+def _is_blob_source_field(field: pa.Field) -> bool:
+    if _field_extension_name(field) == _BLOB_EXTENSION_NAME:
+        return True
+
+    predicates = (
+        "is_binary",
+        "is_large_binary",
+        "is_binary_view",
+        "is_string",
+        "is_large_string",
+        "is_string_view",
+    )
+    if any(
+        predicate(field.type)
+        for name in predicates
+        if (predicate := getattr(pa.types, name, None)) is not None
+    ):
+        return True
+    return pa.types.is_struct(field.type) and any(
+        child.name in {"data", "uri"} for child in field.type
+    )
+
+
 def _align_field_types(
     fields: List[pa.Field],
     target_fields: List[pa.Field],
@@ -508,13 +599,21 @@ 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")
+        target_extension_name = _field_extension_name(target_field)
+        # Preserve accepted blob carriers so Lance can construct the declared
+        # blob struct after optional Python preprocessing.
+        if target_extension_name == _BLOB_EXTENSION_NAME and _is_blob_source_field(
+            field
+        ):
+            new_fields.append(field)
+            continue
         # 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"
+            and target_extension_name in _JSON_EXTENSION_NAMES
         ):
             new_fields.append(field)
             continue

python/python/tests/test_blob.py

@@ -212,6 +212,41 @@ def test_add_all_null_list_to_blob_column():
     assert blobs[0].as_py() is None
 
 
+def test_add_all_null_list_to_blob_column_with_sanitizer():
+    db = lancedb.connect("memory:///")
+    schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
+    table = db.create_table("all_null_sanitized_add", schema=schema)
+
+    table.add([{"id": 1, "image": None}], on_bad_vectors="fill")
+
+    hits = table.search().to_arrow()
+    blobs = table.fetch_blobs("image", hits)
+    assert len(blobs) == 1
+    assert blobs[0].as_py() is None
+
+
+def test_add_all_null_list_to_nested_blob_column():
+    db = lancedb.connect("memory:///")
+    blob_field = lancedb.blob("image")
+    info_field = pa.field("info", pa.struct([blob_field]))
+    info = pa.StructArray.from_arrays(
+        [_blob_array("image", [b"seed"])], fields=[blob_field]
+    )
+    seed = pa.Table.from_arrays(
+        [pa.array([0], type=pa.int64()), info],
+        schema=pa.schema([pa.field("id", pa.int64()), info_field]),
+    )
+    table = db.create_table("nested_null_add", data=seed)
+
+    table.add([{"id": 1, "info": {"image": None}}])
+    table.add([{"id": 2, "info": {"image": None}}], on_bad_vectors="fill")
+
+    hits = table.search().where("id > 0").to_arrow()
+    blobs = table.fetch_blobs("info.image", hits)
+    assert len(blobs) == 2
+    assert all(blob.as_py() is None for blob in blobs)
+
+
 def test_fetch_blob_ranges_aligns_repeated_ranges_and_nulls():
     table = _blob_table(
         "range_alignment",
@@ -688,6 +723,7 @@ def test_add_external_uri_string_round_trips_with_flag(tmp_path):
     table = db.create_table("external_string", schema=schema)
     table.add(
         [{"id": 1, "image": blob_path.as_uri()}],
+        on_bad_vectors="fill",
         allow_external_blob_outside_bases=True,
     )
 

python/python/tests/test_table.py

@@ -792,6 +792,33 @@ async def test_add_list_of_dicts_to_json_column(
     assert [row["value"] for row in rows] == expected
 
 
+@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
+@pytest.mark.asyncio
+async def test_add_list_of_dicts_to_nested_json_column(
+    mem_db_async: AsyncConnection,
+):
+    json_field = pa.field("value", pa.json_())
+    info_field = pa.field("info", pa.struct([json_field]))
+    info = pa.StructArray.from_arrays(
+        [pa.array(['{"seed": 0}'], type=pa.json_())], fields=[json_field]
+    )
+    seed = pa.Table.from_arrays(
+        [pa.array([0], type=pa.int64()), info],
+        schema=pa.schema([pa.field("id", pa.int64()), info_field]),
+    )
+    table = await mem_db_async.create_table("nested_json_list_add", data=seed)
+
+    await table.add([{"id": 1, "info": {"value": '{"k": 1}'}}])
+    await table.add([{"id": 2, "info": {"value": '{"k": 2}'}}], on_bad_vectors="fill")
+
+    rows = (await table.to_arrow()).sort_by("id").to_pylist()
+    assert rows == [
+        {"id": 0, "info": {"value": '{"seed":0}'}},
+        {"id": 1, "info": {"value": '{"k":1}'}},
+        {"id": 2, "info": {"value": '{"k":2}'}},
+    ]
+
+
 def test_add_overwrite_infers_vector_schema(mem_db: DBConnection):
     """Overwrite should infer vector columns the same way create_table does.