chore: import upstream snapshot with attribution

python/ray/data/datasource/util.py

@@ -0,0 +1,109 @@
+from typing import Any, Dict, Iterable, Tuple
+
+import ray
+from ray.data.block import Block
+
+
+def _iter_sliced_blocks(
+    blocks: Iterable[Block], per_task_row_limit: int
+) -> Iterable[Block]:
+    """Iterate over blocks, accumulating rows up to the per-task row limit."""
+    rows_read = 0
+    for block in blocks:
+        if rows_read >= per_task_row_limit:
+            break
+
+        from ray.data.block import BlockAccessor
+
+        accessor = BlockAccessor.for_block(block)
+        block_rows = accessor.num_rows()
+
+        if rows_read + block_rows <= per_task_row_limit:
+            yield block
+            rows_read += block_rows
+        else:
+            # Slice the block to meet the limit exactly
+            remaining_rows = per_task_row_limit - rows_read
+            sliced_block = accessor.slice(0, remaining_rows, copy=True)
+            yield sliced_block
+            break
+
+
+def _validate_head_node_resources_for_local_scheduling(
+    ray_remote_args: Dict[str, Any],
+    *,
+    op_description: str,
+    default_num_cpus: int = 1,
+    default_num_gpus: int = 0,
+    default_memory: int = 0,
+) -> None:
+    """Ensure the head node has enough resources before pinning work there.
+
+    Local paths (``local://``) and other driver-local I/O schedule tasks on the
+    head node via ``NodeAffinitySchedulingStrategy``. If the head node was
+    intentionally started with zero logical resources (a common practice to
+    avoid OOMs), those tasks become unschedulable. Detect this upfront and
+    raise a clear error with remediation steps.
+    """
+
+    # Ray defaults to reserving 1 CPU per task when num_cpus isn't provided.
+    num_cpus = ray_remote_args.get("num_cpus", default_num_cpus)
+    num_gpus = ray_remote_args.get("num_gpus", default_num_gpus)
+    memory = ray_remote_args.get("memory", default_memory)
+
+    # Resource keys follow the Resources map of ray.nodes() (e.g., CPU, GPU, memory).
+    required_resources: Dict[str, float] = {}
+    required_resources["CPU"] = float(num_cpus)
+    required_resources["GPU"] = float(num_gpus)
+    required_resources["memory"] = float(memory)
+
+    # Include any additional custom resources requested.
+    custom_resources = ray_remote_args.get("resources", {})
+    for name, amount in custom_resources.items():
+        if amount is None:
+            continue
+        try:
+            amount = float(amount)
+        except (TypeError, ValueError) as err:
+            raise ValueError(f"Invalid resource amount for '{name}': {amount}") from err
+        required_resources[name] = amount
+
+    head_node = next(
+        (
+            node
+            for node in ray.nodes()
+            if node.get("Alive")
+            and "node:__internal_head__" in node.get("Resources", {})
+        ),
+        None,
+    )
+    if not head_node:
+        # The head node metadata is unavailable (e.g., during shutdown). Fall back
+        # to the default behavior and let Ray surface its own error.
+        return
+
+    # Build a map of required vs available resources on the head node.
+    head_resources: Dict[str, float] = head_node.get("Resources", {})
+    # Map: resource name -> (required, available).
+    insufficient: Dict[str, Tuple[float, float]] = {}
+    for name, req in required_resources.items():
+        avail = head_resources.get(name, 0.0)
+        if avail < req:
+            insufficient[name] = (req, avail)
+
+    # If nothing is below the required amount, we are good to proceed.
+    if not insufficient:
+        return
+
+    details = "; ".join(
+        f"{name} required {req:g} but head has {avail:g}"
+        for name, (req, avail) in insufficient.items()
+    )
+
+    raise ValueError(
+        f"{op_description} must run on the head node (e.g., for local:// paths), "
+        f"but the head node doesn't have enough resources: {details}. "
+        "Add resources to the head node, switch to a shared filesystem instead "
+        "of local://, or set the resource requests on this operation to 0 "
+        "(for example, num_cpus=0) so it can run without head resources."
+    )