lancedb/python/python/tests/test_blob.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors

import io
import subprocess
import sys
import textwrap

import lance
import pyarrow as pa
import pyarrow.compute as pc
import pytest
from lance.blob import BlobType as LanceBlobType

import lancedb
from lancedb._blob import (
    blob_v2_projection_sources,
    read_row_ids_from_hits,
    stash_auto_row_ids,
)
from lancedb.expr import col
from lancedb.index import FTS
from lancedb.schema import blob_column_paths, blob_v2_column_paths


_HIDE_LANCE_BLOB = """\
import importlib.abc
import sys

class _MissingLanceBlob(importlib.abc.MetaPathFinder):
    def find_spec(self, fullname, path, target=None):
        if fullname == "lance.blob" or fullname.startswith("lance.blob."):
            raise ModuleNotFoundError(fullname, name="lance.blob")

sys.modules.pop("lance.blob", None)
sys.meta_path.insert(0, _MissingLanceBlob())
"""


def _blob_table(name, rows):
    db = lancedb.connect("memory:///")
    schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
    table = db.create_table(name, schema=schema)
    table.add(rows)
    return table


def _blob_array(name, values):
    blob_type = lancedb.blob(name).type
    storage_type = blob_type.storage_type
    storage = pa.StructArray.from_arrays(
        [
            pa.array(values, type=pa.large_binary()),
            pa.array([None] * len(values), type=pa.string()),
            pa.array([None] * len(values), type=pa.uint64()),
            pa.array([None] * len(values), type=pa.uint64()),
        ],
        fields=list(storage_type),
    )
    return pa.ExtensionArray.from_storage(blob_type, storage)


def _row_ids_by_id(table):
    hits = table.search().with_row_id(True).limit(1000).to_arrow()
    assert "_rowid" in hits.column_names
    return dict(zip(hits["id"].to_pylist(), hits["_rowid"].to_pylist()))


def test_blob_factory_declares_v2_field():
    field = lancedb.blob("image")
    assert isinstance(field.type, pa.ExtensionType)
    assert field.type.extension_name == "lance.blob.v2"
    assert lancedb.BlobType is LanceBlobType
    assert type(field.type) is LanceBlobType


def test_blob_type_works_without_pylance():
    script = _HIDE_LANCE_BLOB + textwrap.dedent(
        """\
        import lancedb
        import pyarrow as pa

        field = lancedb.blob("image")
        if not isinstance(field.type, pa.ExtensionType):
            raise SystemExit("expected an extension type")
        if field.type.extension_name != "lance.blob.v2":
            raise SystemExit(field.type.extension_name)
        if lancedb.BlobType is not type(field.type):
            raise SystemExit("BlobType is not the field type class")
        if lancedb.BlobType.__module__ != "lancedb.schema":
            raise SystemExit(lancedb.BlobType.__module__)

        db = lancedb.connect("memory:///")
        table = db.create_table(
            "images",
            schema=pa.schema([pa.field("id", pa.int64()), field]),
        )
        table.add([{"id": 1, "image": b"hello"}])
        result = (
            table.merge_insert("id")
            .when_matched_update_all()
            .when_not_matched_insert_all()
            .execute([{"id": 1, "image": b"updated"}, {"id": 2, "image": b"inserted"}])
        )
        if result.num_updated_rows != 1 or result.num_inserted_rows != 1:
            raise SystemExit(
                f"merge_insert rows updated={result.num_updated_rows} "
                f"inserted={result.num_inserted_rows}"
            )
        """
    )
    result = subprocess.run(
        [sys.executable, "-c", script],
        capture_output=True,
        text=True,
        check=False,
    )
    assert result.returncode == 0, result.stderr


def test_blob_resolves_pylance_type_without_eager_import():
    script = textwrap.dedent(
        """\
        import sys
        import lancedb

        if "lance.blob" in sys.modules:
            raise SystemExit("import lancedb imported lance.blob")
        field = lancedb.blob("image")
        from lance.blob import BlobType

        if type(field.type) is not BlobType:
            raise SystemExit(f"{type(field.type)} is not {BlobType}")
        import lance

        image = lance.blob_array([b"x"])
        if type(image.type) is not BlobType:
            raise SystemExit("blob_array used a different class")
        if type(image.type) is not type(field.type):
            raise SystemExit("field and array classes differ")
        """
    )
    result = subprocess.run(
        [sys.executable, "-c", script],
        capture_output=True,
        text=True,
        check=False,
    )
    assert result.returncode == 0, result.stderr


def test_blob_fallback_fails_if_name_already_registered():
    script = _HIDE_LANCE_BLOB + textwrap.dedent(
        """\
        import pyarrow as pa

        class OtherBlobType(pa.ExtensionType):
            def __init__(self):
                super().__init__(
                    pa.struct([pa.field("data", pa.large_binary())]),
                    "lance.blob.v2",
                )

            def __arrow_ext_serialize__(self):
                return b""

            @classmethod
            def __arrow_ext_deserialize__(cls, storage_type, serialized):
                return cls()

        pa.register_extension_type(OtherBlobType())
        import lancedb

        try:
            lancedb.blob("image")
        except ValueError as err:
            if "already registered" not in str(err):
                raise SystemExit(err)
        else:
            raise SystemExit("expected ValueError")
        """
    )
    result = subprocess.run(
        [sys.executable, "-c", script],
        capture_output=True,
        text=True,
        check=False,
    )
    assert result.returncode == 0, result.stderr


def test_blob_type_rejects_competing_registration_with_pylance():
    script = textwrap.dedent(
        """\
        import pyarrow as pa
        import pyarrow.ipc

        class OtherBlobType(pa.ExtensionType):
            def __init__(self):
                super().__init__(
                    pa.struct(
                        [
                            pa.field("data", pa.large_binary()),
                            pa.field("uri", pa.utf8()),
                            pa.field("position", pa.uint64()),
                            pa.field("size", pa.uint64()),
                        ]
                    ),
                    "lance.blob.v2",
                )

            def __arrow_ext_serialize__(self):
                return b""

            @classmethod
            def __arrow_ext_deserialize__(cls, storage_type, serialized):
                return cls()

        pa.register_extension_type(OtherBlobType())

        from lance.blob import BlobType

        if BlobType is OtherBlobType:
            raise SystemExit("pylance BlobType was replaced")
        schema = pa.schema([pa.field("value", BlobType())])
        restored = pa.ipc.read_schema(schema.serialize())
        if type(restored.field("value").type) is not OtherBlobType:
            raise SystemExit(type(restored.field("value").type))

        import lancedb

        try:
            lancedb.blob("image")
        except ValueError as err:
            if "__main__.OtherBlobType" not in str(err):
                raise SystemExit(err)
        else:
            raise SystemExit("expected ValueError")
        """
    )
    result = subprocess.run(
        [sys.executable, "-c", script],
        capture_output=True,
        text=True,
        check=False,
    )
    assert result.returncode == 0, result.stderr


def test_blob_v2_column_paths_include_list_children():
    schema = pa.schema(
        [
            pa.field("id", pa.int64()),
            pa.field("info", pa.struct([lancedb.blob("blob")])),
            pa.field("images", pa.list_(lancedb.blob("image"))),
            pa.field("large_images", pa.large_list(lancedb.blob("large_image"))),
            pa.field(
                "fixed_images",
                pa.list_(lancedb.blob("fixed_image"), list_size=2),
            ),
        ]
    )

    assert blob_v2_column_paths(schema) == [
        "info.blob",
        "images.image",
        "large_images.large_image",
        "fixed_images.fixed_image",
    ]


def test_blob_v2_projection_sources_use_typed_column_name():
    schema = pa.schema([lancedb.blob("blob")])

    assert blob_v2_projection_sources(schema, {"blob_alias": col("blob")}) == {
        "blob_alias": "blob"
    }


def _legacy_v1_table(name):
    db = lancedb.connect("memory:///")
    schema = pa.schema(
        [
            pa.field("id", pa.int64()),
            pa.field(
                "legacy", pa.large_binary(), metadata={"lance-encoding:blob": "true"}
            ),
        ]
    )
    table = db.create_table(name, schema=schema)
    table.add([{"id": 1, "legacy": b"old"}])
    return table


def test_blob_v2_column_paths_exclude_legacy_metadata():
    schema = pa.schema(
        [
            pa.field("id", pa.int64()),
            lancedb.blob("image"),
            pa.field(
                "legacy", pa.large_binary(), metadata={"lance-encoding:blob": "true"}
            ),
        ]
    )
    assert blob_v2_column_paths(schema) == ["image"]
    assert blob_column_paths(schema) == ["image", "legacy"]


def test_blob_v2_paths_match_blob_columns():
    table = _blob_table("paths_match", [{"id": 1, "image": b"x"}])
    assert blob_v2_column_paths(table.schema) == table.blob_columns()

    db = lancedb.connect("memory:///")
    info = pa.StructArray.from_arrays(
        [
            pa.array(["first"], type=pa.string()),
            _blob_array("blob", [b"nested"]),
        ],
        names=["name", "blob"],
    )
    data = pa.Table.from_arrays(
        [pa.array([1], type=pa.int64()), info],
        names=["id", "info"],
    )
    nested = db.create_table("nested_paths", data=data)
    assert blob_v2_column_paths(nested.schema) == nested.blob_columns()


def test_auto_row_id_stash_round_trip():
    table = _blob_table(
        "stash_round_trip",
        [{"id": 1, "image": b"alpha"}, {"id": 2, "image": b"beta"}],
    )
    hits = table.search().with_row_id(True).limit(10).to_arrow()
    row_ids = hits["_rowid"].to_pylist()

    stashed = stash_auto_row_ids(hits, ["image"])

    assert "_rowid" not in stashed.column_names
    assert stashed.schema.field("image").metadata == hits.schema.field("image").metadata
    assert read_row_ids_from_hits(stashed, "image") == row_ids


def test_blob_query_omits_auto_row_id():
    table = _blob_table("rowid", [{"id": 1, "image": b"x"}])
    hits = table.search().limit(10).to_arrow()
    assert "_rowid" not in hits.column_names


def test_blob_query_explicit_row_id_opt_in():
    table = _blob_table("explicit_rowid", [{"id": 1, "image": b"x"}])
    hits = table.search().with_row_id(True).limit(10).to_arrow()
    assert "_rowid" in hits.column_names


def test_table_to_pandas_descriptions_mode_omits_row_id():
    table = _blob_table("descriptions_no_leak", [{"id": 1, "image": b"x"}])
    df = table.to_pandas(blob_mode="descriptions")
    descriptor = df["image"].iloc[0]
    assert "_lance_row_id" not in descriptor
    assert set(descriptor.keys()) == {"kind", "position", "size", "blob_id", "blob_uri"}


@pytest.mark.asyncio
async def test_async_table_to_pandas_descriptions_mode_omits_row_id():
    db = await lancedb.connect_async("memory:///")
    schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
    table = await db.create_table("descriptions_no_leak_async", schema=schema)
    await table.add([{"id": 1, "image": b"x"}])
    df = await table.to_pandas(blob_mode="descriptions")
    descriptor = df["image"].iloc[0]
    assert "_lance_row_id" not in descriptor
    assert set(descriptor.keys()) == {"kind", "position", "size", "blob_id", "blob_uri"}


@pytest.mark.asyncio
async def test_async_typed_blob_projection_preserves_source_column():
    db = await lancedb.connect_async("memory:///typed_blob_projection")
    schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("blob")])
    table = await db.create_table("typed_blob_projection", schema=schema)
    await table.add([{"id": 1, "blob": b"alpha"}])

    hits = await table.query().select({"blob_alias": col("blob")}).to_arrow()

    assert "_lance_row_id" in hits.schema.field("blob_alias").type.names
    blobs = await table.fetch_blobs("blob", hits)
    assert blobs.to_pylist() == [b"alpha"]


def test_fetch_blobs_round_trip():
    table = _blob_table(
        "round_trip",
        [{"id": 1, "image": b"alpha"}, {"id": 2, "image": b"beta"}],
    )
    by_id = _row_ids_by_id(table)
    blobs = table.fetch_blobs("image", [by_id[1], by_id[2]])
    assert [blobs[0].as_py(), blobs[1].as_py()] == [b"alpha", b"beta"]


def test_merge_insert_writes_python_bytes():
    table = _blob_table("merge_bytes", [{"id": 1, "image": b"before"}])
    result = (
        table.merge_insert("id")
        .when_matched_update_all()
        .when_not_matched_insert_all()
        .execute([{"id": 1, "image": b"updated"}, {"id": 2, "image": b"inserted"}])
    )
    assert result.num_updated_rows == 1
    assert result.num_inserted_rows == 1
    by_id = _row_ids_by_id(table)
    blobs = table.fetch_blobs("image", [by_id[1], by_id[2]])
    assert blobs.to_pylist() == [b"updated", b"inserted"]


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")])
    table = db.create_table("images", schema=schema)
    table.add([{"id": 1, "image": b"hello"}])

    script = textwrap.dedent(
        f"""\
        import lancedb

        db = lancedb.connect({str(tmp_path)!r})
        table = db.open_table("images")
        image_type = table.schema.field("image").type
        if type(image_type).__name__ != "StructType":
            raise SystemExit(f"expected StructType, got {{type(image_type)}}")
        result = (
            table.merge_insert("id")
            .when_matched_update_all()
            .when_not_matched_insert_all()
            .execute(
                [{{"id": 1, "image": b"updated"}}, {{"id": 2, "image": b"inserted"}}]
            )
        )
        if result.num_updated_rows != 1 or result.num_inserted_rows != 1:
            raise SystemExit(
                f"rows updated={{result.num_updated_rows}} "
                f"inserted={{result.num_inserted_rows}}"
            )
        hits = table.search().with_row_id(True).limit(10).to_arrow()
        by_id = dict(zip(hits["id"].to_pylist(), hits["_rowid"].to_pylist()))
        blobs = table.fetch_blobs("image", [by_id[1], by_id[2]])
        if blobs.to_pylist() != [b"updated", b"inserted"]:
            raise SystemExit(blobs.to_pylist())
        """
    )
    result = subprocess.run(
        [sys.executable, "-c", script],
        capture_output=True,
        text=True,
        check=False,
    )
    assert result.returncode == 0, result.stderr


def test_merge_insert_bytes_after_reopen_without_pylance(tmp_path):
    db = lancedb.connect(tmp_path)
    schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
    table = db.create_table("images", schema=schema)
    table.add([{"id": 1, "image": b"hello"}])

    script = _HIDE_LANCE_BLOB + textwrap.dedent(
        f"""\
        import lancedb

        db = lancedb.connect({str(tmp_path)!r})
        table = db.open_table("images")
        image_type = table.schema.field("image").type
        if type(image_type).__name__ != "StructType":
            raise SystemExit(f"expected StructType, got {{type(image_type)}}")
        result = (
            table.merge_insert("id")
            .when_matched_update_all()
            .when_not_matched_insert_all()
            .execute(
                [{{"id": 1, "image": b"updated"}}, {{"id": 2, "image": b"inserted"}}]
            )
        )
        if result.num_updated_rows != 1 or result.num_inserted_rows != 1:
            raise SystemExit(
                f"rows updated={{result.num_updated_rows}} "
                f"inserted={{result.num_inserted_rows}}"
            )
        hits = table.search().with_row_id(True).limit(10).to_arrow()
        by_id = dict(zip(hits["id"].to_pylist(), hits["_rowid"].to_pylist()))
        blobs = table.fetch_blobs("image", [by_id[1], by_id[2]])
        if blobs.to_pylist() != [b"updated", b"inserted"]:
            raise SystemExit(blobs.to_pylist())
        """
    )
    result = subprocess.run(
        [sys.executable, "-c", script],
        capture_output=True,
        text=True,
        check=False,
    )
    assert result.returncode == 0, result.stderr


def test_merge_insert_blob_array_into_reopened_unregistered_table(tmp_path):
    db = lancedb.connect(tmp_path)
    schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
    table = db.create_table("images", schema=schema)
    table.add([{"id": 1, "image": b"before"}])

    script = textwrap.dedent(
        f"""\
        import pyarrow as pa
        import lancedb

        db = lancedb.connect({str(tmp_path)!r})
        table = db.open_table("images")
        image_type = table.schema.field("image").type
        if type(image_type).__name__ != "StructType":
            raise SystemExit(
                f"expected StructType before lance import, got {{type(image_type)}}"
            )

        import lance

        updates = pa.Table.from_arrays(
            [
                pa.array([1, 2], type=pa.int64()),
                lance.blob_array([b"updated", b"inserted"]),
            ],
            names=["id", "image"],
        )
        result = (
            table.merge_insert("id")
            .when_matched_update_all()
            .when_not_matched_insert_all()
            .execute(updates)
        )
        if result.num_updated_rows != 1 or result.num_inserted_rows != 1:
            raise SystemExit(
                f"rows updated={{result.num_updated_rows}} "
                f"inserted={{result.num_inserted_rows}}"
            )
        hits = table.search().with_row_id(True).limit(10).to_arrow()
        by_id = dict(zip(hits["id"].to_pylist(), hits["_rowid"].to_pylist()))
        blobs = table.fetch_blobs("image", [by_id[1], by_id[2]])
        if blobs.to_pylist() != [b"updated", b"inserted"]:
            raise SystemExit(blobs.to_pylist())
        """
    )
    result = subprocess.run(
        [sys.executable, "-c", script],
        capture_output=True,
        text=True,
        check=False,
    )
    assert result.returncode == 0, result.stderr


def test_add_all_null_blob_column():
    db = lancedb.connect("memory:///")
    schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
    table = db.create_table("all_null", schema=schema)
    table.add([{"id": 1, "image": None}, {"id": 2, "image": None}])
    by_id = _row_ids_by_id(table)
    blobs = table.fetch_blobs("image", [by_id[1], by_id[2]])
    assert blobs.to_pylist() == [None, None]


def test_create_table_nested_blob_schema_without_rows():
    db = lancedb.connect("memory:///")
    schema = pa.schema(
        [
            pa.field("id", pa.int64()),
            pa.field("info", pa.struct([lancedb.blob("blob")])),
            pa.field("images", pa.list_(lancedb.blob("image"))),
        ]
    )
    table = db.create_table("nested_empty", schema=schema)
    assert table.count_rows() == 0


def test_merge_insert_nested_blob_dicts():
    db = lancedb.connect("memory:///")
    info = pa.StructArray.from_arrays(
        [
            pa.array(["first"], type=pa.string()),
            _blob_array("blob", [b"before"]),
        ],
        names=["name", "blob"],
    )
    data = pa.Table.from_arrays(
        [pa.array([1], type=pa.int64()), info],
        names=["id", "info"],
    )
    table = db.create_table("nested_merge", data=data)
    result = (
        table.merge_insert("id")
        .when_matched_update_all()
        .execute([{"id": 1, "info": {"name": "first", "blob": b"after"}}])
    )
    assert result.num_updated_rows == 1
    by_id = _row_ids_by_id(table)
    blobs = table.fetch_blobs("info.blob", [by_id[1]])
    assert blobs.to_pylist() == [b"after"]


def _list_blob_table(name):
    db = lancedb.connect("memory:///")
    blob_field = lancedb.blob("image")
    images = pa.ListArray.from_arrays(
        pa.array([0, 1], type=pa.int32()), _blob_array("image", [b"before"])
    )
    data = pa.Table.from_arrays(
        [pa.array([1], type=pa.int64()), images],
        schema=pa.schema(
            [pa.field("id", pa.int64()), pa.field("images", pa.list_(blob_field))]
        ),
    )
    return db.create_table(name, data=data)


def test_merge_insert_list_blob_dicts():
    table = _list_blob_table("list_merge")
    result = (
        table.merge_insert("id")
        .when_matched_update_all()
        .when_not_matched_insert_all()
        .execute([{"id": 1, "images": [b"one", b"two"]}, {"id": 2, "images": None}])
    )
    assert result.num_updated_rows == 1
    assert result.num_inserted_rows == 1
    hits = table.search().limit(10).to_arrow()
    sizes = {
        row["id"]: None if row["images"] is None else [d["size"] for d in row["images"]]
        for row in hits.to_pylist()
    }
    assert sizes == {1: [3, 3], 2: None}


def test_list_blob_column_queries_as_raw_descriptors():
    table = _list_blob_table("list_query")
    hits = table.search().limit(10).to_arrow()
    element = hits.schema.field("images").type.value_type
    assert pa.types.is_struct(element)
    assert "_lance_row_id" not in element.names
    with pytest.raises(ValueError, match="expected struct before segment"):
        table.fetch_blobs("images.image", [0])


def test_row_addressable_paths_exclude_list_children():
    from lancedb.schema import row_addressable_blob_v2_paths

    schema = pa.schema(
        [
            pa.field("id", pa.int64()),
            pa.field("info", pa.struct([lancedb.blob("blob")])),
            pa.field("images", pa.list_(lancedb.blob("image"))),
        ]
    )
    assert blob_v2_column_paths(schema) == ["info.blob", "images.image"]
    assert row_addressable_blob_v2_paths(schema) == ["info.blob"]


def test_merge_insert_writes_pylance_blob_array():
    table = _blob_table("merge_pylance", [{"id": 1, "image": b"before"}])
    image = lance.blob_array([b"updated", b"inserted"])
    assert type(image.type) is LanceBlobType
    assert type(image.type) is type(lancedb.BlobType())
    updates = pa.Table.from_arrays(
        [pa.array([1, 2], type=pa.int64()), image], names=["id", "image"]
    )

    result = (
        table.merge_insert("id")
        .when_matched_update_all()
        .when_not_matched_insert_all()
        .execute(updates)
    )

    assert result.num_updated_rows == 1
    assert result.num_inserted_rows == 1
    by_id = _row_ids_by_id(table)
    blobs = table.fetch_blobs("image", [by_id[1], by_id[2]])
    assert blobs.to_pylist() == [b"updated", b"inserted"]


def test_fetch_blobs_accepts_query_result():
    table = _blob_table("from_result", [{"id": 1, "image": b"gamma"}])
    hits = table.search().limit(10).to_arrow()
    assert "_rowid" not in hits.column_names
    blobs = table.fetch_blobs("image", hits)
    assert {blobs[i].as_py() for i in range(len(blobs))} == {b"gamma"}


def test_fetch_blobs_preserves_null_and_empty_values():
    table = _blob_table(
        "nulls",
        [
            {"id": 1, "image": b"present"},
            {"id": 2, "image": None},
            {"id": 3, "image": b""},
        ],
    )
    by_id = _row_ids_by_id(table)
    request = [by_id[1], by_id[2], by_id[3], by_id[1]]
    blobs = table.fetch_blobs("image", request)
    assert len(blobs) == len(request)
    assert blobs[0].as_py() == b"present"
    assert blobs[1].as_py() is None
    assert blobs[2].as_py() == b""
    assert blobs[3].as_py() == b"present"


def test_fetch_blob_ranges_aligns_repeated_ranges_and_nulls():
    table = _blob_table(
        "range_alignment",
        [{"id": 1, "image": b"abcdefghij"}, {"id": 2, "image": None}],
    )
    by_id = _row_ids_by_id(table)
    requests = [
        (by_id[1], 2, 3),
        (by_id[2], 0, 0),
        (by_id[1], 0, 2),
        (by_id[1], 2, 3),
        (by_id[1], 10, 0),
    ]

    ranges = table.fetch_blob_ranges("image", requests)

    assert ranges.to_pylist() == [b"cde", None, b"ab", b"cde", b""]


def test_fetch_blob_ranges_validates_requests():
    table = _blob_table("range_validation", [{"id": 1, "image": b"abc"}])
    row_id = _row_ids_by_id(table)[1]

    with pytest.raises(ValueError, match="exceeds blob size"):
        table.fetch_blob_ranges("image", [(row_id, 2, 2)])

    with pytest.raises(ValueError, match="offset \\+ length overflowed"):
        table.fetch_blob_ranges("image", [(row_id, 2**64 - 1, 1)])

    with pytest.raises(ValueError, match="row IDs"):
        table.fetch_blob_ranges("image", [(2**64 - 1, 0, 1)])


def test_fetch_blob_ranges_empty_requests_returns_empty_array():
    table = _blob_table("range_empty", [{"id": 1, "image": b"x"}])
    assert table.fetch_blob_ranges("image", []).to_pylist() == []


@pytest.mark.asyncio
async def test_async_fetch_blob_ranges():
    db = await lancedb.connect_async("memory:///")
    schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
    table = await db.create_table("range_async", schema=schema)
    await table.add([{"id": 1, "image": b"abcdefghij"}])
    hits = await table.query().with_row_id().to_arrow()
    row_id = hits["_rowid"][0].as_py()

    ranges = await table.fetch_blob_ranges("image", [(row_id, 1, 3), (row_id, 6, 2)])

    assert ranges.to_pylist() == [b"bcd", b"gh"]


def test_fetch_blobs_nested_path():
    db = lancedb.connect("memory:///")
    info = pa.StructArray.from_arrays(
        [
            pa.array(["first", "second"], type=pa.string()),
            _blob_array("blob", [b"nested-alpha", b"nested-beta"]),
        ],
        names=["name", "blob"],
    )
    data = pa.Table.from_arrays(
        [pa.array([1, 2], type=pa.int64()), info],
        names=["id", "info"],
    )
    table = db.create_table("nested", data=data)

    by_id = _row_ids_by_id(table)
    blobs = table.fetch_blobs("info.blob", [by_id[1], by_id[2]])
    assert [blobs[0].as_py(), blobs[1].as_py()] == [b"nested-alpha", b"nested-beta"]


def test_fetch_blob_files_lazy_read():
    payload = b"lazy-read" * 100
    table = _blob_table("lazy", [{"id": 1, "image": payload}])
    by_id = _row_ids_by_id(table)
    handles = table.fetch_blob_files("image", [by_id[1]])
    assert len(handles) == 1
    assert handles[0].read() == payload


def test_fetch_blob_files_null_alignment():
    table = _blob_table(
        "lazy_nulls",
        [{"id": 1, "image": b"here"}, {"id": 2, "image": None}],
    )
    by_id = _row_ids_by_id(table)
    handles = table.fetch_blob_files("image", [by_id[2], by_id[1]])
    assert len(handles) == 2
    assert handles[0] is None
    assert handles[1].read() == b"here"


def test_fetch_blobs_rejects_non_blob_column():
    table = _blob_table("reject", [{"id": 1, "image": b"x"}])
    with pytest.raises(ValueError, match="not a blob column"):
        table.fetch_blobs("id", [0])


def test_legacy_v1_query_omits_auto_row_id():
    table = _legacy_v1_table("legacy_v1")
    hits = table.search().select(["legacy"]).limit(10).to_arrow()
    assert "_rowid" not in hits.column_names


def test_fetch_blobs_rejects_legacy_v1_column():
    table = _legacy_v1_table("legacy_fetch")
    with pytest.raises(ValueError, match="legacy blob column.*blob v2"):
        table.fetch_blobs("legacy", [0])


@pytest.mark.asyncio
async def test_async_fetch_blob_files_lazy_read():
    db = await lancedb.connect_async("memory:///")
    schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
    table = await db.create_table("async_lazy", schema=schema)
    payload = b"async-lazy" * 100
    await table.add([{"id": 1, "image": payload}])
    hits = (
        await table.query().select({"image_alias": "image"}).limit(10).to_arrow()
    ).combine_chunks()
    assert "_rowid" not in hits.column_names
    handles = await table.fetch_blob_files("image", hits)
    assert len(handles) == 1
    assert await handles[0].aread() == payload


def test_fetch_blobs_from_query_result_without_row_id_raises():
    table = _blob_table("no_rowid", [{"id": 1, "image": b"x"}])
    hits = table.search().select(["id"]).to_arrow()
    assert "_rowid" not in hits.column_names
    with pytest.raises(ValueError, match="_rowid"):
        table.fetch_blobs("image", hits)


_HYBRID_BLOB_SCHEMA = pa.schema(
    [
        pa.field("id", pa.int64()),
        pa.field("text", pa.utf8()),
        pa.field("vector", pa.list_(pa.float32(), list_size=2)),
        lancedb.blob("image"),
    ]
)
_HYBRID_BLOB_ROWS = [
    {"id": 1, "text": "hello alpha", "vector": [1.0, 0.0], "image": b"alpha"},
    {"id": 2, "text": "hello beta", "vector": [0.9, 0.1], "image": b"beta"},
    {"id": 3, "text": "other", "vector": [0.0, 1.0], "image": b"other"},
]


def _hybrid_blob_table(db):
    table = db.create_table("hybrid_blob_fetch", schema=_HYBRID_BLOB_SCHEMA)
    table.add(_HYBRID_BLOB_ROWS)
    table.create_index("text", config=FTS(with_position=False))
    return table


async def _hybrid_blob_table_async(db):
    table = await db.create_table("hybrid_blob_fetch_async", schema=_HYBRID_BLOB_SCHEMA)
    await table.add(_HYBRID_BLOB_ROWS)
    await table.create_index("text", config=FTS(with_position=False))
    return table


def test_blob_v2_hybrid_fetch_blobs():
    table = _hybrid_blob_table(lancedb.connect("memory:///"))
    hits = (
        table.search(query_type="hybrid")
        .vector([1.0, 0.0])
        .text("hello")
        .select(["id", "image"])
        .limit(2)
        .to_arrow()
    )

    assert "_rowid" not in hits.column_names
    assert "_lance_row_id" in hits.schema.field("image").type.names
    blobs = table.fetch_blobs("image", hits)
    assert {blobs[i].as_py() for i in range(len(blobs))} == {b"alpha", b"beta"}


@pytest.mark.asyncio
async def test_blob_v2_hybrid_fetch_blobs_async():
    db = await lancedb.connect_async("memory:///hybrid_blob_fetch_async")
    table = await _hybrid_blob_table_async(db)
    hits = await (
        table.query()
        .nearest_to([1.0, 0.0])
        .nearest_to_text("hello")
        .select(["id", "image"])
        .limit(2)
        .to_arrow()
    )

    assert "_rowid" not in hits.column_names
    assert "_lance_row_id" in hits.schema.field("image").type.names
    blobs = await table.fetch_blobs("image", hits)
    assert {blobs[i].as_py() for i in range(len(blobs))} == {b"alpha", b"beta"}


@pytest.mark.asyncio
async def test_async_hybrid_typed_blob_projection_preserves_source_column():
    db = await lancedb.connect_async("memory:///hybrid_typed_blob")
    schema = pa.schema(
        [
            pa.field("id", pa.int64()),
            pa.field("text", pa.utf8()),
            pa.field("vector", pa.list_(pa.float32(), list_size=2)),
            lancedb.blob("blob"),
        ]
    )
    table = await db.create_table("hybrid_typed_blob", schema=schema)
    await table.add(
        [
            {
                "id": 1,
                "text": "hello alpha",
                "vector": [1.0, 0.0],
                "blob": b"alpha",
            },
            {
                "id": 2,
                "text": "hello beta",
                "vector": [0.9, 0.1],
                "blob": b"beta",
            },
        ]
    )
    await table.create_index("text", config=FTS(with_position=False))

    hits = await (
        table.query()
        .nearest_to([1.0, 0.0])
        .nearest_to_text("hello")
        .select({"blob_alias": col("blob")})
        .limit(2)
        .to_arrow()
    )

    assert "_lance_row_id" in hits.schema.field("blob_alias").type.names
    blobs = await table.fetch_blobs("blob", hits)
    assert {blobs[i].as_py() for i in range(len(blobs))} == {b"alpha", b"beta"}


def test_blob_file_seek_read_and_read_range():
    payload = _identifiable_payload(1024)
    table = _blob_table("seek_read", [{"id": 1, "image": payload}])
    by_id = _row_ids_by_id(table)
    handle = table.fetch_blob_files("image", [by_id[1]])[0]

    assert handle.seek(100) == 100
    assert handle.read(16) == payload[100:116]
    handle.seek(100)
    assert handle.read_range(500, 8) == payload[500:508]
    assert handle.tell() == 100

    with pytest.raises(ValueError, match="whence"):
        handle.seek(0, 99)


def test_fetch_blob_files_from_query_partial_read():
    payload = _identifiable_payload(65536)
    table = _blob_table("query_partial", [{"id": 1, "image": payload}])
    hits = table.search().select(["id", "image"]).limit(1).to_arrow()
    assert "_rowid" not in hits.column_names

    handle = table.fetch_blob_files("image", hits)[0]
    assert handle.size() == 65536
    assert handle.read_range(0, 128) == payload[:128]
    assert handle.tell() == 0
    assert handle.seek(40000) == 40000
    assert handle.read(16) == payload[40000:40016]


def test_blob_file_buffered_reader():
    payload = _identifiable_payload(4096)
    table = _blob_table("buffered_reader", [{"id": 1, "image": payload}])
    hits = table.search().select(["id", "image"]).limit(1).to_arrow()
    handle = table.fetch_blob_files("image", hits)[0]
    reader = io.BufferedReader(handle)
    assert reader.read(8) == payload[:8]
    assert reader.read(8) == payload[8:16]
    assert reader.read() == payload[16:]


def test_fetch_blob_files_cross_fragment_nulls_and_dups():
    db = lancedb.connect("memory:///")
    schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
    table = db.create_table("cross_fragment", schema=schema)
    table.add([{"id": 1, "image": b"alpha"}])
    table.add([{"id": 2, "image": None}, {"id": 3, "image": b"beta"}])

    by_id = _row_ids_by_id(table)
    request = [by_id[3], by_id[2], by_id[1], by_id[3]]
    handles = table.fetch_blob_files("image", request)
    assert len(handles) == 4
    assert handles[1] is None
    assert handles[0].read() == b"beta"
    assert handles[2].read() == b"alpha"
    assert handles[3].seek(1) == 1
    assert handles[3].read() == b"eta"


def test_blob_file_pyav_decode_seek(tmp_path):
    av = pytest.importorskip("av")
    import fractions

    clip = tmp_path / "clip.mp4"
    with av.open(str(clip), mode="w") as container:
        stream = container.add_stream("mpeg4", rate=5)
        stream.width, stream.height, stream.pix_fmt = 32, 32, "yuv420p"
        stream.time_base = fractions.Fraction(1, 5)
        for pts in range(5):
            frame = av.VideoFrame(32, 32, "yuv420p")
            frame.pts = pts
            container.mux(stream.encode(frame))
        container.mux(stream.encode(None))

    table = _blob_table("pyav", [{"id": 1, "image": clip.read_bytes()}])
    hits = table.search().select(["image"]).limit(1).to_arrow()
    handle = table.fetch_blob_files("image", hits)[0]

    with av.open(handle) as container:
        stream = container.streams.video[0]
        container.seek(0)
        assert next(container.decode(stream)) is not None


def test_blob_v2_hybrid_fetch_blob_files_seek():
    table = _hybrid_blob_table(lancedb.connect("memory:///"))
    hits = (
        table.search(query_type="hybrid")
        .vector([1.0, 0.0])
        .text("hello")
        .select(["id", "image"])
        .limit(2)
        .to_arrow()
    )
    assert "_rowid" not in hits.column_names

    handles = table.fetch_blob_files("image", hits)
    assert len(handles) == 2
    assert {handle.read_range(0, 2) for handle in handles} == {b"al", b"be"}
    first = handles[0]
    assert first.seek(1) == 1
    assert first.read(2) in {b"lp", b"et"}


def test_blob_file_header_sniff_from_search():
    payload = b"%PDF-1.7\n" + bytes(4096)
    table = _blob_table("header_sniff", [{"id": 1, "image": payload}])
    hits = table.search().select(["id", "image"]).limit(1).to_arrow()
    handle = table.fetch_blob_files("image", hits)[0]
    assert handle.read_range(0, 4) == b"%PDF"
    assert handle.tell() == 0


def test_blob_file_multiple_handles_independent_cursors():
    table = _blob_table(
        "multi_handle",
        [{"id": 1, "image": b"first-payload"}, {"id": 2, "image": b"second-payload"}],
    )
    by_id = _row_ids_by_id(table)
    first, second = table.fetch_blob_files("image", [by_id[1], by_id[2]])
    assert first.seek(6) == 6
    assert second.tell() == 0
    assert first.read(7) == b"payload"
    assert second.read(6) == b"second"


def test_fetch_blob_files_nested_path_seek():
    db = lancedb.connect("memory:///")
    info = pa.StructArray.from_arrays(
        [
            pa.array(["first", "second"], type=pa.string()),
            _blob_array("blob", [b"nested-alpha", b"nested-beta"]),
        ],
        names=["name", "blob"],
    )
    data = pa.Table.from_arrays(
        [pa.array([1, 2], type=pa.int64()), info],
        names=["id", "info"],
    )
    table = db.create_table("nested_seek", data=data)
    by_id = _row_ids_by_id(table)
    handle = table.fetch_blob_files("info.blob", [by_id[2]])[0]
    assert handle.seek(7) == 7
    assert handle.read() == b"beta"


def test_fetch_blobs_survives_sort_after_query():
    table = _blob_table(
        "sort_survives",
        [{"id": i, "image": f"payload-{i}".encode()} for i in range(5)],
    )
    hits = table.search().select(["id", "image"]).to_arrow()
    sort_idx = pc.sort_indices(hits["id"], sort_keys=[("id", "descending")])
    sorted_hits = hits.take(sort_idx)

    blobs = table.fetch_blobs("image", sorted_hits)
    expected = [f"payload-{i}".encode() for i in sorted_hits["id"].to_pylist()]
    assert [blobs[i].as_py() for i in range(len(blobs))] == expected


def test_fetch_blobs_survives_filter_and_sort_after_query():
    table = _blob_table(
        "filter_sort_survives",
        [{"id": i, "image": f"payload-{i}".encode()} for i in range(5)],
    )
    hits = table.search().select(["id", "image"]).to_arrow()
    filtered = hits.filter(pc.field("id") >= 2)
    sort_idx = pc.sort_indices(filtered["id"], sort_keys=[("id", "descending")])
    filtered_sorted = filtered.take(sort_idx)

    blobs = table.fetch_blobs("image", filtered_sorted)
    expected = [f"payload-{i}".encode() for i in filtered_sorted["id"].to_pylist()]
    assert [blobs[i].as_py() for i in range(len(blobs))] == expected


def test_fetch_blob_files_survives_sort_after_query():
    table = _blob_table(
        "lazy_sort_survives",
        [{"id": i, "image": f"payload-{i}".encode()} for i in range(5)],
    )
    hits = table.search().select(["id", "image"]).to_arrow()
    sort_idx = pc.sort_indices(hits["id"], sort_keys=[("id", "descending")])
    sorted_hits = hits.take(sort_idx)

    handles = table.fetch_blob_files("image", sorted_hits)
    expected = [f"payload-{i}".encode() for i in sorted_hits["id"].to_pylist()]
    assert [handle.read() for handle in handles] == expected


def test_fetch_blobs_nested_path_survives_sort_after_query():
    db = lancedb.connect("memory:///")
    values = [f"payload-{i}".encode() for i in range(4)]
    info = pa.StructArray.from_arrays(
        [pa.array(["row"] * 4, type=pa.string()), _blob_array("blob", values)],
        names=["name", "blob"],
    )
    data = pa.Table.from_arrays(
        [pa.array(range(4), type=pa.int64()), info],
        names=["id", "info"],
    )
    table = db.create_table("nested_sort_survives", data=data)

    hits = table.search().to_arrow()
    sort_idx = pc.sort_indices(hits["id"], sort_keys=[("id", "descending")])
    sorted_hits = hits.take(sort_idx)

    blobs = table.fetch_blobs("info.blob", sorted_hits)
    expected = [f"payload-{i}".encode() for i in sorted_hits["id"].to_pylist()]
    assert [blobs[i].as_py() for i in range(len(blobs))] == expected


def _identifiable_payload(size: int) -> bytes:
    block = 256
    return b"".join(bytes([i % 256]) * block for i in range(size // block))


def _external_uri_blob_array(uris):
    blob_type = lancedb.blob("image").type
    storage_type = blob_type.storage_type
    child_names = [field.name for field in storage_type]
    assert "uri" in child_names, "blob layout no longer has a uri child"
    children = [
        pa.array(uris if field.name == "uri" else [None] * len(uris), type=field.type)
        for field in storage_type
    ]
    storage = pa.StructArray.from_arrays(children, fields=list(storage_type))
    return pa.ExtensionArray.from_storage(blob_type, storage)


def _external_uri_table_and_rows(name, uris):
    db = lancedb.connect("memory:///")
    schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
    table = db.create_table(name, schema=schema)
    rows = pa.Table.from_arrays(
        [
            pa.array(range(len(uris)), type=pa.int64()),
            _external_uri_blob_array(uris),
        ],
        schema=schema,
    )
    return table, rows


def test_add_external_uri_struct_round_trips_with_flag(tmp_path):
    payload = b"external-uri-bytes"
    blob_path = tmp_path / "payload.bin"
    blob_path.write_bytes(payload)

    table, rows = _external_uri_table_and_rows("external_struct", [blob_path.as_uri()])
    table.add(rows, allow_external_blob_outside_bases=True)

    hits = table.search().to_arrow()
    blobs = table.fetch_blobs("image", hits)
    assert blobs[0].as_py() == payload


def test_add_external_uri_without_flag_raises(tmp_path):
    blob_path = tmp_path / "payload.bin"
    blob_path.write_bytes(b"unreachable")

    table, rows = _external_uri_table_and_rows("external_no_flag", [blob_path.as_uri()])
    with pytest.raises(ValueError, match="allow_external_blob_outside_bases"):
        table.add(rows)
    assert table.count_rows() == 0


def test_add_external_uri_string_round_trips_with_flag(tmp_path):
    payload = b"external-uri-bytes"
    blob_path = tmp_path / "payload.bin"
    blob_path.write_bytes(payload)

    db = lancedb.connect("memory:///")
    schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
    table = db.create_table("external_string", schema=schema)
    table.add(
        [{"id": 1, "image": blob_path.as_uri()}],
        allow_external_blob_outside_bases=True,
    )

    hits = table.search().to_arrow()
    blobs = table.fetch_blobs("image", hits)
    assert blobs[0].as_py() == payload