diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index 254f46ed65e2..010e65a3c463 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -15,7 +15,11 @@ from cudf_polars.dsl.tracing import nvtx_annotate_cudf_polars from cudf_polars.dsl.traversal import traversal -from cudf_polars.streaming.io import Scan, StreamingScan +from cudf_polars.streaming.io import ( + ParquetSourceInfo, + Scan, + StreamingScan, +) if TYPE_CHECKING: from cudf_polars.dsl.ir import IR @@ -181,14 +185,12 @@ def prefetch_parquet_file_metadata_for_ir( ------- A dictionary mapping each individual path to its cached parquet metadata. """ - from cudf_polars.streaming.io import ParquetSourceInfo, StreamingScan - all_paths: set[str] = set() for node in traversal([root]): if isinstance(node, StreamingScan) and node.base_scan.typ == "parquet": - for scan in node.scans: - for path in scan.paths: + for task in node.tasks: + for path in task.paths: all_paths.add(path) elif isinstance(node, Scan) and node.typ == "parquet": # pragma: no cover raise RuntimeError("Unexpected parquet 'Scan' node in lowered IR graph.") @@ -241,7 +243,7 @@ def attach_cached_parquet_metadata( cached_parquet_info_map: dict[str, CachedParquetInfo], ) -> None: """ - Attach prefetched metadata to scan nodes. + Attach prefetched metadata to parquet scan tasks. This is an optimization only and does not affect IR identity. @@ -254,10 +256,15 @@ def attach_cached_parquet_metadata( """ for node in traversal([root]): if isinstance(node, StreamingScan) and node.base_scan.typ == "parquet": - for scan in node.scans: - if not all(path in cached_parquet_info_map for path in scan.paths): - continue - cached = [cached_parquet_info_map[path] for path in scan.paths] - Scan._validate_cached_parquet_info(scan.paths, cached) - scan.cached_parquet_info = cached - scan._non_child_args = (*scan._non_child_args[:-1], cached) + base_scan = node.base_scan + task_paths = {path for task in node.tasks for path in task.paths} + cached_paths = [ + path + for path in base_scan.paths + if path in task_paths and path in cached_parquet_info_map + ] + cached = [cached_parquet_info_map[path] for path in cached_paths] + if not cached: + continue + Scan._validate_cached_parquet_info(cached_paths, cached) + base_scan.cached_parquet_info = cached 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 9145c688da85..2c3f8c7bf577 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py @@ -39,6 +39,7 @@ send_metadata, ) from cudf_polars.streaming.io import ( + ScanTask, StreamingScan, StreamingSink, _prepare_sink_directory, @@ -60,7 +61,6 @@ IOPartitionPlan, PartitionInfo, ) - from cudf_polars.streaming.io import FusedScan, SplitScan from cudf_polars.utils.config import MaxConcurrentIOTasks @@ -516,7 +516,7 @@ def _( async def read_chunk( context: Context, - scan: IR, + task: IR, seq_num: int, ch_out: Channel[TableChunk], ir_context: IRExecutionContext, @@ -530,8 +530,8 @@ async def read_chunk( ---------- context The rapidsmpf context. - scan - The Scan or DataFrameScan node. + task + The scan task to evaluate. seq_num The sequence number. ch_out @@ -546,7 +546,7 @@ async def read_chunk( """ reservation_bytes = ( estimated_chunk_bytes - if isinstance(scan, DataFrameScan) + if isinstance(task, DataFrameScan) else 2 * estimated_chunk_bytes ) start = time.monotonic_ns() @@ -558,8 +558,8 @@ async def read_chunk( admitted = time.monotonic_ns() with opaque_memory_usage(reservation): df = await ir_context.to_thread( - scan.do_evaluate, - *scan._non_child_args, + task.do_evaluate, + *task._non_child_args, context=ir_context, ) chunk = TableChunk.from_pylibcudf_table( @@ -569,14 +569,17 @@ async def read_chunk( br=context.br(), ) stop = time.monotonic_ns() + ir_type = ( + task.trace_ir_type() if isinstance(task, ScanTask) else type(task).__name__ + ) 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__, + ir_id=task.get_stable_id(), + ir_type=ir_type, sequence_number=seq_num, estimated_output_bytes=estimated_chunk_bytes, reservation_bytes=reservation_bytes, @@ -613,7 +616,7 @@ async def scan_node( 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 + tasks: Sequence[ScanTask] = ir.tasks async with shutdown_on_error( context, ch_out, trace_ir=ir, ir_context=ir_context @@ -622,21 +625,21 @@ async def scan_node( await send_metadata( ch_out, context, - ChannelMetadata(local_count=len(scans)), + ChannelMetadata(local_count=len(tasks)), ) # If there is nothing to scan, drain the channel and return - if len(scans) == 0: + if len(tasks) == 0: await ch_out.drain(context) return - # If there is only one scan or one producer, we can + # If there is only one task or one producer, we can # skip the lineariser and read the chunks directly - if len(scans) == 1 or num_producers == 1: - for seq_num, scan in enumerate(scans): + if len(tasks) == 1 or num_producers == 1: + for seq_num, task in enumerate(tasks): await read_chunk( context, - scan, + task, seq_num, ch_out, ir_context, @@ -647,23 +650,22 @@ async def scan_node( return # Use Lineariser to ensure ordered delivery - num_producers = min(num_producers, len(scans)) + num_producers = min(num_producers, len(tasks)) lineariser = Lineariser(context, ch_out, num_producers) # Assign tasks to producers using round-robin - producer_tasks: list[list[tuple[int, SplitScan | FusedScan]]] = [ + producer_tasks: list[list[tuple[int, ScanTask]]] = [ [] for _ in range(num_producers) ] - for task_idx, scan in enumerate(scans): + for task_idx, task in enumerate(tasks): producer_id = task_idx % num_producers - # mypy resolves __iter__ on union-of-sequences to the common base (IR) - producer_tasks[producer_id].append((task_idx, scan)) # type: ignore[arg-type] + producer_tasks[producer_id].append((task_idx, task)) async def _producer(producer_id: int, ch_out: Channel) -> None: - for task_idx, scan in producer_tasks[producer_id]: + for task_idx, task in producer_tasks[producer_id]: await read_chunk( context, - scan, + task, task_idx, ch_out, ir_context, diff --git a/python/cudf_polars/cudf_polars/streaming/explain.py b/python/cudf_polars/cudf_polars/streaming/explain.py index 011e3be232a7..8d69a9292f90 100644 --- a/python/cudf_polars/cudf_polars/streaming/explain.py +++ b/python/cudf_polars/cudf_polars/streaming/explain.py @@ -563,7 +563,7 @@ def _(ir: Scan) -> dict[str, Serializable]: def _(ir: StreamingScan) -> dict[str, Serializable]: return { "typ": ir.base_scan.typ, - "scan_count": len(ir.scans), + "task_count": len(ir.tasks), "prefix": os.path.commonprefix(ir.base_scan.paths), "predicate": ( _serialize_expr(ir.base_scan.predicate) if ir.base_scan.predicate else None diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 638691f30045..93fb69802ce2 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, Self, overload +from typing import TYPE_CHECKING, Any, ClassVar, Literal, NamedTuple, overload import polars as pl @@ -88,8 +88,8 @@ def scan_partition_plan( if ir.typ == "parquet": blocksize: int = config_options.executor.target_partition_size single_file = len(ir.paths) == 1 - # A single file always uses SplitScan when hybrid scan is enabled, so the - # hybrid reader can be used on it even when it would otherwise not split. + # A single file always uses a split parquet task when hybrid scan is enabled, + # so the hybrid reader can be used even when the file would otherwise not split. # The split factor is still size-based, so a large file is split into many. hybrid_single_file = ( single_file and config_options.parquet_options.use_hybrid_scan @@ -183,38 +183,53 @@ def expand_scan_for_rank( StreamingScan Rank-local streaming scan. """ + local_offset, local_count = _rank_slice(partition_count, rank, nranks) if plan.flavor == IOPartitionFlavor.SPLIT_FILES: - return StreamingScan.for_split_files( - ir, - plan, - partition_count, - rank=rank, - nranks=nranks, - parquet_options=parquet_options, - ) + path_offset = local_offset // plan.factor + path_end = math.ceil((local_offset + local_count) / plan.factor) + local_paths = ir.paths[path_offset:path_end] + sindex = local_offset % plan.factor + tasks: list[ScanTask] = [] + splits_created = 0 + for path in local_paths: + while sindex < plan.factor and splits_created < local_count: + tasks.append( + ParquetScanTask( + ir, + [path], + sindex, + plan.factor, + parquet_options, + ) + ) + sindex += 1 + splits_created += 1 + sindex = 0 else: - return StreamingScan.for_fused_files( - ir, - plan, - partition_count, - rank=rank, - nranks=nranks, - parquet_options=parquet_options, - ) + paths_start = local_offset * plan.factor + paths_end = paths_start + plan.factor * local_count + tasks = [] + for offset in range(paths_start, paths_end, plan.factor): + paths = ir.paths[offset : offset + plan.factor] + if not paths: + continue + if ir.typ == "parquet": + tasks.append(ParquetScanTask(ir, paths, 0, 1, parquet_options)) + else: + tasks.append(ScanTask(ir, paths, 0, 1)) + return StreamingScan(tasks, ir) def hybrid_scan_eligible( parquet_options: ParquetOptions, *, - cached_parquet_info: list[CachedParquetInfo] | None, row_index: tuple[str, int] | None, include_file_paths: str | None, predicate: NamedExpr | None, ) -> bool: - """Whether a parquet split is eligible for the HybridScanReader path.""" + """Whether scan options allow hybrid scan if metadata is available.""" return ( parquet_options.use_hybrid_scan - and cached_parquet_info is not None and row_index is None and include_file_paths is None and predicate is not None @@ -234,10 +249,8 @@ def _read_with_hybrid_scan( total_splits: int = 1, stats_pruning: bool = True, ) -> DataFrame: - """Two-pass parquet read via HybridScanReader for a row-group-aligned split.""" - assert len(paths) == 1, ( - "hybrid scan only supported for SplitScan; one physical file" - ) + """Two-pass parquet read via HybridScanReader for a row-group-aligned task.""" + assert len(paths) == 1, "hybrid scan only supports one physical file" with nvtx_annotate_cudf_polars( message="HybridScan", payload=(split_index + 1, total_splits) ): @@ -287,10 +300,10 @@ def _read_with_hybrid_scan( stream=stream, ) - # TODO: Consider implementing page-index stats pruning. For SplitScans, we can - # reuse the same page index for all splits of the same file, so the overhead of - # reading the page index can be amortized. For FusedScans, we would need to read - # the page index for all files, which may be too expensive. + # TODO: Consider implementing page-index stats pruning. For split tasks, we + # can reuse the same page index for all splits of the same file, so the + # overhead of reading the page index can be amortized. For fused tasks, we + # would need to read the page index for all files, which may be too expensive. row_mask = reader.build_all_true_row_mask(row_group_indices, stream=stream) filter_chunks = plc.io.parquet_io_utils.fetch_byte_ranges_to_device( @@ -345,84 +358,74 @@ def _read_with_hybrid_scan( return DataFrame(columns, stream=stream).select(list(schema.keys())) -class SplitScan(IR): +class ParquetScanTaskBounds(NamedTuple): """ - Input from a split file. + Read bounds for a parquet scan task. - This class wraps a single-file ``Scan`` object. At - IO/evaluation time, this class will only perform - a partial read of the underlying file. The range - (skip_rows and n_rows) is calculated at IO time. + ``row_groups=None`` means the task is not row-group aligned. """ + row_groups: list[list[int]] | None + skip_rows: int + n_rows: int + + +class ScanTask(IR): + """Generic streaming scan task.""" + __slots__ = ( "base_scan", - "cached_parquet_info", - "parquet_options", "paths", "schema", "split_index", "total_splits", ) - _non_child = ( - "schema", + _non_child: ClassVar[tuple[str, ...]] = ( "base_scan", "paths", "split_index", "total_splits", - "parquet_options", ) - _n_non_child_args = 13 + _n_non_child_args = 4 base_scan: Scan - """Scan operation this node is based on.""" + """Scan operation this task is based on.""" paths: list[str] - """File path for this split task.""" + """File paths assigned to this task.""" split_index: int - """Index of the current split.""" + """Index of the current split, or 0 for non-split tasks.""" total_splits: int - """Total number of splits.""" - parquet_options: ParquetOptions - """Parquet-specific options.""" - cached_parquet_info: list[CachedParquetInfo] | None + """Total number of splits for a split file, or 1 for non-split tasks.""" def __init__( self, - schema: Schema, base_scan: Scan, paths: list[str], split_index: int, total_splits: int, - parquet_options: ParquetOptions, - cached_parquet_info: list[CachedParquetInfo] | None = None, ): - self.schema = schema + if total_splits < 1: + raise ValueError(f"Expected at least one split, got: {total_splits}") + if not 0 <= split_index < total_splits: + raise ValueError( + f"Expected split_index in [0, {total_splits}), got: {split_index}" + ) + self.schema = base_scan.schema self.base_scan = base_scan self.paths = paths self.split_index = split_index self.total_splits = total_splits self._non_child_args = ( + base_scan, + paths, split_index, total_splits, - base_scan.schema, - base_scan.typ, - base_scan.reader_options, - paths, - base_scan.with_columns, - base_scan.skip_rows, - base_scan.n_rows, - base_scan.row_index, - base_scan.include_file_paths, - base_scan.predicate, - parquet_options, - cached_parquet_info, ) - self.parquet_options = parquet_options - self.cached_parquet_info = cached_parquet_info self.children = () - if base_scan.typ not in ("parquet",): # pragma: no cover - raise NotImplementedError( - f"Unhandled Scan type for file splitting: {base_scan.typ}" - ) + + @property + def is_split(self) -> bool: + """Whether this task is one of multiple splits of a single file.""" + return self.total_splits > 1 def get_hashable(self) -> Hashable: """Hashable representation of the node.""" @@ -433,204 +436,73 @@ def get_hashable(self) -> Hashable: tuple(self.paths), self.split_index, self.total_splits, - self.parquet_options, ) + def trace_ir_type(self) -> str: + """Return the task type to use for IO-task tracing.""" + return "SplitScan" if self.is_split else type(self).__name__ + @classmethod def do_evaluate( cls, + base_scan: Scan, + paths: list[str], split_index: int, total_splits: int, - schema: Schema, - typ: str, - reader_options: dict[str, Any], - paths: list[str], - with_columns: list[str] | None, - skip_rows: int, - n_rows: int, - row_index: tuple[str, int] | None, - include_file_paths: str | None, - predicate: NamedExpr | None, - parquet_options: ParquetOptions, - cached_parquet_info: list[CachedParquetInfo] | None, *, context: IRExecutionContext, ) -> DataFrame: """Evaluate and return a dataframe.""" - if typ not in ("parquet",): # pragma: no cover - raise NotImplementedError(f"Unhandled Scan type for file splitting: {typ}") - - if len(paths) > 1: # pragma: no cover - raise ValueError(f"Expected a single path, got: {paths}") - - # Parquet logic: - # - We are one of "total_splits" SplitScan nodes - # assigned to the same file. - # - We know our index within this file ("split_index") - # - We can also use parquet metadata to query the - # total number of rows in each row-group of the file. - # - We can use all this information to calculate the - # "skip_rows" and "n_rows" options to use locally. - - if cached_parquet_info is not None: - parquet_metadatas = [info.file_metadata for info in cached_parquet_info] - - row_group_num_rows = [ - num_rows - for metadata in parquet_metadatas - for num_rows in metadata.row_group_num_rows - ] - - else: - row_group_num_rows = [ - rg["num_rows"] - for rg in plc.io.parquet_metadata.read_parquet_metadata( - plc.io.SourceInfo(paths) - ).rowgroup_metadata() - ] - - total_row_groups = len(row_group_num_rows) - if total_splits <= total_row_groups: - # We have enough row-groups in the file to align - # all "total_splits" of our reads with row-group - # boundaries. Calculate which row-groups to include - # in the current read, and use metadata to translate - # the row-group indices to "skip_rows" and "n_rows". - rg_stride = total_row_groups // total_splits - skip_rgs = rg_stride * split_index - skip_rows = sum(row_group_num_rows[:skip_rgs]) - n_rows = sum(row_group_num_rows[skip_rgs : skip_rgs + rg_stride]) - # Hybrid scan reads through the prefetched, shared file metadata, so - # it is only used when footer prefetching is enabled. - # TODO: Investigate re-enabling for some of the excluded paths - # (row_index / include_file_paths). Needs performance investigation. - if hybrid_scan_eligible( - parquet_options, - cached_parquet_info=cached_parquet_info, - row_index=row_index, - include_file_paths=include_file_paths, - predicate=predicate, - ): - assert predicate is not None - assert cached_parquet_info is not None - stream = context.get_cuda_stream() - plc_filter, residual = to_parquet_filter( - _prepare_parquet_predicate( - predicate.value, paths, schema, with_columns - ), - stream=stream, - ) - if plc_filter is not None and residual is None: - end_rg = ( - total_row_groups - if split_index == total_splits - 1 - else skip_rgs + rg_stride - ) - return _read_with_hybrid_scan( - schema, - paths, - with_columns, - plc_filter, - list(range(skip_rgs, end_rg)), - stream, - cached_parquet_info[0], - split_index=split_index, - total_splits=total_splits, - stats_pruning=parquet_options._hybrid_scan_stats_pruning, - ) - else: - # There are not enough row-groups to align - # all "total_splits" of our reads with row-group - # boundaries. Use metadata to directly calculate - # "skip_rows" and "n_rows" for the current read. - total_rows = sum(row_group_num_rows) - n_rows = total_rows // total_splits - skip_rows = n_rows * split_index - - # Last split should always read to end of file - if split_index == (total_splits - 1): - n_rows = -1 - - # Perform the partial read - with nvtx_annotate_cudf_polars( - message=f"SplitScan: {paths[0]} [{split_index + 1}/{total_splits}]" - ): + if total_splits > 1: # pragma: no cover + raise NotImplementedError( + f"File splitting is not implemented for {base_scan.typ} scans." + ) + with nvtx_annotate_cudf_polars(message=f"ScanTask: {', '.join(paths)}"): return Scan.do_evaluate( - schema, - typ, - reader_options, + base_scan.schema, + base_scan.typ, + base_scan.reader_options, paths, - with_columns, - skip_rows, - n_rows, - row_index, - include_file_paths, - predicate, - parquet_options, - cached_parquet_info, + base_scan.with_columns, + base_scan.skip_rows, + base_scan.n_rows, + base_scan.row_index, + base_scan.include_file_paths, + base_scan.predicate, + base_scan.parquet_options, + None, context=context, ) -class FusedScan(IR): - """ - Input from one or more complete files read as a single task. - - Covers both FUSED_FILES (N > 1 small files grouped together) and - SINGLE_FILE (N = 1). - """ +class ParquetScanTask(ScanTask): + """Parquet-specific streaming scan task.""" - __slots__ = ( - "base_scan", - "cached_parquet_info", - "parquet_options", - "paths", - "schema", - ) - _non_child = ( - "schema", - "base_scan", - "paths", + __slots__ = ("parquet_options",) + _non_child: ClassVar[tuple[str, ...]] = ( + *ScanTask._non_child, "parquet_options", ) - _n_non_child_args = 11 - base_scan: Scan - """Scan operation this node is based on.""" - paths: list[str] - """File paths assigned to this task.""" - parquet_options: ParquetOptions - """Parquet-specific options.""" - cached_parquet_info: list[CachedParquetInfo] | None - """Cached parquet metadata.""" + _n_non_child_args = 5 def __init__( self, - schema: Schema, base_scan: Scan, paths: list[str], + split_index: int, + total_splits: int, parquet_options: ParquetOptions, - cached_parquet_info: list[CachedParquetInfo] | None = None, ): - self.schema = schema - self.base_scan = base_scan - self.paths = paths + if base_scan.typ != "parquet": + raise ValueError(f"Expected a parquet scan, got: {base_scan.typ}") + if total_splits > 1 and len(paths) > 1: + raise ValueError(f"Expected a single path for a split task, got: {paths}") + super().__init__(base_scan, paths, split_index, total_splits) self.parquet_options = parquet_options - self.cached_parquet_info = cached_parquet_info self._non_child_args = ( - base_scan.schema, - base_scan.typ, - base_scan.reader_options, - paths, - base_scan.with_columns, - base_scan.skip_rows, - base_scan.n_rows, - base_scan.row_index, - base_scan.include_file_paths, - base_scan.predicate, + *self._non_child_args, parquet_options, - cached_parquet_info, ) - self.children = () def get_hashable(self) -> Hashable: """Hashable representation of the node.""" @@ -639,40 +511,178 @@ def get_hashable(self) -> Hashable: tuple(self.schema.items()), self.base_scan.get_hashable(), tuple(self.paths), + self.split_index, + self.total_splits, self.parquet_options, ) + def get_task_bounds(self) -> ParquetScanTaskBounds | None: + """Return parquet read bounds for this task.""" + return self._task_bounds_from_cached(self._cached_parquet_info()) + + def _cached_parquet_info(self) -> list[CachedParquetInfo] | None: + """Return cached parquet metadata matching this task's paths.""" + cached_parquet_info = self.base_scan.cached_parquet_info + if cached_parquet_info is None or cached_parquet_info == []: + return None + if self.paths == [info.path for info in cached_parquet_info]: + return cached_parquet_info + + cached_by_path = {info.path: info for info in cached_parquet_info} + if not all(path in cached_by_path for path in self.paths): + return None + return [cached_by_path[path] for path in self.paths] + + def _fetch_parquet_info( + self, *, parse_hybrid_metadata: bool = False + ) -> list[CachedParquetInfo]: + """Fetch parquet metadata for this task's paths.""" + from cudf_polars.dsl.utils.io import _prefetch_parquet_footers_for_paths + + return _prefetch_parquet_footers_for_paths( + self.paths, + parse_hybrid_metadata=parse_hybrid_metadata, + ) + + def _split_task_bounds( + self, + cached_parquet_info: list[CachedParquetInfo], + ) -> ParquetScanTaskBounds: + """Return parquet read bounds for a split task.""" + row_group_num_rows = cached_parquet_info[0].file_metadata.row_group_num_rows + total_row_groups = len(row_group_num_rows) + if self.total_splits <= total_row_groups: + row_group_stride = total_row_groups // self.total_splits + row_group_start = row_group_stride * self.split_index + row_group_stop = ( + total_row_groups + if self.split_index == self.total_splits - 1 + else row_group_start + row_group_stride + ) + skip_rows = sum(row_group_num_rows[:row_group_start]) + n_rows = sum(row_group_num_rows[row_group_start:row_group_stop]) + row_groups = [list(range(row_group_start, row_group_stop))] + else: + row_groups = None + total_rows = sum(row_group_num_rows) + n_rows = total_rows // self.total_splits + skip_rows = n_rows * self.split_index + + if self.split_index == self.total_splits - 1: + n_rows = -1 + return ParquetScanTaskBounds(row_groups, skip_rows, n_rows) + + def _task_bounds_from_cached( + self, + cached_parquet_info: list[CachedParquetInfo] | None, + ) -> ParquetScanTaskBounds | None: + if self.is_split: + return ( + None + if cached_parquet_info is None + else self._split_task_bounds(cached_parquet_info) + ) + + base_scan = self.base_scan + row_groups: list[list[int]] | None = None + if ( + cached_parquet_info is not None + and base_scan.skip_rows == 0 + and base_scan.n_rows == -1 + and base_scan.row_index is None + ): + row_groups = [ + list(range(len(info.file_metadata.row_group_num_rows))) + for info in cached_parquet_info + ] + return ParquetScanTaskBounds(row_groups, base_scan.skip_rows, base_scan.n_rows) + @classmethod - def do_evaluate( + def do_evaluate( # type: ignore[override] cls, - schema: Schema, - typ: str, - reader_options: dict[str, Any], + base_scan: Scan, paths: list[str], - with_columns: list[str] | None, - skip_rows: int, - n_rows: int, - row_index: tuple[str, int] | None, - include_file_paths: str | None, - predicate: NamedExpr | None, + split_index: int, + total_splits: int, parquet_options: ParquetOptions, - cached_parquet_info: list[CachedParquetInfo] | None, *, context: IRExecutionContext, ) -> DataFrame: - """Evaluate and return a dataframe.""" - with nvtx_annotate_cudf_polars(message=f"FusedScan: {', '.join(paths)}"): + """Evaluate a parquet scan task.""" + task = cls(base_scan, paths, split_index, total_splits, parquet_options) + base_scan = task.base_scan + paths = task.paths + cached_parquet_info = task._cached_parquet_info() + should_try_hybrid_scan = ( + len(paths) == 1 + and base_scan.skip_rows == 0 + and base_scan.n_rows == -1 + and hybrid_scan_eligible( + parquet_options, + row_index=base_scan.row_index, + include_file_paths=base_scan.include_file_paths, + predicate=base_scan.predicate, + ) + ) + if cached_parquet_info is None and (task.is_split or should_try_hybrid_scan): + cached_parquet_info = task._fetch_parquet_info( + parse_hybrid_metadata=should_try_hybrid_scan + ) + bounds = task._task_bounds_from_cached(cached_parquet_info) + + assert bounds is not None + # Hybrid scan reads through cached parquet metadata, so it is only used + # when the metadata is available to this task. + # TODO: Investigate re-enabling for some of the excluded paths + # (row_index / include_file_paths). Needs performance investigation. + if ( + should_try_hybrid_scan + and bounds.row_groups is not None + and len(bounds.row_groups) == 1 + and cached_parquet_info is not None + ): + assert base_scan.predicate is not None + assert cached_parquet_info is not None + stream = context.get_cuda_stream() + plc_filter, residual = to_parquet_filter( + _prepare_parquet_predicate( + base_scan.predicate.value, + paths, + base_scan.schema, + base_scan.with_columns, + ), + stream=stream, + ) + if plc_filter is not None and residual is None: + return _read_with_hybrid_scan( + base_scan.schema, + paths, + base_scan.with_columns, + plc_filter, + bounds.row_groups[0], + stream, + cached_parquet_info[0], + split_index=split_index, + total_splits=total_splits, + stats_pruning=parquet_options._hybrid_scan_stats_pruning, + ) + + if task.is_split: + nvtx_message = f"SplitScan: {paths[0]} [{split_index + 1}/{total_splits}]" + else: + nvtx_message = f"ParquetScanTask: {', '.join(paths)}" + with nvtx_annotate_cudf_polars(message=nvtx_message): return Scan.do_evaluate( - schema, - typ, - reader_options, + base_scan.schema, + base_scan.typ, + base_scan.reader_options, paths, - with_columns, - skip_rows, - n_rows, - row_index, - include_file_paths, - predicate, + base_scan.with_columns, + bounds.skip_rows, + bounds.n_rows, + base_scan.row_index, + base_scan.include_file_paths, + base_scan.predicate, parquet_options, cached_parquet_info, context=context, @@ -748,113 +758,44 @@ class StreamingScan(IR): __slots__ = ( "base_scan", - "scan_type", - "scans", "schema", + "tasks", ) _non_child = ( - "scans", + "tasks", "base_scan", - "scan_type", ) - _n_non_child_args = 3 - scans: Sequence[SplitScan] | Sequence[FusedScan] + _n_non_child_args = 2 base_scan: Scan + tasks: Sequence[ScanTask] def __init__( self, - scans: Sequence[SplitScan] | Sequence[FusedScan], + tasks: Sequence[ScanTask], base_scan: Scan, - scan_type: Literal["split", "fused"], ): - self.scans = scans self.base_scan = base_scan self.schema = base_scan.schema - self.scan_type = scan_type - self._non_child_args = (scans, base_scan, scan_type) + self.tasks = tasks + self._non_child_args = (tasks, 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, - None, - ) - ) - sindex += 1 - splits_created += 1 - sindex = 0 - return cls(scans, base_scan, "split") - - @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, - None, - ) - for offset in range(paths_start, paths_end, plan.factor) - if base_scan.paths[offset : offset + plan.factor] - ] - return cls(scans, base_scan, "fused") - 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)) + # We don't need to include base_scan / schema, since it's in all the scan tasks. + return (type(self), *tuple(task.get_hashable() for task in self.tasks)) @classmethod def do_evaluate( cls, - scans: Sequence[SplitScan] | Sequence[FusedScan], + tasks: Sequence[ScanTask], 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." + "StreamingScan.do_evaluate should not be called directly. Call Scan.do_evaluate on each scan task instead." ) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 47c8a9f6dfdb..bb506bb2a248 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -354,7 +354,7 @@ class ParquetOptions: improved performance on large datasets with complex filters. Default is False. use_hybrid_scan - Whether to use the two-pass ``HybridScanReader`` for ``SplitScan`` + Whether to use the two-pass ``HybridScanReader`` for split parquet tasks when a predicate can be pushed down to a parquet filter. Default is False. """ diff --git a/python/cudf_polars/tests/streaming/test_explain.py b/python/cudf_polars/tests/streaming/test_explain.py index 5ea6be578ae4..707811cd5a71 100644 --- a/python/cudf_polars/tests/streaming/test_explain.py +++ b/python/cudf_polars/tests/streaming/test_explain.py @@ -558,7 +558,7 @@ def test_scan_properties(tmp_path: Path, predicate: pl.Expr | None): "prefix": f"{root}/", "typ": "parquet", "predicate": None, - "scan_count": 1, + "task_count": 1, } if predicate is not None: q = q.filter(predicate) diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 72389d58e9fd..8510d8a3bba9 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -20,6 +20,7 @@ from cudf_polars.dsl.utils.io import ( CachedParquetInfo, _prefetch_parquet_footers_for_paths, + attach_cached_parquet_metadata, prefetch_parquet_file_metadata_for_ir, ) from cudf_polars.engine.options import StreamingOptions @@ -31,8 +32,8 @@ StatsCollector, ) from cudf_polars.streaming.io import ( - FusedScan, - SplitScan, + ParquetScanTask, + ScanTask, StreamingScan, expand_scan_for_rank, scan_partition_plan, @@ -178,8 +179,8 @@ def recording_prefetch( ) scan = _make_parquet_scan(paths) - fused = FusedScan(scan.schema, scan, paths, scan.parquet_options, None) - streaming_scan = StreamingScan([fused], scan, "fused") + task = ParquetScanTask(scan, paths, 0, 1, scan.parquet_options) + streaming_scan = StreamingScan([task], scan) result = prefetch_parquet_file_metadata_for_ir( streaming_scan, py_executor=None, stats=stats @@ -194,14 +195,16 @@ def test_prefetch_parquet_file_metadata_remote_only(tmp_path, df) -> None: local_path = str(next(tmp_path.glob("*.parquet"))) scan = _make_parquet_scan([local_path]) - fused = FusedScan(scan.schema, scan, scan.paths, scan.parquet_options, []) - streaming_scan = StreamingScan([fused], scan, "fused") + task = ParquetScanTask(scan, scan.paths, 0, 1, scan.parquet_options) + streaming_scan = StreamingScan([task], scan) # Local paths are skipped entirely when remote_only=True. result = prefetch_parquet_file_metadata_for_ir( streaming_scan, py_executor=None, stats=None, remote_only=True ) assert result == {} + attach_cached_parquet_metadata(streaming_scan, result) + assert scan.cached_parquet_info is None # The same local path is prefetched when remote_only=False (the default). result = prefetch_parquet_file_metadata_for_ir( @@ -394,7 +397,12 @@ def test_scan_union(engine: pl.GPUEngine, tmp_path: Path) -> None: def _make_parquet_scan( - paths: list[str], parquet_options: ParquetOptions | None = None + paths: list[str], + parquet_options: ParquetOptions | None = None, + *, + skip_rows: int = 0, + n_rows: int = -1, + row_index: tuple[str, int] | None = None, ) -> Scan: parquet_options = parquet_options or ParquetOptions() return Scan( @@ -404,12 +412,43 @@ def _make_parquet_scan( None, paths, None, + skip_rows, + n_rows, + row_index, + None, + None, + parquet_options, + None, + ) + + +def _make_csv_scan(paths: list[str]) -> Scan: + return Scan( + {"x": DataType(pl.Int64())}, + "csv", + { + "has_header": True, + "schema": None, + "skip_rows": 0, + "skip_rows_after_header": 0, + "parse_options": { + "separator": ord(","), + "quote_char": ord('"'), + "eol_char": ord("\n"), + "null_values": None, + "comment_prefix": None, + "decimal_comma": False, + }, + }, + None, + paths, + None, 0, -1, None, None, None, - parquet_options, + ParquetOptions(), None, ) @@ -458,9 +497,11 @@ def test_expand_scan_for_rank_fused_and_single_read( parquet_options=ParquetOptions(), ) for scan, expected_paths in zip( - streaming_scan.scans, expected_path_groups, strict=True + streaming_scan.tasks, expected_path_groups, strict=True ): - assert isinstance(scan, FusedScan) + assert isinstance(scan, ParquetScanTask) + assert scan.split_index == 0 + assert scan.total_splits == 1 assert scan.paths == expected_paths @@ -486,23 +527,143 @@ def test_expand_scan_for_rank_split_files( nranks=2, parquet_options=ParquetOptions(), ) - assert len(streaming_scan.scans) == len(expected_splits) + assert len(streaming_scan.tasks) == len(expected_splits) for scan, (split_index, total_splits) in zip( - streaming_scan.scans, expected_splits, strict=True + streaming_scan.tasks, expected_splits, strict=True ): - assert isinstance(scan, SplitScan) + assert isinstance(scan, ParquetScanTask) assert scan.split_index == split_index assert scan.total_splits == total_splits assert scan.paths == ["file.parquet"] +def test_attach_cached_parquet_metadata_resolves_row_groups( + tmp_path: Path, +) -> None: + source = tmp_path / "data.parquet" + pl.DataFrame({"x": range(4)}).write_parquet(source, row_group_size=2) + + base = _make_parquet_scan([str(source)]) + streaming_scan = expand_scan_for_rank( + base, + IOPartitionPlan(2, IOPartitionFlavor.SPLIT_FILES), + partition_count=2, + rank=0, + nranks=1, + parquet_options=base.parquet_options, + ) + + cached = prefetch_parquet_file_metadata_for_ir(streaming_scan, None) + attach_cached_parquet_metadata(streaming_scan, cached) + + row_groups = [] + for scan in streaming_scan.tasks: + assert isinstance(scan, ParquetScanTask) + bounds = scan.get_task_bounds() + assert bounds is not None + row_groups.append(bounds.row_groups) + assert row_groups == [[[0]], [[1]]] + + +def test_attach_cached_parquet_metadata_uses_rank_local_tasks( + tmp_path: Path, +) -> None: + paths = [str(tmp_path / f"part-{i}.parquet") for i in range(4)] + for path in paths: + pl.DataFrame({"x": range(4)}).write_parquet(path, row_group_size=2) + + base = _make_parquet_scan(paths) + streaming_scan = expand_scan_for_rank( + base, + IOPartitionPlan(1, IOPartitionFlavor.SINGLE_FILE), + partition_count=4, + rank=0, + nranks=2, + parquet_options=base.parquet_options, + ) + + cached = prefetch_parquet_file_metadata_for_ir(streaming_scan, None) + attach_cached_parquet_metadata(streaming_scan, cached) + + assert base.cached_parquet_info is not None + assert [info.path for info in base.cached_parquet_info] == paths[:2] + for scan in streaming_scan.tasks: + assert isinstance(scan, ParquetScanTask) + bounds = scan.get_task_bounds() + assert bounds is not None + assert bounds.row_groups == [[0, 1]] + + +def test_attach_cached_parquet_metadata_leaves_sub_row_group_split_unaligned( + tmp_path: Path, +) -> None: + source = tmp_path / "data.parquet" + pl.DataFrame({"x": range(4)}).write_parquet(source, row_group_size=2) + + base = _make_parquet_scan([str(source)]) + streaming_scan = expand_scan_for_rank( + base, + IOPartitionPlan(4, IOPartitionFlavor.SPLIT_FILES), + partition_count=4, + rank=0, + nranks=1, + parquet_options=base.parquet_options, + ) + + cached = prefetch_parquet_file_metadata_for_ir(streaming_scan, None) + attach_cached_parquet_metadata(streaming_scan, cached) + + for scan in streaming_scan.tasks: + assert isinstance(scan, ParquetScanTask) + assert scan.is_split + bounds = scan.get_task_bounds() + assert bounds is not None + assert bounds.row_groups is None + + +@pytest.mark.parametrize( + "skip_rows,n_rows,row_index", + [(1, -1, None), (0, 2, None), (0, -1, ("index", 0))], +) +def test_attach_cached_parquet_metadata_leaves_sliced_fused_scan_unaligned( + tmp_path: Path, + skip_rows: int, + n_rows: int, + row_index: tuple[str, int] | None, +) -> None: + source = tmp_path / "data.parquet" + pl.DataFrame({"x": range(4)}).write_parquet(source, row_group_size=2) + + base = _make_parquet_scan( + [str(source)], skip_rows=skip_rows, n_rows=n_rows, row_index=row_index + ) + streaming_scan = expand_scan_for_rank( + base, + IOPartitionPlan(1, IOPartitionFlavor.SINGLE_READ), + partition_count=1, + rank=0, + nranks=1, + parquet_options=base.parquet_options, + ) + + cached = prefetch_parquet_file_metadata_for_ir(streaming_scan, None) + attach_cached_parquet_metadata(streaming_scan, cached) + + for scan in streaming_scan.tasks: + assert isinstance(scan, ParquetScanTask) + assert not scan.is_split + bounds = scan.get_task_bounds() + assert bounds is not None + assert bounds.row_groups is None + + def test_streaming_scan_raises() -> None: # This isn't reachable by normal cudf-polars usage. - scan = _make_parquet_scan(["file.parquet"]) - fused = FusedScan(scan.schema, scan, scan.paths, scan.parquet_options, []) + scan = _make_csv_scan(["file.csv"]) + task = ScanTask(scan, scan.paths, 0, 1) ctx = IRExecutionContext() with pytest.raises(NotImplementedError, match=r"StreamingScan.do_evaluate"): - StreamingScan.do_evaluate([fused], scan, context=ctx) + StreamingScan.do_evaluate([task], scan, context=ctx) @pytest.mark.parametrize( @@ -571,45 +732,24 @@ def test_scan_path_mismatch_raises() -> None: ) -def test_streaming_scan_missing_prefetch_metadata_raises() -> None: - # This isn't reachable by polars' public API, so we test it directly. +def test_parquet_split_task_fetches_missing_metadata(tmp_path: Path) -> None: + source = tmp_path / "data.parquet" + pl.DataFrame({"x": range(4)}).write_parquet(source, row_group_size=2) + scan = _make_parquet_scan( - ["file.parquet"], parquet_options=ParquetOptions(prefetch_file_metadata=True) + [str(source)], parquet_options=ParquetOptions(prefetch_file_metadata=True) ) - fused = FusedScan(scan.schema, scan, scan.paths, scan.parquet_options, []) - - ctx = IRExecutionContext() - with pytest.raises(NotImplementedError, match=r"StreamingScan.do_evaluate"): - StreamingScan.do_evaluate([fused], scan, context=ctx) - -def test_split_scan_do_evaluate_missing_prefetch_metadata() -> None: - paths = ["/some/missing/file.parquet"] - parquet_options = ParquetOptions(prefetch_file_metadata=True) - context = IRExecutionContext() - schema = {"x": DataType(pl.Int64())} + result = ParquetScanTask.do_evaluate( + scan, + scan.paths, + 0, + 2, + scan.parquet_options, + context=IRExecutionContext(), + ) - with pytest.raises( - AssertionError, - match=(r"Paths do not match cached parquet info."), - ): - SplitScan.do_evaluate( - 0, - 4, - schema, - "parquet", - {}, - paths, - None, - 0, - -1, - None, - None, - None, - parquet_options, - [], - context=context, - ) + assert result.to_polars().to_dict(as_series=False) == {"x": [0, 1]} def test_prefetch_file_metadata_join( @@ -667,29 +807,51 @@ def test_prefetch_file_metadata_with_cached_scan_parent_nodes( assert_gpu_result_equal(q, engine=engine) -def test_fused_scan_identity_equality() -> None: - base = _make_parquet_scan(["a.parquet", "b.parquet"]) - paths = ["a.parquet"] - info = _make_cached_parquet_info(paths) +def test_scan_task_identity_equality() -> None: + base = _make_csv_scan(["a.csv", "b.csv"]) + paths = ["a.csv"] - a = FusedScan(base.schema, base, paths, base.parquet_options, info) - b = FusedScan(base.schema, base, paths, base.parquet_options, info.copy()) - c = FusedScan(base.schema, base, ["b.parquet"], base.parquet_options, info) + a = ScanTask(base, paths, 0, 1) + b = ScanTask(base, paths, 0, 1) + c = ScanTask(base, ["b.csv"], 0, 1) assert a == b assert hash(a) == hash(b) assert a != c -def test_split_scan_identity_equality() -> None: +def test_scan_task_validates_split_bounds() -> None: + base = _make_csv_scan(["a.csv"]) + + with pytest.raises(ValueError, match=r"Expected at least one split"): + ScanTask(base, ["a.csv"], 0, 0) + + with pytest.raises(ValueError, match=r"Expected split_index in"): + ScanTask(base, ["a.csv"], 1, 1) + + +def test_parquet_scan_task_validates_inputs() -> None: + csv_scan = _make_csv_scan(["a.csv"]) + with pytest.raises(ValueError, match=r"Expected a parquet scan"): + ParquetScanTask(csv_scan, csv_scan.paths, 0, 1, csv_scan.parquet_options) + + parquet_scan = _make_parquet_scan(["a.parquet", "b.parquet"]) + with pytest.raises(ValueError, match=r"Expected a single path for a split task"): + ParquetScanTask( + parquet_scan, + parquet_scan.paths, + 0, + 2, + parquet_scan.parquet_options, + ) + + +def test_parquet_split_task_identity_equality() -> None: base = _make_parquet_scan(["a.parquet"]) - info = _make_cached_parquet_info(base.paths) - a = SplitScan(base.schema, base, base.paths, 0, 4, base.parquet_options, info) - b = SplitScan( - base.schema, base, base.paths, 0, 4, base.parquet_options, info.copy() - ) - c = SplitScan(base.schema, base, base.paths, 1, 4, base.parquet_options, info) + a = ParquetScanTask(base, base.paths, 0, 4, base.parquet_options) + b = ParquetScanTask(base, base.paths, 0, 4, base.parquet_options) + c = ParquetScanTask(base, base.paths, 1, 4, base.parquet_options) assert a == b assert hash(a) == hash(b) @@ -698,37 +860,31 @@ def test_split_scan_identity_equality() -> None: def test_streaming_scan_identity_equality() -> None: base = _make_parquet_scan(["a.parquet"]) - split = SplitScan( - base.schema, + split = ParquetScanTask( base, base.paths, 0, 2, base.parquet_options, - _make_cached_parquet_info(base.paths, size=10), ) - split_same = SplitScan( - base.schema, + split_same = ParquetScanTask( base, base.paths, 0, 2, base.parquet_options, - _make_cached_parquet_info(base.paths, size=11), ) - split_diff = SplitScan( - base.schema, + split_diff = ParquetScanTask( base, base.paths, 1, 2, base.parquet_options, - _make_cached_parquet_info(base.paths, size=10), ) - a = StreamingScan([split], base, "split") - b = StreamingScan([split_same], base, "split") - c = StreamingScan([split_diff], base, "split") + a = StreamingScan([split], base) + b = StreamingScan([split_same], base) + c = StreamingScan([split_diff], base) assert a == b assert hash(a) == hash(b) @@ -758,20 +914,6 @@ def test_cached_parquet_info_excluded_from_identity() -> None: assert scan_without == scan_with assert hash(scan_without) == hash(scan_with) - split_without = SplitScan( - base.schema, base, base.paths, 0, 4, base.parquet_options, None - ) - split_with = SplitScan( - base.schema, base, base.paths, 0, 4, base.parquet_options, info - ) - assert split_without == split_with - assert hash(split_without) == hash(split_with) - - fused_without = FusedScan(base.schema, base, base.paths, base.parquet_options, None) - fused_with = FusedScan(base.schema, base, base.paths, base.parquet_options, info) - assert fused_without == fused_with - assert hash(fused_without) == hash(fused_with) - class FooSource(DataSourceInfo): def __init__(self, size: int):