# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import io
import json
import os
import shutil
import tempfile
import unittest

import pyarrow as pa
import pypaimon.multimodal as pmm
from pypaimon.multimodal import source_col
from pypaimon.common.predicate_builder import PredicateBuilder
from pypaimon import Schema as PaimonSchema
from pypaimon.globalindex.global_index_result import GlobalIndexResult
from pypaimon.utils.range import Range

try:
    import ray
except ImportError:
    ray = None


_PARQUET_OPTIONS = {
    "row-tracking.enabled": "true",
    "data-evolution.enabled": "true",
    "deletion-vectors.enabled": "true",
    "file.format": "parquet",
    "vector.file.format": "parquet",
}


def _schema(fields):
    return pa.schema([
        pa.field(name, field_type)
        for name, field_type in fields.items()
    ])


def _vector(dim):
    return pa.list_(pa.float32(), dim)


def _raw_schema(options=None, primary_keys=None):
    return PaimonSchema.from_pyarrow_schema(
        _schema({"id": pa.int32(), "name": pa.string()}),
        primary_keys=list(primary_keys or []),
        options=dict(options or {"file.format": "parquet"}),
    )


class MultimodalTableTest(unittest.TestCase):

    def setUp(self):
        self.temp_dir = tempfile.mkdtemp(prefix="pypaimon_mm_")
        self.warehouse = os.path.join(self.temp_dir, "warehouse")
        self.conn = pmm.connect(options={"warehouse": self.warehouse})

    def tearDown(self):
        shutil.rmtree(self.temp_dir, ignore_errors=True)

    def test_connect_accepts_options(self):
        conn = pmm.connect(
            database="analytics",
            options={
                "warehouse": os.path.join(self.temp_dir, "warehouse_options"),
            },
        )

        table = conn.create_table(
            "docs",
            schema=_schema({"id": pa.int32()}),
            options=_PARQUET_OPTIONS,
        )

        self.assertEqual("analytics.docs", table.identifier)

    def test_connect_rejects_positional_warehouse(self):
        with self.assertRaises(TypeError):
            pmm.connect(self.warehouse)

    def test_create_table_defaults_data_evolution_options(self):
        table = self.conn.create_table(
            "docs",
            schema=_schema({
                "id": pa.int32(),
                "content": pa.string(),
                "embedding": _vector(3),
                "payload": pa.large_binary(),
            }),
        )

        options = table.raw_table.table_schema.options
        self.assertEqual("true", options["row-tracking.enabled"])
        self.assertEqual("true", options["data-evolution.enabled"])
        self.assertEqual("true", options["deletion-vectors.enabled"])
        self.assertEqual("true", options["blob-as-descriptor"])
        self.assertNotIn("data-evolution.row-sidecar.enabled", options)
        self.assertEqual("full", options["global-index.search-mode"])
        self.assertEqual("vortex", options["vector.file.format"])
        self.assertEqual("default.docs", self.conn.get_table("docs").identifier)
        self.assertEqual(["id", "content", "embedding", "payload"],
                         [field.name for field in table.raw_table.fields])

    def test_create_table_uses_options_and_partitioned(self):
        table = self.conn.create_table(
            "docs",
            schema=_schema({
                "id": pa.int32(),
                "embedding": _vector(3),
                "dt": pa.string(),
            }),
            options=dict(_PARQUET_OPTIONS, **{
                "deletion-vectors.enabled": "false",
            }),
            partitioned=["dt"],
        )

        options = table.raw_table.table_schema.options
        self.assertEqual(["dt"], table.raw_table.table_schema.partition_keys)
        self.assertEqual("false", options["deletion-vectors.enabled"])
        self.assertEqual("parquet", options["file.format"])
        self.assertEqual("parquet", options["vector.file.format"])

    def test_create_table_rejects_non_arrow_schema(self):
        with self.assertRaisesRegex(ValueError, "pyarrow.Schema"):
            self.conn.create_table("docs", schema={"id": pa.int32()})

    def test_create_table_accepts_pyarrow_schema_types(self):
        table = self.conn.create_table(
            "typed",
            schema=pa.schema([
                pa.field("flag", pa.bool_(), nullable=False),
                pa.field("tiny", pa.int8()),
                pa.field("small", pa.int16()),
                pa.field("id", pa.int64()),
                pa.field("score", pa.float32()),
                pa.field("ratio", pa.float64()),
                pa.field("title", pa.string()),
                pa.field("payload", pa.binary()),
                pa.field("hash", pa.binary(16)),
                pa.field("blob", pa.large_binary()),
                pa.field("amount", pa.decimal128(12, 2)),
                pa.field("dt", pa.date32()),
                pa.field("created_at", pa.timestamp("us")),
                pa.field("created_ltz", pa.timestamp("us", tz="UTC")),
                pa.field("tags", pa.list_(pa.string())),
                pa.field("attrs", pa.map_(pa.string(), pa.int32())),
                pa.field("meta", pa.struct([pa.field("rank", pa.int32())])),
                pa.field("embedding", _vector(3)),
            ]),
        )

        types_by_name = {
            field.name: str(field.type) for field in table.raw_table.fields
        }
        self.assertEqual("BOOLEAN NOT NULL", types_by_name["flag"])
        self.assertEqual("TINYINT", types_by_name["tiny"])
        self.assertEqual("SMALLINT", types_by_name["small"])
        self.assertEqual("BIGINT", types_by_name["id"])
        self.assertEqual("FLOAT", types_by_name["score"])
        self.assertEqual("DOUBLE", types_by_name["ratio"])
        self.assertEqual("STRING", types_by_name["title"])
        self.assertEqual("BYTES", types_by_name["payload"])
        self.assertEqual("BINARY(16)", types_by_name["hash"])
        self.assertEqual("BLOB", types_by_name["blob"])
        self.assertEqual("DECIMAL(12, 2)", types_by_name["amount"])
        self.assertEqual("DATE", types_by_name["dt"])
        self.assertEqual("TIMESTAMP(6)", types_by_name["created_at"])
        self.assertEqual("TIMESTAMP_LTZ(6)", types_by_name["created_ltz"])
        self.assertEqual("ARRAY<STRING>", types_by_name["tags"])
        self.assertEqual("MAP<STRING NOT NULL, INT>", types_by_name["attrs"])
        self.assertEqual("ROW<rank: INT>", types_by_name["meta"])
        self.assertEqual("VECTOR<FLOAT, 3>", types_by_name["embedding"])

    def test_blob_store_put_objects_get_list_and_delete(self):
        table = self.conn.create_table(
            "objects",
            schema=_schema({
                "key": pa.string(),
                "image": pa.large_binary(),
                "content_type": pa.string(),
                "owner": pa.string(),
            }),
            options=_PARQUET_OPTIONS,
        )
        store = table.blobs(column="image")

        results = store.put_objects([
            {
                "key": "images/cat.jpg",
                "body": b"cat-image-old",
                "columns": {"content_type": "image/gif", "owner": "ignored"},
            },
            {
                "key": "images/cat.jpg",
                "body": b"cat-image-v1",
                "columns": {"content_type": "image/jpeg", "owner": "alice"},
            },
            {
                "key": "images/dog.jpg",
                "body": bytearray(b"dog-image"),
                "columns": {"content_type": "image/jpeg", "owner": "bob"},
            },
        ])

        self.assertEqual(["images/cat.jpg", "images/dog.jpg"],
                         [result.key for result in results])
        self.assertEqual([12, 9], [result.size for result in results])

        cat = store.get_object("images/cat.jpg")
        self.assertEqual(b"cat-image-v1", cat.read())
        self.assertEqual(b"at-i", store.get_object(
            "images/cat.jpg", range="bytes=1-4").read())
        clipped = store.get_object("images/cat.jpg", range="bytes=10-999")
        self.assertEqual(2, clipped.content_length)
        self.assertEqual(b"v1", clipped.read())
        with self.assertRaisesRegex(ValueError, "Range start"):
            store.get_object("images/cat.jpg", range="bytes=12-13")
        self.assertEqual("image/jpeg", cat.columns["content_type"])
        self.assertEqual("alice", cat.columns["owner"])
        owner_only = store.get_object(
            "images/cat.jpg",
            columns=["owner"],
        )
        self.assertEqual({"owner": "alice"}, owner_only.columns)

        listed = store.list_objects(prefix="images/")
        self.assertEqual(["images/cat.jpg", "images/dog.jpg"],
                         sorted(obj.key for obj in listed))
        self.assertEqual(
            {
                "images/cat.jpg": "image/jpeg",
                "images/dog.jpg": "image/jpeg",
            },
            {obj.key: obj.columns["content_type"] for obj in listed},
        )
        listed_without_columns = store.list_objects(
            prefix="images/",
            columns=[],
        )
        self.assertEqual(
            {"images/cat.jpg": {}, "images/dog.jpg": {}},
            {obj.key: obj.columns for obj in listed_without_columns},
        )
        self.assertEqual([], store.list_objects(prefix="images/", limit=0))
        with self.assertRaisesRegex(ValueError, "limit"):
            store.list_objects(prefix="images/", limit=-1)

        store.put_object(
            "images/cat.jpg",
            b"cat-image-v2",
            columns={"content_type": "image/png", "owner": "alice"},
        )
        self.assertEqual(b"cat-image-v2", store.get_object("images/cat.jpg").read())
        info = store.head_object("images/cat.jpg")
        self.assertEqual("image/png", info.columns["content_type"])
        self.assertEqual(
            {"content_type": "image/png"},
            store.head_object(
                "images/cat.jpg",
                columns="content_type",
            ).columns,
        )
        self.assertEqual(2, table.scan().to_arrow().num_rows)

        previous_descriptor = info.descriptor
        updated = store.update_object_columns(
            "images/cat.jpg",
            {"content_type": "image/webp"},
        )
        self.assertEqual(previous_descriptor, updated.descriptor)
        self.assertEqual("image/webp", updated.columns["content_type"])
        self.assertEqual("alice", updated.columns["owner"])
        self.assertEqual(b"cat-image-v2", store.get_object("images/cat.jpg").read())

        batch_updates = store.update_objects_columns([
            {"key": "images/cat.jpg", "columns": {"owner": "carol"}},
            {"key": "images/dog.jpg", "columns": {"owner": "dave"}},
        ])
        self.assertEqual(
            ["images/cat.jpg", "images/dog.jpg"],
            [obj.key for obj in batch_updates],
        )
        self.assertEqual("carol", store.head_object("images/cat.jpg").columns["owner"])
        self.assertEqual("dave", store.head_object("images/dog.jpg").columns["owner"])
        self.assertEqual(2, table.scan().to_arrow().num_rows)

        with self.assertRaisesRegex(ValueError, "columns must not be empty"):
            store.update_object_columns("images/cat.jpg", {})
        with self.assertRaises(pmm.NoSuchKey):
            store.update_object_columns("images/missing.jpg", {"owner": "nobody"})
        with self.assertRaisesRegex(ValueError, "columns must not include"):
            store.update_object_columns("images/cat.jpg", {"image": b"new"})

        store.delete_objects(["images/dog.jpg", "images/dog.jpg"])
        with self.assertRaises(pmm.NoSuchKey):
            store.head_object("images/dog.jpg")
        self.assertEqual(["images/cat.jpg"],
                         [obj.key for obj in store.list_objects(prefix="images/")])
        store.delete_object("images/cat.jpg")
        self.assertEqual([], store.list_objects(prefix="images/"))

    def test_blob_store_put_object_accepts_blob_without_materializing(self):
        from pypaimon.table.row.blob import Blob, BlobDescriptor

        table = self.conn.create_table(
            "streamed_objects",
            schema=_schema({
                "key": pa.string(),
                "payload": pa.large_binary(),
            }),
            options=_PARQUET_OPTIONS,
        )
        data = b"streamed-managed-payload"
        stream_uri = "stream://payloads/1"
        stream_calls = []

        class StreamOnlyReader:

            def new_input_stream(self, uri):
                stream_calls.append(("new_input_stream", uri))
                return io.BytesIO(data)

        class StreamOnlyReaderFactory:

            def create(self, uri):
                stream_calls.append(("create", uri))
                return StreamOnlyReader()

        class DescriptorOnlyBlob(Blob):

            def to_data(self):
                raise AssertionError("put_object should not materialize Blob data")

            def to_descriptor(self):
                return BlobDescriptor(stream_uri, 0, len(data))

            def new_input_stream(self):
                raise AssertionError("put_object should use the URI stream")

        store = table.blobs(column="payload")
        original_factory = table.raw_table.file_io.uri_reader_factory
        table.raw_table.file_io.uri_reader_factory = StreamOnlyReaderFactory()
        try:
            result = store.put_object("payloads/1", DescriptorOnlyBlob())
        finally:
            table.raw_table.file_io.uri_reader_factory = original_factory

        self.assertEqual(len(data), result.size)
        self.assertEqual([
            ("create", stream_uri),
            ("new_input_stream", stream_uri),
        ], stream_calls)
        stored = store.get_object("payloads/1")
        self.assertNotEqual(stream_uri, stored.descriptor.uri)
        self.assertEqual(data, stored.read())

    def test_blob_store_put_object_reference_preserves_descriptor_uri(self):
        table = self.conn.create_table(
            "referenced_objects",
            schema=_schema({
                "object_key": pa.string(),
                "payload": pa.large_binary(),
                "media_type": pa.string(),
            }),
            options=dict(_PARQUET_OPTIONS, **{
                "blob-descriptor-field": "payload",
            }),
        )
        data = b"external-video-payload"
        external_path = os.path.join(self.temp_dir, "video.bin")
        with open(external_path, "wb") as f:
            f.write(data)

        store = table.blobs(column="payload")
        result = store.put_object(
            "videos/1",
            uri=external_path,
            length=len(data),
            columns={"media_type": "video/mp4"},
        )
        descriptor = store.head_object("videos/1").descriptor
        more = store.put_objects([
            {
                "key": "videos/2",
                "descriptor": descriptor,
                "columns": {"media_type": "video/mp4"},
            }
        ])

        self.assertEqual("videos/1", result.key)
        self.assertEqual(len(data), result.size)
        self.assertEqual("videos/2", more[0].key)
        obj = store.get_object("videos/1")
        self.assertEqual(data, obj.read())
        self.assertEqual("video/mp4", obj.columns["media_type"])
        self.assertEqual(external_path, store.head_object("videos/1").descriptor.uri)
        self.assertEqual(data, store.get_object("videos/2").read())

    def test_blob_store_put_object_uri_streams_into_managed_blob(self):
        table = self.conn.create_table(
            "managed_only_objects",
            schema=_schema({
                "key": pa.string(),
                "payload": pa.large_binary(),
            }),
            options=_PARQUET_OPTIONS,
        )
        data = b"external-managed-payload"
        external_path = os.path.join(self.temp_dir, "managed.bin")
        with open(external_path, "wb") as f:
            f.write(data)

        store = table.blobs(column="payload")
        result = store.put_object(
            "payloads/1",
            uri=external_path,
            length=len(data),
        )

        self.assertEqual(len(data), result.size)
        stored = store.get_object("payloads/1")
        self.assertNotEqual(external_path, stored.descriptor.uri)
        self.assertEqual(data, stored.read())

    def test_drop_table_can_ignore_missing_table(self):
        self.conn.drop_table("missing", ignore_if_not_exists=True)

    def test_get_table_rejects_non_data_evolution_table(self):
        self.conn.catalog.create_database("default", ignore_if_exists=True)
        self.conn.catalog.create_table(
            "default.raw",
            _raw_schema(),
            False,
        )

        with self.assertRaisesRegex(ValueError, "data-evolution.enabled"):
            self.conn.get_table("raw")

    def test_get_table_rejects_primary_key_table(self):
        self.conn.catalog.create_database("default", ignore_if_exists=True)
        self.conn.catalog.create_table(
            "default.pk",
            _raw_schema(
                options=dict(_PARQUET_OPTIONS, **{"bucket": "1"}),
                primary_keys=["id"],
            ),
            False,
        )

        with self.assertRaisesRegex(ValueError, "primary keys"):
            self.conn.get_table("pk")

    def test_create_table_can_add_initial_data_and_get_by_short_name(self):
        self.conn.create_table(
            "users",
            data=[
                {"id": 1, "name": "Alice", "age": 30},
                {"id": 2, "name": "Bob", "age": 25},
            ],
            schema=_schema({
                "id": pa.int32(),
                "name": pa.string(),
                "age": pa.int32(),
            }),
            options=_PARQUET_OPTIONS,
        )

        users = self.conn.get_table("users")
        result = users.scan().select(["id", "name"]).to_arrow()

        self.assertEqual(["id", "name"], result.column_names)
        self.assertEqual([1, 2], result["id"].to_pylist())

    def test_add_scan_where_select_limit(self):
        users = self.conn.create_table(
            "users",
            schema=_schema({
                "id": pa.int32(),
                "name": pa.string(),
                "age": pa.int32(),
            }),
            options=_PARQUET_OPTIONS,
        )

        users.add([
            {"id": 1, "name": "Alice", "age": 30},
            {"id": 2, "name": "Bob", "age": 25},
            {"id": 3, "name": "Carol", "age": 40},
        ])

        result = (
            users.scan()
            .where("age >= 30")
            .select(["id", "name"])
            .limit(1)
            .to_arrow()
        )

        self.assertEqual(["id", "name"], result.column_names)
        self.assertEqual(1, result.num_rows)
        self.assertEqual([1], result["id"].to_pylist())

    def test_scan_with_row_id_returns_system_column(self):
        users = self.conn.create_table(
            "users",
            data=[
                {"id": 1, "name": "Alice"},
                {"id": 2, "name": "Bob"},
            ],
            schema=_schema({
                "id": pa.int32(),
                "name": pa.string(),
            }),
            options=_PARQUET_OPTIONS,
        )

        result = (
            users.scan()
            .with_row_id()
            .select(["id"])
            .to_arrow()
        )

        self.assertEqual(["id", "_ROW_ID"], result.column_names)
        self.assertEqual([1, 2], result["id"].to_pylist())
        self.assertEqual([0, 1], result["_ROW_ID"].to_pylist())

    def test_take_row_ids_reads_projected_rows(self):
        docs = self.conn.create_table(
            "docs",
            data=[
                {"id": 1, "content": "alpha"},
                {"id": 2, "content": "beta"},
                {"id": 3, "content": "gamma"},
            ],
            schema=_schema({
                "id": pa.int32(),
                "content": pa.string(),
            }),
            options=_PARQUET_OPTIONS,
        )
        manifest = {
            row["id"]: row["_ROW_ID"]
            for row in docs.scan().select(["id"]).with_row_id().to_list()
        }

        rows = sorted(
            docs.take_row_ids([manifest[3], manifest[1]])
            .select(["id", "content"])
            .with_row_id()
            .to_list(),
            key=lambda row: row["id"],
        )

        self.assertEqual(
            [
                {"id": 1, "content": "alpha", "_ROW_ID": manifest[1]},
                {"id": 3, "content": "gamma", "_ROW_ID": manifest[3]},
            ],
            rows,
        )

    def test_take_row_ids_accepts_empty_manifest(self):
        docs = self.conn.create_table(
            "docs",
            data=[{"id": 1}],
            schema=_schema({"id": pa.int32()}),
            options=_PARQUET_OPTIONS,
        )

        self.assertEqual([], docs.take_row_ids([]).select(["id"]).to_list())

    def test_take_row_ids_supports_snapshot_and_tag_time_travel(self):
        docs = self.conn.create_table(
            "docs",
            schema=_schema({
                "id": pa.int32(),
                "content": pa.string(),
                "embedding": _vector(3),
            }),
            options=_PARQUET_OPTIONS,
        )
        docs.add([
            {"id": 1, "content": "first", "embedding": [1.0, 0.0, 0.0]},
        ])
        snapshot_id = docs.raw_table.snapshot_manager().get_latest_snapshot().id
        docs.raw_table.create_tag("v1", snapshot_id=snapshot_id)
        row_id = (
            docs.scan(snapshot_id=snapshot_id)
            .select(["id"])
            .with_row_id()
            .to_list()[0]["_ROW_ID"]
        )

        docs.delete(where="id = 1")

        self.assertEqual(
            [],
            docs.take_row_ids([row_id]).select(["id"]).to_list(),
        )
        self.assertEqual(
            [{"id": 1, "_ROW_ID": row_id}],
            docs.take_row_ids([row_id], snapshot_id=snapshot_id)
            .select(["id"])
            .with_row_id()
            .to_list(),
        )
        self.assertEqual(
            [{"id": 1}],
            docs.take_row_ids([row_id], tag_name="v1")
            .select(["id"])
            .to_list(),
        )
        self.assertEqual(
            snapshot_id,
            docs.search(
                [1.0, 0.0, 0.0],
                column="embedding",
                snapshot_id=snapshot_id,
            )._table.options.scan_snapshot_id(),
        )

    def test_time_travel_rejects_snapshot_id_and_tag_name_together(self):
        docs = self.conn.create_table(
            "docs",
            data=[{"id": 1}],
            schema=_schema({"id": pa.int32()}),
            options=_PARQUET_OPTIONS,
        )

        with self.assertRaisesRegex(ValueError, "cannot be set at the same time"):
            docs.take_row_ids([0], snapshot_id=1, tag_name="v1")

    def test_overwrite_replaces_unpartitioned_table(self):
        users = self.conn.create_table(
            "users",
            data=[
                {"id": 1, "name": "Alice", "age": 30},
                {"id": 2, "name": "Bob", "age": 25},
            ],
            schema=_schema({
                "id": pa.int32(),
                "name": pa.string(),
                "age": pa.int32(),
            }),
            options=_PARQUET_OPTIONS,
        )

        result = users.overwrite([
            {"id": 3, "name": "Carol", "age": 40},
        ])

        self.assertIs(users, result)
        self.assertEqual(
            [
                {"id": 3, "name": "Carol", "age": 40},
            ],
            users.scan().to_list(),
        )

    def test_empty_overwrite_clears_unpartitioned_table(self):
        users = self.conn.create_table(
            "users",
            data=[
                {"id": 1, "name": "Alice"},
                {"id": 2, "name": "Bob"},
            ],
            schema=_schema({
                "id": pa.int32(),
                "name": pa.string(),
            }),
            options=_PARQUET_OPTIONS,
        )

        users.overwrite([])

        self.assertEqual([], users.scan().to_list())

    def test_overwrite_replaces_dynamic_partitions(self):
        users = self.conn.create_table(
            "users",
            data=[
                {"id": 1, "name": "Alice", "dt": "2024-01-01"},
                {"id": 2, "name": "Bob", "dt": "2024-01-02"},
            ],
            schema=_schema({
                "id": pa.int32(),
                "name": pa.string(),
                "dt": pa.string(),
            }),
            options=_PARQUET_OPTIONS,
            partitioned=["dt"],
        )

        users.overwrite([
            {"id": 3, "name": "Carol", "dt": "2024-01-01"},
        ])

        rows = sorted(users.scan().to_list(), key=lambda r: r["id"])
        self.assertEqual(
            [
                {"id": 2, "name": "Bob", "dt": "2024-01-02"},
                {"id": 3, "name": "Carol", "dt": "2024-01-01"},
            ],
            rows,
        )

    def test_overwrite_replaces_static_partition(self):
        users = self.conn.create_table(
            "users",
            data=[
                {"id": 1, "name": "Alice", "dt": "2024-01-01"},
                {"id": 2, "name": "Bob", "dt": "2024-01-02"},
            ],
            schema=_schema({
                "id": pa.int32(),
                "name": pa.string(),
                "dt": pa.string(),
            }),
            options=dict(_PARQUET_OPTIONS, **{
                "dynamic-partition-overwrite": "false",
            }),
            partitioned=["dt"],
        )

        users.overwrite([
            {"id": 3, "name": "Carol", "dt": "2024-01-01"},
        ], partition={"dt": "2024-01-01"})

        rows = sorted(users.scan().to_list(), key=lambda r: r["id"])
        self.assertEqual(
            [
                {"id": 2, "name": "Bob", "dt": "2024-01-02"},
                {"id": 3, "name": "Carol", "dt": "2024-01-01"},
            ],
            rows,
        )

    def test_scan_read_blobs(self):
        obs = self.conn.create_table(
            "obs",
            schema=_schema({
                "clip": pa.string(),
                "idx": pa.int32(),
                "image": pa.large_binary(),
            }),
            options=_PARQUET_OPTIONS,
            partitioned=["clip"],
        )
        payloads = [("blob-%d-" % i).encode() * (i + 1) for i in range(6)]
        obs.add([
            {"clip": "c1", "idx": 0, "image": payloads[0]},
            {"clip": "c1", "idx": 1, "image": payloads[1]},
            {"clip": "c1", "idx": 2, "image": payloads[2]},
            {"clip": "c2", "idx": 0, "image": payloads[3]},
            {"clip": "c2", "idx": 1, "image": payloads[4]},
            {"clip": "c3", "idx": 0, "image": None},
        ])

        # scalar table is row-aligned with the blob list and drops the blob column
        scalar, blobs = obs.scan().where("clip = 'c1'").read_blobs("image")
        images = blobs["image"]
        self.assertEqual(3, len(images))
        self.assertNotIn("image", scalar.column_names)
        got = dict(zip(scalar.column("idx").to_pylist(), images))
        self.assertEqual({0: payloads[0], 1: payloads[1], 2: payloads[2]}, got)

        # blob column auto-detected when columns=None
        _, blobs2 = obs.scan().where("clip = 'c2'").read_blobs()
        self.assertEqual({payloads[3], payloads[4]}, set(blobs2["image"]))

        # null blob -> None
        _, blobs3 = obs.scan().where("clip = 'c3'").read_blobs("image")
        self.assertEqual([None], blobs3["image"])

        # non-BLOB column is rejected
        with self.assertRaisesRegex(ValueError, "not a BLOB column"):
            obs.scan().read_blobs("idx")

        # empty result: 0 matching rows -> empty blob list, schema preserved
        scalar_e, blobs_e = obs.scan().where("clip = 'none'").read_blobs("image")
        self.assertEqual(0, scalar_e.num_rows)
        self.assertEqual([], blobs_e["image"])

    def test_scan_read_blobs_multi_column(self):
        obs = self.conn.create_table(
            "multi",
            schema=_schema({
                "clip": pa.string(),
                "idx": pa.int32(),
                "img": pa.large_binary(),
                "aud": pa.large_binary(),
            }),
            options=_PARQUET_OPTIONS,
            partitioned=["clip"],
        )
        obs.add([
            {"clip": "c1", "idx": i,
             "img": ("img-%d" % i).encode(), "aud": ("aud-%d" % i).encode()}
            for i in range(4)
        ])

        # both BLOB columns fetched in one call, each row-aligned with the scalars
        scalar, blobs = obs.scan().where("clip = 'c1'").read_blobs(["img", "aud"])
        self.assertEqual({"img", "aud"}, set(blobs))
        self.assertNotIn("img", scalar.column_names)
        self.assertNotIn("aud", scalar.column_names)
        idx = scalar.column("idx").to_pylist()
        self.assertEqual({i: ("img-%d" % i).encode() for i in range(4)},
                         dict(zip(idx, blobs["img"])))
        self.assertEqual({i: ("aud-%d" % i).encode() for i in range(4)},
                         dict(zip(idx, blobs["aud"])))

        # reading only one BLOB column must not leak the other into scalar as
        # descriptor bytes
        scalar1, blobs1 = obs.scan().where("clip = 'c1'").read_blobs("img")
        self.assertEqual({"img"}, set(blobs1))
        self.assertEqual(["clip", "idx"], scalar1.column_names)

        # columns=None intersects all BLOBs with select() -> only projected BLOB
        _, blobs2 = obs.scan().where("clip = 'c1'").select(["clip", "idx", "img"]).read_blobs()
        self.assertEqual({"img"}, set(blobs2))

        # duplicate columns are de-duplicated
        _, blobs3 = obs.scan().where("clip = 'c1'").read_blobs(["img", "img"])
        self.assertEqual({"img"}, set(blobs3))

    def test_scan_read_blobs_filter_column_not_selected(self):
        # The row filter must apply even when its column is not in select().
        obs = self.conn.create_table(
            "obs",
            schema=_schema({
                "clip": pa.string(),
                "idx": pa.int32(),
                "image": pa.large_binary(),
            }),
            options=_PARQUET_OPTIONS,
            partitioned=["clip"],
        )
        obs.add([
            {"clip": "c1", "idx": i, "image": ("v-%d" % i).encode()}
            for i in range(4)
        ])

        scalar, blobs = (
            obs.scan().select(["image"]).where("idx = 1").read_blobs("image"))
        self.assertEqual([b"v-1"], blobs["image"])
        self.assertNotIn("idx", scalar.schema.names)

        scalar2, blobs2 = (
            obs.scan().select(["idx", "image"]).where("idx = 2").read_blobs("image"))
        self.assertEqual([b"v-2"], blobs2["image"])
        self.assertEqual([2], scalar2.column("idx").to_pylist())

        streamed = []
        for _, body in (
                obs.scan().select(["image"]).where("idx = 1").stream_blobs("image")):
            streamed.extend(body["image"])
        self.assertEqual([b"v-1"], streamed)

        scalar3, _ = obs.scan().select(["missing", "image"]).read_blobs("image")
        self.assertNotIn("missing", scalar3.schema.names)

    def test_scan_read_blobs_with_row_id(self):
        # with_row_id() must expose _ROW_ID in the scalar table, like to_arrow().
        obs = self.conn.create_table(
            "obs",
            schema=_schema({
                "clip": pa.string(),
                "idx": pa.int32(),
                "image": pa.large_binary(),
            }),
            options=_PARQUET_OPTIONS,
            partitioned=["clip"],
        )
        obs.add([
            {"clip": "c1", "idx": i, "image": ("v-%d" % i).encode()}
            for i in range(3)
        ])
        row_id = "_ROW_ID"

        expected = {r["idx"]: r[row_id]
                    for r in obs.scan().with_row_id().to_arrow().to_pylist()}

        scalar, blobs = (
            obs.scan().with_row_id().select(["idx", "image"]).read_blobs("image"))
        self.assertIn(row_id, scalar.schema.names)
        self.assertEqual(expected, dict(zip(scalar.column("idx").to_pylist(),
                                            scalar.column(row_id).to_pylist())))

        got = {}
        for sb, _ in obs.scan().with_row_id().select(["idx", "image"]).stream_blobs("image"):
            self.assertIn(row_id, sb.schema.names)
            got.update(zip(sb.column("idx").to_pylist(), sb.column(row_id).to_pylist()))
        self.assertEqual(expected, got)

    def test_scan_read_blobs_where_on_blob_column(self):
        # A where() on one BLOB column must still filter when a different BLOB is
        # read (the predicate BLOB column is read as descriptor for filtering).
        obs = self.conn.create_table(
            "obs",
            schema=_schema({
                "clip": pa.string(),
                "idx": pa.int32(),
                "imga": pa.large_binary(),
                "imgb": pa.large_binary(),
            }),
            options=_PARQUET_OPTIONS,
            partitioned=["clip"],
        )
        obs.add([
            {"clip": "c1", "idx": 0, "imga": b"a0", "imgb": b"b0"},
            {"clip": "c1", "idx": 1, "imga": None, "imgb": b"b1"},
            {"clip": "c1", "idx": 2, "imga": b"a2", "imgb": b"b2"},
        ])

        scalar, blobs = obs.scan().where("imga IS NOT NULL").read_blobs("imgb")
        self.assertEqual([b"b0", b"b2"], blobs["imgb"])
        self.assertNotIn("imga", scalar.schema.names)

    def test_search_query_rejects_blob_reads(self):
        t = self.conn.create_table(
            "srch",
            schema=_schema({
                "id": pa.int32(),
                "emb": _vector(3),
                "img": pa.large_binary(),
            }),
            options=_PARQUET_OPTIONS,
        )
        with self.assertRaisesRegex(TypeError, "only supported on scan"):
            t.search([1.0, 0.0, 0.0], column="emb").read_blobs("img")
        with self.assertRaisesRegex(TypeError, "only supported on scan"):
            t.search([1.0, 0.0, 0.0], column="emb").stream_blobs("img")
        with self.assertRaisesRegex(TypeError, "only supported on scan"):
            t.search([1.0, 0.0, 0.0], column="emb").to_ray()

    def test_scan_stream_blobs(self):
        obs = self.conn.create_table(
            "obs",
            schema=_schema({
                "clip": pa.string(),
                "idx": pa.int32(),
                "image": pa.large_binary(),
            }),
            options=_PARQUET_OPTIONS,
            partitioned=["clip"],
        )
        payloads = {i: ("s-%d-" % i).encode() * (i + 1) for i in range(5)}
        obs.add([
            {"clip": "c1", "idx": i, "image": payloads[i]} for i in range(5)
        ])

        # collecting every streamed batch must reproduce the full, row-aligned result
        got = {}
        batches = 0
        for scalar, blobs in obs.scan().where("clip = 'c1'").stream_blobs("image"):
            batches += 1
            self.assertNotIn("image", scalar.schema.names)
            for idx, img in zip(scalar.column("idx").to_pylist(), blobs["image"]):
                got[idx] = img
        self.assertGreaterEqual(batches, 1)
        self.assertEqual({i: payloads[i] for i in range(5)}, got)

        # bad column is rejected eagerly (not deferred to first iteration)
        with self.assertRaisesRegex(ValueError, "not a BLOB column"):
            obs.scan().stream_blobs("idx")
        # breaking out early closes the reader without error
        it = obs.scan().where("clip = 'c1'").stream_blobs("image")
        next(it)
        it.close()
        # empty result streams nothing
        self.assertEqual(
            [], list(obs.scan().where("clip = 'none'").stream_blobs("image")))

    @unittest.skipIf(ray is None, "ray is not installed")
    def test_scan_to_ray_map_with_blobs(self):
        started_ray = False
        if not ray.is_initialized():
            ray.init(ignore_reinit_error=True, num_cpus=2)
            started_ray = True
        obs = self.conn.create_table(
            "ray_obs",
            schema=_schema({
                "clip": pa.string(),
                "idx": pa.int32(),
                "image": pa.large_binary(),
                "audio": pa.large_binary(),
            }),
            options=_PARQUET_OPTIONS,
            partitioned=["clip"],
        )
        obs.add([
            {"clip": "c1", "idx": 0, "image": b"img-0", "audio": b"aud-0"},
            {"clip": "c1", "idx": 1, "image": b"img-1", "audio": b"aud-1"},
            {"clip": "c2", "idx": 0, "image": b"img-x", "audio": b"aud-x"},
        ])

        def collect_batch(scalar, blobs, prefix):
            assert isinstance(scalar, pa.Table)
            assert ["idx"] == scalar.column_names
            idxs = scalar.column("idx").to_pylist()
            rows = []
            for idx, image in zip(idxs, blobs["image"]):
                rows.append({"idx": idx, "image": prefix + image})
            return pa.Table.from_pylist(rows)

        try:
            from pypaimon.ray import map_with_blobs

            ds = (
                obs.scan()
                .where("clip = 'c1'")
                .select(["idx", "image", "audio"])
                .to_ray(concurrency=1, override_num_blocks=1)
            )
            ds = ds.filter(lambda row: row["idx"] >= 0)

            with self.assertRaisesRegex(ValueError, "not a BLOB column"):
                obs.map_with_blobs(ds, ["idx"], collect_batch)
            with self.assertRaisesRegex(ValueError, "FileIO"):
                map_with_blobs(ds, ["image"], collect_batch)

            from pypaimon.table.row.blob import BlobDescriptor

            descriptor = BlobDescriptor("oss://bucket/blob", 0, 1).serialize()
            foreign_ds = ray.data.from_arrow(pa.table({
                "idx": [0],
                "image": [descriptor],
                "foreign_blob": [descriptor],
            }))
            with self.assertRaisesRegex(Exception, "does not own"):
                map_with_blobs(
                    foreign_ds,
                    ["image"],
                    collect_batch,
                    file_io=obs.raw_table.file_io,
                    all_blob_columns=["image"],
                    batch_size=1,
                ).take_all()

            result = obs.map_with_blobs(
                ds,
                ["image"],
                collect_batch,
                parallelism=2,
                batch_size=1,
                fn_kwargs={"prefix": b"got-"},
                ray_remote_args={"num_cpus": 1},
            )
            rows = sorted(result.to_pandas().to_dict("records"), key=lambda row: row["idx"])

            self.assertEqual(
                [
                    {"idx": 0, "image": b"got-img-0"},
                    {"idx": 1, "image": b"got-img-1"},
                ],
                rows,
            )
        finally:
            if started_ray:
                ray.shutdown()

    @unittest.skipIf(ray is None, "ray is not installed")
    def test_scan_to_ray_map_with_blobs_guards(self):
        started_ray = False
        if not ray.is_initialized():
            ray.init(ignore_reinit_error=True, num_cpus=2)
            started_ray = True
        obs = self.conn.create_table(
            "ray_obs_guards",
            schema=_schema({
                "clip": pa.string(),
                "idx": pa.int32(),
                "image": pa.large_binary(),
            }),
            options=_PARQUET_OPTIONS,
            partitioned=["clip"],
        )
        obs.add([{"clip": "c1", "idx": 0, "image": b"img-0"}])

        def return_none(scalar, blobs):
            return None

        try:
            from pypaimon.ray import map_with_blobs

            ds = (
                obs.scan()
                .where("clip = 'c1'")
                .select(["idx", "image"])
                .to_ray(concurrency=1, override_num_blocks=1)
            )

            with self.assertRaisesRegex(ValueError, "all_blob_columns"):
                map_with_blobs(
                    ray.data.from_arrow(pa.table({"image": [b"inline"]})),
                    ["image"],
                    lambda scalar, blobs: pa.table({"rows": [scalar.num_rows]}),
                    file_io=obs.raw_table.file_io,
                )

            with self.assertRaisesRegex(Exception, "must return"):
                obs.map_with_blobs(
                    ds,
                    ["image"],
                    return_none,
                    batch_size=1,
                ).take_all()

            empty_ds = (
                obs.scan()
                .where("clip = 'none'")
                .select(["idx", "image"])
                .to_ray(concurrency=1, override_num_blocks=1)
            )
            result = obs.map_with_blobs(
                empty_ds,
                ["image"],
                lambda scalar, blobs: pa.table({"rows": [scalar.num_rows]}),
                batch_size=1,
            )
            self.assertEqual([], result.take_all())
        finally:
            if started_ray:
                ray.shutdown()

    def test_scan_to_ray_nested_projection_output_names(self):
        obs = self.conn.create_table(
            "ray_nested_obs",
            schema=_schema({
                "id": pa.int32(),
                "tag": pa.string(),
                "payload": pa.struct([("a", pa.int64()), ("b", pa.string())]),
                "image": pa.large_binary(),
            }),
            options=_PARQUET_OPTIONS,
        )

        _, _, visible_columns = (
            obs.scan()
            .where("tag = 'keep'")
            .select(["id", "payload.a"])
            ._blob_descriptor_query_read_builder()
        )

        self.assertEqual(["id", "payload_a"], visible_columns)

    def test_fetch_bodies_decodes_descriptor_inline_and_null(self):
        # Cells may be descriptor bytes (incl. -1 read-to-EOF), inline bytes, or null.
        from pypaimon.multimodal.query import ScanQuery
        from pypaimon.common.file_io import FileIO
        from pypaimon.table.row.blob import BlobDescriptor
        data = bytes(range(64))
        path = os.path.join(self.temp_dir, "blob.bin")
        with open(path, "wb") as f:
            f.write(data)
        file_io = FileIO.get("file://" + self.temp_dir, {})
        cells = [
            BlobDescriptor(path, 0, -1).serialize(),
            BlobDescriptor(path, 10, 8).serialize(),
            b"raw-inline-bytes",
            None,
        ]
        bodies = ScanQuery._fetch_bodies(
            file_io, {"image": cells}, ["image"], 4)
        self.assertEqual(
            [data, data[10:18], b"raw-inline-bytes", None], bodies["image"])

    def test_fetch_bodies_reads_via_file_io_not_uri_reader(self):
        # Blobs must be read through file_io.read_ranges_coalesced (which carries the
        # resolved DLF/OSS token), never via uri_reader_factory -- that would rebuild a
        # FileIO from the raw catalog options and fail in DLF data-token mode.
        from pypaimon.multimodal.query import ScanQuery
        from pypaimon.table.row.blob import BlobDescriptor

        class _FakeIO:
            @property
            def uri_reader_factory(self):
                raise AssertionError("_fetch_bodies must not use uri_reader_factory")

            def read_ranges_coalesced(self, ranges, parallelism):
                return [None if r is None else b"BODY:" + r[0].encode() for r in ranges]

        cells = [BlobDescriptor("oss://bucket/x", 4, 10).serialize(), b"inline-blob", None]
        bodies = ScanQuery._fetch_bodies(_FakeIO(), {"img": cells}, ["img"], 8)
        self.assertEqual([b"BODY:oss://bucket/x", b"inline-blob", None], bodies["img"])

    def test_fetch_bodies_rejects_unresolved_blob_view(self):
        from pypaimon.multimodal.query import ScanQuery
        from pypaimon.table.row.blob import BlobViewStruct
        view_bytes = BlobViewStruct("db.tbl", 1, 2).serialize()
        with self.assertRaisesRegex(ValueError, "blob-view"):
            ScanQuery._fetch_bodies(None, {"image": [view_bytes]}, ["image"], 4)

    def test_scan_does_not_expose_pre_filter(self):
        users = self.conn.create_table(
            "users",
            schema=_schema({"id": pa.int32()}),
            options=_PARQUET_OPTIONS,
        )

        self.assertFalse(hasattr(users.scan(), "pre_filter"))

    def test_where_rejects_predicate_object(self):
        users = self.conn.create_table(
            "users",
            schema=_schema({"id": pa.int32()}),
            options=_PARQUET_OPTIONS,
        )
        predicate = PredicateBuilder(users.raw_table.fields).equal("id", 1)

        with self.assertRaisesRegex(ValueError, "SQL-like string"):
            users.scan().where(predicate)

    def test_update_by_filter(self):
        users = self.conn.create_table(
            "users",
            data=[
                {"id": 1, "name": "Alice", "age": 30},
                {"id": 2, "name": "Bob", "age": 25},
            ],
            schema=_schema({
                "id": pa.int32(),
                "name": pa.string(),
                "age": pa.int32(),
            }),
            options=_PARQUET_OPTIONS,
        )

        users.update(where="id = 2", values={"age": 26})

        rows = sorted(users.scan().to_list(), key=lambda r: r["id"])
        self.assertEqual(
            [
                {"id": 1, "name": "Alice", "age": 30},
                {"id": 2, "name": "Bob", "age": 26},
            ],
            rows,
        )

    def test_delete_by_filter(self):
        users = self.conn.create_table(
            "users",
            data=[
                {"id": 1, "name": "Alice", "age": 30},
                {"id": 2, "name": "Bob", "age": 25},
                {"id": 3, "name": "Carol", "age": 40},
            ],
            schema=_schema({
                "id": pa.int32(),
                "name": pa.string(),
                "age": pa.int32(),
            }),
            options=_PARQUET_OPTIONS,
        )

        users.delete(where="id = 2")

        rows = sorted(users.scan().to_list(), key=lambda r: r["id"])
        self.assertEqual(
            [
                {"id": 1, "name": "Alice", "age": 30},
                {"id": 3, "name": "Carol", "age": 40},
            ],
            rows,
        )

    def test_merge_updates_matches_and_inserts_new_rows(self):
        users = self.conn.create_table(
            "users",
            data=[
                {"id": 1, "name": "Alice", "age": 30},
                {"id": 2, "name": "Bob", "age": 25},
            ],
            schema=_schema({
                "id": pa.int32(),
                "name": pa.string(),
                "age": pa.int32(),
            }),
            options=_PARQUET_OPTIONS,
        )

        users.merge("id") \
            .when_matched_update() \
            .when_not_matched_insert() \
            .execute([
                {"id": 2, "name": "Bob_v2", "age": 26},
                {"id": 3, "name": "Carol", "age": 40},
            ])

        rows = sorted(users.scan().to_list(), key=lambda r: r["id"])
        self.assertEqual(
            [
                {"id": 1, "name": "Alice", "age": 30},
                {"id": 2, "name": "Bob_v2", "age": 26},
                {"id": 3, "name": "Carol", "age": 40},
            ],
            rows,
        )

    def test_merge_deletes_matched_rows(self):
        users = self.conn.create_table(
            "users",
            data=[
                {"id": 1, "name": "Alice", "age": 30},
                {"id": 2, "name": "Bob", "age": 25},
                {"id": 3, "name": "Carol", "age": 40},
            ],
            schema=_schema({
                "id": pa.int32(),
                "name": pa.string(),
                "age": pa.int32(),
            }),
            options=_PARQUET_OPTIONS,
        )

        users.merge("id") \
            .when_matched_delete() \
            .execute([
                {"id": 2},
                {"id": 3},
            ])

        rows = sorted(users.scan().to_list(), key=lambda r: r["id"])
        self.assertEqual(
            [
                {"id": 1, "name": "Alice", "age": 30},
            ],
            rows,
        )

    def test_merge_all_uses_only_source_columns(self):
        users = self.conn.create_table(
            "users",
            data=[
                {"id": 1, "name": "Alice", "age": 30},
                {"id": 2, "name": "Bob", "age": 25},
            ],
            schema=_schema({
                "id": pa.int32(),
                "name": pa.string(),
                "age": pa.int32(),
            }),
            options=_PARQUET_OPTIONS,
        )

        users.merge("id") \
            .when_matched_update() \
            .when_not_matched_insert() \
            .execute([
                {"id": 2, "age": 26},
                {"id": 3, "age": 40},
            ])

        rows = sorted(users.scan().to_list(), key=lambda r: r["id"])
        self.assertEqual(
            [
                {"id": 1, "name": "Alice", "age": 30},
                {"id": 2, "name": "Bob", "age": 26},
                {"id": 3, "name": None, "age": 40},
            ],
            rows,
        )

    def test_merge_supports_source_key_mapping(self):
        users = self.conn.create_table(
            "users",
            data=[
                {"id": 1, "name": "Alice", "age": 30},
                {"id": 2, "name": "Bob", "age": 25},
            ],
            schema=_schema({
                "id": pa.int32(),
                "name": pa.string(),
                "age": pa.int32(),
            }),
            options=_PARQUET_OPTIONS,
        )

        users.merge({"id": "source_id"}) \
            .when_matched_update({"age": source_col("age")}) \
            .when_not_matched_insert() \
            .execute([
                {"source_id": 2, "name": "Bob_v2", "age": 26},
                {"source_id": 3, "name": "Carol", "age": 40},
            ])

        rows = sorted(users.scan().to_list(), key=lambda r: r["id"])
        self.assertEqual(
            [
                {"id": 1, "name": "Alice", "age": 30},
                {"id": 2, "name": "Bob", "age": 26},
                {"id": 3, "name": "Carol", "age": 40},
            ],
            rows,
        )

    def test_merge_where_uses_source_and_target_aliases(self):
        users = self.conn.create_table(
            "users",
            schema=_schema({
                "id": pa.int32(),
                "name": pa.string(),
                "age": pa.int32(),
            }),
            options=_PARQUET_OPTIONS,
        )

        calls = {}

        class FakeUpdate:
            def merge_into(
                    self,
                    source,
                    on,
                    when_matched=None,
                    when_not_matched=None):
                calls["on"] = on
                calls["matched"] = list(when_matched or [])
                calls["not_matched"] = list(when_not_matched or [])
                return []

        class FakeCommit:
            def commit(self, messages):
                calls["messages"] = messages

            def close(self):
                pass

        class FakeWriteBuilder:
            def new_update(self):
                return FakeUpdate()

            def new_commit(self):
                return FakeCommit()

        users.raw_table.new_batch_write_builder = lambda: FakeWriteBuilder()

        (
            users.merge("id")
            .when_matched_update(
                {"age": source_col("age")},
                where="source.age > target.age and source.name != 'target.name'",
            )
            .when_not_matched_insert(where="source.age > 0")
            .execute([
                {"id": 1, "name": "Alice", "age": 31},
            ])
        )

        self.assertEqual({"id": "id"}, calls["on"])
        self.assertEqual(
            "s.age > t.age and s.name != 'target.name'",
            calls["matched"][0].condition,
        )
        self.assertEqual("s.age > 0", calls["not_matched"][0].condition)
        self.assertEqual([], calls["messages"])

    def test_merge_delete_where_uses_source_and_target_aliases(self):
        users = self.conn.create_table(
            "users",
            schema=_schema({
                "id": pa.int32(),
                "name": pa.string(),
                "age": pa.int32(),
            }),
            options=_PARQUET_OPTIONS,
        )

        calls = {}

        class FakeUpdate:
            def merge_into(
                    self,
                    source,
                    on,
                    when_matched=None,
                    when_not_matched=None):
                calls["matched"] = list(when_matched or [])
                return []

        class FakeCommit:
            def commit(self, messages):
                calls["messages"] = messages

            def close(self):
                pass

        class FakeWriteBuilder:
            def new_update(self):
                return FakeUpdate()

            def new_commit(self):
                return FakeCommit()

        users.raw_table.new_batch_write_builder = lambda: FakeWriteBuilder()

        (
            users.merge("id")
            .when_matched_delete(
                where="source.age < target.age and source.name != 'target.name'",
            )
            .execute([
                {"id": 1, "name": "Alice", "age": 29},
            ])
        )

        self.assertTrue(calls["matched"][0].delete)
        self.assertEqual(
            "s.age < t.age and s.name != 'target.name'",
            calls["matched"][0].condition,
        )
        self.assertEqual([], calls["messages"])

    def test_merge_validates_on_columns(self):
        users = self.conn.create_table(
            "users",
            schema=_schema({"id": pa.int32(), "name": pa.string()}),
            options=_PARQUET_OPTIONS,
        )

        with self.assertRaisesRegex(ValueError, "source columns"):
            users.merge({"id": "missing"}) \
                .when_matched_update() \
                .execute([{"id": 1, "name": "Alice"}])

    def test_create_index_normalizes_full_text_alias(self):
        docs = self.conn.create_table(
            "docs",
            schema=_schema({
                "id": pa.int32(),
                "content": pa.string(),
            }),
            options=_PARQUET_OPTIONS,
        )

        calls = []

        def create_global_index(column, index_type, options=None):
            calls.append((column, index_type, options))
            return index_type

        docs.raw_table.create_global_index = create_global_index

        options = {"tokenizer": "default"}
        self.assertEqual(
            "full-text",
            docs.create_index("content", index_type="full-text",
                              options=options),
        )

        self.assertEqual(
            [
                ("content", "full-text", options),
            ],
            calls,
        )

    def test_create_index_requires_index_type(self):
        docs = self.conn.create_table(
            "docs",
            schema=_schema({
                "id": pa.int32(),
                "content": pa.string(),
            }),
            options=_PARQUET_OPTIONS,
        )

        def create_global_index(column, index_type, options=None):
            return index_type

        docs.raw_table.create_global_index = create_global_index

        with self.assertRaises(TypeError):
            docs.create_index("content")

    def test_table_does_not_expose_predicate_builder(self):
        docs = self.conn.create_table(
            "docs",
            schema=_schema({"id": pa.int32()}),
            options=_PARQUET_OPTIONS,
        )

        self.assertFalse(hasattr(docs, "predicate_builder"))

    def test_search_reads_vector_matching_rows(self):
        docs = self.conn.create_table(
            "docs",
            schema=_schema({
                "id": pa.int32(),
                "content": pa.string(),
                "embedding": _vector(3),
            }),
            options=_PARQUET_OPTIONS,
        )
        docs.add([
            {"id": 1, "content": "a", "embedding": [1.0, 0.0, 0.0]},
            {"id": 2, "content": "b", "embedding": [0.0, 1.0, 0.0]},
            {"id": 3, "content": "c", "embedding": [0.0, 0.0, 1.0]},
        ])

        calls = {}

        class FakeVectorBuilder:
            def with_vector_column(self, column):
                calls["column"] = column
                return self

            def with_query_vector(self, vector):
                calls["vector"] = vector
                return self

            def with_limit(self, limit):
                calls["limit"] = limit
                return self

            def with_options(self, options):
                calls["options"] = options
                return self

            def with_filter(self, predicate):
                calls["filter"] = predicate
                return self

            def execute_local(self):
                return GlobalIndexResult.from_range(Range(1, 1))

        docs.raw_table.new_vector_search_builder = lambda: FakeVectorBuilder()

        result = (
            docs.search([0.0, 1.0, 0.0])
            .where("id >= 1")
            .limit(5)
            .to_arrow()
        )

        self.assertEqual("embedding", calls["column"])
        self.assertEqual([0.0, 1.0, 0.0], calls["vector"])
        self.assertEqual(5, calls["limit"])
        self.assertEqual([2], result["id"].to_pylist())

    def test_search_applies_pre_filter_to_vector_builder(self):
        docs = self.conn.create_table(
            "docs",
            schema=_schema({
                "id": pa.int32(),
                "category": pa.string(),
                "embedding": _vector(3),
            }),
            options=_PARQUET_OPTIONS,
        )
        docs.add([
            {"id": 1, "category": "lake", "embedding": [1.0, 0.0, 0.0]},
            {"id": 2, "category": "city", "embedding": [0.0, 1.0, 0.0]},
        ])

        calls = {}

        class FakeVectorBuilder:
            def with_vector_column(self, column):
                return self

            def with_query_vector(self, vector):
                return self

            def with_limit(self, limit):
                return self

            def with_options(self, options):
                return self

            def with_filter(self, predicate):
                calls["pre_filter"] = predicate
                return self

            def execute_local(self):
                return GlobalIndexResult.from_range(Range(0, 0))

        docs.raw_table.new_vector_search_builder = lambda: FakeVectorBuilder()

        docs.search(
            [1.0, 0.0, 0.0],
            pre_filter="category = 'lake'",
        ).limit(1).to_list()

        self.assertEqual("equal", calls["pre_filter"].method)
        self.assertEqual("category", calls["pre_filter"].field)
        self.assertEqual(["lake"], calls["pre_filter"].literals)

    def test_search_pre_filter_rejects_predicate_object(self):
        docs = self.conn.create_table(
            "docs",
            schema=_schema({
                "id": pa.int32(),
                "embedding": _vector(3),
            }),
            options=_PARQUET_OPTIONS,
        )
        predicate = PredicateBuilder(docs.raw_table.fields).equal("id", 1)

        with self.assertRaisesRegex(ValueError, "SQL-like string"):
            docs.search([1.0, 0.0, 0.0]).pre_filter(predicate)

    def test_search_accepts_generator_vector(self):
        docs = self.conn.create_table(
            "docs",
            schema=_schema({
                "id": pa.int32(),
                "embedding": _vector(3),
            }),
            options=_PARQUET_OPTIONS,
        )
        docs.add([{"id": 1, "embedding": [1.0, 0.0, 0.0]}])

        calls = {}

        class FakeVectorBuilder:
            def with_vector_column(self, column):
                calls["column"] = column
                return self

            def with_query_vector(self, vector):
                calls["vector"] = vector
                return self

            def with_limit(self, limit):
                calls["limit"] = limit
                return self

            def with_options(self, options):
                calls["options"] = options
                return self

            def execute_local(self):
                return GlobalIndexResult.from_range(Range(0, 0))

        docs.raw_table.new_vector_search_builder = lambda: FakeVectorBuilder()

        result = docs.search((v for v in [1.0, 0.0, 0.0])).limit(1).to_list()

        self.assertEqual("embedding", calls["column"])
        self.assertEqual([1.0, 0.0, 0.0], calls["vector"])
        self.assertEqual([{"id": 1, "embedding": [1.0, 0.0, 0.0]}], result)

    def test_search_with_row_id_returns_system_column(self):
        docs = self.conn.create_table(
            "docs",
            schema=_schema({
                "id": pa.int32(),
                "embedding": _vector(3),
            }),
            options=_PARQUET_OPTIONS,
        )
        docs.add([{"id": 1, "embedding": [1.0, 0.0, 0.0]}])

        class FakeVectorBuilder:
            def with_vector_column(self, column):
                return self

            def with_query_vector(self, vector):
                return self

            def with_limit(self, limit):
                return self

            def with_options(self, options):
                return self

            def execute_local(self):
                return GlobalIndexResult.from_range(Range(0, 0))

        docs.raw_table.new_vector_search_builder = lambda: FakeVectorBuilder()

        result = (
            docs.search([1.0, 0.0, 0.0], column="embedding")
            .select(["id"])
            .with_row_id()
            .limit(1)
            .to_list()
        )

        self.assertEqual([{"id": 1, "_ROW_ID": 0}], result)

    def test_search_rejects_batch_vectors(self):
        docs = self.conn.create_table(
            "docs",
            schema=_schema({
                "id": pa.int32(),
                "embedding": _vector(3),
            }),
            options=_PARQUET_OPTIONS,
        )

        with self.assertRaisesRegex(ValueError, "use search_vectors"):
            docs.search([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])

    def test_search_vectors_reads_one_result_set_per_query_vector(self):
        docs = self.conn.create_table(
            "docs",
            schema=_schema({
                "id": pa.int32(),
                "content": pa.string(),
                "embedding": _vector(3),
            }),
            options=_PARQUET_OPTIONS,
        )
        docs.add([
            {"id": 1, "content": "a", "embedding": [1.0, 0.0, 0.0]},
            {"id": 2, "content": "b", "embedding": [0.0, 1.0, 0.0]},
            {"id": 3, "content": "c", "embedding": [0.0, 0.0, 1.0]},
        ])

        calls = {}

        class FakeBatchVectorBuilder:
            def with_vector_column(self, column):
                calls["column"] = column
                return self

            def with_query_vectors(self, vectors):
                calls["vectors"] = vectors
                return self

            def with_limit(self, limit):
                calls["limit"] = limit
                return self

            def with_options(self, options):
                calls["options"] = options
                return self

            def with_filter(self, predicate):
                calls["filter"] = predicate
                return self

            def execute_batch_local(self):
                return [
                    GlobalIndexResult.from_range(Range(0, 0)),
                    GlobalIndexResult.from_range(Range(2, 2)),
                ]

        docs.raw_table.new_batch_vector_search_builder = (
            lambda: FakeBatchVectorBuilder())

        result = (
            docs.search_vectors(
                [[1.0, 0.0, 0.0], [0.0, 0.0, 1.0]],
                options={"nprobe": "8"},
                pre_filter="content = 'a'",
            )
            .where("id >= 1")
            .select(["id"])
            .limit(2)
            .to_list()
        )

        self.assertEqual("embedding", calls["column"])
        self.assertEqual(
            [[1.0, 0.0, 0.0], [0.0, 0.0, 1.0]],
            calls["vectors"],
        )
        self.assertEqual(2, calls["limit"])
        self.assertEqual({"nprobe": "8"}, calls["options"])
        self.assertEqual("content", calls["filter"].field)
        self.assertEqual(["a"], calls["filter"].literals)
        self.assertEqual([[{"id": 1}], [{"id": 3}]], result)

    def test_search_vectors_rejects_text_parameter(self):
        docs = self.conn.create_table(
            "docs",
            schema=_schema({
                "id": pa.int32(),
                "content": pa.string(),
                "embedding": _vector(3),
            }),
            options=_PARQUET_OPTIONS,
        )

        with self.assertRaises(TypeError):
            docs.search_vectors([[1.0, 0.0, 0.0]], text="paimon")

    def test_search_reads_text_query(self):
        docs = self.conn.create_table(
            "docs",
            schema=_schema({
                "id": pa.int32(),
                "content": pa.string(),
            }),
            options=_PARQUET_OPTIONS,
        )
        docs.add([
            {"id": 1, "content": "paimon vector"},
            {"id": 2, "content": "lakehouse"},
        ])

        calls = {}

        class FakeFullTextBuilder:
            def with_query(self, field_name, query):
                calls["query"] = (field_name, json.loads(query))
                return self

            def with_limit(self, limit):
                calls["limit"] = limit
                return self

            def execute_local(self):
                return GlobalIndexResult.from_range(Range(0, 0))

        docs.raw_table.new_full_text_search_builder = lambda: FakeFullTextBuilder()

        result = (
            docs.search("paimon vector")
            .limit(1)
            .to_arrow()
        )

        self.assertEqual(1, calls["limit"])
        self.assertEqual("content", calls["query"][0])
        self.assertEqual("paimon vector", calls["query"][1]["match"]["query"])
        self.assertEqual([1], result["id"].to_pylist())

    def test_search_string_requires_unambiguous_text_column(self):
        docs = self.conn.create_table(
            "docs",
            schema=_schema({
                "id": pa.int32(),
                "title": pa.string(),
                "content": pa.string(),
            }),
            options=_PARQUET_OPTIONS,
        )

        with self.assertRaisesRegex(ValueError, "Multiple text columns"):
            docs.search("paimon")

    def test_search_hybrid_rejects_shorthand_arguments(self):
        docs = self.conn.create_table(
            "docs",
            schema=_schema({
                "id": pa.int32(),
                "embedding": _vector(3),
            }),
            options=_PARQUET_OPTIONS,
        )

        with self.assertRaises(TypeError):
            docs.search_hybrid(vector=[1.0, 0.0, 0.0])

    def test_hybrid_route_helpers_use_route_specific_arguments(self):
        with self.assertRaises(TypeError):
            pmm.vector_route([1.0, 0.0, 0.0])
        route = pmm.text_route("paimon")
        self.assertIsNone(route.column)
        self.assertEqual("paimon", route.query)

    def test_search_can_build_hybrid_routes(self):
        docs = self.conn.create_table(
            "docs",
            schema=_schema({
                "id": pa.int32(),
                "content": pa.string(),
                "embedding": _vector(3),
            }),
            options=_PARQUET_OPTIONS,
        )
        docs.add([
            {"id": 1, "content": "paimon", "embedding": [1.0, 0.0, 0.0]},
            {"id": 2, "content": "vector", "embedding": [0.0, 1.0, 0.0]},
        ])

        calls = {}

        class FakeHybridBuilder:
            def with_limit(self, limit):
                calls["limit"] = limit
                return self

            def with_ranker(self, ranker):
                calls["ranker"] = ranker
                return self

            def add_vector_route(
                    self, column, vector, limit, weight=1.0, options=None):
                calls["vector"] = (column, vector, limit, weight, options)
                return self

            def add_full_text_route(
                    self, field_name, query, limit, weight=1.0, options=None):
                calls["text"] = (
                    field_name, json.loads(query), limit, weight, options)
                return self

            def with_filter(self, predicate):
                calls["filter"] = predicate
                return self

            def execute_local(self):
                return GlobalIndexResult.from_range(Range(0, 1))

        docs.raw_table.new_hybrid_search_builder = lambda: FakeHybridBuilder()

        result = (
            docs.search_hybrid(
                [
                    pmm.vector_route("embedding", [1.0, 0.0, 0.0]),
                    pmm.text_route("paimon"),
                ],
            )
            .rerank("rrf")
            .limit(2)
            .to_arrow()
        )

        self.assertEqual(2, calls["limit"])
        self.assertEqual("rrf", calls["ranker"])
        self.assertEqual(
            ("embedding", [1.0, 0.0, 0.0], 2, 1.0, {}),
            calls["vector"],
        )
        self.assertEqual("content", calls["text"][0])
        self.assertEqual("paimon", calls["text"][1]["match"]["query"])
        self.assertEqual([1, 2], result["id"].to_pylist())

    def test_search_hybrid_applies_pre_filter_to_vector_routes(self):
        docs = self.conn.create_table(
            "docs",
            schema=_schema({
                "id": pa.int32(),
                "category": pa.string(),
                "embedding": _vector(3),
            }),
            options=_PARQUET_OPTIONS,
        )
        docs.add([
            {"id": 1, "category": "lake", "embedding": [1.0, 0.0, 0.0]},
            {"id": 2, "category": "city", "embedding": [0.0, 1.0, 0.0]},
        ])

        calls = {}

        class FakeHybridBuilder:
            def with_limit(self, limit):
                calls["limit"] = limit
                return self

            def with_ranker(self, ranker):
                calls["ranker"] = ranker
                return self

            def add_vector_route(
                    self, column, vector, limit, weight=1.0, options=None):
                calls["vector"] = (column, vector, limit, weight, options)
                return self

            def with_filter(self, predicate):
                calls["pre_filter"] = predicate
                return self

            def execute_local(self):
                return GlobalIndexResult.from_range(Range(0, 0))

        docs.raw_table.new_hybrid_search_builder = lambda: FakeHybridBuilder()

        result = (
            docs.search_hybrid(
                [pmm.vector_route("embedding", [1.0, 0.0, 0.0])],
                pre_filter="category = 'lake'",
            )
            .limit(1)
            .to_list()
        )

        self.assertEqual(1, calls["limit"])
        self.assertEqual(
            ("embedding", [1.0, 0.0, 0.0], 1, 1.0, {}),
            calls["vector"],
        )
        self.assertEqual("category", calls["pre_filter"].field)
        self.assertEqual(["lake"], calls["pre_filter"].literals)
        self.assertEqual(
            [{"id": 1, "category": "lake", "embedding": [1.0, 0.0, 0.0]}],
            result,
        )

    def test_module_does_not_export_vectors_route(self):
        self.assertFalse(hasattr(pmm, "vectors_route"))
        self.assertFalse(hasattr(pmm, "VectorsRoute"))

    def test_search_hybrid_can_build_multiple_vector_routes(self):
        docs = self.conn.create_table(
            "docs",
            schema=_schema({
                "id": pa.int32(),
                "content": pa.string(),
                "image_embedding": _vector(3),
                "text_embedding": _vector(3),
            }),
            options=_PARQUET_OPTIONS,
        )
        docs.add([
            {
                "id": 1,
                "content": "paimon",
                "image_embedding": [1.0, 0.0, 0.0],
                "text_embedding": [0.0, 1.0, 0.0],
            },
            {
                "id": 2,
                "content": "vector",
                "image_embedding": [0.0, 1.0, 0.0],
                "text_embedding": [1.0, 0.0, 0.0],
            },
        ])

        calls = {"vector_routes": [], "text_routes": []}

        class FakeHybridBuilder:
            def with_limit(self, limit):
                calls["limit"] = limit
                return self

            def with_ranker(self, ranker):
                calls["ranker"] = ranker
                return self

            def add_vector_route(
                    self, column, vector, limit, weight=1.0, options=None):
                calls["vector_routes"].append(
                    (column, vector, limit, weight, options))
                return self

            def add_full_text_route(
                    self, field_name, query, limit, weight=1.0, options=None):
                calls["text_routes"].append(
                    (field_name, json.loads(query), limit, weight, options))
                return self

            def execute_local(self):
                return GlobalIndexResult.from_range(Range(0, 0))

        docs.raw_table.new_hybrid_search_builder = lambda: FakeHybridBuilder()

        result = (
            docs.search_hybrid(
                [
                    pmm.vector_route(
                        "image_embedding",
                        [1.0, 0.0, 0.0],
                        weight=0.7,
                        limit=6,
                        options={"nprobe": "8"},
                    ),
                    pmm.vector_route(
                        "text_embedding",
                        [0.0, 1.0, 0.0],
                        weight=0.3,
                        limit=4,
                        options={"nprobe": "4"},
                    ),
                    pmm.text_route("paimon", weight=0.2),
                ],
                ranker="weighted_score",
                route_limit=4,
            )
            .limit(2)
            .to_list()
        )

        self.assertEqual(2, calls["limit"])
        self.assertEqual("weighted_score", calls["ranker"])
        self.assertEqual(
            [
                (
                    "image_embedding",
                    [1.0, 0.0, 0.0],
                    6,
                    0.7,
                    {"nprobe": "8"},
                ),
                (
                    "text_embedding",
                    [0.0, 1.0, 0.0],
                    4,
                    0.3,
                    {"nprobe": "4"},
                ),
            ],
            calls["vector_routes"],
        )
        self.assertEqual("content", calls["text_routes"][0][0])
        self.assertEqual("paimon", calls["text_routes"][0][1]["match"]["query"])
        self.assertEqual(4, calls["text_routes"][0][2])
        self.assertEqual(0.2, calls["text_routes"][0][3])
        self.assertEqual(
            [{
                "id": 1,
                "content": "paimon",
                "image_embedding": [1.0, 0.0, 0.0],
                "text_embedding": [0.0, 1.0, 0.0],
            }],
            result,
        )


if __name__ == "__main__":
    unittest.main()
