# 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 os
import random
import string
import threading
import unittest
from unittest import mock

import pyarrow as pa

from pypaimon.tests.data_evolution_test_helpers import (
    BatchModeMixin,
    DataEvolutionTestBase,
    StreamModeMixin,
)


# ======================================================================
# Shared base for batch & stream table-update tests
# ======================================================================

class _TableUpdateTestBase(DataEvolutionTestBase):
    """Shared tests for ``TableUpdate.update_by_arrow_with_row_id``.

    Concrete subclasses must inherit from :class:`unittest.TestCase` AND one
    of :class:`_BatchModeMixin` / :class:`_StreamModeMixin`, which add the
    operation-specific primitive ``_apply_update(tu, data, cid)`` on top of
    the framework primitives provided by
    :class:`DataEvolutionTestBase`/:class:`BatchModeMixin`/:class:`StreamModeMixin`.
    """

    # ------------------------------------------------------------------
    # Operation-specific primitive (overridden by mixins)
    # ------------------------------------------------------------------

    def _apply_update(self, table_update, data, cid):
        raise NotImplementedError

    def _apply_update_by_predicate(
            self, table_update, predicate, assignments, cid):
        raise NotImplementedError

    def _apply_delete_by_predicate(self, table_update, predicate, cid):
        raise NotImplementedError

    def _apply_delete_by_row_id(self, table_update, row_ids, cid):
        raise NotImplementedError

    # ------------------------------------------------------------------
    # Helpers built on the primitives
    # ------------------------------------------------------------------

    def _create_seeded_table(self, partition_keys=None, options=None):
        """Create the canonical 5-row / 2-file table used by most tests.

        Layout (row_id → row):
            0: (1, Alice,   25, NYC)
            1: (2, Bob,     30, LA)
            2: (3, Charlie, 35, Chicago)
            3: (4, David,   40, Houston)
            4: (5, Eve,     45, Phoenix)
        """
        table = self._create_table(
            partition_keys=partition_keys,
            options=options,
        )
        self._write_arrow(table, pa.Table.from_pydict({
            'id': [1, 2],
            'name': ['Alice', 'Bob'],
            'age': [25, 30],
            'city': ['NYC', 'LA'],
        }, schema=self.pa_schema))
        self._write_arrow(table, pa.Table.from_pydict({
            'id': [3, 4, 5],
            'name': ['Charlie', 'David', 'Eve'],
            'age': [35, 40, 45],
            'city': ['Chicago', 'Houston', 'Phoenix'],
        }, schema=self.pa_schema))
        return table

    def _create_global_indexed_table_for_predicate_update(self, extra_options=None):
        options = dict(self.table_options)
        options.update({
            'global-index.enabled': 'true',
            'bucket': '-1',
            'file.format': 'parquet',
        })
        options.update(extra_options or {})
        table = self._create_table(options=options)
        self._write_arrow(table, pa.Table.from_pydict({
            'id': [1, 2],
            'name': ['old', 'indexed'],
            'age': [10, 15],
            'city': ['NYC', 'LA'],
        }, schema=self.pa_schema))

        self.assertEqual(
            1,
            table.create_global_index(
                'name',
                options={'sorted-index.records-per-range': '1000'},
            ),
        )
        self._write_arrow(table, pa.Table.from_pydict({
            'id': [3, 4],
            'name': ['new', 'other'],
            'age': [20, 30],
            'city': ['LA', 'SF'],
        }, schema=self.pa_schema))
        return table

    def _do_update(self, table, data, columns):
        """End-to-end ``update_by_arrow_with_row_id`` + commit. Returns the
        commit messages so callers can inspect produced files."""
        wb = self._make_write_builder(table)
        tu = wb.new_update().with_update_type(columns)
        cid = self._next_commit_id()
        msgs = self._apply_update(tu, data, cid)
        tc = wb.new_commit()
        self._apply_commit(tc, msgs, cid)
        tc.close()
        return msgs

    def _do_update_by_predicate(self, table, predicate, assignments):
        wb = self._make_write_builder(table)
        tu = wb.new_update()
        cid = self._next_commit_id()
        msgs = self._apply_update_by_predicate(
            tu,
            predicate,
            assignments,
            cid,
        )
        tc = wb.new_commit()
        self._apply_commit(tc, msgs, cid)
        tc.close()
        return msgs

    def _do_delete_by_predicate(self, table, predicate):
        wb = self._make_write_builder(table)
        tu = wb.new_update()
        cid = self._next_commit_id()
        msgs = self._apply_delete_by_predicate(tu, predicate, cid)
        tc = wb.new_commit()
        self._apply_commit(tc, msgs, cid)
        tc.close()
        return msgs

    def _do_delete_by_row_id(self, table, row_ids):
        wb = self._make_write_builder(table)
        tu = wb.new_update()
        cid = self._next_commit_id()
        msgs = self._apply_delete_by_row_id(tu, row_ids, cid)
        tc = wb.new_commit()
        self._apply_commit(tc, msgs, cid)
        tc.close()
        return msgs

    def _create_seeded_deletion_vector_table(self, partition_keys=None):
        options = dict(self.table_options)
        options['deletion-vectors.enabled'] = 'true'
        table = self._create_seeded_table(
            partition_keys=partition_keys,
            options=options,
        )
        return table

    # ==================================================================
    # Shared tests (run under both batch and stream modes)
    # ==================================================================

    def test_update_existing_column(self):
        """Update a single column across both files (row_ids unordered)."""
        table = self._create_seeded_table()
        self._do_update(table, pa.Table.from_pydict({
            '_ROW_ID': [1, 0, 2, 3, 4],
            'age': [31, 26, 36, 39, 42],
        }), ['age'])
        self.assertEqual(
            [26, 31, 36, 39, 42],
            self._read_all(table)['age'].to_pylist(),
        )

    def test_update_by_predicate(self):
        table = self._create_seeded_table()
        wb = self._make_write_builder(table)
        tu = wb.new_update()
        pb = tu.new_predicate_builder()
        predicate = pb.greater_or_equal('age', 35)

        cid = self._next_commit_id()
        msgs = self._apply_update_by_predicate(
            tu,
            predicate,
            {'age': 99, 'city': 'Updated'},
            cid,
        )
        tc = wb.new_commit()
        self._apply_commit(tc, msgs, cid)
        tc.close()

        result = self._read_all(table)
        self.assertEqual(
            [25, 30, 99, 99, 99],
            result['age'].to_pylist(),
        )
        self.assertEqual(
            ['NYC', 'LA', 'Updated', 'Updated', 'Updated'],
            result['city'].to_pylist(),
        )
        self.assertEqual(
            ['Alice', 'Bob', 'Charlie', 'David', 'Eve'],
            result['name'].to_pylist(),
        )

    def test_update_by_predicate_no_match_is_noop(self):
        table = self._create_seeded_table()
        pb = table.new_read_builder().new_predicate_builder()
        msgs = self._do_update_by_predicate(
            table,
            pb.greater_than('age', 100),
            {'age': 1},
        )

        self.assertEqual([], msgs)
        self.assertEqual(
            [25, 30, 35, 40, 45],
            self._read_all(table)['age'].to_pylist(),
        )

    def test_update_by_predicate_accepts_array_chunked_and_scalar_values(self):
        table = self._create_seeded_table()
        pb = table.new_read_builder().new_predicate_builder()
        self._do_update_by_predicate(
            table,
            pb.greater_or_equal('age', 35),
            {
                'age': pa.array([101, 102, 103], type=pa.int64()),
                'city': pa.chunked_array([
                    pa.array(['Chicago_v2']),
                    pa.array(['Houston_v2', 'Phoenix_v2']),
                ]),
                'name': pa.scalar('patched'),
            },
        )

        result = self._read_all(table)
        rows = {
            row_id: (name, age, city)
            for row_id, name, age, city in zip(
                result['id'].to_pylist(),
                result['name'].to_pylist(),
                result['age'].to_pylist(),
                result['city'].to_pylist(),
            )
        }
        self.assertEqual(
            {
                1: ('Alice', 25, 'NYC'),
                2: ('Bob', 30, 'LA'),
                3: ('patched', 101, 'Chicago_v2'),
                4: ('patched', 102, 'Houston_v2'),
                5: ('patched', 103, 'Phoenix_v2'),
            },
            rows,
        )

    def test_update_by_predicate_updates_all_rows_when_predicate_is_none(self):
        table = self._create_seeded_table()
        self._do_update_by_predicate(
            table,
            None,
            {'age': pa.scalar(7, type=pa.int64()), 'city': None},
        )

        result = self._read_all(table)
        self.assertEqual([7, 7, 7, 7, 7], result['age'].to_pylist())
        self.assertEqual([None, None, None, None, None],
                         result['city'].to_pylist())
        self.assertEqual(
            ['Alice', 'Bob', 'Charlie', 'David', 'Eve'],
            result['name'].to_pylist(),
        )

    def test_update_by_predicate_rejects_assignment_array_length_mismatch(self):
        table = self._create_seeded_table()
        pb = table.new_read_builder().new_predicate_builder()
        with self.assertRaises(ValueError) as ctx:
            self._do_update_by_predicate(
                table,
                pb.greater_or_equal('age', 35),
                {'age': pa.array([1, 2], type=pa.int32())},
            )

        self.assertIn('Assignment array length', str(ctx.exception))
        self.assertEqual(
            [25, 30, 35, 40, 45],
            self._read_all(table)['age'].to_pylist(),
        )

    def test_update_by_predicate_rejects_uncastable_assignment(self):
        table = self._create_seeded_table()
        pb = table.new_read_builder().new_predicate_builder()
        with self.assertRaises((pa.ArrowInvalid, pa.ArrowTypeError, ValueError)):
            self._do_update_by_predicate(
                table,
                pb.equal('id', 1),
                {'age': 'not-an-int'},
            )

        self.assertEqual(
            [25, 30, 35, 40, 45],
            self._read_all(table)['age'].to_pylist(),
        )

    def test_update_by_predicate_with_global_index_updates_unindexed_rows(self):
        table = self._create_global_indexed_table_for_predicate_update({
            'scalar-index.search-mode': 'fast',
        })

        pb = table.new_read_builder().new_predicate_builder()
        self._do_update_by_predicate(
            table,
            pb.equal('name', 'new'),
            {'age': 21},
        )

        result = self._read_all(table)
        ages_by_id = dict(zip(
            result['id'].to_pylist(),
            result['age'].to_pylist(),
        ))
        self.assertEqual({1: 10, 2: 15, 3: 21, 4: 30}, ages_by_id)

    def test_delete_by_predicate_with_global_index_deletes_unindexed_rows(self):
        table = self._create_global_indexed_table_for_predicate_update({
            'deletion-vectors.enabled': 'true',
            'scalar-index.search-mode': 'fast',
        })

        pb = table.new_read_builder().new_predicate_builder()
        self._do_delete_by_predicate(table, pb.equal('name', 'new'))

        result = self._read_all(table).sort_by('id')
        self.assertEqual([1, 2, 4], result['id'].to_pylist())

    def test_update_by_predicate_with_global_index_falls_back_to_full_scan(self):
        table = self._create_global_indexed_table_for_predicate_update()

        pb = table.new_read_builder().new_predicate_builder()
        self._do_update_by_predicate(
            table,
            pb.equal('city', 'LA'),
            {'age': 88},
        )

        result = self._read_all(table)
        ages_by_id = dict(zip(
            result['id'].to_pylist(),
            result['age'].to_pylist(),
        ))
        self.assertEqual({1: 10, 2: 88, 3: 88, 4: 30}, ages_by_id)

    def test_update_by_predicate_with_global_index_handles_compound_predicate(self):
        table = self._create_global_indexed_table_for_predicate_update()

        pb = table.new_read_builder().new_predicate_builder()
        predicate = pb.or_predicates([
            pb.equal('name', 'old'),
            pb.equal('city', 'SF'),
        ])
        self._do_update_by_predicate(table, predicate, {'age': 77})

        result = self._read_all(table)
        ages_by_id = dict(zip(
            result['id'].to_pylist(),
            result['age'].to_pylist(),
        ))
        self.assertEqual({1: 77, 2: 15, 3: 20, 4: 77}, ages_by_id)

    def test_update_by_predicate_resolves_time_travel_scan_snapshot(self):
        table = self._create_table()
        self._write_arrow(table, pa.Table.from_pydict({
            'id': [1],
            'name': ['old'],
            'age': [10],
            'city': ['NYC'],
        }, schema=self.pa_schema))
        table.create_tag('before_new')
        self._write_arrow(table, pa.Table.from_pydict({
            'id': [2],
            'name': ['new'],
            'age': [20],
            'city': ['LA'],
        }, schema=self.pa_schema))

        travel_table = table.copy({'scan.tag-name': 'before_new'})
        pb = travel_table.new_read_builder().new_predicate_builder()
        msgs = self._do_update_by_predicate(
            travel_table,
            pb.equal('name', 'new'),
            {'age': 21},
        )

        self.assertEqual([], msgs)
        result = self._read_all(table)
        ages_by_id = dict(zip(
            result['id'].to_pylist(),
            result['age'].to_pylist(),
        ))
        self.assertEqual({1: 10, 2: 20}, ages_by_id)

    def test_update_by_predicate_resolves_scan_snapshot_id(self):
        table = self._create_table()
        self._write_arrow(table, pa.Table.from_pydict({
            'id': [1],
            'name': ['old'],
            'age': [10],
            'city': ['NYC'],
        }, schema=self.pa_schema))
        snapshot = table.snapshot_manager().get_latest_snapshot()
        self._write_arrow(table, pa.Table.from_pydict({
            'id': [2],
            'name': ['new'],
            'age': [20],
            'city': ['LA'],
        }, schema=self.pa_schema))

        travel_table = table.copy({'scan.snapshot-id': str(snapshot.id)})
        pb = travel_table.new_read_builder().new_predicate_builder()
        msgs = self._do_update_by_predicate(
            travel_table,
            pb.equal('name', 'new'),
            {'age': 21},
        )

        self.assertEqual([], msgs)
        result = self._read_all(table)
        ages_by_id = dict(zip(
            result['id'].to_pylist(),
            result['age'].to_pylist(),
        ))
        self.assertEqual({1: 10, 2: 20}, ages_by_id)

    def test_update_by_predicate_resets_explicit_scan_mode_after_travel(self):
        table = self._create_table()
        self._write_arrow(table, pa.Table.from_pydict({
            'id': [1],
            'name': ['old'],
            'age': [10],
            'city': ['NYC'],
        }, schema=self.pa_schema))
        snapshot = table.snapshot_manager().get_latest_snapshot()
        self._write_arrow(table, pa.Table.from_pydict({
            'id': [2],
            'name': ['new'],
            'age': [20],
            'city': ['LA'],
        }, schema=self.pa_schema))

        travel_table = table.copy({
            'scan.mode': 'from-timestamp',
            'scan.timestamp-millis': str(snapshot.time_millis),
        })
        pb = travel_table.new_read_builder().new_predicate_builder()
        msgs = self._do_update_by_predicate(
            travel_table,
            pb.equal('name', 'new'),
            {'age': 21},
        )

        self.assertEqual([], msgs)
        result = self._read_all(table)
        ages_by_id = dict(zip(
            result['id'].to_pylist(),
            result['age'].to_pylist(),
        ))
        self.assertEqual({1: 10, 2: 20}, ages_by_id)

    def test_update_by_predicate_rejects_empty_assignments(self):
        table = self._create_seeded_table()
        pb = table.new_read_builder().new_predicate_builder()

        with self.assertRaises(ValueError) as ctx:
            self._do_update_by_predicate(
                table,
                pb.equal('id', 1),
                {},
            )
        self.assertIn('assignments must not be empty', str(ctx.exception))

    def test_update_by_predicate_rejects_unknown_column(self):
        table = self._create_seeded_table()
        pb = table.new_read_builder().new_predicate_builder()

        with self.assertRaises(ValueError) as ctx:
            self._do_update_by_predicate(
                table,
                pb.equal('id', 1),
                {'unknown': 1},
            )
        self.assertIn('Column unknown is not in table schema',
                      str(ctx.exception))

    def test_update_by_predicate_rejects_partition_column(self):
        table = self._create_table(partition_keys=['city'])
        self._write_arrow(table, pa.Table.from_pydict({
            'id': [1],
            'name': ['Alice'],
            'age': [25],
            'city': ['NYC'],
        }, schema=self.pa_schema))
        pb = table.new_read_builder().new_predicate_builder()

        with self.assertRaises(ValueError) as ctx:
            self._do_update_by_predicate(
                table,
                pb.equal('id', 1),
                {'city': 'LA'},
            )
        self.assertIn('partition column', str(ctx.exception))

    def test_delete_by_predicate(self):
        table = self._create_seeded_deletion_vector_table()
        pb = table.new_read_builder().new_predicate_builder()

        msgs = self._do_delete_by_predicate(
            table,
            pb.greater_or_equal('age', 35),
        )

        self.assertEqual(1, sum(len(m.index_adds) for m in msgs))
        result = self._read_all(table).sort_by('id')
        self.assertEqual([1, 2], result['id'].to_pylist())
        self.assertEqual(['Alice', 'Bob'], result['name'].to_pylist())

    def test_delete_by_predicate_requires_deletion_vectors(self):
        table = self._create_seeded_table()
        pb = table.new_read_builder().new_predicate_builder()

        with self.assertRaises(ValueError) as ctx:
            self._do_delete_by_predicate(table, pb.equal('id', 1))

        self.assertIn('deletion-vectors.enabled', str(ctx.exception))

    def test_delete_by_row_id(self):
        table = self._create_seeded_deletion_vector_table()

        msgs = self._do_delete_by_row_id(table, [0, 2, 4])

        self.assertEqual(1, sum(len(m.index_adds) for m in msgs))
        result = self._read_all(table).sort_by('id')
        self.assertEqual([2, 4], result['id'].to_pylist())
        self.assertEqual(['Bob', 'David'], result['name'].to_pylist())

    def test_delete_by_row_id_requires_deletion_vectors(self):
        table = self._create_seeded_table()

        with self.assertRaises(ValueError) as ctx:
            self._do_delete_by_row_id(table, [0])

        self.assertIn('deletion-vectors.enabled', str(ctx.exception))

    def test_concurrent_row_level_deletes_conflict_on_same_dv_file(self):
        table = self._create_seeded_deletion_vector_table()
        pb = table.new_read_builder().new_predicate_builder()

        first_wb = self._make_write_builder(table)
        first_update = first_wb.new_update()
        first_cid = self._next_commit_id()
        first_msgs = self._apply_delete_by_predicate(
            first_update,
            pb.equal('id', 1),
            first_cid,
        )

        self._do_delete_by_predicate(table, pb.equal('id', 2))
        result = self._read_all(table).sort_by('id')
        self.assertEqual([1, 3, 4, 5], result['id'].to_pylist())

        first_commit = first_wb.new_commit()
        with self.assertRaises(RuntimeError) as ctx:
            self._apply_commit(first_commit, first_msgs, first_cid)
        first_commit.close()
        self.assertIn('Deletion vector index conflict', str(ctx.exception))

        result = self._read_all(table).sort_by('id')
        self.assertEqual([1, 3, 4, 5], result['id'].to_pylist())

    def test_row_level_delete_conflicts_when_target_file_removed(self):
        table = self._create_seeded_deletion_vector_table(partition_keys=['city'])
        pb = table.new_read_builder().new_predicate_builder()

        first_wb = self._make_write_builder(table)
        first_update = first_wb.new_update()
        first_cid = self._next_commit_id()
        first_msgs = self._apply_delete_by_predicate(
            first_update,
            pb.equal('id', 2),
            first_cid,
        )

        overwrite_wb = table.new_batch_write_builder().overwrite({'city': 'LA'})
        overwrite_write = overwrite_wb.new_write()
        overwrite_commit = overwrite_wb.new_commit()
        overwrite_write.write_arrow(pa.Table.from_pydict({
            'id': [6],
            'name': ['Frank'],
            'age': [28],
            'city': ['LA'],
        }, schema=self.pa_schema))
        overwrite_commit.commit(overwrite_write.prepare_commit())
        overwrite_write.close()
        overwrite_commit.close()

        first_commit = first_wb.new_commit()
        with self.assertRaises(RuntimeError) as ctx:
            self._apply_commit(first_commit, first_msgs, first_cid)
        first_commit.close()
        self.assertIn('Deletion vector index conflict', str(ctx.exception))

        result = self._read_all(table).sort_by('id')
        self.assertEqual([1, 3, 4, 5, 6], result['id'].to_pylist())

    def test_delete_by_partition_predicate_drops_partition_without_dv(self):
        table = self._create_seeded_table(partition_keys=['city'])
        pb = table.new_read_builder().new_predicate_builder()

        msgs = self._do_delete_by_predicate(table, pb.equal('city', 'LA'))

        self.assertGreater(sum(len(m.deleted_files) for m in msgs), 0)
        self.assertEqual(0, sum(len(m.index_adds) for m in msgs))
        result = self._read_all(table).sort_by('id')
        self.assertEqual([1, 3, 4, 5], result['id'].to_pylist())
        self.assertEqual(
            ['NYC', 'Chicago', 'Houston', 'Phoenix'],
            result['city'].to_pylist(),
        )

    def test_delete_by_partition_predicate_deletes_matching_indexes(self):
        table = self._create_seeded_deletion_vector_table(partition_keys=['city'])
        pb = table.new_read_builder().new_predicate_builder()

        self._do_delete_by_predicate(table, pb.equal('id', 2))
        msgs = self._do_delete_by_predicate(table, pb.equal('city', 'LA'))

        self.assertGreater(sum(len(m.deleted_files) for m in msgs), 0)
        self.assertGreater(sum(len(m.index_deletes) for m in msgs), 0)
        result = self._read_all(table).sort_by('id')
        self.assertEqual([1, 3, 4, 5], result['id'].to_pylist())
        self.assertEqual(
            ['NYC', 'Chicago', 'Houston', 'Phoenix'],
            result['city'].to_pylist(),
        )

    def test_delete_by_partition_predicate_conflicts_when_partition_changes(self):
        table = self._create_seeded_table(partition_keys=['city'])
        wb = self._make_write_builder(table)
        tu = wb.new_update()
        pb = tu.new_predicate_builder()
        cid = self._next_commit_id()
        msgs = self._apply_delete_by_predicate(tu, pb.equal('city', 'LA'), cid)

        self._write_arrow(table, pa.Table.from_pydict({
            'id': [6],
            'name': ['Frank'],
            'age': [28],
            'city': ['LA'],
        }, schema=self.pa_schema))

        tc = wb.new_commit()
        with self.assertRaises(RuntimeError) as ctx:
            self._apply_commit(tc, msgs, cid)
        tc.close()
        self.assertIn('Overwrite conflict', str(ctx.exception))

    def test_delete_by_mixed_partition_predicate_still_requires_dv(self):
        table = self._create_seeded_table(partition_keys=['city'])
        pb = table.new_read_builder().new_predicate_builder()
        predicate = pb.and_predicates([
            pb.equal('city', 'LA'),
            pb.equal('age', 30),
        ])

        with self.assertRaises(ValueError) as ctx:
            self._do_delete_by_predicate(table, predicate)

        self.assertIn('deletion-vectors.enabled', str(ctx.exception))

    def test_update_multiple_columns(self):
        """Update ``age`` + ``city`` together; other columns are untouched."""
        table = self._create_seeded_table()
        self._do_update(table, pa.Table.from_pydict({
            '_ROW_ID': [1, 0, 2, 3, 4],
            'age': [31, 26, 36, 39, 42],
            'city': ['Los Angeles', 'New York', 'Chicago', 'Phoenix', 'Houston'],
        }), ['age', 'city'])

        result = self._read_all(table)
        self.assertEqual([26, 31, 36, 39, 42], result['age'].to_pylist())
        self.assertEqual(
            ['New York', 'Los Angeles', 'Chicago', 'Phoenix', 'Houston'],
            result['city'].to_pylist(),
        )

    def test_update_columns_fall_back_to_data_when_unset(self):
        table = self._create_seeded_table()

        self._do_update(table, pa.Table.from_pydict({
            '_ROW_ID': [0, 1, 2, 3, 4],
            'id': [1, 2, 3, 4, 5],
            'name': ['A', 'B', 'C', 'D', 'E'],
            'age': [1, 2, 3, 4, 5],
            'city': ['c0', 'c1', 'c2', 'c3', 'c4'],
        }), ['id', 'name', 'age', 'city'])
        result = self._read_all(table)
        self.assertEqual(['A', 'B', 'C', 'D', 'E'], result['name'].to_pylist())
        self.assertEqual([1, 2, 3, 4, 5], result['age'].to_pylist())
        self.assertEqual(['c0', 'c1', 'c2', 'c3', 'c4'], result['city'].to_pylist())

        wb = self._make_write_builder(table)
        tu = wb.new_update()
        cid = self._next_commit_id()
        msgs = self._apply_update(tu, pa.Table.from_pydict({
            '_ROW_ID': [0, 1],
            'age': [99, 98],
        }), cid)
        tc = wb.new_commit()
        self._apply_commit(tc, msgs, cid)
        tc.close()
        result = self._read_all(table)
        self.assertEqual([99, 98, 3, 4, 5], result['age'].to_pylist())
        self.assertEqual(['A', 'B', 'C', 'D', 'E'], result['name'].to_pylist())

    def test_update_with_only_row_id_raises(self):
        table = self._create_seeded_table()
        wb = self._make_write_builder(table)
        tu = wb.new_update()
        cid = self._next_commit_id()
        with self.assertRaises(ValueError):
            self._apply_update(tu, pa.Table.from_pydict({
                '_ROW_ID': [0, 1],
            }), cid)

    def test_partitioned_table_update(self):
        """Updates work on a partitioned table the same as a flat one."""
        table = self._create_table(partition_keys=['city'])
        self._write_arrow(table, pa.Table.from_pydict({
            'id': [1, 2, 3, 4, 5],
            'name': ['Alice', 'Bob', 'Charlie', 'David', 'Eve'],
            'age': [25, 30, 35, 40, 45],
            'city': ['NYC', 'NYC', 'LA', 'LA', 'Chicago'],
        }, schema=self.pa_schema))

        self._do_update(table, pa.Table.from_pydict({
            '_ROW_ID': [1, 0, 2, 3, 4],
            'age': [31, 26, 36, 41, 46],
        }), ['age'])

        self.assertEqual(
            [26, 31, 36, 41, 46],
            self._read_all(table)['age'].to_pylist(),
        )

    def test_update_preserves_other_columns(self):
        """Updating ``age`` leaves all other columns untouched."""
        table = self._create_seeded_table()
        self._do_update(table, pa.Table.from_pydict({
            '_ROW_ID': [1],
            'age': [999],
        }), ['age'])

        result = self._read_all(table)
        self.assertEqual([25, 999, 35, 40, 45], result['age'].to_pylist())
        self.assertEqual(
            ['Alice', 'Bob', 'Charlie', 'David', 'Eve'],
            result['name'].to_pylist(),
        )
        self.assertEqual(
            ['NYC', 'LA', 'Chicago', 'Houston', 'Phoenix'],
            result['city'].to_pylist(),
        )

    def test_partial_row_updates(self):
        """All partial-row patterns share one parameterised test."""
        cases = [
            # name, row_ids, ages, expected ages after update
            ('single_first_file',  [0],       [100],            [100, 30,  35,  40, 45]),
            ('single_second_file', [1],       [31],             [25,  31,  35,  40, 45]),
            ('one_per_file',       [1, 2],    [200, 300],       [25,  200, 300, 40, 45]),
            ('non_consecutive',    [0, 2, 4], [100, 300, 500],  [100, 30,  300, 40, 500]),
        ]
        for name, row_ids, ages, expected in cases:
            with self.subTest(case=name):
                table = self._create_seeded_table()
                self._do_update(table, pa.Table.from_pydict({
                    '_ROW_ID': row_ids,
                    'age': ages,
                }), ['age'])
                self.assertEqual(
                    expected,
                    self._read_all(table)['age'].to_pylist(),
                )

    def test_multiple_sequential_single_column_updates(self):
        """Two sequential updates on different single columns compose."""
        table = self._create_seeded_table()
        self._do_update(table, pa.Table.from_pydict({
            '_ROW_ID': [1, 0, 2, 3, 4],
            'age': [31, 26, 36, 41, 46],
        }), ['age'])
        self._do_update(table, pa.Table.from_pydict({
            '_ROW_ID': [1, 0, 2, 3, 4],
            'city': ['Los Angeles', 'New York', 'Chicago', 'Phoenix', 'Houston'],
        }), ['city'])

        result = self._read_all(table)
        self.assertEqual([26, 31, 36, 41, 46], result['age'].to_pylist())
        self.assertEqual(
            ['New York', 'Los Angeles', 'Chicago', 'Phoenix', 'Houston'],
            result['city'].to_pylist(),
        )

    def test_sequential_partial_updates(self):
        """Sequential single-row updates accumulate correctly."""
        table = self._create_seeded_table()
        for row_id, age in [(0, 100), (2, 300), (4, 500)]:
            self._do_update(table, pa.Table.from_pydict({
                '_ROW_ID': [row_id],
                'age': [age],
            }), ['age'])
        self.assertEqual(
            [100, 30, 300, 40, 500],
            self._read_all(table)['age'].to_pylist(),
        )

    def test_update_partial_rows_across_two_files(self):
        """A single update modifies partial rows from both data files."""
        table = self._create_table()
        self._write_arrow(table, pa.Table.from_pydict({
            'id': [1, 2, 3, 4, 5],
            'name': ['Alice', 'Bob', 'Charlie', 'David', 'Eve'],
            'age': [20, 25, 30, 35, 40],
            'city': ['NYC', 'LA', 'Chicago', 'Houston', 'Phoenix'],
        }, schema=self.pa_schema))
        self._write_arrow(table, pa.Table.from_pydict({
            'id': [6, 7, 8, 9, 10],
            'name': ['Frank', 'Grace', 'Henry', 'Ivy', 'Jack'],
            'age': [45, 50, 55, 60, 65],
            'city': ['Seattle', 'Boston', 'Denver', 'Miami', 'Atlanta'],
        }, schema=self.pa_schema))

        self._do_update(table, pa.Table.from_pydict({
            '_ROW_ID': [1, 3, 6, 8],
            'age': [100, 200, 300, 400],
            'name': ['Updated_Bob', 'Updated_David', 'Updated_Grace', 'Updated_Ivy'],
        }), ['age', 'name'])

        result = self._read_all(table)
        self.assertEqual(
            [20, 100, 30, 200, 40, 45, 300, 55, 400, 65],
            result['age'].to_pylist(),
        )
        self.assertEqual(
            ['Alice', 'Updated_Bob', 'Charlie', 'Updated_David', 'Eve',
             'Frank', 'Updated_Grace', 'Henry', 'Updated_Ivy', 'Jack'],
            result['name'].to_pylist(),
        )
        # Untouched columns
        self.assertEqual([1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
                         result['id'].to_pylist())
        self.assertEqual(
            ['NYC', 'LA', 'Chicago', 'Houston', 'Phoenix',
             'Seattle', 'Boston', 'Denver', 'Miami', 'Atlanta'],
            result['city'].to_pylist(),
        )

    def test_large_table_partial_column_updates(self):
        """4-step update sequence on a 2000-row / 2-file table.

        Covers: single-column / multi-column / both-files updates while
        verifying that untouched columns stay intact.
        """
        num_row = 1000
        table = self._create_table()

        # 2 commits, num_row rows each
        self._write_arrow(table, pa.Table.from_pydict({
            'id': list(range(num_row)),
            'name': [f'Name_{i}' for i in range(num_row)],
            'age': [20 + i for i in range(num_row)],
            'city': [f'City_{i}' for i in range(num_row)],
        }, schema=self.pa_schema))
        self._write_arrow(table, pa.Table.from_pydict({
            'id': list(range(num_row, num_row * 2)),
            'name': [f'Name_{i}' for i in range(num_row, num_row * 2)],
            'age': [20 + num_row + i for i in range(num_row)],
            'city': [f'City_{i}' for i in range(num_row, num_row * 2)],
        }, schema=self.pa_schema))

        expected_ids = list(range(num_row * 2))
        expected_names = [f'Name_{i}' for i in range(num_row * 2)]
        expected_ages = ([20 + i for i in range(num_row)]
                         + [20 + num_row + i for i in range(num_row)])

        # 1) update id col, 1 row
        self._do_update(table, pa.Table.from_pydict({
            '_ROW_ID': [5], 'id': [999],
        }), ['id'])
        expected_ids[5] = 999

        # 2) update id+name cols, 2 rows
        self._do_update(table, pa.Table.from_pydict({
            '_ROW_ID': [3, 15],
            'id': [888, 777],
            'name': ['Updated_Name_3', 'Updated_Name_15'],
        }), ['id', 'name'])
        expected_ids[3], expected_ids[15] = 888, 777
        expected_names[3] = 'Updated_Name_3'
        expected_names[15] = 'Updated_Name_15'

        # 3) update name col, 1 row
        self._do_update(table, pa.Table.from_pydict({
            '_ROW_ID': [12], 'name': ['NewName_12'],
        }), ['name'])
        expected_names[12] = 'NewName_12'

        # 4) update age col, 4 rows across both files
        self._do_update(table, pa.Table.from_pydict({
            '_ROW_ID': [0, 5, 10, 15],
            'age': [100, 105, 110, 115],
        }), ['age'])
        for r, a in [(0, 100), (5, 105), (10, 110), (15, 115)]:
            expected_ages[r] = a

        result = self._read_all(table)
        self.assertEqual(expected_ids, result['id'].to_pylist())
        self.assertEqual(expected_names, result['name'].to_pylist())
        self.assertEqual(expected_ages, result['age'].to_pylist())
        # city was never touched
        self.assertEqual(
            [f'City_{i}' for i in range(num_row * 2)],
            result['city'].to_pylist(),
        )

    def test_update_with_large_file(self):
        """Updates disable both size- and row-based rolling."""
        from pypaimon.schema.schema_change import SetOption

        N = 5000
        schema = pa.schema([('id', pa.int64()), ('name', pa.string())])
        opts = {
            'row-tracking.enabled': 'true',
            'data-evolution.enabled': 'true',
            'write-only': 'true',
        }
        table = self._create_table(pa_schema=schema, options=opts)
        table_identifier = table.identifier

        self._write_arrow(table, pa.table({
            'id': list(range(N)),
            'name': [''.join(random.choices(string.ascii_letters, k=200))
                     for _ in range(N)],
        }))

        self.catalog.alter_table(
            table_identifier, [
                SetOption('target-file-size', '10kb'),
                SetOption('target-file-row-num', '1'),
            ]
        )
        table = self.catalog.get_table(table_identifier)

        msgs = self._do_update(table, pa.table({
            '_ROW_ID': pa.array(list(range(N)), type=pa.int64()),
            'name': [''.join(random.choices(string.ascii_letters, k=200))
                     for _ in range(N)],
        }), ['name'])

        all_files = [f for m in msgs for f in m.new_files]
        self.assertEqual(1, len(all_files),
                         "Update should produce exactly one file per group")
        self.assertEqual(0, all_files[0].first_row_id)
        self.assertEqual(N, all_files[0].row_count)

    # ------------------------------------------------------------------
    # Validation tests
    # ------------------------------------------------------------------

    def test_nonexistent_column_raises(self):
        table = self._create_seeded_table()
        bad = pa.Table.from_pydict({
            '_ROW_ID': [0, 1, 2, 3, 4],
            'nonexistent_column': [100, 200, 300, 400, 500],
        })
        with self.assertRaises(ValueError) as ctx:
            wb = self._make_write_builder(table)
            tu = wb.new_update().with_update_type(['nonexistent_column'])
            self._apply_update(tu, bad, self._next_commit_id())
        self.assertIn('not in table schema', str(ctx.exception))

    def test_missing_row_id_column_raises(self):
        table = self._create_seeded_table()
        bad = pa.Table.from_pydict({'age': [26, 27, 28, 29, 30]})
        with self.assertRaises(ValueError) as ctx:
            wb = self._make_write_builder(table)
            tu = wb.new_update().with_update_type(['age'])
            self._apply_update(tu, bad, self._next_commit_id())
        self.assertIn('_ROW_ID column', str(ctx.exception))

    def test_invalid_row_id_raises(self):
        """row_id outside valid row_id ranges raises."""
        table = self._create_seeded_table()
        cases = [
            ('out_of_range_high', [0, 10], [26, 100]),
            ('negative',          [-1, 0], [100, 26]),
        ]
        for name, row_ids, ages in cases:
            with self.subTest(case=name):
                wb = self._make_write_builder(table)
                tu = wb.new_update().with_update_type(['age'])
                bad = pa.Table.from_pydict({'_ROW_ID': row_ids, 'age': ages})
                with self.assertRaises(ValueError) as ctx:
                    self._apply_update(tu, bad, self._next_commit_id())
                self.assertIn('does not belong to any valid range', str(ctx.exception))

    def test_duplicate_row_id_raises(self):
        table = self._create_seeded_table()
        wb = self._make_write_builder(table)
        tu = wb.new_update().with_update_type(['age'])
        with self.assertRaises(ValueError) as ctx:
            self._apply_update(
                tu,
                pa.Table.from_pydict({
                    '_ROW_ID': [0, 0, 1],
                    'age': [100, 200, 300],
                }),
                self._next_commit_id(),
            )
        self.assertIn('duplicate _ROW_ID', str(ctx.exception))

    def test_update_deleted_row_id_raises(self):
        """Updating a row_id that fell into a hole after truncate raises."""
        partitioned_schema = pa.schema([
            ('id', pa.int32()),
            ('name', pa.string()),
            ('age', pa.int32()),
            ('region', pa.string()),
        ])
        table = self._create_table(
            pa_schema=partitioned_schema,
            partition_keys=['region'],
        )
        self._write_arrow(table, pa.Table.from_pydict({
            'id': pa.array([1, 2, 3], type=pa.int32()),
            'name': ['A', 'B', 'C'],
            'age': pa.array([10, 20, 30], type=pa.int32()),
            'region': ['US', 'US', 'US'],
        }, schema=partitioned_schema))

        self._write_arrow(table, pa.Table.from_pydict({
            'id': pa.array([4, 5], type=pa.int32()),
            'name': ['D', 'E'],
            'age': pa.array([40, 50], type=pa.int32()),
            'region': ['EU', 'EU'],
        }, schema=partitioned_schema))

        wb = table.new_batch_write_builder()
        tc = wb.new_commit()
        tc.truncate_partitions([{'region': 'US'}])

        wb = self._make_write_builder(table)
        tu = wb.new_update().with_update_type(['age'])
        with self.assertRaises(ValueError) as ctx:
            self._apply_update(
                tu,
                pa.Table.from_pydict({'_ROW_ID': [0], 'age': [99]}),
                self._next_commit_id(),
            )
        self.assertIn('does not belong to any valid range', str(ctx.exception))

    # ------------------------------------------------------------------
    # Concurrency tests
    # ------------------------------------------------------------------

    def _run_concurrent_updates(self, table, thread_specs, max_retries):
        """Run a batch of concurrent updates with conflict-retry; return the
        commit order (``thread_index`` of the winning commit appended last)."""
        errors = []
        completion_order = []
        lock = threading.Lock()
        retry_marker = (
            "multiple 'MERGE INTO' operations have encountered conflicts"
        )

        def worker(idx, spec):
            # Tag each thread's commits so the durable winner can be read back from the
            # latest snapshot's commit_user (the order threads return in is not the commit order).
            worker_table = table.copy({'commit.user-prefix': 'w%d' % idx})
            for _ in range(max_retries):
                try:
                    self._do_update(worker_table, pa.Table.from_pydict({
                        '_ROW_ID': spec['row_ids'],
                        'age': spec['ages'],
                    }), ['age'])
                    with lock:
                        completion_order.append(idx)
                    return
                except Exception as e:
                    if retry_marker not in str(e):
                        errors.append(f"Thread-{idx} unexpected error: {e}")
                        return
            errors.append(f"Thread-{idx} did not succeed in {max_retries} retries")

        threads = [threading.Thread(target=worker, args=(i, s))
                   for i, s in enumerate(thread_specs)]
        for t in threads:
            t.start()
        for t in threads:
            t.join(timeout=120)

        self.assertEqual([], errors)
        self.assertEqual(len(thread_specs), len(completion_order),
                         "Not all threads committed successfully")
        return completion_order

    def test_concurrent_updates_disjoint_rows(self):
        """Disjoint per-thread updates all materialise into the final state."""
        table = self._create_seeded_table()
        specs = [
            {'row_ids': [0, 1], 'ages': [100, 200]},
            {'row_ids': [2, 3], 'ages': [300, 400]},
            {'row_ids': [4],    'ages': [500]},
        ]
        self._run_concurrent_updates(table, specs, max_retries=20)
        self.assertEqual(
            [100, 200, 300, 400, 500],
            self._read_all(table)['age'].to_pylist(),
        )

    def test_concurrent_updates_overlapping_rows_last_writer_wins(self):
        """When threads compete on the same rows, the final commit wins."""
        table = self._create_seeded_table()
        specs = [
            {'row_ids': [0, 1, 2], 'ages': [101, 201, 301]},
            {'row_ids': [0, 1, 2], 'ages': [102, 202, 302]},
            {'row_ids': [0, 1, 2], 'ages': [103, 203, 303]},
        ]
        self._run_concurrent_updates(table, specs, max_retries=30)
        ages = self._read_all(table)['age'].to_pylist()
        # The real winner is the thread whose commit produced the latest snapshot.
        winner = int(table.snapshot_manager().get_latest_snapshot().commit_user[1:].split('_')[0])
        self.assertEqual(specs[winner]['ages'], ages[:3])
        self.assertEqual([40, 45], ages[3:])

    def test_update_list_and_map_columns(self):
        list_map_schema = pa.schema([
            ('id', pa.int32()),
            ('tags', pa.list_(pa.string())),
            ('meta', pa.map_(pa.string(), pa.string())),
        ])
        table = self._create_table(pa_schema=list_map_schema)
        self._write_arrow(table, pa.Table.from_pydict({
            'id': [1, 2, 3],
            'tags': [['a', 'b'], ['c'], ['d', 'e']],
            'meta': [[('k1', 'v1')], [('k2', 'v2')], [('k3', 'v3')]],
        }, schema=list_map_schema))

        rb = table.new_read_builder().with_projection(
            ['id', '_ROW_ID'])
        rid_result = rb.new_read().to_arrow(
            rb.new_scan().plan().splits()).sort_by('id')
        row_ids = rid_result['_ROW_ID'].to_pylist()

        self._do_update(table, pa.Table.from_pydict({
            '_ROW_ID': pa.array([row_ids[0], row_ids[2]],
                                type=pa.int64()),
            'tags': [['x', 'y'], ['z']],
            'meta': [[('k1', 'new1')], [('k3', 'new3')]],
        }), ['tags', 'meta'])

        result = self._read_all(table).sort_by('id')
        self.assertEqual(
            [['x', 'y'], ['c'], ['z']],
            result['tags'].to_pylist())
        self.assertEqual(
            [[('k1', 'new1')], [('k2', 'v2')], [('k3', 'new3')]],
            result['meta'].to_pylist())

        self._do_update(table, pa.Table.from_pydict({
            '_ROW_ID': pa.array([row_ids[1]], type=pa.int64()),
            'meta': [{'k2': 'dict_val'}],
        }), ['meta'])
        result2 = self._read_all(table).sort_by('id')
        self.assertEqual(
            [[('k1', 'new1')], [('k2', 'dict_val')], [('k3', 'new3')]],
            result2['meta'].to_pylist())

        with self.assertRaisesRegex(ValueError, "schema-less dict"):
            self._do_update(table, pa.Table.from_pydict({
                '_ROW_ID': pa.array([row_ids[0], row_ids[2]],
                                    type=pa.int64()),
                'meta': [{'a': '1'}, {'b': '2'}],
            }), ['meta'])

        self._do_update(table, pa.Table.from_pydict({
            '_ROW_ID': pa.array([row_ids[0], row_ids[2]],
                                type=pa.int64()),
            'meta': pa.array(
                [[('a', '1')], [('b', '2')]],
                type=pa.map_(pa.string(), pa.string())),
        }), ['meta'])
        result3 = self._read_all(table).sort_by('id')
        self.assertEqual(
            [[('a', '1')], [('k2', 'dict_val')], [('b', '2')]],
            result3['meta'].to_pylist())

        with self.assertRaisesRegex(ValueError, "schema-less dict"):
            self._do_update(table, pa.Table.from_pydict({
                '_ROW_ID': pa.array([row_ids[0]], type=pa.int64()),
                'meta': [{'a': None}],
            }), ['meta'])

        self._do_update(table, pa.Table.from_pydict({
            '_ROW_ID': pa.array([row_ids[0]], type=pa.int64()),
            'meta': [[('a', None)]],
        }, schema=pa.schema([
            ('_ROW_ID', pa.int64()),
            ('meta', pa.map_(pa.string(), pa.string())),
        ])), ['meta'])
        result4 = self._read_all(table).sort_by('id')
        self.assertEqual(
            [[('a', None)], [('k2', 'dict_val')], [('b', '2')]],
            result4['meta'].to_pylist())


# ======================================================================
# Mode-specific mixins (add the ``update_by_arrow_with_row_id`` primitive)
# ======================================================================

class _BatchModeMixin(BatchModeMixin):
    def _apply_update(self, table_update, data, cid):
        return table_update.update_by_arrow_with_row_id(data)

    def _apply_update_by_predicate(
            self, table_update, predicate, assignments, cid):
        return table_update.update_by_predicate(predicate, assignments)

    def _apply_delete_by_predicate(self, table_update, predicate, cid):
        return table_update.delete_by_predicate(predicate)

    def _apply_delete_by_row_id(self, table_update, row_ids, cid):
        return table_update.delete_by_row_id(row_ids)


class _StreamModeMixin(StreamModeMixin):
    def _apply_update(self, table_update, data, cid):
        return table_update.update_by_arrow_with_row_id(data, cid)

    def _apply_update_by_predicate(
            self, table_update, predicate, assignments, cid):
        return table_update.update_by_predicate(
            predicate,
            assignments,
            cid,
        )

    def _apply_delete_by_predicate(self, table_update, predicate, cid):
        return table_update.delete_by_predicate(predicate, cid)

    def _apply_delete_by_row_id(self, table_update, row_ids, cid):
        return table_update.delete_by_row_id(row_ids, cid)


# ======================================================================
# Concrete test classes
# ======================================================================

class TableUpdateBatchTest(_BatchModeMixin, _TableUpdateTestBase, unittest.TestCase):
    """All shared update tests under batch (``BatchWriteBuilder``) semantics."""

    def test_update_by_row_id_aborts_files_after_prepare_commit_failure(self):
        from pypaimon.write.table_update_by_row_id import TableUpdateByRowId

        table_schema = pa.schema([
            ('id', pa.int32()),
            ('age', pa.int32()),
            ('picture', pa.large_binary()),
        ])
        table = self._create_table(pa_schema=table_schema)
        self._write_arrow(table, pa.Table.from_pydict({
            'id': [1, 2],
            'age': [10, 20],
            'picture': [b'blob-1', b'blob-2'],
        }, schema=table_schema))

        rb = table.new_read_builder().with_projection(['id', '_ROW_ID'])
        row_ids = rb.new_read().to_arrow(
            rb.new_scan().plan().splits()).sort_by('id')['_ROW_ID']
        before_files = self._list_table_files(table)

        def fail_after_prepare_commit(
                new_files, first_row_id, column_names, blob_columns):
            raise RuntimeError("forced failure after prepare_commit")

        wb = self._make_write_builder(table)
        tu = wb.new_update().with_update_type(['age', 'picture'])
        with mock.patch.object(
                TableUpdateByRowId,
                '_assign_update_file_metadata',
                new=staticmethod(fail_after_prepare_commit)):
            with self.assertRaisesRegex(
                    RuntimeError, "forced failure after prepare_commit"):
                self._apply_update(tu, pa.Table.from_pydict({
                    '_ROW_ID': [row_ids[0].as_py()],
                    'age': [99],
                    'picture': [b'updated-blob'],
                }, schema=pa.schema([
                    ('_ROW_ID', pa.int64()),
                    ('age', pa.int32()),
                    ('picture', pa.large_binary()),
                ])), self._next_commit_id())

        self.assertEqual(before_files, self._list_table_files(table))

    @staticmethod
    def _list_table_files(table):
        return {
            os.path.relpath(os.path.join(root, name), table.table_path)
            for root, _dirs, files in os.walk(table.table_path)
            for name in files
        }


class TableUpdateStreamTest(_StreamModeMixin, _TableUpdateTestBase, unittest.TestCase):
    """All shared update tests under stream (``StreamWriteBuilder``) semantics,
    plus stream-only multi-commit scenarios that are impossible to express
    under :class:`BatchTableCommit` (which forbids re-committing).
    """

    # ------------------------------------------------------------------
    # Stream-only helpers
    # ------------------------------------------------------------------

    def _stream_commit_age_updates_by_row_id(self, tu, tc, commit_ids, row_age_pairs):
        """Apply one ``age`` update per ``commit_id``; ``tu`` is already
        ``with_update_type(['age'])``."""
        for cid, (row_id, age) in zip(commit_ids, row_age_pairs):
            msgs = tu.update_by_arrow_with_row_id(
                pa.Table.from_pydict({'_ROW_ID': [row_id], 'age': [age]}),
                cid,
            )
            tc.commit(msgs, cid)

    # ------------------------------------------------------------------
    # Stream-only tests
    # ------------------------------------------------------------------

    def test_stream_multi_commit_on_one_write_builder(self):
        """A single ``StreamWriteBuilder`` drives many update+commit cycles
        with reused ``tu`` and ``tc`` instances. Each commit produces its own
        snapshot tagged with the caller-supplied ``commit_identifier`` under
        a stable ``commit_user`` — the core contract distinguishing stream
        from batch mode.

        Parameterised over both contiguous and sparse identifier sequences
        to catch any accidental coupling between ``commit_identifier`` and
        ``snapshot_id`` ordering.
        """
        # (row_id, age) for each of the 3 commits — same triplet for both runs.
        per_commit_rows = [(0, 100), (2, 300), (4, 500)]
        expected_ages = [100, 30, 300, 40, 500]

        for case_name, commit_ids in [
            ('contiguous', [1, 2, 3]),
            ('sparse',     [42, 100, 1000]),
        ]:
            with self.subTest(case=case_name):
                table = self._create_seeded_table()
                wb, tc, base_snapshot_id = self._stream_commit_session(table)
                tu = wb.new_update().with_update_type(['age'])
                self._stream_commit_age_updates_by_row_id(
                    tu, tc, commit_ids, per_commit_rows,
                )
                tc.close()

                result = self._read_all(table)
                self.assertEqual(expected_ages, result['age'].to_pylist())
                self.assertEqual([1, 2, 3, 4, 5], result['id'].to_pylist())
                self.assertEqual(
                    ['Alice', 'Bob', 'Charlie', 'David', 'Eve'],
                    result['name'].to_pylist(),
                )
                self.assertEqual(
                    ['NYC', 'LA', 'Chicago', 'Houston', 'Phoenix'],
                    result['city'].to_pylist(),
                )
                self._assert_stream_builder_snapshots(
                    table, wb, base_snapshot_id, commit_ids,
                )

    def test_stream_commit_instance_accepts_messages_from_different_updates(self):
        """``StreamTableCommit`` may be committed many times, and each
        commit may carry messages from a *different* ``StreamTableUpdate``
        instance (different ``update_cols``). The one-shot restriction
        that applies to :class:`BatchTableCommit` does not apply here.
        """
        table = self._create_seeded_table()
        wb, tc, base_snapshot_id = self._stream_commit_session(table)

        # First update: only 'age' projection, fresh tu.
        tu1 = wb.new_update().with_update_type(['age'])
        msgs1 = tu1.update_by_arrow_with_row_id(
            pa.Table.from_pydict({'_ROW_ID': [0], 'age': [111]}), 1,
        )
        tc.commit(msgs1, 1)

        # Second update: different 'city' projection, fresh tu, same tc.
        tu2 = wb.new_update().with_update_type(['city'])
        msgs2 = tu2.update_by_arrow_with_row_id(
            pa.Table.from_pydict({'_ROW_ID': [0], 'city': ['Beijing']}), 2,
        )
        tc.commit(msgs2, 2)
        tc.close()

        result = self._read_all(table)
        self.assertEqual([111, 30, 35, 40, 45], result['age'].to_pylist())
        self.assertEqual(
            ['Beijing', 'LA', 'Chicago', 'Houston', 'Phoenix'],
            result['city'].to_pylist(),
        )
        self.assertEqual([1, 2, 3, 4, 5], result['id'].to_pylist())
        self.assertEqual(
            ['Alice', 'Bob', 'Charlie', 'David', 'Eve'],
            result['name'].to_pylist(),
        )
        self._assert_stream_builder_snapshots(
            table, wb, base_snapshot_id, [1, 2],
        )

    def test_stream_interleaved_write_and_update_on_same_builder(self):
        """A single stream builder may perform an initial write followed by
        an update on the same ``tc``, each tagged with its own
        ``commit_identifier`` that propagates to its own snapshot.
        """
        table = self._create_table()
        wb, tc, base_snapshot_id = self._stream_commit_session(table)
        tw = wb.new_write()

        tw.write_arrow(pa.Table.from_pydict({
            'id': [1, 2, 3],
            'name': ['A', 'B', 'C'],
            'age': [10, 20, 30],
            'city': ['X', 'Y', 'Z'],
        }, schema=self.pa_schema))
        tc.commit(tw.prepare_commit(1), 1)
        tw.close()

        tu = wb.new_update().with_update_type(['age'])
        msgs = tu.update_by_arrow_with_row_id(
            pa.Table.from_pydict({'_ROW_ID': [0, 2], 'age': [111, 333]}), 2,
        )
        tc.commit(msgs, 2)
        tc.close()

        result = self._read_all(table)
        self.assertEqual([111, 20, 333], result['age'].to_pylist())
        self.assertEqual([1, 2, 3], result['id'].to_pylist())
        self.assertEqual(['A', 'B', 'C'], result['name'].to_pylist())
        self.assertEqual(['X', 'Y', 'Z'], result['city'].to_pylist())
        self._assert_stream_builder_snapshots(
            table, wb, base_snapshot_id, [1, 2],
        )


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