Skip to content

refactor(parquet): introduce level_state_scan struct for preprocess_levels_kernel #23479

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rapids-bot merged 2 commits into from
Jul 31, 2026
Merged

refactor(parquet): introduce level_state_scan struct for preprocess_levels_kernel #23479

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
86856a1
refactor(parquet): extract page_decode_stream_state and narrow prepro…
vyasr Jul 29, 2026
5a5e241
Comments
vyasr Jul 30, 2026
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
22 changes: 13 additions & 9 deletions cpp/src/io/parquet/decode_fixed.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -184,7 +184,7 @@ __device__ inline void decode_fixed_width_split_values(
auto const data_out = s->nesting_info[leaf_level_index].data_out;

Type const dtype = s->setup.col.physical_type;
auto const data_len = cuda::std::distance(s->data_start, s->data_end);
auto const data_len = cuda::std::distance(s->stream.data_start, s->stream.data_end);

// Check malformed BYTE_STREAM_SPLIT pages
if (s->dtype_len_in <= 0 or data_len <= 0) {
Expand Down Expand Up @@ -223,7 +223,7 @@ __device__ inline void decode_fixed_width_split_values(
}();

uint32_t const dtype_len = s->dtype_len;
uint8_t const* const src = s->data_start + src_pos;
uint8_t const* const src = s->stream.data_start + src_pos;
uint8_t* const dst = data_out + static_cast<size_t>(dst_pos) * dtype_len;
auto const is_decimal = s->setup.col.logical_type.has_value() and
s->setup.col.logical_type->type == LogicalType::DECIMAL;
Expand Down Expand Up @@ -806,7 +806,7 @@ inline __device__ void bool_plain_decode(page_state_s* s,
int const byte_offset = bit_pos >> 3;
int const bit_in_byte_index = bit_pos & 7;

uint8_t const* const read_from = s->data_start + byte_offset;
uint8_t const* const read_from = s->stream.data_start + byte_offset;
bool const read_bit = (*read_from) & (1 << bit_in_byte_index);

int const write_to_index = rolling_index<state_buf::dict_buf_size>(bit_pos);
Expand Down Expand Up @@ -1057,20 +1057,24 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size_t, 8)
rle_stream<uint32_t, decode_block_size_t, rolling_buf_size> dict_stream{dict_runs};
if constexpr (has_dict_t) {
dict_stream.init(block,
s->dict_bits,
s->data_start,
s->data_end,
s->stream.dict_bits,
s->stream.data_start,
s->stream.data_end,
sb->dict_idx,
s->setup.page.num_input_values);
}

// Use dictionary stream memory for bools
rle_stream<uint32_t, decode_block_size_t, rolling_buf_size> bool_stream{bool_runs};
bool bools_are_rle_stream = (s->dict_run == 0);
bool bools_are_rle_stream = (s->stream.dict_run == 0);
if constexpr (has_bools_t) {
if (bools_are_rle_stream) {
bool_stream.init(
block, 1, s->data_start, s->data_end, sb->dict_idx, s->setup.page.num_input_values);
bool_stream.init(block,
1,
s->stream.data_start,
s->stream.data_end,
sb->dict_idx,
s->setup.page.num_input_values);
}
}
block.sync();
Expand Down
21 changes: 11 additions & 10 deletions cpp/src/io/parquet/decode_preprocess.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
#include "delta_binary.cuh"
#include "io/utilities/column_buffer.hpp"
#include "page_decode.cuh"
#include "page_state_composed.cuh"
#include "reader_impl_chunking_utils.cuh"

#include <cudf/detail/nvtx/ranges.hpp>
Expand Down Expand Up @@ -390,13 +391,13 @@ CUDF_KERNEL void __launch_bounds__(level_decode_block_size)
size_t min_row,
size_t num_rows)
{
__shared__ __align__(16) page_state_s state_g;
__shared__ __align__(16) level_scan_state state_g;

page_state_s* const s = &state_g;
auto const block = cg::this_thread_block();
int const page_idx = cg::this_grid().block_rank();
int const t = block.thread_rank();
PageInfo* pp = &pages[page_idx];
level_scan_state* const s = &state_g;
auto const block = cg::this_thread_block();
int const page_idx = cg::this_grid().block_rank();
int const t = block.thread_rank();
PageInfo* pp = &pages[page_idx];

// Return early if this page is pruned
if (not page_mask.empty() and not page_mask[page_idx]) { return; }
Expand Down Expand Up @@ -448,8 +449,8 @@ CUDF_KERNEL void __launch_bounds__(level_decode_block_size)
block.sync();
decoders[level_type::REPETITION].init(block,
s->setup.col.level_bits[level_type::REPETITION],
s->abs_lvl_start[level_type::REPETITION],
s->abs_lvl_end[level_type::REPETITION],
s->stream.abs_lvl_start[level_type::REPETITION],
s->stream.abs_lvl_end[level_type::REPETITION],
rep,
num_to_decode,
stage,
Expand All @@ -471,8 +472,8 @@ CUDF_KERNEL void __launch_bounds__(level_decode_block_size)
block.sync();
decoders[level_type::DEFINITION].init(block,
s->setup.col.level_bits[level_type::DEFINITION],
s->abs_lvl_start[level_type::DEFINITION],
s->abs_lvl_end[level_type::DEFINITION],
s->stream.abs_lvl_start[level_type::DEFINITION],
s->stream.abs_lvl_end[level_type::DEFINITION],
def,
num_to_decode,
stage,
Expand Down
10 changes: 5 additions & 5 deletions cpp/src/io/parquet/page_data.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size)
bool const has_repetition = s->setup.col.max_level[level_type::REPETITION] > 0;
bool const process_nulls = should_process_nulls(s);

auto const data_len = cuda::std::distance(s->data_start, s->data_end);
auto const data_len = cuda::std::distance(s->stream.data_start, s->stream.data_end);
auto const num_values = data_len / s->dtype_len_in;

// Check malformed BYTE_STREAM_SPLIT pages
Expand Down Expand Up @@ -162,7 +162,7 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size)
int leaf_level_index = s->setup.col.max_nesting_depth - 1;

uint32_t dtype_len = s->dtype_len;
uint8_t const* src = s->data_start + val_src_pos;
uint8_t const* src = s->stream.data_start + val_src_pos;
uint8_t* dst =
nesting_info_base[leaf_level_index].data_out + static_cast<size_t>(dst_pos) * dtype_len;
auto const is_decimal = s->setup.col.logical_type.has_value() and
Expand Down Expand Up @@ -296,8 +296,8 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size)
int const init_valid_map_offset =
s->nesting_info[s->setup.col.max_nesting_depth - 1].valid_map_offset;

if (s->dict_base) {
out_warp_id = (s->dict_bits > 0) ? 2 : 1;
if (s->stream.dict_base) {
out_warp_id = (s->stream.dict_bits > 0) ? 2 : 1;
} else {
switch (s->setup.col.physical_type) {
case Type::BOOLEAN: [[fallthrough]];
Expand Down Expand Up @@ -351,7 +351,7 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size)
// This is likely a false positive in practice, but could be solved by wrapping the next
// 9 lines in `if (s->dict_pos < src_target_pos) {}`. If that change is made here, it will
// be needed in the other DecodeXXX kernels.
if (s->dict_base) {
if (s->stream.dict_base) {
src_target_pos =
decode_dictionary_indices<is_calc_sizes_only::NO>(s, sb, src_target_pos, warp).first;
} else if (s->setup.col.physical_type == Type::BOOLEAN) {
Expand Down
61 changes: 33 additions & 28 deletions cpp/src/io/parquet/page_data.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,17 +120,18 @@ inline __device__ void read_int96_timestamp(page_state_s* s,
using cuda::std::chrono::duration_cast;

uint8_t const* src8;
uint32_t dict_pos, dict_size = s->dict_size, ofs;
uint32_t dict_pos, dict_size = s->stream.dict_size, ofs;

if (s->dict_base) {
if (s->stream.dict_base) {
// Dictionary
dict_pos =
(s->dict_bits > 0) ? sb->dict_idx[rolling_index<state_buf::dict_buf_size>(src_pos)] : 0;
src8 = s->dict_base;
dict_pos = (s->stream.dict_bits > 0)
? sb->dict_idx[rolling_index<state_buf::dict_buf_size>(src_pos)]
: 0;
src8 = s->stream.dict_base;
} else {
// Plain
dict_pos = src_pos;
src8 = s->data_start;
src8 = s->stream.data_start;
}
dict_pos *= (uint32_t)s->dtype_len_in;
ofs = 3 & reinterpret_cast<size_t>(src8);
Expand Down Expand Up @@ -193,18 +194,19 @@ inline __device__ void read_int64_timestamp(page_state_s* s,
int64_t* dst)
{
uint8_t const* src8;
uint32_t dict_pos, dict_size = s->dict_size, ofs;
uint32_t dict_pos, dict_size = s->stream.dict_size, ofs;
int64_t ts;

if (s->dict_base) {
if (s->stream.dict_base) {
// Dictionary
dict_pos =
(s->dict_bits > 0) ? sb->dict_idx[rolling_index<state_buf::dict_buf_size>(src_pos)] : 0;
src8 = s->dict_base;
dict_pos = (s->stream.dict_bits > 0)
? sb->dict_idx[rolling_index<state_buf::dict_buf_size>(src_pos)]
: 0;
src8 = s->stream.dict_base;
} else {
// Plain
dict_pos = src_pos;
src8 = s->data_start;
src8 = s->stream.data_start;
}
dict_pos *= (uint32_t)s->dtype_len_in;
ofs = 3 & reinterpret_cast<size_t>(src8);
Expand Down Expand Up @@ -268,13 +270,14 @@ __device__ void read_fixed_width_byte_array_as_int(page_state_s* s,
T* dst)
{
uint32_t const dtype_len_in = s->dtype_len_in;
uint8_t const* data = s->dict_base ? s->dict_base : s->data_start;
uint8_t const* data = s->stream.dict_base ? s->stream.dict_base : s->stream.data_start;
uint32_t const pos =
(s->dict_base
? ((s->dict_bits > 0) ? sb->dict_idx[rolling_index<state_buf::dict_buf_size>(src_pos)] : 0)
(s->stream.dict_base
? ((s->stream.dict_bits > 0) ? sb->dict_idx[rolling_index<state_buf::dict_buf_size>(src_pos)]
: 0)
: src_pos) *
dtype_len_in;
uint32_t const dict_size = s->dict_size;
uint32_t const dict_size = s->stream.dict_size;

T unscaled = 0;
for (unsigned int i = 0; i < dtype_len_in; i++) {
Expand Down Expand Up @@ -305,17 +308,18 @@ inline __device__ void read_fixed_width_value_fast(page_state_s* s,
T* dst)
{
uint8_t const* dict;
uint32_t dict_pos, dict_size = s->dict_size;
uint32_t dict_pos, dict_size = s->stream.dict_size;

if (s->dict_base) {
if (s->stream.dict_base) {
// Dictionary
dict_pos =
(s->dict_bits > 0) ? sb->dict_idx[rolling_index<state_buf::dict_buf_size>(src_pos)] : 0;
dict = s->dict_base;
dict_pos = (s->stream.dict_bits > 0)
? sb->dict_idx[rolling_index<state_buf::dict_buf_size>(src_pos)]
: 0;
dict = s->stream.dict_base;
} else {
// Plain
dict_pos = src_pos;
dict = s->data_start;
dict = s->stream.data_start;
}
dict_pos *= (uint32_t)s->dtype_len_in;
gpuStoreOutput(dst, dict, dict_pos, dict_size);
Expand All @@ -335,17 +339,18 @@ inline __device__ void read_nbyte_fixed_width_value(
page_state_s* s, state_buf* sb, int src_pos, uint8_t* dst8, int len)
{
uint8_t const* dict;
uint32_t dict_pos, dict_size = s->dict_size;
uint32_t dict_pos, dict_size = s->stream.dict_size;

if (s->dict_base) {
if (s->stream.dict_base) {
// Dictionary
dict_pos =
(s->dict_bits > 0) ? sb->dict_idx[rolling_index<state_buf::dict_buf_size>(src_pos)] : 0;
dict = s->dict_base;
dict_pos = (s->stream.dict_bits > 0)
? sb->dict_idx[rolling_index<state_buf::dict_buf_size>(src_pos)]
: 0;
dict = s->stream.dict_base;
} else {
// Plain
dict_pos = src_pos;
dict = s->data_start;
dict = s->stream.data_start;
}
dict_pos *= (uint32_t)s->dtype_len_in;
if (len & 3) {
Expand Down
Loading
Loading