ray-project--ray/python/ray/data/tests/test_hash_shuffle_aggregator.py

"""Unit tests for HashShuffleAggregator."""

from typing import Dict, Iterator, List

import pyarrow as pa
import pytest

from ray.data._internal.arrow_ops import transform_pyarrow
from ray.data._internal.execution.operators.hash_shuffle import (
    HashShuffleAggregator,
    ShuffleAggregation,
)
from ray.data._internal.planner.exchange.sort_task_spec import SortKey
from ray.data.block import Block

# Access underlying class for direct instantiation (bypassing Ray actor)
_HashShuffleAggregatorClass = HashShuffleAggregator.__ray_actor_class__


def make_block(n: int = 10, offset: int = 0) -> pa.Table:
    return pa.table({"x": list(range(offset, offset + n))})


def split_block(block: pa.Table, chunk_size: int) -> List[pa.Table]:
    """Split block into chunks of given size."""
    return [block.slice(i, chunk_size) for i in range(0, block.num_rows, chunk_size)]


class MockCompactingAggregation(ShuffleAggregation):
    """Tracks compact/finalize calls and input blocks."""

    def __init__(self):
        self.compact_calls: List[int] = []
        self.finalize_input: Dict[int, List[Block]] = {}

    @classmethod
    def is_compacting(cls):
        return True

    def compact(self, shards: List[Block]) -> Block:
        self.compact_calls.append(len(shards))
        return pa.concat_tables(shards) if shards else make_block(0)

    def finalize(self, shards: Dict[int, List[Block]]) -> Iterator[Block]:
        self.finalize_input = dict(shards)
        blocks = [b for bs in shards.values() for b in bs]
        yield pa.concat_tables(blocks) if blocks else make_block(0)


class MockNonCompactingAggregation(ShuffleAggregation):
    """Tracks finalize input blocks."""

    def __init__(self):
        self.finalize_input: Dict[int, List[Block]] = {}

    @classmethod
    def is_compacting(cls):
        return False

    def compact(self, shards: List[Block]) -> Block:
        raise RuntimeError("Should not be called")

    def finalize(self, shards: Dict[int, List[Block]]) -> Iterator[Block]:
        self.finalize_input = dict(shards)
        blocks = [b for bs in shards.values() for b in bs]
        yield pa.concat_tables(blocks) if blocks else make_block(0)


class TestHashShuffleAggregator:
    def test_compacting_workflow(self, ray_start_regular_shared):
        """Tests compaction triggers, threshold doubling, multi-partition/sequence."""
        agg = MockCompactingAggregation()
        aggregator = _HashShuffleAggregatorClass(
            aggregator_id=0,
            num_input_seqs=2,
            target_partition_ids=[0, 1, 2],
            agg_factory=lambda: agg,
            target_max_block_size=None,
            min_max_shards_compaction_thresholds=(3, 2000),
        )

        # Pre-generate blocks: split a 100-row block into 10 chunks of 10 rows
        full_block = make_block(80)
        input_seq0_part0 = split_block(full_block, 10)

        def get_compaction_thresholds():
            """Helper to extract compaction thresholds from partition buckets."""
            # Thresholds are now per-partition in PartitionBucket
            return {
                part_id: bucket.compaction_threshold
                for part_id, bucket in aggregator._input_seq_partition_buckets[
                    0
                ].items()
                if bucket.compaction_threshold is not None
            }

        # Submit 2 blocks (below threshold=3) - no compaction
        for b in input_seq0_part0[:2]:
            aggregator.submit(0, 0, b)
        assert agg.compact_calls == []
        assert get_compaction_thresholds() == {0: 3, 1: 3, 2: 3}

        # Submit 3rd block - triggers compaction, threshold doubles
        aggregator.submit(0, 0, input_seq0_part0[2])
        assert agg.compact_calls == [3]
        assert get_compaction_thresholds() == {0: 6, 1: 3, 2: 3}

        # Submit 5 more (queue: 1+5=6) - triggers at new threshold
        for b in input_seq0_part0[3:8]:
            aggregator.submit(0, 0, b)

        assert agg.compact_calls == [3, 6]
        assert get_compaction_thresholds() == {0: 12, 1: 3, 2: 3}

        # Partition 1 has independent threshold (starts at 3)
        for b in split_block(make_block(30, offset=1000), 10):
            aggregator.submit(0, 1, b)

        assert agg.compact_calls == [3, 6, 3]
        assert get_compaction_thresholds() == {0: 12, 1: 6, 2: 3}

        # Multiple sequences (join scenario) - seq_id=1 for partition 0
        input_seq1_part0 = split_block(make_block(20, offset=2000), 10)
        for b in input_seq1_part0:
            aggregator.submit(1, 0, b)

        # Finalize partition 0 - receives blocks from both sequences
        results = list(aggregator.finalize(0))
        block, metadata = results
        assert len(agg.finalize_input) == 2  # dict with 2 sequences

        # Verify output equals concatenation of seq0 (first 8 chunks) + seq1
        expected = transform_pyarrow.sort(
            pa.concat_tables(tables=[*input_seq0_part0, *input_seq1_part0]),
            sort_key=SortKey("x"),
        )
        assert transform_pyarrow.sort(block, sort_key=SortKey("x")) == expected

        # Empty partition
        results = list(aggregator.finalize(2))
        assert results[0] == make_block(0)

    def test_non_compacting_workflow(self, ray_start_regular_shared):
        """Tests non-compacting aggregation with and without block splitting."""
        # Without splitting
        full_block = make_block(50)
        input_seq = split_block(full_block, 10)

        aggregator = _HashShuffleAggregatorClass(
            aggregator_id=1,
            num_input_seqs=1,
            target_partition_ids=[0],
            agg_factory=MockNonCompactingAggregation,
            target_max_block_size=None,
        )
        for b in input_seq:
            aggregator.submit(0, 0, b)

        results = list(aggregator.finalize(0))
        block, metadata = results
        assert block == full_block

        # With splitting - output blocks should reconstruct to original
        full_block = make_block(500)
        input_seq = split_block(full_block, 100)

        aggregator = _HashShuffleAggregatorClass(
            aggregator_id=2,
            num_input_seqs=1,
            target_partition_ids=[0],
            agg_factory=MockNonCompactingAggregation,
            target_max_block_size=50,
        )
        for b in input_seq:
            aggregator.submit(0, 0, b)

        results = list(aggregator.finalize(0))
        output_blocks = [results[i] for i in range(0, len(results), 2)]
        assert pa.concat_tables(output_blocks) == full_block


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", __file__]))