diff --git a/python/cudf_polars/cudf_polars/dsl/tracing.py b/python/cudf_polars/cudf_polars/dsl/tracing.py index 2c1dd1f26f0c..bb676c610990 100644 --- a/python/cudf_polars/cudf_polars/dsl/tracing.py +++ b/python/cudf_polars/cudf_polars/dsl/tracing.py @@ -56,6 +56,7 @@ class Scope(enum.StrEnum): PLAN = "plan" ACTOR = "actor" + IO_TASK = "io_task" EVALUATE_IR_NODE = "evaluate_ir_node" diff --git a/python/cudf_polars/cudf_polars/engine/options.py b/python/cudf_polars/cudf_polars/engine/options.py index 295a081307fa..611129cf2c01 100644 --- a/python/cudf_polars/cudf_polars/engine/options.py +++ b/python/cudf_polars/cudf_polars/engine/options.py @@ -223,6 +223,11 @@ class StreamingOptions: Env: ``CUDF_POLARS__EXECUTOR__NUM_PY_EXECUTORS``. Default: ``8``. Category: executor. + max_concurrent_io_tasks + Maximum concurrent IO tasks for each scan node. + Env: ``CUDF_POLARS__EXECUTOR__MAX_CONCURRENT_IO_TASKS``. + Default: ``2``. + Category: executor. fallback_mode Fallback behavior (``"warn"``, ``"raise"``, ``"silent"``). Env: ``CUDF_POLARS__EXECUTOR__FALLBACK_MODE``. @@ -340,6 +345,9 @@ class StreamingOptions: num_py_executors: int | Unspecified = _opt( "executor", "CUDF_POLARS__EXECUTOR__NUM_PY_EXECUTORS", int ) + max_concurrent_io_tasks: int | Unspecified = _opt( + "executor", "CUDF_POLARS__EXECUTOR__MAX_CONCURRENT_IO_TASKS", int + ) fallback_mode: str | Unspecified = _opt( "executor", "CUDF_POLARS__EXECUTOR__FALLBACK_MODE" ) @@ -532,6 +540,7 @@ def _get(attr: str) -> Any: unbounded_file_read_cache=_get("unbounded_file_read_cache"), hardware_binding=_get("hardware_binding"), num_py_executors=_get("num_py_executors"), + max_concurrent_io_tasks=_get("max_concurrent_io_tasks"), fallback_mode=_get("fallback_mode"), max_rows_per_partition=_get("max_rows_per_partition"), broadcast_limit=_get("broadcast_limit"), @@ -695,6 +704,16 @@ def _add_cli_args(parser: argparse.ArgumentParser) -> None: Env: CUDF_POLARS__EXECUTOR__NUM_PY_EXECUTORS. Built-in default: 8."""), ) + g.add_argument( + "--max-concurrent-io-tasks", + dest="max_concurrent_io_tasks", + default=None, + type=int, + help=textwrap.dedent("""\ + Maximum concurrent IO tasks for each scan node. + Env: CUDF_POLARS__EXECUTOR__MAX_CONCURRENT_IO_TASKS. + Built-in default: 2."""), + ) g.add_argument( "--raise-on-fail", dest="raise_on_fail", 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 a1e78f43d42c..fbec424d3104 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py @@ -13,7 +13,6 @@ import cudf_polars.dsl.tracing from cudf_polars.dsl.ir import ( - DataFrameScan, Join, Union, ) @@ -23,7 +22,6 @@ 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 @@ -250,22 +248,15 @@ def generate_network( ------- The network nodes and output hook. """ - # Count the number of IO nodes and the number of IR dependencies - num_io_nodes: int = 0 + # Count the number of IR dependencies ir_dep_count: defaultdict[IR, int] = defaultdict(int) for node in traversal([ir]): - if isinstance(node, (DataFrameScan, StreamingScan)): - num_io_nodes += 1 for child in node.children: ir_dep_count[child] += 1 # Determine which nodes need fanout fanout_nodes = determine_fanout_nodes(ir, partition_info, ir_dep_count) - # Get max_io_threads from config (default: 2) - max_io_threads_global = config_options.executor.max_io_threads - max_io_threads_local = max(1, max_io_threads_global // max(1, num_io_nodes)) - # Generate the network state: GenState = { "context": context, @@ -274,7 +265,7 @@ def generate_network( "partition_info": partition_info, "fanout_nodes": fanout_nodes, "ir_context": ir_context, - "max_io_threads": max_io_threads_local, + "max_concurrent_io_tasks": config_options.executor.max_concurrent_io_tasks, "stats": stats, "collective_id_map": collective_id_map, } diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py index 2554d95fe750..88d6784b0940 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Dispatching for the RapidsMPF streaming runtime.""" @@ -51,9 +51,8 @@ class GenState(TypedDict): Dictionary mapping IR nodes to fanout information. ir_context The execution context for the IR node. - max_io_threads - The maximum number of IO threads to use for - a single IO node. + max_concurrent_io_tasks + The maximum number of concurrent IO tasks to use for a single IO node. stats Statistics collector. collective_id_map @@ -66,7 +65,7 @@ class GenState(TypedDict): partition_info: MutableMapping[IR, PartitionInfo] fanout_nodes: dict[IR, FanoutInfo] ir_context: IRExecutionContext - max_io_threads: int + max_concurrent_io_tasks: int stats: StatsCollector collective_id_map: dict[IR, list[int]] 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 0a6a568f0e1e..e89035e12e2a 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py @@ -8,36 +8,22 @@ import functools import io import math +import time from typing import TYPE_CHECKING, Any, cast import polars as pl -import pylibcudf as plc from cudf_streaming.channel_metadata import ChannelMetadata from cudf_streaming.table_chunk import TableChunk from rapidsmpf.memory.memory_reservation import opaque_memory_usage -from rapidsmpf.streaming.core.memory_reserve_or_wait import ( - reserve_memory, -) +from rapidsmpf.streaming.core.memory_reserve_or_wait import reserve_memory from rapidsmpf.streaming.core.message import Message from cudf_polars.containers import DataFrame -from cudf_polars.dsl.ir import ( - IR, - DataFrameScan, - PythonScan, - Sink, - _prepare_parquet_predicate, -) -from cudf_polars.dsl.to_ast import to_parquet_filter -from cudf_polars.streaming.actor_graph.dispatch import ( - generate_ir_sub_network, -) -from cudf_polars.streaming.actor_graph.nodes import ( - define_actor, - metadata_feeder_node, - shutdown_on_error, -) +from cudf_polars.dsl.ir import IR, DataFrameScan, PythonScan, Sink +from cudf_polars.dsl.tracing import Scope, log +from cudf_polars.streaming.actor_graph.dispatch import generate_ir_sub_network +from cudf_polars.streaming.actor_graph.nodes import define_actor, shutdown_on_error from cudf_polars.streaming.actor_graph.tracing import send_chunk from cudf_polars.streaming.actor_graph.utils import ( ChannelManager, @@ -53,7 +39,6 @@ StreamingSink, _prepare_sink_directory, _sink_to_file, - can_use_native_parquet_node, ) from cudf_polars.streaming.rank_aware_source import RankAwareSource @@ -64,16 +49,14 @@ from rapidsmpf.streaming.core.channel import Channel from rapidsmpf.streaming.core.context import Context - from cudf_polars.dsl.ir import IR, IRExecutionContext, Scan + from cudf_polars.dsl.ir import IR, IRExecutionContext from cudf_polars.streaming.actor_graph.core import SubNetGenerator from cudf_polars.streaming.actor_graph.tracing import ActorTracer from cudf_polars.streaming.base import ( IOPartitionPlan, PartitionInfo, - StatsCollector, ) from cudf_polars.streaming.io import FusedScan, SplitScan - from cudf_polars.utils.config import ParquetOptions class Lineariser: @@ -297,7 +280,7 @@ def _( ) -> tuple[dict[IR, list[Any]], dict[IR, ChannelManager]]: config_options = rec.state["config_options"] rows_per_partition = config_options.executor.max_rows_per_partition - num_producers = rec.state["max_io_threads"] + num_producers = rec.state["max_concurrent_io_tasks"] # Use target_partition_size as the estimated chunk size estimated_chunk_bytes = config_options.executor.target_partition_size @@ -539,26 +522,47 @@ async def read_chunk( ir_context The execution context for the IR node. estimated_chunk_bytes - Estimated size of the chunk in bytes. Used for memory reservation - with block spilling to avoid thrashing. + Estimated retained output size in bytes. Used to estimate peak memory + for admission before launching the read. tracer The actor tracer for collecting runtime statistics. """ - with opaque_memory_usage( - await reserve_memory( - context, size=estimated_chunk_bytes, net_memory_delta=estimated_chunk_bytes - ) - ): + reservation_bytes = ( + estimated_chunk_bytes + if isinstance(scan, DataFrameScan) + else 2 * estimated_chunk_bytes + ) + start = time.monotonic_ns() + reservation = await reserve_memory( + context, + size=reservation_bytes, + net_memory_delta=estimated_chunk_bytes, + ) + admitted = time.monotonic_ns() + with opaque_memory_usage(reservation): df = await ir_context.to_thread( scan.do_evaluate, *scan._non_child_args, context=ir_context, ) - chunk = TableChunk.from_pylibcudf_table( - df.table, - df.stream, - exclusive_view=True, - br=context.br(), + chunk = TableChunk.from_pylibcudf_table( + df.table, + df.stream, + exclusive_view=True, + br=context.br(), + ) + stop = time.monotonic_ns() + log( + "IO Task", + scope=Scope.IO_TASK.value, + start=start, + admitted=admitted, + stop=stop, + ir_id=scan.get_stable_id(), + ir_type=type(scan).__name__, + sequence_number=seq_num, + estimated_output_bytes=estimated_chunk_bytes, + reservation_bytes=reservation_bytes, ) await send_chunk(context, ch_out, chunk, seq_num, tracer=tracer) @@ -589,8 +593,8 @@ async def scan_node( num_producers The number of producers to use for the scan node. estimated_chunk_bytes - Estimated size of each chunk in bytes. Used for memory reservation - with block spilling to avoid thrashing. + Estimated retained output size of each chunk in bytes. Used to estimate + peak memory for admission before launching each read. """ scans: Sequence[SplitScan] | Sequence[FusedScan] = ir.scans @@ -663,184 +667,34 @@ async def _producer(producer_id: int, ch_out: Channel) -> None: ) -def make_rapidsmpf_read_parquet_node( - context: Context, - comm: Communicator, - ir: Scan, - num_producers: int, - ch_out: Channel[TableChunk], - stats: StatsCollector, - partition_info: PartitionInfo, - parquet_options: ParquetOptions, -) -> Any | None: - """ - Make a RapidsMPF read parquet node. - - Parameters - ---------- - context - The rapidsmpf context. - comm - The communicator. - ir - The Scan node. - num_producers - The number of producers to use for the scan node. - ch_out - The output Channel[TableChunk]. - stats - The statistics collector. - partition_info - The partition information. - parquet_options - The Parquet options. - - Returns - ------- - The RapidsMPF read parquet node, or None if the predicate cannot be - converted to a parquet filter (caller should fall back to scan_node). - """ - from cudf_streaming.parquet import Filter, read_parquet - - # Build ParquetReaderOptions - try: - stream = context.br().stream_pool.get_stream() - builder = plc.io.parquet.ParquetReaderOptions.builder( - plc.io.SourceInfo(ir.paths) - ) - if ( - ir.predicate is not None and parquet_options.use_jit_filter - ): # pragma: no cover; no test yet - builder.use_jit_filter(use_jit_filter=True) - parquet_reader_options = builder.decimal_width(plc.TypeId.DECIMAL128).build() - - if ir.with_columns is not None: - parquet_reader_options.set_column_names(ir.with_columns) - - # Build predicate filter if present (passed separately to read_parquet) - filter_obj = None - if ir.predicate is not None: - filter_expr = to_parquet_filter( - _prepare_parquet_predicate( - ir.predicate.value, ir.paths, ir.schema, ir.with_columns - ), - stream=stream, - ) - if filter_expr is None: - # Predicate cannot be converted to parquet filter - # Return None to signal fallback to scan_node - return None - filter_obj = Filter(stream, filter_expr) - except Exception as e: - raise ValueError(f"Failed to build ParquetReaderOptions: {e}") from e - - # Calculate num_rows_per_chunk from statistics - # Default to a reasonable chunk size if statistics are unavailable - source = stats.scan_stats.get(ir) - estimated_row_count = source.row_count if source is not None else None - if estimated_row_count is not None: - num_rows_per_chunk = int(max(1, estimated_row_count // partition_info.count)) - else: - # Fallback: use a default chunk size if statistics are not available - num_rows_per_chunk = 1_000_000 # 1 million rows as default - - # Validate inputs - if num_rows_per_chunk <= 0: - raise ValueError(f"Invalid num_rows_per_chunk: {num_rows_per_chunk}") - if num_producers <= 0: - raise ValueError(f"Invalid num_producers: {num_producers}") - - try: - return read_parquet( - context, - comm, - ch_out, - num_producers, - parquet_reader_options, - num_rows_per_chunk, - filter=filter_obj, - ) - except Exception as e: - raise RuntimeError( - f"Failed to create read_parquet node: {e}\n" - f" paths: {ir.paths}\n" - f" num_producers: {num_producers}\n" - f" num_rows_per_chunk: {num_rows_per_chunk}\n" - f" partition_count: {partition_info.count}\n" - f" filter: {filter_obj}" - ) from e - - @generate_ir_sub_network.register(StreamingScan) def _( ir: StreamingScan, rec: SubNetGenerator ) -> tuple[dict[IR, list[Any]], dict[IR, ChannelManager]]: config_options = rec.state["config_options"] executor = config_options.executor - parquet_options = config_options.parquet_options partition_info = rec.state["partition_info"][ir] - num_producers = rec.state["max_io_threads"] + num_producers = rec.state["max_concurrent_io_tasks"] channels: dict[IR, ChannelManager] = {ir: ChannelManager(rec.state["context"])} assert partition_info.io_plan is not None, "Scan node must have a partition plan" plan: IOPartitionPlan = partition_info.io_plan - # 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 - - 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.base_scan, - num_producers, - ch_in, - rec.state["stats"], - partition_info, - parquet_options, - ) - # Need metadata node, because the native read_parquet - # node does not send metadata. - metadata_node = metadata_feeder_node( + nodes[ir] = [ + scan_node( rec.state["context"], ir, - ch_in, + rec.state["ir_context"], ch_out, - ChannelMetadata( - # partition_info.count is the estimated "global" count. - # Just estimate the local count as well. - local_count=math.ceil(partition_info.count / rec.state["comm"].nranks), + num_producers=num_producers, + estimated_chunk_bytes=( + plan.estimated_chunk_bytes or executor.target_partition_size ), - rec.state["ir_context"], ) - nodes[ir] = [native_node, metadata_node] - else: - nodes[ir] = [ - scan_node( - rec.state["context"], - ir, - rec.state["ir_context"], - ch_out, - num_producers=num_producers, - estimated_chunk_bytes=( - plan.estimated_chunk_bytes or executor.target_partition_size - ), - ) - ] + ] return nodes, channels diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index ce4f7c0d4b81..b07d400d35d0 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -480,8 +480,6 @@ class RunConfig: iterations: int io_mode: Literal["cold", "lukewarm", "hot"] = "lukewarm" collect_traces: bool = False - native_parquet: bool = True - max_io_threads: int = 2 # All streaming/rapidsmpf/engine knobs streaming_options: StreamingOptions = dataclasses.field( default_factory=lambda: __import__( @@ -629,8 +627,6 @@ def from_args(cls, args: argparse.Namespace) -> RunConfig: iterations=args.iterations, io_mode=args.io_mode, collect_traces=args.collect_traces, - native_parquet=args.native_parquet, - max_io_threads=args.max_io_threads, streaming_options=streaming_options, connect=args.connect, num_gpus=args.num_gpus, @@ -662,8 +658,6 @@ def serialize(self, engine: StreamingEngine | None) -> dict: "iterations": self.iterations, "io_mode": self.io_mode, "collect_traces": self.collect_traces, - "native_parquet": self.native_parquet, - "max_io_threads": self.max_io_threads, "n_workers": self.n_workers, "extra_info": self.extra_info, "run_id": str(self.run_id), @@ -713,7 +707,6 @@ def summarize(self) -> None: print(f"frontend: {self.frontend}") if self.frontend in _STREAMING_FRONTENDS: opts = self.streaming_options.to_executor_options() - print(f"native_parquet: {self.native_parquet}") print(f"n_workers: {self.n_workers}") print(f"target_partition_size: {opts.get('target_partition_size')}") print(f"broadcast_limit: {opts.get('broadcast_limit')}") @@ -744,7 +737,6 @@ def get_executor_options( executor_options: dict[str, Any] = ( run_config.streaming_options.to_executor_options() ) - executor_options["max_io_threads"] = run_config.max_io_threads executor_options["quent_context"] = cudf_polars.quent.QuentContext( engine=cudf_polars.quent.Engine(id=run_config.run_id) ) @@ -2040,18 +2032,6 @@ def build_parser(num_queries: int = 22) -> argparse.ArgumentParser: action="store_true", help="Debug run.", ) - parser.add_argument( - "--max-io-threads", - default=4, - type=int, - help="Sets cudf_polars.utils.config.StreamingExecutor.max_io_threads.", - ) - parser.add_argument( - "--native-parquet", - action=argparse.BooleanOptionalAction, - default=False, - help="Sets cudf_polars.utils.config.ParquetOptions.use_rapidsmpf_native.", - ) parser.add_argument( "-o", "--output", @@ -2278,7 +2258,7 @@ def run_polars(benchmark: Any, args: argparse.Namespace) -> None: "Unset CUDA_VISIBLE_DEVICES or use it directly to control GPU visibility." ) - parquet_options = {"use_rapidsmpf_native": run_config.native_parquet} + parquet_options: dict[str, Any] = {} numeric_type, date_type = check_input_data_type(run_config) match args.frontend: case "dask": diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 5464b586f661..4e061e6a3cb2 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -33,7 +33,6 @@ 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 @@ -500,67 +499,6 @@ def _( return ir, {ir: PartitionInfo(count=2)} -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 @@ -590,15 +528,7 @@ def _( ) 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) + parquet_options = dataclasses.replace(parquet_options, chunked=False) new_ir = expand_scan_for_rank( ir, diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 63fc734a0ca7..185f43b7142f 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -219,10 +219,6 @@ class ParquetOptions: Set to 0 to avoid row-group sampling. Note that row-group sampling will also be skipped if ``max_footer_samples`` is 0. - use_rapidsmpf_native - Whether to use the native rapidsmpf node for parquet reading. - This option is only used by the streaming executor. - Default is False. prefetch_file_metadata Whether to prefetch parquet file metadata and pass it through `parquet_metadatas` to avoid rereading file footers. @@ -265,13 +261,6 @@ class ParquetOptions: f"{_env_prefix}__MAX_ROW_GROUP_SAMPLES", int, default=1 ) ) - use_rapidsmpf_native: bool = dataclasses.field( - default_factory=_make_default_factory( - f"{_env_prefix}__USE_RAPIDSMPF_NATIVE", - _bool_converter, - default=False, - ) - ) prefetch_file_metadata: bool = dataclasses.field( default_factory=_make_default_factory( f"{_env_prefix}__PREFETCH_FILE_METADATA", @@ -300,15 +289,8 @@ def __post_init__(self) -> None: # noqa: D105 raise TypeError("max_footer_samples must be an int") if not isinstance(self.max_row_group_samples, int): raise TypeError("max_row_group_samples must be an int") - if not isinstance(self.use_rapidsmpf_native, bool): - raise TypeError("use_rapidsmpf_native must be a bool") if not isinstance(self.prefetch_file_metadata, bool): raise TypeError("prefetch_file_metadata must be a bool") - - if self.use_rapidsmpf_native and self.prefetch_file_metadata: - raise NotImplementedError( - "'use_rapidsmpf_native=True' does not currently support 'prefetch_file_metadata=True'" - ) if not isinstance(self.use_jit_filter, bool): raise TypeError("use_jit_filter must be a bool") @@ -666,7 +648,7 @@ class StreamingExecutor: This can be set via - - keyword argument to ``polars.GPUEngine`` + - ``executor_options`` passed to ``polars.GPUEngine`` - the ``CUDF_POLARS__EXECUTOR__TARGET_PARTITION_SIZE`` environment variable By default, cudf-polars uses the minimum of 1.5GB or 2.5% of the minimum @@ -694,9 +676,12 @@ class StreamingExecutor: Enable through environment variables with ``CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN=1``. - max_io_threads - Maximum number of IO threads. Default is 4. - This controls the parallelism of IO operations when reading data. + max_concurrent_io_tasks + Maximum number of concurrent IO tasks for each scan node. Default is 2. + This can be set via + + - ``executor_options`` passed to ``polars.GPUEngine`` + - the ``CUDF_POLARS__EXECUTOR__MAX_CONCURRENT_IO_TASKS`` environment variable num_py_executors Maximum number of workers for the Python ThreadPoolExecutor. Default is 8. @@ -761,9 +746,9 @@ class StreamingExecutor: join_filter_pushdown: JoinFilterPushdownOptions | None = dataclasses.field( default_factory=JoinFilterPushdownOptions ) - max_io_threads: int = dataclasses.field( + max_concurrent_io_tasks: int = dataclasses.field( default_factory=_make_default_factory( - f"{_env_prefix}__MAX_IO_THREADS", int, default=4 + f"{_env_prefix}__MAX_CONCURRENT_IO_TASKS", int, default=2 ) ) num_py_executors: int = dataclasses.field( @@ -851,8 +836,8 @@ def __post_init__(self) -> None: # noqa: D105 raise TypeError("sink_to_directory must be bool") if not isinstance(self.client_device_threshold, float): raise TypeError("client_device_threshold must be a float") - if not isinstance(self.max_io_threads, int): - raise TypeError("max_io_threads must be an int") + if not isinstance(self.max_concurrent_io_tasks, int): + raise TypeError("max_concurrent_io_tasks must be an int") if not isinstance(self.num_py_executors, int): raise TypeError("num_py_executors must be an int") diff --git a/python/cudf_polars/docs/cudf-polars-mp.md b/python/cudf_polars/docs/cudf-polars-mp.md index 657903acd9e0..0aaf27d073d8 100644 --- a/python/cudf_polars/docs/cudf-polars-mp.md +++ b/python/cudf_polars/docs/cudf-polars-mp.md @@ -749,7 +749,7 @@ with SPMDEngine( rapidsmpf_options=Options(num_streaming_threads=8), executor_options={"num_py_executors": 2}, executor_options={"max_rows_per_partition": 500_000}, - engine_options={"parquet_options": {"use_rapidsmpf_native": True}}, + engine_options={"parquet_options": {}}, ) as engine: ... ``` @@ -771,9 +771,7 @@ to `None` (uses RapidsMPF defaults). `executor_options` is forwarded directly to `pl.GPUEngine` as its `executor_options` argument; user-supplied keys are merged with reserved entries set by `SPMDEngine`. -`engine_options` is forwarded as keyword arguments to `pl.GPUEngine`. For example, -pass `engine_options={"parquet_options": {"use_rapidsmpf_native": True}}` to enable -native Parquet reads. +`engine_options` is forwarded as keyword arguments to `pl.GPUEngine`. [dask-cli]: https://docs.dask.org/en/latest/deploying-cli.html diff --git a/python/cudf_polars/docs/overview.md b/python/cudf_polars/docs/overview.md index cd7911f82120..1a5a51d206f8 100644 --- a/python/cudf_polars/docs/overview.md +++ b/python/cudf_polars/docs/overview.md @@ -410,6 +410,23 @@ engine = pl.GPUEngine( ) ``` +Each scan node may run up to `max_concurrent_io_tasks` reads concurrently. The +limit applies independently to each scan node, each corresponding to a +single `pl.scan_parquet` call in the query. Configure it through +`executor_options` or +`CUDF_POLARS__EXECUTOR__MAX_CONCURRENT_IO_TASKS`: + +```python +engine = pl.GPUEngine( + executor="streaming", + executor_options={"max_concurrent_io_tasks": 8}, +) +``` + +Before each read is submitted, it waits for a device-memory reservation. +This makes aggregate read concurrency respond to memory pressure across all +scan nodes on the rank. + Internally, `collect_statistics` walks the IR graph, groups Parquet `Scan` nodes that share the same file paths (unioning projected columns for sampling), and builds one `DataSourceInfo` per path group. It then diff --git a/python/cudf_polars/tests/streaming/test_options.py b/python/cudf_polars/tests/streaming/test_options.py index d1a5a54aafa2..d422c92d5a65 100644 --- a/python/cudf_polars/tests/streaming/test_options.py +++ b/python/cudf_polars/tests/streaming/test_options.py @@ -73,6 +73,11 @@ def test_executor_options_num_py_executors() -> None: assert result["num_py_executors"] == 4 +def test_executor_options_max_concurrent_io_tasks() -> None: + result = StreamingOptions(max_concurrent_io_tasks=6).to_executor_options() + assert result["max_concurrent_io_tasks"] == 6 + + @pytest.mark.parametrize("value", [True, False]) def test_executor_options_sink_to_directory(*, value: bool) -> None: result = StreamingOptions(sink_to_directory=value).to_executor_options() @@ -305,6 +310,8 @@ def test_add_cli_args_then_from_argparse_roundtrip() -> None: "4GiB", "--unbounded-file-read-cache", "host", + "--max-concurrent-io-tasks", + "6", ] ) opts = StreamingOptions._from_argparse(args) @@ -313,6 +320,7 @@ def test_add_cli_args_then_from_argparse_roundtrip() -> None: assert opts.raise_on_fail is True assert opts.pinned_max_pool_size == "4GiB" assert opts.unbounded_file_read_cache == "host" + assert opts.max_concurrent_io_tasks == 6 # Unprovided args default to None → UNSPECIFIED assert isinstance(opts.fallback_mode, Unspecified) diff --git a/python/cudf_polars/tests/streaming/test_parallel.py b/python/cudf_polars/tests/streaming/test_parallel.py index e32fab2ce3a3..7940bcda0a59 100644 --- a/python/cudf_polars/tests/streaming/test_parallel.py +++ b/python/cudf_polars/tests/streaming/test_parallel.py @@ -81,7 +81,7 @@ def test_evaluate_streaming(streaming_engine): df = pl.LazyFrame({"a": [1, 2, 3], "b": [3, 4, 5], "c": [5, 6, 7], "d": [7, 9, 8]}) q = df.select(pl.col("a") - (pl.col("b") + pl.col("c") * 2), pl.col("d")).sort("d") - expected = q.collect(engine="cpu") + expected = q.collect(engine="in-memory") # Use the in-memory executor for the GPU comparison: a vanilla # ``pl.GPUEngine(raise_on_fail=True)`` defaults to ``executor="streaming"``, # which routes through ``DefaultSingletonEngine`` and would fail while the @@ -116,11 +116,6 @@ def test_optimize_removes_cache_nodes() -> None: assert not any(isinstance(node, Cache) for node in traversal([optimized])) -# --------------------------------------------------------------------------- -# Tests migrated from tests/streaming/test_parallel.py (round 3) -# --------------------------------------------------------------------------- - - @pytest.mark.parametrize( "agg", [ diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 9cb44ce6c0a7..0fcf3c663a9f 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -92,17 +92,6 @@ def test_parallel_scan( assert_gpu_result_equal(q, engine=streaming_engine) -def test_scan_parquet_use_rapidsmpf_native(tmp_path, df, streaming_engine_factory): - streaming_engine = streaming_engine_factory( - StreamingOptions( - target_partition_size=1_000, - parquet_options={"use_rapidsmpf_native": True}, - ), - ) - make_partitioned_source(df, tmp_path, "parquet", n_files=1) - assert_gpu_result_equal(pl.scan_parquet(tmp_path), engine=streaming_engine) - - @pytest.mark.parametrize( "target_partition_size_and_n_files", [(1_000, 1), (1_000, 2), (1_000_000, 5)] ) diff --git a/python/cudf_polars/tests/streaming/test_tracing.py b/python/cudf_polars/tests/streaming/test_tracing.py index f584ceff6528..a14ff016e985 100644 --- a/python/cudf_polars/tests/streaming/test_tracing.py +++ b/python/cudf_polars/tests/streaming/test_tracing.py @@ -5,6 +5,7 @@ from __future__ import annotations import asyncio +import json import os import subprocess import sys @@ -21,6 +22,8 @@ from cudf_polars.streaming.actor_graph.tracing import ActorTracer, send_chunk if TYPE_CHECKING: + import pathlib + from cudf_polars.engine.spmd import SPMDEngine @@ -69,10 +72,8 @@ def test_structlog_streaming_node_events(timeout_seconds: int): """Test that structlog emits 'Streaming Actor' events when tracing is enabled.""" pytest.importorskip("structlog") code = textwrap.dedent("""\ - import rmm import polars as pl - rmm.mr.set_current_device_resource(rmm.mr.ManagedMemoryResource()) from cudf_polars.engine.spmd import SPMDEngine df = pl.DataFrame({"x": range(100), "y": ["a", "b"] * 50}) @@ -103,10 +104,8 @@ def test_structlog_contains_expected_ir_types(timeout_seconds: int): """Test that structlog output contains expected IR types for a query.""" pytest.importorskip("structlog") code = textwrap.dedent("""\ - import rmm import polars as pl - rmm.mr.set_current_device_resource(rmm.mr.ManagedMemoryResource()) from cudf_polars.engine.spmd import SPMDEngine df = pl.DataFrame({"x": range(100), "y": ["a", "b"] * 50}) @@ -131,14 +130,87 @@ def test_structlog_contains_expected_ir_types(timeout_seconds: int): assert b"ir_type=GroupBy" in result +def test_io_tasks_wait_for_memory_admission( + tmp_path: pathlib.Path, timeout_seconds: int +) -> None: + pytest.importorskip("structlog") + + source = tmp_path / "data.parquet" + pl.DataFrame({"x": range(5_000)}).write_parquet( + source, + compression="uncompressed", + row_group_size=2_500, + ) + + code = textwrap.dedent(f"""\ + import structlog + import polars as pl + + structlog.configure( + processors=[ + structlog.contextvars.merge_contextvars, + structlog.processors.JSONRenderer(), + ] + ) + from cudf_polars.engine.options import StreamingOptions + from cudf_polars.engine.spmd import SPMDEngine + + q = pl.scan_parquet("{source}").select(pl.col("x").sum()) + options = StreamingOptions( + allow_overbooking_by_default=False, + max_concurrent_io_tasks=2, + memory_reserve_timeout="10s", + spill_device_limit="65000", + target_partition_size=21_000, + ) + with SPMDEngine.from_options(options) as engine: + q.collect(engine=engine) + """) + + env = os.environ.copy() + env["CUDF_POLARS_LOG_TRACES"] = "1" + env["CUDF_POLARS_LOG_TRACES_MEMORY"] = "0" + + with subprocess.Popen( + [sys.executable, "-c", code], + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + ) as proc: + result, _ = proc.communicate(timeout=timeout_seconds) + returncode = proc.returncode + + assert returncode == 0, result.decode(errors="replace") + + events = [] + for line in result.splitlines(): + try: + event = json.loads(line) + except json.JSONDecodeError: + continue + if event.get("event") == "IO Task": + events.append(event) + + assert len(events) == 2, result.decode(errors="replace") + assert all(event["scope"] == "io_task" for event in events) + assert all(event["ir_type"] == "SplitScan" for event in events) + assert all( + event["reservation_bytes"] == 2 * event["estimated_output_bytes"] + for event in events + ) + + first, second = sorted(events, key=lambda event: event["admitted"]) + assert first["start"] <= first["admitted"] <= first["stop"] + assert second["start"] <= second["admitted"] <= second["stop"] + assert second["admitted"] >= first["stop"] + + def test_structlog_disabled_by_default(timeout_seconds: int): """Test that structlog does NOT emit events when CUDF_POLARS_LOG_TRACES is not set.""" pytest.importorskip("structlog") code = textwrap.dedent("""\ - import rmm import polars as pl - rmm.mr.set_current_device_resource(rmm.mr.ManagedMemoryResource()) from cudf_polars.engine.spmd import SPMDEngine df = pl.DataFrame({"x": range(10), "y": ["a", "b"] * 5}) diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index aa01207e7e51..dc726970101a 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -323,7 +323,7 @@ def test_validate_cluster() -> None: "broadcast_limit", "sink_to_directory", "client_device_threshold", - "max_io_threads", + "max_concurrent_io_tasks", "num_py_executors", ], ) @@ -359,7 +359,6 @@ def test_parquet_options_from_env(monkeypatch: pytest.MonkeyPatch) -> None: m.setenv("CUDF_POLARS__PARQUET_OPTIONS__PASS_READ_LIMIT", "200") m.setenv("CUDF_POLARS__PARQUET_OPTIONS__MAX_FOOTER_SAMPLES", "0") m.setenv("CUDF_POLARS__PARQUET_OPTIONS__MAX_ROW_GROUP_SAMPLES", "0") - m.setenv("CUDF_POLARS__PARQUET_OPTIONS__USE_RAPIDSMPF_NATIVE", "0") m.setenv("CUDF_POLARS__PARQUET_OPTIONS__PREFETCH_FILE_METADATA", "1") m.setenv("CUDF_POLARS__PARQUET_OPTIONS__USE_JIT_FILTER", "1") @@ -372,7 +371,6 @@ def test_parquet_options_from_env(monkeypatch: pytest.MonkeyPatch) -> None: assert config.parquet_options.pass_read_limit == 200 assert config.parquet_options.max_footer_samples == 0 assert config.parquet_options.max_row_group_samples == 0 - assert config.parquet_options.use_rapidsmpf_native is False assert config.parquet_options.prefetch_file_metadata is True assert config.parquet_options.use_jit_filter is True @@ -390,6 +388,7 @@ def test_config_option_from_env(monkeypatch: pytest.MonkeyPatch) -> None: m.setenv("CUDF_POLARS__EXECUTOR__MAX_ROWS_PER_PARTITION", "42") m.setenv("CUDF_POLARS__EXECUTOR__TARGET_PARTITION_SIZE", "100") m.setenv("CUDF_POLARS__EXECUTOR__BROADCAST_LIMIT", "44") + m.setenv("CUDF_POLARS__EXECUTOR__MAX_CONCURRENT_IO_TASKS", "6") m.setenv("CUDF_POLARS__EXECUTOR__QUENT_CONTEXT", "1") engine = pl.GPUEngine() @@ -400,6 +399,7 @@ def test_config_option_from_env(monkeypatch: pytest.MonkeyPatch) -> None: assert config.executor.max_rows_per_partition == 42 assert config.executor.target_partition_size == 100 assert config.executor.broadcast_limit == 44 + assert config.executor.max_concurrent_io_tasks == 6 assert config.executor.quent_context is not None @@ -478,7 +478,6 @@ def test_fallback_mode_default(monkeypatch: pytest.MonkeyPatch) -> None: "pass_read_limit", "max_footer_samples", "max_row_group_samples", - "use_rapidsmpf_native", "prefetch_file_metadata", "use_jit_filter", ], @@ -493,22 +492,6 @@ def test_validate_parquet_options(option: str) -> None: ) -def test_prefetch_and_use_rapidsmpf_native_raises() -> None: - with pytest.raises( - NotImplementedError, - match="'use_rapidsmpf_native=True' does not currently support 'prefetch_file_metadata=True'", - ): - ConfigOptions.from_polars_engine( - pl.GPUEngine( - executor="streaming", - parquet_options={ - "use_rapidsmpf_native": True, - "prefetch_file_metadata": True, - }, - ) - ) - - def test_validate_raise_on_fail() -> None: with pytest.raises(TypeError, match="'raise_on_fail' must be"): ConfigOptions.from_polars_engine( diff --git a/python/cudf_streaming/CMakeLists.txt b/python/cudf_streaming/CMakeLists.txt index 8eac3cbcd91b..bb1dcf94fa4b 100644 --- a/python/cudf_streaming/CMakeLists.txt +++ b/python/cudf_streaming/CMakeLists.txt @@ -40,7 +40,6 @@ set(cython_sources cudf_streaming/approx_distinct_count.pyx cudf_streaming/bloom_filter.pyx cudf_streaming/channel_metadata.pyx - cudf_streaming/parquet.pyx cudf_streaming/partition.pyx cudf_streaming/partition_utils.pyx cudf_streaming/table_chunk.pyx diff --git a/python/cudf_streaming/cudf_streaming/__init__.py b/python/cudf_streaming/cudf_streaming/__init__.py index 876b197b4a37..e9279370ba2e 100644 --- a/python/cudf_streaming/cudf_streaming/__init__.py +++ b/python/cudf_streaming/cudf_streaming/__init__.py @@ -25,7 +25,6 @@ OrderScheme, Partitioning, ) -from cudf_streaming.parquet import Filter, read_parquet from cudf_streaming.partition import ( partition_and_pack as actor_partition_and_pack, unpack_and_concat as actor_unpack_and_concat, @@ -46,7 +45,6 @@ "CardinalityEstimate", "CardinalityEstimator", "ChannelMetadata", - "Filter", "HashScheme", "OrderKey", "OrderScheme", @@ -58,7 +56,6 @@ "make_table_chunks_available_or_wait", "packed_data_from_cudf_packed_columns", "partition_and_pack", - "read_parquet", "split_and_pack", "unpack_and_concat", ] diff --git a/python/cudf_streaming/cudf_streaming/parquet.pyi b/python/cudf_streaming/cudf_streaming/parquet.pyi deleted file mode 100644 index ebfb59555bae..000000000000 --- a/python/cudf_streaming/cudf_streaming/parquet.pyi +++ /dev/null @@ -1,25 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from pylibcudf.expressions import Expression -from pylibcudf.io.parquet import ParquetReaderOptions - -from cudf_streaming.table_chunk import TableChunk -from rapidsmpf.communicator.communicator import Communicator -from rapidsmpf.streaming.core.actor import CppActor -from rapidsmpf.streaming.core.channel import Channel -from rapidsmpf.streaming.core.context import Context -from rmm.pylibrmm.stream import Stream - -class Filter: - def __init__(self, stream: Stream, filter: Expression) -> None: ... - -def read_parquet( - ctx: Context, - comm: Communicator, - ch_out: Channel[TableChunk], - num_producers: int, - options: ParquetReaderOptions, - num_rows_per_chunk: int, - filter: Filter | None = None, -) -> CppActor: ... diff --git a/python/cudf_streaming/cudf_streaming/parquet.pyx b/python/cudf_streaming/cudf_streaming/parquet.pyx deleted file mode 100644 index af66e45e5f86..000000000000 --- a/python/cudf_streaming/cudf_streaming/parquet.pyx +++ /dev/null @@ -1,145 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from cpython.object cimport PyObject -from cpython.ref cimport Py_INCREF -from cython.operator cimport dereference as deref -from libc.stddef cimport size_t -from libcpp.memory cimport make_unique, shared_ptr, unique_ptr -from libcpp.utility cimport move -from pylibcudf.expressions cimport Expression -from pylibcudf.io.parquet cimport ParquetReaderOptions -from pylibcudf.libcudf.expressions cimport expression -from pylibcudf.libcudf.io.parquet cimport parquet_reader_options -from pylibcudf.libcudf.types cimport size_type -from rmm.librmm.cuda_stream_view cimport cuda_stream_view -from rmm.pylibrmm.stream cimport Stream - -from rapidsmpf._detail.exception_handling cimport ex_handler -from rapidsmpf.communicator.communicator cimport Communicator, cpp_Communicator -from rapidsmpf.streaming.chunks.arbitrary cimport cpp_OwningWrapper -from rapidsmpf.streaming.chunks.utils cimport py_deleter -from rapidsmpf.streaming.core.actor cimport CppActor, cpp_Actor -from rapidsmpf.streaming.core.channel cimport Channel, cpp_Channel -from rapidsmpf.streaming.core.context cimport Context, cpp_Context - - -cdef extern from "" nogil: - cdef cppclass cpp_Filter "cudf_streaming::filter": - cpp_Filter(cuda_stream_view, expression, cpp_OwningWrapper) - - cdef cpp_Actor cpp_read_parquet \ - "cudf_streaming::actor::read_parquet"( - shared_ptr[cpp_Context] ctx, - shared_ptr[cpp_Communicator] comm, - shared_ptr[cpp_Channel] ch_out, - size_t num_producers, - parquet_reader_options options, - size_type num_rows_per_chunk, - unique_ptr[cpp_Filter], - ) except +ex_handler - - -cdef class Filter: - """ - A filter expression for parquet reads. - - Parameters - ---------- - stream - The stream any scalars in the expression are valid on. - expression - The filter expression - - Notes - ----- - The object safely manages the lifetime of the expressions when called - from C++ coroutines, so it is safe to drop the expression passed in on - the python side. - """ - cdef unique_ptr[cpp_Filter] _handle - - def __init__(self, Stream stream not None, Expression filter not None): - Py_INCREF(filter) - self._handle = make_unique[cpp_Filter]( - stream.view(), - deref(filter.c_obj), - cpp_OwningWrapper( - filter, py_deleter - ) - ) - - cdef unique_ptr[cpp_Filter] release_handle(self): - """ - Move the owning C++ handle out of the object. - - Returns - ------- - unique_ptr to the C++ Filter object. - - Raises - ------ - ValueError - If this Filter has already been used and the handle is already released. - """ - if not self._handle: - raise ValueError("Filter is uninitialized, has it been released?") - return move(self._handle) - - def __dealloc__(self): - with nogil: - self._handle.reset() - - -def read_parquet( - Context ctx not None, - Communicator comm not None, - Channel ch_out not None, - size_t num_producers, - ParquetReaderOptions options not None, - size_type num_rows_per_chunk, - Filter filter = None, -): - """ - Create a streaming actor to read from parquet. - - Parameters - ---------- - ctx - Streaming execution context. - comm - The communicator. - ch_out - Output channel to receive the TableChunks. - num_producers - Number of concurrent producers of output chunks. - options - Reader options. - num_rows_per_chunk - Target (maximum) number of rows per output chunk. - filter - Optional filter object. If provided, is consumed by this function - and not subsequently usable. - - Notes - ----- - This is a collective operation, all ranks participating via the - communicator must call it with the same options. - """ - cdef cpp_Actor _ret - cdef unique_ptr[cpp_Filter] c_filter - if filter is not None: - c_filter = move(filter.release_handle()) - with nogil: - _ret = cpp_read_parquet( - ctx._handle, - comm._handle, - ch_out._handle, - num_producers, - options.c_obj, - num_rows_per_chunk, - move(c_filter) - ) - return CppActor.from_handle( - make_unique[cpp_Actor](move(_ret)), owner=None - ) diff --git a/python/cudf_streaming/cudf_streaming/tests/test_read_parquet.py b/python/cudf_streaming/cudf_streaming/tests/test_read_parquet.py deleted file mode 100644 index 422e43d2061c..000000000000 --- a/python/cudf_streaming/cudf_streaming/tests/test_read_parquet.py +++ /dev/null @@ -1,183 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import itertools -from typing import TYPE_CHECKING - -import numpy as np -import pylibcudf as plc -import pytest - -from cudf_streaming.parquet import Filter, read_parquet -from cudf_streaming.table_chunk import TableChunk -from rapidsmpf.streaming.core.actor import run_actor_network -from rapidsmpf.streaming.core.leaf_actor import pull_from_channel - -if TYPE_CHECKING: - from typing import Literal - - from rapidsmpf.communicator.communicator import Communicator - from rapidsmpf.streaming.core.actor import CppActor - from rapidsmpf.streaming.core.channel import Channel - from rapidsmpf.streaming.core.context import Context - from rmm.pylibrmm.stream import Stream - - -@pytest.fixture(scope="module") -def source( - tmp_path_factory: pytest.TempPathFactory, -) -> plc.io.SourceInfo: - path = tmp_path_factory.mktemp("read_parquet") - - nrows = 10 - start = 0 - sources = [] - for i in range(10): - table = plc.Table( - [ - plc.Column.from_array( - np.arange(start, start + nrows, dtype="int32") - ) - ] - ) - # gaps in the column numbering we produce - start += nrows + nrows // 2 - filename = path / f"{i:3d}.pq" - sink = plc.io.SinkInfo([filename]) - options = plc.io.parquet.ParquetWriterOptions.builder( - sink, table - ).build() - plc.io.parquet.write_parquet(options) - sources.append(filename) - return plc.io.SourceInfo(sources) - - -def make_filter(stream: Stream) -> plc.expressions.Expression: - return plc.expressions.Operation( - plc.expressions.ASTOperator.LESS, - plc.expressions.ColumnReference(0), - plc.expressions.Literal( - plc.Scalar.from_py( - 15, dtype=plc.DataType(plc.TypeId.INT32), stream=stream - ) - ), - ) - - -def make_producer( - context: Context, - comm: Communicator, - ch: Channel[TableChunk], - options: plc.io.parquet.ParquetReaderOptions, - *, - use_filter: bool, -) -> CppActor: - if use_filter: - fstream = context.br().stream_pool.get_stream() - return read_parquet( - context, - comm, - ch, - 4, - options, - 3, - Filter(fstream, make_filter(fstream)), - ) - else: - return read_parquet(context, comm, ch, 4, options, 3) - - -def get_expected( - ctx: Context, - source: plc.io.SourceInfo, - skip_rows: int | Literal["none"], - num_rows: int | Literal["all"], - *, - use_filter: bool, -) -> plc.Table: - options = plc.io.parquet.ParquetReaderOptions.builder(source).build() - - if skip_rows != "none": - options.set_skip_rows(skip_rows) - if num_rows != "all": - options.set_num_rows(num_rows) - if use_filter: - fstream = ctx.br().stream_pool.get_stream() - filter = make_filter(fstream) - fstream.synchronize() - options.set_filter(filter) - - expected = plc.io.parquet.read_parquet(options).tbl - - if use_filter: - fstream.synchronize() - return expected - - -@pytest.mark.parametrize( - "skip_rows", ["none", 7, 19, 113], ids=lambda s: f"skip_rows_{s}" -) -@pytest.mark.parametrize( - "num_rows", ["all", 0, 3, 31, 83], ids=lambda s: f"nrows_{s}" -) -@pytest.mark.parametrize("use_filter", [False, True]) -def test_read_parquet( - context: Context, - comm: Communicator, - source: plc.io.SourceInfo, - skip_rows: int | Literal["none"], - num_rows: int | Literal["all"], - use_filter: bool, -) -> None: - if comm.nranks != 1: - pytest.skip("Only support single-rank runs") - - ch: Channel[TableChunk] = context.create_channel() - - options = plc.io.parquet.ParquetReaderOptions.builder(source).build() - - if skip_rows != "none": - options.set_skip_rows(skip_rows) - if num_rows != "all": - options.set_num_rows(num_rows) - - producer = make_producer(context, comm, ch, options, use_filter=use_filter) - - consumer, deferred_messages = pull_from_channel(context, ch) - - run_actor_network(context, actors=[producer, consumer]) - - messages = deferred_messages.release() - assert all( - m1.sequence_number < m2.sequence_number - for m1, m2 in itertools.pairwise(messages) - ) - chunks = [TableChunk.from_message(m, br=context.br()) for m in messages] - for chunk in chunks: - chunk.stream.synchronize() - - got = plc.concatenate.concatenate([chunk.table_view() for chunk in chunks]) - for chunk in chunks: - chunk.stream.synchronize() - - expected = get_expected( - context, source, skip_rows, num_rows, use_filter=use_filter - ) - - assert got.num_rows() == expected.num_rows() - assert got.num_columns() == expected.num_columns() - assert got.num_columns() == 1 - - all_equal = plc.reduce.reduce( - plc.binaryop.binary_operation( - got.columns()[0], - expected.columns()[0], - plc.binaryop.BinaryOperator.EQUAL, - plc.DataType(plc.TypeId.BOOL8), - ), - plc.aggregation.all(), - plc.DataType(plc.TypeId.BOOL8), - ) - assert all_equal.to_py()