diff --git a/python/pylibcudf/pylibcudf/io/parquet.pxd b/python/pylibcudf/pylibcudf/io/parquet.pxd index fd3cc3f16676..58b186647286 100644 --- a/python/pylibcudf/pylibcudf/io/parquet.pxd +++ b/python/pylibcudf/pylibcudf/io/parquet.pxd @@ -1,4 +1,4 @@ -# 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 libc.stdint cimport int64_t, uint8_t @@ -49,6 +49,7 @@ cdef class ParquetReaderOptions: cpdef void set_columns(self, list col_names) cpdef void set_column_names(self, list col_names) cpdef void set_column_indices(self, list col_indices) + cpdef void set_column_field_ids(self, list column_field_ids) cpdef void set_filter(self, Expression filter) cpdef void set_source(self, SourceInfo src) cpdef bool is_enabled_use_jit_filter(self) @@ -68,6 +69,7 @@ cdef class ParquetReaderOptionsBuilder: cpdef ParquetReaderOptionsBuilder columns(self, list col_names) cpdef ParquetReaderOptionsBuilder column_names(self, list col_names) cpdef ParquetReaderOptionsBuilder column_indices(self, list col_indices) + cpdef ParquetReaderOptionsBuilder column_field_ids(self, list column_field_ids) cpdef ParquetReaderOptionsBuilder use_jit_filter(self, bool use_jit_filter) cpdef ParquetReaderOptionsBuilder case_sensitive_names(self, bool val) cpdef ParquetReaderOptionsBuilder decimal_width(self, type_id width) diff --git a/python/pylibcudf/pylibcudf/io/parquet.pyi b/python/pylibcudf/pylibcudf/io/parquet.pyi index a79e9cf3f5b1..8564b270e583 100644 --- a/python/pylibcudf/pylibcudf/io/parquet.pyi +++ b/python/pylibcudf/pylibcudf/io/parquet.pyi @@ -1,4 +1,4 @@ -# 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 collections.abc import Mapping, Sequence @@ -30,6 +30,7 @@ class ParquetReaderOptions: def set_columns(self, col_names: list[str]): ... def set_column_names(self, col_names: list[str]): ... def set_column_indices(self, col_indices: list[int]): ... + def set_column_field_ids(self, column_field_ids: list[int]): ... def set_filter(self, filter: Expression): ... def set_source(self, src: SourceInfo) -> None: ... def is_enabled_use_jit_filter(self) -> bool: ... @@ -45,6 +46,10 @@ class ParquetReaderOptionsBuilder: def allow_mismatched_pq_schemas(self, val: bool) -> Self: ... def ignore_missing_columns(self, val: bool) -> Self: ... def use_arrow_schema(self, val: bool) -> Self: ... + def columns(self, col_names: list[str]) -> Self: ... + def column_names(self, col_names: list[str]) -> Self: ... + def column_indices(self, col_indices: list[int]) -> Self: ... + def column_field_ids(self, column_field_ids: list[int]) -> Self: ... def use_jit_filter(self, use_jit_filter: bool) -> Self: ... def case_sensitive_names(self, val: bool) -> Self: ... def decimal_width(self, width: TypeId) -> Self: ... diff --git a/python/pylibcudf/pylibcudf/io/parquet.pyx b/python/pylibcudf/pylibcudf/io/parquet.pyx index fdffb02b0f32..d287d10900ec 100644 --- a/python/pylibcudf/pylibcudf/io/parquet.pyx +++ b/python/pylibcudf/pylibcudf/io/parquet.pyx @@ -3,7 +3,7 @@ from cython.operator cimport dereference import warnings -from libc.stdint cimport int64_t, uint8_t +from libc.stdint cimport int32_t, int64_t, uint8_t from libcpp cimport bool from libcpp.memory cimport unique_ptr, make_unique @@ -251,7 +251,7 @@ cdef class ParquetReaderOptions: Parameters ---------- - col_names : list + col_indices : list List of top-level column indices Returns @@ -263,6 +263,24 @@ cdef class ParquetReaderOptions: vec.push_back(idx) self.c_obj.set_column_indices(vec) + cpdef void set_column_field_ids(self, list column_field_ids): + """ + Sets Parquet field IDs of the columns/fields to be read. + + Parameters + ---------- + column_field_ids : list + List of Parquet field IDs + + Returns + ------- + None + """ + cdef vector[int32_t] vec + for field_id in column_field_ids: + vec.push_back(field_id) + self.c_obj.set_column_field_ids(vec) + cpdef void set_filter(self, Expression filter): """ Sets AST based filter for predicate pushdown. @@ -465,7 +483,7 @@ cdef class ParquetReaderOptionsBuilder: Parameters ---------- - col_names : list[int] + col_indices : list[int] List of top-level column indices Returns @@ -478,6 +496,25 @@ cdef class ParquetReaderOptionsBuilder: self.c_obj.column_indices(vec) return self + cpdef ParquetReaderOptionsBuilder column_field_ids(self, list column_field_ids): + """ + Sets Parquet field IDs of the columns/fields to be read. + + Parameters + ---------- + column_field_ids : list[int] + List of Parquet field IDs + + Returns + ------- + ParquetReaderOptionsBuilder + """ + cdef vector[int32_t] vec + for field_id in column_field_ids: + vec.push_back(field_id) + self.c_obj.column_field_ids(vec) + return self + cpdef ParquetReaderOptionsBuilder use_jit_filter(self, bool use_jit_filter): """ Sets whether to use JIT compilation for filtering. diff --git a/python/pylibcudf/pylibcudf/libcudf/io/parquet.pxd b/python/pylibcudf/pylibcudf/libcudf/io/parquet.pxd index a6a1ca3e0f91..0f5a531dea8a 100644 --- a/python/pylibcudf/pylibcudf/libcudf/io/parquet.pxd +++ b/python/pylibcudf/pylibcudf/libcudf/io/parquet.pxd @@ -1,6 +1,6 @@ -# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from libc.stdint cimport int64_t, uint8_t +from libc.stdint cimport int32_t, int64_t, uint8_t from libcpp cimport bool from libcpp.functional cimport reference_wrapper from libcpp.map cimport map @@ -53,6 +53,9 @@ cdef extern from "cudf/io/parquet.hpp" namespace "cudf::io" nogil: void set_column_indices( vector[size_type] col_indices ) except +libcudf_exception_handler + void set_column_field_ids( + vector[int32_t] column_field_ids + ) except +libcudf_exception_handler void set_num_rows(int64_t val) except +libcudf_exception_handler void set_row_groups( vector[vector[size_type]] row_grp @@ -89,6 +92,9 @@ cdef extern from "cudf/io/parquet.hpp" namespace "cudf::io" nogil: parquet_reader_options_builder& column_indices( vector[size_type] col_indices ) except +libcudf_exception_handler + parquet_reader_options_builder& column_field_ids( + vector[int32_t] column_field_ids + ) except +libcudf_exception_handler parquet_reader_options_builder& row_groups( vector[vector[size_type]] row_grp ) except +libcudf_exception_handler diff --git a/python/pylibcudf/tests/io/test_parquet.py b/python/pylibcudf/tests/io/test_parquet.py index 88ba587d8e88..a2ba752dbba1 100644 --- a/python/pylibcudf/tests/io/test_parquet.py +++ b/python/pylibcudf/tests/io/test_parquet.py @@ -1,4 +1,4 @@ -# 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 import io import os @@ -106,6 +106,53 @@ def test_read_parquet_basic( assert res.num_row_groups_after_bloom_filter is None +def test_read_parquet_column_field_ids(binary_source_or_sink): + schema = pa.schema( + [ + pa.field( + "col_int64", + pa.int64(), + metadata={b"PARQUET:field_id": b"10"}, + ), + pa.field( + "col_string", + pa.string(), + metadata={b"PARQUET:field_id": b"20"}, + ), + pa.field( + "col_bool", + pa.bool_(), + metadata={b"PARQUET:field_id": b"30"}, + ), + ] + ) + pa_table = pa.Table.from_arrays( + [ + pa.array([1, 2, 3], type=pa.int64()), + pa.array(["a", "b", "c"], type=pa.string()), + pa.array([True, False, True], type=pa.bool_()), + ], + schema=schema, + ) + source = make_source( + binary_source_or_sink, pa_table, **_COMMON_PARQUET_SOURCE_KWARGS + ) + source_info = plc.io.SourceInfo([source]) + options = ( + plc.io.parquet.ParquetReaderOptions.builder(source_info) + .column_field_ids([30, 10]) + .build() + ) + + res = plc.io.parquet.read_parquet(options) + + assert_table_and_meta_eq( + pa_table.select(["col_bool", "col_int64"]), + res, + check_field_nullability=False, + ) + + @pytest.mark.parametrize("if_prune_rowgroup,result", [(True, 0), (False, 1)]) def test_read_parquet_filters_metadata(tmp_path, if_prune_rowgroup, result): col_list = list(range(1, 10))