diff --git a/cpp/include/cudf/io/experimental/hybrid_scan.hpp b/cpp/include/cudf/io/experimental/hybrid_scan.hpp index 759d76100161..e1b65d2d65ef 100644 --- a/cpp/include/cudf/io/experimental/hybrid_scan.hpp +++ b/cpp/include/cudf/io/experimental/hybrid_scan.hpp @@ -26,6 +26,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 @@ -49,6 +54,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 @@ -297,6 +369,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/include/cudf/io/parquet_io_utils.hpp b/cpp/include/cudf/io/parquet_io_utils.hpp index 63066065ec2d..57f23ddb1d13 100644 --- a/cpp/include/cudf/io/parquet_io_utils.hpp +++ b/cpp/include/cudf/io/parquet_io_utils.hpp @@ -152,6 +152,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 `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/experimental/hybrid_scan.cpp b/cpp/src/io/parquet/experimental/hybrid_scan.cpp index 36bc5de06dc4..3ded187fe997 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 08341df6c314..2b83ebaed804 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp @@ -98,7 +98,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()); @@ -109,13 +109,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 db9dd9acfe41..6a4190c9a72b 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp @@ -59,6 +59,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::parquet::experimental::hybrid_scan_multifile::parquet_metadatas */ 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 9740ae71a435..99741dbf3a0a 100644 --- a/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp +++ b/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp @@ -471,4 +471,17 @@ 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/cpp/src/io/parquet/reader_impl.cpp b/cpp/src/io/parquet/reader_impl.cpp index 1155628159bf..066f0345a83c 100644 --- a/cpp/src/io/parquet/reader_impl.cpp +++ b/cpp/src/io/parquet/reader_impl.cpp @@ -522,12 +522,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 81ae58ab05da..03777e13e88c 100644 --- a/cpp/src/io/parquet/reader_impl.hpp +++ b/cpp/src/io/parquet/reader_impl.hpp @@ -489,7 +489,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_common.cpp b/cpp/tests/io/experimental/hybrid_scan_common.cpp index 99ee37e6c333..f65808b4f721 100644 --- a/cpp/tests/io/experimental/hybrid_scan_common.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_common.cpp @@ -328,6 +328,7 @@ std::pair, std::vector> create_parquet_with_s INSTANTIATE_CREATE_PARQUET_WITH_STATS(T, 1, true, true) INSTANTIATE_CREATE_PARQUET_WITH_STATS(uint32_t, 4, true, false); +INSTANTIATE_CREATE_PARQUET_WITH_STATS(int32_t, 2, true, false); INSTANTIATE_CREATE_PARQUET_WITH_STATS(cudf::timestamp_ms, 2, true, false); INSTANTIATE_CREATE_PARQUET_WITH_STATS(cudf::duration_ms, 2, true, false); diff --git a/cpp/tests/io/experimental/hybrid_scan_test.cpp b/cpp/tests/io/experimental/hybrid_scan_test.cpp index e580994b6d34..c7fde4db53e0 100644 --- a/cpp/tests/io/experimental/hybrid_scan_test.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_test.cpp @@ -878,6 +878,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, AllRowsPrunedReportsInputRowGroups) { using cudf::io::parquet::experimental::use_data_page_mask; diff --git a/cpp/tests/strings/contains_tests.cpp b/cpp/tests/strings/contains_tests.cpp index 89ed9756d00a..481d788ad3cb 100644 --- a/cpp/tests/strings/contains_tests.cpp +++ b/cpp/tests/strings/contains_tests.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/tests/strings/findall_tests.cpp b/cpp/tests/strings/findall_tests.cpp index e4999a41bcc7..5076f6f7f50f 100644 --- a/cpp/tests/strings/findall_tests.cpp +++ b/cpp/tests/strings/findall_tests.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/tests/strings/replace_regex_tests.cpp b/cpp/tests/strings/replace_regex_tests.cpp index 8422264d4c36..826d0e3cea41 100644 --- a/cpp/tests/strings/replace_regex_tests.cpp +++ b/cpp/tests/strings/replace_regex_tests.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ 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 15a87175325f..45b4def70051 100644 --- a/docs/cudf/source/pylibcudf/api_docs/io/index.rst +++ b/docs/cudf/source/pylibcudf/api_docs/io/index.rst @@ -17,9 +17,11 @@ I/O Functions avro csv + experimental json orc parquet + parquet_io_utils parquet_metadata text timezone diff --git a/docs/cudf/source/pylibcudf/api_docs/io/parquet_io_utils.rst b/docs/cudf/source/pylibcudf/api_docs/io/parquet_io_utils.rst new file mode 100644 index 000000000000..3f01b0493fd4 --- /dev/null +++ b/docs/cudf/source/pylibcudf/api_docs/io/parquet_io_utils.rst @@ -0,0 +1,6 @@ +================ +Parquet IO Utils +================ + +.. automodule:: pylibcudf.io.parquet_io_utils + :members: diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index f61a738d0bbd..9bfcb5721e05 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -49,7 +49,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/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index 288dcd79f460..58100f83b17a 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -6,7 +6,7 @@ import concurrent.futures import contextlib -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING import pylibcudf as plc @@ -46,6 +46,23 @@ class CachedParquetInfo: path: str size: int | None file_metadata: plc.io.parquet_metadata.FileMetaData + # Shared, pre-parsed hybrid-scan metadata, built once per file and borrowed by + # all of the file's SplitScans. Excluded from identity so it never affects hashing. + _hybrid_scan_metadata: list[plc.io.experimental.HybridScanMetadata] = field( + default_factory=list, compare=False, repr=False + ) + + def hybrid_scan_metadata( + self, options: plc.io.parquet.ParquetReaderOptions + ) -> plc.io.experimental.HybridScanMetadata: + """Return the shared hybrid-scan metadata, parsing it once per file.""" + if not self._hybrid_scan_metadata: + self._hybrid_scan_metadata.append( + plc.io.experimental.HybridScanMetadata.from_parquet_metadata( + self.file_metadata, options + ) + ) + return self._hybrid_scan_metadata[0] @nvtx_annotate_cudf_polars(message="fetch_parquet_footers_for_paths") 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 4c5773874aa5..284505100c29 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py @@ -30,6 +30,7 @@ _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.actor_graph.dispatch import ( generate_ir_sub_network, ) @@ -530,23 +531,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 b3b4438812c5..8631a8b94f01 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -17,6 +17,7 @@ import pylibcudf as plc +from cudf_polars.containers import Column, DataFrame from cudf_polars.dsl.ir import ( IR, DataFrameScan, @@ -24,7 +25,9 @@ PythonScan, Scan, Sink, + _prepare_parquet_predicate, ) +from cudf_polars.dsl.to_ast import to_parquet_filter from cudf_polars.dsl.tracing import nvtx_annotate_cudf_polars from cudf_polars.streaming.base import ( IOPartitionFlavor, @@ -41,7 +44,10 @@ if TYPE_CHECKING: from collections.abc import Hashable, MutableMapping, Sequence - from cudf_polars.containers import DataFrame, DataType + import pylibcudf.expressions as plc_expr + from rmm.pylibrmm.stream import Stream + + from cudf_polars.containers import DataType from cudf_polars.dsl.expr import NamedExpr from cudf_polars.dsl.ir import CachedParquetInfo, IRExecutionContext from cudf_polars.streaming.base import ( @@ -82,6 +88,13 @@ def scan_partition_plan( """Extract the partitioning plan of a Scan operation.""" if ir.typ == "parquet": blocksize: int = config_options.executor.target_partition_size + single_file = len(ir.paths) == 1 + # A single file always uses SplitScan when hybrid scan is enabled, so the + # hybrid reader can be used on it even when it would otherwise not split. + # The split factor is still size-based, so a large file is split into many. + hybrid_single_file = ( + single_file and config_options.parquet_options.use_hybrid_scan + ) if source := stats.scan_stats.get(ir): column_sizes = [ sz @@ -98,12 +111,18 @@ def scan_partition_plan( <= abs(file_size / k_hi - blocksize) else k_hi ) - if factor >= 2: + if factor >= 2 or hybrid_single_file: return IOPartitionPlan( factor, IOPartitionFlavor.SPLIT_FILES, estimated_chunk_bytes=file_size // factor, ) + elif hybrid_single_file: + return IOPartitionPlan( + 1, + IOPartitionFlavor.SPLIT_FILES, + estimated_chunk_bytes=file_size, + ) else: k_lo = min(blocksize // int(file_size), len(ir.paths)) k_hi = k_lo + 1 @@ -120,6 +139,9 @@ def scan_partition_plan( estimated_chunk_bytes=file_size * factor, ) + if hybrid_single_file: + return IOPartitionPlan(1, IOPartitionFlavor.SPLIT_FILES) + # TODO: Use file sizes for csv and json return IOPartitionPlan(1, IOPartitionFlavor.SINGLE_FILE) @@ -182,6 +204,139 @@ def expand_scan_for_rank( ) +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], + with_columns: list[str] | None, + plc_filter: plc_expr.Expression, + row_group_indices: list[int], + stream: Stream, + cached_info: CachedParquetInfo, + *, + split_index: int = 0, + total_splits: int = 1, + stats_pruning: bool = True, +) -> DataFrame: + """Two-pass parquet read via HybridScanReader for a row-group-aligned split.""" + 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]} [{split_index + 1}/{total_splits}]" + ): + options = ( + plc.io.parquet.ParquetReaderOptions.builder(plc.io.SourceInfo(paths)) + .decimal_width(plc.TypeId.DECIMAL128) + .build() + ) + if with_columns is not None: + options.set_column_names(with_columns) + options.set_filter(plc_filter) + + # Borrow the shared, pre-parsed metadata (built once per file) so each + # split does not re-parse and copy the file metadata. + reader = plc.io.experimental.HybridScanReader.from_metadata( + cached_info.hybrid_scan_metadata(options) + ) + + if stats_pruning: + row_group_indices = reader.filter_row_groups_with_stats( + row_group_indices, options, stream=stream + ) + + if row_group_indices: + bloom_ranges, _ = reader.secondary_filters_byte_ranges( + row_group_indices, options + ) + if bloom_ranges: + bloom_chunks = _fetch_byte_ranges(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, + ) + + 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, + 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, + ) + # Ensure the decode kernels are finished before + # filter_chunks and payload_chunks go out of scope + stream.synchronize() + return DataFrame( + [*filter_df.columns, *payload_df.columns], stream=stream + ).select(list(schema.keys())) + + class SplitScan(IR): """ Input from a split file. @@ -337,6 +492,43 @@ def do_evaluate( skip_rgs = rg_stride * split_index skip_rows = sum(row_group_num_rows[:skip_rgs]) n_rows = sum(row_group_num_rows[skip_rgs : skip_rgs + rg_stride]) + # Hybrid scan reads through the prefetched, shared file metadata, so + # it is only used when footer prefetching is enabled. + # TODO: Investigate re-enabling for some of the excluded paths + # (row_index / include_file_paths). Needs performance investigation. + if ( + parquet_options.use_hybrid_scan + and cached_parquet_info is not None + and row_index is None + and include_file_paths is None + and predicate is not None + ): + stream = context.get_cuda_stream() + plc_filter = to_parquet_filter( + _prepare_parquet_predicate( + predicate.value, paths, schema, with_columns + ), + stream=stream, + ) + if plc_filter is not None: + end_rg = ( + total_row_groups + if split_index == total_splits - 1 + else skip_rgs + rg_stride + ) + return _read_with_hybrid_scan( + schema, + paths, + with_columns, + plc_filter, + list(range(skip_rgs, end_rg)), + stream, + cached_parquet_info[0], + split_index=split_index, + total_splits=total_splits, + stats_pruning=parquet_options.hybrid_scan_stats_pruning, + ) + else: # There are not enough row-groups to align # all "total_splits" of our reads with row-group diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index e81355411f85..3a90891dd9fe 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -232,6 +232,19 @@ 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. + 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``. prefetch_file_metadata Whether to prefetch parquet file metadata and pass it through `parquet_metadatas` to avoid rereading file footers. @@ -281,6 +294,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, + ) + ) prefetch_file_metadata: bool = dataclasses.field( default_factory=_make_default_factory( f"{_env_prefix}__PREFETCH_FILE_METADATA", @@ -288,6 +308,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, + ) + ) use_jit_filter: bool = dataclasses.field( default_factory=_make_default_factory( f"{_env_prefix}__USE_JIT_FILTER", @@ -311,6 +338,10 @@ 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") + if not isinstance(self.hybrid_scan_stats_pruning, bool): + raise TypeError("hybrid_scan_stats_pruning must be a bool") if not isinstance(self.prefetch_file_metadata, bool): raise TypeError("prefetch_file_metadata must be a bool") diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 13e88ead7731..c1efebb40b43 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -359,6 +359,42 @@ def test_streaming_scan_raises() -> None: StreamingScan.do_evaluate([fused], scan, context=ctx) +@pytest.mark.parametrize( + "predicate,use_columns", + [ + # uses hybrid scan reader + (pl.col("x") < 1_000, None), + (pl.col("x") < 1_000, ["x", "z"]), + # fallsback to default parquet reader + (pl.col("y").str.contains("cat"), None), + (None, None), + ], +) +def test_split_scan_hybrid( + tmp_path: Path, + df: pl.DataFrame, + predicate: pl.Expr | None, + use_columns: list[str] | None, + streaming_engine_factory: Callable[..., StreamingEngine], +) -> None: + streaming_engine = streaming_engine_factory( + StreamingOptions( + target_partition_size=1_000, + parquet_options={ + "use_hybrid_scan": True, + "prefetch_file_metadata": True, + }, + ), + ) + make_partitioned_source(df, tmp_path, "parquet", n_files=1, row_group_size=100) + q = pl.scan_parquet(tmp_path) + if 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) + + def test_scan_missing_prefetch_metadata_raises() -> None: # This isn't reachable by polars' public API, so we test it directly. scan = _make_parquet_scan( diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index e5deb48cb14d..5b7e3a3b6a7d 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -332,6 +332,8 @@ def test_parquet_options_from_env(monkeypatch: pytest.MonkeyPatch) -> None: m.setenv("CUDF_POLARS__PARQUET_OPTIONS__MAX_FOOTER_SAMPLES", "0") m.setenv("CUDF_POLARS__PARQUET_OPTIONS__MAX_ROW_GROUP_SAMPLES", "0") m.setenv("CUDF_POLARS__PARQUET_OPTIONS__USE_RAPIDSMPF_NATIVE", "0") + m.setenv("CUDF_POLARS__PARQUET_OPTIONS__USE_HYBRID_SCAN", "0") + m.setenv("CUDF_POLARS__PARQUET_OPTIONS__HYBRID_SCAN_STATS_PRUNING", "0") m.setenv("CUDF_POLARS__PARQUET_OPTIONS__PREFETCH_FILE_METADATA", "1") m.setenv("CUDF_POLARS__PARQUET_OPTIONS__USE_JIT_FILTER", "1") @@ -345,6 +347,8 @@ def test_parquet_options_from_env(monkeypatch: pytest.MonkeyPatch) -> None: assert config.parquet_options.max_footer_samples == 0 assert config.parquet_options.max_row_group_samples == 0 assert config.parquet_options.use_rapidsmpf_native is False + assert config.parquet_options.use_hybrid_scan is False + assert config.parquet_options.hybrid_scan_stats_pruning is False assert config.parquet_options.prefetch_file_metadata is True assert config.parquet_options.use_jit_filter is True diff --git a/python/pylibcudf/pylibcudf/io/CMakeLists.txt b/python/pylibcudf/pylibcudf/io/CMakeLists.txt index 089ea8d0e8d9..65bb31d908ff 100644 --- a/python/pylibcudf/pylibcudf/io/CMakeLists.txt +++ b/python/pylibcudf/pylibcudf/io/CMakeLists.txt @@ -1,12 +1,12 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= 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 a6a0ebad3a1e..1f0a0a218199 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, @@ -30,6 +31,7 @@ "json", "orc", "parquet", + "parquet_io_utils", "parquet_metadata", "text", "timezone", diff --git a/python/pylibcudf/pylibcudf/io/experimental/__init__.py b/python/pylibcudf/pylibcudf/io/experimental/__init__.py index 6c64231eb1e9..ef1ca25b7cff 100644 --- a/python/pylibcudf/pylibcudf/io/experimental/__init__.py +++ b/python/pylibcudf/pylibcudf/io/experimental/__init__.py @@ -1,7 +1,8 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # 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..41255f91b5ff 100644 --- a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pxd +++ b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pxd @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from libc.stdint cimport uint8_t @@ -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 f95dc8b054d3..498a458b9f7d 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 & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from collections.abc import Sequence from enum import IntEnum from rmm.pylibrmm.memory_resource import DeviceMemoryResource @@ -10,6 +11,7 @@ from pylibcudf.io.parquet import ParquetReaderOptions from pylibcudf.io.parquet_metadata import FileMetaData from pylibcudf.io.text import ByteRangeInfo from pylibcudf.io.types import TableWithMetadata +from pylibcudf.span import Span from pylibcudf.utils import CudaStreamLike try: @@ -21,6 +23,16 @@ class UseDataPageMask(IntEnum): YES = 1 NO = 0 +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 @@ -29,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: ... @@ -48,18 +62,24 @@ class HybridScanReader: ) -> tuple[list[ByteRangeInfo], list[ByteRangeInfo]]: ... def filter_row_groups_with_dictionary_pages( self, - dictionary_page_data: list, + 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, + bloom_filter_data: Sequence[Span], row_group_indices: list[int], 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], @@ -73,7 +93,7 @@ class HybridScanReader: def materialize_filter_columns( self, row_group_indices: list[int], - column_chunk_data: list, + column_chunk_data: Sequence[Span], row_mask: Column, mask_data_pages: UseDataPageMask, options: ParquetReaderOptions, @@ -86,7 +106,7 @@ class HybridScanReader: def materialize_payload_columns( self, row_group_indices: list[int], - column_chunk_data: list, + column_chunk_data: Sequence[Span], row_mask: Column, mask_data_pages: UseDataPageMask, options: ParquetReaderOptions, @@ -99,7 +119,7 @@ class HybridScanReader: def materialize_all_columns( self, row_group_indices: list[int], - column_chunk_data: list, + column_chunk_data: Sequence[Span], options: ParquetReaderOptions, stream: CudaStreamLike | None = None, mr: DeviceMemoryResource | None = None, @@ -111,7 +131,7 @@ class HybridScanReader: row_group_indices: list[int], row_mask: Column, mask_data_pages: UseDataPageMask, - column_chunk_data: list, + column_chunk_data: Sequence[Span], options: ParquetReaderOptions, stream: CudaStreamLike | None = None, mr: DeviceMemoryResource | None = None, @@ -127,7 +147,7 @@ class HybridScanReader: row_group_indices: list[int], row_mask: Column, mask_data_pages: UseDataPageMask, - column_chunk_data: list, + column_chunk_data: Sequence[Span], options: ParquetReaderOptions, stream: CudaStreamLike | None = None, mr: DeviceMemoryResource | None = None, diff --git a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx index 664eb489428c..1b5e9fa792a0 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 & AFFILIATES. All rights reserved. # 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.span cimport span as std_span @@ -15,6 +17,7 @@ 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 @@ -23,6 +26,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, ) @@ -39,7 +43,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 *: @@ -53,6 +57,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. @@ -82,10 +153,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): @@ -103,10 +175,35 @@ 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 + + @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): @@ -117,7 +214,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. @@ -127,7 +227,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): @@ -138,9 +240,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. @@ -155,9 +258,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 +277,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( - std_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( + std_span[const_size_type](indices_vec.data(), indices_vec.size()) + ) + return result def reset_column_selection(self): """Reset the column selection state. @@ -184,7 +290,8 @@ 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 filter_row_groups_with_stats( self, @@ -210,15 +317,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( std_span[const_size_type]( indices_vec.data(), indices_vec.size() ), options.c_obj, _stream.view().value() - ) - ) + )) return list(filtered) def secondary_filters_byte_ranges( @@ -241,11 +348,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( std_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 @@ -287,15 +395,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( std_span[const_device_span_const_uint8_t]( spans_vec.data(), spans_vec.size() ), std_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( @@ -330,17 +439,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( std_span[const_device_span_const_uint8_t]( spans_vec.data(), spans_vec.size() ), std_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, @@ -369,13 +513,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( std_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( @@ -398,11 +543,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( std_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( @@ -448,8 +594,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( std_span[const_size_type](indices_vec.data(), indices_vec.size()), std_span[const_device_span_const_uint8_t]( spans_vec.data(), spans_vec.size() @@ -459,7 +606,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( @@ -482,11 +629,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( std_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( @@ -532,8 +680,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( std_span[const_size_type](indices_vec.data(), indices_vec.size()), std_span[const_device_span_const_uint8_t]( spans_vec.data(), spans_vec.size() @@ -543,7 +692,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( @@ -566,11 +715,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( std_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( @@ -608,8 +758,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( std_span[const_size_type](indices_vec.data(), indices_vec.size()), std_span[const_device_span_const_uint8_t]( spans_vec.data(), spans_vec.size() @@ -617,7 +768,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( @@ -665,19 +816,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, - std_span[const_size_type](indices_vec.data(), indices_vec.size()), - mask_view, - mask_data_pages, - std_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, + std_span[const_size_type](indices_vec.data(), indices_vec.size()), + mask_view, + mask_data_pages, + std_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, @@ -695,10 +847,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 ) @@ -748,19 +901,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, - std_span[const_size_type](indices_vec.data(), indices_vec.size()), - mask_view, - mask_data_pages, - std_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, + std_span[const_size_type](indices_vec.data(), indices_vec.size()), + mask_view, + mask_data_pages, + std_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, @@ -778,10 +932,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 ) @@ -817,12 +972,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( - std_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( + std_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. @@ -832,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/io/parquet_io_utils.pxd b/python/pylibcudf/pylibcudf/io/parquet_io_utils.pxd new file mode 100644 index 000000000000..e8d872e08f61 --- /dev/null +++ b/python/pylibcudf/pylibcudf/io/parquet_io_utils.pxd @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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 + +cpdef list fetch_byte_ranges_to_device( + SourceInfo source_info, + list byte_ranges, + 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 new file mode 100644 index 000000000000..1a18ab72f5d5 --- /dev/null +++ b/python/pylibcudf/pylibcudf/io/parquet_io_utils.pyi @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from rmm.pylibrmm.memory_resource import DeviceMemoryResource + +from pylibcudf.gpumemoryview import gpumemoryview +from pylibcudf.io.text import ByteRangeInfo +from pylibcudf.io.types import SourceInfo +from pylibcudf.utils import CudaStreamLike + +__all__ = ["fetch_byte_ranges_to_device", "fetch_page_index_to_host"] + +def fetch_byte_ranges_to_device( + source_info: SourceInfo, + byte_ranges: list[ByteRangeInfo], + 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 new file mode 100644 index 000000000000..0f872780845a --- /dev/null +++ b/python/pylibcudf/pylibcudf/io/parquet_io_utils.pyx @@ -0,0 +1,163 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""IO utilities for the Parquet.""" + +from libc.stddef cimport size_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 +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, + 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", "fetch_page_index_to_host"] + + +cpdef list fetch_byte_ranges_to_device( + SourceInfo source_info, + list byte_ranges, + object stream=None, + DeviceMemoryResource mr=None, +): + """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 + + +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..8807e536f85e 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-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # 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 + 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 36201d545de6..d8b9255a9dab 100644 --- a/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan.pxd +++ b/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan.pxd @@ -30,6 +30,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, @@ -41,6 +52,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 @@ -86,6 +101,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( std_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 new file mode 100644 index 000000000000..f00ff47de471 --- /dev/null +++ b/python/pylibcudf/pylibcudf/libcudf/io/parquet_io_utils.pxd @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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 + +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 + + unique_ptr[datasource.buffer] fetch_page_index_to_host( + datasource& ds, + byte_range_info page_index_bytes, + ) except +libcudf_exception_handler diff --git a/python/pylibcudf/pylibcudf/libcudf/utilities/span.pxd b/python/pylibcudf/pylibcudf/libcudf/utilities/span.pxd index f2bf388e4d4c..c33ec94144be 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-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # 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