From 9db0749d0e86a1145f21cc965e84fd5bcf96ac03 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Wed, 10 Jun 2026 16:12:05 +0000 Subject: [PATCH 01/20] Add FusedScan and nxtx annotations for FusedScan and SplitScan --- .../cudf_polars/streaming/actor_graph/io.py | 4 +- .../cudf_polars/cudf_polars/streaming/io.py | 161 ++++++++++++++---- .../cudf_polars/tests/streaming/test_scan.py | 12 +- 3 files changed, 136 insertions(+), 41 deletions(-) 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 822a84d38c75..547f1e6be53c 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py @@ -63,7 +63,7 @@ PartitionInfo, StatsCollector, ) - from cudf_polars.streaming.io import SplitScan + from cudf_polars.streaming.io import FusedScan, SplitScan class Lineariser: @@ -408,7 +408,7 @@ async def scan_node( lineariser = Lineariser(context, ch_out, num_producers) # Assign tasks to producers using round-robin - producer_tasks: list[list[tuple[int, Scan | SplitScan]]] = [ + producer_tasks: list[list[tuple[int, SplitScan | FusedScan]]] = [ [] for _ in range(num_producers) ] for task_idx, scan in enumerate(scans): diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index d56c31ce379b..74006e601620 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -112,9 +112,9 @@ def expand_scan_for_rank( rank: int, nranks: int, parquet_options: ParquetOptions, -) -> list[Scan | SplitScan]: +) -> list[SplitScan | FusedScan]: """ - Expand a Scan node into rank-local Scan and SplitScan operations. + Expand a Scan node into rank-local SplitScan and FusedScan operations. Parameters ---------- @@ -131,10 +131,10 @@ def expand_scan_for_rank( Returns ------- - list[Scan | SplitScan] + list[SplitScan | FusedScan] Rank-local scan operations. """ - scans: list[Scan | SplitScan] = [] + scans: list[SplitScan | FusedScan] = [] if plan.flavor == IOPartitionFlavor.SPLIT_FILES: count = plan.factor * len(ir.paths) local_count = math.ceil(count / nranks) @@ -182,23 +182,8 @@ def expand_scan_for_rank( paths_offset_end = paths_offset_start + plan.factor * local_count for offset in range(paths_offset_start, paths_offset_end, plan.factor): local_paths = ir.paths[offset : offset + plan.factor] - if len(local_paths) > 0: # Only add scan if there are paths - scans.append( - Scan( - ir.schema, - ir.typ, - ir.reader_options, - ir.cloud_options, - local_paths, - ir.with_columns, - ir.skip_rows, - ir.n_rows, - ir.row_index, - ir.include_file_paths, - ir.predicate, - parquet_options, - ) - ) + if len(local_paths) > 0: + scans.append(FusedScan(ir.schema, ir, local_paths, parquet_options)) return scans @@ -207,7 +192,7 @@ class SplitScan(IR): """ Input from a split file. - This class wraps a single-file `Scan` object. At + 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. @@ -337,20 +322,124 @@ def do_evaluate( n_rows = -1 # Perform the partial read - return Scan.do_evaluate( - schema, - typ, - reader_options, + with nvtx_annotate_cudf_polars( + message=f"SplitScan: {Path(paths[0]).name} [{split_index + 1}/{total_splits}]" + ): + return Scan.do_evaluate( + schema, + typ, + reader_options, + paths, + with_columns, + skip_rows, + n_rows, + row_index, + include_file_paths, + predicate, + parquet_options, + 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). The ``paths`` attribute holds the file group + assigned to this task; the remaining read options come from + ``base_scan``. + """ + + __slots__ = ( + "base_scan", + "parquet_options", + "paths", + "schema", + ) + _non_child = ( + "schema", + "base_scan", + "paths", + "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.""" + + def __init__( + self, + schema: Schema, + base_scan: Scan, + paths: list[str], + parquet_options: ParquetOptions, + ): + self.schema = schema + self.base_scan = base_scan + self.paths = paths + self.parquet_options = parquet_options + self._non_child_args = ( + 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, + base_scan.skip_rows, + base_scan.n_rows, + base_scan.row_index, + base_scan.include_file_paths, + base_scan.predicate, parquet_options, - context=context, ) + self.children = () + + def get_hashable(self) -> Hashable: + """Hashable representation of the node.""" + return ( + type(self), + tuple(self.schema.items()), + self.base_scan.get_hashable(), + tuple(self.paths), + self.parquet_options, + ) + + @classmethod + def do_evaluate( + cls, + 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, + *, + context: IRExecutionContext, + ) -> DataFrame: + """Evaluate and return a dataframe.""" + with nvtx_annotate_cudf_polars(message=f"FusedScan: {len(paths)} files"): + return Scan.do_evaluate( + schema, + typ, + reader_options, + paths, + with_columns, + skip_rows, + n_rows, + row_index, + include_file_paths, + predicate, + parquet_options, + context=context, + ) @lower_ir_node.register(Empty) @@ -484,10 +573,10 @@ class StreamingScan(IR): "base_scan", ) _n_non_child_args = 2 - scans: list[Scan | SplitScan] + scans: list[SplitScan | FusedScan] base_scan: Scan - def __init__(self, scans: list[Scan | SplitScan], base_scan: Scan): + def __init__(self, scans: list[SplitScan | FusedScan], base_scan: Scan): self.scans = scans self.base_scan = base_scan self.schema = base_scan.schema @@ -502,7 +591,7 @@ def get_hashable(self) -> Hashable: @classmethod def do_evaluate( cls, - scans: list[Scan | SplitScan], + scans: list[SplitScan | FusedScan], base_scan: Scan, *, context: IRExecutionContext, diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index c9ccb13202dc..7ea69d4296d9 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -14,7 +14,12 @@ from cudf_polars.dsl.ir import IRExecutionContext, Scan from cudf_polars.engine.options import StreamingOptions from cudf_polars.streaming.base import IOPartitionFlavor, IOPartitionPlan -from cudf_polars.streaming.io import SplitScan, StreamingScan, expand_scan_for_rank +from cudf_polars.streaming.io import ( + FusedScan, + SplitScan, + StreamingScan, + expand_scan_for_rank, +) from cudf_polars.streaming.parallel import lower_ir_graph from cudf_polars.streaming.statistics import collect_statistics from cudf_polars.testing.asserts import assert_gpu_result_equal @@ -217,7 +222,7 @@ def test_expand_scan_for_rank_fused_and_single_read( parquet_options=ParquetOptions(), ) for scan, expected_paths in zip(scans, expected_path_groups, strict=True): - assert isinstance(scan, Scan) + assert isinstance(scan, FusedScan) assert scan.paths == expected_paths @@ -251,6 +256,7 @@ def test_expand_scan_for_rank_split_files( 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) ctx = IRExecutionContext() with pytest.raises(NotImplementedError, match=r"StreamingScan.do_evaluate"): - StreamingScan.do_evaluate([scan], scan, context=ctx) + StreamingScan.do_evaluate([fused], scan, context=ctx) From b6bc1c13e0456af617fcb18d821ff1b3f6b33851 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Wed, 10 Jun 2026 17:32:10 +0000 Subject: [PATCH 02/20] Have both streaming scan class take paths --- .../cudf_polars/cudf_polars/streaming/io.py | 34 +++++++------------ .../cudf_polars/tests/streaming/test_scan.py | 6 ++-- 2 files changed, 16 insertions(+), 24 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 74006e601620..dccb77d887d0 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -146,25 +146,12 @@ def expand_scan_for_rank( sindex = local_offset % plan.factor splits_created = 0 for path in local_paths: - base_scan = Scan( - ir.schema, - ir.typ, - ir.reader_options, - ir.cloud_options, - [path], - ir.with_columns, - ir.skip_rows, - ir.n_rows, - ir.row_index, - ir.include_file_paths, - ir.predicate, - parquet_options, - ) while sindex < plan.factor and splits_created < local_count: scans.append( SplitScan( ir.schema, - base_scan, + ir, + [path], sindex, plan.factor, parquet_options, @@ -201,6 +188,7 @@ class SplitScan(IR): __slots__ = ( "base_scan", "parquet_options", + "paths", "schema", "split_index", "total_splits", @@ -208,6 +196,7 @@ class SplitScan(IR): _non_child = ( "schema", "base_scan", + "paths", "split_index", "total_splits", "parquet_options", @@ -215,6 +204,8 @@ class SplitScan(IR): _n_non_child_args = 13 base_scan: Scan """Scan operation this node is based on.""" + paths: list[str] + """File path for this split task.""" split_index: int """Index of the current split.""" total_splits: int @@ -226,12 +217,14 @@ def __init__( self, schema: Schema, base_scan: Scan, + paths: list[str], split_index: int, total_splits: int, parquet_options: ParquetOptions, ): self.schema = schema self.base_scan = base_scan + self.paths = paths self.split_index = split_index self.total_splits = total_splits self._non_child_args = ( @@ -240,14 +233,14 @@ def __init__( base_scan.schema, base_scan.typ, base_scan.reader_options, - base_scan.paths, + 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, - base_scan.parquet_options, + parquet_options, ) self.parquet_options = parquet_options self.children = () @@ -346,9 +339,7 @@ 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). The ``paths`` attribute holds the file group - assigned to this task; the remaining read options come from - ``base_scan``. + SINGLE_FILE (N = 1). """ __slots__ = ( @@ -425,7 +416,8 @@ def do_evaluate( context: IRExecutionContext, ) -> DataFrame: """Evaluate and return a dataframe.""" - with nvtx_annotate_cudf_polars(message=f"FusedScan: {len(paths)} files"): + names = ", ".join(Path(p).name for p in paths) + with nvtx_annotate_cudf_polars(message=f"FusedScan: {names}"): return Scan.do_evaluate( schema, typ, diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 7ea69d4296d9..5ddf0868513c 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -223,7 +223,7 @@ def test_expand_scan_for_rank_fused_and_single_read( ) for scan, expected_paths in zip(scans, expected_path_groups, strict=True): assert isinstance(scan, FusedScan) - assert scan.paths == expected_paths + assert scan.base_scan.paths == expected_paths @pytest.mark.parametrize( @@ -250,13 +250,13 @@ def test_expand_scan_for_rank_split_files( assert isinstance(scan, SplitScan) assert scan.split_index == split_index assert scan.total_splits == total_splits - assert scan.base_scan.paths == ["file.parquet"] + assert scan.paths == ["file.parquet"] def test_streaming_scan_raises() -> None: # This isn't reachable by normal cudf-polars usage. scan = _make_parquet_scan(["file.parquet"]) - fused = FusedScan(scan.schema, scan, scan.paths, scan.parquet_options) + fused = FusedScan(scan.schema, scan, scan.parquet_options) ctx = IRExecutionContext() with pytest.raises(NotImplementedError, match=r"StreamingScan.do_evaluate"): StreamingScan.do_evaluate([fused], scan, context=ctx) From c30e48829a029d69ca7488eb760bd14d1066b4f3 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Wed, 10 Jun 2026 17:34:25 +0000 Subject: [PATCH 03/20] missed --- python/cudf_polars/tests/streaming/test_scan.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 5ddf0868513c..db8492b2c6ee 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -256,7 +256,7 @@ def test_expand_scan_for_rank_split_files( 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.parquet_options) + 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) From d3f4086a32eccb0414b4ed95d51d28a912c92da6 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Wed, 10 Jun 2026 22:58:26 +0000 Subject: [PATCH 04/20] address review, add absolute path prefix to nvtx annotation --- python/cudf_polars/cudf_polars/streaming/io.py | 17 ++++++++++++++--- python/cudf_polars/tests/streaming/test_scan.py | 2 +- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index dccb77d887d0..4162a08791d1 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -249,6 +249,18 @@ def __init__( f"Unhandled Scan type for file splitting: {base_scan.typ}" ) + def get_hashable(self) -> Hashable: + """Hashable representation of the node.""" + return ( + type(self), + tuple(self.schema.items()), + self.base_scan.get_hashable(), + tuple(self.paths), + self.split_index, + self.total_splits, + self.parquet_options, + ) + @classmethod def do_evaluate( cls, @@ -316,7 +328,7 @@ def do_evaluate( # Perform the partial read with nvtx_annotate_cudf_polars( - message=f"SplitScan: {Path(paths[0]).name} [{split_index + 1}/{total_splits}]" + message=f"SplitScan: {paths[0]} [{split_index + 1}/{total_splits}]" ): return Scan.do_evaluate( schema, @@ -416,8 +428,7 @@ def do_evaluate( context: IRExecutionContext, ) -> DataFrame: """Evaluate and return a dataframe.""" - names = ", ".join(Path(p).name for p in paths) - with nvtx_annotate_cudf_polars(message=f"FusedScan: {names}"): + with nvtx_annotate_cudf_polars(message=f"FusedScan: {', '.join(paths)}"): return Scan.do_evaluate( schema, typ, diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index db8492b2c6ee..cf6e14d31bca 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -223,7 +223,7 @@ def test_expand_scan_for_rank_fused_and_single_read( ) for scan, expected_paths in zip(scans, expected_path_groups, strict=True): assert isinstance(scan, FusedScan) - assert scan.base_scan.paths == expected_paths + assert scan.paths == expected_paths @pytest.mark.parametrize( From 9f805eeb14eb7edbdf05d87ae1a02c7126b92751 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 11 Jun 2026 17:48:22 +0000 Subject: [PATCH 05/20] Dispatch SplitScan to HybridScanReader in the streaming engine --- cpp/include/cudf/io/parquet_io_utils.hpp | 22 +++ .../io/parquet/io_utils/parquet_io_utils.cpp | 12 ++ python/cudf_polars/cudf_polars/dsl/ir.py | 5 +- .../cudf_polars/streaming/actor_graph/io.py | 23 ++- .../cudf_polars/cudf_polars/streaming/io.py | 187 ++++++++++++++++-- .../cudf_polars/streaming/select.py | 9 +- .../cudf_polars/cudf_polars/utils/config.py | 13 ++ .../cudf_polars/tests/streaming/test_scan.py | 47 ++++- python/cudf_polars/tests/test_select.py | 1 + python/pylibcudf/pylibcudf/io/CMakeLists.txt | 3 +- python/pylibcudf/pylibcudf/io/__init__.py | 2 + .../pylibcudf/io/parquet_io_utils.pyx | 118 +++++++++++ .../pylibcudf/libcudf/io/parquet_io_utils.pxd | 29 +++ .../pylibcudf/libcudf/utilities/span.pxd | 2 + 14 files changed, 444 insertions(+), 29 deletions(-) create mode 100644 python/pylibcudf/pylibcudf/io/parquet_io_utils.pyx create mode 100644 python/pylibcudf/pylibcudf/libcudf/io/parquet_io_utils.pxd diff --git a/cpp/include/cudf/io/parquet_io_utils.hpp b/cpp/include/cudf/io/parquet_io_utils.hpp index 97e2fe97b5db..d3801795244d 100644 --- a/cpp/include/cudf/io/parquet_io_utils.hpp +++ b/cpp/include/cudf/io/parquet_io_utils.hpp @@ -131,6 +131,28 @@ fetch_byte_ranges_to_device_async( rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr); +/** + * @brief Fetches byte ranges from a single datasource into device buffers, blocking until complete + * + * @ingroup io_utils + * + * Convenience wrapper around @ref fetch_byte_ranges_to_device_async that waits for the + * returned future before returning. + * + * @param datasource Input datasource + * @param byte_ranges Byte ranges to fetch + * @param stream CUDA stream + * @param mr Device memory resource + * + * @return A pair containing the device buffers and the device spans of the fetched data + */ +[[nodiscard]] std::pair, + std::vector>> +fetch_byte_ranges_to_device(cudf::io::datasource& datasource, + cudf::host_span byte_ranges, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + /** @} */ // end of group } // namespace io::parquet } // namespace CUDF_EXPORT cudf diff --git a/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp b/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp index 7170d818cd87..6c72ef935430 100644 --- a/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp +++ b/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp @@ -424,4 +424,16 @@ fetch_byte_ranges_to_device_async( mr); } +std::pair, std::vector>> +fetch_byte_ranges_to_device(cudf::io::datasource& datasource, + cudf::host_span byte_ranges, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + CUDF_FUNC_RANGE(); + auto [buffers, spans, fut] = fetch_byte_ranges_to_device_async(datasource, byte_ranges, stream, mr); + fut.get(); + return {std::move(buffers), std::move(spans)}; +} + } // namespace cudf::io::parquet diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index cb55d9c01308..65d89882355d 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -47,7 +47,10 @@ from cudf_polars.dsl.expressions.base import ExecutionContext from cudf_polars.dsl.nodebase import Node from cudf_polars.dsl.to_ast import _DECIMAL_IDS, to_ast, to_parquet_filter -from cudf_polars.dsl.tracing import log_do_evaluate, nvtx_annotate_cudf_polars +from cudf_polars.dsl.tracing import ( + log_do_evaluate, + nvtx_annotate_cudf_polars, +) from cudf_polars.dsl.utils.reshape import broadcast from cudf_polars.dsl.utils.windows import ( offsets_to_windows, 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 547f1e6be53c..e4b29b38247c 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py @@ -23,6 +23,7 @@ Sink, _prepare_parquet_predicate, ) +from cudf_polars.dsl.tracing import nvtx_annotate_cudf_polars from cudf_polars.dsl.to_ast import to_parquet_filter from cudf_polars.streaming.actor_graph.dispatch import ( generate_ir_sub_network, @@ -322,23 +323,25 @@ async def read_chunk( tracer The actor tracer for collecting runtime statistics. """ - with opaque_memory_usage( - await reserve_memory( + with nvtx_annotate_cudf_polars(message="reserve_memory"): + reservation = await reserve_memory( context, size=estimated_chunk_bytes, net_memory_delta=estimated_chunk_bytes ) - ): + 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(), - ) - await send_chunk(context, ch_out, chunk, seq_num, tracer=tracer) + with nvtx_annotate_cudf_polars(message="TableChunk.from_pylibcudf_table"): + chunk = TableChunk.from_pylibcudf_table( + df.table, + df.stream, + exclusive_view=True, + br=context.br(), + ) + with nvtx_annotate_cudf_polars(message="send_chunk"): + await send_chunk(context, ch_out, chunk, seq_num, tracer=tracer) @define_actor() diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 4162a08791d1..87371f0d74f3 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -17,13 +17,16 @@ import pylibcudf as plc +from cudf_polars.containers import DataFrame from cudf_polars.dsl.ir import ( IR, DataFrameScan, Empty, 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,9 @@ if TYPE_CHECKING: from collections.abc import Hashable, MutableMapping - from cudf_polars.containers import DataFrame + import pylibcudf.expressions as plc_expr + from rmm.pylibrmm.stream import Stream + from cudf_polars.dsl.expr import NamedExpr from cudf_polars.dsl.ir import IRExecutionContext from cudf_polars.streaming.base import ( @@ -175,6 +180,129 @@ def expand_scan_for_rank( return scans +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, + file_metadata: plc.io.parquet_metadata.FileMetaData, +) -> DataFrame: + """Two-pass parquet read via HybridScanReader for a row-group-aligned split.""" + 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=f"HybridScan: {paths[0]}"): + source_info = plc.io.SourceInfo(paths) + + options = ( + plc.io.parquet.ParquetReaderOptions.builder(source_info) + .decimal_width(plc.TypeId.DECIMAL128) + .build() + ) + if with_columns is not None: + options.set_column_names(with_columns) + options.set_filter(plc_filter) + + reader = plc.io.experimental.HybridScanReader.from_parquet_metadata( + file_metadata, options + ) + + row_group_indices = reader.filter_row_groups_with_stats( + row_group_indices, options, stream=stream + ) + + if row_group_indices: + bloom_ranges, _dict_ranges = reader.secondary_filters_byte_ranges( + row_group_indices, options + ) + if bloom_ranges: + bloom_chunks = plc.io.parquet_io_utils.fetch_byte_ranges_to_device( + source_info, bloom_ranges, stream=stream + ) + row_group_indices = reader.filter_row_groups_with_bloom_filters( + bloom_chunks, row_group_indices, options, stream=stream + ) + + if not row_group_indices: + byte_ranges = reader.all_column_chunks_byte_ranges(row_group_indices, options) + chunks = plc.io.parquet_io_utils.fetch_byte_ranges_to_device( + source_info, byte_ranges, stream=stream + ) + tbl_w_meta = reader.materialize_all_columns( + row_group_indices, chunks, options, stream=stream + ) + col_names = tbl_w_meta.column_names(include_children=False) + num_rows = tbl_w_meta.num_rows_per_source[0] if not col_names else None + stream.synchronize() + return DataFrame.from_table( + tbl_w_meta.tbl, + col_names, + [schema[name] for name in col_names], + stream=stream, + num_rows=num_rows, + ) + + n_rows = reader.total_rows_in_row_groups(row_group_indices) + row_mask = plc.Column.from_scalar( + plc.Scalar.from_py(True, dtype=plc.DataType(plc.TypeId.BOOL8), stream=stream), + n_rows, + stream=stream, + ) + + filter_ranges = reader.filter_column_chunks_byte_ranges( + row_group_indices, options + ) + filter_chunks = plc.io.parquet_io_utils.fetch_byte_ranges_to_device( + source_info, filter_ranges, stream=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_ranges = reader.payload_column_chunks_byte_ranges( + row_group_indices, options + ) + # PERFORMANCE!! payload_column_chunks_byte_ranges does not need the row mask, so + # for local NVMe/GDS this fetch could be submitted async before + # materialize_filter_columns to overlap I/O with GPU decode. + payload_chunks = plc.io.parquet_io_utils.fetch_byte_ranges_to_device( + source_info, payload_ranges, stream=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. @@ -296,10 +424,14 @@ def do_evaluate( # - We can use all this information to calculate the # "skip_rows" and "n_rows" options to use locally. - rowgroup_metadata = plc.io.parquet_metadata.read_parquet_metadata( - plc.io.SourceInfo(paths) - ).rowgroup_metadata() - total_row_groups = len(rowgroup_metadata) + 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 @@ -308,17 +440,48 @@ def do_evaluate( # 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(rg["num_rows"] for rg in rowgroup_metadata[:skip_rgs]) - n_rows = sum( - rg["num_rows"] - for rg in rowgroup_metadata[skip_rgs : skip_rgs + rg_stride] - ) + skip_rows = sum(row_group_num_rows[:skip_rgs]) + n_rows = sum(row_group_num_rows[skip_rgs : skip_rgs + rg_stride]) + # TODO: Investigate re-enabling for some of these + # paths. Needs perfromance investigation. + if ( + parquet_options.use_hybrid_scan + 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 + ) + [file_metadata] = plc.io.parquet_metadata.read_parquet_footers( + plc.io.SourceInfo(paths) + ) + return _read_with_hybrid_scan( + schema, + paths, + with_columns, + plc_filter, + list(range(skip_rgs, end_rg)), + stream, + file_metadata, + ) + 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(rg["num_rows"] for rg in rowgroup_metadata) + total_rows = sum(row_group_num_rows) n_rows = total_rows // total_splits skip_rows = n_rows * split_index @@ -516,7 +679,7 @@ def can_use_native_parquet_node( @lower_ir_node.register(Scan) def _( ir: Scan, rec: LowerIRTransformer -) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: +) -> tuple[StreamingScan, MutableMapping[IR, PartitionInfo]]: config_options = rec.state["config_options"] parquet_options = config_options.parquet_options if ( diff --git a/python/cudf_polars/cudf_polars/streaming/select.py b/python/cudf_polars/cudf_polars/streaming/select.py index e1466f0b6610..81bfec030bc9 100644 --- a/python/cudf_polars/cudf_polars/streaming/select.py +++ b/python/cudf_polars/cudf_polars/streaming/select.py @@ -431,8 +431,15 @@ def _( if scan_child and scan_child.predicate is None and scan_child.typ == "parquet": # Special Case: Fast count. + # We can't use prefetched file metadata here, because we're in lowering, + # not execution, so we don't have an IRExecutionContext with the prefetched + # file metadata yet. count = Scan._get_parquet_row_count_from_metadata( - scan_child.paths, scan_child.skip_rows, scan_child.n_rows + scan_child.paths, + scan_child.skip_rows, + scan_child.n_rows, + scan_child.parquet_options, + context=None, ) dtype = ir.exprs[0].value.dtype diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 5786a5351cc4..c41e3fc3c871 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -206,6 +206,10 @@ class ParquetOptions: Whether to use the native rapidsmpf node for parquet reading. This option is only used by the streaming executor. 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" @@ -247,6 +251,13 @@ class ParquetOptions: default=False, ) ) + use_hybrid_scan: bool = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__USE_HYBRID_SCAN", + _bool_converter, + default=False, + ) + ) def __post_init__(self) -> None: # noqa: D105 if not isinstance(self.chunked, bool): @@ -263,6 +274,8 @@ def __post_init__(self) -> None: # noqa: D105 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.use_hybrid_scan, bool): + raise TypeError("use_hybrid_scan must be a bool") def default_target_partition_size(min_device_size: int | None) -> int: diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index cf6e14d31bca..ec51ef28bd4f 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -11,7 +11,10 @@ from cudf_polars import Translator from cudf_polars.containers import DataType -from cudf_polars.dsl.ir import IRExecutionContext, Scan +from cudf_polars.dsl.ir import ( + IRExecutionContext, + Scan, +) from cudf_polars.engine.options import StreamingOptions from cudf_polars.streaming.base import IOPartitionFlavor, IOPartitionPlan from cudf_polars.streaming.io import ( @@ -28,8 +31,11 @@ if TYPE_CHECKING: import concurrent.futures + from collections.abc import Callable from pathlib import Path + from cudf_polars.engine.core import StreamingEngine + @pytest.fixture(scope="module") def df(): @@ -254,9 +260,42 @@ def test_expand_scan_for_rank_split_files( 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) ctx = IRExecutionContext() with pytest.raises(NotImplementedError, match=r"StreamingScan.do_evaluate"): - StreamingScan.do_evaluate([fused], scan, context=ctx) + StreamingScan.do_evaluate([scan], scan, context=ctx) + + +@pytest.mark.parametrize( + "predicate,use_columns", + [ + # pushdown-able predicate + (pl.col("x") < 1_000, None), + # predicate on all columns, with column selection + (pl.col("x") < 1_000, ["x", "z"]), + # non-pushdown predicate falls back to normal scan (no error) + (pl.col("y").str.contains("cat"), None), + # no predicate — hybrid scan disabled, normal read + (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}, + ), + ) + make_partitioned_source(df, tmp_path, "parquet", n_files=1, row_group_size=100) + q = pl.scan_parquet(tmp_path) + if use_columns is not None: + q = q.select(use_columns) + if predicate is not None: + q = q.filter(predicate) + assert_gpu_result_equal(q, engine=streaming_engine) diff --git a/python/cudf_polars/tests/test_select.py b/python/cudf_polars/tests/test_select.py index f37c2d195d92..7369de844708 100644 --- a/python/cudf_polars/tests/test_select.py +++ b/python/cudf_polars/tests/test_select.py @@ -147,3 +147,4 @@ def test_select_fast_count_parquet_skip_rows( q = pl.scan_parquet(file).slice(1, 5).select(pl.len()) assert_gpu_result_equal(q, engine=engine) + diff --git a/python/pylibcudf/pylibcudf/io/CMakeLists.txt b/python/pylibcudf/pylibcudf/io/CMakeLists.txt index 089ea8d0e8d9..ac023c71a425 100644 --- a/python/pylibcudf/pylibcudf/io/CMakeLists.txt +++ b/python/pylibcudf/pylibcudf/io/CMakeLists.txt @@ -6,7 +6,8 @@ # ============================================================================= set(cython_sources avro.pyx csv.pyx datasource.pyx json.pyx orc.pyx parquet.pyx - parquet_metadata.pyx text.pyx timezone.pyx types.pyx + parquet_io_utils.pyx parquet_metadata.pyx text.pyx timezone.pyx + types.pyx ) set(linked_libraries cudf::cudf) diff --git a/python/pylibcudf/pylibcudf/io/__init__.py b/python/pylibcudf/pylibcudf/io/__init__.py index 2162b50e963c..d4bbb5665bcc 100644 --- a/python/pylibcudf/pylibcudf/io/__init__.py +++ b/python/pylibcudf/pylibcudf/io/__init__.py @@ -9,6 +9,7 @@ json, orc, parquet, + parquet_io_utils, parquet_metadata, text, timezone, @@ -29,6 +30,7 @@ "json", "orc", "parquet", + "parquet_io_utils", "parquet_metadata", "text", "timezone", diff --git a/python/pylibcudf/pylibcudf/io/parquet_io_utils.pyx b/python/pylibcudf/pylibcudf/io/parquet_io_utils.pyx new file mode 100644 index 000000000000..7267cbf26505 --- /dev/null +++ b/python/pylibcudf/pylibcudf/io/parquet_io_utils.pyx @@ -0,0 +1,118 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +from libc.stddef cimport size_t +from libc.stdint cimport uintptr_t +from libcpp.memory cimport make_unique, unique_ptr +from libcpp.pair cimport pair +from libcpp.utility cimport move +from libcpp.vector cimport vector +from cython.operator cimport dereference + +from rmm.librmm.device_buffer cimport device_buffer +from rmm.pylibrmm.device_buffer cimport DeviceBuffer +from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource +from rmm.pylibrmm.stream cimport Stream + +from pylibcudf.gpumemoryview cimport gpumemoryview +from pylibcudf.io.text cimport ByteRangeInfo +from pylibcudf.io.types cimport SourceInfo +from pylibcudf.libcudf.io.datasource cimport datasource, make_datasources +from pylibcudf.libcudf.io.parquet_io_utils cimport ( + const_byte_range_info, + const_uint8_t, + fetch_byte_ranges_to_device as cpp_fetch_byte_ranges_to_device, +) +from pylibcudf.libcudf.io.text cimport byte_range_info +from pylibcudf.libcudf.utilities.span cimport device_span, host_span +from pylibcudf.utils cimport _get_memory_resource, _get_stream + +__all__ = ["fetch_byte_ranges_to_device"] + + +def fetch_byte_ranges_to_device( + SourceInfo source_info, + list byte_ranges, + object stream=None, + object mr=None, +) -> list[gpumemoryview]: + """Fetch byte ranges from a Parquet source into device memory. + + Parameters + ---------- + source_info : SourceInfo + Source describing a single Parquet file. + byte_ranges : list[ByteRangeInfo] + Byte ranges to fetch, as returned by + :meth:`~pylibcudf.io.experimental.HybridScanReader.filter_column_chunks_byte_ranges`, + :meth:`~pylibcudf.io.experimental.HybridScanReader.payload_column_chunks_byte_ranges`, + or + :meth:`~pylibcudf.io.experimental.HybridScanReader.all_column_chunks_byte_ranges`. + stream : Stream, optional + CUDA stream. + mr : DeviceMemoryResource, optional + Device memory resource. + + Returns + ------- + list[gpumemoryview] + One view per byte range. Each view holds a reference to the + :class:`~rmm.DeviceBuffer` that owns its memory, keeping the + allocation alive for as long as the view is referenced. + + Raises + ------ + ValueError + If ``source_info`` does not describe exactly one source. + """ + cdef Stream _stream = _get_stream(stream) + cdef DeviceMemoryResource _mr = _get_memory_resource(mr) + cdef vector[unique_ptr[datasource]] sources = make_datasources(source_info.c_obj) + if sources.size() != 1: + raise ValueError( + f"fetch_byte_ranges_to_device requires exactly one source, " + f"got {sources.size()}" + ) + + cdef vector[byte_range_info] ranges_vec + cdef ByteRangeInfo bri + for bri in byte_ranges: + ranges_vec.push_back(bri.c_obj) + + cdef pair[vector[device_buffer], vector[device_span[const_uint8_t]]] fetched + with nogil: + fetched = cpp_fetch_byte_ranges_to_device( + dereference(sources[0]), + host_span[const_byte_range_info](ranges_vec.data(), ranges_vec.size()), + _stream.view(), + _mr.get_mr(), + ) + + # Wrap the device_buffer as a Python DeviceBuffer that owns the allocation. + # All views share a reference to it, keeping the memory alive. + cdef DeviceBuffer owner = DeviceBuffer.c_from_unique_ptr( + make_unique[device_buffer](move(fetched.first[0])), + _stream, + _mr, + ) + + cdef gpumemoryview gmv + cdef uintptr_t ptr + cdef size_t n + result = [] + for i in range(fetched.second.size()): + ptr = fetched.second[i].data() + n = fetched.second[i].size() + gmv = gpumemoryview.__new__(gpumemoryview) + gmv.ptr = ptr + gmv.nbytes = n + gmv.obj = owner + gmv.cai = { + "shape": (n,), + "strides": None, + "typestr": "|u1", + "data": (ptr, False), + "version": 3, + } + result.append(gmv) + return result diff --git a/python/pylibcudf/pylibcudf/libcudf/io/parquet_io_utils.pxd b/python/pylibcudf/pylibcudf/libcudf/io/parquet_io_utils.pxd new file mode 100644 index 000000000000..fe613c5c3da2 --- /dev/null +++ b/python/pylibcudf/pylibcudf/libcudf/io/parquet_io_utils.pxd @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +from libc.stdint cimport uint8_t +from libcpp.pair cimport pair +from libcpp.vector cimport vector + +from rmm.librmm.cuda_stream_view cimport cuda_stream_view +from rmm.librmm.device_buffer cimport device_buffer +from rmm.librmm.memory_resource cimport device_async_resource_ref + +from pylibcudf.exception_handler cimport libcudf_exception_handler +from pylibcudf.libcudf.io.datasource cimport datasource +from pylibcudf.libcudf.io.text cimport byte_range_info +from pylibcudf.libcudf.utilities.span cimport device_span, host_span + +ctypedef const uint8_t const_uint8_t +ctypedef const byte_range_info const_byte_range_info + +cdef extern from "cudf/io/parquet_io_utils.hpp" \ + namespace "cudf::io::parquet" nogil: + + pair[vector[device_buffer], vector[device_span[const_uint8_t]]] \ + fetch_byte_ranges_to_device( + datasource& source, + host_span[const_byte_range_info] byte_ranges, + cuda_stream_view stream, + device_async_resource_ref mr, + ) except +libcudf_exception_handler diff --git a/python/pylibcudf/pylibcudf/libcudf/utilities/span.pxd b/python/pylibcudf/pylibcudf/libcudf/utilities/span.pxd index f2bf388e4d4c..d800786100ea 100644 --- a/python/pylibcudf/pylibcudf/libcudf/utilities/span.pxd +++ b/python/pylibcudf/pylibcudf/libcudf/utilities/span.pxd @@ -1,5 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2021-2025, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 +from libc.stddef cimport size_t from libcpp.vector cimport vector from pylibcudf.exception_handler cimport libcudf_exception_handler from pylibcudf.libcudf.types cimport size_type @@ -15,3 +16,4 @@ cdef extern from "cudf/utilities/span.hpp" namespace "cudf" nogil: device_span() noexcept device_span(T *data, size_type size) noexcept T *data() noexcept + size_t size() noexcept From cf4451049848912381d507b22418e6424ccfd920 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 11 Jun 2026 17:55:27 +0000 Subject: [PATCH 06/20] docs and cython boilerplate --- .../cudf/source/pylibcudf/api_docs/io/index.rst | 1 + .../pylibcudf/pylibcudf/io/parquet_io_utils.pxd | 12 ++++++++++++ .../pylibcudf/pylibcudf/io/parquet_io_utils.pyi | 17 +++++++++++++++++ .../pylibcudf/pylibcudf/io/parquet_io_utils.pyx | 1 + 4 files changed, 31 insertions(+) create mode 100644 python/pylibcudf/pylibcudf/io/parquet_io_utils.pxd create mode 100644 python/pylibcudf/pylibcudf/io/parquet_io_utils.pyi diff --git a/docs/cudf/source/pylibcudf/api_docs/io/index.rst b/docs/cudf/source/pylibcudf/api_docs/io/index.rst index 15a87175325f..dc485c07de14 100644 --- a/docs/cudf/source/pylibcudf/api_docs/io/index.rst +++ b/docs/cudf/source/pylibcudf/api_docs/io/index.rst @@ -20,6 +20,7 @@ I/O Functions json orc parquet + parquet_io_utils parquet_metadata text timezone diff --git a/python/pylibcudf/pylibcudf/io/parquet_io_utils.pxd b/python/pylibcudf/pylibcudf/io/parquet_io_utils.pxd new file mode 100644 index 000000000000..f0bf640c640a --- /dev/null +++ b/python/pylibcudf/pylibcudf/io/parquet_io_utils.pxd @@ -0,0 +1,12 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +from pylibcudf.io.types cimport SourceInfo +from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource + +cpdef list fetch_byte_ranges_to_device( + SourceInfo source_info, + list byte_ranges, + object stream=*, + DeviceMemoryResource mr=*, +) diff --git a/python/pylibcudf/pylibcudf/io/parquet_io_utils.pyi b/python/pylibcudf/pylibcudf/io/parquet_io_utils.pyi new file mode 100644 index 000000000000..c57031655fe2 --- /dev/null +++ b/python/pylibcudf/pylibcudf/io/parquet_io_utils.pyi @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +from pylibcudf.gpumemoryview import gpumemoryview +from pylibcudf.io.text import ByteRangeInfo +from pylibcudf.io.types import SourceInfo +from pylibcudf.utils import CudaStreamLike +from rmm.pylibrmm.memory_resource import DeviceMemoryResource + +__all__ = ["fetch_byte_ranges_to_device"] + +def fetch_byte_ranges_to_device( + source_info: SourceInfo, + byte_ranges: list[ByteRangeInfo], + stream: CudaStreamLike | None = None, + mr: DeviceMemoryResource | None = None, +) -> list[gpumemoryview]: ... diff --git a/python/pylibcudf/pylibcudf/io/parquet_io_utils.pyx b/python/pylibcudf/pylibcudf/io/parquet_io_utils.pyx index 7267cbf26505..4fe084f21fbc 100644 --- a/python/pylibcudf/pylibcudf/io/parquet_io_utils.pyx +++ b/python/pylibcudf/pylibcudf/io/parquet_io_utils.pyx @@ -1,5 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 +"""IO utilities for the Parquet.""" from libc.stddef cimport size_t from libc.stdint cimport uintptr_t From b7ecc3220a7115f834775497141c8ec597e96cc4 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 11 Jun 2026 18:00:11 +0000 Subject: [PATCH 07/20] clean up from mixed branch --- python/cudf_polars/cudf_polars/streaming/select.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/select.py b/python/cudf_polars/cudf_polars/streaming/select.py index 81bfec030bc9..83dd1fcb6266 100644 --- a/python/cudf_polars/cudf_polars/streaming/select.py +++ b/python/cudf_polars/cudf_polars/streaming/select.py @@ -430,16 +430,10 @@ def _( scan_child = child if scan_child and scan_child.predicate is None and scan_child.typ == "parquet": - # Special Case: Fast count. - # We can't use prefetched file metadata here, because we're in lowering, - # not execution, so we don't have an IRExecutionContext with the prefetched - # file metadata yet. count = Scan._get_parquet_row_count_from_metadata( scan_child.paths, scan_child.skip_rows, scan_child.n_rows, - scan_child.parquet_options, - context=None, ) dtype = ir.exprs[0].value.dtype From 4e53ba596208b2a08e64b28e9dfda5b014182a74 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 11 Jun 2026 18:01:27 +0000 Subject: [PATCH 08/20] remove stream sync --- python/cudf_polars/cudf_polars/streaming/io.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 87371f0d74f3..c1bbf883d4f6 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -233,8 +233,7 @@ def _read_with_hybrid_scan( row_group_indices, chunks, options, stream=stream ) col_names = tbl_w_meta.column_names(include_children=False) - num_rows = tbl_w_meta.num_rows_per_source[0] if not col_names else None - stream.synchronize() + num_rows = tbl_w_meta.num_rows_per_source[0] if not col_names else Nones return DataFrame.from_table( tbl_w_meta.tbl, col_names, From 17682d294b85eaf5c11fed67f07db24da7c3a161 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 11 Jun 2026 18:03:56 +0000 Subject: [PATCH 09/20] remove newline --- python/cudf_polars/tests/test_select.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/cudf_polars/tests/test_select.py b/python/cudf_polars/tests/test_select.py index 7369de844708..f37c2d195d92 100644 --- a/python/cudf_polars/tests/test_select.py +++ b/python/cudf_polars/tests/test_select.py @@ -147,4 +147,3 @@ def test_select_fast_count_parquet_skip_rows( q = pl.scan_parquet(file).slice(1, 5).select(pl.len()) assert_gpu_result_equal(q, engine=engine) - From ff366edc7f8210728e63d1fbd17d25046556f25a Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 11 Jun 2026 18:07:40 +0000 Subject: [PATCH 10/20] add changes back from mixed branch --- python/cudf_polars/tests/streaming/test_scan.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index ec51ef28bd4f..ce6ecc8652c8 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -260,6 +260,7 @@ def test_expand_scan_for_rank_split_files( def test_streaming_scan_raises() -> None: + # This isn't reachable by normal cudf-polars usage. scan = _make_parquet_scan(["file.parquet"]) ctx = IRExecutionContext() with pytest.raises(NotImplementedError, match=r"StreamingScan.do_evaluate"): @@ -269,13 +270,11 @@ def test_streaming_scan_raises() -> None: @pytest.mark.parametrize( "predicate,use_columns", [ - # pushdown-able predicate + # uses hybrid scan reader (pl.col("x") < 1_000, None), - # predicate on all columns, with column selection (pl.col("x") < 1_000, ["x", "z"]), - # non-pushdown predicate falls back to normal scan (no error) + # fallsback to default parquet reader (pl.col("y").str.contains("cat"), None), - # no predicate — hybrid scan disabled, normal read (None, None), ], ) From 9e9ead8dbe7f199e782de5e0e8546697ae234ec6 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 11 Jun 2026 19:02:32 +0000 Subject: [PATCH 11/20] mypy fix --- .../cudf_polars/cudf_polars/streaming/io.py | 25 +++++++------------ .../pylibcudf/io/experimental/hybrid_scan.pyi | 15 +++++------ 2 files changed, 17 insertions(+), 23 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 3ba901f346d8..d665e5aab53e 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -41,7 +41,7 @@ from cudf_polars.utils.versions import POLARS_VERSION_LT_137 if TYPE_CHECKING: - from collections.abc import Hashable, MutableMapping, Sequence + from collections.abc import Hashable, MutableMapping import pylibcudf.expressions as plc_expr from rmm.pylibrmm.stream import Stream @@ -222,19 +222,16 @@ def _read_with_hybrid_scan( bloom_chunks = plc.io.parquet_io_utils.fetch_byte_ranges_to_device( source_info, bloom_ranges, stream=stream ) - bloom_chunks_seq: Sequence[object] = bloom_chunks row_group_indices = reader.filter_row_groups_with_bloom_filters( - bloom_chunks_seq, row_group_indices, options, stream=stream + bloom_chunks, row_group_indices, options, stream=stream ) if not row_group_indices: byte_ranges = reader.all_column_chunks_byte_ranges( row_group_indices, options ) - chunks: Sequence[object] = ( - plc.io.parquet_io_utils.fetch_byte_ranges_to_device( - source_info, byte_ranges, stream=stream - ) + chunks = plc.io.parquet_io_utils.fetch_byte_ranges_to_device( + source_info, byte_ranges, stream=stream ) tbl_w_meta = reader.materialize_all_columns( row_group_indices, chunks, options, stream=stream @@ -252,7 +249,7 @@ def _read_with_hybrid_scan( n_rows = reader.total_rows_in_row_groups(row_group_indices) row_mask = plc.Column.from_scalar( plc.Scalar.from_py( - value=True, dtype=plc.DataType(plc.TypeId.BOOL8), stream=stream + py_val=True, dtype=plc.DataType(plc.TypeId.BOOL8), stream=stream ), n_rows, stream=stream, @@ -261,10 +258,8 @@ def _read_with_hybrid_scan( filter_ranges = reader.filter_column_chunks_byte_ranges( row_group_indices, options ) - filter_chunks: Sequence[object] = ( - plc.io.parquet_io_utils.fetch_byte_ranges_to_device( - source_info, filter_ranges, stream=stream - ) + filter_chunks = plc.io.parquet_io_utils.fetch_byte_ranges_to_device( + source_info, filter_ranges, stream=stream ) filter_tbl_w_meta = reader.materialize_filter_columns( row_group_indices, @@ -281,10 +276,8 @@ def _read_with_hybrid_scan( # PERFORMANCE!! payload_column_chunks_byte_ranges does not need the row mask, so # for local NVMe/GDS this fetch could be submitted async before # materialize_filter_columns to overlap I/O with GPU decode. - payload_chunks: Sequence[object] = ( - plc.io.parquet_io_utils.fetch_byte_ranges_to_device( - source_info, payload_ranges, stream=stream - ) + payload_chunks = plc.io.parquet_io_utils.fetch_byte_ranges_to_device( + source_info, payload_ranges, stream=stream ) payload_tbl_w_meta = reader.materialize_payload_columns( row_group_indices, diff --git a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyi b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyi index a4ff0805320a..b6a7ebafb907 100644 --- a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyi +++ b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyi @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 +from collections.abc import Sequence from enum import IntEnum from rmm.pylibrmm.memory_resource import DeviceMemoryResource @@ -49,14 +50,14 @@ class HybridScanReader: ) -> tuple[list[ByteRangeInfo], list[ByteRangeInfo]]: ... def filter_row_groups_with_dictionary_pages( self, - dictionary_page_data: list[Span], + dictionary_page_data: Sequence[Span], row_group_indices: list[int], options: ParquetReaderOptions, stream: CudaStreamLike | None = None, ) -> list[int]: ... def filter_row_groups_with_bloom_filters( self, - bloom_filter_data: list[Span], + bloom_filter_data: Sequence[Span], row_group_indices: list[int], options: ParquetReaderOptions, stream: CudaStreamLike | None = None, @@ -74,7 +75,7 @@ class HybridScanReader: def materialize_filter_columns( self, row_group_indices: list[int], - column_chunk_data: list[Span], + column_chunk_data: Sequence[Span], row_mask: Column, mask_data_pages: UseDataPageMask, options: ParquetReaderOptions, @@ -87,7 +88,7 @@ class HybridScanReader: def materialize_payload_columns( self, row_group_indices: list[int], - column_chunk_data: list[Span], + column_chunk_data: Sequence[Span], row_mask: Column, mask_data_pages: UseDataPageMask, options: ParquetReaderOptions, @@ -100,7 +101,7 @@ class HybridScanReader: def materialize_all_columns( self, row_group_indices: list[int], - column_chunk_data: list[Span], + column_chunk_data: Sequence[Span], options: ParquetReaderOptions, stream: CudaStreamLike | None = None, mr: DeviceMemoryResource | None = None, @@ -112,7 +113,7 @@ class HybridScanReader: row_group_indices: list[int], row_mask: Column, mask_data_pages: UseDataPageMask, - column_chunk_data: list[Span], + column_chunk_data: Sequence[Span], options: ParquetReaderOptions, stream: CudaStreamLike | None = None, mr: DeviceMemoryResource | None = None, @@ -128,7 +129,7 @@ class HybridScanReader: row_group_indices: list[int], row_mask: Column, mask_data_pages: UseDataPageMask, - column_chunk_data: list[Span], + column_chunk_data: Sequence[Span], options: ParquetReaderOptions, stream: CudaStreamLike | None = None, mr: DeviceMemoryResource | None = None, From d97901002417a056fbca05cc4f440005983fb2b4 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 11 Jun 2026 20:16:27 +0000 Subject: [PATCH 12/20] prepare hybrid-scan path for future async-prefetching pinned host memory datasource --- .../cudf_polars/cudf_polars/streaming/io.py | 62 +++++++++---------- 1 file changed, 31 insertions(+), 31 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index d665e5aab53e..b9ea2c1fda2a 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -17,7 +17,7 @@ import pylibcudf as plc -from cudf_polars.containers import DataFrame +from cudf_polars.containers import Column, DataFrame from cudf_polars.dsl.ir import ( IR, DataFrameScan, @@ -180,6 +180,18 @@ def expand_scan_for_rank( return scans +def _fetch_byte_ranges( + paths: list[str], + byte_ranges: list[plc.io.text.ByteRangeInfo], + stream: Stream, +) -> list[plc.gpumemoryview]: + # TODO: Accept a pinned-host Datasource pre-fetched by the caller so the + # storage I/O overlaps with GPU work for better pipelining. + return plc.io.parquet_io_utils.fetch_byte_ranges_to_device( + plc.io.SourceInfo(paths), byte_ranges, stream=stream + ) + + def _read_with_hybrid_scan( schema: Schema, paths: list[str], @@ -195,10 +207,8 @@ def _read_with_hybrid_scan( "hybrid scan only supported for SplitScan; one physical file" ) with nvtx_annotate_cudf_polars(message=f"HybridScan: {paths[0]}"): - source_info = plc.io.SourceInfo(paths) - options = ( - plc.io.parquet.ParquetReaderOptions.builder(source_info) + plc.io.parquet.ParquetReaderOptions.builder(plc.io.SourceInfo(paths)) .decimal_width(plc.TypeId.DECIMAL128) .build() ) @@ -219,31 +229,25 @@ def _read_with_hybrid_scan( row_group_indices, options ) if bloom_ranges: - bloom_chunks = plc.io.parquet_io_utils.fetch_byte_ranges_to_device( - source_info, bloom_ranges, stream=stream - ) + bloom_chunks = _fetch_byte_ranges(paths, 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: - byte_ranges = reader.all_column_chunks_byte_ranges( - row_group_indices, options - ) - chunks = plc.io.parquet_io_utils.fetch_byte_ranges_to_device( - source_info, byte_ranges, stream=stream - ) - tbl_w_meta = reader.materialize_all_columns( - row_group_indices, chunks, options, stream=stream - ) - col_names = tbl_w_meta.column_names(include_children=False) - num_rows = tbl_w_meta.num_rows_per_source[0] if not col_names else None - return DataFrame.from_table( - tbl_w_meta.tbl, - col_names, - [schema[name] for name in col_names], + 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, - num_rows=num_rows, ) n_rows = reader.total_rows_in_row_groups(row_group_indices) @@ -258,9 +262,7 @@ def _read_with_hybrid_scan( filter_ranges = reader.filter_column_chunks_byte_ranges( row_group_indices, options ) - filter_chunks = plc.io.parquet_io_utils.fetch_byte_ranges_to_device( - source_info, filter_ranges, stream=stream - ) + filter_chunks = _fetch_byte_ranges(paths, filter_ranges, stream) filter_tbl_w_meta = reader.materialize_filter_columns( row_group_indices, filter_chunks, @@ -273,12 +275,10 @@ def _read_with_hybrid_scan( payload_ranges = reader.payload_column_chunks_byte_ranges( row_group_indices, options ) - # PERFORMANCE!! payload_column_chunks_byte_ranges does not need the row mask, so - # for local NVMe/GDS this fetch could be submitted async before + # PERFORMANCE: payload_column_chunks_byte_ranges does not need the row mask, + # so with async prefetch this fetch could be submitted concurrently with # materialize_filter_columns to overlap I/O with GPU decode. - payload_chunks = plc.io.parquet_io_utils.fetch_byte_ranges_to_device( - source_info, payload_ranges, stream=stream - ) + payload_chunks = _fetch_byte_ranges(paths, payload_ranges, stream) payload_tbl_w_meta = reader.materialize_payload_columns( row_group_indices, payload_chunks, From 64142d0445ae615cf1cfd77246fbffa6320d03ac Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 11 Jun 2026 20:22:40 +0000 Subject: [PATCH 13/20] use SplitScan for SINGLE_FILE even if too small so we can use hybrid scan, see https://github.com/rapidsai/cudf/pull/22857#issuecomment-4684017256 --- python/cudf_polars/cudf_polars/streaming/io.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index b9ea2c1fda2a..b8a62543f9f4 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -86,6 +86,7 @@ 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 if source := stats.scan_stats.get(ir): column_sizes = [ sz @@ -93,8 +94,9 @@ def scan_partition_plan( if (sz := source.column_storage_size(col)) is not None ] if (file_size := sum(column_sizes)) > 0: - if file_size > blocksize: - # Split large files + if file_size > blocksize or single_file: + # A single file always uses SplitScan even if it is smaller than + # the blocksize, so that hybrid scan can be used on it. return IOPartitionPlan( math.ceil(file_size / blocksize), IOPartitionFlavor.SPLIT_FILES, @@ -106,6 +108,9 @@ def scan_partition_plan( IOPartitionFlavor.FUSED_FILES, ) + if single_file: + return IOPartitionPlan(1, IOPartitionFlavor.SPLIT_FILES) + # TODO: Use file sizes for csv and json return IOPartitionPlan(1, IOPartitionFlavor.SINGLE_FILE) From 6627c55b8dc7106c41538278e689350dc5c33c27 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Fri, 12 Jun 2026 13:32:19 +0000 Subject: [PATCH 14/20] prune more rows if page index bytes, release gil in hybrid scan cython APIs --- .../cudf_polars/cudf_polars/streaming/io.py | 19 +- .../pylibcudf/io/experimental/hybrid_scan.pyi | 6 + .../pylibcudf/io/experimental/hybrid_scan.pyx | 202 +++++++++++------- .../pylibcudf/io/parquet_io_utils.pxd | 6 + .../pylibcudf/io/parquet_io_utils.pyi | 6 +- .../pylibcudf/io/parquet_io_utils.pyx | 54 ++++- .../pylibcudf/libcudf/io/datasource.pxd | 5 +- .../pylibcudf/libcudf/io/hybrid_scan.pxd | 6 + .../pylibcudf/libcudf/io/parquet_io_utils.pxd | 6 + 9 files changed, 220 insertions(+), 90 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index b8a62543f9f4..798e4ae7e99e 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -255,14 +255,17 @@ def _read_with_hybrid_scan( stream=stream, ) - n_rows = reader.total_rows_in_row_groups(row_group_indices) - row_mask = plc.Column.from_scalar( - plc.Scalar.from_py( - py_val=True, dtype=plc.DataType(plc.TypeId.BOOL8), stream=stream - ), - n_rows, - stream=stream, - ) + pi_range = reader.page_index_byte_range() + if pi_range.size > 0: + page_index_bytes = plc.io.parquet_io_utils.fetch_page_index_to_host( + plc.io.SourceInfo(paths), pi_range + ) + reader.setup_page_index(page_index_bytes) + row_mask = reader.build_row_mask_with_page_index_stats( + row_group_indices, options, stream=stream + ) + else: + row_mask = reader.build_all_true_row_mask(row_group_indices, stream=stream) filter_ranges = reader.filter_column_chunks_byte_ranges( row_group_indices, options diff --git a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyi b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyi index b6a7ebafb907..490d7b8cbd15 100644 --- a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyi +++ b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyi @@ -62,6 +62,12 @@ class HybridScanReader: options: ParquetReaderOptions, stream: CudaStreamLike | None = None, ) -> list[int]: ... + def build_all_true_row_mask( + self, + row_group_indices: list[int], + stream: CudaStreamLike | None = None, + mr: DeviceMemoryResource | None = None, + ) -> Column: ... def build_row_mask_with_page_index_stats( self, row_group_indices: list[int], diff --git a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx index 0f26440fce26..943f62f727b4 100644 --- a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx +++ b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx @@ -137,9 +137,10 @@ cdef class HybridScanReader: page_index_bytes : Buffer Parquet page index buffer bytes """ - self.c_obj.get()[0].setup_page_index( - host_span[const_uint8_t](&page_index_bytes[0], len(page_index_bytes)) - ) + with nogil: + self.c_obj.get()[0].setup_page_index( + host_span[const_uint8_t](&page_index_bytes[0], len(page_index_bytes)) + ) def all_row_groups(self, ParquetReaderOptions options): """Get all available row groups from the parquet file. @@ -209,15 +210,15 @@ cdef class HybridScanReader: """ cdef Stream _stream = _get_stream(stream) cdef vector[size_type] indices_vec = row_group_indices - cdef vector[size_type] filtered = ( - self.c_obj.get()[0].filter_row_groups_with_stats( + cdef vector[size_type] filtered + with nogil: + filtered = move(self.c_obj.get()[0].filter_row_groups_with_stats( host_span[const_size_type]( indices_vec.data(), indices_vec.size() ), options.c_obj, _stream.view().value() - ) - ) + )) return list(filtered) def secondary_filters_byte_ranges( @@ -240,11 +241,12 @@ cdef class HybridScanReader: Tuple of (bloom_filter_ranges, dictionary_page_ranges) """ cdef vector[size_type] indices_vec = row_group_indices - cdef pair[vector[byte_range_info], vector[byte_range_info]] ranges = \ - self.c_obj.get()[0].secondary_filters_byte_ranges( + cdef pair[vector[byte_range_info], vector[byte_range_info]] ranges + with nogil: + ranges = move(self.c_obj.get()[0].secondary_filters_byte_ranges( host_span[const_size_type](indices_vec.data(), indices_vec.size()), options.c_obj - ) + )) bloom_ranges = [ ByteRangeInfo(r.offset(), r.size()) for r in ranges.first @@ -286,15 +288,16 @@ cdef class HybridScanReader: cdef vector[size_type] indices_vec = row_group_indices - cdef vector[size_type] filtered = \ - self.c_obj.get()[0].filter_row_groups_with_dictionary_pages( + cdef vector[size_type] filtered + with nogil: + filtered = move(self.c_obj.get()[0].filter_row_groups_with_dictionary_pages( host_span[const_device_span_const_uint8_t]( spans_vec.data(), spans_vec.size() ), host_span[const_size_type](indices_vec.data(), indices_vec.size()), options.c_obj, _stream.view().value() - ) + )) return list(filtered) def filter_row_groups_with_bloom_filters( @@ -329,17 +332,52 @@ cdef class HybridScanReader: cdef vector[size_type] indices_vec = row_group_indices - cdef vector[size_type] filtered = \ - self.c_obj.get()[0].filter_row_groups_with_bloom_filters( + cdef vector[size_type] filtered + with nogil: + filtered = move(self.c_obj.get()[0].filter_row_groups_with_bloom_filters( host_span[const_device_span_const_uint8_t]( spans_vec.data(), spans_vec.size() ), host_span[const_size_type](indices_vec.data(), indices_vec.size()), options.c_obj, _stream.view().value() - ) + )) return list(filtered) + def build_all_true_row_mask( + self, + list row_group_indices, + object stream=None, + DeviceMemoryResource mr=None + ): + """Build an all-true boolean survival column for the given row groups. + + Parameters + ---------- + row_group_indices : list[int] + Input row group indices + stream : Stream, optional + CUDA stream + mr : DeviceMemoryResource, optional + Device memory resource + + Returns + ------- + Column + All-true boolean column with one entry per row across all row groups + """ + cdef vector[size_type] indices_vec = row_group_indices + cdef Stream _stream = _get_stream(stream) + mr = _get_memory_resource(mr) + cdef unique_ptr[column] c_result + with nogil: + c_result = move(self.c_obj.get()[0].build_all_true_row_mask( + host_span[const_size_type](indices_vec.data(), indices_vec.size()), + _stream.view().value(), + mr.get_mr() + )) + return Column.from_libcudf(move(c_result), _stream, mr) + def build_row_mask_with_page_index_stats( self, list row_group_indices, @@ -368,13 +406,14 @@ cdef class HybridScanReader: cdef vector[size_type] indices_vec = row_group_indices cdef Stream _stream = _get_stream(stream) mr = _get_memory_resource(mr) - cdef unique_ptr[column] c_result = \ - self.c_obj.get()[0].build_row_mask_with_page_index_stats( + cdef unique_ptr[column] c_result + with nogil: + c_result = move(self.c_obj.get()[0].build_row_mask_with_page_index_stats( host_span[const_size_type](indices_vec.data(), indices_vec.size()), options.c_obj, _stream.view().value(), mr.get_mr() - ) + )) return Column.from_libcudf(move(c_result), _stream, mr) def filter_column_chunks_byte_ranges( @@ -397,11 +436,12 @@ cdef class HybridScanReader: Byte ranges to column chunks of filter columns """ cdef vector[size_type] indices_vec = row_group_indices - cdef vector[byte_range_info] ranges = \ - self.c_obj.get()[0].filter_column_chunks_byte_ranges( + cdef vector[byte_range_info] ranges + with nogil: + ranges = move(self.c_obj.get()[0].filter_column_chunks_byte_ranges( host_span[const_size_type](indices_vec.data(), indices_vec.size()), options.c_obj - ) + )) return [ByteRangeInfo(r.offset(), r.size()) for r in ranges] def materialize_filter_columns( @@ -447,8 +487,9 @@ cdef class HybridScanReader: spans_vec.push_back(_get_device_span(span)) cdef mutable_column_view mask_view = row_mask.mutable_view() - cdef table_with_metadata c_result = \ - self.c_obj.get()[0].materialize_filter_columns( + cdef table_with_metadata c_result + with nogil: + c_result = move(self.c_obj.get()[0].materialize_filter_columns( host_span[const_size_type](indices_vec.data(), indices_vec.size()), host_span[const_device_span_const_uint8_t]( spans_vec.data(), spans_vec.size() @@ -458,7 +499,7 @@ cdef class HybridScanReader: options.c_obj, _stream.view().value(), mr.get_mr() - ) + )) return TableWithMetadata.from_libcudf(c_result, _stream, mr) def payload_column_chunks_byte_ranges( @@ -481,11 +522,12 @@ cdef class HybridScanReader: Byte ranges to column chunks of payload columns """ cdef vector[size_type] indices_vec = row_group_indices - cdef vector[byte_range_info] ranges = \ - self.c_obj.get()[0].payload_column_chunks_byte_ranges( + cdef vector[byte_range_info] ranges + with nogil: + ranges = move(self.c_obj.get()[0].payload_column_chunks_byte_ranges( host_span[const_size_type](indices_vec.data(), indices_vec.size()), options.c_obj - ) + )) return [ByteRangeInfo(r.offset(), r.size()) for r in ranges] def materialize_payload_columns( @@ -531,8 +573,9 @@ cdef class HybridScanReader: spans_vec.push_back(_get_device_span(span)) cdef column_view mask_view = row_mask.view() - cdef table_with_metadata c_result = \ - self.c_obj.get()[0].materialize_payload_columns( + cdef table_with_metadata c_result + with nogil: + c_result = move(self.c_obj.get()[0].materialize_payload_columns( host_span[const_size_type](indices_vec.data(), indices_vec.size()), host_span[const_device_span_const_uint8_t]( spans_vec.data(), spans_vec.size() @@ -542,7 +585,7 @@ cdef class HybridScanReader: options.c_obj, _stream.view().value(), mr.get_mr() - ) + )) return TableWithMetadata.from_libcudf(c_result, _stream, mr) def all_column_chunks_byte_ranges( @@ -565,11 +608,12 @@ cdef class HybridScanReader: Byte ranges to column chunks of all columns """ cdef vector[size_type] indices_vec = row_group_indices - cdef vector[byte_range_info] ranges = \ - self.c_obj.get()[0].all_column_chunks_byte_ranges( + cdef vector[byte_range_info] ranges + with nogil: + ranges = move(self.c_obj.get()[0].all_column_chunks_byte_ranges( host_span[const_size_type](indices_vec.data(), indices_vec.size()), options.c_obj - ) + )) return [ByteRangeInfo(r.offset(), r.size()) for r in ranges] def materialize_all_columns( @@ -607,8 +651,9 @@ cdef class HybridScanReader: mr = _get_memory_resource(mr) for span in column_chunk_data: spans_vec.push_back(_get_device_span(span)) - cdef table_with_metadata c_result = \ - self.c_obj.get()[0].materialize_all_columns( + cdef table_with_metadata c_result + with nogil: + c_result = move(self.c_obj.get()[0].materialize_all_columns( host_span[const_size_type](indices_vec.data(), indices_vec.size()), host_span[const_device_span_const_uint8_t]( spans_vec.data(), spans_vec.size() @@ -616,7 +661,7 @@ cdef class HybridScanReader: options.c_obj, _stream.view().value(), mr.get_mr() - ) + )) return TableWithMetadata.from_libcudf(c_result, _stream, mr) def setup_chunking_for_filter_columns( @@ -664,19 +709,20 @@ cdef class HybridScanReader: self.mr = _get_memory_resource(mr) cdef column_view mask_view = row_mask.view() - self.c_obj.get()[0].setup_chunking_for_filter_columns( - chunk_read_limit, - pass_read_limit, - host_span[const_size_type](indices_vec.data(), indices_vec.size()), - mask_view, - mask_data_pages, - host_span[const_device_span_const_uint8_t]( - spans_vec.data(), spans_vec.size() - ), - options.c_obj, - self._stream.view().value(), - self.mr.get_mr() - ) + with nogil: + self.c_obj.get()[0].setup_chunking_for_filter_columns( + chunk_read_limit, + pass_read_limit, + host_span[const_size_type](indices_vec.data(), indices_vec.size()), + mask_view, + mask_data_pages, + host_span[const_device_span_const_uint8_t]( + spans_vec.data(), spans_vec.size() + ), + options.c_obj, + self._stream.view().value(), + self.mr.get_mr() + ) def materialize_filter_columns_chunk( self, @@ -694,10 +740,11 @@ cdef class HybridScanReader: Table chunk of materialized filter columns and metadata """ cdef mutable_column_view mask_view = row_mask.mutable_view() - cdef table_with_metadata c_result = \ - self.c_obj.get()[0].materialize_filter_columns_chunk( + cdef table_with_metadata c_result + with nogil: + c_result = move(self.c_obj.get()[0].materialize_filter_columns_chunk( mask_view - ) + )) return TableWithMetadata.from_libcudf( c_result, self._stream, self.mr ) @@ -747,19 +794,20 @@ cdef class HybridScanReader: self.mr = _get_memory_resource(mr) cdef column_view mask_view = row_mask.view() - self.c_obj.get()[0].setup_chunking_for_payload_columns( - chunk_read_limit, - pass_read_limit, - host_span[const_size_type](indices_vec.data(), indices_vec.size()), - mask_view, - mask_data_pages, - host_span[const_device_span_const_uint8_t]( - spans_vec.data(), spans_vec.size() - ), - options.c_obj, - self._stream.view().value(), - self.mr.get_mr() - ) + with nogil: + self.c_obj.get()[0].setup_chunking_for_payload_columns( + chunk_read_limit, + pass_read_limit, + host_span[const_size_type](indices_vec.data(), indices_vec.size()), + mask_view, + mask_data_pages, + host_span[const_device_span_const_uint8_t]( + spans_vec.data(), spans_vec.size() + ), + options.c_obj, + self._stream.view().value(), + self.mr.get_mr() + ) def materialize_payload_columns_chunk( self, @@ -777,10 +825,11 @@ cdef class HybridScanReader: Table chunk of materialized payload columns and metadata """ cdef column_view mask_view = row_mask.view() - cdef table_with_metadata c_result = \ - self.c_obj.get()[0].materialize_payload_columns_chunk( + cdef table_with_metadata c_result + with nogil: + c_result = move(self.c_obj.get()[0].materialize_payload_columns_chunk( mask_view - ) + )) return TableWithMetadata.from_libcudf( c_result, self._stream, self.mr ) @@ -816,12 +865,15 @@ cdef class HybridScanReader: If ``row_group_indices`` is empty. """ cdef vector[size_type] indices_vec = row_group_indices - return self.c_obj.get()[0].construct_row_group_passes( - host_span[const_size_type]( - indices_vec.data(), indices_vec.size() - ), - pass_read_limit - ) + cdef vector[vector[size_type]] passes + with nogil: + passes = move(self.c_obj.get()[0].construct_row_group_passes( + host_span[const_size_type]( + indices_vec.data(), indices_vec.size() + ), + pass_read_limit + )) + return passes def has_next_table_chunk(self): """Check if there is any parquet data left to read. diff --git a/python/pylibcudf/pylibcudf/io/parquet_io_utils.pxd b/python/pylibcudf/pylibcudf/io/parquet_io_utils.pxd index f0bf640c640a..2c6417a332d3 100644 --- a/python/pylibcudf/pylibcudf/io/parquet_io_utils.pxd +++ b/python/pylibcudf/pylibcudf/io/parquet_io_utils.pxd @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 +from pylibcudf.io.text cimport ByteRangeInfo from pylibcudf.io.types cimport SourceInfo from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource @@ -10,3 +11,8 @@ cpdef list fetch_byte_ranges_to_device( object stream=*, DeviceMemoryResource mr=*, ) + +cpdef bytes fetch_page_index_to_host( + SourceInfo source_info, + ByteRangeInfo page_index_range, +) diff --git a/python/pylibcudf/pylibcudf/io/parquet_io_utils.pyi b/python/pylibcudf/pylibcudf/io/parquet_io_utils.pyi index 98feaf34504b..fc29d955b42f 100644 --- a/python/pylibcudf/pylibcudf/io/parquet_io_utils.pyi +++ b/python/pylibcudf/pylibcudf/io/parquet_io_utils.pyi @@ -8,7 +8,7 @@ from pylibcudf.io.text import ByteRangeInfo from pylibcudf.io.types import SourceInfo from pylibcudf.utils import CudaStreamLike -__all__ = ["fetch_byte_ranges_to_device"] +__all__ = ["fetch_byte_ranges_to_device", "fetch_page_index_to_host"] def fetch_byte_ranges_to_device( source_info: SourceInfo, @@ -16,3 +16,7 @@ def fetch_byte_ranges_to_device( stream: CudaStreamLike | None = None, mr: DeviceMemoryResource | None = None, ) -> list[gpumemoryview]: ... +def fetch_page_index_to_host( + source_info: SourceInfo, + page_index_range: ByteRangeInfo, +) -> bytes: ... diff --git a/python/pylibcudf/pylibcudf/io/parquet_io_utils.pyx b/python/pylibcudf/pylibcudf/io/parquet_io_utils.pyx index 4fe084f21fbc..bf562d0f0579 100644 --- a/python/pylibcudf/pylibcudf/io/parquet_io_utils.pyx +++ b/python/pylibcudf/pylibcudf/io/parquet_io_utils.pyx @@ -3,7 +3,7 @@ """IO utilities for the Parquet.""" from libc.stddef cimport size_t -from libc.stdint cimport uintptr_t +from libc.stdint cimport uint8_t, uintptr_t from libcpp.memory cimport make_unique, unique_ptr from libcpp.pair cimport pair from libcpp.utility cimport move @@ -23,20 +23,21 @@ from pylibcudf.libcudf.io.parquet_io_utils cimport ( const_byte_range_info, const_uint8_t, fetch_byte_ranges_to_device as cpp_fetch_byte_ranges_to_device, + fetch_page_index_to_host as cpp_fetch_page_index_to_host, ) from pylibcudf.libcudf.io.text cimport byte_range_info from pylibcudf.libcudf.utilities.span cimport device_span, host_span from pylibcudf.utils cimport _get_memory_resource, _get_stream -__all__ = ["fetch_byte_ranges_to_device"] +__all__ = ["fetch_byte_ranges_to_device", "fetch_page_index_to_host"] -def fetch_byte_ranges_to_device( +cpdef list fetch_byte_ranges_to_device( SourceInfo source_info, list byte_ranges, object stream=None, - object mr=None, -) -> list[gpumemoryview]: + DeviceMemoryResource mr=None, +): """Fetch byte ranges from a Parquet source into device memory. Parameters @@ -117,3 +118,46 @@ def fetch_byte_ranges_to_device( } result.append(gmv) return result + + +cpdef bytes fetch_page_index_to_host( + SourceInfo source_info, + ByteRangeInfo page_index_range, +): + """Fetch parquet page index bytes to host memory. + + Parameters + ---------- + source_info : SourceInfo + Source describing a single Parquet file. + page_index_range : ByteRangeInfo + Byte range of the page index, as returned by + :meth:`~pylibcudf.io.experimental.HybridScanReader.page_index_byte_range`. + + Returns + ------- + bytes + Raw page index bytes copied to Python host memory. + + Raises + ------ + ValueError + If ``source_info`` does not describe exactly one source. + """ + cdef vector[unique_ptr[datasource]] sources = make_datasources(source_info.c_obj) + if sources.size() != 1: + raise ValueError( + f"fetch_page_index_to_host requires exactly one source, " + f"got {sources.size()}" + ) + + cdef unique_ptr[datasource.buffer] buf + with nogil: + buf = move(cpp_fetch_page_index_to_host( + dereference(sources[0]), + (page_index_range).c_obj, + )) + + cdef const uint8_t* ptr = buf.get().data() + cdef size_t n = buf.get().size() + return bytes(ptr[:n]) diff --git a/python/pylibcudf/pylibcudf/libcudf/io/datasource.pxd b/python/pylibcudf/pylibcudf/libcudf/io/datasource.pxd index 36c5bf928850..4c582e57e7a8 100644 --- a/python/pylibcudf/pylibcudf/libcudf/io/datasource.pxd +++ b/python/pylibcudf/pylibcudf/libcudf/io/datasource.pxd @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 from libc.stddef cimport size_t +from libc.stdint cimport uint8_t from libcpp.memory cimport unique_ptr from libcpp.vector cimport vector from pylibcudf.libcudf.io.types cimport source_info @@ -11,7 +12,9 @@ cdef extern from "cudf/io/datasource.hpp" \ namespace "cudf::io" nogil: cdef cppclass datasource: - pass + cdef cppclass buffer: + size_t size() const nogil + const uint8_t* data() const nogil cdef vector[unique_ptr[datasource]] make_datasources( source_info info diff --git a/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan.pxd b/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan.pxd index 8578908fc43c..5dc33a5be1d4 100644 --- a/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan.pxd +++ b/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan.pxd @@ -85,6 +85,12 @@ cdef extern from "cudf/io/experimental/hybrid_scan.hpp" \ cudaStream_t stream ) except +libcudf_exception_handler + unique_ptr[column] build_all_true_row_mask( + host_span[const_size_type] row_group_indices, + cudaStream_t stream, + device_async_resource_ref mr + ) except +libcudf_exception_handler + unique_ptr[column] build_row_mask_with_page_index_stats( host_span[const_size_type] row_group_indices, const parquet_reader_options& options, diff --git a/python/pylibcudf/pylibcudf/libcudf/io/parquet_io_utils.pxd b/python/pylibcudf/pylibcudf/libcudf/io/parquet_io_utils.pxd index fe613c5c3da2..ed1fa90e5695 100644 --- a/python/pylibcudf/pylibcudf/libcudf/io/parquet_io_utils.pxd +++ b/python/pylibcudf/pylibcudf/libcudf/io/parquet_io_utils.pxd @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 from libc.stdint cimport uint8_t +from libcpp.memory cimport unique_ptr from libcpp.pair cimport pair from libcpp.vector cimport vector @@ -27,3 +28,8 @@ cdef extern from "cudf/io/parquet_io_utils.hpp" \ cuda_stream_view stream, device_async_resource_ref mr, ) except +libcudf_exception_handler + + unique_ptr[datasource.buffer] fetch_page_index_to_host( + datasource& ds, + byte_range_info page_index_bytes, + ) except +libcudf_exception_handler From 106d6b59b4aa46e492926659105fd8d35426c482 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Fri, 12 Jun 2026 21:17:02 +0000 Subject: [PATCH 15/20] add C++ API read for cudf_polars to avoid repeated GIL acqusition/release --- .../cudf/io/experimental/hybrid_scan.hpp | 63 +++++ .../io/parquet/experimental/hybrid_scan.cpp | 221 ++++++++++++++++++ .../io/experimental/hybrid_scan_test.cpp | 77 ++++++ .../cudf_polars/cudf_polars/streaming/io.py | 168 ++++++------- .../cudf_polars/cudf_polars/utils/config.py | 18 ++ .../pylibcudf/io/experimental/hybrid_scan.pyx | 146 ++++++++++-- .../pylibcudf/libcudf/io/datasource.pxd | 6 +- .../pylibcudf/libcudf/io/hybrid_scan.pxd | 19 ++ 8 files changed, 599 insertions(+), 119 deletions(-) diff --git a/cpp/include/cudf/io/experimental/hybrid_scan.hpp b/cpp/include/cudf/io/experimental/hybrid_scan.hpp index 980ab9644d3b..bdf1383bbe4c 100644 --- a/cpp/include/cudf/io/experimental/hybrid_scan.hpp +++ b/cpp/include/cudf/io/experimental/hybrid_scan.hpp @@ -5,6 +5,7 @@ #pragma once +#include #include #include #include @@ -47,6 +48,27 @@ enum class use_data_page_mask : bool { NO = false ///< Do not compute or use a data page mask }; +/** + * @brief Options controlling a fused hybrid scan read via `hybrid_scan_reader::read` + * + * Each row-group pruning stage is optional and is additionally skipped at runtime when the + * parquet file does not contain the corresponding structure (e.g. no bloom filters) or when a + * prior stage has already pruned all row groups. + */ +struct hybrid_scan_read_options { + bool use_stats_filter = true; ///< Prune row groups using column chunk statistics + bool use_dictionary_filter = true; ///< Prune row groups using column chunk dictionary pages + bool use_bloom_filter = true; ///< Prune row groups using column chunk bloom filters + /// Whether to build and use a data page mask to prune filter column pages + use_data_page_mask prune_filter_column_pages = use_data_page_mask::NO; + /// Whether to build and use a data page mask to prune payload column pages + use_data_page_mask prune_payload_column_pages = use_data_page_mask::YES; + /// Host span of page index bytes. If empty and the file contains a page index, the bytes are + /// fetched from the data source. Pass cached bytes to avoid a redundant read across reads of the + /// same file. + cudf::host_span page_index_bytes = {}; +}; + /** * @brief The experimental parquet reader class to optimally read parquet files subject to * highly selective filters, called a Hybrid Scan operation @@ -300,6 +322,47 @@ class hybrid_scan_reader { */ ~hybrid_scan_reader(); + /** + * @brief Reads the Parquet file in a single fused operation + * + * Performs the complete hybrid scan in one call: optional row group pruning (statistics, + * dictionary pages and bloom filters), row mask construction, and the two-pass filter then + * payload column materialization. When `options` carries no filter expression, all selected + * columns are read in a single pass instead. + * + * This is equivalent in result to constructing the reader and invoking the individual step + * methods in sequence (see the class documentation), but issues all byte range fetches and GPU + * work from within this call. Callers that drive the reader from a managed runtime therefore + * cross the language boundary once per read rather than once per step, which avoids repeated + * lock acquisition (e.g. the Python GIL) between steps. + * + * The page index required for the row mask is taken from `read_options.page_index_bytes` when + * provided, and otherwise fetched from `source`. The row mask is built from page-level + * statistics when a page index is available and all-true otherwise. + * + * Output columns are returned in the projection order given by `options.get_column_names()`, + * or in Parquet file schema order when no column selection is set, matching + * `cudf::io::read_parquet`. + * + * @note `read` reads from a single data source. The reader holds per-read state, so a single + * reader instance must not be driven concurrently from multiple threads. + * + * @param source Data source for the Parquet file backing this reader + * @param row_group_indices Candidate row group indices to read, before pruning + * @param options Parquet reader options, including the optional filter expression and column + * selection + * @param read_options Options controlling row group pruning, data page masking and the page index + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource used to allocate the device memory for the output table + * @return Table of materialized columns and metadata, in projection order + */ + [[nodiscard]] table_with_metadata read(cudf::io::datasource& source, + cudf::host_span row_group_indices, + parquet_reader_options const& options, + hybrid_scan_read_options const& read_options, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const; + /** * @brief Get the Parquet file footer metadata * diff --git a/cpp/src/io/parquet/experimental/hybrid_scan.cpp b/cpp/src/io/parquet/experimental/hybrid_scan.cpp index 868813a1b4ed..f20957f1c928 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan.cpp @@ -7,12 +7,117 @@ #include #include +#include +#include #include +#include +#include + #include +#include +#include +#include +#include + namespace cudf::io::parquet::experimental { +namespace { + +/** + * @brief Collects the top-level column names of a Parquet file in schema order + * + * Walks the flattened Parquet schema tree depth-first and records the names of the immediate + * children of the root, skipping each child's descendants. This yields the column order that + * `cudf::io::read_parquet` produces when no column selection is set. + * + * @param metadata Parquet file footer metadata + * @return Top-level column names in file schema order + */ +[[nodiscard]] std::vector top_level_column_names(FileMetaData const& metadata) +{ + auto const& schema = metadata.schema; + if (schema.empty()) { return {}; } + + std::vector names; + names.reserve(schema.front().num_children); + + // Depth-first walk tracking unvisited children remaining at each open level. The root (index 0) + // is skipped; a column is top-level when the only open level is the root's. + std::vector remaining_children{schema.front().num_children}; + for (std::size_t i = 1; i < schema.size() and not remaining_children.empty(); ++i) { + auto const& element = schema[i]; + if (remaining_children.size() == 1) { names.push_back(element.name); } + + --remaining_children.back(); + if (element.num_children > 0) { + remaining_children.push_back(element.num_children); + } else { + while (not remaining_children.empty() and remaining_children.back() == 0) { + remaining_children.pop_back(); + } + } + } + + return names; +} + +/** + * @brief Reassembles the filter and payload tables into a single table in `output_order` + * + * The two passes of a hybrid scan materialize disjoint sets of columns (filter columns and + * payload columns). This combines them into one table whose columns follow `output_order`, + * matching the projection order of `cudf::io::read_parquet`. + * + * @param filter Materialized filter columns and metadata + * @param payload Materialized payload columns and metadata + * @param output_order Desired output column names in order + * @return Combined table and metadata in `output_order` + */ +[[nodiscard]] table_with_metadata assemble_output(table_with_metadata filter, + table_with_metadata payload, + std::vector const& output_order) +{ + enum class source : bool { filter, payload }; + std::unordered_map> location; + location.reserve(filter.metadata.schema_info.size() + payload.metadata.schema_info.size()); + for (std::size_t i = 0; i < filter.metadata.schema_info.size(); ++i) { + location.emplace(filter.metadata.schema_info[i].name, std::pair{source::filter, i}); + } + for (std::size_t i = 0; i < payload.metadata.schema_info.size(); ++i) { + location.emplace(payload.metadata.schema_info[i].name, std::pair{source::payload, i}); + } + + auto filter_columns = filter.tbl->release(); + auto payload_columns = payload.tbl->release(); + + std::vector> output_columns; + output_columns.reserve(output_order.size()); + table_metadata out_metadata; + out_metadata.schema_info.reserve(output_order.size()); + out_metadata.num_rows_per_source = filter.metadata.num_rows_per_source; + + for (auto const& name : output_order) { + auto const it = location.find(name); + CUDF_EXPECTS(it != location.end(), + "Projected column not found in materialized hybrid scan output: " + name); + auto const [tbl, pos] = it->second; + if (tbl == source::filter) { + output_columns.push_back(std::move(filter_columns[pos])); + out_metadata.schema_info.push_back(std::move(filter.metadata.schema_info[pos])); + } else { + output_columns.push_back(std::move(payload_columns[pos])); + out_metadata.schema_info.push_back(std::move(payload.metadata.schema_info[pos])); + } + } + + return table_with_metadata{std::make_unique(std::move(output_columns)), + std::move(out_metadata)}; +} + +} // namespace + hybrid_scan_reader::hybrid_scan_reader(cudf::host_span footer_bytes, parquet_reader_options const& options) : _impl{std::make_unique( @@ -29,6 +134,122 @@ hybrid_scan_reader::hybrid_scan_reader(FileMetaData const& parquet_metadata, hybrid_scan_reader::~hybrid_scan_reader() = default; +table_with_metadata hybrid_scan_reader::read(cudf::io::datasource& source, + cudf::host_span row_group_indices, + parquet_reader_options const& options, + hybrid_scan_read_options const& read_options, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const +{ + CUDF_FUNC_RANGE(); + + // Transient byte-range fetches use the current device resource; output tables use `mr`. + auto temp_mr = cudf::get_current_device_resource_ref(); + + auto const output_order = options.get_column_names().has_value() + ? options.get_column_names().value() + : top_level_column_names(_impl->parquet_metadatas().front()); + + std::vector current_row_groups(row_group_indices.begin(), row_group_indices.end()); + + // No filter: read all selected columns in a single pass. + if (not options.get_filter().has_value()) { + auto const byte_ranges = all_column_chunks_byte_ranges(current_row_groups, options); + auto [buffers, data, tasks] = + parquet::fetch_byte_ranges_to_device_async(source, byte_ranges, stream, temp_mr); + tasks.get(); + return materialize_all_columns(current_row_groups, data, options, stream, mr); + } + + // Use caller-provided page index bytes when available, otherwise fetch them from the source. + auto const page_index_range = page_index_byte_range(); + auto const has_page_index = not page_index_range.is_empty(); + if (has_page_index) { + if (not read_options.page_index_bytes.empty()) { + setup_page_index(read_options.page_index_bytes); + } else { + auto const page_index_buffer = parquet::fetch_page_index_to_host(source, page_index_range); + setup_page_index( + cudf::host_span{page_index_buffer->data(), page_index_buffer->size()}); + } + } + + if (read_options.use_stats_filter and not current_row_groups.empty()) { + current_row_groups = filter_row_groups_with_stats(current_row_groups, options, stream); + } + + if ((read_options.use_dictionary_filter or read_options.use_bloom_filter) and + not current_row_groups.empty()) { + auto const [bloom_filter_ranges, dictionary_page_ranges] = + secondary_filters_byte_ranges(current_row_groups, options); + + if (read_options.use_dictionary_filter and not dictionary_page_ranges.empty()) { + auto [buffers, data, tasks] = + parquet::fetch_byte_ranges_to_device_async(source, dictionary_page_ranges, stream, temp_mr); + tasks.get(); + current_row_groups = + filter_row_groups_with_dictionary_pages(data, current_row_groups, options, stream); + } + + if (read_options.use_bloom_filter and not bloom_filter_ranges.empty() and + not current_row_groups.empty()) { + // Bloom filter data buffers must be allocated on 32-byte aligned addresses. + auto aligned_mr = rmm::mr::aligned_resource_adaptor{temp_mr, rmm::CUDA_ALLOCATION_ALIGNMENT}; + auto [buffers, data, tasks] = + parquet::fetch_byte_ranges_to_device_async(source, bloom_filter_ranges, stream, aligned_mr); + tasks.get(); + current_row_groups = + filter_row_groups_with_bloom_filters(data, current_row_groups, options, stream); + } + } + + // All row groups pruned: return a correctly typed, zero-row table in projection order. The + // materialization path requires a non-empty row mask, so defer to the main reader here. + if (current_row_groups.empty()) { + auto empty_options = options; + empty_options.set_num_rows(0); + return cudf::io::read_parquet(empty_options, stream, mr); + } + + // Use page-level statistics for the row mask when a page index is available; otherwise all-true. + auto row_mask = has_page_index + ? build_row_mask_with_page_index_stats(current_row_groups, options, stream, mr) + : build_all_true_row_mask(current_row_groups, stream, mr); + + // Filter pass: materialize filter columns and narrow the row mask to surviving rows. + auto row_mask_view = row_mask->mutable_view(); + auto filter_table = [&] { + auto const byte_ranges = filter_column_chunks_byte_ranges(current_row_groups, options); + auto [buffers, data, tasks] = + parquet::fetch_byte_ranges_to_device_async(source, byte_ranges, stream, temp_mr); + tasks.get(); + return materialize_filter_columns(current_row_groups, + data, + row_mask_view, + read_options.prune_filter_column_pages, + options, + stream, + mr); + }(); + + // Payload pass: materialize payload columns under the surviving row mask. + auto payload_table = [&] { + auto const byte_ranges = payload_column_chunks_byte_ranges(current_row_groups, options); + auto [buffers, data, tasks] = + parquet::fetch_byte_ranges_to_device_async(source, byte_ranges, stream, temp_mr); + tasks.get(); + return materialize_payload_columns(current_row_groups, + data, + row_mask->view(), + read_options.prune_payload_column_pages, + options, + stream, + mr); + }(); + + return assemble_output(std::move(filter_table), std::move(payload_table), output_order); +} + [[nodiscard]] text::byte_range_info hybrid_scan_reader::page_index_byte_range() const { return _impl->page_index_byte_ranges().front(); diff --git a/cpp/tests/io/experimental/hybrid_scan_test.cpp b/cpp/tests/io/experimental/hybrid_scan_test.cpp index 05e7fba3d844..b6b019082b05 100644 --- a/cpp/tests/io/experimental/hybrid_scan_test.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_test.cpp @@ -28,6 +28,8 @@ #include +#include + namespace { /** @@ -923,6 +925,81 @@ TEST_F(HybridScanTest, StructChildFilterColumn) std::invalid_argument); } +namespace { + +// Drives the fused `hybrid_scan_reader::read` and compares against `cudf::io::read_parquet` for the +// same filter and column selection. `selected_columns` deliberately reorders columns and places the +// filter column away from the front, exercising the projection-order reassembly inside `read`. +void test_read_matches_read_parquet(std::vector const& parquet_buffer, + cudf::ast::operation const& filter_expression, + std::vector const& selected_columns) +{ + auto const stream = cudf::get_default_stream(); + auto const mr = cudf::get_current_device_resource_ref(); + + auto const make_options = [&] { + auto options = cudf::io::parquet_reader_options::builder(cudf::io::source_info( + cudf::host_span( + parquet_buffer.data(), + parquet_buffer.size()))) + .filter(filter_expression) + .build(); + options.set_column_names(selected_columns); + return options; + }; + + auto const expected = cudf::io::read_parquet(make_options(), stream, mr).tbl; + + auto const options = make_options(); + auto datasource = cudf::io::datasource::create(cudf::host_span( + reinterpret_cast(parquet_buffer.data()), parquet_buffer.size())); + auto const footer_buffer = cudf::io::parquet::fetch_footer_to_host(*datasource); + auto reader = + std::make_unique(*footer_buffer, options); + + auto const row_groups = reader->all_row_groups(options); + auto const read_options = + cudf::io::parquet::experimental::hybrid_scan_read_options{}; // page index fetched internally + auto const actual = reader->read(*datasource, row_groups, options, read_options, stream, mr); + + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected->view(), actual.tbl->view()); +} + +} // namespace + +TEST_F(HybridScanTest, ReadMatchesReadParquetReorderedProjection) +{ + using T = int32_t; + auto constexpr num_concat = 2; // multiple row groups + pages so pruning engages + auto [written_table, parquet_buffer] = create_parquet_with_stats(); + + // Filter on the middle column; project columns in a different order with the filter column last. + auto literal_value = cudf::numeric_scalar(T{500}); + auto literal = cudf::ast::literal(literal_value); + auto col_ref_1 = cudf::ast::column_name_reference("col1"); + auto filter_expression = cudf::ast::operation(cudf::ast::ast_operator::LESS, col_ref_1, literal); + + test_read_matches_read_parquet( + parquet_buffer, filter_expression, std::vector{"col2", "col0", "col1"}); +} + +TEST_F(HybridScanTest, ReadMatchesReadParquetNoSurvivingRows) +{ + using T = int32_t; + auto constexpr num_concat = 2; + auto [written_table, parquet_buffer] = create_parquet_with_stats(); + + // A predicate no row can satisfy, exercising the empty-output path of `read`. + auto literal_value = cudf::numeric_scalar(std::numeric_limits::min()); + auto literal = cudf::ast::literal(literal_value); + auto col_ref_0 = cudf::ast::column_name_reference("col0"); + auto filter_expression = + cudf::ast::operation(cudf::ast::ast_operator::LESS, col_ref_0, literal); + + test_read_matches_read_parquet( + parquet_buffer, filter_expression, std::vector{"col2", "col0"}); +} + TEST_F(HybridScanTest, ChunkedReadRowMaskPerPass) { using T = uint32_t; diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 1c094c5332e6..cace8524cfef 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -189,18 +189,30 @@ def expand_scan_for_rank( return scans -def _fetch_byte_ranges( - paths: list[str], - byte_ranges: list[plc.io.text.ByteRangeInfo], - stream: Stream, -) -> list[plc.gpumemoryview]: - # TODO: Accept a pinned-host Datasource pre-fetched by the caller so the - # storage I/O overlaps with GPU work for better pipelining. - return plc.io.parquet_io_utils.fetch_byte_ranges_to_device( - plc.io.SourceInfo(paths), byte_ranges, stream=stream +@functools.lru_cache(maxsize=256) +def _fetch_page_index_bytes(path: str, offset: int, size: int) -> bytes: + """Fetch parquet page index bytes, cached per (path, offset, size). + + Multiple splits of the same file share an identical page index byte range, + so caching here avoids a redundant host-side file read per extra split. + """ + return plc.io.parquet_io_utils.fetch_page_index_to_host( + plc.io.SourceInfo([path]), + plc.io.text.ByteRangeInfo(offset=offset, size=size), ) +# Per-file cache: row-group row counts, keyed by file path. +# Avoids a redundant read_parquet_metadata call for every split of the same +# file in SplitScan.do_evaluate; populated on first encounter. +_row_group_num_rows_cache: dict[str, list[int]] = {} + +# Per-file cache: plain FileMetaData (footer only, no page index). +# Avoids a redundant read_parquet_footers call for every split of the same +# file; populated on first encounter. +_parquet_footer_cache: dict[str, Any] = {} + + def _read_with_hybrid_scan( schema: Schema, paths: list[str], @@ -209,13 +221,18 @@ def _read_with_hybrid_scan( row_group_indices: list[int], stream: Stream, file_metadata: plc.io.parquet_metadata.FileMetaData, + 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.""" 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=f"HybridScan: {paths[0]}"): + with nvtx_annotate_cudf_polars( + message=f"HybridScan: {paths[0]} [{split_index + 1}/{total_splits}]" + ): options = ( plc.io.parquet.ParquetReaderOptions.builder(plc.io.SourceInfo(paths)) .decimal_width(plc.TypeId.DECIMAL128) @@ -226,98 +243,42 @@ def _read_with_hybrid_scan( options.set_filter(plc_filter) reader = plc.io.experimental.HybridScanReader.from_parquet_metadata( - file_metadata, options - ) - - row_group_indices = reader.filter_row_groups_with_stats( - row_group_indices, options, stream=stream + file_metadata, + options, ) - if row_group_indices: - bloom_ranges, _dict_ranges = reader.secondary_filters_byte_ranges( - row_group_indices, options - ) - if bloom_ranges: - bloom_chunks = _fetch_byte_ranges(paths, 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, - ) - + # Pass the (cached) page index bytes so ``read`` does not re-read them per split. pi_range = reader.page_index_byte_range() - if pi_range.size > 0: - page_index_bytes = plc.io.parquet_io_utils.fetch_page_index_to_host( - plc.io.SourceInfo(paths), pi_range - ) - reader.setup_page_index(page_index_bytes) - row_mask = reader.build_row_mask_with_page_index_stats( - row_group_indices, options, stream=stream - ) - else: - row_mask = reader.build_all_true_row_mask(row_group_indices, stream=stream) - - filter_ranges = reader.filter_column_chunks_byte_ranges( - row_group_indices, options - ) - filter_chunks = _fetch_byte_ranges(paths, filter_ranges, stream) - filter_tbl_w_meta = reader.materialize_filter_columns( - row_group_indices, - filter_chunks, - row_mask, - plc.io.experimental.UseDataPageMask.YES, - options, - stream=stream, + page_index_bytes = ( + _fetch_page_index_bytes(paths[0], pi_range.offset, pi_range.size) + if pi_range.size > 0 + else None ) - payload_ranges = reader.payload_column_chunks_byte_ranges( - row_group_indices, options - ) - # PERFORMANCE: payload_column_chunks_byte_ranges does not need the row mask, - # so with async prefetch this fetch could be submitted concurrently with - # materialize_filter_columns to overlap I/O with GPU decode. - payload_chunks = _fetch_byte_ranges(paths, payload_ranges, stream) - payload_tbl_w_meta = reader.materialize_payload_columns( + # One fused C++ call performs pruning, row-mask construction, and the two-pass + # read, crossing into C++ (and the GIL) once per split rather than once per step. + tbl_w_meta = reader.read( + plc.io.SourceInfo(paths), row_group_indices, - payload_chunks, - row_mask, - plc.io.experimental.UseDataPageMask.YES, options, + page_index_bytes=page_index_bytes, + use_stats_filter=stats_pruning, + use_dictionary_filter=stats_pruning, + use_bloom_filter=stats_pruning, + prune_filter_column_pages=True, + prune_payload_column_pages=True, 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], + names = tbl_w_meta.column_names(include_children=False) + df = DataFrame.from_table( + tbl_w_meta.tbl, + names, + [schema[n] for n in names], stream=stream, ) stream.synchronize() - return DataFrame( - [*filter_df.columns, *payload_df.columns], stream=stream - ).select(list(schema.keys())) + return df.select(list(schema.keys())) class SplitScan(IR): @@ -441,12 +402,15 @@ def do_evaluate( # - We can use all this information to calculate the # "skip_rows" and "n_rows" options to use locally. - row_group_num_rows = [ - rg["num_rows"] - for rg in plc.io.parquet_metadata.read_parquet_metadata( - plc.io.SourceInfo(paths) - ).rowgroup_metadata() - ] + row_group_num_rows = _row_group_num_rows_cache.get(paths[0]) + if row_group_num_rows is None: + row_group_num_rows = [ + rg["num_rows"] + for rg in plc.io.parquet_metadata.read_parquet_metadata( + plc.io.SourceInfo(paths) + ).rowgroup_metadata() + ] + _row_group_num_rows_cache[paths[0]] = row_group_num_rows total_row_groups = len(row_group_num_rows) if total_splits <= total_row_groups: @@ -480,9 +444,14 @@ def do_evaluate( if split_index == total_splits - 1 else skip_rgs + rg_stride ) - [file_metadata] = plc.io.parquet_metadata.read_parquet_footers( - plc.io.SourceInfo(paths) - ) + # Reuse the cached plain footer, reading from disk only on + # the first split that encounters this file. + file_metadata = _parquet_footer_cache.get(paths[0]) + if file_metadata is None: + [file_metadata] = plc.io.parquet_metadata.read_parquet_footers( + plc.io.SourceInfo(paths) + ) + _parquet_footer_cache[paths[0]] = file_metadata return _read_with_hybrid_scan( schema, paths, @@ -491,6 +460,9 @@ def do_evaluate( list(range(skip_rgs, end_rg)), stream, file_metadata, + split_index=split_index, + total_splits=total_splits, + stats_pruning=parquet_options.hybrid_scan_stats_pruning, ) else: diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index c41e3fc3c871..4f13cf5372a0 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -210,6 +210,15 @@ class ParquetOptions: Whether to use the two-pass ``HybridScanReader`` for ``SplitScan`` tasks when a predicate can be pushed down to a parquet filter. Default is False. + hybrid_scan_stats_pruning + Whether to apply row-group stats and bloom-filter pruning before the + first pass of a hybrid scan. When ``True`` (default), row groups are + filtered via ``filter_row_groups_with_stats`` and + ``filter_row_groups_with_bloom_filters`` before any data is read. + Set to ``False`` to skip all pre-first-pass pruning and read every + row group assigned to this split, which is useful for benchmarking + the two-pass read overhead in isolation. + Only has effect when ``use_hybrid_scan`` is ``True``. """ _env_prefix = "CUDF_POLARS__PARQUET_OPTIONS" @@ -258,6 +267,13 @@ class ParquetOptions: default=False, ) ) + hybrid_scan_stats_pruning: bool = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__HYBRID_SCAN_STATS_PRUNING", + _bool_converter, + default=True, + ) + ) def __post_init__(self) -> None: # noqa: D105 if not isinstance(self.chunked, bool): @@ -276,6 +292,8 @@ def __post_init__(self) -> None: # noqa: D105 raise TypeError("use_rapidsmpf_native must be a bool") 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") def default_target_partition_size(min_device_size: int | None) -> int: diff --git a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx index 943f62f727b4..0082295c0375 100644 --- a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx +++ b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx @@ -1,8 +1,10 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 +from cython.operator cimport dereference from libc.stdint cimport uint8_t, uintptr_t from libc.stddef cimport size_t +from libcpp cimport bool from libcpp.memory cimport make_unique, unique_ptr from libcpp.pair cimport pair from libcpp.utility cimport move @@ -14,18 +16,22 @@ from rmm.pylibrmm.stream cimport Stream from pylibcudf.column cimport Column from pylibcudf.io.parquet cimport ParquetReaderOptions from pylibcudf.io.parquet_metadata cimport FileMetaData as c_FileMetaData +from pylibcudf.libcudf.io.parquet_schema cimport FileMetaData as cpp_FileMetaData from pylibcudf.io.text cimport ByteRangeInfo from pylibcudf.io.types cimport TableWithMetadata from pylibcudf.libcudf.column.column cimport column from pylibcudf.libcudf.column.column_view cimport column_view, mutable_column_view +from pylibcudf.libcudf.io.datasource cimport datasource, make_datasources from pylibcudf.libcudf.io.hybrid_scan cimport ( const_device_span_const_uint8_t, const_size_type, const_uint8_t, + hybrid_scan_read_options as cpp_hybrid_scan_read_options, hybrid_scan_reader as cpp_hybrid_scan_reader, use_data_page_mask as cpp_use_data_page_mask, ) from pylibcudf.libcudf.io.text cimport byte_range_info +from pylibcudf.io.types cimport SourceInfo from pylibcudf.libcudf.io.types cimport table_with_metadata from pylibcudf.libcudf.types cimport size_type from pylibcudf.libcudf.utilities.span cimport device_span, host_span @@ -81,10 +87,11 @@ cdef class HybridScanReader: """ def __init__(self, const uint8_t[::1] footer_bytes, ParquetReaderOptions options): - self.c_obj = make_unique[cpp_hybrid_scan_reader]( - host_span[const_uint8_t](&footer_bytes[0], len(footer_bytes)), - options.c_obj - ) + with nogil: + self.c_obj = make_unique[cpp_hybrid_scan_reader]( + host_span[const_uint8_t](&footer_bytes[0], len(footer_bytes)), + options.c_obj + ) @staticmethod def from_parquet_metadata(c_FileMetaData metadata, ParquetReaderOptions options): @@ -102,10 +109,11 @@ cdef class HybridScanReader: HybridScanReader """ cdef HybridScanReader reader = HybridScanReader.__new__(HybridScanReader) - reader.c_obj = make_unique[cpp_hybrid_scan_reader]( - metadata.c_obj, - options.c_obj - ) + with nogil: + reader.c_obj = make_unique[cpp_hybrid_scan_reader]( + metadata.c_obj, + options.c_obj + ) return reader def parquet_metadata(self): @@ -116,7 +124,10 @@ cdef class HybridScanReader: FileMetaData Parquet file footer metadata """ - return c_FileMetaData.from_cpp(self.c_obj.get()[0].parquet_metadata()) + cdef cpp_FileMetaData c_result + with nogil: + c_result = self.c_obj.get()[0].parquet_metadata() + return c_FileMetaData.from_cpp(c_result) def page_index_byte_range(self): """Get the byte range of the page index. @@ -126,7 +137,9 @@ cdef class HybridScanReader: ByteRangeInfo Byte range of the page index """ - cdef byte_range_info info = self.c_obj.get()[0].page_index_byte_range() + cdef byte_range_info info + with nogil: + info = self.c_obj.get()[0].page_index_byte_range() return ByteRangeInfo(info.offset(), info.size()) def setup_page_index(self, const uint8_t[::1] page_index_bytes): @@ -155,9 +168,9 @@ cdef class HybridScanReader: list[int] List of row group indices """ - cdef vector[size_type] row_groups = self.c_obj.get()[0].all_row_groups( - options.c_obj - ) + cdef vector[size_type] row_groups + with nogil: + row_groups = self.c_obj.get()[0].all_row_groups(options.c_obj) return list(row_groups) def total_rows_in_row_groups(self, list row_group_indices): @@ -174,9 +187,12 @@ cdef class HybridScanReader: Total number of top-level rows """ cdef vector[size_type] indices_vec = row_group_indices - return self.c_obj.get()[0].total_rows_in_row_groups( - host_span[const_size_type](indices_vec.data(), indices_vec.size()) - ) + cdef size_type result + with nogil: + result = self.c_obj.get()[0].total_rows_in_row_groups( + host_span[const_size_type](indices_vec.data(), indices_vec.size()) + ) + return result def reset_column_selection(self): """Reset the column selection state. @@ -184,7 +200,98 @@ cdef class HybridScanReader: Resets the internal column selection state forcing re-selection of columns in subsequent filter and read operations """ - self.c_obj.get()[0].reset_column_selection() + with nogil: + self.c_obj.get()[0].reset_column_selection() + + def read( + self, + SourceInfo source_info, + list row_group_indices, + ParquetReaderOptions options, + const uint8_t[::1] page_index_bytes=None, + bool use_stats_filter=True, + bool use_dictionary_filter=True, + bool use_bloom_filter=True, + bool prune_filter_column_pages=False, + bool prune_payload_column_pages=True, + object stream=None, + DeviceMemoryResource mr=None, + ): + """Read the Parquet file in a single fused hybrid scan operation. + + Performs row group pruning, row mask construction, and the two-pass + filter/payload materialization in one C++ call, so the entire per-chunk + read crosses into C++ (and acquires the GIL) once rather than once per + step. + + Parameters + ---------- + source_info : SourceInfo + Source describing the single Parquet file backing this reader. + row_group_indices : list[int] + Candidate row group indices to read, before pruning. + options : ParquetReaderOptions + Parquet reader options, including the optional filter expression and + column selection. + page_index_bytes : memoryview, optional + Host bytes of the page index. If ``None`` and the file contains a + page index, the bytes are fetched from the source. Pass cached bytes + to avoid a redundant read across reads of the same file. + use_stats_filter : bool, default True + Prune row groups using column chunk statistics. + use_dictionary_filter : bool, default True + Prune row groups using column chunk dictionary pages. + use_bloom_filter : bool, default True + Prune row groups using column chunk bloom filters. + prune_filter_column_pages : bool, default False + Whether to build and use a data page mask to prune filter column pages. + prune_payload_column_pages : bool, default True + Whether to build and use a data page mask to prune payload column pages. + stream : Stream, optional + CUDA stream. + mr : DeviceMemoryResource, optional + Device memory resource. + + Returns + ------- + TableWithMetadata + Materialized columns and metadata, in projection order. + """ + cdef Stream _stream = _get_stream(stream) + mr = _get_memory_resource(mr) + + cdef vector[unique_ptr[datasource]] sources = make_datasources(source_info.c_obj) + cdef vector[size_type] indices_vec = row_group_indices + + cdef cpp_hybrid_scan_read_options read_options + read_options.use_stats_filter = use_stats_filter + read_options.use_dictionary_filter = use_dictionary_filter + read_options.use_bloom_filter = use_bloom_filter + read_options.prune_filter_column_pages = ( + cpp_use_data_page_mask.YES if prune_filter_column_pages + else cpp_use_data_page_mask.NO + ) + read_options.prune_payload_column_pages = ( + cpp_use_data_page_mask.YES if prune_payload_column_pages + else cpp_use_data_page_mask.NO + ) + if page_index_bytes is not None and page_index_bytes.shape[0] > 0: + read_options.page_index_bytes = host_span[const_uint8_t]( + &page_index_bytes[0], page_index_bytes.shape[0] + ) + + cdef datasource* source_ptr = sources[0].get() + cdef table_with_metadata c_result + with nogil: + c_result = move(self.c_obj.get()[0].read( + dereference(source_ptr), + host_span[const_size_type](indices_vec.data(), indices_vec.size()), + options.c_obj, + read_options, + _stream.view().value(), + mr.get_mr(), + )) + return TableWithMetadata.from_libcudf(c_result, _stream, mr) def filter_row_groups_with_stats( self, @@ -883,7 +990,10 @@ cdef class HybridScanReader: bool True if there is data left to read """ - return self.c_obj.get()[0].has_next_table_chunk() + cdef bool result + with nogil: + result = self.c_obj.get()[0].has_next_table_chunk() + return result UseDataPageMask.__str__ = UseDataPageMask.__repr__ diff --git a/python/pylibcudf/pylibcudf/libcudf/io/datasource.pxd b/python/pylibcudf/pylibcudf/libcudf/io/datasource.pxd index 4c582e57e7a8..d3d8ecc227b5 100644 --- a/python/pylibcudf/pylibcudf/libcudf/io/datasource.pxd +++ b/python/pylibcudf/pylibcudf/libcudf/io/datasource.pxd @@ -12,9 +12,9 @@ cdef extern from "cudf/io/datasource.hpp" \ namespace "cudf::io" nogil: cdef cppclass datasource: - cdef cppclass buffer: - size_t size() const nogil - const uint8_t* data() const nogil + cppclass buffer: + size_t size() const + const uint8_t* data() const cdef vector[unique_ptr[datasource]] make_datasources( source_info info diff --git a/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan.pxd b/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan.pxd index 5dc33a5be1d4..967697569c0e 100644 --- a/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan.pxd +++ b/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan.pxd @@ -9,6 +9,7 @@ from libcpp.vector cimport vector from pylibcudf.exception_handler cimport libcudf_exception_handler from pylibcudf.libcudf.column.column cimport column from pylibcudf.libcudf.column.column_view cimport column_view, mutable_column_view +from pylibcudf.libcudf.io.datasource cimport datasource from pylibcudf.libcudf.io.parquet cimport parquet_reader_options from pylibcudf.libcudf.io.parquet_schema cimport FileMetaData from pylibcudf.libcudf.io.text cimport byte_range_info @@ -29,12 +30,30 @@ cdef extern from "cudf/io/experimental/hybrid_scan.hpp" \ YES NO + cdef cppclass hybrid_scan_read_options: + hybrid_scan_read_options() except +libcudf_exception_handler + bool use_stats_filter + bool use_dictionary_filter + bool use_bloom_filter + use_data_page_mask prune_filter_column_pages + use_data_page_mask prune_payload_column_pages + host_span[const_uint8_t] page_index_bytes + cdef cppclass hybrid_scan_reader: hybrid_scan_reader( host_span[const_uint8_t] footer_bytes, const parquet_reader_options& options ) except +libcudf_exception_handler + table_with_metadata read( + datasource& source, + host_span[const_size_type] row_group_indices, + const parquet_reader_options& options, + const hybrid_scan_read_options& read_options, + cudaStream_t stream, + device_async_resource_ref mr + ) except +libcudf_exception_handler + hybrid_scan_reader( const FileMetaData& parquet_metadata, const parquet_reader_options& options From 9b5781664689628ad9e4efa3163f42fabafa81c7 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Fri, 12 Jun 2026 23:00:58 +0000 Subject: [PATCH 16/20] remove page pruning and leave TODO --- .../cudf/io/experimental/hybrid_scan.hpp | 63 ----- .../io/parquet/experimental/hybrid_scan.cpp | 221 ------------------ .../io/experimental/hybrid_scan_test.cpp | 77 ------ .../cudf_polars/cudf_polars/streaming/io.py | 119 +++++++--- .../pylibcudf/io/experimental/hybrid_scan.pyx | 94 -------- .../pylibcudf/libcudf/io/hybrid_scan.pxd | 19 -- 6 files changed, 85 insertions(+), 508 deletions(-) diff --git a/cpp/include/cudf/io/experimental/hybrid_scan.hpp b/cpp/include/cudf/io/experimental/hybrid_scan.hpp index bdf1383bbe4c..980ab9644d3b 100644 --- a/cpp/include/cudf/io/experimental/hybrid_scan.hpp +++ b/cpp/include/cudf/io/experimental/hybrid_scan.hpp @@ -5,7 +5,6 @@ #pragma once -#include #include #include #include @@ -48,27 +47,6 @@ enum class use_data_page_mask : bool { NO = false ///< Do not compute or use a data page mask }; -/** - * @brief Options controlling a fused hybrid scan read via `hybrid_scan_reader::read` - * - * Each row-group pruning stage is optional and is additionally skipped at runtime when the - * parquet file does not contain the corresponding structure (e.g. no bloom filters) or when a - * prior stage has already pruned all row groups. - */ -struct hybrid_scan_read_options { - bool use_stats_filter = true; ///< Prune row groups using column chunk statistics - bool use_dictionary_filter = true; ///< Prune row groups using column chunk dictionary pages - bool use_bloom_filter = true; ///< Prune row groups using column chunk bloom filters - /// Whether to build and use a data page mask to prune filter column pages - use_data_page_mask prune_filter_column_pages = use_data_page_mask::NO; - /// Whether to build and use a data page mask to prune payload column pages - use_data_page_mask prune_payload_column_pages = use_data_page_mask::YES; - /// Host span of page index bytes. If empty and the file contains a page index, the bytes are - /// fetched from the data source. Pass cached bytes to avoid a redundant read across reads of the - /// same file. - cudf::host_span page_index_bytes = {}; -}; - /** * @brief The experimental parquet reader class to optimally read parquet files subject to * highly selective filters, called a Hybrid Scan operation @@ -322,47 +300,6 @@ class hybrid_scan_reader { */ ~hybrid_scan_reader(); - /** - * @brief Reads the Parquet file in a single fused operation - * - * Performs the complete hybrid scan in one call: optional row group pruning (statistics, - * dictionary pages and bloom filters), row mask construction, and the two-pass filter then - * payload column materialization. When `options` carries no filter expression, all selected - * columns are read in a single pass instead. - * - * This is equivalent in result to constructing the reader and invoking the individual step - * methods in sequence (see the class documentation), but issues all byte range fetches and GPU - * work from within this call. Callers that drive the reader from a managed runtime therefore - * cross the language boundary once per read rather than once per step, which avoids repeated - * lock acquisition (e.g. the Python GIL) between steps. - * - * The page index required for the row mask is taken from `read_options.page_index_bytes` when - * provided, and otherwise fetched from `source`. The row mask is built from page-level - * statistics when a page index is available and all-true otherwise. - * - * Output columns are returned in the projection order given by `options.get_column_names()`, - * or in Parquet file schema order when no column selection is set, matching - * `cudf::io::read_parquet`. - * - * @note `read` reads from a single data source. The reader holds per-read state, so a single - * reader instance must not be driven concurrently from multiple threads. - * - * @param source Data source for the Parquet file backing this reader - * @param row_group_indices Candidate row group indices to read, before pruning - * @param options Parquet reader options, including the optional filter expression and column - * selection - * @param read_options Options controlling row group pruning, data page masking and the page index - * @param stream CUDA stream used for device memory operations and kernel launches - * @param mr Device memory resource used to allocate the device memory for the output table - * @return Table of materialized columns and metadata, in projection order - */ - [[nodiscard]] table_with_metadata read(cudf::io::datasource& source, - cudf::host_span row_group_indices, - parquet_reader_options const& options, - hybrid_scan_read_options const& read_options, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr) const; - /** * @brief Get the Parquet file footer metadata * diff --git a/cpp/src/io/parquet/experimental/hybrid_scan.cpp b/cpp/src/io/parquet/experimental/hybrid_scan.cpp index f20957f1c928..868813a1b4ed 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan.cpp @@ -7,117 +7,12 @@ #include #include -#include -#include #include -#include -#include - #include -#include -#include -#include -#include - namespace cudf::io::parquet::experimental { -namespace { - -/** - * @brief Collects the top-level column names of a Parquet file in schema order - * - * Walks the flattened Parquet schema tree depth-first and records the names of the immediate - * children of the root, skipping each child's descendants. This yields the column order that - * `cudf::io::read_parquet` produces when no column selection is set. - * - * @param metadata Parquet file footer metadata - * @return Top-level column names in file schema order - */ -[[nodiscard]] std::vector top_level_column_names(FileMetaData const& metadata) -{ - auto const& schema = metadata.schema; - if (schema.empty()) { return {}; } - - std::vector names; - names.reserve(schema.front().num_children); - - // Depth-first walk tracking unvisited children remaining at each open level. The root (index 0) - // is skipped; a column is top-level when the only open level is the root's. - std::vector remaining_children{schema.front().num_children}; - for (std::size_t i = 1; i < schema.size() and not remaining_children.empty(); ++i) { - auto const& element = schema[i]; - if (remaining_children.size() == 1) { names.push_back(element.name); } - - --remaining_children.back(); - if (element.num_children > 0) { - remaining_children.push_back(element.num_children); - } else { - while (not remaining_children.empty() and remaining_children.back() == 0) { - remaining_children.pop_back(); - } - } - } - - return names; -} - -/** - * @brief Reassembles the filter and payload tables into a single table in `output_order` - * - * The two passes of a hybrid scan materialize disjoint sets of columns (filter columns and - * payload columns). This combines them into one table whose columns follow `output_order`, - * matching the projection order of `cudf::io::read_parquet`. - * - * @param filter Materialized filter columns and metadata - * @param payload Materialized payload columns and metadata - * @param output_order Desired output column names in order - * @return Combined table and metadata in `output_order` - */ -[[nodiscard]] table_with_metadata assemble_output(table_with_metadata filter, - table_with_metadata payload, - std::vector const& output_order) -{ - enum class source : bool { filter, payload }; - std::unordered_map> location; - location.reserve(filter.metadata.schema_info.size() + payload.metadata.schema_info.size()); - for (std::size_t i = 0; i < filter.metadata.schema_info.size(); ++i) { - location.emplace(filter.metadata.schema_info[i].name, std::pair{source::filter, i}); - } - for (std::size_t i = 0; i < payload.metadata.schema_info.size(); ++i) { - location.emplace(payload.metadata.schema_info[i].name, std::pair{source::payload, i}); - } - - auto filter_columns = filter.tbl->release(); - auto payload_columns = payload.tbl->release(); - - std::vector> output_columns; - output_columns.reserve(output_order.size()); - table_metadata out_metadata; - out_metadata.schema_info.reserve(output_order.size()); - out_metadata.num_rows_per_source = filter.metadata.num_rows_per_source; - - for (auto const& name : output_order) { - auto const it = location.find(name); - CUDF_EXPECTS(it != location.end(), - "Projected column not found in materialized hybrid scan output: " + name); - auto const [tbl, pos] = it->second; - if (tbl == source::filter) { - output_columns.push_back(std::move(filter_columns[pos])); - out_metadata.schema_info.push_back(std::move(filter.metadata.schema_info[pos])); - } else { - output_columns.push_back(std::move(payload_columns[pos])); - out_metadata.schema_info.push_back(std::move(payload.metadata.schema_info[pos])); - } - } - - return table_with_metadata{std::make_unique(std::move(output_columns)), - std::move(out_metadata)}; -} - -} // namespace - hybrid_scan_reader::hybrid_scan_reader(cudf::host_span footer_bytes, parquet_reader_options const& options) : _impl{std::make_unique( @@ -134,122 +29,6 @@ hybrid_scan_reader::hybrid_scan_reader(FileMetaData const& parquet_metadata, hybrid_scan_reader::~hybrid_scan_reader() = default; -table_with_metadata hybrid_scan_reader::read(cudf::io::datasource& source, - cudf::host_span row_group_indices, - parquet_reader_options const& options, - hybrid_scan_read_options const& read_options, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr) const -{ - CUDF_FUNC_RANGE(); - - // Transient byte-range fetches use the current device resource; output tables use `mr`. - auto temp_mr = cudf::get_current_device_resource_ref(); - - auto const output_order = options.get_column_names().has_value() - ? options.get_column_names().value() - : top_level_column_names(_impl->parquet_metadatas().front()); - - std::vector current_row_groups(row_group_indices.begin(), row_group_indices.end()); - - // No filter: read all selected columns in a single pass. - if (not options.get_filter().has_value()) { - auto const byte_ranges = all_column_chunks_byte_ranges(current_row_groups, options); - auto [buffers, data, tasks] = - parquet::fetch_byte_ranges_to_device_async(source, byte_ranges, stream, temp_mr); - tasks.get(); - return materialize_all_columns(current_row_groups, data, options, stream, mr); - } - - // Use caller-provided page index bytes when available, otherwise fetch them from the source. - auto const page_index_range = page_index_byte_range(); - auto const has_page_index = not page_index_range.is_empty(); - if (has_page_index) { - if (not read_options.page_index_bytes.empty()) { - setup_page_index(read_options.page_index_bytes); - } else { - auto const page_index_buffer = parquet::fetch_page_index_to_host(source, page_index_range); - setup_page_index( - cudf::host_span{page_index_buffer->data(), page_index_buffer->size()}); - } - } - - if (read_options.use_stats_filter and not current_row_groups.empty()) { - current_row_groups = filter_row_groups_with_stats(current_row_groups, options, stream); - } - - if ((read_options.use_dictionary_filter or read_options.use_bloom_filter) and - not current_row_groups.empty()) { - auto const [bloom_filter_ranges, dictionary_page_ranges] = - secondary_filters_byte_ranges(current_row_groups, options); - - if (read_options.use_dictionary_filter and not dictionary_page_ranges.empty()) { - auto [buffers, data, tasks] = - parquet::fetch_byte_ranges_to_device_async(source, dictionary_page_ranges, stream, temp_mr); - tasks.get(); - current_row_groups = - filter_row_groups_with_dictionary_pages(data, current_row_groups, options, stream); - } - - if (read_options.use_bloom_filter and not bloom_filter_ranges.empty() and - not current_row_groups.empty()) { - // Bloom filter data buffers must be allocated on 32-byte aligned addresses. - auto aligned_mr = rmm::mr::aligned_resource_adaptor{temp_mr, rmm::CUDA_ALLOCATION_ALIGNMENT}; - auto [buffers, data, tasks] = - parquet::fetch_byte_ranges_to_device_async(source, bloom_filter_ranges, stream, aligned_mr); - tasks.get(); - current_row_groups = - filter_row_groups_with_bloom_filters(data, current_row_groups, options, stream); - } - } - - // All row groups pruned: return a correctly typed, zero-row table in projection order. The - // materialization path requires a non-empty row mask, so defer to the main reader here. - if (current_row_groups.empty()) { - auto empty_options = options; - empty_options.set_num_rows(0); - return cudf::io::read_parquet(empty_options, stream, mr); - } - - // Use page-level statistics for the row mask when a page index is available; otherwise all-true. - auto row_mask = has_page_index - ? build_row_mask_with_page_index_stats(current_row_groups, options, stream, mr) - : build_all_true_row_mask(current_row_groups, stream, mr); - - // Filter pass: materialize filter columns and narrow the row mask to surviving rows. - auto row_mask_view = row_mask->mutable_view(); - auto filter_table = [&] { - auto const byte_ranges = filter_column_chunks_byte_ranges(current_row_groups, options); - auto [buffers, data, tasks] = - parquet::fetch_byte_ranges_to_device_async(source, byte_ranges, stream, temp_mr); - tasks.get(); - return materialize_filter_columns(current_row_groups, - data, - row_mask_view, - read_options.prune_filter_column_pages, - options, - stream, - mr); - }(); - - // Payload pass: materialize payload columns under the surviving row mask. - auto payload_table = [&] { - auto const byte_ranges = payload_column_chunks_byte_ranges(current_row_groups, options); - auto [buffers, data, tasks] = - parquet::fetch_byte_ranges_to_device_async(source, byte_ranges, stream, temp_mr); - tasks.get(); - return materialize_payload_columns(current_row_groups, - data, - row_mask->view(), - read_options.prune_payload_column_pages, - options, - stream, - mr); - }(); - - return assemble_output(std::move(filter_table), std::move(payload_table), output_order); -} - [[nodiscard]] text::byte_range_info hybrid_scan_reader::page_index_byte_range() const { return _impl->page_index_byte_ranges().front(); diff --git a/cpp/tests/io/experimental/hybrid_scan_test.cpp b/cpp/tests/io/experimental/hybrid_scan_test.cpp index b6b019082b05..05e7fba3d844 100644 --- a/cpp/tests/io/experimental/hybrid_scan_test.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_test.cpp @@ -28,8 +28,6 @@ #include -#include - namespace { /** @@ -925,81 +923,6 @@ TEST_F(HybridScanTest, StructChildFilterColumn) std::invalid_argument); } -namespace { - -// Drives the fused `hybrid_scan_reader::read` and compares against `cudf::io::read_parquet` for the -// same filter and column selection. `selected_columns` deliberately reorders columns and places the -// filter column away from the front, exercising the projection-order reassembly inside `read`. -void test_read_matches_read_parquet(std::vector const& parquet_buffer, - cudf::ast::operation const& filter_expression, - std::vector const& selected_columns) -{ - auto const stream = cudf::get_default_stream(); - auto const mr = cudf::get_current_device_resource_ref(); - - auto const make_options = [&] { - auto options = cudf::io::parquet_reader_options::builder(cudf::io::source_info( - cudf::host_span( - parquet_buffer.data(), - parquet_buffer.size()))) - .filter(filter_expression) - .build(); - options.set_column_names(selected_columns); - return options; - }; - - auto const expected = cudf::io::read_parquet(make_options(), stream, mr).tbl; - - auto const options = make_options(); - auto datasource = cudf::io::datasource::create(cudf::host_span( - reinterpret_cast(parquet_buffer.data()), parquet_buffer.size())); - auto const footer_buffer = cudf::io::parquet::fetch_footer_to_host(*datasource); - auto reader = - std::make_unique(*footer_buffer, options); - - auto const row_groups = reader->all_row_groups(options); - auto const read_options = - cudf::io::parquet::experimental::hybrid_scan_read_options{}; // page index fetched internally - auto const actual = reader->read(*datasource, row_groups, options, read_options, stream, mr); - - CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected->view(), actual.tbl->view()); -} - -} // namespace - -TEST_F(HybridScanTest, ReadMatchesReadParquetReorderedProjection) -{ - using T = int32_t; - auto constexpr num_concat = 2; // multiple row groups + pages so pruning engages - auto [written_table, parquet_buffer] = create_parquet_with_stats(); - - // Filter on the middle column; project columns in a different order with the filter column last. - auto literal_value = cudf::numeric_scalar(T{500}); - auto literal = cudf::ast::literal(literal_value); - auto col_ref_1 = cudf::ast::column_name_reference("col1"); - auto filter_expression = cudf::ast::operation(cudf::ast::ast_operator::LESS, col_ref_1, literal); - - test_read_matches_read_parquet( - parquet_buffer, filter_expression, std::vector{"col2", "col0", "col1"}); -} - -TEST_F(HybridScanTest, ReadMatchesReadParquetNoSurvivingRows) -{ - using T = int32_t; - auto constexpr num_concat = 2; - auto [written_table, parquet_buffer] = create_parquet_with_stats(); - - // A predicate no row can satisfy, exercising the empty-output path of `read`. - auto literal_value = cudf::numeric_scalar(std::numeric_limits::min()); - auto literal = cudf::ast::literal(literal_value); - auto col_ref_0 = cudf::ast::column_name_reference("col0"); - auto filter_expression = - cudf::ast::operation(cudf::ast::ast_operator::LESS, col_ref_0, literal); - - test_read_matches_read_parquet( - parquet_buffer, filter_expression, std::vector{"col2", "col0"}); -} - TEST_F(HybridScanTest, ChunkedReadRowMaskPerPass) { using T = uint32_t; diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index cace8524cfef..0f3683c22d1e 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -189,18 +189,6 @@ def expand_scan_for_rank( return scans -@functools.lru_cache(maxsize=256) -def _fetch_page_index_bytes(path: str, offset: int, size: int) -> bytes: - """Fetch parquet page index bytes, cached per (path, offset, size). - - Multiple splits of the same file share an identical page index byte range, - so caching here avoids a redundant host-side file read per extra split. - """ - return plc.io.parquet_io_utils.fetch_page_index_to_host( - plc.io.SourceInfo([path]), - plc.io.text.ByteRangeInfo(offset=offset, size=size), - ) - # Per-file cache: row-group row counts, keyed by file path. # Avoids a redundant read_parquet_metadata call for every split of the same @@ -213,6 +201,18 @@ def _fetch_page_index_bytes(path: str, offset: int, size: int) -> bytes: _parquet_footer_cache: dict[str, Any] = {} +def _fetch_byte_ranges( + paths: list[str], + byte_ranges: list[plc.io.text.ByteRangeInfo], + stream: Stream, +) -> list[plc.gpumemoryview]: + # TODO: Accept a pinned-host Datasource pre-fetched by the caller so the + # storage I/O overlaps with GPU work for better pipelining. + return plc.io.parquet_io_utils.fetch_byte_ranges_to_device( + plc.io.SourceInfo(paths), byte_ranges, stream=stream + ) + + def _read_with_hybrid_scan( schema: Schema, paths: list[str], @@ -247,38 +247,89 @@ def _read_with_hybrid_scan( options, ) - # Pass the (cached) page index bytes so ``read`` does not re-read them per split. - pi_range = reader.page_index_byte_range() - page_index_bytes = ( - _fetch_page_index_bytes(paths[0], pi_range.offset, pi_range.size) - if pi_range.size > 0 - else None + 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(paths, 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_ranges = reader.filter_column_chunks_byte_ranges( + row_group_indices, options + ) + filter_chunks = _fetch_byte_ranges(paths, filter_ranges, stream) + filter_tbl_w_meta = reader.materialize_filter_columns( + row_group_indices, + filter_chunks, + row_mask, + plc.io.experimental.UseDataPageMask.YES, + options, + stream=stream, ) - # One fused C++ call performs pruning, row-mask construction, and the two-pass - # read, crossing into C++ (and the GIL) once per split rather than once per step. - tbl_w_meta = reader.read( - plc.io.SourceInfo(paths), + payload_ranges = reader.payload_column_chunks_byte_ranges( + row_group_indices, options + ) + payload_chunks = _fetch_byte_ranges(paths, payload_ranges, stream) + payload_tbl_w_meta = reader.materialize_payload_columns( row_group_indices, + payload_chunks, + row_mask, + plc.io.experimental.UseDataPageMask.YES, options, - page_index_bytes=page_index_bytes, - use_stats_filter=stats_pruning, - use_dictionary_filter=stats_pruning, - use_bloom_filter=stats_pruning, - prune_filter_column_pages=True, - prune_payload_column_pages=True, stream=stream, ) - names = tbl_w_meta.column_names(include_children=False) - df = DataFrame.from_table( - tbl_w_meta.tbl, - names, - [schema[n] for n in names], + 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, ) + # Ensure the decode kernels are finished before + # filter_chunks and payload_chunks go out of scope stream.synchronize() - return df.select(list(schema.keys())) + return DataFrame( + [*filter_df.columns, *payload_df.columns], stream=stream + ).select(list(schema.keys())) class SplitScan(IR): diff --git a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx index 0082295c0375..362f382f0abb 100644 --- a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx +++ b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx @@ -1,7 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 -from cython.operator cimport dereference from libc.stdint cimport uint8_t, uintptr_t from libc.stddef cimport size_t from libcpp cimport bool @@ -21,17 +20,14 @@ from pylibcudf.io.text cimport ByteRangeInfo from pylibcudf.io.types cimport TableWithMetadata from pylibcudf.libcudf.column.column cimport column from pylibcudf.libcudf.column.column_view cimport column_view, mutable_column_view -from pylibcudf.libcudf.io.datasource cimport datasource, make_datasources from pylibcudf.libcudf.io.hybrid_scan cimport ( const_device_span_const_uint8_t, const_size_type, const_uint8_t, - hybrid_scan_read_options as cpp_hybrid_scan_read_options, hybrid_scan_reader as cpp_hybrid_scan_reader, use_data_page_mask as cpp_use_data_page_mask, ) from pylibcudf.libcudf.io.text cimport byte_range_info -from pylibcudf.io.types cimport SourceInfo from pylibcudf.libcudf.io.types cimport table_with_metadata from pylibcudf.libcudf.types cimport size_type from pylibcudf.libcudf.utilities.span cimport device_span, host_span @@ -203,96 +199,6 @@ cdef class HybridScanReader: with nogil: self.c_obj.get()[0].reset_column_selection() - def read( - self, - SourceInfo source_info, - list row_group_indices, - ParquetReaderOptions options, - const uint8_t[::1] page_index_bytes=None, - bool use_stats_filter=True, - bool use_dictionary_filter=True, - bool use_bloom_filter=True, - bool prune_filter_column_pages=False, - bool prune_payload_column_pages=True, - object stream=None, - DeviceMemoryResource mr=None, - ): - """Read the Parquet file in a single fused hybrid scan operation. - - Performs row group pruning, row mask construction, and the two-pass - filter/payload materialization in one C++ call, so the entire per-chunk - read crosses into C++ (and acquires the GIL) once rather than once per - step. - - Parameters - ---------- - source_info : SourceInfo - Source describing the single Parquet file backing this reader. - row_group_indices : list[int] - Candidate row group indices to read, before pruning. - options : ParquetReaderOptions - Parquet reader options, including the optional filter expression and - column selection. - page_index_bytes : memoryview, optional - Host bytes of the page index. If ``None`` and the file contains a - page index, the bytes are fetched from the source. Pass cached bytes - to avoid a redundant read across reads of the same file. - use_stats_filter : bool, default True - Prune row groups using column chunk statistics. - use_dictionary_filter : bool, default True - Prune row groups using column chunk dictionary pages. - use_bloom_filter : bool, default True - Prune row groups using column chunk bloom filters. - prune_filter_column_pages : bool, default False - Whether to build and use a data page mask to prune filter column pages. - prune_payload_column_pages : bool, default True - Whether to build and use a data page mask to prune payload column pages. - stream : Stream, optional - CUDA stream. - mr : DeviceMemoryResource, optional - Device memory resource. - - Returns - ------- - TableWithMetadata - Materialized columns and metadata, in projection order. - """ - cdef Stream _stream = _get_stream(stream) - mr = _get_memory_resource(mr) - - cdef vector[unique_ptr[datasource]] sources = make_datasources(source_info.c_obj) - cdef vector[size_type] indices_vec = row_group_indices - - cdef cpp_hybrid_scan_read_options read_options - read_options.use_stats_filter = use_stats_filter - read_options.use_dictionary_filter = use_dictionary_filter - read_options.use_bloom_filter = use_bloom_filter - read_options.prune_filter_column_pages = ( - cpp_use_data_page_mask.YES if prune_filter_column_pages - else cpp_use_data_page_mask.NO - ) - read_options.prune_payload_column_pages = ( - cpp_use_data_page_mask.YES if prune_payload_column_pages - else cpp_use_data_page_mask.NO - ) - if page_index_bytes is not None and page_index_bytes.shape[0] > 0: - read_options.page_index_bytes = host_span[const_uint8_t]( - &page_index_bytes[0], page_index_bytes.shape[0] - ) - - cdef datasource* source_ptr = sources[0].get() - cdef table_with_metadata c_result - with nogil: - c_result = move(self.c_obj.get()[0].read( - dereference(source_ptr), - host_span[const_size_type](indices_vec.data(), indices_vec.size()), - options.c_obj, - read_options, - _stream.view().value(), - mr.get_mr(), - )) - return TableWithMetadata.from_libcudf(c_result, _stream, mr) - def filter_row_groups_with_stats( self, list row_group_indices, diff --git a/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan.pxd b/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan.pxd index 967697569c0e..5dc33a5be1d4 100644 --- a/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan.pxd +++ b/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan.pxd @@ -9,7 +9,6 @@ from libcpp.vector cimport vector from pylibcudf.exception_handler cimport libcudf_exception_handler from pylibcudf.libcudf.column.column cimport column from pylibcudf.libcudf.column.column_view cimport column_view, mutable_column_view -from pylibcudf.libcudf.io.datasource cimport datasource from pylibcudf.libcudf.io.parquet cimport parquet_reader_options from pylibcudf.libcudf.io.parquet_schema cimport FileMetaData from pylibcudf.libcudf.io.text cimport byte_range_info @@ -30,30 +29,12 @@ cdef extern from "cudf/io/experimental/hybrid_scan.hpp" \ YES NO - cdef cppclass hybrid_scan_read_options: - hybrid_scan_read_options() except +libcudf_exception_handler - bool use_stats_filter - bool use_dictionary_filter - bool use_bloom_filter - use_data_page_mask prune_filter_column_pages - use_data_page_mask prune_payload_column_pages - host_span[const_uint8_t] page_index_bytes - cdef cppclass hybrid_scan_reader: hybrid_scan_reader( host_span[const_uint8_t] footer_bytes, const parquet_reader_options& options ) except +libcudf_exception_handler - table_with_metadata read( - datasource& source, - host_span[const_size_type] row_group_indices, - const parquet_reader_options& options, - const hybrid_scan_read_options& read_options, - cudaStream_t stream, - device_async_resource_ref mr - ) except +libcudf_exception_handler - hybrid_scan_reader( const FileMetaData& parquet_metadata, const parquet_reader_options& options From 8d92fb02f75bfaad0b2e93e690462eef08b6b726 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Sat, 13 Jun 2026 03:14:40 +0000 Subject: [PATCH 17/20] share file metadata without copying when contructing hybrid scan readers --- .../cudf/io/experimental/hybrid_scan.hpp | 84 ++++++++++++++++ .../io/parquet/experimental/hybrid_scan.cpp | 25 +++++ .../parquet/experimental/hybrid_scan_impl.cpp | 12 ++- .../parquet/experimental/hybrid_scan_impl.hpp | 11 +++ cpp/src/io/parquet/reader_impl.cpp | 4 +- cpp/src/io/parquet/reader_impl.hpp | 3 +- .../io/experimental/hybrid_scan_test.cpp | 46 +++++++++ .../cudf_polars/cudf_polars/streaming/io.py | 21 ++-- .../pylibcudf/io/experimental/__init__.py | 2 + .../pylibcudf/io/experimental/hybrid_scan.pxd | 5 + .../pylibcudf/io/experimental/hybrid_scan.pyi | 12 +++ .../pylibcudf/io/experimental/hybrid_scan.pyx | 95 ++++++++++++++++++- .../pylibcudf/libcudf/io/hybrid_scan.pxd | 15 +++ 13 files changed, 323 insertions(+), 12 deletions(-) diff --git a/cpp/include/cudf/io/experimental/hybrid_scan.hpp b/cpp/include/cudf/io/experimental/hybrid_scan.hpp index 980ab9644d3b..014a6458ce39 100644 --- a/cpp/include/cudf/io/experimental/hybrid_scan.hpp +++ b/cpp/include/cudf/io/experimental/hybrid_scan.hpp @@ -25,6 +25,11 @@ namespace cudf::io::parquet::experimental::detail { * Hybrid Scan operation. */ class hybrid_scan_reader_impl; + +/** + * @brief Internal parsed Parquet file metadata for the Hybrid Scan reader. + */ +class aggregate_reader_metadata; } // namespace cudf::io::parquet::experimental::detail //! Using `byte_range_info` from cudf::io::text @@ -47,6 +52,73 @@ enum class use_data_page_mask : bool { NO = false ///< Do not compute or use a data page mask }; +/** + * @brief Shareable, pre-parsed Parquet file metadata for the Hybrid Scan reader + * + * Parses the Parquet file metadata once so that multiple `hybrid_scan_reader` instances reading the + * same file can borrow it rather than each re-parsing and copying the (potentially large) row group + * metadata. The intended use is to read disjoint row-group ranges of a single file: construct one + * `hybrid_scan_metadata` per file and pass it to as many readers as there are ranges. + * + * @code{.cpp} + * // Parse the metadata once + * auto metadata = std::make_shared(*footer_buffer, + * options); + * // Construct lightweight readers that share it + * auto reader_a = std::make_unique(*metadata); + * auto reader_b = std::make_unique(*metadata); + * @endcode + * + * @note The metadata is immutable once constructed. Readers sharing one instance must read disjoint + * row-group ranges of the same single file; such reads do not mutate the shared metadata, so they + * may run concurrently. This handle does not support multi-source (multi-file) metadata. + */ +class hybrid_scan_metadata { + public: + /** + * @brief Parse and own Parquet file metadata from a span of footer bytes + * + * @param footer_bytes Host span of Parquet file footer bytes + * @param options Parquet reader options + */ + hybrid_scan_metadata(cudf::host_span footer_bytes, + parquet_reader_options const& options); + + /** + * @brief Own Parquet file metadata from a pre-populated `FileMetaData` + * + * @param parquet_metadata Pre-populated Parquet file metadata + * @param options Parquet reader options + */ + hybrid_scan_metadata(FileMetaData const& parquet_metadata, parquet_reader_options const& options); + + /** + * @brief Destructor for the shared Parquet metadata + */ + ~hybrid_scan_metadata(); + + hybrid_scan_metadata(hybrid_scan_metadata const&) = default; ///< Copy constructor + hybrid_scan_metadata(hybrid_scan_metadata&&) = default; ///< Move constructor + + /** + * @brief Copy assignment operator + * + * @return Reference to this object + */ + hybrid_scan_metadata& operator=(hybrid_scan_metadata const&) = default; + + /** + * @brief Move assignment operator + * + * @return Reference to this object + */ + hybrid_scan_metadata& operator=(hybrid_scan_metadata&&) = default; + + private: + std::shared_ptr _metadata; + friend class hybrid_scan_reader; +}; + /** * @brief The experimental parquet reader class to optimally read parquet files subject to * highly selective filters, called a Hybrid Scan operation @@ -295,6 +367,18 @@ class hybrid_scan_reader { explicit hybrid_scan_reader(FileMetaData const& parquet_metadata, parquet_reader_options const& options); + /** + * @brief Constructor that borrows shared, pre-parsed Parquet file metadata + * + * Constructs a reader that shares `metadata` instead of parsing and copying the file metadata + * again. Use this to read disjoint row-group ranges of a single file without paying the metadata + * copy per reader. The reader options that govern reading (column selection, filter, ...) are + * supplied per call to the individual read methods. + * + * @param metadata Shared, pre-parsed Parquet file metadata + */ + explicit hybrid_scan_reader(hybrid_scan_metadata const& metadata); + /** * @brief Destructor for the experimental parquet reader class */ diff --git a/cpp/src/io/parquet/experimental/hybrid_scan.cpp b/cpp/src/io/parquet/experimental/hybrid_scan.cpp index 868813a1b4ed..24e0504d7e2c 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan.cpp @@ -13,6 +13,26 @@ namespace cudf::io::parquet::experimental { +hybrid_scan_metadata::hybrid_scan_metadata(cudf::host_span footer_bytes, + parquet_reader_options const& options) + : _metadata{std::make_shared( + std::vector>{footer_bytes}, + options.is_enabled_use_arrow_schema(), + options.get_column_names().has_value() and options.is_enabled_allow_mismatched_pq_schemas())} +{ +} + +hybrid_scan_metadata::hybrid_scan_metadata(FileMetaData const& parquet_metadata, + parquet_reader_options const& options) + : _metadata{std::make_shared( + std::vector{parquet_metadata}, + options.is_enabled_use_arrow_schema(), + options.get_column_names().has_value() and options.is_enabled_allow_mismatched_pq_schemas())} +{ +} + +hybrid_scan_metadata::~hybrid_scan_metadata() = default; + hybrid_scan_reader::hybrid_scan_reader(cudf::host_span footer_bytes, parquet_reader_options const& options) : _impl{std::make_unique( @@ -27,6 +47,11 @@ hybrid_scan_reader::hybrid_scan_reader(FileMetaData const& parquet_metadata, { } +hybrid_scan_reader::hybrid_scan_reader(hybrid_scan_metadata const& metadata) + : _impl{std::make_unique(metadata._metadata)} +{ +} + hybrid_scan_reader::~hybrid_scan_reader() = default; [[nodiscard]] text::byte_range_info hybrid_scan_reader::page_index_byte_range() const diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp index 0f0fd7de6d96..816e10516a0a 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp @@ -68,7 +68,7 @@ hybrid_scan_reader_impl::hybrid_scan_reader_impl( cudf::host_span const> footer_bytes, parquet_reader_options const& options) { - _metadata = std::make_unique( + _metadata = std::make_shared( footer_bytes, options.is_enabled_use_arrow_schema(), options.get_column_names().has_value() and options.is_enabled_allow_mismatched_pq_schemas()); @@ -79,13 +79,21 @@ hybrid_scan_reader_impl::hybrid_scan_reader_impl( hybrid_scan_reader_impl::hybrid_scan_reader_impl( cudf::host_span parquet_metadatas, parquet_reader_options const& options) { - _metadata = std::make_unique( + _metadata = std::make_shared( parquet_metadatas, options.is_enabled_use_arrow_schema(), options.get_column_names().has_value() and options.is_enabled_allow_mismatched_pq_schemas()); _extended_metadata = static_cast(_metadata.get()); } +hybrid_scan_reader_impl::hybrid_scan_reader_impl( + std::shared_ptr metadata) +{ + CUDF_EXPECTS(metadata != nullptr, "Shared parquet metadata must not be null"); + _metadata = std::move(metadata); + _extended_metadata = static_cast(_metadata.get()); +} + std::vector hybrid_scan_reader_impl::parquet_metadatas() const { return _extended_metadata->parquet_metadatas(); diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp b/cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp index d11ae1e8ddb9..9b88be20e0b7 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp @@ -57,6 +57,17 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl { explicit hybrid_scan_reader_impl(cudf::host_span parquet_metadatas, parquet_reader_options const& options); + /** + * @brief Constructor that shares pre-parsed Parquet metadata + * + * Borrows an already-constructed `aggregate_reader_metadata` instead of parsing and copying the + * file metadata again. Multiple single-file readers can share one metadata object, avoiding a + * per-reader copy of the (potentially large) row group metadata. + * + * @param metadata Shared, pre-parsed Parquet file metadata. Must not be null. + */ + explicit hybrid_scan_reader_impl(std::shared_ptr metadata); + /** * @copydoc cudf::io::experimental::hybrid_scan_multifile::parquet_metadatas */ diff --git a/cpp/src/io/parquet/reader_impl.cpp b/cpp/src/io/parquet/reader_impl.cpp index eb00fafa6896..a31b13514359 100644 --- a/cpp/src/io/parquet/reader_impl.cpp +++ b/cpp/src/io/parquet/reader_impl.cpp @@ -520,12 +520,12 @@ reader_impl::reader_impl(std::size_t chunk_read_limit, // Open and parse the source dataset metadata CUDF_EXPECTS(file_metadatas.empty() or file_metadatas.size() == _sources.size(), "Encountered a mismatch in the number of provided data sources and metadatas"); - _metadata = file_metadatas.empty() ? std::make_unique( + _metadata = file_metadatas.empty() ? std::make_shared( _sources, options.is_enabled_use_arrow_schema(), options.get_column_names().has_value() and options.is_enabled_allow_mismatched_pq_schemas()) - : std::make_unique( + : std::make_shared( std::forward>(file_metadatas), options.is_enabled_use_arrow_schema(), options.get_column_names().has_value() and diff --git a/cpp/src/io/parquet/reader_impl.hpp b/cpp/src/io/parquet/reader_impl.hpp index 4413185e926e..f7c794dc1273 100644 --- a/cpp/src/io/parquet/reader_impl.hpp +++ b/cpp/src/io/parquet/reader_impl.hpp @@ -464,7 +464,8 @@ class reader_impl { named_to_reference_converter _expr_conv{std::nullopt, table_metadata{}, true}; std::vector> _sources; - std::unique_ptr _metadata; + // shared so experimental hybrid scan readers can share one copy across single-file readers + std::shared_ptr _metadata; // Number of sources size_t _num_sources{0}; diff --git a/cpp/tests/io/experimental/hybrid_scan_test.cpp b/cpp/tests/io/experimental/hybrid_scan_test.cpp index 05e7fba3d844..6e3c7c18aa8e 100644 --- a/cpp/tests/io/experimental/hybrid_scan_test.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_test.cpp @@ -923,6 +923,52 @@ TEST_F(HybridScanTest, StructChildFilterColumn) std::invalid_argument); } +TEST_F(HybridScanTest, SharedMetadataReaderMatchesReadParquet) +{ + using T = int32_t; + auto constexpr num_concat = 2; + auto [written_table, parquet_buffer] = create_parquet_with_stats(); + + auto const stream = cudf::get_default_stream(); + auto const mr = cudf::get_current_device_resource_ref(); + auto const options = cudf::io::parquet_reader_options::builder().build(); + + auto datasource = cudf::io::datasource::create(cudf::host_span( + reinterpret_cast(parquet_buffer.data()), parquet_buffer.size())); + auto const footer_buffer = cudf::io::parquet::fetch_footer_to_host(*datasource); + + // Parse the file metadata once and share it across independent readers. + auto const metadata = std::make_shared( + *footer_buffer, options); + + // Read all columns (single step) through a reader that borrows the shared metadata. + auto const read_all_columns = [&] { + auto const reader = + std::make_unique(*metadata); + auto const row_groups = reader->all_row_groups(options); + auto const chunk_ranges = reader->all_column_chunks_byte_ranges(row_groups, options); + auto [buffers, data, tasks] = + cudf::io::parquet::fetch_byte_ranges_to_device_async(*datasource, chunk_ranges, stream, mr); + tasks.get(); + return reader->materialize_all_columns(row_groups, data, options, stream, mr).tbl; + }; + + // Two readers sharing one metadata instance each produce the same table as the main reader. + auto const table_a = read_all_columns(); + auto const table_b = read_all_columns(); + + auto const expected = + cudf::io::read_parquet( + cudf::io::parquet_reader_options::builder( + cudf::io::source_info(cudf::host_span(parquet_buffer.data(), parquet_buffer.size()))) + .build(), + stream) + .tbl; + + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected->view(), table_a->view()); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected->view(), table_b->view()); +} + TEST_F(HybridScanTest, ChunkedReadRowMaskPerPass) { using T = uint32_t; diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 0f3683c22d1e..c1b79f08dde6 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -189,7 +189,6 @@ def expand_scan_for_rank( return scans - # Per-file cache: row-group row counts, keyed by file path. # Avoids a redundant read_parquet_metadata call for every split of the same # file in SplitScan.do_evaluate; populated on first encounter. @@ -200,6 +199,10 @@ def expand_scan_for_rank( # file; populated on first encounter. _parquet_footer_cache: dict[str, Any] = {} +# TODO: Once footer prefetch (#22700) lands, we'll get +# the metadata from IRExecutionContext footer cache +_hybrid_scan_metadata_cache: dict[str, Any] = {} + def _fetch_byte_ranges( paths: list[str], @@ -221,6 +224,7 @@ def _read_with_hybrid_scan( row_group_indices: list[int], stream: Stream, file_metadata: plc.io.parquet_metadata.FileMetaData, + *, split_index: int = 0, total_splits: int = 1, stats_pruning: bool = True, @@ -242,10 +246,15 @@ def _read_with_hybrid_scan( options.set_column_names(with_columns) options.set_filter(plc_filter) - reader = plc.io.experimental.HybridScanReader.from_parquet_metadata( - file_metadata, - options, - ) + # Parse the file metadata once per file and share it across the file's + # splits, rather than re-parsing and copying it per split. + metadata = _hybrid_scan_metadata_cache.get(paths[0]) + if metadata is None: + metadata = plc.io.experimental.HybridScanMetadata.from_parquet_metadata( + file_metadata, options + ) + _hybrid_scan_metadata_cache[paths[0]] = metadata + reader = plc.io.experimental.HybridScanReader.from_metadata(metadata) if stats_pruning: row_group_indices = reader.filter_row_groups_with_stats( @@ -324,7 +333,7 @@ def _read_with_hybrid_scan( [schema[n] for n in payload_names], stream=stream, ) - # Ensure the decode kernels are finished before + # Ensure the decode kernels are finished before # filter_chunks and payload_chunks go out of scope stream.synchronize() return DataFrame( diff --git a/python/pylibcudf/pylibcudf/io/experimental/__init__.py b/python/pylibcudf/pylibcudf/io/experimental/__init__.py index 6c64231eb1e9..b3bfd1f3bdf6 100644 --- a/python/pylibcudf/pylibcudf/io/experimental/__init__.py +++ b/python/pylibcudf/pylibcudf/io/experimental/__init__.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 from pylibcudf.io.experimental.hybrid_scan import ( + HybridScanMetadata, HybridScanReader, UseDataPageMask, ) @@ -9,6 +10,7 @@ __all__ = [ "FileMetaData", # backwards compatibility + "HybridScanMetadata", "HybridScanReader", "UseDataPageMask", ] diff --git a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pxd b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pxd index a19cd7db8bf8..23a90702bace 100644 --- a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pxd +++ b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pxd @@ -13,6 +13,7 @@ from pylibcudf.io.parquet cimport ParquetReaderOptions from pylibcudf.io.parquet_metadata cimport FileMetaData as c_FileMetaData from pylibcudf.io.types cimport TableWithMetadata from pylibcudf.libcudf.io.hybrid_scan cimport ( + hybrid_scan_metadata as cpp_hybrid_scan_metadata, hybrid_scan_reader as cpp_hybrid_scan_reader, use_data_page_mask, ) @@ -23,6 +24,10 @@ from pylibcudf.libcudf.utilities.span cimport device_span cdef device_span[const_uint8_t] _get_device_span(object obj) except * +cdef class HybridScanMetadata: + cdef unique_ptr[cpp_hybrid_scan_metadata] c_obj + + cdef class HybridScanReader: cdef unique_ptr[cpp_hybrid_scan_reader] c_obj cdef Stream _stream diff --git a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyi b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyi index 490d7b8cbd15..b20b3cb263bf 100644 --- a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyi +++ b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyi @@ -23,6 +23,16 @@ class UseDataPageMask(IntEnum): YES: int NO: int +class HybridScanMetadata: + @staticmethod + def from_footer_bytes( + footer_bytes: Buffer, options: ParquetReaderOptions + ) -> HybridScanMetadata: ... + @staticmethod + def from_parquet_metadata( + metadata: FileMetaData, options: ParquetReaderOptions + ) -> HybridScanMetadata: ... + class HybridScanReader: def __init__( self, footer_bytes: Buffer, options: ParquetReaderOptions @@ -31,6 +41,8 @@ class HybridScanReader: def from_parquet_metadata( metadata: FileMetaData, options: ParquetReaderOptions ) -> HybridScanReader: ... + @staticmethod + def from_metadata(metadata: HybridScanMetadata) -> HybridScanReader: ... def parquet_metadata(self) -> FileMetaData: ... def page_index_byte_range(self) -> ByteRangeInfo: ... def setup_page_index(self, page_index_bytes: Buffer) -> None: ... diff --git a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx index 362f382f0abb..b90eb027c3c5 100644 --- a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx +++ b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 +from cython.operator cimport dereference from libc.stdint cimport uint8_t, uintptr_t from libc.stddef cimport size_t from libcpp cimport bool @@ -24,6 +25,7 @@ from pylibcudf.libcudf.io.hybrid_scan cimport ( const_device_span_const_uint8_t, const_size_type, const_uint8_t, + hybrid_scan_metadata as cpp_hybrid_scan_metadata, hybrid_scan_reader as cpp_hybrid_scan_reader, use_data_page_mask as cpp_use_data_page_mask, ) @@ -40,7 +42,7 @@ import pylibcudf.libcudf.io.hybrid_scan UseDataPageMask = pylibcudf.libcudf.io.hybrid_scan.use_data_page_mask -__all__ = ["FileMetaData", "HybridScanReader", "UseDataPageMask"] +__all__ = ["FileMetaData", "HybridScanMetadata", "HybridScanReader", "UseDataPageMask"] cdef device_span[const_uint8_t] _get_device_span(object obj) except *: @@ -54,6 +56,73 @@ cdef device_span[const_uint8_t] _get_device_span(object obj) except *: obj.size) +cdef class HybridScanMetadata: + """Shareable, pre-parsed Parquet file metadata for the hybrid scan reader. + + Parse the metadata of a single Parquet file once, then construct multiple + :class:`HybridScanReader` instances that share it (one per row-group range of the + file) instead of each re-parsing and copying the metadata. + + For details, see :cpp:class:`cudf::io::parquet::experimental::hybrid_scan_metadata` + + Examples + -------- + >>> import pylibcudf as plc + >>> metadata = plc.io.experimental.HybridScanMetadata.from_parquet_metadata( + ... file_metadata, options) + >>> reader = plc.io.experimental.HybridScanReader.from_metadata(metadata) + """ + + @staticmethod + def from_footer_bytes( + const uint8_t[::1] footer_bytes, + ParquetReaderOptions options + ): + """Parse shareable metadata from Parquet footer bytes. + + Parameters + ---------- + footer_bytes : Buffer + Parquet file footer bytes + options : ParquetReaderOptions + Parquet reader options + + Returns + ------- + HybridScanMetadata + """ + cdef HybridScanMetadata self = HybridScanMetadata.__new__(HybridScanMetadata) + with nogil: + self.c_obj = make_unique[cpp_hybrid_scan_metadata]( + host_span[const_uint8_t](&footer_bytes[0], len(footer_bytes)), + options.c_obj + ) + return self + + @staticmethod + def from_parquet_metadata(c_FileMetaData metadata, ParquetReaderOptions options): + """Build shareable metadata from a pre-populated ``FileMetaData``. + + Parameters + ---------- + metadata : FileMetaData + Pre-populated Parquet file metadata + options : ParquetReaderOptions + Parquet reader options + + Returns + ------- + HybridScanMetadata + """ + cdef HybridScanMetadata self = HybridScanMetadata.__new__(HybridScanMetadata) + with nogil: + self.c_obj = make_unique[cpp_hybrid_scan_metadata]( + metadata.c_obj, + options.c_obj + ) + return self + + cdef class HybridScanReader: """Experimental Parquet reader optimized for highly selective filters. @@ -112,6 +181,30 @@ cdef class HybridScanReader: ) return reader + @staticmethod + def from_metadata(HybridScanMetadata metadata): + """Create a HybridScanReader that shares pre-parsed metadata. + + Constructs a lightweight reader that borrows ``metadata`` instead of + re-parsing and copying the file metadata. Use one shared + :class:`HybridScanMetadata` to read disjoint row-group ranges of a single file. + + Parameters + ---------- + metadata : HybridScanMetadata + Shared, pre-parsed Parquet file metadata + + Returns + ------- + HybridScanReader + """ + cdef HybridScanReader reader = HybridScanReader.__new__(HybridScanReader) + with nogil: + reader.c_obj = make_unique[cpp_hybrid_scan_reader]( + dereference(metadata.c_obj.get()) + ) + return reader + def parquet_metadata(self): """Get the Parquet file footer metadata. diff --git a/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan.pxd b/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan.pxd index 5dc33a5be1d4..4fa2b2cb4217 100644 --- a/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan.pxd +++ b/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan.pxd @@ -29,6 +29,17 @@ cdef extern from "cudf/io/experimental/hybrid_scan.hpp" \ YES NO + cdef cppclass hybrid_scan_metadata: + hybrid_scan_metadata( + host_span[const_uint8_t] footer_bytes, + const parquet_reader_options& options + ) except +libcudf_exception_handler + + hybrid_scan_metadata( + const FileMetaData& parquet_metadata, + const parquet_reader_options& options + ) except +libcudf_exception_handler + cdef cppclass hybrid_scan_reader: hybrid_scan_reader( host_span[const_uint8_t] footer_bytes, @@ -40,6 +51,10 @@ cdef extern from "cudf/io/experimental/hybrid_scan.hpp" \ const parquet_reader_options& options ) except +libcudf_exception_handler + hybrid_scan_reader( + const hybrid_scan_metadata& metadata + ) except +libcudf_exception_handler + FileMetaData parquet_metadata() except +libcudf_exception_handler byte_range_info page_index_byte_range() except +libcudf_exception_handler From 7f2cba1632add6b9e18aca7a58b19c891e1a9f80 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Mon, 15 Jun 2026 19:09:15 +0000 Subject: [PATCH 18/20] pre-commit check --- python/cudf_polars/tests/streaming/test_scan.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index a01c2998c6de..b15237af9988 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -38,7 +38,6 @@ from typing import Any, Literal import cudf_polars.engine.core - from cudf_polars.engine.core import StreamingEngine From 4bf25a84ec6427f9397962844399e6b615b128c0 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Thu, 9 Jul 2026 19:31:53 +0000 Subject: [PATCH 19/20] copyright --- ci/check_style.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ci/check_style.sh b/ci/check_style.sh index facdcca72ca4..c8d1f450aaac 100755 --- a/ci/check_style.sh +++ b/ci/check_style.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail From 2362e1d16a6c2f36bcd02d2ecaf2ec44ef8dd97f Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Fri, 10 Jul 2026 17:00:57 +0000 Subject: [PATCH 20/20] docs --- docs/cudf/source/pylibcudf/api_docs/io/experimental.rst | 8 ++++++++ docs/cudf/source/pylibcudf/api_docs/io/index.rst | 1 + 2 files changed, 9 insertions(+) create mode 100644 docs/cudf/source/pylibcudf/api_docs/io/experimental.rst diff --git a/docs/cudf/source/pylibcudf/api_docs/io/experimental.rst b/docs/cudf/source/pylibcudf/api_docs/io/experimental.rst new file mode 100644 index 000000000000..92e9360fdb42 --- /dev/null +++ b/docs/cudf/source/pylibcudf/api_docs/io/experimental.rst @@ -0,0 +1,8 @@ +============ +Experimental +============ + +APIs in this namespace are experimental and may change without warning in the future. + +.. automodule:: pylibcudf.io.experimental + :members: diff --git a/docs/cudf/source/pylibcudf/api_docs/io/index.rst b/docs/cudf/source/pylibcudf/api_docs/io/index.rst index dc485c07de14..45b4def70051 100644 --- a/docs/cudf/source/pylibcudf/api_docs/io/index.rst +++ b/docs/cudf/source/pylibcudf/api_docs/io/index.rst @@ -17,6 +17,7 @@ I/O Functions avro csv + experimental json orc parquet