Skip to content

refactor(parquet): introduce string_offset_scan_state for preprocess_string_offsets_kernel #23495

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

Open
vyasr wants to merge 3 commits into
base: main
Choose a base branch
from
+199 −166
Open

refactor(parquet): introduce string_offset_scan_state for preprocess_string_offsets_kernel #23495

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4f6848d
refactor(parquet): extract page_decode_progress_state and narrow prep…
vyasr Jul 29, 2026
43ca9a2
refactor(parquet): address string offset scan review comments
vyasr Jul 31, 2026
03f9c78
Update struct description
vyasr Jul 31, 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
65 changes: 33 additions & 32 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 @@ -340,7 +340,7 @@ __device__ int update_validity_and_row_indices_nested(
constexpr int max_batch_size = num_warps * cudf::detail::warp_size;

// how many (input) values we've processed in the page so far
int value_count = s->input_value_count;
int value_count = s->progress.input_value_count;

// cap by last row so that we don't process any rows past what we want to output.
int const first_row = s->setup.first_row;
Expand All @@ -363,7 +363,7 @@ __device__ int update_validity_and_row_indices_nested(
int const block_value_count = batch_size;

// compute our row index, whether we're in row bounds, and validity
// This ASSUMES that s->row_index_lower_bound is always -1!
// This ASSUMES that s->progress.row_index_lower_bound is always -1!
// Its purpose is to handle rows than span page boundaries, which only happen for lists.
int const row_index = thread_value_count + value_count;
int const in_row_bounds = (row_index < last_row);
Expand Down Expand Up @@ -428,11 +428,11 @@ __device__ int update_validity_and_row_indices_nested(

if (t == 0) {
// update valid value count for decoding and total # of values we've processed
max_depth_ni.valid_count = max_depth_valid_count;
max_depth_ni.value_count = value_count; // Needed AT LEAST for strings!
s->nz_count = max_depth_valid_count;
s->input_value_count = value_count;
s->input_row_count = value_count;
max_depth_ni.valid_count = max_depth_valid_count;
max_depth_ni.value_count = value_count; // Needed AT LEAST for strings!
s->progress.nz_count = max_depth_valid_count;
s->progress.input_value_count = value_count;
s->progress.input_row_count = value_count;
}

return max_depth_valid_count;
Expand Down Expand Up @@ -463,7 +463,7 @@ __device__ int update_validity_and_row_indices_flat(
auto& ni = s->nesting_info[0];

// how many (input) values we've processed in the page so far
int value_count = s->input_value_count;
int value_count = s->progress.input_value_count;
int valid_count = ni.valid_count;

// cap by last row so that we don't process any rows past what we want to output.
Expand All @@ -482,7 +482,7 @@ __device__ int update_validity_and_row_indices_flat(
int const block_value_count = batch_size;

// compute our row index, whether we're in row bounds, and validity
// This ASSUMES that s->row_index_lower_bound is always -1!
// This ASSUMES that s->progress.row_index_lower_bound is always -1!
// Its purpose is to handle rows than span page boundaries, which only happen for lists.
int const row_index = thread_value_count + value_count;
int const in_row_bounds = (row_index < last_row);
Expand Down Expand Up @@ -538,11 +538,11 @@ __device__ int update_validity_and_row_indices_flat(

if (t == 0) {
// update valid value count for decoding and total # of values we've processed
ni.valid_count = valid_count;
ni.value_count = value_count;
s->nz_count = valid_count;
s->input_value_count = value_count;
s->input_row_count = value_count;
ni.valid_count = valid_count;
ni.value_count = value_count;
s->progress.nz_count = valid_count;
s->progress.input_value_count = value_count;
s->progress.input_row_count = value_count;
}

return valid_count;
Expand Down Expand Up @@ -576,16 +576,16 @@ __device__ int update_validity_and_row_indices_lists(int32_t target_value_count,
constexpr int max_batch_size = num_warps * cudf::detail::warp_size;

// how many (input) values we've processed in the page so far, prior to this loop iteration
int value_count = s->input_value_count;
int value_count = s->progress.input_value_count;

// how many rows we've processed in the page so far
int input_row_count = s->input_row_count;
int input_row_count = s->progress.input_row_count;

// cap by last row so that we don't process any rows past what we want to output.
int const first_row = s->setup.first_row;
int const last_row = first_row + s->setup.num_rows;

int const row_index_lower_bound = s->row_index_lower_bound;
int const row_index_lower_bound = s->progress.row_index_lower_bound;
int const max_depth = s->setup.col.max_nesting_depth - 1;
int max_depth_valid_count = s->nesting_info[max_depth].valid_count;

Expand Down Expand Up @@ -781,11 +781,11 @@ __device__ int update_validity_and_row_indices_lists(int32_t target_value_count,
if (t == 0) {
// update valid value count for decoding and total # of values we've processed
s->nesting_info[max_depth].valid_count = max_depth_valid_count;
s->nz_count = max_depth_valid_count;
s->input_value_count = value_count;
s->progress.nz_count = max_depth_valid_count;
s->progress.input_value_count = value_count;

// If we have lists # rows != # values
s->input_row_count = input_row_count;
s->progress.input_row_count = input_row_count;
}

return max_depth_valid_count;
Expand All @@ -797,7 +797,7 @@ inline __device__ void bool_plain_decode(page_state_s* s,
int target_pos,
thread_group const& group)
{
int const pos = s->dict_pos;
int const pos = s->progress.dict_pos;
int const t = group.thread_rank();
// Ensure all threads have the dict_pos
group.sync();
Expand Down Expand Up @@ -856,7 +856,7 @@ __device__ void skip_ahead_in_decoding(page_state_s* s,
if (bools_are_rle_stream) {
skip_decode<rolling_buf_size>(bool_stream, num_to_skip, t);
} else {
if (t == 0) { s->dict_pos = num_to_skip; }
if (t == 0) { s->progress.dict_pos = num_to_skip; }
// For non-lists don't sync: we'll sync at the end of the lambda instead
if constexpr (has_lists_t) { block.sync(); }
}
Expand Down Expand Up @@ -897,11 +897,11 @@ __device__ void skip_ahead_in_decoding(page_state_s* s,
auto& ni = s->nesting_info[max_depth];

// update valid value count for decoding and total # of values we've processed
ni.valid_count = valid_count;
ni.value_count = processed_count;
s->nz_count = valid_count;
s->input_value_count = processed_count;
s->input_row_count = processed_count;
ni.valid_count = valid_count;
ni.value_count = processed_count;
s->progress.nz_count = valid_count;
s->progress.input_value_count = processed_count;
s->progress.input_row_count = processed_count;
}
block.sync();
}
Expand Down Expand Up @@ -1114,7 +1114,7 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size_t, 8)
int const first_row = s->setup.first_row;
int const last_row = first_row + s->setup.num_rows;
while ((s->setup.error == 0) && (processed_count < s->setup.page.num_input_values) &&
(s->input_row_count <= last_row)) {
(s->progress.input_row_count <= last_row)) {
int next_valid_count;
block.sync();
processed_count += min(rolling_buf_size, s->setup.page.num_input_values - processed_count);
Expand Down Expand Up @@ -1143,10 +1143,10 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size_t, 8)
processed_count, s, sb, nullptr, rep, t);
} else {
// direct copy: no nulls, no lists, no need to update validity or row indices
// This ASSUMES that s->row_index_lower_bound is always -1!
// This ASSUMES that s->progress.row_index_lower_bound is always -1!
// Its purpose is to handle rows than span page boundaries, which only happen for lists.
int const capped_target_value_count = min(processed_count, last_row);
if (t == 0) { s->input_row_count = capped_target_value_count; }
if (t == 0) { s->progress.input_row_count = capped_target_value_count; }
next_valid_count = capped_target_value_count;
}
}
Expand All @@ -1164,7 +1164,7 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size_t, 8)
} else {
auto const target_pos = next_valid_count + s->setup.page.skipped_leaf_values;
bool_plain_decode(s, sb, target_pos, block);
if (t == 0) { s->dict_pos = target_pos; }
if (t == 0) { s->progress.dict_pos = target_pos; }
}
block.sync();
}
Expand Down Expand Up @@ -1222,7 +1222,8 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size_t, 8)
if constexpr (!has_lists_t) {
if (!process_nulls) {
if (t == 0) {
s->nesting_info[s->setup.col.max_nesting_depth - 1].value_count = s->input_row_count;
s->nesting_info[s->setup.col.max_nesting_depth - 1].value_count =
s->progress.input_row_count;
}
block.sync();
}
Expand Down
30 changes: 14 additions & 16 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 @@ -69,11 +69,11 @@ __device__ void update_page_sizes(page_state_s* s,
} temp_storage;

// how many input level values we've processed in the page so far
int value_count = s->input_value_count;
int value_count = s->progress.input_value_count;
// how many rows we've processed in the page so far
int row_count = s->input_row_count;
int row_count = s->progress.input_row_count;
// how many leaf values we've processed in the page so far
int leaf_count = s->input_leaf_count;
int leaf_count = s->progress.input_leaf_count;
// whether or not we need to continue checking for the first row
bool skipped_values_set = s->setup.page.skipped_values >= 0;

Expand Down Expand Up @@ -108,7 +108,7 @@ __device__ void update_page_sizes(page_state_s* s,

// if this thread is in row bounds
int const row_index = (thread_row_count + row_count) - 1;
in_row_bounds = (row_index >= s->row_index_lower_bound) &&
in_row_bounds = (row_index >= s->progress.row_index_lower_bound) &&
(row_index < (s->setup.first_row + s->setup.num_rows));

// if we have not set skipped values yet, see if we found the first in-bounds row
Expand Down Expand Up @@ -150,11 +150,11 @@ __device__ void update_page_sizes(page_state_s* s,

// update final outputs
if (!t) {
s->input_value_count = value_count;
s->progress.input_value_count = value_count;

// only used in the skip_rows/num_rows case
s->input_leaf_count = leaf_count;
s->input_row_count = row_count;
s->progress.input_leaf_count = leaf_count;
s->progress.input_row_count = row_count;
}

block.sync();
Expand Down Expand Up @@ -288,17 +288,15 @@ CUDF_KERNEL void __launch_bounds__(preprocess_block_size)

// in the trim pass, for anything with lists, we only need to fully process bounding pages (those
// at the beginning or the end of the row bounds)
if (!is_base_pass &&
!is_bounds_page(s->setup.page, s->setup.col.start_row, min_row, num_rows, has_repetition)) {
if (!is_base_pass && !is_bounds_page(s, min_row, num_rows, has_repetition)) {
int depth = 0;
while (depth < s->setup.page.num_output_nesting_levels) {
auto const thread_depth = depth + t;
if (thread_depth < s->setup.page.num_output_nesting_levels) {
// if we are not a bounding page (as checked above) then we are either
// returning all rows/values from this page, or 0 of them
pp->nesting[thread_depth].batch_size =
(s->setup.num_rows == 0 &&
!is_page_contained(s->setup.page, s->setup.col.start_row, min_row, num_rows))
(s->setup.num_rows == 0 && !is_page_contained(s, min_row, num_rows))
? 0
: pp->nesting[thread_depth].size;
}
Expand Down Expand Up @@ -326,15 +324,15 @@ CUDF_KERNEL void __launch_bounds__(preprocess_block_size)
if (!t) {
s->setup.page.skipped_values = -1;
s->setup.page.skipped_leaf_values = 0;
s->input_row_count = 0;
s->input_value_count = 0;
s->progress.input_row_count = 0;
s->progress.input_value_count = 0;

// in the base pass, we're computing the number of rows, make sure we visit absolutely
// everything
if (is_base_pass) {
s->setup.first_row = 0;
s->setup.num_rows = cuda::std::numeric_limits<int32_t>::max();
s->row_index_lower_bound = -1;
s->setup.first_row = 0;
s->setup.num_rows = cuda::std::numeric_limits<int32_t>::max();
s->progress.row_index_lower_bound = -1;
}
}

Expand Down
46 changes: 26 additions & 20 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 @@ -108,16 +108,17 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size)

// skipped_leaf_values will always be 0 for flat hierarchies.
uint32_t skipped_leaf_values = s->setup.page.skipped_leaf_values;
while (s->setup.error == 0 &&
(s->input_value_count < s->setup.num_input_values || s->src_pos < s->nz_count)) {
while (s->setup.error == 0 && (s->progress.input_value_count < s->setup.num_input_values ||
s->progress.src_pos < s->progress.nz_count)) {
int target_pos;
int src_pos = s->src_pos;
int src_pos = s->progress.src_pos;

if (warp.meta_group_rank() == 0) {
target_pos = cuda::std::min(src_pos + 2 * (decode_block_size - warp.size()),
s->nz_count + (decode_block_size - warp.size()));
s->progress.nz_count + (decode_block_size - warp.size()));
} else {
target_pos = cuda::std::min<int32_t>(s->nz_count, src_pos + decode_block_size - warp.size());
target_pos =
cuda::std::min<int32_t>(s->progress.nz_count, src_pos + decode_block_size - warp.size());
}
// This needs to be here to prevent warp 1 modifying src_pos before all threads have read it
block.sync();
Expand Down Expand Up @@ -213,7 +214,9 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size)
}
}
// Only the first thread in the warp 1 updates src_pos
if (warp.meta_group_rank() == 1 and warp.thread_rank() == 0) { s->src_pos = target_pos; }
if (warp.meta_group_rank() == 1 and warp.thread_rank() == 0) {
s->progress.src_pos = target_pos;
}
}
block.sync();
}
Expand Down Expand Up @@ -321,18 +324,21 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size)
auto const first_out_thread_id = out_warp_id * warp.size();
// skipped_leaf_values will always be 0 for flat hierarchies.
uint32_t skipped_leaf_values = s->setup.page.skipped_leaf_values;
while (s->setup.error == 0 &&
(s->input_value_count < s->setup.num_input_values || s->src_pos < s->nz_count)) {
while (s->setup.error == 0 && (s->progress.input_value_count < s->setup.num_input_values ||
s->progress.src_pos < s->progress.nz_count)) {
int target_pos;
int src_pos = s->src_pos;
int src_pos = s->progress.src_pos;

if (warp.meta_group_rank() < out_warp_id) {
target_pos = cuda::std::min<int32_t>(src_pos + 2 * (decode_block_size - first_out_thread_id),
s->nz_count + (decode_block_size - first_out_thread_id));
} else {
target_pos =
cuda::std::min<int32_t>(s->nz_count, src_pos + decode_block_size - first_out_thread_id);
if (out_warp_id > 1) { target_pos = cuda::std::min<int32_t>(target_pos, s->dict_pos); }
cuda::std::min<int32_t>(src_pos + 2 * (decode_block_size - first_out_thread_id),
s->progress.nz_count + (decode_block_size - first_out_thread_id));
} else {
target_pos = cuda::std::min<int32_t>(s->progress.nz_count,
src_pos + decode_block_size - first_out_thread_id);
if (out_warp_id > 1) {
target_pos = cuda::std::min<int32_t>(target_pos, s->progress.dict_pos);
}
}
// this needs to be here to prevent warp 3 modifying src_pos before all threads have read it
block.sync();
Expand All @@ -347,10 +353,10 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size)
uint32_t src_target_pos = target_pos + skipped_leaf_values;

// WARP1: Decode dictionary indices, booleans or string positions
// NOTE: racecheck complains of a RAW error involving the s->dict_pos assignment below.
// 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.
// NOTE: racecheck complains of a RAW error involving the s->progress.dict_pos assignment
// below. This is likely a false positive in practice, but could be solved by wrapping the
// next 9 lines in `if (s->progress.dict_pos < src_target_pos) {}`. If that change is made
// here, it will be needed in the other DecodeXXX kernels.
if (s->stream.dict_base) {
src_target_pos =
decode_dictionary_indices<is_calc_sizes_only::NO>(s, sb, src_target_pos, warp).first;
Expand All @@ -360,7 +366,7 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size)
s->setup.col.physical_type == Type::FIXED_LEN_BYTE_ARRAY) {
initialize_string_descriptors<is_calc_sizes_only::NO>(s, sb, src_target_pos, warp);
}
if (warp.thread_rank() == 0) { s->dict_pos = src_target_pos; }
if (warp.thread_rank() == 0) { s->progress.dict_pos = src_target_pos; }
} else {
// WARP1..WARP3: Decode values
src_pos += block.thread_rank() - first_out_thread_id;
Expand Down Expand Up @@ -457,7 +463,7 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size)
}
}

if (block.thread_rank() == first_out_thread_id) { s->src_pos = target_pos; }
if (block.thread_rank() == first_out_thread_id) { s->progress.src_pos = target_pos; }
}
__syncthreads();
}
Expand Down
Loading
Loading