# 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.

"""Tests for ChunkShuffleSplitGenerator and TableScan.with_chunk_shuffle.

Algorithmic tests use Mock entries so they don't touch disk; the
end-to-end test writes a real append table and validates that all
workers together cover the data exactly once.
"""

import os
import shutil
import tempfile
import unittest
from unittest.mock import Mock, patch

import pyarrow as pa

from pypaimon import CatalogFactory, Schema
from pypaimon.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector
from pypaimon.globalindex.indexed_split import IndexedSplit
from pypaimon.manifest.schema.data_file_meta import DataFileMeta
from pypaimon.read.scanner.chunk_shuffle_split_generator import (
    AppendChunkShuffleSplitGenerator,
    DataEvolutionChunkShuffleSplitGenerator,
    _LiveRowRangeSlicer,
)
from pypaimon.read.sliced_split import SlicedSplit
from pypaimon.read.split import DataSplit
from pypaimon.table.source.deletion_file import DeletionFile
from pypaimon.utils.range import Range


def _mock_table(table_path='/tmp/_chunk_shuffle_test_path'):
    table = Mock()
    table.table_path = table_path
    table.options = Mock()
    return table


def _mock_entry(partition_values, bucket, file_name, row_count, file_size=1024):
    entry = Mock()
    entry.partition = Mock()
    entry.partition.values = partition_values
    entry.bucket = bucket
    entry.file = Mock()
    entry.file.file_name = file_name
    entry.file.file_size = file_size
    entry.file.row_count = row_count
    # Swallow set_file_path so we don't need to mock partition path encoding.
    entry.file.set_file_path = Mock()
    return entry


def _make_generator(seed, chunk_size, table=None, deletion_files_map=None):
    if table is None:
        table = _mock_table()
    return AppendChunkShuffleSplitGenerator(
        table,
        target_split_size=128 * 1024 * 1024,
        open_file_cost=4 * 1024 * 1024,
        deletion_files_map=deletion_files_map,
        seed=seed,
        chunk_size=chunk_size,
    )


def _make_de_generator(seed, chunk_size, table=None, deletion_files_map=None):
    if table is None:
        table = _mock_table()
    return DataEvolutionChunkShuffleSplitGenerator(
        table,
        target_split_size=128 * 1024 * 1024,
        open_file_cost=4 * 1024 * 1024,
        deletion_files_map=deletion_files_map,
        seed=seed,
        chunk_size=chunk_size,
    )


def _mock_de_entry(
    partition_values,
    bucket,
    file_name,
    first_row_id,
    row_count,
    file_size=1024,
    max_sequence_number=0,
):
    """A DE-flavoured mock entry: file carries first_row_id and a real
    Range so :meth:`row_id_range` and ``Range.overlaps`` work."""
    entry = Mock()
    entry.partition = Mock()
    entry.partition.values = partition_values
    entry.bucket = bucket
    file = Mock(spec=DataFileMeta)
    file.file_name = file_name
    file.file_size = file_size
    file.row_count = row_count
    file.first_row_id = first_row_id
    file.max_sequence_number = max_sequence_number
    file.row_id_range = lambda f=first_row_id, c=row_count: Range(f, f + c - 1)
    file.set_file_path = Mock()
    entry.file = file
    return entry


def _de_file(file_name, first_row_id, row_count):
    """A DE file mock; row_id_range() is a real Range, or None when first_row_id is None."""
    file = Mock(spec=DataFileMeta)
    file.file_name = file_name
    rng = Range(first_row_id, first_row_id + row_count - 1) if first_row_id is not None else None
    file.row_id_range = lambda r=rng: r
    return file


def _split_signature(split):
    """A stable, comparable identity for a split — what the worker would actually read."""
    if isinstance(split, SlicedSplit):
        underlying = split.data_split()
        files = tuple(f.file_name for f in underlying.files)
        idx_map = tuple(sorted(split.shard_file_idx_map().items()))
        return (tuple(underlying.partition.values), underlying.bucket, files, idx_map)
    if isinstance(split, IndexedSplit):
        underlying = split.data_split()
        files = tuple(sorted(f.file_name for f in underlying.files))
        ranges = tuple((r.from_, r.to) for r in split.row_ranges())
        return (tuple(underlying.partition.values), underlying.bucket, files, ranges)
    if isinstance(split, DataSplit):
        files = tuple(f.file_name for f in split.files)
        return (tuple(split.partition.values), split.bucket, files, ())
    raise AssertionError("unexpected split type: %r" % type(split))


def _split_rows(split):
    """Effective row count this split actually exposes."""
    return split.row_count


def _bitmap_deletion_vector(*positions):
    deletion_vector = BitmapDeletionVector()
    for position in positions:
        deletion_vector.delete(position)
    return deletion_vector


class LiveRowRangeSlicerTest(unittest.TestCase):

    def test_slices_by_live_rows_and_absorbs_deletion_runs(self):
        slicer = _LiveRowRangeSlicer(10, iter([0, 3, 4, 9]))

        first = slicer.take(3)
        second = slicer.take(3)

        self.assertEqual(
            (
                first.start_inclusive,
                first.end_exclusive,
                first.live_row_count,
            ),
            (0, 6, 3),
        )
        self.assertEqual(
            (
                second.start_inclusive,
                second.end_exclusive,
                second.live_row_count,
            ),
            (6, 10, 3),
        )
        self.assertEqual(first.to_closed_row_id_range(100), Range(100, 105))
        self.assertIsNone(slicer.take(3))

    def test_returns_smaller_tail_and_skips_fully_deleted_source(self):
        tail_slicer = _LiveRowRangeSlicer(6, iter([1, 4]))
        tail = tail_slicer.take(10)
        self.assertEqual(
            (
                tail.start_inclusive,
                tail.end_exclusive,
                tail.live_row_count,
            ),
            (0, 6, 4),
        )
        self.assertIsNone(tail_slicer.take(1))

        deleted_slicer = _LiveRowRangeSlicer(3, iter([0, 1, 2]))
        self.assertIsNone(deleted_slicer.take(1))

    def test_rejects_invalid_deletion_positions(self):
        with self.assertRaisesRegex(ValueError, "outside physical row range"):
            _LiveRowRangeSlicer(3, iter([-1]))

        slicer = _LiveRowRangeSlicer(3, iter([0, 0]))
        with self.assertRaisesRegex(ValueError, "strictly increasing"):
            slicer.take(1)


class ChunkShuffleSplitGeneratorAlgoTest(unittest.TestCase):

    def test_no_entries_returns_empty(self):
        gen = _make_generator(seed=1, chunk_size=100)
        self.assertEqual(gen.create_splits([]), [])

    def test_full_files_no_truncation(self):
        entries = [
            _mock_entry([], 0, 'f1', 100),
            _mock_entry([], 0, 'f2', 100),
            _mock_entry([], 0, 'f3', 100),
        ]
        gen = _make_generator(seed=1, chunk_size=100)
        splits = gen.create_splits(entries)
        # 3 chunks, each holding exactly one whole file → all DataSplit, no SlicedSplit
        self.assertEqual(len(splits), 3)
        for s in splits:
            self.assertIsInstance(s, DataSplit)
            self.assertEqual(s.row_count, 100)

    def test_chunk_truncates_inside_file(self):
        # one file of 250 rows, chunk_size 100 → 3 chunks: 100, 100, 50
        entries = [_mock_entry([], 0, 'f1', 250)]
        gen = _make_generator(seed=1, chunk_size=100)
        splits = gen.create_splits(entries)
        self.assertEqual(len(splits), 3)
        # All three chunks slice the same file → all SlicedSplit
        for s in splits:
            self.assertIsInstance(s, SlicedSplit)
        # union of (start, end) intervals must cover [0, 250)
        intervals = sorted(s.shard_file_idx_map()['f1'] for s in splits)
        self.assertEqual(intervals, [(0, 100), (100, 200), (200, 250)])
        total = sum(end - start for start, end in intervals)
        self.assertEqual(total, 250)

    def test_chunk_spans_multiple_files(self):
        # f1=30, f2=30, f3=30, chunk_size=50 → chunks: [f1(30)+f2(0,20)], [f2(20,30)+f3(0,40 cap 30=30)] ...
        entries = [
            _mock_entry([], 0, 'f1', 30),
            _mock_entry([], 0, 'f2', 30),
            _mock_entry([], 0, 'f3', 30),
        ]
        gen = _make_generator(seed=1, chunk_size=50)
        splits = gen.create_splits(entries)
        # total 90 rows, chunk_size 50 → 2 chunks (50 + 40)
        self.assertEqual(len(splits), 2)
        total_rows = sum(_split_rows(s) for s in splits)
        self.assertEqual(total_rows, 90)

    def test_chunk_size_larger_than_total(self):
        entries = [
            _mock_entry([], 0, 'f1', 30),
            _mock_entry([], 0, 'f2', 30),
        ]
        gen = _make_generator(seed=1, chunk_size=1000)
        splits = gen.create_splits(entries)
        self.assertEqual(len(splits), 1)
        # No truncation — full files inside one chunk → DataSplit not SlicedSplit
        self.assertIsInstance(splits[0], DataSplit)
        self.assertEqual(_split_rows(splits[0]), 60)

    def test_deletion_vector_slices_by_live_rows_and_is_attached(self):
        entry = _mock_entry([], 0, 'f1', 10)
        deletion_file = DeletionFile('dv.index', 10, 20, cardinality=4)
        deletion_files_map = {((), 0): {'f1': deletion_file}}
        deletion_vector = _bitmap_deletion_vector(0, 3, 4, 9)
        gen = _make_generator(
            seed=1,
            chunk_size=3,
            deletion_files_map=deletion_files_map,
        )

        with patch(
            'pypaimon.read.scanner.chunk_shuffle_split_generator.DeletionVector.read',
            return_value=deletion_vector,
        ) as read:
            splits = gen.create_splits([entry])

        self.assertEqual(
            sorted(s.shard_file_idx_map()['f1'] for s in splits),
            [(0, 6), (6, 10)],
        )
        self.assertEqual(
            sorted(s.merged_row_count() for s in splits),
            [3, 3],
        )
        for split in splits:
            self.assertEqual(
                split.data_split().data_deletion_files,
                [deletion_file],
            )
        read.assert_called_once_with(gen.table.file_io, deletion_file)

    def test_whole_file_unknown_dv_cardinality_preserves_live_row_count(self):
        entry = _mock_entry([], 0, 'f1', 10)
        deletion_file = DeletionFile('dv.index', 10, 20, cardinality=None)
        gen = _make_generator(
            seed=1,
            chunk_size=6,
            deletion_files_map={((), 0): {'f1': deletion_file}},
        )

        with patch(
            'pypaimon.read.scanner.chunk_shuffle_split_generator.DeletionVector.read',
            return_value=_bitmap_deletion_vector(0, 3, 4, 9),
        ) as read:
            splits = gen.create_splits([entry])

        self.assertEqual(len(splits), 1)
        self.assertIsInstance(splits[0], SlicedSplit)
        self.assertEqual(splits[0].shard_file_idx_map(), {})
        self.assertEqual(splits[0].row_count, 10)
        self.assertEqual(splits[0].merged_row_count(), 6)
        read.assert_called_once_with(gen.table.file_io, deletion_file)

    def test_two_deletion_vector_files_share_one_live_row_chunk(self):
        # Each file contributes 5 - 2 = 3 live rows, so both should fit
        # exactly in one six-row chunk.
        entries = [
            _mock_entry([], 0, 'f1', 5),
            _mock_entry([], 0, 'f2', 5),
        ]
        first_deletion_file = DeletionFile(
            'dv.index',
            10,
            20,
            cardinality=2,
        )
        second_deletion_file = DeletionFile(
            'dv.index',
            30,
            20,
            cardinality=2,
        )
        deletion_files_map = {
            ((), 0): {
                'f1': first_deletion_file,
                'f2': second_deletion_file,
            }
        }
        deletion_vectors = {
            first_deletion_file: _bitmap_deletion_vector(1, 3),
            second_deletion_file: _bitmap_deletion_vector(0, 4),
        }
        gen = _make_generator(
            seed=1,
            chunk_size=6,
            deletion_files_map=deletion_files_map,
        )

        with patch(
            'pypaimon.read.scanner.chunk_shuffle_split_generator.DeletionVector.read',
            side_effect=lambda _, deletion_file: deletion_vectors[deletion_file],
        ) as read:
            splits = gen.create_splits(entries)

        self.assertEqual(len(splits), 1)
        self.assertIsInstance(splits[0], DataSplit)
        self.assertEqual(
            [file.file_name for file in splits[0].files],
            ['f1', 'f2'],
        )
        self.assertEqual(
            splits[0].data_deletion_files,
            [first_deletion_file, second_deletion_file],
        )
        self.assertEqual(read.call_count, 2)
        self.assertEqual(
            {mock_call.args[1] for mock_call in read.call_args_list},
            {first_deletion_file, second_deletion_file},
        )

    def test_fully_deleted_file_does_not_create_a_segment(self):
        entry = _mock_entry([], 0, 'f1', 5)
        deletion_file = DeletionFile('dv.index', 10, 20, cardinality=5)
        gen = _make_generator(
            seed=1,
            chunk_size=2,
            deletion_files_map={((), 0): {'f1': deletion_file}},
        )

        with patch(
            'pypaimon.read.scanner.chunk_shuffle_split_generator.DeletionVector.read',
        ) as read:
            self.assertEqual(gen.create_splits([entry]), [])

        read.assert_not_called()

    def test_deterministic_same_seed_same_order(self):
        entries = [_mock_entry([], 0, f'f{i}', 100) for i in range(20)]
        gen1 = _make_generator(seed=42, chunk_size=50)
        gen2 = _make_generator(seed=42, chunk_size=50)
        splits1 = gen1.create_splits(entries)
        splits2 = gen2.create_splits(entries)
        self.assertEqual(
            [_split_signature(s) for s in splits1],
            [_split_signature(s) for s in splits2],
        )

    def test_different_seed_different_order(self):
        entries = [_mock_entry([], 0, f'f{i}', 100) for i in range(50)]
        gen1 = _make_generator(seed=1, chunk_size=100)
        gen2 = _make_generator(seed=2, chunk_size=100)
        sigs1 = [_split_signature(s) for s in gen1.create_splits(entries)]
        sigs2 = [_split_signature(s) for s in gen2.create_splits(entries)]
        # Same set of chunks, different order — high probability they differ on 50 items
        self.assertEqual(sorted(sigs1), sorted(sigs2))
        self.assertNotEqual(sigs1, sigs2)

    def test_shuffle_actually_reorders(self):
        # 20 files in scan order f0..f19. After shuffle the file order should not be sorted.
        entries = [_mock_entry([], 0, f'f{i:02d}', 100) for i in range(20)]
        gen = _make_generator(seed=42, chunk_size=100)
        splits = gen.create_splits(entries)
        file_names = [s.files[0].file_name for s in splits]
        self.assertNotEqual(file_names, sorted(file_names))

    def test_shard_round_trip_no_overlap_no_loss(self):
        # 13 files × 100 rows = 1300 rows. 4 workers.
        entries = [_mock_entry([], 0, f'f{i:02d}', 100) for i in range(13)]
        num_workers = 4
        all_sigs = []
        total_rows = 0
        for worker in range(num_workers):
            gen = _make_generator(seed=7, chunk_size=100)
            gen.with_shard(worker, num_workers)
            splits = gen.create_splits(list(entries))  # copy: shuffle is in-place on chunks list
            for s in splits:
                all_sigs.append(_split_signature(s))
                total_rows += _split_rows(s)
        self.assertEqual(total_rows, 13 * 100)
        # No duplicate chunks across workers
        self.assertEqual(len(all_sigs), len(set(all_sigs)))
        # All chunks together equal an unsharded run
        unsharded = _make_generator(seed=7, chunk_size=100).create_splits(list(entries))
        self.assertEqual(
            sorted(all_sigs),
            sorted(_split_signature(s) for s in unsharded),
        )

    def test_shard_balanced_distribution(self):
        # 10 chunks across 3 workers → 4, 3, 3 (front-loaded by _compute_shard_range)
        entries = [_mock_entry([], 0, f'f{i:02d}', 100) for i in range(10)]
        counts = []
        for worker in range(3):
            gen = _make_generator(seed=0, chunk_size=100)
            gen.with_shard(worker, 3)
            counts.append(len(gen.create_splits(list(entries))))
        self.assertEqual(sorted(counts, reverse=True), [4, 3, 3])

    def test_chunks_fewer_than_workers(self):
        # 2 chunks, 5 workers → 3 workers get nothing
        entries = [_mock_entry([], 0, f'f{i}', 100) for i in range(2)]
        empties = 0
        non_empties = 0
        for worker in range(5):
            gen = _make_generator(seed=0, chunk_size=100)
            gen.with_shard(worker, 5)
            n = len(gen.create_splits(list(entries)))
            if n == 0:
                empties += 1
            else:
                non_empties += 1
                self.assertEqual(n, 1)
        self.assertEqual(empties, 3)
        self.assertEqual(non_empties, 2)

    def test_multi_partition_no_chunk_crosses_partition(self):
        entries = [
            _mock_entry(['p1'], 0, 'f1', 100),
            _mock_entry(['p1'], 0, 'f2', 100),
            _mock_entry(['p2'], 0, 'f3', 100),
            _mock_entry(['p2'], 0, 'f4', 100),
        ]
        gen = _make_generator(seed=0, chunk_size=100)
        splits = gen.create_splits(entries)
        # Each split's underlying files come from one partition only
        for s in splits:
            partitions_in_files = set()
            data_split = s.data_split() if isinstance(s, SlicedSplit) else s
            partitions_in_files.add(tuple(data_split.partition.values))
            self.assertEqual(len(partitions_in_files), 1)

    def test_null_and_non_null_partitions_sort_safely(self):
        # Mixing null and non-null partition values used to raise
        # ``TypeError: '<' not supported between instances of 'NoneType' and 'str'``
        # before _null_safe_partition_key. Validate planning succeeds and
        # both partitions produce splits.
        entries = [
            _mock_entry(['p1'], 0, 'f1', 100),
            _mock_entry([None], 0, 'f2', 100),
            _mock_entry(['p2'], 0, 'f3', 100),
        ]
        gen = _make_generator(seed=1, chunk_size=100)
        splits = gen.create_splits(entries)
        self.assertEqual(len(splits), 3)
        partitions = {tuple(_split_signature(s)[0]) for s in splits}
        self.assertEqual(partitions, {('p1',), ('p2',), (None,)})

    def test_input_order_does_not_affect_output_when_same_files(self):
        """Manifest read parallelism shouldn't bleed through — sorting is internal."""
        a = _mock_entry([], 0, 'f1', 100)
        b = _mock_entry([], 0, 'f2', 100)
        c = _mock_entry([], 0, 'f3', 100)
        gen1 = _make_generator(seed=99, chunk_size=100)
        gen2 = _make_generator(seed=99, chunk_size=100)
        sigs1 = [_split_signature(s) for s in gen1.create_splits([a, b, c])]
        sigs2 = [_split_signature(s) for s in gen2.create_splits([c, a, b])]
        self.assertEqual(sigs1, sigs2)


class ChunkShuffleEndToEndTest(unittest.TestCase):
    """Real append table → with_chunk_shuffle → multiple workers → union == original."""

    @classmethod
    def setUpClass(cls):
        cls.tempdir = tempfile.mkdtemp()
        cls.warehouse = os.path.join(cls.tempdir, 'warehouse')
        cls.catalog = CatalogFactory.create({'warehouse': cls.warehouse})
        cls.catalog.create_database('default', True)

    @classmethod
    def tearDownClass(cls):
        shutil.rmtree(cls.tempdir, ignore_errors=True)

    def _create_append_table(self, name, partition_keys=None):
        pa_schema = pa.schema([
            ('id', pa.int64()),
            ('value', pa.string()),
            ('part', pa.string()),
        ])
        schema = Schema.from_pyarrow_schema(
            pa_schema, partition_keys=partition_keys or [])
        identifier = f'default.{name}'
        self.catalog.create_table(identifier, schema, False)
        return self.catalog.get_table(identifier), pa_schema

    def _write_n_batches(self, table, pa_schema, batches):
        wb = table.new_batch_write_builder()
        for batch in batches:
            tw = wb.new_write()
            tc = wb.new_commit()
            tw.write_arrow(pa.Table.from_pydict(batch, schema=pa_schema))
            tc.commit(tw.prepare_commit())
            tw.close()
            tc.close()

    def test_workers_union_equals_full_table(self):
        table, pa_schema = self._create_append_table('cs_union')
        # 4 commits × 50 rows = 200 rows across several files
        batches = []
        for c in range(4):
            base = c * 50
            batches.append({
                'id': list(range(base, base + 50)),
                'value': [f'v{i}' for i in range(base, base + 50)],
                'part': ['p1'] * 50,
            })
        self._write_n_batches(table, pa_schema, batches)

        read_builder = table.new_read_builder()
        table_read = read_builder.new_read()

        num_workers = 3
        worker_tables = []
        for w in range(num_workers):
            scan = read_builder.new_scan() \
                .with_chunk_shuffle(seed=123, chunk_size=37) \
                .with_shard(w, num_workers)
            splits = scan.plan().splits()
            if splits:
                worker_tables.append(table_read.to_arrow(splits))

        actual = pa.concat_tables(worker_tables).sort_by('id') if worker_tables else None
        self.assertIsNotNone(actual)
        self.assertEqual(actual.num_rows, 200)
        self.assertEqual(actual.column('id').to_pylist(), list(range(200)))

    def test_deterministic_plan_across_calls(self):
        table, pa_schema = self._create_append_table('cs_determinism')
        self._write_n_batches(table, pa_schema, [{
            'id': list(range(100)),
            'value': [f'v{i}' for i in range(100)],
            'part': ['p'] * 100,
        }])

        def plan_files(worker):
            scan = table.new_read_builder().new_scan() \
                .with_chunk_shuffle(seed=42, chunk_size=20) \
                .with_shard(worker, 3)
            return [_split_signature(s) for s in scan.plan().splits()]

        for worker in range(3):
            self.assertEqual(plan_files(worker), plan_files(worker))


class ChunkShuffleCompatibilityTest(unittest.TestCase):
    """Validates the reject-on-incompatible-combination matrix."""

    @classmethod
    def setUpClass(cls):
        cls.tempdir = tempfile.mkdtemp()
        cls.warehouse = os.path.join(cls.tempdir, 'warehouse')
        cls.catalog = CatalogFactory.create({'warehouse': cls.warehouse})
        cls.catalog.create_database('default', True)

    @classmethod
    def tearDownClass(cls):
        shutil.rmtree(cls.tempdir, ignore_errors=True)

    def _append_table(self, name, options=None, partition_keys=None):
        if partition_keys:
            pa_schema = pa.schema([
                ('id', pa.int64()),
                ('value', pa.string()),
                ('part', pa.string()),
            ])
        else:
            pa_schema = pa.schema([('id', pa.int64()), ('value', pa.string())])
        schema = Schema.from_pyarrow_schema(
            pa_schema, partition_keys=partition_keys, options=options or {})
        self.catalog.create_table(f'default.{name}', schema, False)
        return self.catalog.get_table(f'default.{name}')

    def _pk_table(self, name):
        pa_schema = pa.schema([
            pa.field('id', pa.int64(), nullable=False),
            ('value', pa.string()),
        ])
        schema = Schema.from_pyarrow_schema(
            pa_schema, primary_keys=['id'], options={'bucket': '1'})
        self.catalog.create_table(f'default.{name}', schema, False)
        return self.catalog.get_table(f'default.{name}')

    def test_pk_table_rejected(self):
        table = self._pk_table('cs_pk')
        scan = table.new_read_builder().new_scan()
        scan.with_chunk_shuffle(seed=1, chunk_size=100)
        with self.assertRaisesRegex(ValueError, "only supports append tables"):
            scan.plan()

    def test_dv_table_supported(self):
        table = self._append_table('cs_dv', options={'deletion-vectors.enabled': 'true'})
        scan = table.new_read_builder().new_scan()
        scan.with_chunk_shuffle(seed=1, chunk_size=100)
        self.assertEqual(scan.plan().splits(), [])

    def test_with_slice_then_chunk_shuffle_rejected(self):
        table = self._append_table('cs_slice')
        scan = table.new_read_builder().new_scan()
        scan.with_slice(0, 100).with_chunk_shuffle(seed=1, chunk_size=100)
        with self.assertRaisesRegex(ValueError, "with_slice"):
            scan.plan()

    def test_limit_with_chunk_shuffle_rejected(self):
        table = self._append_table('cs_limit')
        scan = table.new_read_builder().with_limit(50).new_scan()
        scan.with_chunk_shuffle(seed=1, chunk_size=100)
        with self.assertRaisesRegex(ValueError, "limit"):
            scan.plan()

    def test_invalid_chunk_size(self):
        table = self._append_table('cs_invalid')
        scan = table.new_read_builder().new_scan()
        with self.assertRaisesRegex(ValueError, "chunk_size"):
            scan.with_chunk_shuffle(seed=1, chunk_size=0)
        with self.assertRaisesRegex(ValueError, "chunk_size"):
            scan.with_chunk_shuffle(seed=1, chunk_size=-5)

    def test_column_predicate_rejected(self):
        # Non-partition predicate would silently shrink effective chunk
        # row counts inside the reader → not allowed.
        table = self._append_table('cs_col_pred', partition_keys=['part'])
        rb = table.new_read_builder()
        col_pred = rb.new_predicate_builder().equal('id', 5)
        rb = rb.with_filter(col_pred)
        scan = rb.new_scan().with_chunk_shuffle(seed=1, chunk_size=10)
        with self.assertRaisesRegex(ValueError, "partition keys"):
            scan.plan()

    def test_partition_predicate_allowed(self):
        # Filter is partition-only → must succeed and read only the
        # matching partition.
        table, pa_schema = self._partitioned_table_with_data('cs_part_pred')

        rb = table.new_read_builder()
        pred = rb.new_predicate_builder().equal('part', 'p1')
        scan = rb.with_filter(pred).new_scan() \
            .with_chunk_shuffle(seed=1, chunk_size=10)
        plan = scan.plan()
        # All splits should be from partition 'p1'
        for split in plan.splits():
            partition_values = split.partition.values
            self.assertEqual(tuple(partition_values), ('p1',))

    def _partitioned_table_with_data(self, name):
        pa_schema = pa.schema([
            ('id', pa.int64()),
            ('value', pa.string()),
            ('part', pa.string()),
        ])
        schema = Schema.from_pyarrow_schema(pa_schema, partition_keys=['part'])
        identifier = f'default.{name}'
        self.catalog.create_table(identifier, schema, False)
        table = self.catalog.get_table(identifier)
        wb = table.new_batch_write_builder()
        for part, ids in [('p1', range(50)), ('p2', range(50, 100))]:
            tw = wb.new_write()
            tc = wb.new_commit()
            tw.write_arrow(pa.Table.from_pydict(
                {'id': list(ids),
                 'value': [f'v{i}' for i in ids],
                 'part': [part] * 50},
                schema=pa_schema,
            ))
            tc.commit(tw.prepare_commit())
            tw.close()
            tc.close()
        return table, pa_schema


class DataEvolutionChunkShuffleAlgoTest(unittest.TestCase):
    """Mock-based tests for the DE chunk slicer."""

    def test_no_entries_returns_empty(self):
        gen = _make_de_generator(seed=1, chunk_size=100)
        self.assertEqual(gen.create_splits([]), [])

    def test_split_by_row_id_with_range_groups_and_merges(self):
        # a[0,4] & c[2,3] overlap -> one merged group [0,4]; b[10,14] -> another. Input unsorted.
        files = [_de_file("a", 0, 5), _de_file("b", 10, 5), _de_file("c", 2, 2)]
        out = DataEvolutionChunkShuffleSplitGenerator._split_by_row_id_with_range(files)
        self.assertEqual([(r.from_, r.to) for r, _ in out], [(0, 4), (10, 14)])
        self.assertEqual([[f.file_name for f in fs] for _, fs in out], [["a", "c"], ["b"]])

    def test_split_by_row_id_with_range_raises_on_missing_first_row_id(self):
        files = [_de_file("a", 0, 5), _de_file("b", None, 5)]
        with self.assertRaisesRegex(ValueError, "requires row tracking"):
            DataEvolutionChunkShuffleSplitGenerator._split_by_row_id_with_range(files)

    def test_full_aligned_groups_one_per_chunk(self):
        # Three commits of 100 rows each → three aligned groups.
        # chunk_size = 100 → 3 chunks, each holding one group whole.
        entries = [
            _mock_de_entry([], 0, 'g0.parquet', 0, 100),
            _mock_de_entry([], 0, 'g1.parquet', 100, 100),
            _mock_de_entry([], 0, 'g2.parquet', 200, 100),
        ]
        gen = _make_de_generator(seed=1, chunk_size=100)
        splits = gen.create_splits(entries)
        self.assertEqual(len(splits), 3)
        for s in splits:
            self.assertIsInstance(s, IndexedSplit)
            self.assertEqual(s.row_count, 100)
            self.assertEqual(len(s.row_ranges()), 1)

    def test_aligned_group_split_across_chunks(self):
        # One 250-row group, chunk_size=100 → 3 chunks (100, 100, 50).
        # All three chunks reference the SAME aligned group's files but
        # each carries a distinct row_range slice.
        entries = [_mock_de_entry([], 0, 'g0.parquet', 1000, 250)]
        gen = _make_de_generator(seed=1, chunk_size=100)
        splits = gen.create_splits(entries)
        self.assertEqual(len(splits), 3)

        # Union of the three chunks' row_ranges must cover the whole group [1000, 1249].
        ranges = []
        for s in splits:
            self.assertIsInstance(s, IndexedSplit)
            ranges.extend((r.from_, r.to) for r in s.row_ranges())
        ranges.sort()
        self.assertEqual(ranges, [(1000, 1099), (1100, 1199), (1200, 1249)])
        total = sum(r[1] - r[0] + 1 for r in ranges)
        self.assertEqual(total, 250)

    def test_chunk_pulls_in_blob_siblings(self):
        # One aligned group with a main parquet and a blob sibling sharing the
        # row_id range. A single chunk must include BOTH files so the reader
        # can union the columns.
        entries = [
            _mock_de_entry([], 0, 'g0.parquet', 0, 100),
            _mock_de_entry([], 0, 'g0.blob', 0, 100),  # .blob ext → is_blob_file
        ]
        gen = _make_de_generator(seed=1, chunk_size=100)
        splits = gen.create_splits(entries)
        self.assertEqual(len(splits), 1)
        files = sorted(f.file_name for f in splits[0].files)
        self.assertEqual(files, ['g0.blob', 'g0.parquet'])

    def test_blob_propagates_when_group_split(self):
        # Same scenario but chunk_size halves the group → the blob sibling
        # must appear in BOTH chunk splits.
        entries = [
            _mock_de_entry([], 0, 'g0.parquet', 0, 100),
            _mock_de_entry([], 0, 'g0.blob', 0, 100),
        ]
        gen = _make_de_generator(seed=1, chunk_size=50)
        splits = gen.create_splits(entries)
        self.assertEqual(len(splits), 2)
        for s in splits:
            files = sorted(f.file_name for f in s.files)
            self.assertEqual(files, ['g0.blob', 'g0.parquet'])

    def test_deletion_vector_uses_anchor_and_slices_by_live_rows(self):
        anchor = _mock_de_entry(
            [],
            0,
            'anchor.parquet',
            100,
            10,
            max_sequence_number=1,
        )
        blob = _mock_de_entry(
            [],
            0,
            'field.blob',
            100,
            10,
            max_sequence_number=2,
        )
        deletion_file = DeletionFile('dv.index', 10, 20, cardinality=4)
        deletion_files_map = {((), 0): {'anchor.parquet': deletion_file}}
        gen = _make_de_generator(
            seed=1,
            chunk_size=3,
            deletion_files_map=deletion_files_map,
        )

        with patch(
            'pypaimon.read.scanner.chunk_shuffle_split_generator.DeletionVector.read',
            return_value=_bitmap_deletion_vector(1, 2, 7, 9),
        ) as read:
            splits = gen.create_splits([blob, anchor])

        ranges = sorted(
            (row_range.from_, row_range.to)
            for split in splits
            for row_range in split.row_ranges()
        )
        self.assertEqual(ranges, [(100, 104), (105, 109)])
        self.assertEqual(
            sorted(split.merged_row_count() for split in splits),
            [3, 3],
        )
        for split in splits:
            self.assertEqual(
                sorted(f.file_name for f in split.files),
                ['anchor.parquet', 'field.blob'],
            )
            deletion_by_name = dict(
                (file.file_name, dv)
                for file, dv in zip(split.files, split.data_deletion_files)
            )
            self.assertEqual(deletion_by_name['anchor.parquet'], deletion_file)
            self.assertIsNone(deletion_by_name['field.blob'])
        read.assert_called_once_with(gen.table.file_io, deletion_file)

    def test_multiple_deletion_vector_groups_share_one_live_row_chunk(self):
        # The groups contribute 3 live rows each despite spanning 5 and 6
        # physical rows, so both should become segments of the same chunk.
        entries = [
            _mock_de_entry([], 0, 'g0.parquet', 100, 5),
            _mock_de_entry([], 0, 'g1.parquet', 200, 6),
        ]
        first_deletion_file = DeletionFile(
            'dv.index',
            10,
            20,
            cardinality=2,
        )
        second_deletion_file = DeletionFile(
            'dv.index',
            30,
            20,
            cardinality=3,
        )
        deletion_files_map = {
            ((), 0): {
                'g0.parquet': first_deletion_file,
                'g1.parquet': second_deletion_file,
            }
        }
        deletion_vectors = {
            first_deletion_file: _bitmap_deletion_vector(1, 3),
            second_deletion_file: _bitmap_deletion_vector(0, 2, 5),
        }
        gen = _make_de_generator(
            seed=1,
            chunk_size=6,
            deletion_files_map=deletion_files_map,
        )

        with patch(
            'pypaimon.read.scanner.chunk_shuffle_split_generator.DeletionVector.read',
            side_effect=lambda _, deletion_file: deletion_vectors[deletion_file],
        ) as read:
            splits = gen.create_splits(entries)

        self.assertEqual(len(splits), 1)
        self.assertIsInstance(splits[0], IndexedSplit)
        self.assertEqual(
            [
                (row_range.from_, row_range.to)
                for row_range in splits[0].row_ranges()
            ],
            [(100, 104), (200, 205)],
        )
        self.assertEqual(
            [file.file_name for file in splits[0].files],
            ['g0.parquet', 'g1.parquet'],
        )
        self.assertEqual(
            splits[0].data_deletion_files,
            [first_deletion_file, second_deletion_file],
        )
        self.assertEqual(splits[0].merged_row_count(), 6)
        self.assertEqual(read.call_count, 2)
        self.assertEqual(
            {mock_call.args[1] for mock_call in read.call_args_list},
            {first_deletion_file, second_deletion_file},
        )

    def test_deterministic_same_seed(self):
        entries = [_mock_de_entry([], 0, f'g{i:02d}.parquet', i * 100, 100) for i in range(20)]
        gen1 = _make_de_generator(seed=42, chunk_size=100)
        gen2 = _make_de_generator(seed=42, chunk_size=100)
        sigs1 = [_split_signature(s) for s in gen1.create_splits(entries)]
        sigs2 = [_split_signature(s) for s in gen2.create_splits(entries)]
        self.assertEqual(sigs1, sigs2)

    def test_different_seed_reorders(self):
        entries = [_mock_de_entry([], 0, f'g{i:02d}.parquet', i * 100, 100) for i in range(50)]
        gen1 = _make_de_generator(seed=1, chunk_size=100)
        gen2 = _make_de_generator(seed=2, chunk_size=100)
        sigs1 = [_split_signature(s) for s in gen1.create_splits(entries)]
        sigs2 = [_split_signature(s) for s in gen2.create_splits(entries)]
        self.assertEqual(sorted(sigs1), sorted(sigs2))
        self.assertNotEqual(sigs1, sigs2)

    def test_input_order_does_not_affect_output(self):
        a = _mock_de_entry([], 0, 'g0.parquet', 0, 100)
        b = _mock_de_entry([], 0, 'g1.parquet', 100, 100)
        c = _mock_de_entry([], 0, 'g2.parquet', 200, 100)
        gen1 = _make_de_generator(seed=99, chunk_size=100)
        gen2 = _make_de_generator(seed=99, chunk_size=100)
        sigs1 = [_split_signature(s) for s in gen1.create_splits([a, b, c])]
        sigs2 = [_split_signature(s) for s in gen2.create_splits([c, a, b])]
        self.assertEqual(sigs1, sigs2)

    def test_shard_round_trip_no_overlap_no_loss(self):
        # 13 aligned groups × 100 rows = 1300 rows. Shard across 4 workers.
        entries = [_mock_de_entry([], 0, f'g{i:02d}.parquet', i * 100, 100) for i in range(13)]
        num_workers = 4

        unsharded = _make_de_generator(seed=7, chunk_size=100).create_splits(list(entries))
        unsharded_sigs = sorted(_split_signature(s) for s in unsharded)

        sharded_sigs = []
        total_rows = 0
        for w in range(num_workers):
            gen = _make_de_generator(seed=7, chunk_size=100)
            gen.with_shard(w, num_workers)
            for s in gen.create_splits(list(entries)):
                sharded_sigs.append(_split_signature(s))
                total_rows += s.row_count
        self.assertEqual(total_rows, 13 * 100)
        # No duplicate splits across workers
        self.assertEqual(len(sharded_sigs), len(set(sharded_sigs)))
        self.assertEqual(sorted(sharded_sigs), unsharded_sigs)

    def test_multi_partition_no_chunk_crosses_partition(self):
        entries = [
            _mock_de_entry(['p1'], 0, 'g0.parquet', 0, 100),
            _mock_de_entry(['p1'], 0, 'g1.parquet', 100, 100),
            _mock_de_entry(['p2'], 0, 'g2.parquet', 200, 100),
            _mock_de_entry(['p2'], 0, 'g3.parquet', 300, 100),
        ]
        gen = _make_de_generator(seed=0, chunk_size=100)
        splits = gen.create_splits(entries)
        for s in splits:
            data_split = s.data_split() if isinstance(s, IndexedSplit) else s
            self.assertEqual(len({tuple(data_split.partition.values)}), 1)

    def test_null_and_non_null_partitions_sort_safely(self):
        # Same null-vs-non-null sort guard, exercised on the DE path.
        entries = [
            _mock_de_entry(['p1'], 0, 'g0.parquet', 0, 100),
            _mock_de_entry([None], 0, 'g1.parquet', 100, 100),
            _mock_de_entry(['p2'], 0, 'g2.parquet', 200, 100),
        ]
        gen = _make_de_generator(seed=1, chunk_size=100)
        splits = gen.create_splits(entries)
        self.assertEqual(len(splits), 3)
        partitions = {_split_signature(s)[0] for s in splits}
        self.assertEqual(partitions, {('p1',), ('p2',), (None,)})


class DataEvolutionChunkShuffleEndToEndTest(unittest.TestCase):
    """Real DE table → with_chunk_shuffle → multi-worker → union == full table."""

    @classmethod
    def setUpClass(cls):
        cls.tempdir = tempfile.mkdtemp()
        cls.warehouse = os.path.join(cls.tempdir, 'warehouse')
        cls.catalog = CatalogFactory.create({'warehouse': cls.warehouse})
        cls.catalog.create_database('default', True)

    @classmethod
    def tearDownClass(cls):
        shutil.rmtree(cls.tempdir, ignore_errors=True)

    def _create_de_table(self, name, deletion_vectors_enabled=False):
        pa_schema = pa.schema([
            ('id', pa.int32()),
            ('value', pa.string()),
            ('payload', pa.large_binary()),
        ])
        options = {
            'row-tracking.enabled': 'true',
            'data-evolution.enabled': 'true',
            'blob.target-file-size': '1 b',
        }
        if deletion_vectors_enabled:
            options['deletion-vectors.enabled'] = 'true'
        schema = Schema.from_pyarrow_schema(
            pa_schema,
            options=options,
        )
        identifier = f'default.{name}'
        self.catalog.create_table(identifier, schema, False)
        return self.catalog.get_table(identifier), pa_schema

    @staticmethod
    def _payloads(ids):
        return [f'payload-{i:03d}'.encode('utf-8') for i in ids]

    def _commit_full_rows(self, table, pa_schema, ids):
        wb = table.new_batch_write_builder()
        tw = wb.new_write()
        tc = wb.new_commit()
        tw.write_arrow(pa.Table.from_pydict(
            {
                'id': ids,
                'value': [f'v{i}' for i in ids],
                'payload': self._payloads(ids),
            },
            schema=pa_schema))
        commit_messages = tw.prepare_commit()
        tc.commit(commit_messages)
        tw.close()
        tc.close()
        return commit_messages

    @staticmethod
    def _delete_by_row_id(table, row_ids):
        write_builder = table.new_batch_write_builder()
        commit_messages = write_builder.new_update().delete_by_row_id(row_ids)
        table_commit = write_builder.new_commit()
        table_commit.commit(commit_messages)
        table_commit.close()

    def _assert_commit_has_main_and_multiple_blob_files(self, commit_messages):
        all_files = [f for msg in commit_messages for f in msg.new_files]
        main_files = [f for f in all_files if not DataFileMeta.is_blob_file(f.file_name)]
        blob_files = [f for f in all_files if DataFileMeta.is_blob_file(f.file_name)]
        self.assertGreaterEqual(len(main_files), 1)
        self.assertGreater(
            len(blob_files), 1,
            "DE chunk-shuffle tests should exercise one row-id group with multiple blob files",
        )

    def _assert_splits_include_blob_files(self, splits):
        self.assertGreater(len(splits), 0)
        for split in splits:
            data_split = split.data_split() if isinstance(split, IndexedSplit) else split
            blob_files = [
                f for f in data_split.files
                if DataFileMeta.is_blob_file(f.file_name)
            ]
            self.assertGreater(
                len(blob_files), 0,
                "Each DE chunk should keep blob sidecar files with its aligned row-id group",
            )

    def test_workers_union_equals_full_table(self):
        table, pa_schema = self._create_de_table('cs_de_union')
        # 4 commits → 4 aligned groups. Each group has one normal file and
        # multiple blob sidecar files because blob.target-file-size is 1 byte.
        for c in range(4):
            base = c * 50
            commit_messages = self._commit_full_rows(
                table, pa_schema, list(range(base, base + 50)))
            self._assert_commit_has_main_and_multiple_blob_files(commit_messages)

        read_builder = table.new_read_builder()
        table_read = read_builder.new_read()

        num_workers = 3
        worker_tables = []
        for w in range(num_workers):
            scan = read_builder.new_scan() \
                .with_chunk_shuffle(seed=123, chunk_size=37) \
                .with_shard(w, num_workers)
            splits = scan.plan().splits()
            if splits:
                self._assert_splits_include_blob_files(splits)
                worker_tables.append(table_read.to_arrow(splits))

        actual = pa.concat_tables(worker_tables).sort_by('id')
        self.assertEqual(actual.num_rows, 200)
        self.assertEqual(actual.column('id').to_pylist(), list(range(200)))
        self.assertEqual(actual.column('payload').to_pylist(), self._payloads(range(200)))

    def test_deterministic_plan_across_calls(self):
        table, pa_schema = self._create_de_table('cs_de_determinism')
        for c in range(3):
            base = c * 40
            commit_messages = self._commit_full_rows(
                table, pa_schema, list(range(base, base + 40)))
            self._assert_commit_has_main_and_multiple_blob_files(commit_messages)

        def plan_sigs(worker):
            scan = table.new_read_builder().new_scan() \
                .with_chunk_shuffle(seed=42, chunk_size=15) \
                .with_shard(worker, 4)
            splits = scan.plan().splits()
            if splits:
                self._assert_splits_include_blob_files(splits)
            return [_split_signature(s) for s in splits]

        for worker in range(4):
            self.assertEqual(plan_sigs(worker), plan_sigs(worker))

    def test_python_delete_then_chunk_shuffle_read(self):
        table, pa_schema = self._create_de_table(
            'cs_de_python_delete',
            deletion_vectors_enabled=True,
        )
        commit_messages = self._commit_full_rows(
            table,
            pa_schema,
            list(range(12)),
        )
        self._assert_commit_has_main_and_multiple_blob_files(commit_messages)
        self._delete_by_row_id(table, [0, 2, 3, 7, 11])

        read_builder = table.new_read_builder()
        splits = read_builder.new_scan() \
            .with_chunk_shuffle(seed=42, chunk_size=3) \
            .plan() \
            .splits()
        self._assert_splits_include_blob_files(splits)

        table_read = read_builder.new_read()
        rows_per_split = [
            table_read.to_arrow([split]).num_rows
            for split in splits
        ]
        self.assertEqual(sorted(rows_per_split), [1, 3, 3])

        actual = table_read.to_arrow(splits).sort_by('id')
        expected_ids = [1, 4, 5, 6, 8, 9, 10]
        self.assertEqual(actual.column('id').to_pylist(), expected_ids)
        self.assertEqual(
            actual.column('payload').to_pylist(),
            self._payloads(expected_ids),
        )


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