Merge origin/main into gatekeeper/fix-3759-1

python/python/lancedb/table.py

@@ -559,12 +559,18 @@ 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):
+    if _is_string_like(values.type):
-        data = pa.nulls(len(values), type=pa.large_binary())
+        carrier_name = "uri"
+        carrier = values
+    elif pa.types.is_null(values.type):
+        carrier_name = None
+        carrier = None
     elif pa.types.is_large_binary(values.type):
-        data = values
+        carrier_name = "data"
+        carrier = values
     else:
-        data = values.cast(pa.large_binary())
+        carrier_name = "data"
+        carrier = values.cast(pa.large_binary())
     length = len(values)
     storage_type = target_field.type
     if isinstance(storage_type, pa.ExtensionType):
@@ -572,8 +578,8 @@ 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":
+        if storage_field.name == carrier_name:
-            children.append(data)
+            children.append(carrier.cast(storage_field.type))
         else:
             children.append(pa.nulls(length, type=storage_field.type))
     storage = pa.StructArray.from_arrays(
@@ -593,7 +599,11 @@ 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_binary_like(data_type: pa.DataType) -> bool:
@@ -604,6 +614,15 @@ def _is_binary_like(data_type: pa.DataType) -> bool:
     )
 
 
+def _is_string_like(data_type: pa.DataType) -> bool:
+    predicates = ("is_string", "is_large_string", "is_string_view")
+    return any(
+        predicate(data_type)
+        for name in predicates
+        if (predicate := getattr(pa.types, name, None)) is not None
+    )
+
+
 def _field_extension_name(field: pa.Field) -> Optional[str]:
     extension_name = getattr(field.type, "extension_name", None)
     if extension_name is not None:
@@ -618,6 +637,187 @@ 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 _with_list_value_field(
+    data_type: pa.DataType, value_field: pa.Field
+) -> pa.DataType:
+    if pa.types.is_list(data_type):
+        return pa.list_(value_field)
+    if pa.types.is_large_list(data_type):
+        return pa.large_list(value_field)
+    return pa.list_(value_field, data_type.list_size)
+
+
+def _extension_storage_field(field: pa.Field) -> pa.Field:
+    """Return a from-pylist-compatible field for nested write extensions."""
+    extension_name = _field_extension_name(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)
+    if extension_name == _BLOB_EXTENSION_NAME:
+        metadata = dict(field.metadata or {})
+        metadata[b"ARROW:extension:name"] = _BLOB_EXTENSION_NAME.encode()
+        metadata[b"ARROW:extension:metadata"] = b""
+        storage_type = getattr(field.type, "storage_type", field.type)
+        return _with_field_type(field, storage_type, metadata=metadata)
+    if pa.types.is_struct(field.type):
+        children = [_extension_storage_field(child) for child in field.type]
+        return _with_field_type(field, pa.struct(children))
+    if _is_list_like(field.type):
+        value_field = _extension_storage_field(field.type.value_field)
+        return _with_field_type(field, _with_list_value_field(field.type, value_field))
+    return field
+
+
+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):
+        target_value_field = target_field.type.value_field
+        if _field_contains_write_extension(target_value_field):
+            prepared = _extension_storage_field(target_value_field)
+            data_type = _with_list_value_field(target_field.type, prepared)
+            return _with_field_type(field, data_type), True
+
+    return field, False
+
+
+def _prepare_extension_value(
+    value: Any, target_field: pa.Field, *, within_list: bool = False
+) -> Any:
+    """Shape raw nested blob values for PyArrow's struct construction."""
+    if value is None:
+        return None
+
+    extension_name = _field_extension_name(target_field)
+    if extension_name == _BLOB_EXTENSION_NAME and within_list:
+        if isinstance(value, (bytes, bytearray, memoryview)):
+            return {"data": value}
+        if isinstance(value, str):
+            return {"uri": value}
+        return value
+
+    if pa.types.is_struct(target_field.type) and isinstance(value, dict):
+        target_children = {child.name: child for child in target_field.type}
+        return {
+            name: _prepare_extension_value(
+                child_value, target_children[name], within_list=within_list
+            )
+            if name in target_children
+            else child_value
+            for name, child_value in value.items()
+        }
+
+    if _is_list_like(target_field.type) and isinstance(value, (list, tuple)):
+        return [
+            _prepare_extension_value(
+                item, target_field.type.value_field, within_list=True
+            )
+            for item in value
+        ]
+
+    return value
+
+
+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}
+    if not any(
+        _field_contains_write_extension(field) for field in target_fields.values()
+    ):
+        return data
+
+    inferred = pa.Table.from_pylist(data)
+    fields = []
+    changed = False
+    for field in inferred.schema:
+        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
+
+    insert_schema = pa.schema(fields, metadata=inferred.schema.metadata)
+    prepared_data = [
+        {
+            name: _prepare_extension_value(value, target_fields[name])
+            if name in target_fields
+            else value
+            for name, value in row.items()
+        }
+        for row in data
+    ]
+    return pa.Table.from_pylist(prepared_data, schema=insert_schema)
+
+
 def _align_field_types(
     fields: List[pa.Field],
     target_fields: List[pa.Field],
@@ -5624,6 +5824,9 @@ class AsyncTable:
         if fill_value is None:
             fill_value = 0.0
 
+        if mode != "overwrite":
+            data = _prepare_extension_list(data, schema)
+
         # _santitize_data is an old code path, but we will use it until the
         # new code path is ready.
         if mode == "overwrite":

python/python/tests/test_blob.py

@@ -710,6 +710,80 @@ def test_fetch_blobs_preserves_null_and_empty_values():
     assert blobs[3].as_py() == b"present"
 
 
+def test_add_all_null_list_to_blob_column():
+    table = _blob_table("all_null_add", [{"id": 1, "image": None}])
+
+    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_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)
+
+
+@pytest.mark.parametrize("large_list", [False, True], ids=["list", "large_list"])
+def test_add_list_of_dicts_to_blob_list_column(large_list):
+    db = lancedb.connect("memory:///")
+    blob_field = lancedb.blob("image")
+    blob_values = _blob_array("image", [b"seed"])
+    if large_list:
+        items_field = pa.field("items", pa.large_list(blob_field))
+        items = pa.LargeListArray.from_arrays(
+            pa.array([0, 1], type=pa.int64()), blob_values
+        )
+    else:
+        items_field = pa.field("items", pa.list_(blob_field))
+        items = pa.ListArray.from_arrays(pa.array([0, 1], type=pa.int32()), blob_values)
+    seed = pa.Table.from_arrays(
+        [pa.array([0], type=pa.int64()), items],
+        schema=pa.schema([pa.field("id", pa.int64()), items_field]),
+    )
+    table = db.create_table(f"blob_{large_list}_list_add", data=seed)
+
+    table.add([{"id": 1, "items": [None]}])
+    table.add(
+        [{"id": 2, "items": [b"a", None]}],
+        on_bad_vectors="fill",
+    )
+
+    ids = table.search().select(["id"]).to_arrow()["id"].to_pylist()
+    assert sorted(ids) == [0, 1, 2]
+    assert pa.types.is_large_list(table.schema.field("items").type) is large_list
+
+
 def test_fetch_blob_ranges_aligns_repeated_ranges_and_nulls():
     table = _blob_table(
         "range_alignment",
@@ -1230,6 +1304,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

@@ -786,6 +786,55 @@ async def test_add_async(mem_db_async: AsyncConnection):
     assert await table.count_rows() == 3
 
 
+@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("values", "expected"),
+    [
+        ([None], [None]),
+        ([None, '{"k": 1}'], [None, '{"k":1}']),
+        (['{"k": 2}'], ['{"k":2}']),
+    ],
+)
+async def test_add_list_of_dicts_to_json_column(
+    mem_db_async: AsyncConnection, values, expected
+):
+    schema = pa.schema([pa.field("id", pa.int64()), pa.field("value", pa.json_())])
+    table = await mem_db_async.create_table("json_list_add", schema=schema)
+
+    await table.add([{"id": idx, "value": value} for idx, value in enumerate(values)])
+
+    rows = (await table.to_arrow()).sort_by("id").to_pylist()
+    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.