diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index 8fe25132bbb3..73a6d853b6e8 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -6,7 +6,7 @@ import concurrent.futures import contextlib -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING import kvikio @@ -48,10 +48,46 @@ class CachedParquetInfo: path: str size: int | None file_metadata: plc.io.parquet_metadata.FileMetaData + # Pre-created during footer prefetch; shared across all splits and scans of this file. + # HybridScanReader is not cached: it holds mutable per-read state so each worker + # creates its own from the shared metadata. + _hybrid_scan_metadata: list[plc.io.experimental.HybridScanMetadata] = field( + default_factory=list, compare=False, repr=False + ) + + def hybrid_scan_reader( # pragma: no cover; only called from thread pool workers where coverage.py does not trace + self, + options: plc.io.parquet.ParquetReaderOptions, + ) -> plc.io.experimental.HybridScanReader: + """Return a fresh HybridScanReader backed by shared pre-parsed file metadata.""" + if not self._hybrid_scan_metadata: + self._hybrid_scan_metadata.append( + plc.io.experimental.HybridScanMetadata.from_parquet_metadata( + self.file_metadata, options + ) + ) + return plc.io.experimental.HybridScanReader.from_metadata( + self._hybrid_scan_metadata[0] + ) + + +def _default_reader_options( + info: CachedParquetInfo, +) -> plc.io.parquet.ParquetReaderOptions: + """Return baseline ``ParquetReaderOptions`` for a cached parquet file.""" + return ( + plc.io.parquet.ParquetReaderOptions.builder( + plc.io.SourceInfo([plc.io.types.FilepathSource(info.path, info.size)]) + ) + .decimal_width(plc.TypeId.DECIMAL128) + .build() + ) @nvtx_annotate_cudf_polars(message="fetch_parquet_footers_for_paths") -def _prefetch_parquet_footers_for_paths(paths: list[str]) -> list[CachedParquetInfo]: +def _prefetch_parquet_footers_for_paths( + paths: list[str], *, parse_hybrid_metadata: bool = False +) -> list[CachedParquetInfo]: """ Prefetch parquet footers for a list of paths. @@ -62,6 +98,8 @@ def _prefetch_parquet_footers_for_paths(paths: list[str]) -> list[CachedParquetI ---------- paths The paths to prefetch. + parse_hybrid_metadata + Whether to eagerly parse ``HybridScanMetadata`` for each path. Returns ------- @@ -94,10 +132,14 @@ def _prefetch_parquet_footers_for_paths(paths: list[str]) -> list[CachedParquetI ) ) - return [ + infos = [ CachedParquetInfo(path, size, file_metadata) for path, size, file_metadata in zip(paths, sizes, metadata, strict=True) ] + if parse_hybrid_metadata: + for info in infos: + info.hybrid_scan_reader(_default_reader_options(info)) + return infos @nvtx_annotate_cudf_polars(message="prefetch_parquet_file_metadata_for_ir") @@ -107,6 +149,7 @@ def prefetch_parquet_file_metadata_for_ir( stats: StatsCollector | None = None, *, remote_only: bool = False, + parse_hybrid_metadata: bool = False, ) -> dict[str, CachedParquetInfo]: """ Prefetch parquet metadata for all parquet scans in an IR graph. @@ -125,6 +168,9 @@ def prefetch_parquet_file_metadata_for_ir( remote_only If ``True``, only prefetch metadata for remote URIs (e.g. ``s3://``), skipping local paths. + parse_hybrid_metadata + Whether to eagerly parse ``HybridScanMetadata`` for newly-prefetched + paths. Only useful when ``ParquetOptions.use_hybrid_scan`` is enabled. Returns ------- @@ -171,7 +217,11 @@ def prefetch_parquet_file_metadata_for_ir( with cm: futures = [ - py_executor.submit(_prefetch_parquet_footers_for_paths, [path]) + py_executor.submit( + _prefetch_parquet_footers_for_paths, + [path], + parse_hybrid_metadata=parse_hybrid_metadata, + ) for path in missing_paths ] diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index 3a58ef3ac1c3..5f1c31709b87 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -792,6 +792,7 @@ def evaluate_on_rank( ir_context.py_executor, stats=stats, remote_only=isinstance(prefetch_file_metadata, Unspecified), + parse_hybrid_metadata=config_options.parquet_options.use_hybrid_scan, ) attach_cached_parquet_metadata(ir, cached_parquet_info_map) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 4e061e6a3cb2..4eff80e39ad6 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -17,6 +17,7 @@ import pylibcudf as plc +from cudf_polars.containers import Column, DataFrame from cudf_polars.dsl.ir import ( IR, DataFrameScan, @@ -24,7 +25,9 @@ PythonScan, Scan, Sink, + _prepare_parquet_predicate, ) +from cudf_polars.dsl.to_ast import to_parquet_filter from cudf_polars.dsl.tracing import nvtx_annotate_cudf_polars from cudf_polars.streaming.base import ( IOPartitionFlavor, @@ -40,7 +43,10 @@ if TYPE_CHECKING: from collections.abc import Hashable, MutableMapping, Sequence - from cudf_polars.containers import DataFrame, DataType + import pylibcudf.expressions as plc_expr + from rmm.pylibrmm.stream import Stream + + from cudf_polars.containers import DataType from cudf_polars.dsl.expr import NamedExpr from cudf_polars.dsl.ir import CachedParquetInfo, IRExecutionContext from cudf_polars.streaming.base import ( @@ -81,6 +87,13 @@ def scan_partition_plan( """Extract the partitioning plan of a Scan operation.""" 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. + # 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 + ) if source := stats.scan_stats.get(ir): column_sizes = [ sz @@ -97,12 +110,18 @@ def scan_partition_plan( <= abs(file_size / k_hi - blocksize) else k_hi ) - if factor >= 2: + if factor >= 2 or hybrid_single_file: return IOPartitionPlan( factor, IOPartitionFlavor.SPLIT_FILES, estimated_chunk_bytes=file_size // factor, ) + elif hybrid_single_file: + return IOPartitionPlan( + 1, + IOPartitionFlavor.SPLIT_FILES, + estimated_chunk_bytes=file_size, + ) else: k_lo = min(blocksize // int(file_size), len(ir.paths)) k_hi = k_lo + 1 @@ -119,6 +138,9 @@ def scan_partition_plan( estimated_chunk_bytes=file_size * factor, ) + if hybrid_single_file: + return IOPartitionPlan(1, IOPartitionFlavor.SPLIT_FILES) + # TODO: Use file sizes for csv and json return IOPartitionPlan(1, IOPartitionFlavor.SINGLE_FILE) @@ -181,6 +203,134 @@ def expand_scan_for_rank( ) +def _fetch_byte_ranges( + source_info: plc.io.SourceInfo, + byte_ranges: list[plc.io.text.ByteRangeInfo], + stream: Stream, +) -> list[plc.gpumemoryview]: + return plc.io.parquet_io_utils.fetch_byte_ranges_to_device( + source_info, byte_ranges, stream=stream + ) + + +def _read_with_hybrid_scan( + schema: Schema, + paths: list[str], + with_columns: list[str] | None, + plc_filter: plc_expr.Expression, + row_group_indices: list[int], + stream: Stream, + cached_info: CachedParquetInfo, + *, + split_index: int = 0, + total_splits: int = 1, + stats_pruning: bool = True, +) -> DataFrame: + """Two-pass parquet read via HybridScanReader for a row-group-aligned split.""" + from cudf_polars.dsl.utils.io import _default_reader_options + + assert plc_filter is not None + assert len(paths) == 1, ( + "hybrid scan only supported for SplitScan; one physical file" + ) + with nvtx_annotate_cudf_polars( + message="HybridScan", payload=(split_index + 1, total_splits) + ): + source_info = plc.io.SourceInfo( + [plc.io.types.FilepathSource(cached_info.path, cached_info.size)] + ) + options = _default_reader_options(cached_info) + if with_columns is not None: + options.set_column_names(with_columns) + options.set_filter(plc_filter) + + reader = cached_info.hybrid_scan_reader(options) + + if stats_pruning: + row_group_indices = reader.filter_row_groups_with_stats( + row_group_indices, options, stream=stream + ) + + if row_group_indices: + bloom_ranges, _ = reader.secondary_filters_byte_ranges( + row_group_indices, options + ) + if bloom_ranges: + bloom_chunks = _fetch_byte_ranges(source_info, bloom_ranges, stream) + row_group_indices = reader.filter_row_groups_with_bloom_filters( + bloom_chunks, row_group_indices, options, stream=stream + ) + + if not row_group_indices: + col_names = with_columns if with_columns is not None else list(schema) + return DataFrame( + [ + Column( + plc.column_factories.make_empty_column( + schema[name].plc_type, stream=stream + ), + dtype=schema[name], + name=name, + ) + for name in col_names + ], + 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. + row_mask = reader.build_all_true_row_mask(row_group_indices, stream=stream) + + filter_chunks = _fetch_byte_ranges( + source_info, + reader.filter_column_chunks_byte_ranges(row_group_indices, options), + stream, + ) + filter_tbl_w_meta = reader.materialize_filter_columns( + row_group_indices, + filter_chunks, + row_mask, + plc.io.experimental.UseDataPageMask.YES, + options, + stream=stream, + ) + + payload_chunks = _fetch_byte_ranges( + source_info, + reader.payload_column_chunks_byte_ranges(row_group_indices, options), + stream, + ) + payload_tbl_w_meta = reader.materialize_payload_columns( + row_group_indices, + payload_chunks, + row_mask, + plc.io.experimental.UseDataPageMask.YES, + options, + stream=stream, + ) + + filter_names = filter_tbl_w_meta.column_names(include_children=False) + payload_names = payload_tbl_w_meta.column_names(include_children=False) + filter_df = DataFrame.from_table( + filter_tbl_w_meta.tbl, + filter_names, + [schema[n] for n in filter_names], + stream=stream, + ) + payload_df = DataFrame.from_table( + payload_tbl_w_meta.tbl, + payload_names, + [schema[n] for n in payload_names], + stream=stream, + ) + stream.synchronize() + return DataFrame( + [*filter_df.columns, *payload_df.columns], stream=stream + ).select(list(schema.keys())) + + class SplitScan(IR): """ Input from a split file. @@ -336,6 +486,42 @@ def do_evaluate( 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 ( + 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 + ): + stream = context.get_cuda_stream() + plc_filter = to_parquet_filter( + _prepare_parquet_predicate( + predicate.value, paths, schema, with_columns + ), + stream=stream, + ) + if plc_filter is not 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 @@ -840,6 +1026,9 @@ class ParquetMetadata: Parquet-dataset paths. max_footer_samples Maximum number of file footers to sample metadata from. + parse_hybrid_metadata + Whether to eagerly parse ``HybridScanMetadata`` for sampled paths. + Only useful when ``ParquetOptions.use_hybrid_scan`` is enabled. """ __slots__ = ( @@ -873,7 +1062,13 @@ class ParquetMetadata: """Cached parquet info for the sampled paths. Only set if all files were sampled.""" @nvtx_annotate_cudf_polars(message="ParquetMetadata") - def __init__(self, paths: tuple[str, ...], max_footer_samples: int): + def __init__( + self, + paths: tuple[str, ...], + max_footer_samples: int, + *, + parse_hybrid_metadata: bool = False, + ): from cudf_polars.dsl.utils.io import _prefetch_parquet_footers_for_paths self.paths = paths @@ -900,7 +1095,7 @@ def __init__(self, paths: tuple[str, ...], max_footer_samples: int): sampled_file_count = len(self.sample_paths) sample_parquet_info = _prefetch_parquet_footers_for_paths( - list(self.sample_paths) + list(self.sample_paths), parse_hybrid_metadata=parse_hybrid_metadata ) sample_footers = [info.file_metadata for info in sample_parquet_info] @@ -1031,9 +1226,13 @@ def from_paths( schema: tuple[tuple[str, DataType], ...], max_footer_samples: int, max_row_group_samples: int, + *, + parse_hybrid_metadata: bool = False, ) -> ParquetSourceInfo: """Build a ParquetSourceInfo from a list of paths.""" - metadata = ParquetMetadata(paths, max_footer_samples) + metadata = ParquetMetadata( + paths, max_footer_samples, parse_hybrid_metadata=parse_hybrid_metadata + ) row_count = metadata.row_count file_count = len(paths) @@ -1159,10 +1358,17 @@ def _build_parquet_source( schema: tuple[tuple[str, DataType], ...], max_footer_samples: int, max_row_group_samples: int, + *, + parse_hybrid_metadata: bool = False, ) -> ParquetSourceInfo: """Return cached, fully-computed Parquet datasource information.""" return ParquetSourceInfo.from_paths( - paths, needed_cols, schema, max_footer_samples, max_row_group_samples + paths, + needed_cols, + schema, + max_footer_samples, + max_row_group_samples, + parse_hybrid_metadata=parse_hybrid_metadata, ) @@ -1182,7 +1388,15 @@ def _build_source_info( needed_cols = frozenset(ir.schema) if needed_cols is None else needed_cols schema = tuple(ir.schema.items()) if schema is None else schema paths = tuple(ir.paths) - return _build_parquet_source(paths, needed_cols, schema, max_footer, max_rg) + use_hybrid_scan = config_options.parquet_options.use_hybrid_scan + return _build_parquet_source( + paths, + needed_cols, + schema, + max_footer, + max_rg, + parse_hybrid_metadata=use_hybrid_scan, + ) else: # pragma: no cover raise ValueError(f"Unsupported Scan type: {ir.typ}") diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 315aa0c8f8cd..df349578f316 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -280,6 +280,10 @@ class ParquetOptions: When enabled, filter predicates are JIT-compiled to CUDA kernels for improved performance on large datasets with complex filters. Default is False. + use_hybrid_scan + Whether to use the two-pass ``HybridScanReader`` for ``SplitScan`` + tasks when a predicate can be pushed down to a parquet filter. + Default is False. """ _env_prefix = "CUDF_POLARS__PARQUET_OPTIONS" @@ -321,6 +325,24 @@ class ParquetOptions: default=UNSPECIFIED, ) ) + use_hybrid_scan: bool = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__USE_HYBRID_SCAN", + _bool_converter, + default=False, + ) + ) + # Internal benchmarking flag. When False, skips stats and bloom-filter pruning + # before the first pass of a hybrid scan so you can measure two-pass read + # overhead in isolation. No reason to set this to False in production. + _hybrid_scan_stats_pruning: bool = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__HYBRID_SCAN_STATS_PRUNING", + _bool_converter, + default=True, + ), + init=False, + ) use_jit_filter: bool = dataclasses.field( default_factory=_make_default_factory( f"{_env_prefix}__USE_JIT_FILTER", @@ -344,6 +366,10 @@ def __post_init__(self) -> None: # noqa: D105 raise TypeError("max_row_group_samples must be an int") if not isinstance(self.prefetch_file_metadata, (bool, Unspecified)): raise TypeError("prefetch_file_metadata must be a bool when specified") + if not isinstance(self.use_hybrid_scan, bool): + raise TypeError("use_hybrid_scan must be a bool") + if not isinstance(self._hybrid_scan_stats_pruning, bool): + raise TypeError("_hybrid_scan_stats_pruning must be a bool") if not isinstance(self.use_jit_filter, bool): raise TypeError("use_jit_filter must be a bool") diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index a44af628ea1c..61b223d8e82f 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -370,6 +370,43 @@ def test_streaming_scan_raises() -> None: StreamingScan.do_evaluate([fused], scan, context=ctx) +@pytest.mark.parametrize( + "predicate,use_columns", + [ + # uses hybrid scan reader + (pl.col("x") < 1_000, None), + (pl.col("x") < 1_000, ["x", "z"]), + (pl.col("x") < 1_000, ["z"]), + # falls back to default parquet reader + (pl.col("y").str.contains("cat"), None), + (None, None), + ], +) +def test_split_scan_hybrid( + tmp_path: Path, + df: pl.DataFrame, + predicate: pl.Expr | None, + use_columns: list[str] | None, + streaming_engine_factory: Callable[..., StreamingEngine], +) -> None: + streaming_engine = streaming_engine_factory( + StreamingOptions( + target_partition_size=1_000, + parquet_options={ + "use_hybrid_scan": True, + "prefetch_file_metadata": True, + }, + ), + ) + make_partitioned_source(df, tmp_path, "parquet", n_files=1, row_group_size=100) + q = pl.scan_parquet(tmp_path) + if predicate is not None: + q = q.filter(predicate) + if use_columns is not None: + q = q.select(use_columns) + assert_gpu_result_equal(q, engine=streaming_engine) + + def test_scan_path_mismatch_raises() -> None: # This isn't reachable by polars' public API, so we test it directly. scan = _make_parquet_scan( diff --git a/python/cudf_polars/tests/streaming/test_stats.py b/python/cudf_polars/tests/streaming/test_stats.py index 22c7a1795ee1..fa574c621eeb 100644 --- a/python/cudf_polars/tests/streaming/test_stats.py +++ b/python/cudf_polars/tests/streaming/test_stats.py @@ -161,7 +161,13 @@ class FakeParquetMetadata: } num_row_groups_per_file = (1, 1) - def __init__(self, paths: tuple[str, ...], max_footer_samples: int) -> None: + def __init__( + self, + paths: tuple[str, ...], + max_footer_samples: int, + *, + parse_hybrid_metadata: bool = False, + ) -> None: self.paths = paths self.max_footer_samples = max_footer_samples self.sampled_file_count = 1 diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index fd184cf2f2c3..9fc20ecf9e70 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -373,6 +373,8 @@ 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_HYBRID_SCAN", "0") + m.setenv("CUDF_POLARS__PARQUET_OPTIONS__HYBRID_SCAN_STATS_PRUNING", "0") m.setenv("CUDF_POLARS__PARQUET_OPTIONS__PREFETCH_FILE_METADATA", "1") m.setenv("CUDF_POLARS__PARQUET_OPTIONS__USE_JIT_FILTER", "1") @@ -385,6 +387,8 @@ 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_hybrid_scan is False + assert config.parquet_options._hybrid_scan_stats_pruning is False assert config.parquet_options.prefetch_file_metadata is True assert config.parquet_options.use_jit_filter is True