chore: import upstream snapshot with attribution

python/ray/data/tests/test_util.py

@@ -0,0 +1,620 @@
+from typing import Any, Dict, List, Optional, Type
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pytest
+from typing_extensions import Hashable
+
+import ray
+from ray._common.retry import matches_error
+from ray.data._internal.datasource.parquet_datasource import ParquetDatasource
+from ray.data._internal.execution.interfaces import ExecutionResources
+from ray.data._internal.execution.util import merge_label_selector
+from ray.data._internal.logical.interfaces import LogicalPlan
+from ray.data._internal.logical.interfaces.logical_operator import LogicalOperator
+from ray.data._internal.logical.operators import Read
+from ray.data._internal.memory_tracing import (
+    leak_report,
+    trace_allocation,
+    trace_deallocation,
+)
+from ray.data._internal.planner.exchange.sort_task_spec import SortKey
+from ray.data._internal.remote_fn import _make_hashable, cached_remote_fn
+from ray.data._internal.usage.util import (
+    _recorded_operators,
+    _recorded_operators_lock,
+)
+from ray.data._internal.util import (
+    NULL_SENTINEL,
+    find_partition_index,
+    get_max_task_capacity,
+    iterate_with_retry,
+    merge_resources_to_ray_remote_args,
+    rows_same,
+)
+from ray.data.tests.conftest import *  # noqa: F401, F403
+
+
+def _check_usage_record(op_names: List[str], clear_after_check: Optional[bool] = True):
+    """Check if operators with given names in `op_names` have been used.
+    If `clear_after_check` is True, we clear the list of recorded operators
+    (so that subsequent checks do not use existing records of operator usage)."""
+    for op_name in op_names:
+        assert op_name not in ("Unknown", "ReadCustom", "WriteCustom"), op_name
+        with _recorded_operators_lock:
+            assert _recorded_operators.get(op_name, 0) > 0, (
+                op_name,
+                _recorded_operators,
+            )
+    if clear_after_check:
+        with _recorded_operators_lock:
+            _recorded_operators.clear()
+
+
+# Utilities for structural logical plan inspection
+# These provide type-safe alternatives to string-based plan matching
+
+
+def plan_has_operator(plan: LogicalPlan, op_type: Type[LogicalOperator]) -> bool:
+    """Check if plan contains at least one operator of given type.
+
+    Args:
+        plan: The logical plan to inspect
+        op_type: The operator type to search for
+
+    Returns:
+        True if at least one operator of the given type exists in the plan
+    """
+    return any(isinstance(op, op_type) for op in plan.dag.post_order_iter())
+
+
+def plan_operator_comes_before(
+    plan: LogicalPlan,
+    first_type: Type[LogicalOperator],
+    second_type: Type[LogicalOperator],
+) -> bool:
+    """Check if any operator of first_type comes before any operator of second_type.
+
+    Args:
+        plan: The logical plan to inspect
+        first_type: The operator type that should come first
+        second_type: The operator type that should come second
+
+    Returns:
+        True if at least one operator of first_type appears before at least one
+        operator of second_type in post-order traversal, False otherwise.
+    """
+    operators = list(plan.dag.post_order_iter())
+    first_indices = [i for i, op in enumerate(operators) if isinstance(op, first_type)]
+    second_indices = [
+        i for i, op in enumerate(operators) if isinstance(op, second_type)
+    ]
+
+    if not first_indices or not second_indices:
+        return False
+
+    # Check if the earliest first_type comes before the earliest second_type
+    return min(first_indices) < min(second_indices)
+
+
+def get_operators_of_type(
+    plan: LogicalPlan, op_type: Type[LogicalOperator]
+) -> List[LogicalOperator]:
+    """Get all operators of a specific type from the plan.
+
+    Args:
+        plan: The logical plan to inspect
+        op_type: The operator type to search for
+
+    Returns:
+        List of all operators of the given type in post-order traversal
+    """
+    return [op for op in plan.dag.post_order_iter() if isinstance(op, op_type)]
+
+
+def get_operator_types(plan: LogicalPlan) -> List[str]:
+    """Get list of operator type names in post-order traversal.
+
+    Args:
+        plan: The logical plan to inspect
+
+    Returns:
+        List of operator class names in the order they appear in post-order traversal
+    """
+    return [type(op).__name__ for op in plan.dag.post_order_iter()]
+
+
+def test_cached_remote_fn():
+    def foo():
+        pass
+
+    cpu_only_foo = cached_remote_fn(foo, num_cpus=1)
+    cached_cpu_only_foo = cached_remote_fn(foo, num_cpus=1)
+
+    assert cpu_only_foo == cached_cpu_only_foo
+
+    gpu_only_foo = cached_remote_fn(foo, num_gpus=1)
+
+    assert cpu_only_foo != gpu_only_foo
+
+
+def test_null_sentinel():
+    """Check that NULL_SENTINEL sorts greater than any other value."""
+
+    assert NULL_SENTINEL != NULL_SENTINEL
+    assert NULL_SENTINEL < NULL_SENTINEL
+    assert NULL_SENTINEL <= NULL_SENTINEL
+    assert not NULL_SENTINEL > NULL_SENTINEL
+    assert not NULL_SENTINEL >= NULL_SENTINEL
+
+    # With NoneType
+    assert None > NULL_SENTINEL
+    assert None >= NULL_SENTINEL
+    assert NULL_SENTINEL < None
+    assert NULL_SENTINEL <= None
+    assert NULL_SENTINEL != None  # noqa: E711
+
+    # With np.nan
+    assert np.nan > NULL_SENTINEL
+    assert np.nan >= NULL_SENTINEL
+    assert NULL_SENTINEL < np.nan
+    assert NULL_SENTINEL <= np.nan
+    assert NULL_SENTINEL != np.nan
+
+    # Rest
+    assert NULL_SENTINEL > 1000
+    assert NULL_SENTINEL > "abc"
+    assert NULL_SENTINEL != 1000
+    assert NULL_SENTINEL != "abc"
+    assert not NULL_SENTINEL < 1000
+    assert not NULL_SENTINEL < "abc"
+    assert not NULL_SENTINEL <= 1000
+    assert not NULL_SENTINEL <= "abc"
+    assert NULL_SENTINEL >= 1000
+    assert NULL_SENTINEL >= "abc"
+
+
+def test_make_hashable():
+    valid_args = {
+        "int": 0,
+        "float": 1.2,
+        "str": "foo",
+        "dict": {
+            0: 0,
+            1.2: 1.2,
+        },
+        "list": list(range(10)),
+        "tuple": tuple(range(3)),
+        "type": Hashable,
+    }
+
+    hashable_args = _make_hashable(valid_args)
+
+    assert hash(hashable_args) == hash(
+        (
+            ("dict", ((0, 0), (1.2, 1.2))),
+            ("float", 1.2),
+            ("int", 0),
+            ("list", (0, 1, 2, 3, 4, 5, 6, 7, 8, 9)),
+            ("str", "foo"),
+            ("tuple", (0, 1, 2)),
+            ("type", Hashable),
+        )
+    )
+
+    # Invalid case # 1: can't mix up key types
+    invalid_args = {0: 1, "bar": "baz"}
+
+    with pytest.raises(TypeError) as exc_info:
+        _make_hashable(invalid_args)
+
+    assert (
+        str(exc_info.value) == "'<' not supported between instances of 'str' and 'int'"
+    )
+
+
+@pytest.mark.parametrize("enabled", [False, True])
+def test_memory_tracing(enabled):
+    ctx = ray.data.context.DataContext.get_current()
+    ctx.trace_allocations = enabled
+    ref1 = ray.put(np.zeros(1024 * 1024))
+    ref2 = ray.put(np.zeros(1024 * 1024))
+    ref3 = ray.put(np.zeros(1024 * 1024))
+    trace_allocation(ref1, "test1")
+    trace_allocation(ref2, "test2")
+    trace_allocation(ref3, "test5")
+    trace_deallocation(ref1, "test3", free=False)
+    trace_deallocation(ref2, "test4", free=True)
+    ray.get(ref1)
+    with pytest.raises(ray.exceptions.ObjectFreedError):
+        ray.get(ref2)
+    report = leak_report()
+    print(report)
+
+    if enabled:
+        assert "Leaked object, created at test1" in report, report
+        assert "Leaked object, created at test5" in report, report
+        assert "Freed object from test2 at test4" in report, report
+        assert "skipped dealloc at test3" in report, report
+    else:
+        assert "test1" not in report, report
+        assert "test2" not in report, report
+        assert "test3" not in report, report
+        assert "test4" not in report, report
+        assert "test5" not in report, report
+
+
+def get_parquet_read_logical_op(
+    ray_remote_args: Optional[Dict[str, Any]] = None,
+    **read_kwargs,
+) -> Read:
+    datasource = ParquetDatasource(paths="example://iris.parquet")
+    if "parallelism" not in read_kwargs:
+        read_kwargs["parallelism"] = 10
+    read_op = Read(
+        datasource=datasource,
+        datasource_or_legacy_reader=datasource,
+        ray_remote_args=ray_remote_args,
+        **read_kwargs,
+    )
+    return read_op
+
+
+@ray.remote(num_cpus=0)
+class ConcurrencyCounter:
+    def __init__(self):
+        self.concurrency = 0
+        self.max_concurrency = 0
+
+    def inc(self):
+        self.concurrency += 1
+        if self.concurrency > self.max_concurrency:
+            self.max_concurrency = self.concurrency
+        return self.concurrency
+
+    def decr(self):
+        self.concurrency -= 1
+        return self.concurrency
+
+    def get_max_concurrency(self):
+        return self.max_concurrency
+
+
+def test_iterate_with_retry():
+    has_raised_error = False
+
+    class MockIterable:
+        """Iterate over the numbers 0, 1, 2, and raise an error on the first iteration
+        attempt.
+        """
+
+        def __init__(self, fail_at_index=3):
+            self._index = -1
+            self._fail_at_index = fail_at_index
+
+        def __iter__(self):
+            return self
+
+        def __next__(self):
+            self._index += 1
+
+            if self._index >= 10:
+                raise StopIteration
+
+            nonlocal has_raised_error
+            if self._index == self._fail_at_index and not has_raised_error:
+                has_raised_error = True
+                raise RuntimeError("Transient error")
+
+            return self._index
+
+    expected = list(range(10))
+    assert list(iterate_with_retry(MockIterable, description="get item")) == expected
+
+    has_raised_error = False
+    assert (
+        list(iterate_with_retry(MockIterable, description="get item", max_attempts=2))
+        == expected
+    )
+
+
+def test_iterate_with_retry_unwrap_cause():
+    """unwrap_cause=True makes `match` patterns search e.__cause__ as well."""
+    attempts = 0
+
+    class MockIterable:
+        def __init__(self):
+            nonlocal attempts
+            attempts += 1
+
+        def __iter__(self):
+            return self
+
+        def __next__(self):
+            try:
+                raise RuntimeError("transient inner")
+            except RuntimeError as inner:
+                raise ValueError("outer wrapper") from inner
+
+    # unwrap_cause=True: pattern matches the cause → all attempts consumed.
+    attempts = 0
+    with pytest.raises(ValueError, match="outer wrapper"):
+        list(
+            iterate_with_retry(
+                MockIterable,
+                description="get item",
+                match=["transient inner"],
+                max_attempts=2,
+                unwrap_cause=True,
+            )
+        )
+    assert attempts == 2
+
+    # unwrap_cause=False: cause invisible → not retryable → only one attempt.
+    attempts = 0
+    with pytest.raises(ValueError, match="outer wrapper"):
+        list(
+            iterate_with_retry(
+                MockIterable,
+                description="get item",
+                match=["transient inner"],
+                max_attempts=2,
+                unwrap_cause=False,
+            )
+        )
+    assert attempts == 1
+
+
+def test_iterate_with_retry_matches_class_name():
+    """Patterns can match the exception class name (e.g., 'RateLimit')."""
+
+    class RateLimitError(Exception):
+        pass
+
+    attempts = 0
+
+    class MockIterable:
+        def __init__(self):
+            nonlocal attempts
+            attempts += 1
+
+        def __iter__(self):
+            return self
+
+        def __next__(self):
+            raise RateLimitError("Error code: 429")
+
+    attempts = 0
+    with pytest.raises(RateLimitError):
+        list(
+            iterate_with_retry(
+                MockIterable,
+                description="get item",
+                match=["RateLimit"],
+                max_attempts=2,
+            )
+        )
+    assert attempts == 2
+
+
+@pytest.mark.parametrize(
+    "pattern, error_message, expected",
+    [
+        # Plain substring match.
+        ("transient", "transient network error", True),
+        # Regex match when substring check fails.
+        ("40[0-9]", "HTTP 404 not found", True),
+        # Substring takes priority: literal "(503)" found before regex is tried.
+        ("(503)", "error (503) returned", True),
+        # Invalid regex falls back to False, not re.error.
+        ("[unclosed", "some error message", False),
+        # No match at all.
+        ("rate limit", "connection refused", False),
+    ],
+)
+def test_matches_error(pattern, error_message, expected):
+    """Retry helper matches substring first, then regex; invalid patterns do not raise."""
+    assert matches_error(pattern, error_message) is expected
+
+
+def test_find_partition_index_single_column_ascending():
+    table = pa.table({"value": [1, 2, 2, 3, 5]})
+    sort_key = SortKey(key=["value"], descending=[False])
+    assert find_partition_index(table, (0,), sort_key) == 0  # all entries > 0
+    assert find_partition_index(table, (2,), sort_key) == 1  # first match index
+    assert find_partition_index(table, (4,), sort_key) == 4  # belongs after 3, before 5
+    assert find_partition_index(table, (6,), sort_key) == 5  # all entries < 6
+
+
+def test_find_partition_index_single_column_descending():
+    table = pa.table({"value": [5, 3, 2, 2, 1]})
+    sort_key = SortKey(key=["value"], descending=[True])
+    assert find_partition_index(table, (6,), sort_key) == 0  # belongs before 5
+    assert find_partition_index(table, (3,), sort_key) == 2  # after the last 3
+    assert find_partition_index(table, (2,), sort_key) == 4  # after the last 2
+    assert find_partition_index(table, (0,), sort_key) == 5  # all entries > 0
+
+
+def test_find_partition_index_multi_column():
+    # Table sorted by col1 asc, then col2 desc.
+    table = pa.table({"col1": [1, 1, 1, 2, 2], "col2": [3, 2, 1, 2, 1]})
+    sort_key = SortKey(key=["col1", "col2"], descending=[False, True])
+    # Insert value (1,3) -> belongs before (1,2)
+    assert find_partition_index(table, (1, 3), sort_key) == 0
+    # Insert value (1,2) -> belongs after the first (1,3) and before (1,2)
+    # because col1 ties, col2 descending
+    assert find_partition_index(table, (1, 2), sort_key) == 1
+    # Insert value (2,2) -> belongs right before (2,2) that starts at index 3
+    assert find_partition_index(table, (2, 2), sort_key) == 3
+    # Insert value (0, 4) -> belongs at index 0 (all col1 > 0)
+    assert find_partition_index(table, (0, 4), sort_key) == 0
+    # Insert value (2,0) -> belongs after (2,1)
+    assert find_partition_index(table, (2, 0), sort_key) == 5
+
+
+def test_find_partition_index_with_nulls():
+    # _NullSentinel is sorted greater, so they appear after all real values.
+    table = pa.table({"value": [1, 2, 3, None, None]})
+    sort_key = SortKey(key=["value"], descending=[False])
+    # Insert (2,) -> belongs after 1, before 2 => index 1
+    # (But the actual find_partition_index uses the table as-is.)
+    assert find_partition_index(table, (2,), sort_key) == 1
+    # Insert (4,) -> belongs before any null => index 3
+    assert find_partition_index(table, (4,), sort_key) == 3
+    # Insert (None,) -> always belongs at the end
+    assert find_partition_index(table, (None,), sort_key) == 3
+
+
+def test_find_partition_index_with_nan():
+    # NaN sorts after regular values in Arrow (before nulls).
+    table = pa.table({"value": [1.0, 2.0, 3.0, float("nan"), float("nan")]})
+    sort_key = SortKey(key=["value"], descending=[False])
+    assert find_partition_index(table, (2.0,), sort_key) == 1
+    assert find_partition_index(table, (4.0,), sort_key) == 3
+
+
+def test_find_partition_index_with_nan_and_nulls():
+    # NaN sorts after regular values, nulls sort after NaN.
+    table = pa.table({"value": [1.0, 2.0, 3.0, float("nan"), None]})
+    sort_key = SortKey(key=["value"], descending=[False])
+    assert find_partition_index(table, (2.0,), sort_key) == 1
+    assert find_partition_index(table, (4.0,), sort_key) == 3
+    assert find_partition_index(table, (None,), sort_key) == 3
+
+
+def test_find_partition_index_duplicates():
+    table = pa.table({"value": [2, 2, 2, 2, 2]})
+    sort_key = SortKey(key=["value"], descending=[False])
+    # Insert (2,) in a table of all 2's -> first matching index is 0
+    assert find_partition_index(table, (2,), sort_key) == 0
+    # Insert (1,) -> belongs at index 0
+    assert find_partition_index(table, (1,), sort_key) == 0
+    # Insert (3,) -> belongs at index 5
+    assert find_partition_index(table, (3,), sort_key) == 5
+
+
+def test_find_partition_index_duplicates_descending():
+    table = pa.table({"value": [2, 2, 2, 2, 2]})
+    sort_key = SortKey(key=["value"], descending=[True])
+    # Insert (2,) in a table of all 2's -> belongs at index 5
+    assert find_partition_index(table, (2,), sort_key) == 5
+    # Insert (1,) -> belongs at index 5
+    assert find_partition_index(table, (1,), sort_key) == 5
+    # Insert (3,) -> belongs at index 0
+    assert find_partition_index(table, (3,), sort_key) == 0
+
+
+def test_merge_resources_to_ray_remote_args():
+    ray_remote_args = {}
+    ray_remote_args = merge_resources_to_ray_remote_args(1, 1, 1, ray_remote_args)
+    assert ray_remote_args == {"num_cpus": 1, "num_gpus": 1, "memory": 1}
+
+    ray_remote_args = {"other_resource": 1}
+    ray_remote_args = merge_resources_to_ray_remote_args(1, 1, 1, ray_remote_args)
+    assert ray_remote_args == {
+        "num_cpus": 1,
+        "num_gpus": 1,
+        "memory": 1,
+        "other_resource": 1,
+    }
+
+
+@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)
+
+
+@pytest.mark.parametrize(
+    "allocated,min_scheduling,expected",
+    [
+        (None, ExecutionResources(cpu=1, gpu=1), 0),
+        (ExecutionResources(cpu=1, gpu=1), ExecutionResources(cpu=1, gpu=1), 1),
+        (
+            ExecutionResources(cpu=1, gpu=1),
+            ExecutionResources(cpu=0, gpu=0),
+            float("inf"),
+        ),
+        (ExecutionResources(cpu=1, gpu=1), ExecutionResources(cpu=0, gpu=1), 1),
+        (ExecutionResources(cpu=1, gpu=1), ExecutionResources(cpu=1, gpu=0), 1),
+    ],
+)
+def test_get_max_task_capacity(allocated, min_scheduling, expected):
+    assert get_max_task_capacity(allocated, min_scheduling) == expected
+
+
+class TestMergeLabelSelector:
+    """Tests for ``merge_label_selector``.
+
+    The helper merges a DataContext-level label_selector into a ray_remote_args
+    dict. Operator-level entries win on key conflicts.
+    """
+
+    def test_ctx_none_returns_input_unchanged(self):
+        args = {"num_cpus": 1}
+        assert merge_label_selector(args, None) is args
+
+    def test_ctx_empty_returns_input_unchanged(self):
+        args = {"num_cpus": 1}
+        assert merge_label_selector(args, {}) is args
+
+    def test_ctx_only(self):
+        args = {"num_cpus": 1}
+        out = merge_label_selector(args, {"subcluster": "train"})
+        assert out == {"num_cpus": 1, "label_selector": {"subcluster": "train"}}
+        assert args == {"num_cpus": 1}  # input not mutated
+
+    def test_op_only_no_ctx(self):
+        args = {"label_selector": {"node": "X"}}
+        assert merge_label_selector(args, None) is args
+
+    def test_op_and_ctx_no_collision(self):
+        args = {"label_selector": {"node": "X"}}
+        out = merge_label_selector(args, {"subcluster": "train"})
+        assert out["label_selector"] == {"subcluster": "train", "node": "X"}
+
+    def test_op_wins_on_collision(self):
+        args = {"label_selector": {"subcluster": "val"}}
+        out = merge_label_selector(args, {"subcluster": "train"})
+        assert out["label_selector"] == {"subcluster": "val"}
+
+    def test_input_not_mutated(self):
+        args = {"label_selector": {"node": "X"}}
+        ctx = {"subcluster": "train"}
+        merge_label_selector(args, ctx)
+        assert args == {"label_selector": {"node": "X"}}
+        assert ctx == {"subcluster": "train"}
+
+
+def test_execution_options_label_selector_field():
+    """Smoke test that ExecutionOptions exposes label_selector."""
+    from ray.data._internal.execution.interfaces import ExecutionOptions
+
+    options = ExecutionOptions()
+    assert options.label_selector is None
+
+    options = ExecutionOptions(label_selector={"subcluster": "train"})
+    assert options.label_selector == {"subcluster": "train"}
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))