ronantakizawa/codereview-bench huggingface dataset:
Repository: ray-project/ray
File: python/ray/data/tests/test_operators.py (Python)

Code before:
    union_op = UnionOperator(ctx, input_op1, input_op2)

    # Verify that throttling_disabled() returns True
    assert union_op.throttling_disabled() is True


def test_union_operator_throttling_matches_similar_operators(ray_start_regular_shared):
    """Test that UnionOperator's throttling behavior matches other metadata-only operators."""
    from ray.data._internal.execution.operators.limit_operator import LimitOperator
    from ray.data._internal.execution.operators.output_splitter import OutputSplitter

    ctx = DataContext.get_current()

    # Create input buffer
    input_bundles = make_ref_bundles([[1, 2, 3]])
    input_op = InputDataBuffer(ctx, input_bundles)

    # Create operators that only manipulate metadata
    union_op = UnionOperator(ctx, input_op)
    limit_op = LimitOperator(10, input_op, ctx)
    output_splitter = OutputSplitter(input_op, 2, False, ctx)

    # All three operators should have throttling disabled
    assert union_op.throttling_disabled() is True
    assert limit_op.throttling_disabled() is True
    assert output_splitter.throttling_disabled() is True


if __name__ == "__main__":
    import sys

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


Review comment (bug):
IMO this isn't a super useful test:   
  - If LimitOperator or OutputSplitter change their throttling behavior for a legitimate reason, this test fails even though UnionOperator is working correctly                             
  - The first test already verifies UnionOperator.throttling_disabled() is True                                                                                                             
                                                                                     

Code after:
    # Verify that throttling_disabled() returns True
    assert union_op.throttling_disabled() is True


def test_union_operator_not_allocated_resources(
    ray_start_regular_shared, restore_data_context
):
    ctx = DataContext.get_current()
    ctx.op_resource_reservation_enabled = True
    input_op1 = InputDataBuffer(ctx, make_ref_bundles([[1, 2, 3]]))
    input_op2 = InputDataBuffer(ctx, make_ref_bundles([[4, 5, 6]]))
    union_op = UnionOperator(ctx, input_op1, input_op2)

    topo = build_streaming_topology(union_op, ExecutionOptions())
    resource_manager = ResourceManager(topo, ExecutionOptions(), MagicMock(), ctx)
    allocator = create_resource_allocator(resource_manager, ctx)
    assert allocator is not None

    allocator.update_budgets(
        limits=ExecutionResources(cpu=1, gpu=0, object_store_memory=1000)
    )
    allocation = allocator.get_allocation(union_op)
    assert allocation is None


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", __file__]))
Repository: ray-project/ray
File: python/ray/data/tests/unit/test_deduping_schema.py (Python)

Code before:
    incoming_schema = pa.schema(
        [
            pa.field("foo", pa.int64()),
            pa.field("baz", pa.bool8()),
            pa.field("qux", pa.float64()),
        ]
    )

    incoming_bundle = RefBundle([], owns_blocks=False, schema=incoming_schema)
    out_bundle, diverged = dedupe_schemas_with_validation(
        old_schema, incoming_bundle, enforce_schemas=enforce_schemas, warn=True
    )

    assert diverged

    if enforce_schemas:
        assert out_bundle.schema == pa.schema(
            [
                pa.field("foo", pa.int64()),
                pa.field("bar", pa.string()),
                pa.field("baz", pa.bool8()),
                pa.field("qux", pa.float64()),
            ]
        )
    else:
        assert out_bundle.schema == old_schema


if __name__ == "__main__":
    import sys

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


Review comment (suggestion):
The new test `test_dedupe_schema_mismatch` exercises the schema unification behavior but doesn’t actually assert on the new warning message content (or even that a warning is emitted), even though the PR description calls out this test as showcasing the new warning. Consider capturing the log output (e.g. via `caplog`) and asserting that only the mismatched fields appear and that the full schemas are not printed, so regressions in the warning format are caught by tests.

Code after:
)
def test_build_mismatch_warning_empty(incoming_schema: Optional["Schema"]):
    old_schema = pa.schema(
        [
            pa.field("foo", pa.int32()),
            pa.field("bar", pa.string()),
            pa.field("baz", pa.bool_()),
        ]
    )

    msg = _build_schemas_mismatch_warning(old_schema, incoming_schema)

    assert (
        """Operator produced a RefBundle with an empty/unknown schema. (3 total):
    foo: int32
    bar: string
    baz: bool
This may lead to unexpected behavior."""
        == msg
    )


@pytest.mark.parametrize("truncation_length", [20, 4, 3, 2])
def test_build_mismatch_warning_truncation(truncation_length: int):
    old_schema = pa.schema(
        [
            pa.field("foo", pa.int32()),
            pa.field("bar", pa.string()),
            pa.field("baz", pa.bool_()),
            pa.field("quux", pa.uint16()),
            pa.field("corge", pa.int64()),
            pa.field("grault", pa.int32()),
        ]
    )
    incoming_schema = pa.schema(
        [
            pa.field("foo", pa.int64()),
            pa.field("baz", pa.bool_()),
            pa.field("qux", pa.float64()),
            pa.field("quux", pa.uint32()),
            pa.field("corge", pa.int32()),
            pa.field("grault", pa.int64()),
        ]
    )

    msg = _build_schemas_mismatch_warning(
        old_schema, incoming_schema, truncation_length=truncation_length
    )

    if truncation_length >= 4:
        assert (
Repository: ray-project/ray
File: python/ray/data/tests/test_with_column.py (Python)

Code before:
    # Create dataset
    ds = ray.data.range(5, override_num_blocks=1)

    # Use async generator UDF
    udf_instance = AsyncGeneratorMultiYield(5)
    result = ds.with_column("result", udf_instance(col("id")))

    result_df = result.to_pandas()

    # Fixed behavior: last yield is returned
    # Input: [0, 1, 2, 3, 4]
    # First yield: [0+5, 1+5, 2+5, 3+5, 4+5] = [5, 6, 7, 8, 9]
    # Second yield: [0*15, 1*15, 2*15, 3*15, 4*15] = [0, 15, 30, 45, 60]
    # Third yield (RETURNED): [0+10, 1+10, 2+10, 3+10, 4+10] = [10, 11, 12, 13, 14]

    expected_after_fix = pd.DataFrame(
        {
            "id": [0, 1, 2, 3, 4],
            "result": [10, 11, 12, 13, 14],  # Last yield returned: id + (5*2) = id + 10
        }
    )

    assert rows_same(result_df, expected_after_fix)


@pytest.mark.skipif(
    get_pyarrow_version() < parse_version("20.0.0"),
    reason="with_column requires PyArrow >= 20.0.0",
)
@pytest.mark.parametrize(
    "expr, target_type, expected_value, expected_type",
    [
        # Basic type conversions
        (col("id"), "int64", 0, int),  # int64 -> int64 (no change)
        (col("id"), "float64", 0.0, float),  # int64 -> float64
        (col("id"), "string", "0", str),  # int64 -> string
        (col("id") / 2, "int64", 0, int),  # float -> int64
        (col("id") / 2, "float64", 0.0, float),  # float -> float64 (no change)
        # Using PyArrow types
        (col("id"), pa.int64(), 0, int),
        (col("id"), pa.float64(), 0.0, float),
        (col("id"), pa.string(), "0", str),
        # Using DataType objects
        (col("id"), DataType.int64(), 0, int),
        (col("id"), DataType.float64(), 0.0, float),
        (col("id"), DataType.string(), "0", str),
    ],
)
def test_cast_expression_basic(
    ray_start_regular_shared,
    expr,

Review comment (suggestion):
Can you please move these tests to this folder: `python/ray/data/tests/expressions`

You can call it `test_cast.py`

Code after:
    # Create dataset
    ds = ray.data.range(5, override_num_blocks=1)

    # Use async generator UDF
    udf_instance = AsyncGeneratorMultiYield(5)
    result = ds.with_column("result", udf_instance(col("id")))

    result_df = result.to_pandas()

    # Fixed behavior: last yield is returned
    # Input: [0, 1, 2, 3, 4]
    # First yield: [0+5, 1+5, 2+5, 3+5, 4+5] = [5, 6, 7, 8, 9]
    # Second yield: [0*15, 1*15, 2*15, 3*15, 4*15] = [0, 15, 30, 45, 60]
    # Third yield (RETURNED): [0+10, 1+10, 2+10, 3+10, 4+10] = [10, 11, 12, 13, 14]

    expected_after_fix = pd.DataFrame(
        {
            "id": [0, 1, 2, 3, 4],
            "result": [10, 11, 12, 13, 14],  # Last yield returned: id + (5*2) = id + 10
        }
    )

    assert rows_same(result_df, expected_after_fix)


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", __file__]))
Repository: ray-project/ray
File: python/ray/data/_internal/util.py (Python)

Code before:
    num_gpus: Optional[int],
    memory: Optional[int],
    ray_remote_args: Dict[str, Any],
) -> Dict[str, Any]:
    """Convert the given resources to Ray remote args.

    Args:
        num_cpus: The number of CPUs to be added to the Ray remote args.
        num_gpus: The number of GPUs to be added to the Ray remote args.
        memory: The memory to be added to the Ray remote args.
        ray_remote_args: The Ray remote args to be merged.

    Returns:
        The converted arguments.
    """
    ray_remote_args = ray_remote_args.copy()
    if num_cpus is not None:
        ray_remote_args["num_cpus"] = num_cpus
    if num_gpus is not None:
        ray_remote_args["num_gpus"] = num_gpus
    if memory is not None:
        ray_remote_args["memory"] = memory
    return ray_remote_args


def datetime_to_ms(dt: datetime) -> int:
    """Convert a datetime object to milliseconds since epoch (UTC).

    Naive datetime objects (without timezone info) are treated as UTC.
    Timezone-aware datetime objects are converted to UTC automatically.

    Args:
        dt: A datetime object.

    Returns:
        Milliseconds since Unix epoch.
    """
    if dt.tzinfo is None:
        # Naive datetime: treat as UTC
        dt = dt.replace(tzinfo=timezone.utc)
    return int(dt.timestamp() * 1000)


@DeveloperAPI
def infer_compression(path: str) -> Optional[str]:
    import pyarrow as pa

    compression = None
    try:
        # Try to detect compression codec from path.
        compression = pa.Codec.detect(path).name

Review comment (refactor):
Nit: This `util.py` is kind of a kitchen sink file. If this utility function is only used by the Kafka implementation, it might be better located in `kafka_datasource.py`

Code after:
    memory: Optional[int],
    ray_remote_args: Dict[str, Any],
) -> Dict[str, Any]:
    """Convert the given resources to Ray remote args.

    Args:
        num_cpus: The number of CPUs to be added to the Ray remote args.
        num_gpus: The number of GPUs to be added to the Ray remote args.
        memory: The memory to be added to the Ray remote args.
        ray_remote_args: The Ray remote args to be merged.

    Returns:
        The converted arguments.
    """
    ray_remote_args = ray_remote_args.copy()
    if num_cpus is not None:
        ray_remote_args["num_cpus"] = num_cpus
    if num_gpus is not None:
        ray_remote_args["num_gpus"] = num_gpus
    if memory is not None:
        ray_remote_args["memory"] = memory
    return ray_remote_args


@DeveloperAPI
def infer_compression(path: str) -> Optional[str]:
    import pyarrow as pa

    compression = None
    try:
        # Try to detect compression codec from path.
        compression = pa.Codec.detect(path).name
    except (ValueError, TypeError):
        # Arrow's compression inference on the file path doesn't work for Snappy, so we double-check ourselves.
        import pathlib

        suffix = pathlib.Path(path).suffix
        if suffix and suffix[1:] == "snappy":
            compression = "snappy"
    return compression
Repository: ray-project/ray
File: python/ray/data/tests/test_util.py (Python)

Code before:


@pytest.mark.parametrize(
    "actual, expected, expected_equal",
    [
        (pd.DataFrame({"a": [1]}), pd.DataFrame({"a": [1]}), True),
        # Different value.
        (pd.DataFrame({"a": [1]}), pd.DataFrame({"a": [2]}), False),
        # Extra column.
        (pd.DataFrame({"a": [1]}), pd.DataFrame({"a": [1], "b": [2]}), False),
        # Different number of rows.
        (pd.DataFrame({"a": [1]}), pd.DataFrame({"a": [1, 1]}), False),
        # Same rows, but different order.
        (pd.DataFrame({"a": [1, 2]}), pd.DataFrame({"a": [2, 1]}), True),
    ],
)
def test_rows_same(actual: pd.DataFrame, expected: pd.DataFrame, expected_equal: bool):
    if expected_equal:
        assert rows_same(actual, expected)
    else:
        with pytest.raises(AssertionError):
            assert rows_same(actual, expected)


def test_datetime_to_ms_naive():
    """Test that naive datetimes are treated as UTC."""
    # Unix epoch
    dt = datetime(1970, 1, 1, 0, 0, 0)
    assert datetime_to_ms(dt) == 0

    # Known timestamp: 2025-01-01 00:00:00 UTC = 1735689600000 ms
    dt = datetime(2025, 1, 1, 0, 0, 0)
    assert datetime_to_ms(dt) == 1735689600000


def test_datetime_to_ms_aware():
    """Test that timezone-aware datetimes are converted correctly."""
    # UTC timezone-aware
    dt = datetime(1970, 1, 1, 0, 0, 0, tzinfo=timezone.utc)
    assert datetime_to_ms(dt) == 0

    # 2025-01-01 00:00:00 UTC
    dt = datetime(2025, 1, 1, 0, 0, 0, tzinfo=timezone.utc)
    assert datetime_to_ms(dt) == 1735689600000


if __name__ == "__main__":
    import sys

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


Review comment (suggestion):
I have no idea what a naive datetime means. Is that a standard term? If not, it might also confuse other readers

Code after:
@pytest.mark.parametrize(
    "actual, expected, expected_equal",
    [
        (pd.DataFrame({"a": [1]}), pd.DataFrame({"a": [1]}), True),
        # Different value.
        (pd.DataFrame({"a": [1]}), pd.DataFrame({"a": [2]}), False),
        # Extra column.
        (pd.DataFrame({"a": [1]}), pd.DataFrame({"a": [1], "b": [2]}), False),
        # Different number of rows.
        (pd.DataFrame({"a": [1]}), pd.DataFrame({"a": [1, 1]}), False),
        # Same rows, but different order.
        (pd.DataFrame({"a": [1, 2]}), pd.DataFrame({"a": [2, 1]}), True),
    ],
)
def test_rows_same(actual: pd.DataFrame, expected: pd.DataFrame, expected_equal: bool):
    if expected_equal:
        assert rows_same(actual, expected)
    else:
        with pytest.raises(AssertionError):
            assert rows_same(actual, expected)


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", __file__]))
Repository: ray-project/ray
File: python/ray/data/_internal/execution/interfaces/execution_options.py (Python)

Code before:
        """Returns True if resource_limits is the default value."""
        return self._resource_limits == ExecutionResources.for_limits()

    def validate(self) -> None:
        """Validate the options."""
        for attr in ["cpu", "gpu", "object_store_memory"]:
            if (
                getattr(self.resource_limits, attr) != float("inf")
                and getattr(self.exclude_resources, attr, 0) > 0
            ):
                raise ValueError(
                    "resource_limits and exclude_resources cannot "
                    f" both be set for {attr} resource."
                )

    @property
    def locality_with_output(self) -> bool:
        return False

    @locality_with_output.setter
    def locality_with_output(self, value: Union[bool, List[NodeIdStr]]) -> None:
        if value:
            warnings.warn(
                "`ExecutionOptions.locality_with_output` has been removed and is now "
                "a no-op. Use `Dataset.streaming_split(..., locality_hints=...)` "
                "for output locality.",
                DeprecationWarning,
                stacklevel=2,
            )


def safe_or(value: Optional[Any], alt: Any) -> Any:
    return value if value is not None else alt


def safe_round(
    value: Optional[float], ndigits: Optional[int] = None
) -> Optional[float]:
    if value is None:
        return None
    elif math.isinf(value):
        return value
    else:
        return round(value, ndigits)


Review comment (suggestion):
I think this might confuse users because users don't typically call `streaming_split`. It's more often called internally by Ray Train
```suggestion
                "`ExecutionOptions.locality_with_output` has been removed and is now "
                "a no-op.
```

Code after:
    def validate(self) -> None:
        """Validate the options."""
        for attr in ["cpu", "gpu", "object_store_memory"]:
            if (
                getattr(self.resource_limits, attr) != float("inf")
                and getattr(self.exclude_resources, attr, 0) > 0
            ):
                raise ValueError(
                    "resource_limits and exclude_resources cannot "
                    f" both be set for {attr} resource."
                )

    @property
    def locality_with_output(self) -> bool:
        return False

    @locality_with_output.setter
    def locality_with_output(self, value: Union[bool, List[NodeIdStr]]) -> None:
        if value:
            warnings.warn(
                "`ExecutionOptions.locality_with_output` has been removed and is now "
                "a no-op.",
                DeprecationWarning,
                stacklevel=2,
            )


def safe_or(value: Optional[Any], alt: Any) -> Any:
    return value if value is not None else alt


def safe_round(
    value: Optional[float], ndigits: Optional[int] = None
) -> Optional[float]:
    if value is None:
        return None
    elif math.isinf(value):
        return value
    else:
        return round(value, ndigits)
Repository: ray-project/ray
File: ci/ray_ci/supported_images.py (Python)

Code before:
import os
from functools import lru_cache
from typing import List

import yaml


@lru_cache(maxsize=1)
def load_supported_images():
    yaml_path = os.path.join(
        os.path.dirname(os.path.abspath(__file__)),
        "..",
        "..",
        "ray-images.yaml",
    )
    with open(yaml_path) as f:
        return yaml.safe_load(f)


def get_image_config(image_type: str) -> dict:
    return load_supported_images()[image_type]


def get_python_versions(image_type: str) -> List[str]:
    return get_image_config(image_type)["python"]


def get_platforms(image_type: str) -> List[str]:
    return get_image_config(image_type)["platforms"]


def get_architectures(image_type: str) -> List[str]:
    return get_image_config(image_type)["architectures"]


def get_default(image_type: str, key: str) -> str:
    return get_image_config(image_type)["defaults"][key]


Review comment (suggestion):
instead of doing this relative path.. is it possible to search for maybe the `.rayciversion` file for the root? I always feel that `..` is a mysterious way to locate a file in a big repo..

Code after:
from functools import lru_cache
from pathlib import Path
from typing import List

import yaml

_RAYCI_VERSION_FILE = ".rayciversion"


def _find_ray_root() -> Path:
    """Walk up from this file and cwd looking for .rayciversion."""
    start = Path(__file__).resolve()
    for parent in start.parents:
        if (parent / _RAYCI_VERSION_FILE).exists():
            return parent
    if (Path.cwd() / _RAYCI_VERSION_FILE).exists():
        return Path.cwd()
    raise FileNotFoundError("Could not find Ray root (missing .rayciversion).")


@lru_cache(maxsize=1)
def load_supported_images():
    yaml_path = _find_ray_root() / "ray-images.yaml"
    with open(yaml_path) as f:
        return yaml.safe_load(f)


def get_image_config(image_type: str) -> dict:
    return load_supported_images()[image_type]


def get_python_versions(image_type: str) -> List[str]:
    return get_image_config(image_type)["python"]


def get_platforms(image_type: str) -> List[str]:
    return get_image_config(image_type)["platforms"]
Repository: ray-project/ray
File: python/ray/data/_internal/logical/rules/predicate_pushdown.py (Python)

Code before:
        new_inputs = list(conditional_op.input_dependencies)
        branch_filter = Filter(
            new_inputs[branch_idx],
            predicate_expr=filter_op.predicate_expr,
        )
        new_inputs[branch_idx] = cls._try_push_down_predicate(branch_filter)

        # Return operator with updated input
        return cls._clone_op_with_new_inputs(conditional_op, new_inputs)

    @classmethod
    def _clone_op_with_new_inputs(
        cls, op: LogicalOperator, new_inputs: List[LogicalOperator]
    ) -> LogicalOperator:
        """Clone an operator with new inputs.

        Args:
            op: The operator to clone
            new_inputs: List of new input operators (can be single element list)

        Returns:
            A shallow copy of the operator with updated input dependencies
        """
        new_op = copy.copy(op)
        new_op.input_dependencies = new_inputs
        return new_op


Review comment (refactor):
What's the motivation for this change/how does it relate to this refactor? What happens if we remove it?

Code after:
        new_inputs = list(conditional_op.input_dependencies)
        branch_filter = Filter(
            new_inputs[branch_idx],
            predicate_expr=filter_op.predicate_expr,
        )
        new_inputs[branch_idx] = cls._try_push_down_predicate(branch_filter)

        # Return operator with updated input
        return cls._clone_op_with_new_inputs(conditional_op, new_inputs)

    @classmethod
    def _clone_op_with_new_inputs(
        cls, op: LogicalOperator, new_inputs: List[LogicalOperator]
    ) -> LogicalOperator:
        """Clone an operator with new inputs.

        Args:
            op: The operator to clone
            new_inputs: List of new input operators (can be single element list)

        Returns:
            A shallow copy of the operator with updated input dependencies
        """
        if isinstance(op, Limit):
            assert len(new_inputs) == 1, len(new_inputs)
            return Limit(new_inputs[0], op.limit)
        new_op = copy.copy(op)
        new_op.input_dependencies = new_inputs
        return new_op