diff --git a/docs/cudf/source/conf.py b/docs/cudf/source/conf.py index bf89f727e793..e100ae6f96e3 100644 --- a/docs/cudf/source/conf.py +++ b/docs/cudf/source/conf.py @@ -693,6 +693,7 @@ def on_missing_reference(app, env, node, contnode): ("py:class", "SupportsArrayInterface"), ("py:class", "SupportsCudaArrayInterface"), ("py:class", "T"), + ("py:class", "Buffer"), ] # Temporarily disable nitpick warnings for pandas: https://github.com/pandas-dev/pandas/issues/64584 nitpick_ignore_regex = [ diff --git a/python/pylibcudf/pylibcudf/io/experimental/CMakeLists.txt b/python/pylibcudf/pylibcudf/io/experimental/CMakeLists.txt index 4fd71e743a61..fb1891a7fc2a 100644 --- a/python/pylibcudf/pylibcudf/io/experimental/CMakeLists.txt +++ b/python/pylibcudf/pylibcudf/io/experimental/CMakeLists.txt @@ -1,11 +1,11 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= -set(cython_sources hybrid_scan.pyx) +set(cython_sources hybrid_scan.pyx hybrid_scan_multifile.pyx) set(linked_libraries cudf::cudf) rapids_cython_create_modules( diff --git a/python/pylibcudf/pylibcudf/io/experimental/__init__.pxd b/python/pylibcudf/pylibcudf/io/experimental/__init__.pxd index 87cc217ebf94..c84f478e414f 100644 --- a/python/pylibcudf/pylibcudf/io/experimental/__init__.pxd +++ b/python/pylibcudf/pylibcudf/io/experimental/__init__.pxd @@ -1,7 +1,8 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, 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 cimport ( FileMetaData, HybridScanReader, ) +from pylibcudf.io.experimental.hybrid_scan_multifile cimport HybridScanMultiFile diff --git a/python/pylibcudf/pylibcudf/io/experimental/__init__.py b/python/pylibcudf/pylibcudf/io/experimental/__init__.py index ef1ca25b7cff..29b85c09c8b0 100644 --- a/python/pylibcudf/pylibcudf/io/experimental/__init__.py +++ b/python/pylibcudf/pylibcudf/io/experimental/__init__.py @@ -6,11 +6,13 @@ HybridScanReader, UseDataPageMask, ) +from pylibcudf.io.experimental.hybrid_scan_multifile import HybridScanMultiFile from pylibcudf.io.parquet_metadata import FileMetaData __all__ = [ "FileMetaData", # backwards compatibility "HybridScanMetadata", + "HybridScanMultiFile", "HybridScanReader", "UseDataPageMask", ] diff --git a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan_multifile.pxd b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan_multifile.pxd new file mode 100644 index 000000000000..d41397ef14d9 --- /dev/null +++ b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan_multifile.pxd @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from libcpp.memory cimport unique_ptr +from libcpp.vector cimport vector + +from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource +from rmm.pylibrmm.stream cimport Stream + +from pylibcudf.libcudf.io.hybrid_scan_multifile cimport ( + hybrid_scan_multifile as cpp_hybrid_scan_multifile, +) +from pylibcudf.libcudf.types cimport size_type + + +cdef vector[vector[size_type]] _get_row_group_indices(object row_group_indices) except * + + +cdef class HybridScanMultiFile: + cdef unique_ptr[cpp_hybrid_scan_multifile] c_obj + cdef Stream _stream + cdef DeviceMemoryResource mr + cdef object _payload_page_data diff --git a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan_multifile.pyi b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan_multifile.pyi new file mode 100644 index 000000000000..b3593c52f714 --- /dev/null +++ b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan_multifile.pyi @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from collections.abc import Sequence + +from rmm.pylibrmm.memory_resource import DeviceMemoryResource + +from pylibcudf.column import Column +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: + from collections.abc import Buffer +except ImportError: + from typing_extensions import Buffer + +class HybridScanMultiFile: + @staticmethod + def from_parquet_metadatas( + parquet_metadatas: Sequence[FileMetaData], + options: ParquetReaderOptions, + ) -> HybridScanMultiFile: ... + def parquet_metadatas(self) -> list[FileMetaData]: ... + def page_index_byte_ranges(self) -> list[ByteRangeInfo]: ... + def setup_page_indexes( + self, page_index_bytes: Sequence[Buffer] + ) -> None: ... + def total_rows_in_row_groups( + self, row_group_indices: list[list[int]] + ) -> int: ... + def payload_pages_byte_ranges( + self, + row_group_indices: list[list[int]], + row_mask: Column, + options: ParquetReaderOptions, + stream: CudaStreamLike | None = None, + ) -> tuple[list[ByteRangeInfo], list[int]]: ... + def setup_chunking_for_payload_columns( + self, + chunk_read_limit: int, + pass_read_limit: int, + row_group_indices: list[list[int]], + row_mask: Column, + page_data: Sequence[Span | None], + options: ParquetReaderOptions, + stream: CudaStreamLike | None = None, + mr: DeviceMemoryResource | None = None, + ) -> None: ... + def materialize_payload_columns_chunk( + self, + row_mask: Column, + ) -> TableWithMetadata: ... + def construct_row_group_passes( + self, + row_group_indices: list[list[int]], + pass_read_limit: int, + ) -> list[list[list[int]]]: ... + def has_next_table_chunk(self) -> bool: ... diff --git a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan_multifile.pyx b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan_multifile.pyx new file mode 100644 index 000000000000..7f7f72e01e0b --- /dev/null +++ b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan_multifile.pyx @@ -0,0 +1,410 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from libc.stdint cimport uint8_t +from libc.stddef cimport size_t +from libcpp cimport bool +from libcpp.memory cimport make_unique +from libcpp.pair cimport pair +from libcpp.utility cimport move +from libcpp.vector cimport vector + +from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource +from rmm.pylibrmm.stream cimport Stream + +from pylibcudf.column cimport Column +from pylibcudf.io.experimental.hybrid_scan cimport _get_device_span +from pylibcudf.io.parquet cimport ParquetReaderOptions +from pylibcudf.io.parquet_metadata cimport FileMetaData as c_FileMetaData +from pylibcudf.io.text cimport ByteRangeInfo +from pylibcudf.io.types cimport TableWithMetadata +from pylibcudf.libcudf.column.column_view cimport column_view +from pylibcudf.libcudf.io.hybrid_scan cimport const_device_span_const_uint8_t +from pylibcudf.libcudf.io.hybrid_scan_multifile cimport ( + const_host_span_const_uint8_t, + const_uint8_t, + const_vector_size_type, + host_span_const_uint8_t, + hybrid_scan_multifile as cpp_hybrid_scan_multifile, +) +from pylibcudf.libcudf.io.parquet_metadata cimport const_FileMetaData +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 +from pylibcudf.libcudf.utilities.span cimport device_span, host_span +from pylibcudf.utils cimport _get_memory_resource, _get_stream +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Sequence + from typing_extensions import Buffer + from pylibcudf.typing import CudaStreamLike + +from pylibcudf.io.parquet_metadata import FileMetaData + +__all__ = ["HybridScanMultiFile"] + + +cdef vector[vector[size_type]] _get_row_group_indices( + object row_group_indices +) except *: + """Convert per-source row group indices to a vector of vectors.""" + cdef vector[vector[size_type]] indices + cdef vector[size_type] source_indices + for source in row_group_indices: + source_indices = source + indices.push_back(source_indices) + return indices + + +cdef class HybridScanMultiFile: + """Experimental multi-source Parquet reader for highly selective filters. + + Vectorizes the :class:`HybridScanReader` APIs over multiple Parquet sources. + Inputs and outputs are indexed by source order, except for the row mask which + is a single boolean column spanning all rows of all sources concatenated in + source order, then in row-group order within a source. + + For details, see + :cpp:class:`cudf::io::parquet::experimental::hybrid_scan_multifile` + + Examples + -------- + >>> import pylibcudf as plc + >>> reader = plc.io.experimental.HybridScanMultiFile.from_parquet_metadatas( + ... metadatas, options) + >>> byte_ranges, sources = reader.payload_pages_byte_ranges( + ... row_groups, row_mask, options) + """ + + def __init__(self): + raise ValueError( + "HybridScanMultiFile cannot be constructed directly. " + "Use from_parquet_metadatas()." + ) + + @staticmethod + def from_parquet_metadatas( + object parquet_metadatas: Sequence[FileMetaData], + ParquetReaderOptions options, + ) -> HybridScanMultiFile: + """Create a HybridScanMultiFile from pre-populated metadata. + + Parameters + ---------- + parquet_metadatas : Sequence[FileMetaData] + Pre-populated Parquet file metadata, one per source + options : ParquetReaderOptions + Parquet reader options + + Returns + ------- + HybridScanMultiFile + """ + cdef HybridScanMultiFile reader = HybridScanMultiFile.__new__( + HybridScanMultiFile + ) + cdef vector[cpp_FileMetaData] metadatas + cdef c_FileMetaData metadata + for metadata in parquet_metadatas: + metadatas.push_back(metadata.c_obj.get()[0]) + with nogil: + reader.c_obj = make_unique[cpp_hybrid_scan_multifile]( + host_span[const_FileMetaData]( + metadatas.data(), metadatas.size() + ), + options.c_obj + ) + return reader + + def parquet_metadatas(self) -> list[FileMetaData]: + """Get the Parquet file footer metadata of all sources. + + Returns + ------- + list[FileMetaData] + Parquet file footer metadata, one per source + """ + cdef vector[cpp_FileMetaData] c_result + with nogil: + c_result = move(self.c_obj.get()[0].parquet_metadatas()) + return [ + c_FileMetaData.from_libcudf( + make_unique[cpp_FileMetaData](move(c_result[source])) + ) + for source in range(c_result.size()) + ] + + def page_index_byte_ranges(self) -> list[ByteRangeInfo]: + """Get the byte range of the page index of all sources. + + Returns + ------- + list[ByteRangeInfo] + Byte range of the page index, one per source + """ + cdef vector[byte_range_info] ranges + with nogil: + ranges = move(self.c_obj.get()[0].page_index_byte_ranges()) + return [ByteRangeInfo(r.offset(), r.size()) for r in ranges] + + def setup_page_indexes( + self, object page_index_bytes: Sequence[Buffer] + ) -> None: + """Setup the page index within the Parquet file metadata of all sources. + + Parameters + ---------- + page_index_bytes : Sequence[Buffer] + Parquet page index buffer bytes, one per source + """ + cdef vector[host_span_const_uint8_t] spans + cdef const uint8_t[::1] page_index + # The spans point into page_index_bytes, which outlives this call + for page_index in page_index_bytes: + if len(page_index) == 0: + spans.push_back(host_span[const_uint8_t](0, 0)) + else: + spans.push_back( + host_span[const_uint8_t](&page_index[0], len(page_index)) + ) + with nogil: + self.c_obj.get()[0].setup_page_indexes( + host_span[const_host_span_const_uint8_t]( + spans.data(), spans.size() + ) + ) + + def total_rows_in_row_groups( + self, list row_group_indices: list[list[int]] + ) -> int: + """Get the total number of top-level rows in the row groups. + + Parameters + ---------- + row_group_indices : list[list[int]] + Input row group indices, one list per source + + Returns + ------- + int + Total number of top-level rows across all sources + """ + cdef vector[vector[size_type]] indices = _get_row_group_indices( + row_group_indices + ) + cdef size_type result + with nogil: + result = self.c_obj.get()[0].total_rows_in_row_groups( + host_span[const_vector_size_type]( + indices.data(), indices.size() + ) + ) + return result + + def payload_pages_byte_ranges( + self, + list row_group_indices: list[list[int]], + Column row_mask, + ParquetReaderOptions options, + object stream: CudaStreamLike | None = None + ) -> tuple[list[ByteRangeInfo], list[int]]: + """Get byte ranges of the pages of payload columns. + + Byte ranges are flattened in source, row group, column chunk, and page + order. Dictionary pages precede data pages within a column chunk, and + pruned pages are returned as empty byte ranges. + + Parameters + ---------- + row_group_indices : list[list[int]] + Input row group indices, one list per source + row_mask : Column + Boolean column indicating which rows need to be read + options : ParquetReaderOptions + Parquet reader options + stream : Stream, optional + CUDA stream + + Returns + ------- + tuple[list[ByteRangeInfo], list[int]] + Flattened byte ranges to the pages of payload columns and the source + index of each byte range + """ + cdef vector[vector[size_type]] indices = _get_row_group_indices( + row_group_indices + ) + cdef Stream _stream = _get_stream(stream) + cdef column_view mask_view = row_mask.view() + cdef pair[vector[byte_range_info], vector[size_type]] c_result + with nogil: + c_result = move(self.c_obj.get()[0].payload_pages_byte_ranges( + host_span[const_vector_size_type]( + indices.data(), indices.size() + ), + mask_view, + options.c_obj, + _stream.view().value() + )) + return ( + [ByteRangeInfo(r.offset(), r.size()) for r in c_result.first], + list(c_result.second), + ) + + def setup_chunking_for_payload_columns( + self, + size_t chunk_read_limit, + size_t pass_read_limit, + list row_group_indices: list[list[int]], + Column row_mask, + object page_data: Sequence, + ParquetReaderOptions options, + object stream: CudaStreamLike | None = None, + DeviceMemoryResource mr=None + ) -> None: + """Setup chunking information for payload columns read page by page. + + The data page mask is inferred from ``page_data``, which must have the + same shape as the byte ranges returned by + :meth:`payload_pages_byte_ranges`. + + Parameters + ---------- + chunk_read_limit : int + Limit on bytes returned per chunk (0 for no limit) + pass_read_limit : int + Limit on memory for reading/decompressing (0 for no limit) + row_group_indices : list[list[int]] + Input row group indices, one list per source + row_mask : Column + Boolean column indicating which rows need to be read + page_data : Sequence + Span-like objects containing the page data of payload columns, in + the same order as the byte ranges returned by + :meth:`payload_pages_byte_ranges`. ``None`` indicates a pruned page + options : ParquetReaderOptions + Parquet reader options + stream : Stream, optional + CUDA stream + mr : DeviceMemoryResource, optional + Device memory resource + """ + cdef vector[vector[size_type]] indices = _get_row_group_indices( + row_group_indices + ) + + cdef vector[device_span[const_uint8_t]] spans_vec + for page in page_data: + if page is None: + spans_vec.push_back(device_span[const_uint8_t]()) + else: + spans_vec.push_back(_get_device_span(page)) + + self._stream = _get_stream(stream) + self.mr = _get_memory_resource(mr) + # keep reference to avoid use-after-free of device spans + self._payload_page_data = page_data + + cdef column_view mask_view = row_mask.view() + with nogil: + self.c_obj.get()[0].setup_chunking_for_payload_columns( + chunk_read_limit, + pass_read_limit, + host_span[const_vector_size_type]( + indices.data(), indices.size() + ), + mask_view, + host_span[const_device_span_const_uint8_t]( + spans_vec.data(), + spans_vec.size() + ), + options.c_obj, + self._stream.view().value(), + self.mr.get_mr() + ) + + def materialize_payload_columns_chunk( + self, + Column row_mask, + ) -> TableWithMetadata: + """Materialize a chunk of payload columns. + + Parameters + ---------- + row_mask : Column + Boolean column indicating which rows need to be read + + Returns + ------- + TableWithMetadata + Table chunk of materialized payload columns and metadata + """ + cdef column_view mask_view = row_mask.view() + cdef table_with_metadata c_result + cdef bool more_chunks + with nogil: + c_result = move(self.c_obj.get()[0].materialize_payload_columns_chunk( + mask_view + )) + more_chunks = self.c_obj.get()[0].has_next_table_chunk() + if not more_chunks: + self._payload_page_data = None + return TableWithMetadata.from_libcudf(c_result, self._stream, self.mr) + + def construct_row_group_passes( + self, + list row_group_indices: list[list[int]], + size_t pass_read_limit, + ) -> list[list[list[int]]]: + """Partition row groups into passes such that the GPU memory required to + materialize a pass is bounded by the specified limit. + + Note that ``pass_read_limit`` is a hint, not an absolute limit. i.e. if + a row group cannot fit within the limit, it will still constitute a valid + pass. + + Parameters + ---------- + row_group_indices : list[list[int]] + Input row group indices, one list per source + pass_read_limit : int + Limit on the amount of memory used for reading and decompressing data + or 0 if there is no limit + + Returns + ------- + list[list[list[int]]] + Per-source row group indices, one list per pass + + Raises + ------ + ValueError + If ``row_group_indices`` is empty. + """ + cdef vector[vector[size_type]] indices = _get_row_group_indices( + row_group_indices + ) + cdef vector[vector[vector[size_type]]] passes + with nogil: + passes = move(self.c_obj.get()[0].construct_row_group_passes( + host_span[const_vector_size_type]( + indices.data(), indices.size() + ), + pass_read_limit + )) + return passes + + def has_next_table_chunk(self) -> bool: + """Check if there is any parquet data left to read. + + Returns + ------- + bool + True if there is data left to read + """ + cdef bool result + with nogil: + result = self.c_obj.get()[0].has_next_table_chunk() + return result diff --git a/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan_multifile.pxd b/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan_multifile.pxd new file mode 100644 index 000000000000..894a687cc321 --- /dev/null +++ b/python/pylibcudf/pylibcudf/libcudf/io/hybrid_scan_multifile.pxd @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from libcpp cimport bool +from libcpp.pair cimport pair +from libcpp.vector cimport vector +from pylibcudf.exception_handler cimport libcudf_exception_handler +from pylibcudf.libcudf.column.column_view cimport column_view +from pylibcudf.libcudf.io.hybrid_scan cimport ( + const_device_span_const_uint8_t, + const_uint8_t, +) +from pylibcudf.libcudf.io.parquet cimport parquet_reader_options +from pylibcudf.libcudf.io.parquet_metadata cimport const_FileMetaData +from pylibcudf.libcudf.io.parquet_schema cimport 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 +from pylibcudf.libcudf.utilities.span cimport host_span +from cuda.bindings.cyruntime cimport cudaStream_t +from rmm.librmm.memory_resource cimport device_async_resource_ref + +ctypedef const vector[size_type] const_vector_size_type +ctypedef host_span[const_uint8_t] host_span_const_uint8_t +ctypedef const host_span_const_uint8_t const_host_span_const_uint8_t + +cdef extern from "cudf/io/experimental/hybrid_scan_multifile.hpp" \ + namespace "cudf::io::parquet::experimental" nogil: + + cdef cppclass hybrid_scan_multifile: + hybrid_scan_multifile( + host_span[const_FileMetaData] parquet_metadata, + const parquet_reader_options& options + ) except +libcudf_exception_handler + + vector[FileMetaData] parquet_metadatas() except +libcudf_exception_handler + + vector[byte_range_info] page_index_byte_ranges() \ + except +libcudf_exception_handler + + void setup_page_indexes( + host_span[const_host_span_const_uint8_t] page_index_bytes + ) except +libcudf_exception_handler + + size_type total_rows_in_row_groups( + host_span[const_vector_size_type] row_group_indices + ) except +libcudf_exception_handler + + pair[vector[byte_range_info], vector[size_type]] payload_pages_byte_ranges( + host_span[const_vector_size_type] row_group_indices, + const column_view& row_mask, + const parquet_reader_options& options, + cudaStream_t stream + ) except +libcudf_exception_handler + + void setup_chunking_for_payload_columns( + size_t chunk_read_limit, + size_t pass_read_limit, + host_span[const_vector_size_type] row_group_indices, + const column_view& row_mask, + host_span[const_device_span_const_uint8_t] page_data, + const parquet_reader_options& options, + cudaStream_t stream, + device_async_resource_ref mr + ) except +libcudf_exception_handler + + table_with_metadata materialize_payload_columns_chunk( + const column_view& row_mask + ) except +libcudf_exception_handler + + vector[vector[vector[size_type]]] construct_row_group_passes( + host_span[const_vector_size_type] row_group_indices, + size_t pass_read_limit, + ) except +libcudf_exception_handler + + bool has_next_table_chunk() except +libcudf_exception_handler diff --git a/python/pylibcudf/tests/io/test_experimental_hybrid_scan_multifile.py b/python/pylibcudf/tests/io/test_experimental_hybrid_scan_multifile.py new file mode 100644 index 000000000000..ec9cab33e7b4 --- /dev/null +++ b/python/pylibcudf/tests/io/test_experimental_hybrid_scan_multifile.py @@ -0,0 +1,263 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import io + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +from utils import synchronize_stream + +import rmm +from rmm.pylibrmm.stream import Stream + +import pylibcudf as plc +from pylibcudf.io.experimental import HybridScanMultiFile + + +@pytest.fixture(scope="module") +def num_rows() -> int: + """Number of rows in each parquet source.""" + return 500 + + +@pytest.fixture(scope="module") +def num_row_groups() -> int: + """Number of row groups in each parquet source.""" + return 2 + + +@pytest.fixture(scope="module") +def parquet_table(num_rows: int) -> pa.Table: + """Create a simple PyArrow table for testing.""" + return pa.table( + { + "col0": pa.array(list(range(num_rows)), type=pa.uint32()), + "col1": [float(row) * 1.5 for row in range(num_rows)], + } + ) + + +@pytest.fixture(scope="module") +def parquet_bytes( + parquet_table: pa.Table, num_rows: int, num_row_groups: int +) -> list[bytes]: + """Create two identical parquet sources with a page index.""" + buf = io.BytesIO() + pq.write_table( + parquet_table, + buf, + row_group_size=num_rows // num_row_groups, + # Small, non-dictionary pages so a sparse row mask prunes some of them + data_page_size=256, + write_batch_size=32, + use_dictionary=False, + write_statistics=True, + write_page_index=True, + ) + return [buf.getvalue()] * 2 + + +def footer_bytes(parquet_bytes: bytes) -> memoryview: + """Extract the footer bytes of a parquet file. + + According to Parquet file format specification: + https://parquet.apache.org/docs/file-format/ + """ + PARQUET_FOOTER_SIZE_BYTES = 4 # Number of bytes encoding footer length + PARQUET_MAGIC_BYTES = 4 # Number of bytes for "PAR1" magic number + PARQUET_SUFFIX_BYTES = PARQUET_FOOTER_SIZE_BYTES + PARQUET_MAGIC_BYTES + + parquet_mv = memoryview(parquet_bytes) + + footer_size = int.from_bytes( + parquet_mv[-PARQUET_SUFFIX_BYTES:-PARQUET_MAGIC_BYTES], + byteorder="little", + ) + footer_start = len(parquet_mv) - PARQUET_SUFFIX_BYTES - footer_size + footer_end = len(parquet_mv) - PARQUET_SUFFIX_BYTES + return parquet_mv[footer_start:footer_end] + + +@pytest.fixture +def parquet_options() -> plc.io.parquet.ParquetReaderOptions: + """Create ParquetReaderOptions for the hybrid scan reader. + + The reader never reads through a datasource, so the options carry no source. + + Note: This is function-scoped (not module-scoped) because tests may modify + the options, so each test needs its own independent copy. + """ + return plc.io.parquet.ParquetReaderOptions.builder( + plc.io.SourceInfo([]) + ).build() + + +@pytest.fixture +def row_groups(num_row_groups: int) -> list[list[int]]: + """Row group indices of both parquet sources.""" + return [list(range(num_row_groups))] * 2 + + +@pytest.fixture +def hybrid_scan_multifile_reader( + parquet_bytes: list[bytes], + parquet_options: plc.io.parquet.ParquetReaderOptions, +) -> HybridScanMultiFile: + """Create a HybridScanMultiFile with the page index of both sources.""" + # Create the reader from the footer bytes of each source + reader = HybridScanMultiFile.from_parquet_metadatas( + [ + plc.io.parquet_metadata.FileMetaData.from_bytes( + footer_bytes(source) + ) + for source in parquet_bytes + ], + parquet_options, + ) + # Fetch the page index of each source and set it up within the metadata + reader.setup_page_indexes( + [ + memoryview(source)[ + byte_range.offset : byte_range.offset + byte_range.size + ] + for source, byte_range in zip( + parquet_bytes, + reader.page_index_byte_ranges(), + strict=True, + ) + ] + ) + return reader + + +def test_hybrid_scan_multifile_construct_directly_raises() -> None: + """Test that a HybridScanMultiFile cannot be constructed directly.""" + with pytest.raises(ValueError, match="cannot be constructed directly"): + HybridScanMultiFile() + + +def test_hybrid_scan_multifile_metadata( + hybrid_scan_multifile_reader: HybridScanMultiFile, + row_groups: list[list[int]], + num_rows: int, +) -> None: + """Test the metadata of a reader built from pre-populated metadata.""" + # One metadata object per source, in source order + assert [ + metadata.num_rows + for metadata in hybrid_scan_multifile_reader.parquet_metadatas() + ] == [num_rows, num_rows] + + # Row counts are totalled across all sources + assert ( + hybrid_scan_multifile_reader.total_rows_in_row_groups(row_groups) + == 2 * num_rows + ) + + # Every source was written with a page index + assert all( + byte_range.size > 0 + for byte_range in hybrid_scan_multifile_reader.page_index_byte_ranges() + ) + + +def test_hybrid_scan_multifile_construct_row_group_passes( + hybrid_scan_multifile_reader: HybridScanMultiFile, + row_groups: list[list[int]], +) -> None: + """Test partitioning the input row groups into passes.""" + # No read limit yields a single pass spanning all sources + assert hybrid_scan_multifile_reader.construct_row_group_passes( + row_groups, 0 + ) == [row_groups] + + # A tiny read limit splits the row groups across multiple passes + assert ( + len( + hybrid_scan_multifile_reader.construct_row_group_passes( + row_groups, 1 + ) + ) + > 1 + ) + + +@pytest.mark.parametrize("stream", [None, Stream()]) +def test_hybrid_scan_multifile_materialize_payload_pages( + parquet_bytes: list[bytes], + hybrid_scan_multifile_reader: HybridScanMultiFile, + parquet_options: plc.io.parquet.ParquetReaderOptions, + row_groups: list[list[int]], + parquet_table: pa.Table, + num_rows: int, + stream: Stream | None, +) -> None: + """Test reading payload columns page by page from multiple sources.""" + # Without a filter, all selected columns are payload columns + parquet_options.set_column_names(["col0", "col1"]) + + # Select a sparse set of rows across both sources + mask = [row % 100 == 0 for row in range(2 * num_rows)] + row_mask = plc.Column.from_arrow(pa.array(mask), stream=stream) + + # Get the byte ranges of the payload pages surviving the row mask + page_ranges, source_indices = ( + hybrid_scan_multifile_reader.payload_pages_byte_ranges( + row_groups, row_mask, parquet_options, stream + ) + ) + + # Byte ranges are flattened, with one source index per byte range + assert len(page_ranges) == len(source_indices) + assert set(source_indices) == set(range(len(parquet_bytes))) + + # Pruned pages are reported as empty byte ranges + empty_ranges = sum(byte_range.size == 0 for byte_range in page_ranges) + assert 0 < empty_ranges < len(page_ranges) + + # Fetch the surviving pages from their source, passing None for pruned ones + page_data = [ + None + if byte_range.size == 0 + else plc.gpumemoryview( + rmm.DeviceBuffer.to_device( + memoryview(parquet_bytes[source])[ + byte_range.offset : byte_range.offset + byte_range.size + ], + plc.utils._get_stream(stream), + ) + ) + for byte_range, source in zip(page_ranges, source_indices, strict=True) + ] + synchronize_stream(stream) + + # The data page mask is inferred from the fetched page data + hybrid_scan_multifile_reader.setup_chunking_for_payload_columns( + 1024, # chunk_read_limit + 0, # pass_read_limit + row_groups, + row_mask, + page_data, + parquet_options, + stream, + ) + + # Read the output chunks, applying the row mask to each + chunks = [] + while hybrid_scan_multifile_reader.has_next_table_chunk(): + chunks.append( + hybrid_scan_multifile_reader.materialize_payload_columns_chunk( + row_mask + ) + ) + synchronize_stream(stream) + + # The chunk read limit is small enough to split the output + assert len(chunks) > 1 + + # The chunks reassemble into the masked rows of both sources + result = pa.concat_tables([chunk.tbl.to_arrow() for chunk in chunks]) + expected = pa.concat_tables([parquet_table] * 2).filter(pa.array(mask)) + assert result.equals( + expected.rename_columns(result.schema.names).cast(result.schema) + )