diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py index 07ba5de31d00..ce7c2692c400 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py @@ -51,6 +51,8 @@ ) if TYPE_CHECKING: + from collections.abc import Sequence + from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.streaming.core.channel import Channel from rapidsmpf.streaming.core.context import Context @@ -370,7 +372,7 @@ async def scan_node( Estimated size of each chunk in bytes. Used for memory reservation with block spilling to avoid thrashing. """ - scans = ir.scans + scans: Sequence[SplitScan] | Sequence[FusedScan] = ir.scans async with shutdown_on_error( context, ch_out, trace_ir=ir, ir_context=ir_context @@ -413,7 +415,8 @@ async def scan_node( ] for task_idx, scan in enumerate(scans): producer_id = task_idx % num_producers - producer_tasks[producer_id].append((task_idx, scan)) + # mypy resolves __iter__ on union-of-sequences to the common base (IR) + producer_tasks[producer_id].append((task_idx, scan)) # type: ignore[arg-type] async def _producer(producer_id: int, ch_out: Channel) -> None: for task_idx, scan in producer_tasks[producer_id]: diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index e70d28d73863..e24226d7846b 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -11,7 +11,7 @@ import statistics from collections import defaultdict from pathlib import Path -from typing import TYPE_CHECKING, Any, Literal, overload +from typing import TYPE_CHECKING, Any, Literal, Self, overload import polars as pl @@ -38,7 +38,7 @@ from cudf_polars.utils.versions import POLARS_VERSION_LT_137 if TYPE_CHECKING: - from collections.abc import Hashable, MutableMapping + from collections.abc import Hashable, MutableMapping, Sequence from cudf_polars.containers import DataFrame from cudf_polars.dsl.expr import NamedExpr @@ -109,16 +109,23 @@ def scan_partition_plan( return IOPartitionPlan(1, IOPartitionFlavor.SINGLE_FILE) +def _rank_slice(total: int, rank: int, nranks: int) -> tuple[int, int]: + """Return the partition range owned by this rank.""" + count = math.ceil(total / nranks) + return count * rank, count + + def expand_scan_for_rank( ir: Scan, plan: IOPartitionPlan, + partition_count: int, *, rank: int, nranks: int, parquet_options: ParquetOptions, -) -> list[SplitScan | FusedScan]: +) -> StreamingScan: """ - Expand a Scan node into rank-local SplitScan and FusedScan operations. + Expand a Scan node into a rank-local StreamingScan. Parameters ---------- @@ -126,6 +133,8 @@ def expand_scan_for_rank( The Scan node to expand. plan The IO partitioning plan for the scan. + partition_count + Total number of partitions across all ranks. rank Rank of the current worker. nranks @@ -135,48 +144,27 @@ def expand_scan_for_rank( Returns ------- - list[SplitScan | FusedScan] - Rank-local scan operations. + StreamingScan + Rank-local streaming scan. """ - scans: list[SplitScan | FusedScan] = [] if plan.flavor == IOPartitionFlavor.SPLIT_FILES: - count = plan.factor * len(ir.paths) - local_count = math.ceil(count / nranks) - local_offset = local_count * rank - path_offset = local_offset // plan.factor - path_end = math.ceil((local_offset + local_count) / plan.factor) - path_count = path_end - path_offset - local_paths = ir.paths[path_offset : path_offset + path_count] - sindex = local_offset % plan.factor - splits_created = 0 - for path in local_paths: - while sindex < plan.factor and splits_created < local_count: - scans.append( - SplitScan( - ir.schema, - ir, - [path], - sindex, - plan.factor, - parquet_options, - ) - ) - sindex += 1 - splits_created += 1 - sindex = 0 - + return StreamingScan.for_split_files( + ir, + plan, + partition_count, + rank=rank, + nranks=nranks, + parquet_options=parquet_options, + ) else: - count = math.ceil(len(ir.paths) / plan.factor) - local_count = math.ceil(count / nranks) - local_offset = local_count * rank - paths_offset_start = local_offset * plan.factor - paths_offset_end = paths_offset_start + plan.factor * local_count - for offset in range(paths_offset_start, paths_offset_end, plan.factor): - local_paths = ir.paths[offset : offset + plan.factor] - if len(local_paths) > 0: - scans.append(FusedScan(ir.schema, ir, local_paths, parquet_options)) - - return scans + return StreamingScan.for_fused_files( + ir, + plan, + partition_count, + rank=rank, + nranks=nranks, + parquet_options=parquet_options, + ) class SplitScan(IR): @@ -556,14 +544,14 @@ def _( ): parquet_options = dataclasses.replace(parquet_options, chunked=False) - scans = expand_scan_for_rank( + new_ir = expand_scan_for_rank( ir, plan, + count, rank=rec.state["rank"], nranks=rec.state["nranks"], parquet_options=parquet_options, ) - new_ir = StreamingScan(scans, ir) return new_ir, {new_ir: PartitionInfo(count=count, io_plan=plan)} @@ -580,16 +568,81 @@ class StreamingScan(IR): "base_scan", ) _n_non_child_args = 2 - scans: list[SplitScan | FusedScan] + scans: Sequence[SplitScan] | Sequence[FusedScan] base_scan: Scan - def __init__(self, scans: list[SplitScan | FusedScan], base_scan: Scan): + def __init__( + self, scans: Sequence[SplitScan] | Sequence[FusedScan], base_scan: Scan + ): self.scans = scans self.base_scan = base_scan self.schema = base_scan.schema self._non_child_args = (scans, base_scan) self.children = () + @classmethod + def for_split_files( + cls, + base_scan: Scan, + plan: IOPartitionPlan, + partition_count: int, + *, + rank: int, + nranks: int, + parquet_options: ParquetOptions, + ) -> Self: + """Construct a StreamingScan where each file is split into factor partitions.""" + local_offset, local_count = _rank_slice(partition_count, rank, nranks) + path_offset = local_offset // plan.factor + path_end = math.ceil((local_offset + local_count) / plan.factor) + local_paths = base_scan.paths[path_offset:path_end] + sindex = local_offset % plan.factor + scans: list[SplitScan] = [] + splits_created = 0 + for path in local_paths: + while sindex < plan.factor and splits_created < local_count: + scans.append( + SplitScan( + base_scan.schema, + base_scan, + [path], + sindex, + plan.factor, + parquet_options, + ) + ) + sindex += 1 + splits_created += 1 + sindex = 0 + return cls(scans, base_scan) + + @classmethod + def for_fused_files( + cls, + base_scan: Scan, + plan: IOPartitionPlan, + partition_count: int, + *, + rank: int, + nranks: int, + parquet_options: ParquetOptions, + ) -> Self: + """Construct a StreamingScan where factor files are grouped into one partition.""" + local_offset, local_count = _rank_slice(partition_count, rank, nranks) + paths_start = local_offset * plan.factor + paths_end = paths_start + plan.factor * local_count + scans = [ + FusedScan( + base_scan.schema, + base_scan, + base_scan.paths[offset : offset + plan.factor], + parquet_options, + ) + for offset in range(paths_start, paths_end, plan.factor) + if base_scan.paths[offset : offset + plan.factor] + ] + return cls(scans, base_scan) + def get_hashable(self) -> Hashable: """Hashable representation of the node.""" # We don't need to include base_scan / schema, since it's in all the scan nodes. @@ -598,7 +651,7 @@ def get_hashable(self) -> Hashable: @classmethod def do_evaluate( cls, - scans: list[SplitScan | FusedScan], + scans: Sequence[SplitScan] | Sequence[FusedScan], base_scan: Scan, *, context: IRExecutionContext, diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index c40ce183c640..c1773a0e7bb1 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -3,6 +3,7 @@ from __future__ import annotations +import math from typing import TYPE_CHECKING import pytest @@ -231,14 +232,18 @@ def test_expand_scan_for_rank_fused_and_single_read( nranks: int, expected_path_groups: list[list[str]], ) -> None: - scans = expand_scan_for_rank( + partition_count = math.ceil(len(paths) / plan.factor) + streaming_scan = expand_scan_for_rank( _make_parquet_scan(paths), plan, + partition_count, rank=rank, nranks=nranks, parquet_options=ParquetOptions(), ) - for scan, expected_paths in zip(scans, expected_path_groups, strict=True): + for scan, expected_paths in zip( + streaming_scan.scans, expected_path_groups, strict=True + ): assert isinstance(scan, FusedScan) assert scan.paths == expected_paths @@ -255,15 +260,20 @@ def test_expand_scan_for_rank_split_files( expected_splits: list[tuple[int, int]], ) -> None: plan = IOPartitionPlan(4, IOPartitionFlavor.SPLIT_FILES) - scans = expand_scan_for_rank( - _make_parquet_scan(["file.parquet"]), + paths = ["file.parquet"] + partition_count = plan.factor * len(paths) + streaming_scan = expand_scan_for_rank( + _make_parquet_scan(paths), plan, + partition_count, rank=rank, nranks=2, parquet_options=ParquetOptions(), ) - assert len(scans) == len(expected_splits) - for scan, (split_index, total_splits) in zip(scans, expected_splits, strict=True): + assert len(streaming_scan.scans) == len(expected_splits) + for scan, (split_index, total_splits) in zip( + streaming_scan.scans, expected_splits, strict=True + ): assert isinstance(scan, SplitScan) assert scan.split_index == split_index assert scan.total_splits == total_splits