diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index 5fddabcf78df..79914d12fe8d 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -685,7 +685,9 @@ def evaluate_on_rank( Collected channel metadata. """ stats = allgather_stats(comm, ctx.br(), ir, config_options, py_executor) - ir, partition_info = lower_ir_graph(ir, config_options, stats) + ir, partition_info = lower_ir_graph( + ir, config_options, stats, rank=comm.rank, nranks=comm.nranks + ) if comm.rank == 0: # At least for now, the query plan is identical on all ranks, diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py index 837384a1561f..fb7cc4e1439a 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py @@ -15,7 +15,6 @@ from cudf_polars.dsl.ir import ( DataFrameScan, Join, - Scan, Union, ) from cudf_polars.dsl.traversal import CachingVisitor, traversal @@ -24,6 +23,7 @@ generate_ir_sub_network_wrapper, metadata_drain_node, ) +from cudf_polars.streaming.io import StreamingScan from cudf_polars.streaming.over import Over from cudf_polars.utils.config import SPMDContext @@ -246,7 +246,7 @@ def generate_network( num_io_nodes: int = 0 ir_dep_count: defaultdict[IR, int] = defaultdict(int) for node in traversal([ir]): - if isinstance(node, (DataFrameScan, Scan)): + if isinstance(node, (DataFrameScan, StreamingScan)): num_io_nodes += 1 for child in node.children: ir_dep_count[child] += 1 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 78672bb0cab0..7f4db5e38bb3 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py @@ -5,7 +5,6 @@ from __future__ import annotations import asyncio -import dataclasses import math from typing import TYPE_CHECKING, Any @@ -22,7 +21,6 @@ from cudf_polars.dsl.ir import ( IR, DataFrameScan, - Scan, Sink, _prepare_parquet_predicate, ) @@ -44,23 +42,20 @@ recv_metadata, send_metadata, ) -from cudf_polars.streaming.base import ( - IOPartitionFlavor, -) from cudf_polars.streaming.io import ( - SplitScan, + StreamingScan, StreamingSink, _prepare_sink_directory, _sink_to_file, + can_use_native_parquet_node, ) -from cudf_polars.streaming.utils import _dynamic_planning_on if TYPE_CHECKING: from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.streaming.core.channel import Channel from rapidsmpf.streaming.core.context import Context - from cudf_polars.dsl.ir import IR, IRExecutionContext + from cudf_polars.dsl.ir import IR, IRExecutionContext, Scan from cudf_polars.streaming.actor_graph.core import SubNetGenerator from cudf_polars.streaming.actor_graph.tracing import ActorTracer from cudf_polars.streaming.base import ( @@ -68,7 +63,7 @@ PartitionInfo, StatsCollector, ) - from cudf_polars.utils.config import ParquetOptions + from cudf_polars.streaming.io import SplitScan class Lineariser: @@ -356,14 +351,11 @@ async def read_chunk( @define_actor() async def scan_node( context: Context, - comm: Communicator, - ir: Scan, + ir: StreamingScan, ir_context: IRExecutionContext, ch_out: Channel[TableChunk], *, num_producers: int, - plan: IOPartitionPlan, - parquet_options: ParquetOptions, estimated_chunk_bytes: int, ) -> None: """ @@ -373,8 +365,6 @@ async def scan_node( ---------- context The rapidsmpf context. - comm - The communicator. ir The Scan node. ir_context @@ -383,84 +373,15 @@ async def scan_node( The output Channel[TableChunk]. num_producers The number of producers to use for the scan node. - plan - The partitioning plan. - parquet_options - The Parquet options. estimated_chunk_bytes Estimated size of each chunk in bytes. Used for memory reservation with block spilling to avoid thrashing. """ + scans = ir.scans + async with shutdown_on_error( context, ch_out, trace_ir=ir, ir_context=ir_context ) as tracer: - # Build a list of local Scan operations - scans: list[Scan | SplitScan] = [] - if plan.flavor == IOPartitionFlavor.SPLIT_FILES: - count = plan.factor * len(ir.paths) - local_count = math.ceil(count / comm.nranks) - local_offset = local_count * comm.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: - base_scan = Scan( - ir.schema, - ir.typ, - ir.reader_options, - ir.cloud_options, - [path], - ir.with_columns, - ir.skip_rows, - ir.n_rows, - ir.row_index, - ir.include_file_paths, - ir.predicate, - parquet_options, - ) - while sindex < plan.factor and splits_created < local_count: - scans.append( - SplitScan( - ir.schema, - base_scan, - sindex, - plan.factor, - parquet_options, - ) - ) - sindex += 1 - splits_created += 1 - sindex = 0 - - else: - count = math.ceil(len(ir.paths) / plan.factor) - local_count = math.ceil(count / comm.nranks) - local_offset = local_count * comm.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: # Only add scan if there are paths - scans.append( - Scan( - ir.schema, - ir.typ, - ir.reader_options, - ir.cloud_options, - local_paths, - ir.with_columns, - ir.skip_rows, - ir.n_rows, - ir.row_index, - ir.include_file_paths, - ir.predicate, - parquet_options, - ) - ) - # Send basic metadata await send_metadata( ch_out, @@ -628,9 +549,9 @@ def make_rapidsmpf_read_parquet_node( ) from e -@generate_ir_sub_network.register(Scan) +@generate_ir_sub_network.register(StreamingScan) def _( - ir: Scan, rec: SubNetGenerator + ir: StreamingScan, rec: SubNetGenerator ) -> tuple[dict[IR, list[Any]], dict[IR, ChannelManager]]: config_options = rec.state["config_options"] executor = config_options.executor @@ -642,39 +563,33 @@ def _( assert partition_info.io_plan is not None, "Scan node must have a partition plan" plan: IOPartitionPlan = partition_info.io_plan - # Native node cannot split large files in distributed mode yet - distributed_split_files = ( - plan.flavor == IOPartitionFlavor.SPLIT_FILES and rec.state["comm"].nranks > 1 - ) - # Use rapidsmpf native read_parquet node if possible ch_in: Channel[TableChunk] | None = None ch_out = channels[ir].reserve_input_slot() nodes: dict[IR, list[Any]] = {} native_node: Any = None - if ( - parquet_options.use_rapidsmpf_native - and (partition_info.count > 1 or _dynamic_planning_on(config_options)) - and ir.typ == "parquet" - and ir.row_index is None - and ir.include_file_paths is None - and ir.n_rows == -1 - and ir.skip_rows == 0 - and not distributed_split_files - ): + + use_native = can_use_native_parquet_node( + ir.base_scan, + plan=plan, + count=partition_info.count, + nranks=rec.state["comm"].nranks, + parquet_options=parquet_options, + config_options=config_options, + ) + if use_native: # Create new channel to so ch_out can be used to add metadata ch_in = rec.state["context"].create_channel() native_node = make_rapidsmpf_read_parquet_node( rec.state["context"], rec.state["comm"], - ir, + ir.base_scan, num_producers, ch_in, rec.state["stats"], partition_info, ) - if native_node is not None and ch_in is not None: # Need metadata node, because the native read_parquet # node does not send metadata. metadata_node = metadata_feeder_node( @@ -691,19 +606,13 @@ def _( ) nodes[ir] = [native_node, metadata_node] else: - # Fall back to scan_node (predicate not convertible, or other constraint) - parquet_options = dataclasses.replace(parquet_options, chunked=False) - nodes[ir] = [ scan_node( rec.state["context"], - rec.state["comm"], ir, rec.state["ir_context"], ch_out, num_producers=num_producers, - plan=plan, - parquet_options=parquet_options, estimated_chunk_bytes=executor.target_partition_size, ) ] diff --git a/python/cudf_polars/cudf_polars/streaming/dispatch.py b/python/cudf_polars/cudf_polars/streaming/dispatch.py index 81937e6ce6fd..3dd612f3de02 100644 --- a/python/cudf_polars/cudf_polars/streaming/dispatch.py +++ b/python/cudf_polars/cudf_polars/streaming/dispatch.py @@ -31,10 +31,17 @@ class State(TypedDict): GPUEngine configuration options. stats Statistics collector. + rank + Rank of the current worker for IO sharding. Always + 0 for non-streaming engines. + nranks + Number of workers for IO sharding. Always 1 for non-streaming engines. """ config_options: ConfigOptions[StreamingExecutor] stats: StatsCollector + rank: int + nranks: int LowerIRTransformer: TypeAlias = GenericTransformer[ diff --git a/python/cudf_polars/cudf_polars/streaming/explain.py b/python/cudf_polars/cudf_polars/streaming/explain.py index 92a76fa55209..67cdf311ba16 100644 --- a/python/cudf_polars/cudf_polars/streaming/explain.py +++ b/python/cudf_polars/cudf_polars/streaming/explain.py @@ -28,6 +28,7 @@ ) from cudf_polars.dsl.translate import Translator from cudf_polars.dsl.traversal import traversal +from cudf_polars.streaming.io import StreamingScan from cudf_polars.streaming.parallel import lower_ir_graph from cudf_polars.streaming.shuffle import Shuffle from cudf_polars.streaming.statistics import ( @@ -288,6 +289,18 @@ def _(ir: Scan) -> dict[str, Serializable]: } +@_serialize_properties.register +def _(ir: StreamingScan) -> dict[str, Serializable]: + return { + "typ": ir.base_scan.typ, + "scan_count": len(ir.scans), + "prefix": os.path.commonprefix(ir.base_scan.paths), + "predicate": ( + _serialize_expr(ir.base_scan.predicate) if ir.base_scan.predicate else None + ), + } + + @_serialize_properties.register def _(ir: Join) -> dict[str, Serializable]: return { diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 4472735ac356..d56c31ce379b 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -4,6 +4,7 @@ from __future__ import annotations +import dataclasses import functools import itertools import math @@ -31,6 +32,7 @@ SerializedDataSourceInfo, ) from cudf_polars.streaming.dispatch import lower_ir_node +from cudf_polars.streaming.utils import _dynamic_planning_on from cudf_polars.utils.config import Cluster from cudf_polars.utils.cuda_stream import get_cuda_stream from cudf_polars.utils.versions import POLARS_VERSION_LT_137 @@ -103,6 +105,104 @@ def scan_partition_plan( return IOPartitionPlan(1, IOPartitionFlavor.SINGLE_FILE) +def expand_scan_for_rank( + ir: Scan, + plan: IOPartitionPlan, + *, + rank: int, + nranks: int, + parquet_options: ParquetOptions, +) -> list[Scan | SplitScan]: + """ + Expand a Scan node into rank-local Scan and SplitScan operations. + + Parameters + ---------- + ir + The Scan node to expand. + plan + The IO partitioning plan for the scan. + rank + Rank of the current worker. + nranks + Number of workers. + parquet_options + Parquet reader options. + + Returns + ------- + list[Scan | SplitScan] + Rank-local scan operations. + """ + scans: list[Scan | SplitScan] = [] + 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: + base_scan = Scan( + ir.schema, + ir.typ, + ir.reader_options, + ir.cloud_options, + [path], + ir.with_columns, + ir.skip_rows, + ir.n_rows, + ir.row_index, + ir.include_file_paths, + ir.predicate, + parquet_options, + ) + while sindex < plan.factor and splits_created < local_count: + scans.append( + SplitScan( + ir.schema, + base_scan, + sindex, + plan.factor, + parquet_options, + ) + ) + sindex += 1 + splits_created += 1 + sindex = 0 + + 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: # Only add scan if there are paths + scans.append( + Scan( + ir.schema, + ir.typ, + ir.reader_options, + ir.cloud_options, + local_paths, + ir.with_columns, + ir.skip_rows, + ir.n_rows, + ir.row_index, + ir.include_file_paths, + ir.predicate, + parquet_options, + ) + ) + + return scans + + class SplitScan(IR): """ Input from a split file. @@ -260,11 +360,73 @@ def _( return ir, {ir: PartitionInfo(count=1)} # pragma: no cover +def can_use_native_parquet_node( + ir: Scan, + *, + plan: IOPartitionPlan, + count: int, + nranks: int, + parquet_options: ParquetOptions, + config_options: ConfigOptions[StreamingExecutor], +) -> bool: + """ + Determine whether we should use rapidsmpf's native parquet node. + + Parameters + ---------- + ir + The Scan node that might need to fall back. + plan + The IO partitioning plan. + count + The number of partitions associated with this Scan node. + nranks + The number of ranks. + parquet_options + The parquet options. + config_options + The configuration options. + + Returns + ------- + bool + Whether to use rapidsmpf's native parquet node. + + Notes + ----- + Native parquet node is used under the following conditions: + + - Our plan indicates we should split the file into multiple partitions + - We have more than one rank + - There's more than one partition or dynamic planning is enabled + - The file type is parquet + - The row index is not set + - File paths are not included + - The number of rows is not set + - The skip rows is not set + """ + distributed_split_files = ( + plan.flavor == IOPartitionFlavor.SPLIT_FILES and nranks > 1 + ) + + return ( + parquet_options.use_rapidsmpf_native + and (count > 1 or _dynamic_planning_on(config_options)) + and ir.typ == "parquet" + and ir.row_index is None + and ir.include_file_paths is None + and ir.n_rows == -1 + and ir.skip_rows == 0 + and not distributed_split_files + ) + + @lower_ir_node.register(Scan) def _( ir: Scan, rec: LowerIRTransformer ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: config_options = rec.state["config_options"] + parquet_options = config_options.parquet_options if ( ir.typ in ("csv", "parquet", "ndjson") and ir.n_rows == -1 @@ -282,13 +444,73 @@ def _( count = plan.factor * len(paths) else: count = math.ceil(len(paths) / plan.factor) - - return ir, {ir: PartitionInfo(count=count, io_plan=plan)} else: plan = IOPartitionPlan( flavor=IOPartitionFlavor.SINGLE_READ, factor=len(ir.paths) ) - return ir, {ir: PartitionInfo(count=1, io_plan=plan)} + count = 1 + + if not can_use_native_parquet_node( + ir, + plan=plan, + count=count, + nranks=rec.state["nranks"], + parquet_options=parquet_options, + config_options=config_options, + ): + parquet_options = dataclasses.replace(parquet_options, chunked=False) + + scans = expand_scan_for_rank( + ir, + plan, + 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)} + + +class StreamingScan(IR): + """A streaming scan node.""" + + __slots__ = ( + "base_scan", + "scans", + "schema", + ) + _non_child = ( + "scans", + "base_scan", + ) + _n_non_child_args = 2 + scans: list[Scan | SplitScan] + base_scan: Scan + + def __init__(self, scans: list[Scan | SplitScan], 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 = () + + 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. + return (type(self), *tuple(x.get_hashable() for x in self.scans)) + + @classmethod + def do_evaluate( + cls, + scans: list[Scan | SplitScan], + base_scan: Scan, + *, + context: IRExecutionContext, + ) -> DataFrame: + """Raises NotImplementedError for StreamingScan nodes.""" + raise NotImplementedError( + "StreamingScan.do_evaluate should not be called directly. Call Scan.do_evaluate on each scan node instead." + ) class StreamingSink(IR): diff --git a/python/cudf_polars/cudf_polars/streaming/parallel.py b/python/cudf_polars/cudf_polars/streaming/parallel.py index e5cdfcfb87cd..8d8d162f1111 100644 --- a/python/cudf_polars/cudf_polars/streaming/parallel.py +++ b/python/cudf_polars/cudf_polars/streaming/parallel.py @@ -68,6 +68,9 @@ def lower_ir_graph( ir: IR, config_options: ConfigOptions[StreamingExecutor], stats: StatsCollector, + *, + rank: int = 0, + nranks: int = 1, ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: """ Rewrite an IR graph and extract partitioning information. @@ -80,6 +83,10 @@ def lower_ir_graph( GPUEngine configuration options. stats Pre-computed statistics collector. + rank + Rank of the current worker. + nranks + Number of workers in the current cluster. Returns ------- @@ -99,6 +106,8 @@ def lower_ir_graph( state: State = { "config_options": config_options, "stats": stats, + "rank": rank, + "nranks": nranks, } mapper: LowerIRTransformer = CachingVisitor(lower_ir_node, state=state) return mapper(ir) diff --git a/python/cudf_polars/cudf_polars/streaming/select.py b/python/cudf_polars/cudf_polars/streaming/select.py index e0a97fdbafbb..e1466f0b6610 100644 --- a/python/cudf_polars/cudf_polars/streaming/select.py +++ b/python/cudf_polars/cudf_polars/streaming/select.py @@ -20,6 +20,7 @@ decompose_expr_graph, make_expr_decomposer, ) +from cudf_polars.streaming.io import StreamingScan from cudf_polars.streaming.over import _fuse_over_nodes from cudf_polars.streaming.repartition import Repartition from cudf_polars.streaming.utils import ( @@ -421,8 +422,11 @@ def _( ): # Task engine case scan_child = child.children[0] - elif isinstance(child, Scan): # pragma: no cover; Requires rapidsmpf runtime + elif isinstance(child, StreamingScan): # pragma: no cover # RapidsMPF case + scan_child = child.base_scan + elif isinstance(child, Scan): # pragma: no cover; Requires rapidsmpf runtime + # Legacy task-engine case scan_child = child if scan_child and scan_child.predicate is None and scan_child.typ == "parquet": diff --git a/python/cudf_polars/tests/streaming/test_explain.py b/python/cudf_polars/tests/streaming/test_explain.py index abbbc90ec616..60eb48bb310c 100644 --- a/python/cudf_polars/tests/streaming/test_explain.py +++ b/python/cudf_polars/tests/streaming/test_explain.py @@ -535,6 +535,7 @@ def test_scan_properties(tmp_path: Path, predicate: pl.Expr | None): "prefix": f"{root}/", "typ": "parquet", "predicate": None, + "scan_count": 1, } if predicate is not None: q = q.filter(predicate) @@ -551,7 +552,7 @@ def test_scan_properties(tmp_path: Path, predicate: pl.Expr | None): dag = serialize_query(q, engine) node = dag.nodes[dag.roots[0]] - assert node.type == "Scan" + assert node.type == "StreamingScan" assert node.properties == expected_properties diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 11afaac838ac..c9ccb13202dc 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -10,15 +10,20 @@ import polars as pl from cudf_polars import Translator +from cudf_polars.containers import DataType +from cudf_polars.dsl.ir import IRExecutionContext, Scan from cudf_polars.engine.options import StreamingOptions +from cudf_polars.streaming.base import IOPartitionFlavor, IOPartitionPlan +from cudf_polars.streaming.io import SplitScan, StreamingScan, expand_scan_for_rank from cudf_polars.streaming.parallel import lower_ir_graph from cudf_polars.streaming.statistics import collect_statistics from cudf_polars.testing.asserts import assert_gpu_result_equal from cudf_polars.testing.io import make_partitioned_source -from cudf_polars.utils.config import ConfigOptions +from cudf_polars.utils.config import ConfigOptions, ParquetOptions if TYPE_CHECKING: import concurrent.futures + from pathlib import Path @pytest.fixture(scope="module") @@ -125,3 +130,127 @@ def test_target_partition_size( assert count > n_files else: assert count < n_files + + +def test_scan_join(engine: pl.GPUEngine, tmp_path: Path) -> None: + # This test exercises some logic on nodes with multiple children (join) + # where one or more of the children are Scan nodes. + left = pl.DataFrame({"a": ["a", "b", "c", "d"], "b": [1, 2, 3, 4]}) + right = pl.DataFrame({"a": ["a", "b", "c", "d"], "c": [10, 20, 30, 40]}) + + left.write_parquet(tmp_path / "left.parquet") + right.write_parquet(tmp_path / "right.parquet") + + left_q = pl.scan_parquet(tmp_path / "left.parquet") + right_q = pl.scan_parquet(tmp_path / "right.parquet") + q = left_q.join(right_q, on="a", how="inner") + assert_gpu_result_equal(q, engine=engine) + + +def test_scan_union(engine: pl.GPUEngine, tmp_path: Path) -> None: + # This test exercises some logic on nodes with a Union[Scan, ...] + df = pl.DataFrame({"a": ["a", "b", "c", "d"], "b": [1, 2, 3, 4]}) + df.write_parquet(tmp_path / "data.parquet") + + df_q = pl.scan_parquet(tmp_path / "data.parquet") + + q = pl.concat([df_q, df_q]) + assert_gpu_result_equal(q, engine=engine) + + +def _make_parquet_scan(paths: list[str]) -> Scan: + return Scan( + {"x": DataType(pl.Int64())}, + "parquet", + {}, + None, + paths, + None, + 0, + -1, + None, + None, + None, + ParquetOptions(), + ) + + +@pytest.mark.parametrize( + "plan,paths,rank,nranks,expected_path_groups", + [ + ( + IOPartitionPlan(2, IOPartitionFlavor.FUSED_FILES), + [f"f{i}" for i in range(6)], + 0, + 1, + [["f0", "f1"], ["f2", "f3"], ["f4", "f5"]], + ), + ( + IOPartitionPlan(2, IOPartitionFlavor.FUSED_FILES), + [f"f{i}" for i in range(6)], + 0, + 2, + [["f0", "f1"], ["f2", "f3"]], + ), + ( + IOPartitionPlan(2, IOPartitionFlavor.FUSED_FILES), + [f"f{i}" for i in range(6)], + 1, + 2, + [["f4", "f5"]], + ), + (IOPartitionPlan(3, IOPartitionFlavor.SINGLE_READ), ["a", "b", "c"], 1, 2, []), + ], +) +def test_expand_scan_for_rank_fused_and_single_read( + plan: IOPartitionPlan, + paths: list[str], + rank: int, + nranks: int, + expected_path_groups: list[list[str]], +) -> None: + scans = expand_scan_for_rank( + _make_parquet_scan(paths), + plan, + rank=rank, + nranks=nranks, + parquet_options=ParquetOptions(), + ) + for scan, expected_paths in zip(scans, expected_path_groups, strict=True): + assert isinstance(scan, Scan) + assert scan.paths == expected_paths + + +@pytest.mark.parametrize( + "rank,expected_splits", + [ + (0, [(0, 4), (1, 4)]), + (1, [(2, 4), (3, 4)]), + ], +) +def test_expand_scan_for_rank_split_files( + rank: int, + expected_splits: list[tuple[int, int]], +) -> None: + plan = IOPartitionPlan(4, IOPartitionFlavor.SPLIT_FILES) + scans = expand_scan_for_rank( + _make_parquet_scan(["file.parquet"]), + plan, + 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 isinstance(scan, SplitScan) + assert scan.split_index == split_index + assert scan.total_splits == total_splits + assert scan.base_scan.paths == ["file.parquet"] + + +def test_streaming_scan_raises() -> None: + # This isn't reachable by normal cudf-polars usage. + scan = _make_parquet_scan(["file.parquet"]) + ctx = IRExecutionContext() + with pytest.raises(NotImplementedError, match=r"StreamingScan.do_evaluate"): + StreamingScan.do_evaluate([scan], scan, context=ctx) diff --git a/python/cudf_polars/tests/test_scan.py b/python/cudf_polars/tests/test_scan.py index 360116aace09..7a9a4f2bb10f 100644 --- a/python/cudf_polars/tests/test_scan.py +++ b/python/cudf_polars/tests/test_scan.py @@ -41,10 +41,10 @@ @pytest.fixture( - params=[(None, None), ("row-index", 0), ("index", 10)], + params=[(None, 0), ("row-index", 0), ("index", 10)], ids=["no_row_index", "zero_offset_row_index", "offset_row_index"], ) -def row_index(request): +def row_index(request) -> tuple[str | None, int]: return request.param @@ -475,13 +475,14 @@ def test_select_arbitrary_order_with_row_index_column(engine: pl.GPUEngine, tmp_ ) def test_scan_csv_with_and_without_header( engine: pl.GPUEngine, - df, - tmp_path, - has_header, - new_columns, - row_index, - columns, - zlice, + df: pl.DataFrame, + tmp_path: Path, + *, + has_header: bool, + new_columns: list[str] | None, + row_index: tuple[str | None, int], + columns: list[str] | None, + zlice: tuple[int, int] | None, ): path = tmp_path / "test.csv" make_partitioned_source(