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

from typing import Optional

import numpy as np
import pandas as pd
import pytest
from packaging.version import parse as parse_version

import ray
from ray.data._internal.logical.operators import JoinType
from ray.data._internal.util import MiB, rows_same
from ray.data._internal.utils.arrow_utils import get_pyarrow_version
from ray.data.context import DataContext
from ray.data.dataset import Dataset
from ray.exceptions import RayTaskError
from ray.tests.conftest import *  # noqa


@pytest.mark.parametrize(
    "num_rows_left,num_rows_right,partition_size_hint",
    [
        (32, 32, 1 * MiB),
        (32, 16, None),
        (16, 32, None),
        # "Degenerate" cases with mostly empty partitions
        (32, 1, None),
        (1, 32, None),
    ],
)
def test_simple_inner_join(
    ray_start_regular_shared_2_cpus,
    num_rows_left: int,
    num_rows_right: int,
    partition_size_hint: Optional[int],
):
    # NOTE: We override max-block size to make sure that in cases when a partition
    #       size hint is not provided, we're not over-estimating amount of memory
    #       required for the aggregators
    DataContext.get_current().target_max_block_size = 1 * MiB

    doubles = ray.data.range(num_rows_left).map(
        lambda row: {"id": row["id"], "double": int(row["id"]) * 2}
    )

    squares = ray.data.range(num_rows_right).map(
        lambda row: {"id": row["id"], "square": int(row["id"]) ** 2}
    )

    doubles_pd = doubles.to_pandas()
    squares_pd = squares.to_pandas()

    # Join using Pandas (to assert against)
    expected_pd = doubles_pd.join(squares_pd.set_index("id"), on="id", how="inner")
    expected_pd_sorted = expected_pd.sort_values(by=["id"]).reset_index(drop=True)

    # Join using Ray Data
    joined: Dataset = doubles.join(
        squares,
        join_type="inner",
        num_partitions=16,
        on=("id",),
        partition_size_hint=partition_size_hint,
    )

    # TODO use native to_pandas() instead
    joined_pd = pd.DataFrame(joined.take_all())

    # Sort resulting frame and reset index (to be able to compare with expected one)
    joined_pd_sorted = joined_pd.sort_values(by=["id"]).reset_index(drop=True)
    expected_pd_sorted = expected_pd_sorted.astype(joined_pd_sorted.dtypes.to_dict())

    pd.testing.assert_frame_equal(expected_pd_sorted, joined_pd_sorted)


@pytest.mark.parametrize(
    "join_type",
    [
        "left_outer",
        "right_outer",
        "left_semi",
        "right_semi",
        "left_anti",
        "right_anti",
    ],
)
@pytest.mark.parametrize(
    "num_rows_left,num_rows_right",
    [
        (32, 32),
        (32, 16),
        (16, 32),
        # "Degenerate" cases with mostly empty partitions
        (1, 32),
        (32, 1),
    ],
)
def test_simple_left_right_outer_semi_anti_join(
    ray_start_regular_shared_2_cpus,
    join_type,
    num_rows_left,
    num_rows_right,
):
    # NOTE: We override max-block size to make sure that in cases when a partition
    #       size hint is not provided, we're not over-estimating amount of memory
    #       required for the aggregators
    DataContext.get_current().target_max_block_size = 1 * MiB

    doubles = ray.data.range(num_rows_left).map(
        lambda row: {"id": row["id"], "double": int(row["id"]) * 2}
    )

    squares = ray.data.range(num_rows_right).map(
        lambda row: {"id": row["id"], "square": int(row["id"]) ** 2}
    )

    doubles_pd = doubles.to_pandas()
    squares_pd = squares.to_pandas()

    # Join using Pandas (to assert against)
    if join_type == "left_outer":
        expected_pd = doubles_pd.join(
            squares_pd.set_index("id"), on="id", how="left"
        ).reset_index(drop=True)
    elif join_type == "right_outer":
        expected_pd = (
            doubles_pd.set_index("id")
            .join(squares_pd, on="id", how="right")
            .reset_index(drop=True)
        )
    elif join_type == "left_semi":
        # Left semi: left rows that have matches in right (left columns only)
        merged = doubles_pd.merge(squares_pd, on="id", how="inner")
        expected_pd = merged[["id", "double"]].drop_duplicates().reset_index(drop=True)
    elif join_type == "right_semi":
        # Right semi: right rows that have matches in left (right columns only)
        merged = doubles_pd.merge(squares_pd, on="id", how="inner")
        expected_pd = merged[["id", "square"]].drop_duplicates().reset_index(drop=True)
    elif join_type == "left_anti":
        # Left anti: left rows that don't have matches in right
        merged = doubles_pd.merge(squares_pd, on="id", how="left", indicator=True)
        expected_pd = merged[merged["_merge"] == "left_only"][
            ["id", "double"]
        ].reset_index(drop=True)
    elif join_type == "right_anti":
        # Right anti: right rows that don't have matches in left
        merged = doubles_pd.merge(squares_pd, on="id", how="right", indicator=True)
        expected_pd = merged[merged["_merge"] == "right_only"][
            ["id", "square"]
        ].reset_index(drop=True)
    else:
        raise ValueError(f"Unsupported join type: {join_type}")

    # Join using Ray Data
    joined: Dataset = doubles.join(
        squares,
        join_type=join_type,
        num_partitions=16,
        on=("id",),
    )

    joined_pd = pd.DataFrame(joined.take_all())

    # Handle empty results from Ray Data which may not preserve schema
    if len(joined_pd) == 0 and len(expected_pd) == 0:
        pass
    else:
        # Sort resulting frame and reset index (to be able to compare with expected one)
        joined_pd_sorted = joined_pd.sort_values(by=["id"]).reset_index(drop=True)
        expected_pd_sorted = expected_pd.sort_values(by=["id"]).reset_index(drop=True)
        expected_pd_sorted = expected_pd_sorted.astype(
            joined_pd_sorted.dtypes.to_dict()
        )

        pd.testing.assert_frame_equal(expected_pd_sorted, joined_pd_sorted)


@pytest.mark.parametrize(
    "num_rows_left,num_rows_right",
    [
        (32, 32),
        (32, 16),
        (16, 32),
        # # "Degenerate" cases with mostly empty partitions
        (1, 32),
        (32, 1),
    ],
)
def test_simple_full_outer_join(
    ray_start_regular_shared_2_cpus,
    num_rows_left,
    num_rows_right,
):
    # NOTE: We override max-block size to make sure that in cases when a partition
    #       size hint is not provided, we're not over-estimating amount of memory
    #       required for the aggregators
    DataContext.get_current().target_max_block_size = 1 * MiB

    doubles = ray.data.range(num_rows_left).map(
        lambda row: {"id": row["id"], "double": int(row["id"]) * 2}
    )

    squares = ray.data.range(num_rows_right).map(
        lambda row: {"id": row["id"] + num_rows_left, "square": int(row["id"]) ** 2}
    )

    doubles_pd = doubles.to_pandas()
    squares_pd = squares.to_pandas()

    # Join using Pandas (to assert against)
    expected_pd = doubles_pd.join(
        squares_pd.set_index("id"), on="id", how="outer"
    ).reset_index(drop=True)

    # Join using Ray Data
    joined: Dataset = doubles.join(
        squares,
        join_type="full_outer",
        num_partitions=16,
        on=("id",),
        # NOTE: We override this to reduce hardware requirements
        #       for every aggregator (by default requiring 1 logical CPU)
        aggregator_ray_remote_args={"num_cpus": 0.01},
    )

    joined_pd = pd.DataFrame(joined.take_all())

    # Handle empty results from Ray Data which may not preserve schema
    if len(joined_pd) == 0 and len(expected_pd) == 0:
        pass
    else:
        # Sort resulting frame and reset index (to be able to compare with expected one)
        joined_pd_sorted = joined_pd.sort_values(by=["id"]).reset_index(drop=True)
        expected_pd_sorted = expected_pd.sort_values(by=["id"]).reset_index(drop=True)
        expected_pd_sorted = expected_pd_sorted.astype(
            joined_pd_sorted.dtypes.to_dict()
        )

        pd.testing.assert_frame_equal(expected_pd_sorted, joined_pd_sorted)


@pytest.mark.parametrize("left_suffix", [None, "_left"])
@pytest.mark.parametrize("right_suffix", [None, "_right"])
def test_simple_self_join(ray_start_regular_shared_2_cpus, left_suffix, right_suffix):
    # NOTE: We override max-block size to make sure that in cases when a partition
    #       size hint is not provided, we're not over-estimating amount of memory
    #       required for the aggregators
    DataContext.get_current().target_max_block_size = 1 * MiB

    doubles = ray.data.range(100).map(
        lambda row: {"id": row["id"], "double": int(row["id"]) * 2}
    )

    doubles_pd = doubles.to_pandas()

    # Self-join
    joined: Dataset = doubles.join(
        doubles,
        join_type="inner",
        num_partitions=16,
        on=("id",),
        left_suffix=left_suffix,
        right_suffix=right_suffix,
        # NOTE: We override this to reduce hardware requirements
        #       for every aggregator (by default requiring 1 logical CPU)
        aggregator_ray_remote_args={"num_cpus": 0.01},
    )

    if left_suffix is None and right_suffix is None:
        with pytest.raises(RayTaskError) as exc_info:
            joined.count()

        assert "Left and right columns suffixes cannot be both None" in str(
            exc_info.value.cause
        )
    else:
        joined_pd = joined.to_pandas()

        # Join using Pandas (to assert against)
        expected_pd = doubles_pd.join(
            doubles_pd.set_index("id"),
            on="id",
            how="inner",
            lsuffix=left_suffix,
            rsuffix=right_suffix,
        ).reset_index(drop=True)

        assert rows_same(expected_pd, joined_pd), "Expected contents to be same"


def test_invalid_join_config(ray_start_regular_shared_2_cpus):
    ds = ray.data.range(32)

    with pytest.raises(ValueError) as exc_info:
        ds.join(
            ds,
            "inner",
            num_partitions=16,
            on="id",  # has to be tuple/list
            validate_schemas=True,
        )

    assert str(exc_info.value) == "Expected tuple or list as `on` (got str)"

    with pytest.raises(ValueError) as exc_info:
        ds.join(
            ds,
            "inner",
            num_partitions=16,
            on=("id",),
            right_on="id",  # has to be tuple/list
            validate_schemas=True,
        )

    assert str(exc_info.value) == "Expected tuple or list as `right_on` (got str)"


@pytest.mark.parametrize("join_type", [jt for jt in JoinType])  # noqa: C416
def test_invalid_join_not_matching_key_columns(
    ray_start_regular_shared_2_cpus, join_type
):
    # Case 1: Check on missing key column
    empty_ds = ray.data.range(0)

    non_empty_ds = ray.data.range(32)

    with pytest.raises(ValueError) as exc_info:
        empty_ds.join(
            non_empty_ds,
            join_type,
            num_partitions=16,
            on=("id",),
            validate_schemas=True,
        )

    assert (
        str(exc_info.value)
        == "Key columns are expected to be present and have the same types in both "
        "left and right operands of the join operation: left has None, but right "
        "has Column  Type\n------  ----\nid      int64"
    )

    # Case 2: Check mismatching key column
    id_int_type_ds = ray.data.range(32).map(lambda row: {"id": int(row["id"])})

    id_float_type_ds = ray.data.range(32).map(lambda row: {"id": float(row["id"])})

    with pytest.raises(ValueError) as exc_info:
        id_int_type_ds.join(
            id_float_type_ds,
            join_type,
            num_partitions=16,
            on=("id",),
            validate_schemas=True,
        )

    assert (
        str(exc_info.value)
        == "Key columns are expected to be present and have the same types in both "
        "left and right operands of the join operation: left has "
        "Column  Type\n------  ----\nid      int64, but right has "
        "Column  Type\n------  ----\nid      double"
    )


@pytest.mark.parametrize("join_type", ["left_anti", "right_anti"])
def test_anti_join_no_matches(
    ray_start_regular_shared_2_cpus,
    join_type,
):
    """Test anti-join when there are no matches - should return all rows from respective side"""
    DataContext.get_current().target_max_block_size = 1 * MiB

    doubles = ray.data.range(32).map(
        lambda row: {"id": row["id"], "double": int(row["id"]) * 2}
    )

    # Create squares with completely different keys
    squares = ray.data.range(32).map(
        lambda row: {"id": row["id"] + 100, "square": int(row["id"]) ** 2}
    )

    # Anti-join should return all rows from respective side
    joined: Dataset = doubles.join(
        squares,
        join_type=join_type,
        num_partitions=4,
        on=("id",),
    )

    joined_pd = pd.DataFrame(joined.take_all())

    if join_type == "left_anti":
        expected_pd = doubles.to_pandas()
    else:  # right_anti
        expected_pd = squares.to_pandas()

    # Should get all rows from the respective table
    joined_pd_sorted = joined_pd.sort_values(by=["id"]).reset_index(drop=True)
    expected_pd_sorted = expected_pd.sort_values(by=["id"]).reset_index(drop=True)
    expected_pd_sorted = expected_pd_sorted.astype(joined_pd_sorted.dtypes.to_dict())

    pd.testing.assert_frame_equal(expected_pd_sorted, joined_pd_sorted)


@pytest.mark.parametrize("join_type", ["left_anti", "right_anti"])
def test_anti_join_all_matches(
    ray_start_regular_shared_2_cpus,
    join_type,
):
    """Test anti-join when all rows match - should return empty result"""
    DataContext.get_current().target_max_block_size = 1 * MiB

    doubles = ray.data.range(32).map(
        lambda row: {"id": row["id"], "double": int(row["id"]) * 2}
    )

    squares = ray.data.range(32).map(
        lambda row: {"id": row["id"], "square": int(row["id"]) ** 2}
    )

    # Anti-join should return no rows since all keys match
    joined: Dataset = doubles.join(
        squares,
        join_type=join_type,
        num_partitions=4,
        on=("id",),
    )

    joined_pd = pd.DataFrame(joined.take_all())

    # Should get empty result
    assert len(joined_pd) == 0


@pytest.mark.parametrize("join_type", ["left_anti", "right_anti"])
def test_anti_join_multi_key(
    ray_start_regular_shared_2_cpus,
    join_type,
):
    """Test anti-join with multiple join keys"""
    DataContext.get_current().target_max_block_size = 1 * MiB

    # Create left dataset using ray.data.range for consistency
    left_ds = ray.data.range(32).map(
        lambda row: {
            "id": row["id"],
            "oddness": row["id"] % 2,  # Even
            "10x": row["id"] * 10,
        }
    )

    # Create right dataset with partial matches (16 vs 32 for partial overlap)
    right_ds = ray.data.range(16).map(
        lambda row: {
            "id": row["id"] % 2,
            "oddness": row["id"] % 2 + 1,  # odd
            "100x": row["id"] * 100,
        }
    )

    # Anti-join should return rows that don't have matching key1,key2 in the other dataset
    joined: Dataset = left_ds.join(
        right_ds,
        join_type=join_type,
        num_partitions=4,
        on=("id", "oddness"),
    )

    joined_pd = pd.DataFrame(joined.take_all())

    # Create expected data for pandas comparison
    left_pd = left_ds.to_pandas()
    right_pd = right_ds.to_pandas()

    # Calculate expected result using pandas
    if join_type == "left_anti":
        expected_cols = ["id", "oddness", "10x"]

        merged = left_pd.merge(
            right_pd, on=["id", "oddness"], how="left", indicator=True
        )
        expected_pd = merged[merged["_merge"] == "left_only"][expected_cols]
    else:
        expected_cols = ["id", "oddness", "100x"]

        merged = left_pd.merge(
            right_pd, on=["id", "oddness"], how="right", indicator=True
        )
        expected_pd = merged[merged["_merge"] == "right_only"][expected_cols]

    # Sort resulting frames and reset index (to be able to compare with expected one)
    expected_pd_sorted = expected_pd.sort_values(by=expected_cols).reset_index(
        drop=True
    )
    joined_pd_sorted = joined_pd.sort_values(by=expected_cols).reset_index(drop=True)
    expected_pd_sorted = expected_pd_sorted.astype(joined_pd_sorted.dtypes.to_dict())

    pd.testing.assert_frame_equal(expected_pd_sorted, joined_pd_sorted)


# Helper functions to reduce test code bloat
def _assert_columns_match(result, expected_columns):
    """Assert that result has the expected column schema."""
    actual_columns = set(result[0].keys())
    assert actual_columns == expected_columns


def _assert_list_values(result_by_id, expected_values):
    """Assert list column values match expected values."""
    for row_id, expected_list in expected_values.items():
        assert result_by_id[row_id]["list_col"] == expected_list


def _assert_tensor_values(result_by_id, expected_values):
    """Assert tensor column values match expected tensor data."""
    for row_id, expected_tensor in expected_values.items():
        assert np.array_equal(result_by_id[row_id]["tensor_col"], expected_tensor)


def _assert_none_values(result_by_id, none_checks):
    """Assert that specified columns are None for specified row IDs."""
    for row_id, columns in none_checks.items():
        for column in columns:
            assert result_by_id[row_id][column] is None


def _assert_scalar_values(result_by_id, expected_values):
    """Assert scalar column values match expected values."""
    for row_id, column_values in expected_values.items():
        for column, expected_value in column_values.items():
            assert result_by_id[row_id][column] == expected_value


def test_should_not_index_empty_schema_tables():
    import pyarrow as pa

    from ray.data._internal.execution.operators.join import _should_index_side

    supported_table = pa.table({"id": pa.array([1])})
    unsupported_table = pa.table({"unsupported": pa.array([[1]])})
    empty_schema_table = pa.table({})

    assert not _should_index_side(
        "left", empty_schema_table, unsupported_table, JoinType.LEFT_OUTER
    )
    assert not _should_index_side(
        "left", supported_table, empty_schema_table, JoinType.LEFT_OUTER
    )
    assert _should_index_side(
        "left", supported_table, unsupported_table, JoinType.LEFT_OUTER
    )


@pytest.mark.skipif(
    get_pyarrow_version() < parse_version("10.0.0"),
    reason="""Joins use empty arrays with type coercion. This pyarrow
    version does not support type coercion of extension types, which
    are needed for tensors.""",
)
@pytest.mark.parametrize(
    "join_type",
    [
        "inner",
        "left_outer",
        "right_outer",
        "full_outer",
        "left_semi",
        "right_semi",
        "left_anti",
        "right_anti",
    ],
)
def test_join_with_unjoinable_non_key_columns(
    ray_start_regular_shared_2_cpus, join_type
):
    """Test that joins work correctly when non-key columns have unjoinable types."""
    # Left dataset with joinable key but unjoinable non-key columns

    # Create test data - centralized for clarity and maintainability
    list_data = [
        [1, 2, 3],  # list for id=0
        [4, 5, 6],  # list for id=1
        [7, 8, 9],  # list for id=2
    ]

    tensor_data = [
        np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32),  # 2x2 tensor for id=0
        np.array([[5.0, 6.0], [7.0, 8.0]], dtype=np.float32),  # 2x2 tensor for id=1
        np.array([[9.0, 10.0], [11.0, 12.0]], dtype=np.float32),  # 2x2 tensor for id=2
    ]

    scalar_data = ["a", "b", "c"]  # scalar data for id=0,1,2

    left_ds = ray.data.from_items(
        [
            {
                "id": 0,
                "list_col": list_data[0],
                "tensor_col": tensor_data[0],
                "data": scalar_data[0],
            },
            {
                "id": 1,
                "list_col": list_data[1],
                "tensor_col": tensor_data[1],
                "data": scalar_data[1],
            },
            {
                "id": 2,
                "list_col": list_data[2],
                "tensor_col": tensor_data[2],
                "data": scalar_data[2],
            },
        ]
    )

    # Right dataset with joinable key and columns
    # ids: 0, 1, 3 (so id=2 from left won't match, id=3 from right won't match)
    right_ds = ray.data.from_items(
        [
            {"id": 0, "value": "x", "score": 10},
            {"id": 1, "value": "y", "score": 20},
            {"id": 3, "value": "z", "score": 30},
        ]
    )

    # Verify the join worked and includes unjoinable columns
    joined = left_ds.join(right_ds, join_type=join_type, on=("id",), num_partitions=2)
    result = joined.take_all()
    result_by_id = {row["id"]: row for row in result}

    # Basic validation - join should succeed with unjoinable non-key columns
    if join_type == "inner":
        # Should have 2 rows (id=0 and id=1 match)
        assert len(result) == 2
        # Verify unjoinable columns are preserved
        _assert_list_values(result_by_id, {i: list_data[i] for i in [0, 1]})
        _assert_tensor_values(result_by_id, {i: tensor_data[i] for i in [0, 1]})

    elif join_type == "left_outer":
        # Should have 3 rows (all from left: id=0, 1, 2)
        assert len(result) == 3
        # All left unjoinable columns preserved
        _assert_list_values(result_by_id, {i: list_data[i] for i in [0, 1, 2]})
        _assert_tensor_values(result_by_id, {i: tensor_data[i] for i in [0, 1, 2]})
        # Unmatched left row (id=2) should have None for right columns
        _assert_none_values(result_by_id, {2: ["value"]})

    elif join_type == "right_outer":
        # Should have 3 rows (all from right: id=0, 1, 3)
        assert len(result) == 3
        # Matched rows should have unjoinable columns from left
        _assert_list_values(result_by_id, {i: list_data[i] for i in [0, 1]})
        _assert_tensor_values(result_by_id, {i: tensor_data[i] for i in [0, 1]})
        _assert_scalar_values(result_by_id, {3: {"value": "z"}})
        # Unmatched right row (id=3) should have None for left unjoinable columns
        _assert_none_values(result_by_id, {3: ["list_col", "tensor_col"]})

    elif join_type == "full_outer":
        # Should have 4 rows (all from both sides: id=0, 1, 2, 3)
        assert len(result) == 4
        # Matched rows (id=0, 1) should have data from both sides
        _assert_list_values(result_by_id, {i: list_data[i] for i in [0, 1, 2]})
        _assert_tensor_values(result_by_id, {i: tensor_data[i] for i in [0, 1, 2]})
        _assert_scalar_values(
            result_by_id,
            {
                0: {"value": "x"},
                1: {"value": "y"},
                2: {"data": scalar_data[2]},
                3: {"value": "z", "score": 30},
            },
        )
        # Unmatched rows should have None for columns from the other side
        _assert_none_values(
            result_by_id, {2: ["value", "score"], 3: ["list_col", "tensor_col", "data"]}
        )

    elif join_type == "left_semi":
        # Should return left rows that have matches in right (id=0, 1)
        assert len(result) == 2
        _assert_columns_match(result, {"id", "list_col", "tensor_col", "data"})
        _assert_list_values(result_by_id, {i: list_data[i] for i in [0, 1]})
        _assert_tensor_values(result_by_id, {i: tensor_data[i] for i in [0, 1]})

    elif join_type == "left_anti":
        # Should return left rows that DON'T have matches in right (id=2)
        assert len(result) == 1
        _assert_columns_match(result, {"id", "list_col", "tensor_col", "data"})
        _assert_list_values(result_by_id, {2: list_data[2]})
        _assert_tensor_values(result_by_id, {2: tensor_data[2]})
        _assert_scalar_values(result_by_id, {2: {"data": scalar_data[2]}})

    elif join_type == "right_semi":
        # Should return right rows that have matches in left (id=0, 1)
        assert len(result) == 2
        _assert_columns_match(result, {"id", "value", "score"})
        _assert_scalar_values(result_by_id, {0: {"value": "x"}, 1: {"value": "y"}})

    elif join_type == "right_anti":
        # Should return right rows that DON'T have matches in left (id=3)
        assert len(result) == 1
        _assert_columns_match(result, {"id", "value", "score"})
        _assert_scalar_values(result_by_id, {3: {"value": "z", "score": 30}})

    # For outer joins, ensure unjoinable columns are present
    if join_type in ["inner", "left_outer", "right_outer", "full_outer"]:
        _assert_columns_match(
            result, {"id", "list_col", "tensor_col", "data", "value", "score"}
        )


@pytest.mark.parametrize(
    "join_type,filter_side,should_push",
    [
        ("inner", "left", True),
        ("inner", "right", True),
        ("left_outer", "left", True),
        ("left_outer", "right", False),
    ],
    ids=["inner_left", "inner_right", "left_outer_left", "left_outer_right"],
)
def test_join_with_predicate_pushdown(
    ray_start_regular_shared_2_cpus, join_type, filter_side, should_push
):
    """Test that predicate pushdown works correctly with different join types.

    Filters on single-side predicates should push past the join when appropriate:
    - Inner join: can push to either side
    - Left outer: can push to left (preserved) side only
    - Right outer: can push to right (preserved) side only
    """
    from ray.data._internal.logical.optimizers import LogicalOptimizer
    from ray.data._internal.util import MiB
    from ray.data.expressions import col

    DataContext.get_current().target_max_block_size = 1 * MiB

    # Create datasets directly without map to allow filter pushdown through join
    # Both have ids 0-31 with different value columns
    left_data = [{"id": i, "left_val": i * 10} for i in range(32)]
    right_data = [{"id": i, "right_val": i * 100} for i in range(32)]

    left_ds = ray.data.from_items(left_data)
    right_ds = ray.data.from_items(right_data)

    # Join then filter
    joined = left_ds.join(
        right_ds,
        join_type=join_type,
        num_partitions=4,
        on=("id",),
        aggregator_ray_remote_args={"num_cpus": 0.01},
    )

    # Filter on column from specified side
    if filter_side == "left":
        filtered_ds = joined.filter(expr=col("left_val") < 100)
    else:
        filtered_ds = joined.filter(expr=col("right_val") < 1000)

    left_pd = left_ds.to_pandas()
    right_pd = right_ds.to_pandas()

    # Compute expected join result
    if join_type == "inner":
        expected_pd = left_pd.merge(right_pd, on="id", how="inner")
    elif join_type == "left_outer":
        expected_pd = left_pd.merge(right_pd, on="id", how="left")
    else:
        raise ValueError(f"Unsupported join type for this test: {join_type}")

    # Apply filter (must match what we filtered in Ray Data)
    if filter_side == "left":
        # For left-side filter, use notna() to include NaN rows from outer joins
        expected_pd = expected_pd[expected_pd["left_val"] < 100]
    else:
        # For right-side filter in outer joins, NaN values fail the comparison
        # and are filtered out (matching Ray Data behavior)
        expected_pd = expected_pd[expected_pd["right_val"] < 1000]

    actual_df = filtered_ds.to_pandas()
    expected_df = expected_pd.reset_index(drop=True)

    assert rows_same(actual_df, expected_df), (
        f"Results don't match for {join_type} join with {filter_side} filter:\n"
        f"Actual:\n{actual_df}\n\nExpected:\n{expected_df}"
    )

    # Check plan to verify pushdown behavior
    logical_plan = filtered_ds._logical_plan
    optimized_plan = LogicalOptimizer().optimize(logical_plan)
    plan_str = optimized_plan.dag.dag_str

    join_idx = plan_str.find("Join[Join]")
    filter_idx = plan_str.find("Filter[Filter(")

    if should_push:
        # Filter should be pushed before join
        assert filter_idx != -1, f"Filter should exist in plan: {plan_str}"
        assert filter_idx < join_idx, (
            f"Filter should be pushed before Join for {join_type} with {filter_side} "
            f"predicate. Plan: {plan_str}"
        )
    else:
        # Filter should remain after join
        if filter_idx != -1:
            assert filter_idx > join_idx, (
                f"Filter should stay after Join for {join_type} with {filter_side} "
                f"predicate. Plan: {plan_str}"
            )


def test_join_cross_side_column_comparison_no_pushdown(ray_start_regular_shared_2_cpus):
    """Test PR bug: comparing differently-named join keys from both sides.

    When join keys have different names on left and right
    sides (e.g., left.id and right.user_id), a predicate like col("id") > col("user_id")
    references both sides but cannot be pushed to either side alone since each side
    only has one of these columns.

    Setup:
    - Left has columns: {id, user_id, left_val} - join on "id"
    - Right has columns: {id, user_id, right_val} - join on "user_id"
    - Join: left.id = right.user_id
    - Filter: col("id") > col("user_id") (with suffixes to avoid collision)
    """
    from ray.data._internal.logical.operators import Filter, Join
    from ray.data._internal.logical.optimizers import LogicalOptimizer
    from ray.data._internal.util import MiB
    from ray.data.expressions import col
    from ray.data.tests.test_util import plan_operator_comes_before

    DataContext.get_current().target_max_block_size = 1 * MiB

    # Left: has both id and user_id as columns, joins on "id"
    left_data = [{"id": i, "user_id": i + 5, "left_val": i * 10} for i in range(10)]
    # Right: has both id and user_id as columns, joins on "user_id"
    right_data = [{"id": i + 20, "user_id": i, "right_val": i * 5} for i in range(10)]

    left_ds = ray.data.from_items(left_data)
    right_ds = ray.data.from_items(right_data)

    # Join on left.id = right.user_id (different column names used as keys)
    # Need suffixes to avoid column name collision
    joined = left_ds.join(
        right_ds,
        join_type="inner",
        num_partitions=2,
        on=("id",),
        right_on=("user_id",),
        left_suffix="_l",
        right_suffix="_r",
        aggregator_ray_remote_args={"num_cpus": 0.01},
    )

    # Filter comparing non-join-key columns from both sides
    # left_val exists only on left, right_val exists only on right
    # Neither side can evaluate this alone
    filtered_ds = joined.filter(expr=col("left_val") > col("right_val"))

    # Verify correctness
    result = filtered_ds.take_all()
    # left.id = right.user_id means they match (both 0-9)
    # left_val = id * 10, right_val = user_id * 5 = id * 5
    # So left_val > right_val means id*10 > id*5, true for all id > 0
    assert len(result) == 9, f"Should have 9 rows (id 1-9), got {len(result)}"
    assert all(row["left_val"] > row["right_val"] for row in result)

    # Check plan: filter should NOT be pushed down (should stay after join)
    logical_plan = filtered_ds._logical_plan
    optimized_plan = LogicalOptimizer().optimize(logical_plan)

    # Filter should come AFTER Join (not pushed down)
    # Before join: left has left_val but not right_val, right has right_val but not left_val
    assert not plan_operator_comes_before(optimized_plan, Filter, Join), (
        "Filter comparing columns from both sides should NOT be pushed before Join "
        "since neither side has both columns"
    )


def test_chained_left_outer_join_with_empty_blocks(ray_start_regular_shared_2_cpus):
    """Regression test for https://github.com/ray-project/ray/issues/60013.

    The bug
    -------
    When a hash-shuffle join receives an **empty-row** block as the very first
    block for an input sequence, _shuffle_block() returns early (num_rows == 0)
    without sending any data to any aggregator.  The caller, however, marks
    _has_schemas_broadcasted[input_index] = True immediately after submitting
    the task. Every subsequent block for that sequence uses
    send_empty_blocks=False.  Aggregators that receive no non-empty rows from
    those subsequent blocks end up with an empty queue.  When finalize() is
    called, _combine([]) builds an ArrowBlockBuilder with zero blocks and
    returns a (0 rows, 0 columns) table.  The downstream join then fails with
    ColumnNotFoundError because the join key column is absent.

    We bypass the first join entirely and use ray.data.from_blocks() to build a
    dataset whose very first block is an explicit zero-row Arrow table that
    carries the full column schema.  With num_partitions=20 and only 10 data
    rows the second join has at least 10 aggregator partitions that receive no
    non-empty data.  Before the fix those partitions produce (0, 0) tables and
    the join raises ColumnNotFoundError.  After the fix schema-carrier blocks
    are broadcast even for the empty first block, so every aggregator can
    finalize correctly.
    """
    import pyarrow as pa

    # Build a dataset that simulates the output of a first left-outer join:
    #   - block 0: explicitly empty (0 rows) but carries the full schema
    #   - blocks 1-10: one row each, with b_val populated for id >= 5
    #
    # from_blocks() preserves block order and count exactly, so the empty block
    # is guaranteed to be the first block the second join's shuffle sees.
    schema = pa.schema(
        [
            pa.field("id", pa.int64()),
            pa.field("a_val", pa.string()),
            pa.field("b_val", pa.string()),
        ]
    )
    empty_block = schema.empty_table()  # shape (0, 3), schema but no rows

    data_blocks = [
        pa.table(
            {
                "id": pa.array([i], type=pa.int64()),
                "a_val": pa.array([f"a_{i}"], type=pa.string()),
                "b_val": pa.array([f"b_{i}" if i >= 5 else None], type=pa.string()),
            }
        )
        for i in range(10)
    ]

    # The first block must be the empty one so the bug fires.
    # from_blocks guarantees block order and count are preserved as-is.
    joined_1 = ray.data.from_blocks([empty_block] + data_blocks)

    # Second dataset for the chained join
    ds_c = ray.data.from_arrow(
        pa.table(
            {
                "id": pa.array(range(10), type=pa.int64()),
                "c_val": pa.array([f"c_{i}" for i in range(10)], type=pa.string()),
            }
        )
    )

    # num_partitions=20 with only 10 data rows means at least 10 aggregator
    # partitions receive no non-empty left-side data.
    joined_2 = joined_1.join(
        ds_c,
        join_type="left_outer",
        on=("id",),
        num_partitions=20,
    )

    result = joined_2.to_pandas()

    expected = pd.DataFrame(
        {
            "id": list(range(10)),
            "a_val": [f"a_{i}" for i in range(10)],
            "b_val": [f"b_{i}" if i >= 5 else None for i in range(10)],
            "c_val": [f"c_{i}" for i in range(10)],
        }
    )

    assert rows_same(result, expected)


@pytest.mark.parametrize(
    "join_type, expected_row_count",
    [
        ("inner", None),
        ("left_outer", None),
        ("right_outer", None),
        ("full_outer", None),
        ("left_semi", 1),
        ("right_semi", 1),
        ("left_anti", 1),
        ("right_anti", 0),
    ],
)
def test_overlapping_non_key_columns_without_suffixes(
    ray_start_regular_shared_2_cpus, join_type, expected_row_count
):
    """When both sides share a non-key column and no suffixes are provided,
    inner/outer joins must raise a clear ValueError (expected_row_count=None),
    while semi/anti joins should succeed because only one side's columns
    appear in the result."""
    left = ray.data.from_items([{"id": 1, "value": 10}, {"id": 2, "value": 20}])
    right = ray.data.from_items([{"id": 1, "value": 99}])

    joined = left.join(right, join_type=join_type, on=("id",), num_partitions=1)

    if expected_row_count is not None:
        assert len(joined.take_all()) == expected_row_count
    else:
        with pytest.raises(RayTaskError) as exc_info:
            joined.count()
        assert "Left and right columns suffixes cannot be both None" in str(
            exc_info.value.cause
        )


if __name__ == "__main__":
    import sys

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