chore: import upstream snapshot with attribution

python/ray/data/tests/util.py

@@ -0,0 +1,215 @@
+import functools
+import os
+import tempfile
+from contextlib import contextmanager
+from typing import Any, Callable, Iterable, List, Optional
+
+import pandas as pd
+
+import ray
+from ray.data._internal.execution.interfaces.physical_operator import (
+    DataOpTask,
+    MetadataOpTask,
+    PhysicalOperator,
+    RefBundle,
+)
+from ray.data._internal.execution.metadata_fetcher import (
+    InlineMetadataFetcher,
+    ThreadedMetadataFetcher,
+)
+from ray.data._internal.execution.operators.map_transformer import (
+    BlockMapTransformFn,
+    MapTransformCallable,
+    MapTransformer,
+)
+from ray.data._internal.output_buffer import OutputBlockSizeOption
+from ray.data.block import Block
+from ray.data.expressions import Expr
+
+
+@ray.remote
+class Counter:
+    def __init__(self):
+        self.count = 0
+
+    def increment(self):
+        self.count += 1
+
+    def get(self):
+        return self.count
+
+    def reset(self):
+        self.count = 0
+
+
+@contextmanager
+def gen_bin_files(n):
+    with tempfile.TemporaryDirectory() as temp_dir:
+        paths = []
+        for i in range(n):
+            path = os.path.join(temp_dir, f"{i}.bin")
+            paths.append(path)
+            with open(path, "wb") as fp:
+                to_write = str(i) * 500
+                fp.write(to_write.encode())
+        yield (temp_dir, paths)
+
+
+def column_udf(col, udf):
+    @functools.wraps(udf)
+    def wraps(row):
+        return {col: udf(row[col])}
+
+    return wraps
+
+
+def column_udf_class(col, udf):
+    class UDFClass:
+        def __call__(self, row):
+            return {col: udf(row[col])}
+
+    return UDFClass
+
+
+# Ex: named_values("id", [1, 2, 3])
+# Ex: named_values(["id", "id2"], [(1, 1), (2, 2), (3, 3)])
+def named_values(col_names, tuples):
+    output = []
+    if isinstance(col_names, list):
+        for t in tuples:
+            output.append(dict(zip(col_names, t)))
+    else:
+        for t in tuples:
+            output.append({col_names: t})
+    return output
+
+
+def extract_values(col_name, tuples):
+    return [t[col_name] for t in tuples]
+
+
+def assert_exprs_equal(actual: List[Expr], expected: List[Expr]):
+    """Assert two expression lists match element-wise.
+
+    ``Expr`` overloads ``==`` to build a comparison expression (e.g.
+    ``col("a") == 5``), so it can't be used to compare exprs for equality;
+    use ``structurally_equals`` instead.
+    """
+    actual_names = [e.name for e in actual]
+    expected_names = [e.name for e in expected]
+    assert len(actual) == len(expected), (actual_names, expected_names)
+    assert all(a.structurally_equals(b) for a, b in zip(actual, expected)), (
+        actual_names,
+        expected_names,
+    )
+
+
+def fetcher_has_pending_work(fetcher: ThreadedMetadataFetcher) -> bool:
+    """Whether a ``ThreadedMetadataFetcher`` still has a submitted pair to emit
+    or a postponed done-callback to fire. Test-only poll helper (peeks at the
+    fetcher's internals); call from the same thread that drives the fetcher."""
+    return any(fetcher._fifos.values()) or bool(fetcher._drained_tasks)
+
+
+def run_op_tasks_sync(op: PhysicalOperator, only_existing=False):
+    """Run tasks of a PhysicalOperator synchronously.
+
+    By default, this function will run until the op no longer has any active tasks.
+    If only_existing is True, this function will only run the currently existing tasks.
+    """
+    tasks = op.get_active_tasks()
+    while tasks:
+        ref_to_task = {task.get_waitable(): task for task in tasks}
+        ready, _ = ray.wait(
+            [task.get_waitable() for task in tasks],
+            num_returns=len(tasks),
+            fetch_local=False,
+            timeout=0.1,
+        )
+
+        for ref in ready:
+            task = ref_to_task[ref]
+            if isinstance(task, DataOpTask):
+                # Read all currently available output from the streaming generator
+                task.on_data_ready(None, InlineMetadataFetcher())
+                # Only remove the task when the generator has been fully exhausted
+                if task.has_finished:
+                    tasks.remove(task)
+            else:
+                assert isinstance(task, MetadataOpTask)
+                task.on_task_finished()
+                tasks.remove(task)
+
+        # NOTE: If only existing tasks need to be handled skip refreshing list
+        #       of outstanding tasks
+        if only_existing:
+            pass
+        else:
+            tasks = op.get_active_tasks()
+
+
+def run_one_op_task(op):
+    """Run one task of a PhysicalOperator."""
+    tasks = op.get_active_tasks()
+
+    while tasks:
+        waitable_to_tasks = {task.get_waitable(): task for task in tasks}
+
+        # Block, until 1 task is ready
+        ready, _ = ray.wait(
+            list(waitable_to_tasks.keys()), num_returns=1, fetch_local=False
+        )
+
+        task = waitable_to_tasks[ready[0]]
+        # Reset tasks to track just 1 task
+        tasks = [task]
+
+        if isinstance(task, DataOpTask):
+            task.on_data_ready(None, InlineMetadataFetcher())
+            if task.has_finished:
+                tasks.remove(task)
+        else:
+            assert isinstance(task, MetadataOpTask)
+            task.on_task_finished()
+            tasks.remove(task)
+
+
+def _get_blocks(bundle: RefBundle, output_list: List[Block]):
+    for block_ref in bundle.block_refs:
+        output_list.append(list(ray.get(block_ref)["id"]))
+
+
+def _mul2_transform(block_iter: Iterable[Block], ctx) -> Iterable[Block]:
+    for block in block_iter:
+        yield pd.DataFrame({"id": [b * 2 for b in block["id"]]})
+
+
+def create_map_transformer_from_block_fn(
+    block_fn: MapTransformCallable[Block, Block],
+    init_fn: Optional[Callable[[], None]] = None,
+    output_block_size_option: Optional[OutputBlockSizeOption] = None,
+    disable_block_shaping: bool = False,
+):
+    """Create a MapTransformer from a single block-based transform function.
+
+    This method should only be used for testing and legacy compatibility.
+    """
+    return MapTransformer(
+        [
+            BlockMapTransformFn(
+                block_fn,
+                output_block_size_option=output_block_size_option,
+                disable_block_shaping=disable_block_shaping,
+            ),
+        ],
+        init_fn=init_fn,
+    )
+
+
+def _take_outputs(op: PhysicalOperator) -> List[Any]:
+    output = []
+    while op.has_next():
+        ref = op.get_next()
+        assert ref.owns_blocks, ref
+        _get_blocks(ref, output)
+    return output