Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions cpp/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
31 changes: 3 additions & 28 deletions cpp/src/io/parquet/decode_fixed.cu
Original file line number Diff line number Diff line change
Expand Up @@ -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; }
Comment thread
mhaseeb123 marked this conversation as resolved.

// 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;
Comment thread
mhaseeb123 marked this conversation as resolved.
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,
Expand All @@ -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<decode_block_size_t>(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<decode_block_size_t, has_lists_t>(
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
Expand Down
116 changes: 116 additions & 0 deletions cpp/src/io/parquet/decode_pruned_pages.cu
Original file line number Diff line number Diff line change
@@ -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 <cudf/detail/utilities/cuda.cuh>

#include <cooperative_groups.h>
#include <cuda/atomic>
#include <cuda/std/algorithm>

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<PageInfo> pages,
Comment thread
mhaseeb123 marked this conversation as resolved.
device_span<ColumnChunkDesc const> chunks,
device_span<bool const> page_mask,
device_span<size_t> 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<size_type>(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<size_type*>(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<size_t, cuda::std::thread_scope_device> initial_str_offset{
initial_str_offsets[input_col_idx]};
initial_str_offset.fetch_min(page.str_offset, cuda::std::memory_order_relaxed);
}
Comment thread
mhaseeb123 marked this conversation as resolved.

// 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<size_type>(page.str_offset);
for (auto row = begin + t; row < end; row += block.size()) {
data[row - skip_rows] = value;
}
return;
}
Comment thread
mhaseeb123 marked this conversation as resolved.

// 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<size_type*>(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<size_type>(block.size())) {
offsets[output_begin + offset_idx] = offset;
}
}
}
}
Comment thread
mhaseeb123 marked this conversation as resolved.

} // namespace

void fill_pruned_offsets(cudf::device_span<PageInfo> pages,
cudf::device_span<ColumnChunkDesc const> chunks,
cudf::device_span<bool const> page_mask,
cudf::device_span<size_t> 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.size(), block_size, 0, stream.value()>>>(
pages, chunks, page_mask, initial_str_offsets, skip_rows, num_rows);
CUDF_CUDA_TRY(cudaGetLastError());
}

} // namespace cudf::io::parquet::detail
2 changes: 1 addition & 1 deletion cpp/src/io/parquet/experimental/hybrid_scan_chunking.cu
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Comment thread
mhaseeb123 marked this conversation as resolved.
generate_list_column_row_counts(is_estimate_row_counts::NO);
} else {
generate_list_column_row_counts(is_estimate_row_counts::YES);
Expand Down
2 changes: 1 addition & 1 deletion cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
17 changes: 9 additions & 8 deletions cpp/src/io/parquet/experimental/hybrid_scan_preprocess.cu
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ void decode_dictionary_page_headers(cudf::detail::hostdevice_span<ColumnChunkDes

parquet::kernel_error error_code(stream);

parquet::detail::decode_page_headers(chunks, chunk_page_info.begin(), error_code.data(), stream);
parquet::detail::decode_page_headers(chunks, chunk_page_info, error_code.data(), stream);
Comment thread
mhaseeb123 marked this conversation as resolved.

if (auto const error = error_code.value_sync(stream); error != 0) {
CUDF_FAIL("Parquet header parsing failed with code(s) " +
Expand Down Expand Up @@ -115,10 +115,11 @@ void hybrid_scan_reader_impl::prepare_row_groups(
_file_itm_data.num_rows_per_source.cend(),
_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()) {
Expand Down Expand Up @@ -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<PageInfo> 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<size_t>(_input_columns.size()),
"Encountered page_offsets / num_columns mismatch");
}
Expand Down
65 changes: 8 additions & 57 deletions cpp/src/io/parquet/page_data.cu
Original file line number Diff line number Diff line change
@@ -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
*/

Expand Down Expand Up @@ -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<cudf::detail::warp_size>(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<int>(block.thread_rank())};

// Setup local page info
Expand All @@ -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<decode_block_size>(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;

Expand Down Expand Up @@ -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<cudf::detail::warp_size>(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<int>(block.thread_rank())};

// Setup local page info
Expand All @@ -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<decode_block_size>(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<size_type*>(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
Expand Down
32 changes: 0 additions & 32 deletions cpp/src/io/parquet/page_decode.cuh
Original file line number Diff line number Diff line change
Expand Up @@ -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 <int block_size>
static __device__ void update_list_offsets_for_pruned_pages(page_state_s* state)
Comment thread
mhaseeb123 marked this conversation as resolved.
{
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<cudf::size_type*>(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.
Expand Down
Loading
Loading