diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 727d608c19bc..15ceb7fff93f 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -818,6 +818,7 @@ add_library( src/io/parquet/writer_impl.cu src/io/parquet/writer_impl_helpers.cpp src/io/parquet/decode_fixed.cu + src/io/parquet/decode_pruned_pages.cu src/io/statistics/orc_column_statistics.cu src/io/statistics/parquet_column_statistics.cu src/io/text/byte_range_info.cpp diff --git a/cpp/src/io/parquet/decode_fixed.cu b/cpp/src/io/parquet/decode_fixed.cu index 45a7aeb4ec6d..5e3122d9fa37 100644 --- a/cpp/src/io/parquet/decode_fixed.cu +++ b/cpp/src/io/parquet/decode_fixed.cu @@ -1015,20 +1015,12 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size_t, 8) if (!(BitAnd(pages[page_idx].kernel_mask, kernel_mask_t))) { return; } + // Exit early if the page is pruned + if (page_mask.size() > 0 and not page_mask[page_idx]) { return; } + // must come after the kernel mask check [[maybe_unused]] null_count_back_copier _{s, t}; - // Exit super early for simple types if the page does not need to be decoded - if constexpr (not has_lists_t and not has_strings_t and not has_nesting_t) { - if (not page_mask[page_idx]) { - pp->num_nulls = pp->nesting[0].batch_size; - pp->num_valids = 0; - // Set s->nesting info = nullptr to bypass `null_count_back_copier` at return - s->nesting_info = nullptr; - return; - } - } - // Setup local page info if (!setup_local_page_info(s, pp, @@ -1040,23 +1032,6 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size_t, 8) return; } - // Write list and/or string offsets and exit if the page does not need to be decoded - if (not page_mask[page_idx]) { - // Update offsets for all list depth levels - if constexpr (has_lists_t) { update_list_offsets_for_pruned_pages(s); } - // Update string offsets or write string sizes for small and large strings respectively - if constexpr (has_strings_t) { - update_string_offsets_for_pruned_pages( - s, initial_str_offsets, pages[page_idx]); - } - // Must be set after computing above list and string offsets - pp->num_nulls = pp->nesting[s->col.max_nesting_depth - 1].batch_size; - if constexpr (not has_lists_t) { pp->num_nulls -= s->first_row; } - pp->num_valids = 0; - - return; - } - bool const process_nulls = should_process_nulls(s); // shared buffer. all shared memory is suballocated out of here diff --git a/cpp/src/io/parquet/decode_pruned_pages.cu b/cpp/src/io/parquet/decode_pruned_pages.cu new file mode 100644 index 000000000000..c668ebdb8530 --- /dev/null +++ b/cpp/src/io/parquet/decode_pruned_pages.cu @@ -0,0 +1,116 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "parquet_gpu.hpp" + +#include + +#include +#include +#include + +namespace cudf::io::parquet::detail { + +namespace { + +auto constexpr block_size = 4 * cudf::detail::warp_size; + +/** + * @brief Initialize output entries for pruned string and list pages. + * + * String entries receive either the page's initial offset for small strings or a zero size for + * large strings. List offsets receive the start value of their child nesting level. These targeted + * writes avoid zero-initializing entire output buffers. + */ +CUDF_KERNEL void __launch_bounds__(block_size) + fill_pruned_offsets_kernel(device_span pages, + device_span chunks, + device_span page_mask, + device_span initial_str_offsets, + size_t skip_rows, + size_t num_rows) +{ + namespace cg = cooperative_groups; + + auto const block = cg::this_thread_block(); + auto const page_idx = cg::this_grid().block_rank(); + auto const t = static_cast(block.thread_rank()); + if (page_mask[page_idx]) { return; } + + auto const& page = pages[page_idx]; + auto const& chunk = chunks[page.chunk_idx]; + + if (chunk.column_data_base == nullptr) { return; } + + auto const is_list_col = chunk.max_level[level_type::REPETITION] != 0; + + // Write offsets for pruned non-list (flat) string columns. + if (not is_list_col and is_string_col(chunk)) { + auto data = static_cast(chunk.column_data_base[chunk.max_nesting_depth - 1]); + if (data == nullptr) { return; } + + auto const page_begin = chunk.start_row + page.chunk_row; + auto const page_end = page_begin + page.num_rows; + auto const read_end = skip_rows + num_rows; + auto const begin = cuda::std::max(page_begin, skip_rows); + auto const end = cuda::std::min(page_end, read_end); + if (begin >= end) { return; } + + // Large strings needs the first page string offset, including if the page was + // pruned. Record it here. + if (chunk.is_large_string_col and t == 0) { + auto const chunks_per_rowgroup = initial_str_offsets.size(); + auto const input_col_idx = page.chunk_idx % chunks_per_rowgroup; + cuda::atomic_ref initial_str_offset{ + initial_str_offsets[input_col_idx]}; + initial_str_offset.fetch_min(page.str_offset, cuda::std::memory_order_relaxed); + } + + // Write zeros for large strings and the page's initial offset otherwise. + auto const value = + chunk.is_large_string_col ? size_type{0} : static_cast(page.str_offset); + for (auto row = begin + t; row < end; row += block.size()) { + data[row - skip_rows] = value; + } + return; + } + + // Write offsets to list locations at each depth. + if (is_list_col and page.nesting != nullptr and page.nesting_decode != nullptr) { + for (auto depth = 0; depth < chunk.max_nesting_depth - 1; depth++) { + auto offsets = static_cast(chunk.column_data_base[depth]); + auto& nesting_info = page.nesting[depth]; + // Pruned list pages retain rows through the first list level but contribute no child values. + // The preprocessing pass computes the output range and child start value at every list depth. + if (nesting_info.type != type_id::LIST or offsets == nullptr) { continue; } + // Emit an offset for the current nesting level equal to current length of the next nesting + // level + auto const output_begin = page.nesting_decode[depth].page_start_value; + auto const offset = page.nesting_decode[depth + 1].page_start_value; + for (auto offset_idx = t; offset_idx < nesting_info.batch_size; + offset_idx += static_cast(block.size())) { + offsets[output_begin + offset_idx] = offset; + } + } + } +} + +} // namespace + +void fill_pruned_offsets(cudf::device_span pages, + cudf::device_span chunks, + cudf::device_span page_mask, + cudf::device_span initial_str_offsets, + size_t skip_rows, + size_t num_rows, + rmm::cuda_stream_view stream) +{ + CUDF_EXPECTS(pages.size() == page_mask.size(), "Page mask size does not match page count"); + fill_pruned_offsets_kernel<<>>( + pages, chunks, page_mask, initial_str_offsets, skip_rows, num_rows); + CUDF_CUDA_TRY(cudaGetLastError()); +} + +} // namespace cudf::io::parquet::detail diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_chunking.cu b/cpp/src/io/parquet/experimental/hybrid_scan_chunking.cu index 0f846261c35b..7c819d0f5657 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_chunking.cu +++ b/cpp/src/io/parquet/experimental/hybrid_scan_chunking.cu @@ -158,7 +158,7 @@ void hybrid_scan_reader_impl::setup_next_pass( // if we are doing subpass reading, generate more accurate num_row estimates for list columns. // this helps us to generate more accurate subpass splits. if (pass.has_compressed_data && _input_pass_read_limit != 0) { - if (_has_page_index) { + if (_has_offset_index) { generate_list_column_row_counts(is_estimate_row_counts::NO); } else { generate_list_column_row_counts(is_estimate_row_counts::YES); diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp index 98353b88f432..ffacc6796301 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp @@ -877,7 +877,7 @@ void hybrid_scan_reader_impl::reset_internal_state() _row_mask_offset = 0; _file_itm_data = file_intermediate_data{}; _file_preprocessed = false; - _has_page_index = false; + _has_offset_index = false; _pass_itm_data.reset(); _pass_page_mask.clear(); _subpass_page_mask.reset(); diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_preprocess.cu b/cpp/src/io/parquet/experimental/hybrid_scan_preprocess.cu index 64037e0c87df..acfbc2a0dac6 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_preprocess.cu +++ b/cpp/src/io/parquet/experimental/hybrid_scan_preprocess.cu @@ -57,7 +57,7 @@ void decode_dictionary_page_headers(cudf::detail::hostdevice_span 0 && not _file_itm_data.row_groups.empty() && not _input_columns.empty()) { @@ -183,13 +184,13 @@ void hybrid_scan_reader_impl::setup_compressed_data( pass.has_compressed_data = setup_column_chunks(column_chunk_data); // Process dataset chunk pages into output columns - auto const total_pages = _has_page_index ? count_page_headers_with_pgidx(chunks, _stream) - : count_page_headers(chunks, _stream); + auto const total_pages = _has_offset_index ? count_page_headers_with_pgidx(chunks, _stream) + : count_page_headers(chunks, _stream); if (total_pages <= 0) { return; } rmm::device_uvector unsorted_pages(total_pages, _stream); // decoding of column/page information - parquet::detail::decode_page_headers(pass, unsorted_pages, _has_page_index, _stream); + parquet::detail::decode_page_headers(pass, unsorted_pages, _has_offset_index, _stream); CUDF_EXPECTS(pass.page_offsets.size() - 1 == static_cast(_input_columns.size()), "Encountered page_offsets / num_columns mismatch"); } diff --git a/cpp/src/io/parquet/page_data.cu b/cpp/src/io/parquet/page_data.cu index 5b2f854fce08..e5737b790b3c 100644 --- a/cpp/src/io/parquet/page_data.cu +++ b/cpp/src/io/parquet/page_data.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -61,6 +61,9 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size) auto const block = cg::this_thread_block(); auto const warp = cg::tiled_partition(block); + // Exit early if the page is pruned + if (not page_mask.empty() and not page_mask[page_idx]) { return; } + [[maybe_unused]] null_count_back_copier _{s, static_cast(block.thread_rank())}; // Setup local page info @@ -78,21 +81,6 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size) bool const has_repetition = s->col.max_level[level_type::REPETITION] > 0; bool const process_nulls = should_process_nulls(s); - // Write list offsets and exit if the page does not need to be decoded - if (not page_mask[page_idx]) { - auto& page = pages[page_idx]; - // Update offsets for all list depth levels - if (has_repetition) { update_list_offsets_for_pruned_pages(s); } - - // Must be set after computing above list offsets - cg::invoke_one(block, [&]() { - page.num_nulls = page.nesting[s->col.max_nesting_depth - 1].batch_size; - page.num_nulls -= has_repetition ? 0 : s->first_row; - page.num_valids = 0; - }); - return; - } - auto const data_len = cuda::std::distance(s->data_start, s->data_end); auto const num_values = data_len / s->dtype_len_in; @@ -280,6 +268,10 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size) auto const block = cg::this_thread_block(); auto const warp = cg::tiled_partition(block); int out_warp_id; + + // Exit early if the page is pruned + if (not page_mask.empty() and not page_mask[page_idx]) { return; } + [[maybe_unused]] null_count_back_copier _{s, static_cast(block.thread_rank())}; // Setup local page info @@ -297,47 +289,6 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size) bool const has_repetition = s->col.max_level[level_type::REPETITION] > 0; bool const process_nulls = should_process_nulls(s); - // Write list offsets and exit if the page does not need to be decoded - if (not page_mask[page_idx]) { - auto& page = pages[page_idx]; - - // Update offsets for all list depth levels - if (has_repetition) { update_list_offsets_for_pruned_pages(s); } - - // Fill offsets with the initial `str_offset` to indicate empty strings for BYTE_ARRAY and - // FIXED_LEN_BYTE_ARRAY types. These types are now decoded by `decode_page_data_generic()` - // anyway so the following code should never be reached. Also note that this decoder does not - // handle large strings either and should eventually be removed. - Type const dtype = s->col.physical_type; - auto const is_decimal = - s->col.logical_type.has_value() and s->col.logical_type->type == LogicalType::DECIMAL; - if (dtype == Type::FIXED_LEN_BYTE_ARRAY or (dtype == Type::BYTE_ARRAY and not is_decimal)) { - // Initial string offset - auto const initial_value = page.str_offset; - - // We must use the batch size from the nesting info (the size of the page for this batch) - auto value_count = page.nesting[s->col.max_nesting_depth - 1].batch_size; - - // If no repetition we haven't calculated start/end bounds and instead just skipped - // values until we reach first_row. account for that here. - if (not has_repetition) { value_count -= s->first_row; } - - auto& ni = s->nesting_info[s->col.max_nesting_depth - 1]; - auto offptr = reinterpret_cast(ni.data_out); - - // Write the initial string offset at all positions to indicate empty strings - for (int idx = block.thread_rank(); idx < value_count; idx += block.size()) { - offptr[idx] = initial_value; - } - } - - page.num_nulls = page.nesting[s->col.max_nesting_depth - 1].batch_size; - page.num_nulls -= has_repetition ? 0 : s->first_row; - page.num_valids = 0; - - return; - } - PageNestingDecodeInfo* nesting_info_base = s->nesting_info; // Capture initial valid_map_offset before any processing that might modify it diff --git a/cpp/src/io/parquet/page_decode.cuh b/cpp/src/io/parquet/page_decode.cuh index ef0a8f9e502d..31a9f4a05d05 100644 --- a/cpp/src/io/parquet/page_decode.cuh +++ b/cpp/src/io/parquet/page_decode.cuh @@ -717,38 +717,6 @@ inline __device__ void get_nesting_bounds(int& start_depth, } } -/** - * @brief Updates nesting level offsets for pruned pages of a list column - * - * This function iterates through the nesting levels of a column and updates the offsets for a list - * column. The offset for the current nesting level equals the length of the next nesting level - * - * @tparam block_size The size of the block used for decoding. - * @param[in,out] state Pointer to page state containing column and nesting information. - */ -template -static __device__ void update_list_offsets_for_pruned_pages(page_state_s* state) -{ - int const max_depth = state->col.max_nesting_depth - 1; - bool const in_nesting_bounds = max_depth >= 0; - auto const tid = cg::this_thread_block().thread_rank(); - - // Iterate by depth and store offset(s) to the list location(s) - for (int depth = 0; depth < max_depth; depth++) { - auto& nesting_info = state->nesting_info[depth]; - // If we're -not- at a leaf column and we're within nesting/row bounds and we have a valid - // data_out pointer, it implies this is a list column, so emit an offset for the current nesting - // level equal to current length of the next nesting level - if (in_nesting_bounds and nesting_info.data_out != nullptr) { - auto const& next_nesting_info = state->nesting_info[depth + 1]; - auto const offset = next_nesting_info.page_start_value; - for (int idx = tid; idx < state->page.nesting[depth].batch_size; idx += block_size) { - (reinterpret_cast(nesting_info.data_out))[idx] = offset; - } - } - } -} - /** * @brief Process a batch of incoming repetition/definition level values and generate * validity, nested column offsets (where appropriate) and decoding indices. diff --git a/cpp/src/io/parquet/page_delta_decode.cu b/cpp/src/io/parquet/page_delta_decode.cu index 85bf383bb43f..7cb5fbd59d26 100644 --- a/cpp/src/io/parquet/page_delta_decode.cu +++ b/cpp/src/io/parquet/page_delta_decode.cu @@ -317,6 +317,10 @@ CUDF_KERNEL void __launch_bounds__(decode_delta_binary_block_size) auto const block = cg::this_thread_block(); auto const warp = cg::tiled_partition(block); auto* const db = &db_state; + + // Exit early if the page is pruned + if (page_mask.size() > 0 and not page_mask[page_idx]) { return; } + [[maybe_unused]] null_count_back_copier _{s, static_cast(block.thread_rank())}; // Setup local page info @@ -337,17 +341,6 @@ CUDF_KERNEL void __launch_bounds__(decode_delta_binary_block_size) // Capture initial valid_map_offset before any processing that might modify it int const init_valid_map_offset = s->nesting_info[s->col.max_nesting_depth - 1].valid_map_offset; - // Write list offsets and exit if the page does not need to be decoded - if (not page_mask[page_idx]) { - auto& page = pages[page_idx]; - // Update offsets for all list depth levels - if (has_repetition) { update_list_offsets_for_pruned_pages(s); } - page.num_nulls = page.nesting[s->col.max_nesting_depth - 1].batch_size; - page.num_nulls -= has_repetition ? 0 : s->first_row; - page.num_valids = 0; - return; - } - // copying logic from gpuDecodePageData. PageNestingDecodeInfo const* nesting_info_base = s->nesting_info; @@ -478,6 +471,7 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size) auto* const prefix_db = &db_state.prefixes; auto* const suffix_db = &db_state.suffixes; auto* const dba = &db_state; + if (page_mask.size() > 0 and not page_mask[page_idx]) { return; } [[maybe_unused]] null_count_back_copier _{s, static_cast(block.thread_rank())}; if (!setup_local_page_info(s, @@ -504,28 +498,6 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size) // Capture initial valid_map_offset before any processing that might modify it int const init_valid_map_offset = s->nesting_info[s->col.max_nesting_depth - 1].valid_map_offset; - // Write list/string offsets and exit if the page does not need to be decoded - if (not page_mask[page_idx]) { - auto page = &pages[page_idx]; - // Update list offsets and string offsets or sizes depending on the large-string property - if (has_repetition) { - // Update list offsets - update_list_offsets_for_pruned_pages(s); - // Update string offsets or sizes - update_string_offsets_for_pruned_pages( - s, initial_str_offsets, pages[page_idx]); - } else { - // Update string offsets or sizes - update_string_offsets_for_pruned_pages( - s, initial_str_offsets, pages[page_idx]); - } - page->num_nulls = page->nesting[s->col.max_nesting_depth - 1].batch_size; - page->num_nulls -= has_repetition ? 0 : s->first_row; - page->num_valids = 0; - - return; - } - // choose a character parallel string copy when the average string is longer than a warp auto const use_char_ll = (s->page.str_bytes / s->page.num_valids) > cudf::detail::warp_size; @@ -704,6 +676,7 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size) auto const block = cg::this_thread_block(); auto const warp = cg::tiled_partition(block); auto* const db = &db_state; + if (page_mask.size() > 0 and not page_mask[page_idx]) { return; } [[maybe_unused]] null_count_back_copier _{s, static_cast(block.thread_rank())}; auto const mask = decode_kernel_mask::DELTA_LENGTH_BA; @@ -731,28 +704,6 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size) // Capture initial valid_map_offset before any processing that might modify it int const init_valid_map_offset = s->nesting_info[s->col.max_nesting_depth - 1].valid_map_offset; - // Write list/string offsets and exit if the page does not need to be decoded - if (not page_mask[page_idx]) { - auto page = &pages[page_idx]; - // Update list offsets and string offsets or sizes depending on the large-string property - if (has_repetition) { - // Update list offsets - update_list_offsets_for_pruned_pages(s); - // Update string offsets or sizes - update_string_offsets_for_pruned_pages( - s, initial_str_offsets, pages[page_idx]); - } else { - // Update string offsets or sizes - update_string_offsets_for_pruned_pages( - s, initial_str_offsets, pages[page_idx]); - } - page->num_nulls = page->nesting[s->col.max_nesting_depth - 1].batch_size; - page->num_nulls -= has_repetition ? 0 : s->first_row; - page->num_valids = 0; - - return; - } - // copying logic from gpuDecodePageData. PageNestingDecodeInfo const* nesting_info_base = s->nesting_info; diff --git a/cpp/src/io/parquet/page_hdr.cu b/cpp/src/io/parquet/page_hdr.cu index 2565dab3ae31..b69c26a4d5eb 100644 --- a/cpp/src/io/parquet/page_hdr.cu +++ b/cpp/src/io/parquet/page_hdr.cu @@ -3,6 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include "cuda/std/__utility/cmp.h" #include "error.hpp" #include "io/utilities/block_utils.cuh" #include "parquet_gpu.hpp" @@ -530,7 +531,7 @@ void __forceinline__ __device__ zero_out_page_header_info(byte_stream_s* bs) CUDF_KERNEL void __launch_bounds__(decode_page_headers_block_size) decode_page_headers_kernel(device_span chunks, - chunk_page_info* chunk_pages, + device_span chunk_pages, kernel_error::pointer error_code) { auto constexpr num_warps_per_block = decode_page_headers_block_size / cudf::detail::warp_size; @@ -737,53 +738,55 @@ CUDF_KERNEL void __launch_bounds__(count_page_headers_block_size) } /** - * @brief Functor to decode page headers from specified page locations + * @brief Functor to decode specified page headers from corresponding page data spans */ -struct decode_page_headers_with_pgidx_fn { +struct decode_from_page_data_fn { cudf::device_span colchunks; cudf::device_span pages; - uint8_t** page_locations; - size_type* chunk_page_offsets; + cudf::device_span const> page_data; + cudf::device_span chunk_page_offsets; kernel_error::pointer error_code; __device__ void operator()(size_type page_idx) const noexcept { auto const num_chunks = static_cast(colchunks.size()); - // Binary search the the column chunk index for this page + // Binary search the column chunk index for this page auto const chunk_idx = static_cast( cuda::std::distance( - chunk_page_offsets, + chunk_page_offsets.begin(), thrust::upper_bound( - thrust::seq, chunk_page_offsets, chunk_page_offsets + num_chunks + 1, page_idx)) - + thrust::seq, chunk_page_offsets.begin(), chunk_page_offsets.end(), page_idx)) - 1); - // Check if the chunk index is valid + // Check if the chunk index is valid. if (chunk_idx < 0 or chunk_idx >= num_chunks) { set_error(static_cast(decode_error::DATA_STREAM_OVERRUN), error_code); return; } + auto const page_span = page_data[page_idx]; + byte_stream_s bs{}; bs.ck = colchunks[chunk_idx]; - bs.base = bs.cur = page_locations[page_idx]; - bs.end = bs.ck.compressed_data + bs.ck.compressed_size; - // Check if byte stream pointers are valid. - if (bs.end < bs.cur) { - set_error(static_cast(decode_error::DATA_STREAM_OVERRUN), - error_code); - return; - } - // Clear page header info before writing known fields + bs.base = bs.cur = page_span.data(); + bs.end = page_span.data() + page_span.size(); + + // Clear the logical page descriptor. zero_out_page_header_info(&bs); bs.page.chunk_idx = chunk_idx; bs.page.src_col_schema = bs.ck.src_col_schema; - // bs.page.chunk_row not computed here and will be filled in later by // `fill_in_page_info()`. + // Return if empty page span (pruned page) + if (page_span.empty()) { + pages[page_idx] = bs.page; + return; + } + // Parsed page must be valid and not empty if (not parse_valid_page_header(&bs)) { set_error(static_cast(decode_error::INVALID_PAGE_HEADER), @@ -815,6 +818,13 @@ struct decode_page_headers_with_pgidx_fn { return; } + // Ensure we read the entire page. + if (cuda::std::cmp_not_equal(bs.end - bs.cur, bs.page.compressed_page_size)) { + set_error(static_cast(decode_error::DATA_STREAM_OVERRUN), + error_code); + return; + } + bs.page.page_data = const_cast(bs.cur); bs.page.kernel_mask = kernel_mask_for_page(bs.page, bs.ck); @@ -917,12 +927,14 @@ void count_page_headers(cudf::detail::hostdevice_span chunks, } void decode_page_headers(cudf::device_span chunks, - chunk_page_info* chunk_pages, + cudf::device_span chunk_pages, kernel_error::pointer error_code, rmm::cuda_stream_view stream) { static_assert(decode_page_headers_block_size % cudf::detail::warp_size == 0, "Block size for decode page headers kernel must be a multiple of warp size"); + CUDF_EXPECTS(chunk_pages.size() == chunks.size(), + "Chunk page info must contain one entry per chunk"); auto const num_chunks = static_cast(chunks.size()); auto constexpr num_warps_per_block = decode_page_headers_block_size / cudf::detail::warp_size; @@ -937,21 +949,24 @@ void decode_page_headers(cudf::device_span chunks, CUDF_CUDA_TRY(cudaGetLastError()); } -void decode_page_headers_with_pgidx(cudf::device_span chunks, - cudf::device_span pages, - uint8_t** page_locations, - size_type* chunk_page_offsets, - kernel_error::pointer error_code, - rmm::cuda_stream_view stream) +void decode_page_headers_from_page_data( + cudf::device_span chunks, + cudf::device_span pages, + cudf::device_span const> page_data, + cudf::device_span chunk_page_offsets, + kernel_error::pointer error_code, + rmm::cuda_stream_view stream) { + CUDF_EXPECTS(chunk_page_offsets.size() == chunks.size() + 1, + "Chunk page offsets must cover all chunks"); thrust::for_each(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), cuda::counting_iterator{0}, cuda::counting_iterator{static_cast(pages.size())}, - decode_page_headers_with_pgidx_fn{.colchunks = chunks, - .pages = pages, - .page_locations = page_locations, - .chunk_page_offsets = chunk_page_offsets, - .error_code = error_code}); + decode_from_page_data_fn{.colchunks = chunks, + .pages = pages, + .page_data = page_data, + .chunk_page_offsets = chunk_page_offsets, + .error_code = error_code}); } void build_string_dictionary_index(ColumnChunkDesc* chunks, diff --git a/cpp/src/io/parquet/page_string_decode.cu b/cpp/src/io/parquet/page_string_decode.cu index 4a53c717cbcd..adc5c375b8f7 100644 --- a/cpp/src/io/parquet/page_string_decode.cu +++ b/cpp/src/io/parquet/page_string_decode.cu @@ -520,7 +520,7 @@ CUDF_KERNEL void __launch_bounds__(preprocess_block_size) if (t == 0) { // don't clobber these if they're already computed from the index - if (!pp->has_page_index) { + if (!pp->has_value_info) { pp->num_nulls = 0; pp->num_valids = 0; } @@ -553,7 +553,7 @@ CUDF_KERNEL void __launch_bounds__(preprocess_block_size) is_bounds_page(s->page, s->col.start_row, min_row, num_rows, has_repetition); // if we have size info, then we only need to do this for bounds pages - if (pp->has_page_index && !is_bounds_pg) { return; } + if (pp->has_value_info && !is_bounds_pg) { return; } // Zero out everything and return early if the page is pruned if (not page_mask.empty() and not page_mask[page_idx]) { @@ -646,7 +646,7 @@ CUDF_KERNEL void __launch_bounds__(delta_preproc_block_size) is_bounds_page(s->page, s->col.start_row, min_row, num_rows, has_repetition); // if we have size info, then we only need to do this for bounds pages - if (pp->has_page_index && !is_bounds_pg) { + if (pp->has_value_info && !is_bounds_pg) { // check if we need to store values from the index if (t == 0 && is_page_contained(s->page, s->col.start_row, min_row, num_rows)) { pp->str_bytes = pp->str_bytes_from_index; @@ -728,7 +728,7 @@ CUDF_KERNEL void __launch_bounds__(delta_length_block_size) is_bounds_page(s->page, s->col.start_row, min_row, num_rows, has_repetition); // if we have size info, then we only need to do this for bounds pages - if (pp->has_page_index && !is_bounds_pg) { + if (pp->has_value_info && !is_bounds_pg) { // check if we need to store values from the index if (t == 0 && is_page_contained(s->page, s->col.start_row, min_row, num_rows)) { pp->str_bytes = pp->str_bytes_from_index; @@ -842,7 +842,7 @@ CUDF_KERNEL void __launch_bounds__(preprocess_block_size) is_bounds_page(s->page, s->col.start_row, min_row, num_rows, has_repetition); // if we have size info, then we only need to do this for bounds pages - if (pp->has_page_index && !is_bounds_pg) { + if (pp->has_value_info && !is_bounds_pg) { // check if we need to store values from the index if (t == 0 && is_page_contained(s->page, s->col.start_row, min_row, num_rows)) { pp->str_bytes = pp->str_bytes_from_index; diff --git a/cpp/src/io/parquet/page_string_utils.cuh b/cpp/src/io/parquet/page_string_utils.cuh index 084d9334602f..85f8bf06200e 100644 --- a/cpp/src/io/parquet/page_string_utils.cuh +++ b/cpp/src/io/parquet/page_string_utils.cuh @@ -1,5 +1,5 @@ /* - * 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 */ @@ -138,56 +138,6 @@ inline __device__ void compute_initial_large_strings_offset(page_state_s const* } } -/** - * @brief Update offsets with either zeros if this is a large string column, `initial_value` - * otherwise. - * - * For large string columns, fill zeros (sizes) at all offsets and atomically update the initial - * string offset. Otherwise, fill `initial_value` at all offsets. - * - * @tparam block_size Thread block size - * @tparam has_lists Whether the column is a list column - * @param[in,out] state page state - * @param[out] initial_str_offsets Initial string offsets - * @param[in] page Page information - */ -template -__device__ void update_string_offsets_for_pruned_pages( - page_state_s* state, cudf::device_span initial_str_offsets, PageInfo const& page) -{ - namespace cg = cooperative_groups; - - // Initial string offset - auto const initial_value = page.str_offset; - // The value count is either the leaf-level batch size in case of lists or the number of - // effective rows being read by this page - auto const value_count = - has_lists ? page.nesting[state->col.max_nesting_depth - 1].batch_size : state->num_rows; - auto const tid = cg::this_thread_block().thread_rank(); - - // Offsets pointer contains string sizes in case of large strings and actual offsets - // otherwise - auto& ni = state->nesting_info[state->col.max_nesting_depth - 1]; - auto offptr = reinterpret_cast(ni.data_out); - // For large strings, update the initial string buffer offset to be used during large string - // column construction. Otherwise, convert string sizes to final offsets - if (state->col.is_large_string_col) { - // Write zero string sizes - for (int idx = tid; idx < value_count; idx += block_size) { - offptr[idx] = 0; - } - // page.chunk_idx are ordered by input_col_idx and row_group_idx respectively - auto const chunks_per_rowgroup = initial_str_offsets.size(); - auto const input_col_idx = page.chunk_idx % chunks_per_rowgroup; - compute_initial_large_strings_offset(state, initial_str_offsets[input_col_idx]); - } else { - // Write the initial offset at all positions to indicate zero sized strings - for (int idx = tid; idx < value_count; idx += block_size) { - offptr[idx] = initial_value; - } - } -} - template CUDF_HOST_DEVICE constexpr int log2_int() { diff --git a/cpp/src/io/parquet/parquet_gpu.hpp b/cpp/src/io/parquet/parquet_gpu.hpp index 090431058c47..b8324b3f570a 100644 --- a/cpp/src/io/parquet/parquet_gpu.hpp +++ b/cpp/src/io/parquet/parquet_gpu.hpp @@ -374,7 +374,7 @@ struct PageInfo { Encoding definition_level_encoding; // Encoding used for definition levels (data page) Encoding repetition_level_encoding; // Encoding used for repetition levels (data page) bool is_compressed; // Whether the page is compressed (V2 header) - bool has_page_index; // true if str_bytes, num_valids, etc are derivable from page indexes + bool has_value_info; // true if str_bytes, num_valids, etc are derivable from page indexes }; // forward declaration @@ -707,26 +707,29 @@ void count_page_headers(cudf::detail::hostdevice_span chunks, * @param[in] stream CUDA stream to use */ void decode_page_headers(cudf::device_span chunks, - chunk_page_info* chunk_pages, + cudf::device_span chunk_pages, kernel_error::pointer error_code, rmm::cuda_stream_view stream); /** - * @brief Decode page headers from specified page locations from the page index + * @brief Decode page headers from corresponding specified page data spans. + * + * Empty spans initialize the corresponding logical page descriptor but are not decoded. * * @param[in] chunks Device span of column chunks * @param[out] pages Device span of pages - * @param[in] page_locations List of page locations + * @param[in] page_data Device span of page data * @param[in] chunk_page_offsets List of running count of page locations per column chunk * @param[out] error_code Error code for kernel failures * @param[in] stream CUDA stream to use */ -void decode_page_headers_with_pgidx(cudf::device_span chunks, - cudf::device_span pages, - uint8_t** page_locations, - size_type* chunk_page_offsets, - kernel_error::pointer error_code, - rmm::cuda_stream_view stream); +void decode_page_headers_from_page_data( + cudf::device_span chunks, + cudf::device_span pages, + cudf::device_span const> page_data, + cudf::device_span chunk_page_offsets, + kernel_error::pointer error_code, + rmm::cuda_stream_view stream); /** * @brief Launches kernel for building the dictionary index for the column @@ -1017,6 +1020,25 @@ void preprocess_levels(cudf::detail::hostdevice_span pages, int level_type_size, rmm::cuda_stream_view stream); +/** + * @brief Fills output offset entries for pruned string and list pages + * + * @param[in] pages All pages to be processed + * @param[in] chunks All chunks to be processed + * @param[in] page_mask Boolean vector indicating which pages are decoded + * @param[in,out] initial_str_offsets Initial offsets used to construct large nested strings + * @param[in] skip_rows Number of rows to skip + * @param[in] num_rows Number of rows to read + * @param[in] stream CUDA stream to use + */ +void fill_pruned_offsets(cudf::device_span pages, + cudf::device_span chunks, + cudf::device_span page_mask, + cudf::device_span initial_str_offsets, + size_t skip_rows, + size_t num_rows, + rmm::cuda_stream_view stream); + /** * @brief Launches kernel for reading non-dictionary fixed width column data stored in the pages * diff --git a/cpp/src/io/parquet/reader_impl.cpp b/cpp/src/io/parquet/reader_impl.cpp index 38ee9301d03f..02485a65637c 100644 --- a/cpp/src/io/parquet/reader_impl.cpp +++ b/cpp/src/io/parquet/reader_impl.cpp @@ -397,6 +397,9 @@ void reader_impl::decode_page_data(read_mode mode, size_t skip_rows, size_t num_ // Invalidate output buffer nullmasks at row indices spanned by pruned pages update_output_nullmasks_for_pruned_pages(subpass_page_mask_span(), skip_rows, num_rows); + // Fill output offsets for pruned pages before retrieving large string initial offsets. + fill_pruned_offsets(skip_rows, num_rows, initial_str_offsets); + // Copy over initial string offsets from device auto h_initial_str_offsets = cudf::detail::make_pinned_vector_async(initial_str_offsets, _stream); @@ -439,9 +442,15 @@ void reader_impl::decode_page_data(read_mode mode, size_t skip_rows, size_t num_ } // Nested large strings column else if (input_col.nesting_depth() > 0) { - CUDF_EXPECTS(h_initial_str_offsets[idx] != std::numeric_limits::max(), - "Encountered invalid initial offset for large string column"); - out_buf.set_initial_string_offset(h_initial_str_offsets[idx]); + // A fully pruned list may have no string child values and therefore no page from which + // to record an initial offset. + if (out_buf.size == 0) { + out_buf.set_initial_string_offset(0); + } else { + CUDF_EXPECTS(h_initial_str_offsets[idx] != std::numeric_limits::max(), + "Encountered invalid initial offset for large string column"); + out_buf.set_initial_string_offset(h_initial_str_offsets[idx]); + } } } } diff --git a/cpp/src/io/parquet/reader_impl.hpp b/cpp/src/io/parquet/reader_impl.hpp index 61f85e047809..76d88f52e310 100644 --- a/cpp/src/io/parquet/reader_impl.hpp +++ b/cpp/src/io/parquet/reader_impl.hpp @@ -350,6 +350,17 @@ class reader_impl { size_t skip_rows, size_t num_rows); + /** + * @brief Fill in string and list offsets for rows covered by pruned data pages. + * + * @param skip_rows Offset of the first row in the table chunk + * @param num_rows Number of rows in the table chunk + * @param initial_str_offsets Initial offsets used to construct large nested strings + */ + void fill_pruned_offsets(size_t skip_rows, + size_t num_rows, + cudf::device_span initial_str_offsets); + /** * @brief Creates file-wide parquet chunk information. * @@ -568,8 +579,8 @@ class reader_impl { bool _strings_to_categorical = false; - // are there usable page indexes available - bool _has_page_index = false; + // are offset indexes available for selected row groups + bool _has_offset_index = false; std::optional> _reader_column_schema; diff --git a/cpp/src/io/parquet/reader_impl_chunking.cu b/cpp/src/io/parquet/reader_impl_chunking.cu index 9e4957dbea05..c2a53826bb05 100644 --- a/cpp/src/io/parquet/reader_impl_chunking.cu +++ b/cpp/src/io/parquet/reader_impl_chunking.cu @@ -1,5 +1,5 @@ /* - * 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 */ @@ -177,7 +177,7 @@ void reader_impl::setup_next_pass(read_mode mode) // if we are doing subpass reading, generate more accurate num_row estimates for list columns. // this helps us to generate more accurate subpass splits. if (pass.has_compressed_data && _input_pass_read_limit != 0) { - if (not _has_page_index) { + if (not _has_offset_index) { generate_list_column_row_counts(is_estimate_row_counts::YES); } else { generate_list_column_row_counts(is_estimate_row_counts::NO); @@ -295,7 +295,7 @@ void reader_impl::setup_next_subpass(read_mode mode) remaining_read_limit, num_columns, is_first_subpass, - _has_page_index, + _has_offset_index, _stream); }(); @@ -422,7 +422,7 @@ void reader_impl::create_global_chunk_info() // Mapping of input column to page index column std::vector column_mapping; - if (_has_page_index and not row_groups_info.empty()) { + if (_has_offset_index and not row_groups_info.empty()) { // use first row group to define mappings (assumes same schema for each file) auto const& rg = row_groups_info[0]; auto const& columns = _metadata->get_row_group(rg.index, rg.source_index).columns; @@ -480,7 +480,7 @@ void reader_impl::create_global_chunk_info() // grab the column_chunk_info for each chunk (if it exists) column_chunk_info const* const chunk_info = - _has_page_index ? &rg.column_chunks.value()[column_mapping[i]] : nullptr; + _has_offset_index ? &rg.column_chunks.value()[column_mapping[i]] : nullptr; chunks.emplace_back(col_meta.total_compressed_size, nullptr, diff --git a/cpp/src/io/parquet/reader_impl_chunking_utils.cu b/cpp/src/io/parquet/reader_impl_chunking_utils.cu index 96d3f1a67cc2..ebca8918171c 100644 --- a/cpp/src/io/parquet/reader_impl_chunking_utils.cu +++ b/cpp/src/io/parquet/reader_impl_chunking_utils.cu @@ -356,7 +356,7 @@ std::tuple, size_t, size_t> compute_next_subpass( size_t size_limit, size_t num_columns, bool is_first_subpass, - bool has_page_index, + bool has_offset_index, rmm::cuda_stream_view stream) { auto [aggregated_info, page_keys_by_split] = adjust_cumulative_sizes(c_info, pages, stream); @@ -387,13 +387,17 @@ std::tuple, size_t, size_t> compute_next_subpass( auto iter = cuda::counting_iterator{size_t{0}}; auto page_row_index = cudf::detail::make_counting_transform_iterator(0, get_page_end_row_index{c_info}); - thrust::transform( - rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - iter, - iter + num_columns, - page_bounds.begin(), - get_page_span{ - page_offsets, chunks, page_row_index, start_row, end_row, is_first_subpass, has_page_index}); + thrust::transform(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), + iter, + iter + num_columns, + page_bounds.begin(), + get_page_span{page_offsets, + chunks, + page_row_index, + start_row, + end_row, + is_first_subpass, + has_offset_index}); // total page count over all columns auto page_count_iter = cuda::make_transform_iterator(page_bounds.begin(), get_span_size{}); diff --git a/cpp/src/io/parquet/reader_impl_chunking_utils.cuh b/cpp/src/io/parquet/reader_impl_chunking_utils.cuh index 10a172c75c9b..347f874e11a5 100644 --- a/cpp/src/io/parquet/reader_impl_chunking_utils.cuh +++ b/cpp/src/io/parquet/reader_impl_chunking_utils.cuh @@ -1,5 +1,5 @@ /* - * 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 */ @@ -128,7 +128,7 @@ adjust_cumulative_sizes(device_span c_info, * @param size_limit The size limit in bytes of the subpass * @param num_columns The number of columns * @param is_first_subpass Boolean indicating if this is the first subpass - * @param has_page_index Boolean indicating if we have a page index + * @param has_offset_index Boolean indicating if offset indexes are available * @param stream The stream to execute cuda operations on * @returns A tuple containing a vector of page_span structs indicating the page indices to include * for each column to be processed, the total number of pages over all columns, and the total @@ -144,7 +144,7 @@ std::tuple, size_t, size_t> compute_next_subpass( size_t size_limit, size_t num_columns, bool is_first_subpass, - bool has_page_index, + bool has_offset_index, rmm::cuda_stream_view stream); /** @@ -626,7 +626,7 @@ struct get_page_span { size_t const start_row; size_t const end_row; bool const is_first_subpass; - bool const has_page_index; + bool const has_offset_index; get_page_span(device_span _page_offsets, device_span _chunks, @@ -634,14 +634,14 @@ struct get_page_span { size_t _start_row, size_t _end_row, bool _is_first_subpass, - bool _has_page_index) + bool _has_offset_index) : page_offsets(_page_offsets), chunks(_chunks), page_row_index(_page_row_index), start_row(_start_row), end_row(_end_row), is_first_subpass(_is_first_subpass), - has_page_index(_has_page_index) + has_offset_index(_has_offset_index) { } @@ -656,13 +656,13 @@ struct get_page_span { // For list columns, the row counts are estimates so we need all prefix pages to correctly // compute page bounds. For non-list columns, we can get an exact span of pages. auto start_page = first_page_index; - auto const update_start_page = has_page_index or (not is_list) or (not is_first_subpass); + auto const update_start_page = has_offset_index or (not is_list) or (not is_first_subpass); if (update_start_page) { start_page += cuda::std::distance( column_page_start, thrust::lower_bound(thrust::seq, column_page_start, column_page_end, start_row)); } - if (page_row_index[start_page] == start_row and (has_page_index or not is_list)) { + if (page_row_index[start_page] == start_row and (has_offset_index or not is_list)) { start_page++; } diff --git a/cpp/src/io/parquet/reader_impl_helpers.cpp b/cpp/src/io/parquet/reader_impl_helpers.cpp index ce06969cf0da..ed60a4db0e2d 100644 --- a/cpp/src/io/parquet/reader_impl_helpers.cpp +++ b/cpp/src/io/parquet/reader_impl_helpers.cpp @@ -719,14 +719,10 @@ void aggregate_reader_metadata::column_info_for_row_group(row_group_info& rg_inf auto const max_def_level = schema.max_definition_level; auto const max_rep_level = schema.max_repetition_level; - // If any columns lack the page indexes then just return without modifying the - // row_group_info. - if (not col_chunk.offset_index.has_value() or not col_chunk.column_index.has_value()) { - return; - } + // Return early if any columns lack the offset index. + if (not col_chunk.offset_index.has_value()) { return; } auto const& offset_index = col_chunk.offset_index.value(); - auto const& column_index = col_chunk.column_index.value(); auto& chunk_info = chunks[col_idx]; auto const num_pages = offset_index.page_locations.size(); @@ -751,18 +747,24 @@ void aggregate_reader_metadata::column_info_for_row_group(row_group_info& rg_inf } } + // Populate additional value-count metadata if column index is also available. + auto const* column_index = + col_chunk.column_index.has_value() ? &col_chunk.column_index.value() : nullptr; + // Use the definition_level_histogram to get num_valid and num_null. For now, these are // only ever used for byte array columns. The repetition_level_histogram might be // necessary to determine the total number of values in the page if the // definition_level_histogram is absent. // // In the future we might want the full histograms saved in the `column_info` struct. - int64_t const* const def_hist = column_index.definition_level_histogram.has_value() - ? column_index.definition_level_histogram.value().data() - : nullptr; - int64_t const* const rep_hist = column_index.repetition_level_histogram.has_value() - ? column_index.repetition_level_histogram.value().data() - : nullptr; + int64_t const* const def_hist = + column_index != nullptr and column_index->definition_level_histogram.has_value() + ? column_index->definition_level_histogram.value().data() + : nullptr; + int64_t const* const rep_hist = + column_index != nullptr and column_index->repetition_level_histogram.has_value() + ? column_index->repetition_level_histogram.value().data() + : nullptr; for (size_t pg_idx = 0; pg_idx < num_pages; pg_idx++) { auto const& page_loc = offset_index.page_locations[pg_idx]; @@ -777,8 +779,8 @@ void aggregate_reader_metadata::column_info_for_row_group(row_group_info& rg_inf page_info pg_info{.location = page_loc, .num_rows = num_rows}; // check to see if we already have null counts for each page - if (column_index.null_counts.has_value()) { - pg_info.num_nulls = column_index.null_counts.value()[pg_idx]; + if (column_index != nullptr and column_index->null_counts.has_value()) { + pg_info.num_nulls = column_index->null_counts.value()[pg_idx]; } // save variable length byte info if present @@ -820,19 +822,6 @@ void aggregate_reader_metadata::column_info_for_row_group(row_group_info& rg_inf } } - // If none of the ifs above triggered, then we have neither histogram (likely the writer - // doesn't produce them, the r:0 d:1 case should have been handled above). The column index - // doesn't give us value counts, so we'll have to rely on the page headers. If the histogram - // info is missing or insufficient, then just return without modifying the row_group_info. - if (not pg_info.num_nulls.has_value() or not pg_info.num_valid.has_value()) { return; } - - // Like above, if using older page indexes that lack size info, then return without modifying - // the row_group_info. - // TODO: cudf will still set the per-page var_bytes to '0' even for all null pages. Need to - // check the behavior of other implementations (once there are some). Some may not set the - // var bytes for all null pages, so check the `null_pages` field on the column index. - if (schema.type == Type::BYTE_ARRAY and not pg_info.var_bytes_size.has_value()) { return; } - chunk_info.pages.push_back(std::move(pg_info)); } } diff --git a/cpp/src/io/parquet/reader_impl_helpers.hpp b/cpp/src/io/parquet/reader_impl_helpers.hpp index dacaf2bfa22f..40c70babf2b8 100644 --- a/cpp/src/io/parquet/reader_impl_helpers.hpp +++ b/cpp/src/io/parquet/reader_impl_helpers.hpp @@ -77,9 +77,9 @@ struct row_group_info { std::optional> column_chunks; /** - * @brief Indicates the presence of page-level indexes. + * @brief Indicates the presence of page-level offset indexes. */ - [[nodiscard]] bool has_page_index() const { return column_chunks.has_value(); } + [[nodiscard]] bool has_offset_index() const { return column_chunks.has_value(); } }; /** diff --git a/cpp/src/io/parquet/reader_impl_preprocess.cu b/cpp/src/io/parquet/reader_impl_preprocess.cu index 3eddddf8d0c0..5ade86bfdcf2 100644 --- a/cpp/src/io/parquet/reader_impl_preprocess.cu +++ b/cpp/src/io/parquet/reader_impl_preprocess.cu @@ -35,6 +35,7 @@ #include #include #include +#include #include namespace cudf::io::parquet::detail { @@ -51,7 +52,14 @@ inline bool is_treat_fixed_length_as_string(cuda::std::optional con } struct set_str_bytes_all { - __device__ void operator()(PageInfo& p) { p.str_bytes_all = p.str_bytes; } + device_span pages; + device_span page_mask; + + __device__ void operator()(size_type index) const + { + pages[index].str_bytes_all = + page_mask.empty() or page_mask[index] ? pages[index].str_bytes : int32_t{0}; + } }; } // namespace @@ -562,8 +570,8 @@ void reader_impl::read_compressed_data() read_chunks_tasks.get(); // Process dataset chunk pages into output columns - auto const total_pages = _has_page_index ? count_page_headers_with_pgidx(chunks, _stream) - : count_page_headers(chunks, _stream); + auto const total_pages = _has_offset_index ? count_page_headers_with_pgidx(chunks, _stream) + : count_page_headers(chunks, _stream); if (total_pages <= 0) { return; } // Zero out the vector before `decode_page_headers` as it may not write every byte of the buffer, @@ -572,7 +580,7 @@ void reader_impl::read_compressed_data() total_pages, _stream, cudf::get_current_device_resource_ref()); // decoding of column/page information - decode_page_headers(pass, unsorted_pages, _has_page_index, _stream); + decode_page_headers(pass, unsorted_pages, _has_offset_index, _stream); CUDF_EXPECTS(pass.page_offsets.size() - 1 == static_cast(_input_columns.size()), "Encountered page_offsets / num_columns mismatch"); } @@ -624,10 +632,11 @@ void reader_impl::preprocess_file(read_mode mode) _file_itm_data.exclusive_sum_num_rows_per_source.begin()); } - // check for page indexes - _has_page_index = std::all_of(_file_itm_data.row_groups.cbegin(), - _file_itm_data.row_groups.cend(), - [](auto const& row_group) { return row_group.has_page_index(); }); + // Check for offset indexes. + _has_offset_index = + std::all_of(_file_itm_data.row_groups.cbegin(), + _file_itm_data.row_groups.cend(), + [](auto const& row_group) { return row_group.has_offset_index(); }); if (_file_itm_data.global_num_rows > 0 && not _file_itm_data.row_groups.empty() && not _input_columns.empty()) { @@ -833,7 +842,15 @@ void reader_impl::preprocess_subpass_pages(read_mode mode, size_t chunk_read_lim is_treat_fixed_length_as_string(chunk.logical_type); }); - if (!_has_page_index || has_flba) { + // String pages with missing value info do not have the value counts or string byte sizes needed + // for chunking. Scan them so chunk boundaries and output allocations use their actual sizes. + auto const has_string_page_without_info = + std::any_of(subpass.pages.host_begin(), subpass.pages.host_end(), [](auto const& page) { + return (static_cast(page.kernel_mask) & STRINGS_MASK) != 0 and + (page.flags & PAGEINFO_FLAGS_DICTIONARY) == 0 and not page.has_value_info; + }); + + if (has_string_page_without_info || has_flba) { constexpr bool compute_all_string_sizes = true; compute_page_string_sizes_pass1(subpass.pages, pass.chunks, @@ -848,9 +865,9 @@ void reader_impl::preprocess_subpass_pages(read_mode mode, size_t chunk_read_lim } // set str_bytes_all thrust::for_each(rmm::exec_policy_nosync(_stream, cudf::get_current_device_resource_ref()), - subpass.pages.device_begin(), - subpass.pages.device_end(), - set_str_bytes_all{}); + cuda::counting_iterator(0), + cuda::counting_iterator(subpass.pages.size()), + set_str_bytes_all{subpass.pages, subpass_page_mask_span()}); } // retrieve pages back @@ -883,7 +900,7 @@ void reader_impl::preprocess_subpass_pages(read_mode mode, size_t chunk_read_lim // corner case: only decode up to the second-to-last row, except if this is the last page in the // entire pass or if we have the page index. this handles the case where we only have 1 chunk, 1 // page, and potentially even just 1 row. - if (is_list and std::cmp_less(max_col_row, last_pass_row) and not _has_page_index) { + if (is_list and std::cmp_less(max_col_row, last_pass_row) and not _has_offset_index) { // compute min row for this column in the subpass auto const& first_page = subpass.pages[first_page_index]; auto const& first_chunk = pass.chunks[first_page.chunk_idx]; @@ -1092,6 +1109,29 @@ void reader_impl::allocate_columns(read_mode mode, size_t skip_rows, size_t num_ pinned_nullmask_bufs, std::numeric_limits::max(), _stream); } +void reader_impl::fill_pruned_offsets(size_t skip_rows, + size_t num_rows, + cudf::device_span initial_str_offsets) +{ + // Return early if there are no pruned pages + auto const page_mask = subpass_page_mask_span(); + if (page_mask.is_empty() or + std::all_of(page_mask.host_begin(), page_mask.host_end(), cuda::std::identity{})) { + return; + } + + auto const& pass = *_pass_itm_data; + auto const& subpass = *pass.subpass; + auto const pages = device_span{subpass.pages.device_ptr(), subpass.pages.size()}; + auto const chunks = + device_span{pass.chunks.device_ptr(), pass.chunks.size()}; + auto const device_page_mask = static_cast>(page_mask); + + // Set offsets for pruned string and list pages. + parquet::detail::fill_pruned_offsets( + pages, chunks, device_page_mask, initial_str_offsets, skip_rows, num_rows, _stream); +} + cudf::detail::host_vector reader_impl::calculate_page_string_offsets() { auto& pass = *_pass_itm_data; diff --git a/cpp/src/io/parquet/reader_impl_preprocess_utils.cu b/cpp/src/io/parquet/reader_impl_preprocess_utils.cu index a36f62f7af79..9ee4ceecf87d 100644 --- a/cpp/src/io/parquet/reader_impl_preprocess_utils.cu +++ b/cpp/src/io/parquet/reader_impl_preprocess_utils.cu @@ -1,5 +1,5 @@ /* - * 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 */ @@ -272,6 +272,7 @@ void fill_in_page_info(host_span chunks, { auto const num_pages = pages.size(); auto page_indexes = cudf::detail::make_pinned_vector_async(num_pages, stream); + std::fill(page_indexes.begin(), page_indexes.end(), page_index_info{}); for (size_t c = 0, page_count = 0; c < chunks.size(); c++) { auto const& chunk = chunks[c]; @@ -286,6 +287,10 @@ void fill_in_page_info(host_span chunks, page.num_nulls = chunk_info.pages[p].num_nulls.value_or(0); page.num_valids = chunk_info.pages[p].num_valid.value_or(0); page.str_bytes = chunk_info.pages[p].var_bytes_size.value_or(0); + page.has_value_info = + (chunk_info.pages[p].num_nulls.has_value() and chunk_info.pages[p].num_valid.has_value() and + (chunk.physical_type != Type::BYTE_ARRAY or + chunk_info.pages[p].var_bytes_size.has_value())); start_row += page.num_rows; } @@ -406,10 +411,32 @@ cudf::detail::hostdevice_vector sort_pages(device_span return pass_pages; } -void decode_page_headers(pass_intermediate_data& pass, - device_span unsorted_pages, - bool has_page_index, - rmm::cuda_stream_view stream) +namespace { + +/** + * @brief Page data source type + */ +enum class page_data_source_type : uint8_t { + COLUMN_CHUNKS = 0, + OFFSET_INDEX = 1, + PAGE_SPANS = 2, +}; + +/** + * @brief Dispatch decode page headers base on the type of page data source + * + * @tparam data_source_type Type of page data source + * + * @param pass Struct containing pass information + * @param unsorted_pages Device span of page information to decode + * @param page_data Host span of page data spans (only used for PAGE_SPANS source) + * @param stream Stream to use + */ +template +void decode_page_headers_impl(pass_intermediate_data& pass, + device_span unsorted_pages, + host_span const> page_data, + rmm::cuda_stream_view stream) { CUDF_FUNC_RANGE(); @@ -439,11 +466,24 @@ void decode_page_headers(pass_intermediate_data& pass, kernel_error error_code(stream); - // If page index is present, collect data ptrs for all pages and launch the accelerated decode + if constexpr (data_source_type == page_data_source_type::PAGE_SPANS) { + CUDF_EXPECTS(page_data.size() == unsorted_pages.size(), + "Page span count must match the number of logical pages"); + auto device_page_data = cudf::detail::make_device_uvector_async( + page_data, stream, cudf::get_current_device_resource_ref()); + decode_page_headers_from_page_data( + device_span(pass.chunks.device_ptr(), pass.chunks.size()), + unsorted_pages, + device_page_data, + device_span(chunk_page_offsets.data(), chunk_page_offsets.size()), + error_code.data(), + stream); + } + // If offset index is present, collect data spans for all pages and launch the accelerated decode // page headers kernel - if (has_page_index) { - auto host_page_locations = - cudf::detail::make_pinned_vector_async(unsorted_pages.size(), stream); + else if constexpr (data_source_type == page_data_source_type::OFFSET_INDEX) { + auto host_page_data = cudf::detail::make_pinned_vector_async>( + unsorted_pages.size(), stream); auto curr_page_idx = 0; std::for_each(pass.chunks.begin(), pass.chunks.end(), [&](auto const& chunk) { @@ -461,9 +501,11 @@ void decode_page_headers(pass_intermediate_data& pass, CUDF_EXPECTS(std::cmp_less(chunk.h_chunk_info->dictionary_offset.value(), chunk.h_chunk_info->pages.front().location.offset), "Encountered dictionary page located beyond the first data page"); - host_page_locations[curr_page_idx] = data_ptr; + auto const dictionary_size = chunk.h_chunk_info->dictionary_size.value(); + CUDF_EXPECTS(dictionary_size >= 0, "Encountered invalid dictionary page size"); + host_page_data[curr_page_idx] = {data_ptr, static_cast(dictionary_size)}; ++curr_page_idx; - data_ptr += chunk.h_chunk_info->dictionary_size.value(); + data_ptr += dictionary_size; } // Data pages @@ -471,39 +513,42 @@ void decode_page_headers(pass_intermediate_data& pass, std::cmp_equal(chunk.h_chunk_info->pages.size(), chunk.num_data_pages), "Encountered invalid sized data page information in the page index"); auto const num_data_pages = chunk.num_data_pages; - std::for_each(cuda::counting_iterator{0}, - cuda::counting_iterator{num_data_pages}, - [&](auto const page_idx) { - host_page_locations[curr_page_idx] = data_ptr; - ++curr_page_idx; - if (page_idx < num_data_pages - 1) { - data_ptr += chunk.h_chunk_info->pages[page_idx + 1].location.offset - - chunk.h_chunk_info->pages[page_idx].location.offset; - } - }); + std::for_each( + cuda::counting_iterator{0}, + cuda::counting_iterator{num_data_pages}, + [&](auto const page_idx) { + auto const page_size = chunk.h_chunk_info->pages[page_idx].location.compressed_page_size; + CUDF_EXPECTS(page_size >= 0, "Encountered invalid data page size"); + host_page_data[curr_page_idx] = {data_ptr, static_cast(page_size)}; + ++curr_page_idx; + if (page_idx < num_data_pages - 1) { + data_ptr += chunk.h_chunk_info->pages[page_idx + 1].location.offset - + chunk.h_chunk_info->pages[page_idx].location.offset; + } + }); }); - // Check if we have data ptrs for all input pages + // Check if we have data spans for all input pages CUDF_EXPECTS(std::cmp_equal(curr_page_idx, unsorted_pages.size()), "Expected page offsets to match total pages"); - // Copy page data ptrs to device - auto page_locations = cudf::detail::make_device_uvector_async( - host_page_locations, stream, cudf::get_current_device_resource_ref()); + // Copy page data spans to device + auto page_data = cudf::detail::make_device_uvector_async( + host_page_data, stream, cudf::get_current_device_resource_ref()); // Accelerated decode page headers, one thread per page - decode_page_headers_with_pgidx( + decode_page_headers_from_page_data( device_span(pass.chunks.device_ptr(), pass.chunks.size()), unsorted_pages, - page_locations.begin(), - chunk_page_offsets.begin(), + page_data, + device_span(chunk_page_offsets.data(), chunk_page_offsets.size()), error_code.data(), stream); } else { - // (Slow) decode page headers, one warp (lane) per pages of a chunk + // (Slow) decode page headers, one warp (lane) per pages of a column chunk decode_page_headers( device_span(pass.chunks.device_ptr(), pass.chunks.size()), - d_chunk_page_info.begin(), + device_span(d_chunk_page_info.data(), d_chunk_page_info.size()), error_code.data(), stream); } @@ -519,7 +564,10 @@ void decode_page_headers(pass_intermediate_data& pass, } } - if (has_page_index) { fill_in_page_info(pass.chunks, unsorted_pages, stream); } + if (data_source_type == page_data_source_type::OFFSET_INDEX or + data_source_type == page_data_source_type::PAGE_SPANS) { + fill_in_page_info(pass.chunks, unsorted_pages, stream); + } // compute max bytes needed for level data auto level_bit_size = cudf::detail::make_counting_transform_iterator( @@ -576,4 +624,28 @@ void decode_page_headers(pass_intermediate_data& pass, stream.synchronize(); } +} // namespace + +void decode_page_headers(pass_intermediate_data& pass, + device_span unsorted_pages, + bool has_offset_index, + rmm::cuda_stream_view stream) +{ + if (has_offset_index) { + decode_page_headers_impl(pass, unsorted_pages, {}, stream); + } else { + decode_page_headers_impl( + pass, unsorted_pages, {}, stream); + } +} + +void decode_page_headers(pass_intermediate_data& pass, + device_span unsorted_pages, + host_span const> page_data, + rmm::cuda_stream_view stream) +{ + decode_page_headers_impl( + pass, unsorted_pages, page_data, stream); +} + } // namespace cudf::io::parquet::detail diff --git a/cpp/src/io/parquet/reader_impl_preprocess_utils.cuh b/cpp/src/io/parquet/reader_impl_preprocess_utils.cuh index 3a0f54cd5cd6..8f19d7a8d787 100644 --- a/cpp/src/io/parquet/reader_impl_preprocess_utils.cuh +++ b/cpp/src/io/parquet/reader_impl_preprocess_utils.cuh @@ -1,5 +1,5 @@ /* - * 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 */ @@ -128,12 +128,27 @@ std::string encoding_to_string(Encoding encoding); * * @param pass The struct containing pass information * @param unsorted_pages Device span of page information to decode - * @param has_page_index Boolean indicating if the page index is available + * @param has_offset_index Boolean indicating if the offset index is available * @param stream CUDA stream used for device memory operations and kernel launches */ void decode_page_headers(pass_intermediate_data& pass, device_span unsorted_pages, - bool has_page_index, + bool has_offset_index, + rmm::cuda_stream_view stream); + +/** + * @brief Decode page information using one exact span per logical indexed page + * + * Empty data spans represent masked pages and retain their logical page-index metadata. + * + * @param pass Struct containing pass information + * @param unsorted_pages Device span of page information to decode + * @param page_data Host span of page data device spans, one per logical indexed page + * @param stream CUDA stream used for device memory operations and kernel launches + */ +void decode_page_headers(pass_intermediate_data& pass, + device_span unsorted_pages, + host_span const> page_data, rmm::cuda_stream_view stream); /** @@ -157,6 +172,7 @@ struct page_index_info { int32_t num_nulls; int32_t num_valids; int32_t str_bytes; + bool has_value_info; }; /** @@ -168,17 +184,19 @@ struct copy_page_info { __device__ constexpr void operator()(size_type idx) { - auto& pg = pages[idx]; - auto const& pi = page_indexes[idx]; - pg.num_rows = pi.num_rows; - pg.chunk_row = pi.chunk_row; - pg.has_page_index = true; - pg.num_nulls = pi.num_nulls; - pg.num_valids = pi.num_valids; - pg.str_bytes_from_index = pi.str_bytes; - pg.str_bytes = pi.str_bytes; - pg.start_val = 0; - pg.end_val = pg.num_valids; + auto& pg = pages[idx]; + auto const& pi = page_indexes[idx]; + pg.num_rows = pi.num_rows; + pg.chunk_row = pi.chunk_row; + pg.has_value_info = pi.has_value_info; + pg.start_val = 0; + if (pg.has_value_info) { + pg.num_nulls = pi.num_nulls; + pg.num_valids = pi.num_valids; + pg.str_bytes_from_index = pi.str_bytes; + pg.str_bytes = pi.str_bytes; + pg.end_val = pg.num_valids; + } } }; diff --git a/cpp/tests/io/experimental/hybrid_scan_test.cpp b/cpp/tests/io/experimental/hybrid_scan_test.cpp index e580994b6d34..edf43fc6bdff 100644 --- a/cpp/tests/io/experimental/hybrid_scan_test.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_test.cpp @@ -13,11 +13,14 @@ #include #include +#include #include #include #include #include +#include #include +#include #include #include #include @@ -438,6 +441,110 @@ TEST_F(HybridScanTest, MaterializeListsOfStrings) test_hybrid_scan({col0, *col1, *col2, *col3, *col4}, false); } +TEST_F(HybridScanTest, ConsecutivePrunedPageOffsets) +{ + std::mt19937 gen(0x5ca1e); + auto constexpr num_rows = num_ordered_rows; + + auto col0 = testdata::ascending(); + auto col1 = testdata::ascending(); + auto col2 = make_parquet_list_col(gen, num_rows, 3, true); + col2 = cudf::purge_nonempty_nulls(col2->view()); + auto col3 = make_parquet_list_list_col(0, num_rows, 2, 3, true); + col3 = cudf::purge_nonempty_nulls(col3->view()); + auto col4 = make_list_str_column(gen, true, true); + col4 = cudf::purge_nonempty_nulls(col4->view()); + + auto const input = cudf::table_view{{col0, col1, *col2, *col3, *col4}}; + + std::string filepath = "ConsecutivePrunedPageOffsets.parquet"; + { + auto metadata = cudf::io::table_input_metadata(input); + metadata.column_metadata[0].set_name("col0"); + + auto const write_options = + cudf::io::parquet_writer_options::builder(cudf::io::sink_info{filepath}, input) + .metadata(std::move(metadata)) + .row_group_size_rows(num_rows) + .max_page_size_rows(page_size_for_ordered_tests) + .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN) + .build(); + cudf::io::write_parquet(write_options); + } + + auto const stream = cudf::get_default_stream(); + auto const mr = cudf::get_current_device_resource_ref(); + auto aligned_mr = rmm::mr::aligned_resource_adaptor(mr, bloom_filter_alignment); + + // Helper to validate monotonic string and list offsets + auto const expect_monotonic_offsets = [](auto& self, cudf::column_view const& column) -> void { + EXPECT_FALSE(cudf::has_nonempty_nulls(column)); + + if (column.type().id() == cudf::type_id::STRING) { + auto const offsets = cudf::strings_column_view{column}.offsets(); + auto const host_offsets = cudf::test::to_host(offsets).first; + EXPECT_TRUE(std::is_sorted(host_offsets.begin(), host_offsets.end())); + return; + } + + if (column.type().id() == cudf::type_id::LIST) { + auto const lists = cudf::lists_column_view{column}; + auto const host_offsets = cudf::test::to_host(lists.offsets()).first; + EXPECT_TRUE(std::is_sorted(host_offsets.begin(), host_offsets.end())); + self(self, lists.child()); + return; + } + + for (auto index = cudf::size_type{0}; index < column.num_children(); ++index) { + self(self, column.child(index)); + } + }; + + // Helper to validate the `offsets` children of string and list column + auto const validate = [&](cudf::ast::operation const& filter, + std::vector const& expected_slices) { + auto datasource = cudf::io::datasource::create({filepath}); + + auto const expected = cudf::concatenate(cudf::slice(input, expected_slices)); + std::unique_ptr filter_table; + std::unique_ptr payload_table; + ASSERT_NO_THROW(std::tie(filter_table, payload_table) = + hybrid_scan(*datasource, filter, {}, true, stream, mr, aligned_mr)); + + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected->view().select({0}), filter_table->view()); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected->view().select({1, 2, 3, 4}), + payload_table->view()); + for (auto const& column : payload_table->view()) { + expect_monotonic_offsets(expect_monotonic_offsets, column); + } + }; + + auto const col_ref = cudf::ast::column_name_reference{"col0"}; + + // Prune two leading pages. + auto middle_value = cudf::numeric_scalar{2 * page_size_for_ordered_tests / 100}; + auto middle = cudf::ast::literal{middle_value}; + auto const keep_trailing = + cudf::ast::operation{cudf::ast::ast_operator::GREATER_EQUAL, col_ref, middle}; + validate(keep_trailing, {2 * page_size_for_ordered_tests, num_rows}); + + // Prune two trailing pages. + auto const keep_leading = cudf::ast::operation{cudf::ast::ast_operator::LESS, col_ref, middle}; + validate(keep_leading, {0, 2 * page_size_for_ordered_tests}); + + // Prune two middle pages. + auto lower_value = cudf::numeric_scalar{page_size_for_ordered_tests / 100}; + auto upper_value = cudf::numeric_scalar{3 * page_size_for_ordered_tests / 100}; + auto lower = cudf::ast::literal{lower_value}; + auto upper = cudf::ast::literal{upper_value}; + auto const keep_first = cudf::ast::operation{cudf::ast::ast_operator::LESS, col_ref, lower}; + auto const keep_last = + cudf::ast::operation{cudf::ast::ast_operator::GREATER_EQUAL, col_ref, upper}; + auto const keep_outer = + cudf::ast::operation{cudf::ast::ast_operator::LOGICAL_OR, keep_first, keep_last}; + validate(keep_outer, {0, page_size_for_ordered_tests, 3 * page_size_for_ordered_tests, num_rows}); +} + TEST_F(HybridScanTest, MaterializeStructs) { std::mt19937 gen(0xbaLL);