diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 559777f43c7d..7e7ebb6785d3 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -902,12 +902,7 @@ def _get_parquet_row_count_from_metadata( ) -> int: # Zero-width parquet files lose their row count when read through # pylibcudf. See https://github.com/rapidsai/cudf/issues/21428 - if parquet_options.prefetch_file_metadata: - if cached_parquet_info is None: - raise AssertionError( - "Cached parquet info is required when prefetching file metadata is enabled" - ) - + if cached_parquet_info is not None: Scan._validate_cached_parquet_info(paths, cached_parquet_info) parquet_metadatas = [ info.file_metadata for info in cached_parquet_info @@ -1071,11 +1066,7 @@ def read_csv_header( df, ) elif typ == "parquet": - if parquet_options.prefetch_file_metadata: - if cached_parquet_info is None: - raise AssertionError( - "Cached parquet info is required when prefetching file metadata is enabled" - ) + if cached_parquet_info is not None: Scan._validate_cached_parquet_info(paths, cached_parquet_info) filepath_sources = [] parquet_metadatas = [] diff --git a/python/cudf_polars/cudf_polars/dsl/utils/io.py b/python/cudf_polars/cudf_polars/dsl/utils/io.py index 75373afeed96..3ac9db651b01 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/io.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/io.py @@ -110,6 +110,8 @@ def prefetch_parquet_file_metadata_for_ir( root: IR, py_executor: concurrent.futures.Executor | None, stats: StatsCollector | None = None, + *, + remote_only: bool = False, ) -> dict[str, CachedParquetInfo]: """ Prefetch parquet metadata for all parquet scans in an IR graph. @@ -125,6 +127,9 @@ def prefetch_parquet_file_metadata_for_ir( prefetched during statistics collection, when the number of files sampled equals the total number of files. Providing ``stats`` here will skip rereading metadata for those files. + remote_only + If ``True``, only prefetch metadata for remote URIs (e.g. ``s3://``), + skipping local paths. Returns ------- @@ -135,7 +140,7 @@ def prefetch_parquet_file_metadata_for_ir( all_paths: set[str] = set() for node in traversal([root]): - if isinstance(node, StreamingScan): + if isinstance(node, StreamingScan) and node.base_scan.typ == "parquet": for scan in node.scans: for path in scan.paths: all_paths.add(path) @@ -155,6 +160,10 @@ def prefetch_parquet_file_metadata_for_ir( cached_parquet_info[info.path] = info missing_paths = all_paths - set(cached_parquet_info.keys()) + if remote_only: + missing_paths = { + p for p in missing_paths if plc.io.SourceInfo._is_remote_uri(p) + } cm: contextlib.AbstractContextManager[concurrent.futures.Executor | None] if py_executor is None: @@ -194,8 +203,10 @@ def attach_cached_parquet_metadata( Mapping from file paths to cached parquet metadata. """ for node in traversal([root]): - if isinstance(node, StreamingScan): + if isinstance(node, StreamingScan) and node.base_scan.typ == "parquet": for scan in node.scans: + if not all(path in cached_parquet_info_map for path in scan.paths): + continue cached = [cached_parquet_info_map[path] for path in scan.paths] Scan._validate_cached_parquet_info(scan.paths, cached) scan.cached_parquet_info = cached diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index b382f12099a4..771ec1f22177 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -42,7 +42,7 @@ from cudf_polars.streaming.parallel import lower_ir_graph_with_node_map from cudf_polars.streaming.statistics import collect_statistics from cudf_polars.streaming.utils import _concat -from cudf_polars.utils.config import get_total_device_memory +from cudf_polars.utils.config import Unspecified, get_total_device_memory if TYPE_CHECKING: from collections.abc import Callable, MutableMapping @@ -778,11 +778,13 @@ def evaluate_on_rank( py_executor, get_cuda_stream=ctx.br().stream_pool.get_stream, query_id=query_id ) - if config_options.parquet_options.prefetch_file_metadata: + prefetch_file_metadata = config_options.parquet_options.prefetch_file_metadata + if prefetch_file_metadata is not False: cached_parquet_info_map = prefetch_parquet_file_metadata_for_ir( ir, ir_context.py_executor, stats=stats, + remote_only=isinstance(prefetch_file_metadata, Unspecified), ) attach_cached_parquet_metadata(ir, cached_parquet_info_map) diff --git a/python/cudf_polars/cudf_polars/engine/options.py b/python/cudf_polars/cudf_polars/engine/options.py index 611129cf2c01..22cbab1731e8 100644 --- a/python/cudf_polars/cudf_polars/engine/options.py +++ b/python/cudf_polars/cudf_polars/engine/options.py @@ -18,7 +18,7 @@ from cudf_polars.engine.hardware_binding import ( HardwareBindingPolicy, ) -from cudf_polars.utils.config import MemoryResourceConfig +from cudf_polars.utils.config import UNSPECIFIED, MemoryResourceConfig, Unspecified if TYPE_CHECKING: from collections.abc import Callable @@ -37,38 +37,6 @@ ] -class Unspecified: - """ - Sentinel value meaning "fall back to environment variable, then built-in default". - - The singleton instance :data:`UNSPECIFIED` is used as the default for every - :class:`StreamingOptions` field. When a field is still ``UNSPECIFIED`` after - construction (i.e. neither an explicit value nor an environment variable was provided), - the underlying library applies its own built-in default. - """ - - _instance: Unspecified | None = None - - def __new__(cls) -> Unspecified: - """Return the singleton instance.""" - if cls._instance is None: - cls._instance = super().__new__(cls) - return cls._instance - - def __repr__(self) -> str: - """Return ``"UNSPECIFIED"``.""" - return "UNSPECIFIED" - - -UNSPECIFIED = Unspecified() -"""Singleton sentinel for all :class:`StreamingOptions` fields. - -A field set to ``UNSPECIFIED`` after construction means no explicit value and no -matching environment variable was found; the underlying library will apply its own -built-in default. -""" - - def _opt( category: str, env_var: str | None = None, diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 185f43b7142f..ec144fc2fc73 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -49,6 +49,7 @@ __all__ = [ + "UNSPECIFIED", "Cluster", "ConfigOptions", "DaskContext", @@ -60,9 +61,45 @@ "SPMDContext", "StreamingExecutor", "StreamingFallbackMode", + "Unspecified", ] +class Unspecified: + """ + Sentinel value meaning "no value was explicitly provided". + + The singleton instance :data:`UNSPECIFIED` is used as the default for every + :class:`StreamingOptions` field, as well as for + ``ParquetOptions.prefetch_file_metadata``. When a field is still + ``UNSPECIFIED`` after construction (i.e. neither an explicit value nor a + matching environment variable was provided), the consuming component decides + on the semantics. + """ + + _instance: Unspecified | None = None + + def __new__(cls) -> Unspecified: + """Return the singleton instance.""" + if cls._instance is None: + cls._instance = super().__new__(cls) + return cls._instance + + def __repr__(self) -> str: + """Return ``"UNSPECIFIED"``.""" + return "UNSPECIFIED" + + +UNSPECIFIED = Unspecified() +"""Singleton sentinel for all :class:`StreamingOptions` fields, as well as for +``ParquetOptions.prefetch_file_metadata``. + +A field set to ``UNSPECIFIED`` after construction means no explicit value and no +matching environment variable was found; the consuming component decides on the +semantics. +""" + + def _env_get_int(name: str, default: int) -> int: try: return int(os.getenv(name, default)) @@ -221,7 +258,10 @@ class ParquetOptions: will also be skipped if ``max_footer_samples`` is 0. prefetch_file_metadata Whether to prefetch parquet file metadata and pass it through - `parquet_metadatas` to avoid rereading file footers. + `parquet_metadatas` to avoid rereading file footers. Not supported + by the in-memory executor, where it defaults to disabled. For the + streaming executor, it defaults to being enabled for remote URIs + (e.g. ``s3://``) only; pass ``True`` to also prefetch local files. use_jit_filter Whether to use JIT compilation for post-read filtering in Parquet scans. When enabled, filter predicates are JIT-compiled to CUDA kernels for @@ -261,11 +301,11 @@ class ParquetOptions: f"{_env_prefix}__MAX_ROW_GROUP_SAMPLES", int, default=1 ) ) - prefetch_file_metadata: bool = dataclasses.field( + prefetch_file_metadata: bool | Unspecified = dataclasses.field( default_factory=_make_default_factory( f"{_env_prefix}__PREFETCH_FILE_METADATA", _bool_converter, - default=False, + default=UNSPECIFIED, ) ) use_jit_filter: bool = dataclasses.field( @@ -289,8 +329,8 @@ def __post_init__(self) -> None: # noqa: D105 raise TypeError("max_footer_samples must be an int") if not isinstance(self.max_row_group_samples, int): raise TypeError("max_row_group_samples must be an int") - if not isinstance(self.prefetch_file_metadata, bool): - raise TypeError("prefetch_file_metadata must be a bool") + if not isinstance(self.prefetch_file_metadata, (bool, Unspecified)): + raise TypeError("prefetch_file_metadata must be a bool when specified") if not isinstance(self.use_jit_filter, bool): raise TypeError("use_jit_filter must be a bool") @@ -960,9 +1000,31 @@ def from_polars_engine( if user_parquet_options is None: user_parquet_options = {} + # Engine-dependent default: only prefetch for the streaming executor. + # Skipped if the user or the environment has already set a value. + prefetch_default = UNSPECIFIED if user_executor == "streaming" else False + prefetch_env_set = ( + os.environ.get(f"{ParquetOptions._env_prefix}__PREFETCH_FILE_METADATA") + is not None + ) + if isinstance(user_parquet_options, dict): + user_parquet_options = dict(user_parquet_options) + if ( + "prefetch_file_metadata" not in user_parquet_options + and not prefetch_env_set + ): + user_parquet_options["prefetch_file_metadata"] = prefetch_default parquet_options = ParquetOptions(**user_parquet_options) else: + if ( + isinstance(user_parquet_options.prefetch_file_metadata, Unspecified) + and not prefetch_env_set + ): + user_parquet_options = dataclasses.replace( + user_parquet_options, + prefetch_file_metadata=prefetch_default, + ) parquet_options = user_parquet_options # This is set in polars, and so can't be overridden by the environment user_raise_on_fail = engine.config.get("raise_on_fail", False) @@ -991,7 +1053,7 @@ def from_polars_engine( match user_executor: case "in-memory": executor = InMemoryExecutor(**user_executor_options) - if parquet_options.prefetch_file_metadata: + if parquet_options.prefetch_file_metadata is True: raise NotImplementedError( "Prefetching is not supported for the in-memory executor." ) diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 0fcf3c663a9f..a44af628ea1c 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -126,6 +126,27 @@ def test_prefetch_parquet_file_metadata_no_parquet_scans() -> None: assert result == {} +def test_prefetch_parquet_file_metadata_remote_only(tmp_path, df) -> None: + make_partitioned_source(df, tmp_path, "parquet", n_files=1) + local_path = str(next(tmp_path.glob("*.parquet"))) + + scan = _make_parquet_scan([local_path]) + fused = FusedScan(scan.schema, scan, scan.paths, scan.parquet_options, []) + streaming_scan = StreamingScan([fused], scan, "fused") + + # Local paths are skipped entirely when remote_only=True. + result = prefetch_parquet_file_metadata_for_ir( + streaming_scan, py_executor=None, stats=None, remote_only=True + ) + assert result == {} + + # The same local path is prefetched when remote_only=False (the default). + result = prefetch_parquet_file_metadata_for_ir( + streaming_scan, py_executor=None, stats=None + ) + assert set(result) == {local_path} + + def test_prefetch_file_metadata_select_fast_count( df: pl.DataFrame, streaming_engine_factory: Callable[..., StreamingEngine], @@ -349,33 +370,13 @@ def test_streaming_scan_raises() -> None: StreamingScan.do_evaluate([fused], scan, context=ctx) -def test_scan_missing_prefetch_metadata_raises() -> None: +def test_scan_path_mismatch_raises() -> None: # This isn't reachable by polars' public API, so we test it directly. scan = _make_parquet_scan( ["file.parquet"], parquet_options=ParquetOptions(prefetch_file_metadata=True) ) ctx = IRExecutionContext() - with pytest.raises( - AssertionError, - match=r"Cached parquet info is required", - ): - Scan.do_evaluate( - scan.schema, - scan.typ, - scan.reader_options, - scan.paths, - scan.with_columns, - scan.skip_rows, - scan.n_rows, - scan.row_index, - scan.include_file_paths, - scan.predicate, - scan.parquet_options, - None, - context=ctx, - ) - with pytest.raises( AssertionError, match=r"Paths do not match cached parquet info", diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index dc726970101a..883d5cd06374 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -33,7 +33,9 @@ InMemoryExecutor, JoinFilterPushdownOptions, MemoryResourceConfig, + ParquetOptions, StreamingExecutor, + Unspecified, ) from cudf_polars.utils.cuda_stream import get_cuda_stream @@ -374,6 +376,13 @@ def test_parquet_options_from_env(monkeypatch: pytest.MonkeyPatch) -> None: assert config.parquet_options.prefetch_file_metadata is True assert config.parquet_options.use_jit_filter is True + with monkeypatch.context() as m: + # Env must win over the executor-derived default (streaming => True). + m.setenv("CUDF_POLARS__PARQUET_OPTIONS__PREFETCH_FILE_METADATA", "0") + engine = pl.GPUEngine(executor="streaming") + config = ConfigOptions.from_polars_engine(engine) + assert config.parquet_options.prefetch_file_metadata is False + with monkeypatch.context() as m: m.setenv("CUDF_POLARS__PARQUET_OPTIONS__CHUNKED", "foo") engine = pl.GPUEngine() @@ -492,6 +501,47 @@ def test_validate_parquet_options(option: str) -> None: ) +def test_prefetch_file_metadata_default() -> None: + config = ConfigOptions.from_polars_engine(pl.GPUEngine(executor="streaming")) + assert isinstance(config.parquet_options.prefetch_file_metadata, Unspecified) + + config = ConfigOptions.from_polars_engine(pl.GPUEngine(executor="in-memory")) + assert config.parquet_options.prefetch_file_metadata is False + + config = ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", parquet_options={"prefetch_file_metadata": True} + ) + ) + assert config.parquet_options.prefetch_file_metadata is True + + +def test_parquet_options_object_passthrough() -> None: + parquet_options = ParquetOptions(prefetch_file_metadata=False) + config = ConfigOptions.from_polars_engine( + pl.GPUEngine(executor="streaming", parquet_options=parquet_options) + ) + assert config.parquet_options is parquet_options + + +def test_parquet_options_object_engine_default() -> None: + # If a user passes in a ParquetOptions object instead of a plain dict, and + # doesn't set prefetch_file_metadata on it, we still need to fill in the + # right default for the chosen executor. + parquet_options = ParquetOptions() + assert isinstance(parquet_options.prefetch_file_metadata, Unspecified) + + config = ConfigOptions.from_polars_engine( + pl.GPUEngine(executor="in-memory", parquet_options=parquet_options) + ) + assert config.parquet_options.prefetch_file_metadata is False + + config = ConfigOptions.from_polars_engine( + pl.GPUEngine(executor="streaming", parquet_options=parquet_options) + ) + assert isinstance(config.parquet_options.prefetch_file_metadata, Unspecified) + + def test_validate_raise_on_fail() -> None: with pytest.raises(TypeError, match="'raise_on_fail' must be"): ConfigOptions.from_polars_engine( diff --git a/python/cudf_polars/tests/test_select.py b/python/cudf_polars/tests/test_select.py index 435bec5d2031..07405d0d670e 100644 --- a/python/cudf_polars/tests/test_select.py +++ b/python/cudf_polars/tests/test_select.py @@ -176,18 +176,27 @@ def parquet_scan_row_bounds(request) -> dict[str, int | None]: return request.param -def test_get_parquet_row_count_from_metadata_raises() -> None: - paths = ["/some/missing/file.parquet"] +def test_get_parquet_row_count_from_metadata_no_cache_falls_back(tmp_path) -> None: + # If no cached parquet info is available (e.g. because prefetching was + # skipped for this path), we fall back to reading the metadata directly, + # rather than raising. + source = tmp_path / "data.parquet" + pl.DataFrame({"a": range(5)}).write_parquet(source) parquet_options = ParquetOptions(prefetch_file_metadata=True) - with pytest.raises(AssertionError, match=r"Cached parquet info is required"): - Scan._get_parquet_row_count_from_metadata( - paths, - skip_rows=0, - n_rows=-1, - parquet_options=parquet_options, - cached_parquet_info=None, - ) + row_count = Scan._get_parquet_row_count_from_metadata( + [str(source)], + skip_rows=0, + n_rows=-1, + parquet_options=parquet_options, + cached_parquet_info=None, + ) + assert row_count == 5 + + +def test_get_parquet_row_count_from_metadata_path_mismatch_raises() -> None: + paths = ["/some/missing/file.parquet"] + parquet_options = ParquetOptions(prefetch_file_metadata=True) with pytest.raises( AssertionError, diff --git a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx index dfdaabfe5b89..ed878c2647bb 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 & 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.memory cimport make_unique, unique_ptr @@ -26,6 +27,7 @@ from pylibcudf.libcudf.io.hybrid_scan cimport ( hybrid_scan_reader as cpp_hybrid_scan_reader, use_data_page_mask as cpp_use_data_page_mask, ) +from pylibcudf.libcudf.io.parquet_schema cimport FileMetaData as cpp_FileMetaData from pylibcudf.libcudf.io.text cimport byte_range_info from pylibcudf.libcudf.io.types cimport table_with_metadata from pylibcudf.libcudf.types cimport size_type @@ -115,7 +117,7 @@ cdef class HybridScanReader: """ cdef HybridScanReader reader = HybridScanReader.__new__(HybridScanReader) reader.c_obj = make_unique[cpp_hybrid_scan_reader]( - metadata.c_obj, + dereference(metadata.c_obj), options.c_obj ) return reader @@ -128,7 +130,12 @@ cdef class HybridScanReader: FileMetaData Parquet file footer metadata """ - return c_FileMetaData.from_cpp(self.c_obj.get()[0].parquet_metadata()) + cdef unique_ptr[cpp_FileMetaData] metadata + with nogil: + metadata = make_unique[cpp_FileMetaData]( + self.c_obj.get()[0].parquet_metadata() + ) + return c_FileMetaData.from_libcudf(move(metadata)) def page_index_byte_range(self) -> ByteRangeInfo: """Get the byte range of the page index. diff --git a/python/pylibcudf/pylibcudf/io/parquet.pyx b/python/pylibcudf/pylibcudf/io/parquet.pyx index 109454b95120..8ed8ed94981e 100644 --- a/python/pylibcudf/pylibcudf/io/parquet.pyx +++ b/python/pylibcudf/pylibcudf/io/parquet.pyx @@ -82,23 +82,27 @@ def _warn_deprecated(api_name, new_api): ) -cdef vector[cpp_FileMetaData] _build_parquet_metadatas( +cdef vector[cpp_FileMetaData*] _parquet_metadata_ptrs( object parquet_metadatas, size_t num_sources, ) except *: - cdef vector[cpp_FileMetaData] c_metadatas + """Validate Python FileMetaData list and collect non-owning C++ pointers. + + The expensive deep clone must happen later under the same ``nogil`` block as + ``read_parquet`` / the chunked reader ctor. Returning ``vector[FileMetaData]`` + from a cdef helper copies again with the GIL held (no libcudf NVTX). + """ cdef vector[cpp_FileMetaData*] metadata_ptrs cdef object metadata - cdef size_t i if parquet_metadatas is None: - return c_metadatas + return metadata_ptrs for metadata in parquet_metadatas: if not isinstance(metadata, FileMetaData): raise TypeError( "parquet_metadatas must contain only FileMetaData objects" ) - metadata_ptrs.push_back(&(metadata).c_obj) + metadata_ptrs.push_back((metadata).c_obj.get()) if metadata_ptrs.size() != num_sources: raise ValueError( @@ -107,14 +111,7 @@ cdef vector[cpp_FileMetaData] _build_parquet_metadatas( f"({num_sources})" ) - c_metadatas.reserve(metadata_ptrs.size()) - with nogil: - # This copies the (potentially large) metadata object. We don't - # want to hold the GIL for that. - for i in range(metadata_ptrs.size()): - c_metadatas.push_back(dereference(metadata_ptrs[i])) - - return c_metadatas + return metadata_ptrs cdef class ParquetReaderOptions: @@ -617,7 +614,9 @@ cdef class ChunkedParquetReader: self.mr = _get_memory_resource(mr) cdef vector[unique_ptr[datasource]] sources cdef vector[cpp_FileMetaData] c_metadatas + cdef vector[cpp_FileMetaData*] metadata_ptrs cdef cudaStream_t stream_view = self._stream.view().value() + cdef size_t i if parquet_metadatas is None: with nogil: self.reader.reset( @@ -632,10 +631,16 @@ cdef class ChunkedParquetReader: else: with nogil: sources = make_datasources(options.c_obj.get_source()) - c_metadatas = _build_parquet_metadatas( - parquet_metadatas, sources.size() + # Pin wrappers for the nogil clone; do not rely on the caller's + # mutable parquet_metadatas container remaining unchanged. + metadata_holders = tuple(parquet_metadatas) + metadata_ptrs = _parquet_metadata_ptrs( + metadata_holders, sources.size() ) with nogil: + c_metadatas.reserve(metadata_ptrs.size()) + for i in range(metadata_ptrs.size()): + c_metadatas.push_back(dereference(metadata_ptrs[i])) self.reader.reset( new cpp_chunked_parquet_reader( chunk_read_limit, @@ -716,16 +721,26 @@ cpdef TableWithMetadata read_parquet( cdef cudaStream_t _cs = s.view().value() cdef vector[unique_ptr[datasource]] sources cdef vector[cpp_FileMetaData] c_metadatas + cdef vector[cpp_FileMetaData*] metadata_ptrs cdef table_with_metadata c_result + cdef size_t i mr = _get_memory_resource(mr) if parquet_metadatas is None: with nogil: c_result = move(cpp_read_parquet(options.c_obj, _cs, mr.get_mr())) else: + # Collect pointers under GIL; clone + read must share one nogil block so + # Cython does not deep-copy vector[FileMetaData] while holding the GIL. with nogil: sources = make_datasources(options.c_obj.get_source()) - c_metadatas = _build_parquet_metadatas(parquet_metadatas, sources.size()) + # Pin wrappers for the nogil clone; do not rely on the caller's + # mutable parquet_metadatas container remaining unchanged. + metadata_holders = tuple(parquet_metadatas) + metadata_ptrs = _parquet_metadata_ptrs(metadata_holders, sources.size()) with nogil: + c_metadatas.reserve(metadata_ptrs.size()) + for i in range(metadata_ptrs.size()): + c_metadatas.push_back(dereference(metadata_ptrs[i])) c_result = move( cpp_read_parquet( move(sources), diff --git a/python/pylibcudf/pylibcudf/io/parquet_metadata.pxd b/python/pylibcudf/pylibcudf/io/parquet_metadata.pxd index 67fdda5d6907..fedbca1800f2 100644 --- a/python/pylibcudf/pylibcudf/io/parquet_metadata.pxd +++ b/python/pylibcudf/pylibcudf/io/parquet_metadata.pxd @@ -1,6 +1,8 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from libcpp.memory cimport unique_ptr + from pylibcudf.io.types cimport SourceInfo from pylibcudf.libcudf.io.parquet_schema cimport ( ColumnChunk as cpp_ColumnChunk, @@ -65,10 +67,10 @@ cdef class ParquetMetadata: cpdef dict columnchunk_metadata(self) cdef class FileMetaData: - cdef cpp_FileMetaData c_obj + cdef unique_ptr[cpp_FileMetaData] c_obj @staticmethod - cdef FileMetaData from_cpp(cpp_FileMetaData metadata) + cdef FileMetaData from_libcudf(unique_ptr[cpp_FileMetaData] metadata) cdef class SortingColumn: cdef cpp_SortingColumn c_obj diff --git a/python/pylibcudf/pylibcudf/io/parquet_metadata.pyx b/python/pylibcudf/pylibcudf/io/parquet_metadata.pyx index 8a4362153eb6..e6015786173b 100644 --- a/python/pylibcudf/pylibcudf/io/parquet_metadata.pyx +++ b/python/pylibcudf/pylibcudf/io/parquet_metadata.pyx @@ -1,9 +1,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from cython.operator cimport dereference from libc.stdint cimport uint8_t from libcpp.memory cimport make_unique, unique_ptr from libcpp.string cimport string +from libcpp.utility cimport move from libcpp.vector cimport vector from pylibcudf.io.types cimport SourceInfo @@ -472,31 +474,34 @@ cdef class FileMetaData: raise ValueError("FileMetaData cannot be constructed directly") @staticmethod - cdef FileMetaData from_cpp(cpp_FileMetaData metadata): + cdef FileMetaData from_libcudf(unique_ptr[cpp_FileMetaData] metadata): cdef FileMetaData result = FileMetaData.__new__(FileMetaData) - result.c_obj = metadata + result.c_obj = move(metadata) return result @property def version(self) -> int: """Get the file format version.""" - return self.c_obj.version + return dereference(self.c_obj).version @property def num_rows(self) -> int: """Get the total number of rows.""" - return self.c_obj.num_rows + return dereference(self.c_obj).num_rows @property def created_by(self) -> str: """Get the application that created the file.""" - return self.c_obj.created_by.decode("utf-8") + return dereference(self.c_obj).created_by.decode("utf-8") @property def row_groups(self) -> list[RowGroup]: """Get row group metadata in this file.""" cdef cpp_RowGroup row_group - return [RowGroup.from_cpp(row_group) for row_group in self.c_obj.row_groups] + return [ + RowGroup.from_cpp(row_group) + for row_group in dereference(self.c_obj).row_groups + ] @property def row_group_num_rows(self) -> list[int]: @@ -517,8 +522,10 @@ cdef class FileMetaData: >>> [rg.num_rows for rg in file_metadata.row_groups] """ cdef Py_ssize_t i - cdef Py_ssize_t n = self.c_obj.row_groups.size() - return [self.c_obj.row_groups[i].num_rows for i in range(n)] + cdef Py_ssize_t n = dereference(self.c_obj).row_groups.size() + return [ + dereference(self.c_obj).row_groups[i].num_rows for i in range(n) + ] @property def columnchunk_metadata(self) -> dict[str, list[int]]: @@ -549,21 +556,33 @@ cdef class FileMetaData: ... ) """ cdef Py_ssize_t i, j, k, n_path, n_col - cdef Py_ssize_t n_rg = self.c_obj.row_groups.size() + cdef Py_ssize_t n_rg = dereference(self.c_obj).row_groups.size() cdef dict result = {} cdef str name cdef list path_parts for i in range(n_rg): - n_col = self.c_obj.row_groups[i].columns.size() + n_col = dereference(self.c_obj).row_groups[i].columns.size() for j in range(n_col): - n_path = self.c_obj.row_groups[i].columns[j].meta_data.path_in_schema.size() + n_path = ( + dereference(self.c_obj) + .row_groups[i] + .columns[j] + .meta_data.path_in_schema.size() + ) path_parts = [ - self.c_obj.row_groups[i].columns[j].meta_data.path_in_schema[k].decode("utf-8") + dereference(self.c_obj) + .row_groups[i] + .columns[j] + .meta_data.path_in_schema[k] + .decode("utf-8") for k in range(n_path) ] name = ".".join(path_parts) result.setdefault(name, []).append( - self.c_obj.row_groups[i].columns[j].meta_data.total_uncompressed_size + dereference(self.c_obj) + .row_groups[i] + .columns[j] + .meta_data.total_uncompressed_size ) return result @@ -590,7 +609,7 @@ cdef class FileMetaData: """ cdef parquet_reader_options options = parquet_reader_options() cdef unique_ptr[cpp_hybrid_scan_reader] reader - cdef cpp_FileMetaData metadata + cdef unique_ptr[cpp_FileMetaData] metadata cdef const uint8_t* footer_ptr = 0 if len(footer_bytes) > 0: @@ -601,9 +620,11 @@ cdef class FileMetaData: host_span[const_uint8_t](footer_ptr, len(footer_bytes)), options, ) - metadata = reader.get()[0].parquet_metadata() + metadata = make_unique[cpp_FileMetaData]( + reader.get()[0].parquet_metadata() + ) - return FileMetaData.from_cpp(metadata) + return FileMetaData.from_libcudf(move(metadata)) cpdef ParquetMetadata read_parquet_metadata(SourceInfo src_info): @@ -651,7 +672,8 @@ cpdef list read_parquet_footers(SourceInfo src_info): """ cdef vector[unique_ptr[datasource]] sources cdef vector[cpp_FileMetaData] c_result - cdef cpp_FileMetaData metadata + cdef vector[unique_ptr[cpp_FileMetaData]] owned + cdef size_t i, n with nogil: sources = make_datasources(src_info.c_obj) c_result = cpp_parquet_metadata.read_parquet_footers( @@ -660,5 +682,12 @@ cpdef list read_parquet_footers(SourceInfo src_info): sources.size(), ) ) + n = c_result.size() + owned.reserve(n) + for i in range(n): + owned.push_back( + move(make_unique[cpp_FileMetaData](move(c_result[i]))) + ) - return [FileMetaData.from_cpp(metadata) for metadata in c_result] + # GIL held only for Python object allocation + list build + return [FileMetaData.from_libcudf(move(owned[i])) for i in range(n)]