From c39becb7ccc3a35023c1f70cbf04885fa84dcadd Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Thu, 30 Apr 2026 21:46:54 +0000 Subject: [PATCH 01/20] Reduce ORC writer device allocations via per-phase arenas `encode_columns` and `gather_stripes` previously allocated one `device_uvector` per (stripe, stream) pair, producing ~1,000 driver allocations each per write_orc call. Replace both patterns with a single arena per phase plus non-owning `device_span` views, holding the encoded and gathered arenas on the `encoded_data` struct for lifetime management. Each per-(stripe, stream) region inside an arena starts on a 256-byte boundary (`rmm::CUDA_ALLOCATION_ALIGNMENT`) to match the alignment that RMM previously gave each independent allocation: the in-loop alignment fix-up only honors `compress_required_chunk_alignment(compression)`, which is 1 for `compression == NONE` and as low as 8 for nvcomp paths, so the encoder/codec kernels rely on the underlying allocation being naturally aligned. Effect on `orc_write_encode` (INTEGRAL_SIGNED, 512 MB): * `cudaMallocFromPoolAsync` per iter: ~2,135 -> ~88 (-96%). * pool GPU time: 53.95 ms -> 48.23 ms (-10.6%). * async GPU time: 104.65 ms -> 48.52 ms (-53.6%). * async-vs-pool delta: +94% -> +0.6% (within noise). * peak memory: +1 MiB on a 1.97 GiB run (alignment padding). * encoded_file_size unchanged. Made-with: Cursor --- cpp/src/io/orc/writer_impl.cu | 265 +++++++++++++++++++++++---------- cpp/src/io/orc/writer_impl.hpp | 13 +- 2 files changed, 196 insertions(+), 82 deletions(-) diff --git a/cpp/src/io/orc/writer_impl.cu b/cpp/src/io/orc/writer_impl.cu index df29dc85300d..60fb5d25d42b 100644 --- a/cpp/src/io/orc/writer_impl.cu +++ b/cpp/src/io/orc/writer_impl.cu @@ -938,13 +938,15 @@ encoded_data encode_columns(orc_table_view const& orc_table, hostdevice_2dvector chunk_streams( num_columns, segmentation.num_rowgroups(), stream); - // per-stripe, per-stream owning buffers - std::vector>> encoded_data(segmentation.num_stripes()); - for (auto const& stripe : segmentation.stripes) { - std::generate_n(std::back_inserter(encoded_data[stripe.id]), streams.size(), [stream]() { - return rmm::device_uvector(0, stream); - }); + // Pass 1: compute per-rowgroup stream lengths and per-(stripe, strm_id) sizes. + // The encoded buffers for every (stripe, stream) pair are packed into a single + // arena allocation, so accumulate the total here and assign offsets later. + auto const num_streams = streams.size(); + std::vector> stripe_strm_sizes(segmentation.num_stripes(), + std::vector(num_streams, 0)); + size_t encoded_total = 0; + for (auto const& stripe : segmentation.stripes) { for (size_t col_idx = 0; col_idx < num_columns; col_idx++) { for (int strm_type = 0; strm_type < CI_NUM_STREAMS; ++strm_type) { auto const& column = orc_table.column(col_idx); @@ -956,51 +958,97 @@ encoded_data encode_columns(orc_table_view const& orc_table, col_streams[rg_idx].lengths[strm_type] = 0; }); - // Calculate rowgroup sizes and stripe size - if (strm_id >= 0) { - size_t stripe_size = 0; - std::for_each(stripe.cbegin(), stripe.cend(), [&](auto rg_idx) { + if (strm_id < 0) { continue; } + + size_t stripe_size = 0; + std::for_each(stripe.cbegin(), stripe.cend(), [&](auto rg_idx) { #if defined(__GNUC__) && (__GNUC__ >= 14) #pragma GCC diagnostic push #pragma GCC diagnostic ignored "-Wdangling-reference" #endif - auto const& ck = chunks[col_idx][rg_idx]; + auto const& ck = chunks[col_idx][rg_idx]; #if defined(__GNUC__) && (__GNUC__ >= 14) #pragma GCC diagnostic pop #endif - auto& strm = col_streams[rg_idx]; + auto& strm = col_streams[rg_idx]; - if ((strm_type == CI_DICTIONARY) || - (strm_type == CI_DATA2 && ck.encoding_kind == DICTIONARY_V2)) { - if (rg_idx == *stripe.cbegin()) { - auto const stripe_dict = column.host_stripe_dict(stripe.id); - strm.lengths[strm_type] = - (strm_type == CI_DICTIONARY) - ? stripe_dict.char_count - : (((stripe_dict.entry_count + 0x1ff) >> 9) * (512 * 4 + 2)); - } else { - strm.lengths[strm_type] = 0; - } - } else if (strm_type == CI_DATA && ck.type_kind == TypeKind::STRING && - ck.encoding_kind == DIRECT_V2) { - strm.lengths[strm_type] = std::max(column.rowgroup_char_count(rg_idx), 1); - } else if (strm_type == CI_DATA && streams[strm_id].length == 0 && - (ck.type_kind == DOUBLE || ck.type_kind == FLOAT)) { - // Pass-through - strm.lengths[strm_type] = ck.num_rows * ck.dtype_len; - } else if (ck.type_kind == DECIMAL && strm_type == CI_DATA) { - strm.lengths[strm_type] = dec_chunk_sizes.rg_sizes.at(col_idx)[rg_idx]; + if ((strm_type == CI_DICTIONARY) || + (strm_type == CI_DATA2 && ck.encoding_kind == DICTIONARY_V2)) { + if (rg_idx == *stripe.cbegin()) { + auto const stripe_dict = column.host_stripe_dict(stripe.id); + strm.lengths[strm_type] = + (strm_type == CI_DICTIONARY) + ? stripe_dict.char_count + : (((stripe_dict.entry_count + 0x1ff) >> 9) * (512 * 4 + 2)); } else { - strm.lengths[strm_type] = rle_stream_size(streams.type(strm_id), ck.num_rows); + strm.lengths[strm_type] = 0; } - // Allow extra space for alignment - stripe_size += strm.lengths[strm_type] + uncomp_block_align - 1; - }); + } else if (strm_type == CI_DATA && ck.type_kind == TypeKind::STRING && + ck.encoding_kind == DIRECT_V2) { + strm.lengths[strm_type] = std::max(column.rowgroup_char_count(rg_idx), 1); + } else if (strm_type == CI_DATA && streams[strm_id].length == 0 && + (ck.type_kind == DOUBLE || ck.type_kind == FLOAT)) { + // Pass-through + strm.lengths[strm_type] = ck.num_rows * ck.dtype_len; + } else if (ck.type_kind == DECIMAL && strm_type == CI_DATA) { + strm.lengths[strm_type] = dec_chunk_sizes.rg_sizes.at(col_idx)[rg_idx]; + } else { + strm.lengths[strm_type] = rle_stream_size(streams.type(strm_id), ck.num_rows); + } + // Allow extra space for alignment + stripe_size += strm.lengths[strm_type] + uncomp_block_align - 1; + }); - encoded_data[stripe.id][strm_id] = rmm::device_uvector(stripe_size, stream); - } + stripe_strm_sizes[stripe.id][strm_id] = stripe_size; + encoded_total += stripe_size; + } + } + } + + // Build per-(stripe, stream) offsets into a single arena. Each non-empty + // region starts on a 256-byte boundary to match the alignment that RMM + // guaranteed for the original per-region device_uvector allocations -- the + // ORC encoder kernels and downstream compressors rely on natural alignment + // (uncomp_block_align is only 1 for compression == NONE, so we cannot rely + // on the per-rg alignment fix-up alone). + constexpr size_t region_alignment = rmm::CUDA_ALLOCATION_ALIGNMENT; + std::vector> stripe_strm_offsets(segmentation.num_stripes(), + std::vector(num_streams, 0)); + size_t arena_total = 0; + for (size_t s = 0; s < segmentation.num_stripes(); ++s) { + for (size_t strm_id = 0; strm_id < num_streams; ++strm_id) { + auto const sz = stripe_strm_sizes[s][strm_id]; + if (sz == 0) { + stripe_strm_offsets[s][strm_id] = arena_total; // empty span: data() is harmless + continue; + } + arena_total = util::round_up_unsafe(arena_total, region_alignment); + stripe_strm_offsets[s][strm_id] = arena_total; + arena_total += sz; + } + } + + // Single arena allocation for all (stripe, stream) encoded buffers. + rmm::device_uvector encoded_buffer(arena_total, stream); + + std::vector>> encoded_views(segmentation.num_stripes()); + for (size_t s = 0; s < segmentation.num_stripes(); ++s) { + encoded_views[s].resize(num_streams); + for (size_t strm_id = 0; strm_id < num_streams; ++strm_id) { + encoded_views[s][strm_id] = device_span{ + encoded_buffer.data() + stripe_strm_offsets[s][strm_id], stripe_strm_sizes[s][strm_id]}; + } + } + + // Pass 2: write per-chunk data_ptrs and apply alignment fix-up. The lengths + // computed in pass 1 (now stored in `chunk_streams[col_idx][rg_idx].lengths`) + // are read here as-is. + for (auto const& stripe : segmentation.stripes) { + for (size_t col_idx = 0; col_idx < num_columns; col_idx++) { + for (int strm_type = 0; strm_type < CI_NUM_STREAMS; ++strm_type) { + auto col_streams = chunk_streams[col_idx]; + auto const strm_id = streams.id(col_idx * CI_NUM_STREAMS + strm_type); - // Set offsets for (auto rg_idx_it = stripe.cbegin(); rg_idx_it < stripe.cend(); ++rg_idx_it) { auto const rg_idx = *rg_idx_it; #if defined(__GNUC__) && (__GNUC__ >= 14) @@ -1019,10 +1067,10 @@ encoded_data encode_columns(orc_table_view const& orc_table, } else { if ((strm_type == CI_DICTIONARY) || (strm_type == CI_DATA2 && ck.encoding_kind == DICTIONARY_V2)) { - strm.data_ptrs[strm_type] = encoded_data[stripe.id][strm_id].data(); + strm.data_ptrs[strm_type] = encoded_views[stripe.id][strm_id].data(); } else { strm.data_ptrs[strm_type] = (rg_idx_it == stripe.cbegin()) - ? encoded_data[stripe.id][strm_id].data() + ? encoded_views[stripe.id][strm_id].data() : (col_streams[rg_idx - 1].data_ptrs[strm_type] + col_streams[rg_idx - 1].lengths[strm_type]); } @@ -1055,7 +1103,10 @@ encoded_data encode_columns(orc_table_view const& orc_table, } chunk_streams.device_to_host(stream); - return {std::move(encoded_data), std::move(chunk_streams)}; + return {std::move(encoded_buffer), + rmm::device_uvector{0, stream}, // gathered_buffer (filled by gather_stripes) + std::move(encoded_views), + std::move(chunk_streams)}; } // TODO: remove StripeInformation from this function and return strm_desc instead @@ -1078,47 +1129,91 @@ std::vector gather_stripes(size_t num_index_streams, { if (segmentation.num_stripes() == 0) { return {}; } - // gathered stripes - per-stripe, per-stream (same as encoded_data.data) - std::vector>> gathered_stripes(enc_data->data.size()); - for (auto& stripe_data : gathered_stripes) { - std::generate_n(std::back_inserter(stripe_data), enc_data->data[0].size(), [&]() { - return rmm::device_uvector(0, stream); - }); + auto const num_streams_in_data = enc_data->data[0].size(); + + // Pass 1: compute per-(stripe, stream) actual sizes and decide which need a tight + // gathered copy. Don't copy the data to an exactly sized buffer when only one + // chunk is present, to avoid the overhead of the additional copy. When there + // are multiple chunks, they are copied anyway to make them contiguous. + struct gather_info { + size_t actual_size; + bool gathered; + }; + std::vector> gather_meta( + segmentation.num_stripes(), + std::vector(num_streams_in_data, gather_info{0, false})); + + size_t gather_total = 0; + for (auto const& stripe : segmentation.stripes) { + for (size_t col_idx = 0; col_idx < enc_data->streams.size().first; col_idx++) { + auto const& col_streams = (enc_data->streams)[col_idx]; + for (int k = 0; k < CI_INDEX; k++) { + auto const stream_id = col_streams[0].ids[k]; + if (stream_id == -1) { continue; } + + auto const actual_stripe_size = + std::accumulate(col_streams.begin() + stripe.first, + col_streams.begin() + stripe.first + stripe.size, + 0ul, + [&](auto const& sum, auto const& strm) { return sum + strm.lengths[k]; }); + + auto const allocated_stripe_size = enc_data->data[stripe.id][stream_id].size(); + CUDF_EXPECTS(allocated_stripe_size >= actual_stripe_size, + "Internal ORC writer error: insufficient allocation size for encoded data"); + + bool const gathered = (stripe.size > 1 and allocated_stripe_size > actual_stripe_size); + gather_meta[stripe.id][stream_id] = {actual_stripe_size, gathered}; + if (gathered) { gather_total += actual_stripe_size; } + } + } } + + // Lay out gather destinations within a single arena. As with the encoded + // arena above, each region starts on a 256-byte boundary to match the + // alignment RMM provided for the original per-region device_uvectors. + constexpr size_t region_alignment = rmm::CUDA_ALLOCATION_ALIGNMENT; + std::vector> gather_offsets(segmentation.num_stripes(), + std::vector(num_streams_in_data, 0)); + { + size_t cursor = 0; + for (size_t s = 0; s < segmentation.num_stripes(); ++s) { + for (size_t strm_id = 0; strm_id < num_streams_in_data; ++strm_id) { + if (!gather_meta[s][strm_id].gathered) { continue; } + cursor = util::round_up_unsafe(cursor, region_alignment); + gather_offsets[s][strm_id] = cursor; + cursor += gather_meta[s][strm_id].actual_size; + } + } + gather_total = cursor; + } + rmm::device_uvector gather_buffer(gather_total, stream); + + // Pass 2: build strm_desc entries and record gather destination spans. + std::vector>> gather_views( + segmentation.num_stripes(), + std::vector>(num_streams_in_data, device_span{})); std::vector stripes(segmentation.num_stripes()); for (auto const& stripe : segmentation.stripes) { for (size_t col_idx = 0; col_idx < enc_data->streams.size().first; col_idx++) { auto const& col_streams = (enc_data->streams)[col_idx]; - // Assign stream data of column data stream(s) for (int k = 0; k < CI_INDEX; k++) { auto const stream_id = col_streams[0].ids[k]; - if (stream_id != -1) { - auto const actual_stripe_size = std::accumulate( - col_streams.begin() + stripe.first, - col_streams.begin() + stripe.first + stripe.size, - 0ul, - [&](auto const& sum, auto const& strm) { return sum + strm.lengths[k]; }); - - auto const& allocated_stripe_size = enc_data->data[stripe.id][stream_id].size(); - CUDF_EXPECTS(allocated_stripe_size >= actual_stripe_size, - "Internal ORC writer error: insufficient allocation size for encoded data"); - // Allocate buffers of the exact size as encoded data, smaller than the original buffers. - // Don't copying the data to exactly sized buffer when only one chunk is present to avoid - // performance overhead from the additional copy. When there are multiple chunks, they are - // copied anyway, to make them contiguous (i.e. gather them). - if (stripe.size > 1 and allocated_stripe_size > actual_stripe_size) { - gathered_stripes[stripe.id][stream_id] = - rmm::device_uvector(actual_stripe_size, stream); - } + if (stream_id == -1) { continue; } - auto* ss = &(*strm_desc)[stripe.id][stream_id - num_index_streams]; - ss->data_ptr = gathered_stripes[stripe.id][stream_id].data(); - ss->stream_size = actual_stripe_size; - ss->first_chunk_id = stripe.first; - ss->num_chunks = stripe.size; - ss->column_id = col_idx; - ss->stream_type = k; + auto const& meta = gather_meta[stripe.id][stream_id]; + uint8_t* dst_ptr = nullptr; + if (meta.gathered) { + dst_ptr = gather_buffer.data() + gather_offsets[stripe.id][stream_id]; + gather_views[stripe.id][stream_id] = device_span{dst_ptr, meta.actual_size}; } + + auto* ss = &(*strm_desc)[stripe.id][stream_id - num_index_streams]; + ss->data_ptr = dst_ptr; // null when not gathered; init_batched_memcpy_kernel skips + ss->stream_size = meta.actual_size; + ss->first_chunk_id = stripe.first; + ss->num_chunks = stripe.size; + ss->column_id = col_idx; + ss->stream_type = k; } } @@ -1134,14 +1229,21 @@ std::vector gather_stripes(size_t num_index_streams, strm_desc->device_to_host_async(stream); enc_data->streams.device_to_host(stream); - // move the gathered stripes to encoded_data.data for lifetime management - for (auto stripe_id = 0ul; stripe_id < enc_data->data.size(); ++stripe_id) { - for (auto stream_id = 0ul; stream_id < enc_data->data[0].size(); ++stream_id) { - if (not gathered_stripes[stripe_id][stream_id].is_empty()) - enc_data->data[stripe_id][stream_id] = std::move(gathered_stripes[stripe_id][stream_id]); + // Replace data views for gathered (stripe, stream) with the gathered-arena + // spans, so consumers that read enc_data->data observe the post-gather state. + for (size_t stripe_id = 0; stripe_id < enc_data->data.size(); ++stripe_id) { + for (size_t stream_id = 0; stream_id < num_streams_in_data; ++stream_id) { + if (gather_meta[stripe_id][stream_id].gathered) { + enc_data->data[stripe_id][stream_id] = gather_views[stripe_id][stream_id]; + } } } + // Hold the gathered arena for lifetime management. The encoded arena stays + // alive too: compress_orc_data_streams reads encoded data via per-rowgroup + // data_ptrs that for non-gathered (stripe, stream) still point into it. + enc_data->gathered_buffer = std::move(gather_buffer); + return stripes; } @@ -2453,7 +2555,10 @@ auto convert_table_to_orc_data(table_view const& input, comp_results, stream); - // deallocate encoded data as it is not needed anymore + // deallocate encoded data as it is not needed anymore. Free both arenas + // (encoded + gathered) and clear the spans that referenced them. + enc_data.encoded_buffer = rmm::device_uvector{0, stream}; + enc_data.gathered_buffer = rmm::device_uvector{0, stream}; enc_data.data.clear(); strm_descs.device_to_host_async(stream); diff --git a/cpp/src/io/orc/writer_impl.hpp b/cpp/src/io/orc/writer_impl.hpp index 0454cbea61f0..cf1523791449 100644 --- a/cpp/src/io/orc/writer_impl.hpp +++ b/cpp/src/io/orc/writer_impl.hpp @@ -97,10 +97,19 @@ struct file_segmentation { /** * @brief ORC per-chunk streams of encoded data. + * + * The encoded buffers for every (stripe, stream) pair are packed into two arena + * allocations rather than one `device_uvector` per pair: `encoded_buffer` holds + * the raw encoder output, and `gathered_buffer` holds the contiguous gather + * destination produced by `gather_stripes`. The `data` field exposes + * non-owning device_span views into whichever arena currently owns + * each (stripe, stream) entry. */ struct encoded_data { - std::vector>> data; // Owning array of the encoded data - hostdevice_2dvector streams; // streams of encoded data, per chunk + rmm::device_uvector encoded_buffer; // arena for raw encoded streams + rmm::device_uvector gathered_buffer; // arena for gather_stripes output + std::vector>> data; // [stripe][strm_id] views + hostdevice_2dvector streams; // streams of encoded data, per chunk }; /** From 31ec8e2c68a9d0fe5bd169ca8e04950fa882286e Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Mon, 27 Jul 2026 19:48:45 +0000 Subject: [PATCH 02/20] Release the encoded arena as soon as the gather has consumed it Splitting the encoder output into one arena rather than a device_uvector per (stripe, stream) kept the upper-bound-sized allocation alive through compression, which raised peak memory by up to 53% on the ORC writer benchmarks. Partition the output into two arenas instead, by whether a region is expected to be compacted into the gather buffer: a region is, when its stripe spans several rowgroups and its size is a strict upper bound rather than exact. gather_stripes compacts every region in the transient arena as well as any region it measures a gap in, so the transient arena can be released as soon as gathering completes. Predicting which regions carry slack only decides arena placement; whether a region is compacted is still measured, so a misprediction cannot leave non-contiguous chunks to be read as a contiguous stream. It only costs the allocation that would have been released. Peak memory is now within 0.08% of main on all 109 benchmark configs, and timings are unchanged against the previous arena layout under the pool resource. --- cpp/src/io/orc/writer_impl.cu | 151 +++++++++++++++++++++++---------- cpp/src/io/orc/writer_impl.hpp | 19 +++-- 2 files changed, 116 insertions(+), 54 deletions(-) diff --git a/cpp/src/io/orc/writer_impl.cu b/cpp/src/io/orc/writer_impl.cu index 7937128d217f..70004d6055ac 100644 --- a/cpp/src/io/orc/writer_impl.cu +++ b/cpp/src/io/orc/writer_impl.cu @@ -27,6 +27,7 @@ #include #include +#include #include #include #include @@ -943,10 +944,16 @@ encoded_data encode_columns(orc_table_view const& orc_table, // Pass 1: compute per-rowgroup stream lengths and per-(stripe, strm_id) sizes. // The encoded buffers for every (stripe, stream) pair are packed into a single // arena allocation, so accumulate the total here and assign offsets later. + // + // Also record whether a region's size is a strict upper bound on what the encoder will write. + // Some estimates below are exact (dictionary data, string char counts, decimal chunk sizes) + // while `rle_stream_size` is a worst case; `gather_stripes` uses this to decide which arena a + // region belongs to, see the comment on the partitioning below. auto const num_streams = streams.size(); std::vector> stripe_strm_sizes(segmentation.num_stripes(), std::vector(num_streams, 0)); - size_t encoded_total = 0; + std::vector> stripe_strm_has_slack(segmentation.num_stripes(), + std::vector(num_streams, false)); for (auto const& stripe : segmentation.stripes) { for (size_t col_idx = 0; col_idx < num_columns; col_idx++) { for (int strm_type = 0; strm_type < CI_NUM_STREAMS; ++strm_type) { @@ -962,6 +969,8 @@ encoded_data encode_columns(orc_table_view const& orc_table, if (strm_id < 0) { continue; } size_t stripe_size = 0; + // Alignment padding is slack in itself, for every rowgroup of the region. + bool has_slack = uncomp_block_align > 1; std::for_each(stripe.cbegin(), stripe.cend(), [&](auto rg_idx) { #if defined(__GNUC__) && (__GNUC__ >= 14) #pragma GCC diagnostic push @@ -981,6 +990,9 @@ encoded_data encode_columns(orc_table_view const& orc_table, (strm_type == CI_DICTIONARY) ? stripe_dict.char_count : (((stripe_dict.entry_count + 0x1ff) >> 9) * (512 * 4 + 2)); + // The dictionary data is exactly `char_count` bytes; the dictionary lengths are + // RLE-encoded, so their size is a worst case. + has_slack |= strm_type != CI_DICTIONARY; } else { strm.lengths[strm_type] = 0; } @@ -991,53 +1003,76 @@ encoded_data encode_columns(orc_table_view const& orc_table, (ck.type_kind == DOUBLE || ck.type_kind == FLOAT)) { // Pass-through strm.lengths[strm_type] = ck.num_rows * ck.dtype_len; + has_slack = true; } else if (ck.type_kind == DECIMAL && strm_type == CI_DATA) { strm.lengths[strm_type] = dec_chunk_sizes.rg_sizes.at(col_idx)[rg_idx]; } else { strm.lengths[strm_type] = rle_stream_size(streams.type(strm_id), ck.num_rows); + has_slack = true; } // Allow extra space for alignment stripe_size += strm.lengths[strm_type] + uncomp_block_align - 1; }); - stripe_strm_sizes[stripe.id][strm_id] = stripe_size; - encoded_total += stripe_size; + stripe_strm_sizes[stripe.id][strm_id] = stripe_size; + stripe_strm_has_slack[stripe.id][strm_id] = has_slack; } } } - // Build per-(stripe, stream) offsets into a single arena. Each non-empty - // region starts on a 256-byte boundary to match the alignment that RMM - // guaranteed for the original per-region device_uvector allocations -- the - // ORC encoder kernels and downstream compressors rely on natural alignment - // (uncomp_block_align is only 1 for compression == NONE, so we cannot rely - // on the per-rg alignment fix-up alone). + // `gather_stripes` compacts a (stripe, stream) region into a tightly sized buffer when its + // chunks are not already contiguous, and reads it in place otherwise. Regions that are certain + // to be compacted go into a separate arena, so that it can be freed as soon as gathering + // completes instead of staying pinned for the lifetime of the encoded data by the regions that + // are read in place. A region is certain to be compacted when its stripe spans several + // rowgroups (so the chunks of a stream are laid out with gaps between them) and its size is a + // strict upper bound (so at least one of those gaps is non-empty). + // + // Mispredicting is safe in both directions. A region wrongly placed in the transient arena is + // compacted anyway, because that arena is gathered unconditionally. A region wrongly left in + // the persistent arena is compacted if `gather_stripes` measures a gap, and only costs the + // retained upper-bound allocation it would have freed. + auto const is_transient = [&](size_t stripe_id, size_t strm_id) { + return segmentation.stripes[stripe_id].size > 1 and stripe_strm_has_slack[stripe_id][strm_id]; + }; + + // Each non-empty region starts on a 256-byte boundary to match the alignment that RMM + // guaranteed for the original per-region device_uvector allocations -- the ORC encoder kernels + // and downstream compressors rely on natural alignment (uncomp_block_align is only 1 for + // compression == NONE, so we cannot rely on the per-rg alignment fix-up alone). constexpr size_t region_alignment = rmm::CUDA_ALLOCATION_ALIGNMENT; std::vector> stripe_strm_offsets(segmentation.num_stripes(), std::vector(num_streams, 0)); - size_t arena_total = 0; + std::vector> must_gather(segmentation.num_stripes(), + std::vector(num_streams, false)); + size_t persistent_arena_size = 0; + size_t transient_arena_size = 0; for (size_t s = 0; s < segmentation.num_stripes(); ++s) { for (size_t strm_id = 0; strm_id < num_streams; ++strm_id) { auto const sz = stripe_strm_sizes[s][strm_id]; - if (sz == 0) { - stripe_strm_offsets[s][strm_id] = arena_total; // empty span: data() is harmless - continue; - } - arena_total = util::round_up_unsafe(arena_total, region_alignment); - stripe_strm_offsets[s][strm_id] = arena_total; - arena_total += sz; + if (sz == 0) { continue; } + must_gather[s][strm_id] = is_transient(s, strm_id); + auto& arena_size = must_gather[s][strm_id] ? transient_arena_size : persistent_arena_size; + arena_size = util::round_up_unsafe(arena_size, region_alignment); + stripe_strm_offsets[s][strm_id] = arena_size; + arena_size += sz; } } - // Single arena allocation for all (stripe, stream) encoded buffers. - rmm::device_uvector encoded_buffer(arena_total, stream); + rmm::device_uvector persistent_buffer(persistent_arena_size, stream); + rmm::device_uvector transient_buffer(transient_arena_size, stream); - std::vector>> encoded_views(segmentation.num_stripes()); + // Zero-size regions keep a null pointer, matching the empty per-region device_uvector they + // replaced; a null `data_ptrs` entry selects the pass-through path in the encoder kernels. + std::vector>> encoded_views( + segmentation.num_stripes(), std::vector>(num_streams)); for (size_t s = 0; s < segmentation.num_stripes(); ++s) { - encoded_views[s].resize(num_streams); for (size_t strm_id = 0; strm_id < num_streams; ++strm_id) { - encoded_views[s][strm_id] = device_span{ - encoded_buffer.data() + stripe_strm_offsets[s][strm_id], stripe_strm_sizes[s][strm_id]}; + auto const sz = stripe_strm_sizes[s][strm_id]; + if (sz == 0) { continue; } + auto& arena = must_gather[s][strm_id] ? transient_buffer : persistent_buffer; + encoded_views[s][strm_id] = + device_span{arena.data() + stripe_strm_offsets[s][strm_id], sz}; } } @@ -1104,9 +1139,11 @@ encoded_data encode_columns(orc_table_view const& orc_table, } chunk_streams.device_to_host(stream); - return {std::move(encoded_buffer), + return {std::move(persistent_buffer), + std::move(transient_buffer), rmm::device_uvector{0, stream}, // gathered_buffer (filled by gather_stripes) std::move(encoded_views), + std::move(must_gather), std::move(chunk_streams)}; } @@ -1133,9 +1170,9 @@ std::vector gather_stripes(size_t num_index_streams, auto const num_streams_in_data = enc_data->data[0].size(); // Pass 1: compute per-(stripe, stream) actual sizes and decide which need a tight - // gathered copy. Don't copy the data to an exactly sized buffer when only one - // chunk is present, to avoid the overhead of the additional copy. When there - // are multiple chunks, they are copied anyway to make them contiguous. + // gathered copy. Streams of single-rowgroup stripes are read in place, to avoid the + // overhead of the additional copy. When there are multiple rowgroups, the chunks are + // copied anyway to make them contiguous. struct gather_info { size_t actual_size; bool gathered; @@ -1144,7 +1181,9 @@ std::vector gather_stripes(size_t num_index_streams, segmentation.num_stripes(), std::vector(num_streams_in_data, gather_info{0, false})); - size_t gather_total = 0; + // Releasing `transient_buffer` below is only valid if every region it holds was gathered. + bool all_transient_gathered = true; + for (auto const& stripe : segmentation.stripes) { for (size_t col_idx = 0; col_idx < enc_data->streams.size().first; col_idx++) { auto const& col_streams = (enc_data->streams)[col_idx]; @@ -1162,9 +1201,25 @@ std::vector gather_stripes(size_t num_index_streams, CUDF_EXPECTS(allocated_stripe_size >= actual_stripe_size, "Internal ORC writer error: insufficient allocation size for encoded data"); - bool const gathered = (stripe.size > 1 and allocated_stripe_size > actual_stripe_size); + // Compact a region whenever its chunks are not already contiguous, which is the case + // exactly when the encoder wrote less than the upper bound the region was sized for. + // Regions in `transient_buffer` are compacted unconditionally, so that arena can be + // released below; `encode_columns` only places a region there when it is expected to have + // a gap anyway, so this rarely forces a copy that would not have happened. + bool const gathered = (stripe.size > 1 and (enc_data->must_gather[stripe.id][stream_id] or + allocated_stripe_size > actual_stripe_size)); gather_meta[stripe.id][stream_id] = {actual_stripe_size, gathered}; - if (gathered) { gather_total += actual_stripe_size; } + } + } + } + + // Verify the invariant the release relies on, over the arena membership recorded by + // `encode_columns` rather than the loop above, which only visits stream types below CI_INDEX. + for (size_t s = 0; s < segmentation.num_stripes() and all_transient_gathered; ++s) { + for (size_t strm_id = 0; strm_id < num_streams_in_data; ++strm_id) { + if (enc_data->must_gather[s][strm_id] and not gather_meta[s][strm_id].gathered) { + all_transient_gathered = false; + break; } } } @@ -1175,17 +1230,14 @@ std::vector gather_stripes(size_t num_index_streams, constexpr size_t region_alignment = rmm::CUDA_ALLOCATION_ALIGNMENT; std::vector> gather_offsets(segmentation.num_stripes(), std::vector(num_streams_in_data, 0)); - { - size_t cursor = 0; - for (size_t s = 0; s < segmentation.num_stripes(); ++s) { - for (size_t strm_id = 0; strm_id < num_streams_in_data; ++strm_id) { - if (!gather_meta[s][strm_id].gathered) { continue; } - cursor = util::round_up_unsafe(cursor, region_alignment); - gather_offsets[s][strm_id] = cursor; - cursor += gather_meta[s][strm_id].actual_size; - } + size_t gather_total = 0; + for (size_t s = 0; s < segmentation.num_stripes(); ++s) { + for (size_t strm_id = 0; strm_id < num_streams_in_data; ++strm_id) { + if (!gather_meta[s][strm_id].gathered) { continue; } + gather_total = util::round_up_unsafe(gather_total, region_alignment); + gather_offsets[s][strm_id] = gather_total; + gather_total += gather_meta[s][strm_id].actual_size; } - gather_total = cursor; } rmm::device_uvector gather_buffer(gather_total, stream); @@ -1240,10 +1292,13 @@ std::vector gather_stripes(size_t num_index_streams, } } - // Hold the gathered arena for lifetime management. The encoded arena stays - // alive too: compress_orc_data_streams reads encoded data via per-rowgroup - // data_ptrs that for non-gathered (stripe, stream) still point into it. + // Hold the gathered arena for lifetime management, and release the arena whose regions have all + // been copied into it. `persistent_buffer` stays alive: some of its regions are read in place, + // via the per-rowgroup data_ptrs that still point into it. enc_data->gathered_buffer = std::move(gather_buffer); + if (all_transient_gathered) { + enc_data->transient_buffer = rmm::device_uvector{0, stream}; + } return stripes; } @@ -2561,11 +2616,13 @@ auto convert_table_to_orc_data(table_view const& input, comp_results, stream); - // deallocate encoded data as it is not needed anymore. Free both arenas - // (encoded + gathered) and clear the spans that referenced them. - enc_data.encoded_buffer = rmm::device_uvector{0, stream}; - enc_data.gathered_buffer = rmm::device_uvector{0, stream}; + // deallocate encoded data as it is not needed anymore. Free the remaining arenas and clear + // the spans that referenced them. + enc_data.persistent_buffer = rmm::device_uvector{0, stream}; + enc_data.transient_buffer = rmm::device_uvector{0, stream}; + enc_data.gathered_buffer = rmm::device_uvector{0, stream}; enc_data.data.clear(); + enc_data.must_gather.clear(); strm_descs.device_to_host_async(stream); comp_results.device_to_host(stream); diff --git a/cpp/src/io/orc/writer_impl.hpp b/cpp/src/io/orc/writer_impl.hpp index cf1523791449..7b37f8d80296 100644 --- a/cpp/src/io/orc/writer_impl.hpp +++ b/cpp/src/io/orc/writer_impl.hpp @@ -16,6 +16,7 @@ #include #include #include +#include #include #include @@ -98,17 +99,21 @@ struct file_segmentation { /** * @brief ORC per-chunk streams of encoded data. * - * The encoded buffers for every (stripe, stream) pair are packed into two arena - * allocations rather than one `device_uvector` per pair: `encoded_buffer` holds - * the raw encoder output, and `gathered_buffer` holds the contiguous gather - * destination produced by `gather_stripes`. The `data` field exposes - * non-owning device_span views into whichever arena currently owns - * each (stripe, stream) entry. + * The encoded buffers for every (stripe, stream) pair are packed into arena + * allocations rather than one `device_uvector` per pair. The encoder output is + * split across two arenas by whether `gather_stripes` is certain to copy the + * region into `gathered_buffer`: `transient_buffer` holds only such regions, so + * it is released as soon as gathering completes, while `persistent_buffer` holds + * the regions that may be read in place and must outlive the gather. The `data` + * field exposes non-owning device_span views into whichever arena owns + * each (stripe, stream) entry, and `must_gather` records the split. */ struct encoded_data { - rmm::device_uvector encoded_buffer; // arena for raw encoded streams + rmm::device_uvector persistent_buffer; // regions that may be read in place + rmm::device_uvector transient_buffer; // regions always copied out by the gather rmm::device_uvector gathered_buffer; // arena for gather_stripes output std::vector>> data; // [stripe][strm_id] views + std::vector> must_gather; // entry lives in `transient_buffer` hostdevice_2dvector streams; // streams of encoded data, per chunk }; From 681d9a443c291968763b86feaa959aedbe418e31 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Mon, 27 Jul 2026 21:34:12 +0000 Subject: [PATCH 03/20] Trim the arena comments now that the design is in the PR description --- cpp/src/io/orc/writer_impl.cu | 51 ++++++++++++---------------------- cpp/src/io/orc/writer_impl.hpp | 11 +++----- 2 files changed, 22 insertions(+), 40 deletions(-) diff --git a/cpp/src/io/orc/writer_impl.cu b/cpp/src/io/orc/writer_impl.cu index 70004d6055ac..9c8ea7ca4d06 100644 --- a/cpp/src/io/orc/writer_impl.cu +++ b/cpp/src/io/orc/writer_impl.cu @@ -941,14 +941,9 @@ encoded_data encode_columns(orc_table_view const& orc_table, hostdevice_2dvector chunk_streams( num_columns, segmentation.num_rowgroups(), stream); - // Pass 1: compute per-rowgroup stream lengths and per-(stripe, strm_id) sizes. - // The encoded buffers for every (stripe, stream) pair are packed into a single - // arena allocation, so accumulate the total here and assign offsets later. - // - // Also record whether a region's size is a strict upper bound on what the encoder will write. - // Some estimates below are exact (dictionary data, string char counts, decimal chunk sizes) - // while `rle_stream_size` is a worst case; `gather_stripes` uses this to decide which arena a - // region belongs to, see the comment on the partitioning below. + // Pass 1: compute per-rowgroup stream lengths and per-(stripe, strm_id) sizes, along with + // whether a size is a strict upper bound on what the encoder will write. Offsets into the + // arenas are assigned below. auto const num_streams = streams.size(); std::vector> stripe_strm_sizes(segmentation.num_stripes(), std::vector(num_streams, 0)); @@ -969,7 +964,7 @@ encoded_data encode_columns(orc_table_view const& orc_table, if (strm_id < 0) { continue; } size_t stripe_size = 0; - // Alignment padding is slack in itself, for every rowgroup of the region. + // Alignment padding is slack in itself, in every rowgroup of the region. bool has_slack = uncomp_block_align > 1; std::for_each(stripe.cbegin(), stripe.cend(), [&](auto rg_idx) { #if defined(__GNUC__) && (__GNUC__ >= 14) @@ -990,8 +985,7 @@ encoded_data encode_columns(orc_table_view const& orc_table, (strm_type == CI_DICTIONARY) ? stripe_dict.char_count : (((stripe_dict.entry_count + 0x1ff) >> 9) * (512 * 4 + 2)); - // The dictionary data is exactly `char_count` bytes; the dictionary lengths are - // RLE-encoded, so their size is a worst case. + // `char_count` is exact; the RLE-encoded lengths are a worst case. has_slack |= strm_type != CI_DICTIONARY; } else { strm.lengths[strm_type] = 0; @@ -1020,18 +1014,11 @@ encoded_data encode_columns(orc_table_view const& orc_table, } } - // `gather_stripes` compacts a (stripe, stream) region into a tightly sized buffer when its - // chunks are not already contiguous, and reads it in place otherwise. Regions that are certain - // to be compacted go into a separate arena, so that it can be freed as soon as gathering - // completes instead of staying pinned for the lifetime of the encoded data by the regions that - // are read in place. A region is certain to be compacted when its stripe spans several - // rowgroups (so the chunks of a stream are laid out with gaps between them) and its size is a - // strict upper bound (so at least one of those gaps is non-empty). - // - // Mispredicting is safe in both directions. A region wrongly placed in the transient arena is - // compacted anyway, because that arena is gathered unconditionally. A region wrongly left in - // the persistent arena is compacted if `gather_stripes` measures a gap, and only costs the - // retained upper-bound allocation it would have freed. + // Regions that `gather_stripes` is certain to compact go into their own arena, so it can be + // freed as soon as gathering completes rather than staying pinned by the regions read in place. + // A region is certain to be compacted when its stripe spans several rowgroups (so its chunks + // are laid out with gaps) and its size is an upper bound (so a gap is non-empty). Guessing + // wrong either way is safe, and only costs a copy or a retained allocation. auto const is_transient = [&](size_t stripe_id, size_t strm_id) { return segmentation.stripes[stripe_id].size > 1 and stripe_strm_has_slack[stripe_id][strm_id]; }; @@ -1201,11 +1188,9 @@ std::vector gather_stripes(size_t num_index_streams, CUDF_EXPECTS(allocated_stripe_size >= actual_stripe_size, "Internal ORC writer error: insufficient allocation size for encoded data"); - // Compact a region whenever its chunks are not already contiguous, which is the case - // exactly when the encoder wrote less than the upper bound the region was sized for. - // Regions in `transient_buffer` are compacted unconditionally, so that arena can be - // released below; `encode_columns` only places a region there when it is expected to have - // a gap anyway, so this rarely forces a copy that would not have happened. + // Compact when the chunks are not already contiguous, i.e. when the encoder wrote less + // than the region was sized for. Regions in `transient_buffer` are compacted regardless, + // so that arena can be released below. bool const gathered = (stripe.size > 1 and (enc_data->must_gather[stripe.id][stream_id] or allocated_stripe_size > actual_stripe_size)); gather_meta[stripe.id][stream_id] = {actual_stripe_size, gathered}; @@ -1213,8 +1198,8 @@ std::vector gather_stripes(size_t num_index_streams, } } - // Verify the invariant the release relies on, over the arena membership recorded by - // `encode_columns` rather than the loop above, which only visits stream types below CI_INDEX. + // Verify the invariant the release relies on. Uses the membership recorded by `encode_columns`, + // because the loop above only visits stream types below CI_INDEX. for (size_t s = 0; s < segmentation.num_stripes() and all_transient_gathered; ++s) { for (size_t strm_id = 0; strm_id < num_streams_in_data; ++strm_id) { if (enc_data->must_gather[s][strm_id] and not gather_meta[s][strm_id].gathered) { @@ -1292,9 +1277,9 @@ std::vector gather_stripes(size_t num_index_streams, } } - // Hold the gathered arena for lifetime management, and release the arena whose regions have all - // been copied into it. `persistent_buffer` stays alive: some of its regions are read in place, - // via the per-rowgroup data_ptrs that still point into it. + // Hold the gathered arena for lifetime management, and release the arena it copied from. + // `persistent_buffer` stays alive: some of its regions are read in place, via per-rowgroup + // data_ptrs that still point into it. enc_data->gathered_buffer = std::move(gather_buffer); if (all_transient_gathered) { enc_data->transient_buffer = rmm::device_uvector{0, stream}; diff --git a/cpp/src/io/orc/writer_impl.hpp b/cpp/src/io/orc/writer_impl.hpp index 7b37f8d80296..d67c516f1074 100644 --- a/cpp/src/io/orc/writer_impl.hpp +++ b/cpp/src/io/orc/writer_impl.hpp @@ -100,13 +100,10 @@ struct file_segmentation { * @brief ORC per-chunk streams of encoded data. * * The encoded buffers for every (stripe, stream) pair are packed into arena - * allocations rather than one `device_uvector` per pair. The encoder output is - * split across two arenas by whether `gather_stripes` is certain to copy the - * region into `gathered_buffer`: `transient_buffer` holds only such regions, so - * it is released as soon as gathering completes, while `persistent_buffer` holds - * the regions that may be read in place and must outlive the gather. The `data` - * field exposes non-owning device_span views into whichever arena owns - * each (stripe, stream) entry, and `must_gather` records the split. + * allocations rather than one `device_uvector` per pair, and exposed as + * non-owning views in `data`. The encoder output is split across two arenas by + * whether `gather_stripes` is certain to copy the region into `gathered_buffer`, + * so that `transient_buffer` can be released as soon as gathering completes. */ struct encoded_data { rmm::device_uvector persistent_buffer; // regions that may be read in place From 3b57a99d945a3a5fc7b9211769cfd051520a4d60 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Mon, 27 Jul 2026 22:43:30 +0000 Subject: [PATCH 04/20] Address review feedback on the encoded-data arenas Treat the pass-through float/double data stream as an exact size estimate, so those regions stay in the arena that is read in place instead of being copied out for no reason. Hand the transient-region set from `encode_columns` to `gather_stripes` as its own argument rather than parking it in `encoded_data`, where it outlived its only use, and make the release of `transient_buffer` unconditional with the invariant it relies on checked explicitly. Flatten the per-(stripe, stream) bookkeeping into single allocations, and cover stripes that span several rowgroups with compression disabled and with pass-through floats. --- cpp/src/io/orc/writer_impl.cu | 178 ++++++++++++++++----------------- cpp/src/io/orc/writer_impl.hpp | 34 ++++++- cpp/tests/io/orc_test.cpp | 61 +++++++++++ 3 files changed, 182 insertions(+), 91 deletions(-) diff --git a/cpp/src/io/orc/writer_impl.cu b/cpp/src/io/orc/writer_impl.cu index 9c8ea7ca4d06..ed5f96cf4c02 100644 --- a/cpp/src/io/orc/writer_impl.cu +++ b/cpp/src/io/orc/writer_impl.cu @@ -63,6 +63,12 @@ namespace cudf::io::orc::detail { +// Alignment of every non-empty region within the encoded and gathered arenas. Matches what RMM +// guaranteed for the per-region device_uvector allocations the arenas replaced -- the ORC encoder +// kernels and downstream compressors rely on natural alignment, and uncomp_block_align is 1 for +// some codecs, so the per-rowgroup alignment fix-up alone is not enough. +constexpr size_t region_alignment = rmm::CUDA_ALLOCATION_ALIGNMENT; + template [[nodiscard]] CUDF_HOST_DEVICE constexpr int varint_size(T val) { @@ -850,12 +856,12 @@ struct segmented_valid_cnt_input { std::vector indices; }; -encoded_data encode_columns(orc_table_view const& orc_table, - encoder_decimal_info&& dec_chunk_sizes, - file_segmentation const& segmentation, - orc_streams const& streams, - uint32_t uncomp_block_align, - rmm::cuda_stream_view stream) +std::pair encode_columns(orc_table_view const& orc_table, + encoder_decimal_info&& dec_chunk_sizes, + file_segmentation const& segmentation, + orc_streams const& streams, + uint32_t uncomp_block_align, + rmm::cuda_stream_view stream) { auto const num_columns = orc_table.num_columns(); hostdevice_2dvector chunks(num_columns, segmentation.num_rowgroups(), stream); @@ -943,12 +949,13 @@ encoded_data encode_columns(orc_table_view const& orc_table, // Pass 1: compute per-rowgroup stream lengths and per-(stripe, strm_id) sizes, along with // whether a size is a strict upper bound on what the encoder will write. Offsets into the - // arenas are assigned below. + // arenas are assigned below. Per-region data is indexed `stripe_id * num_streams + strm_id`. auto const num_streams = streams.size(); - std::vector> stripe_strm_sizes(segmentation.num_stripes(), - std::vector(num_streams, 0)); - std::vector> stripe_strm_has_slack(segmentation.num_stripes(), - std::vector(num_streams, false)); + auto const region_idx = [num_streams](size_t stripe_id, size_t strm_id) { + return stripe_id * num_streams + strm_id; + }; + std::vector region_sizes(segmentation.num_stripes() * num_streams, 0); + std::vector region_has_slack(segmentation.num_stripes() * num_streams, false); for (auto const& stripe : segmentation.stripes) { for (size_t col_idx = 0; col_idx < num_columns; col_idx++) { for (int strm_type = 0; strm_type < CI_NUM_STREAMS; ++strm_type) { @@ -995,9 +1002,8 @@ encoded_data encode_columns(orc_table_view const& orc_table, strm.lengths[strm_type] = std::max(column.rowgroup_char_count(rg_idx), 1); } else if (strm_type == CI_DATA && streams[strm_id].length == 0 && (ck.type_kind == DOUBLE || ck.type_kind == FLOAT)) { - // Pass-through + // Pass-through. The encoder reports this length back unchanged, so it is exact. strm.lengths[strm_type] = ck.num_rows * ck.dtype_len; - has_slack = true; } else if (ck.type_kind == DECIMAL && strm_type == CI_DATA) { strm.lengths[strm_type] = dec_chunk_sizes.rg_sizes.at(col_idx)[rg_idx]; } else { @@ -1008,8 +1014,8 @@ encoded_data encode_columns(orc_table_view const& orc_table, stripe_size += strm.lengths[strm_type] + uncomp_block_align - 1; }); - stripe_strm_sizes[stripe.id][strm_id] = stripe_size; - stripe_strm_has_slack[stripe.id][strm_id] = has_slack; + region_sizes[region_idx(stripe.id, strm_id)] = stripe_size; + region_has_slack[region_idx(stripe.id, strm_id)] = has_slack; } } } @@ -1019,30 +1025,22 @@ encoded_data encode_columns(orc_table_view const& orc_table, // A region is certain to be compacted when its stripe spans several rowgroups (so its chunks // are laid out with gaps) and its size is an upper bound (so a gap is non-empty). Guessing // wrong either way is safe, and only costs a copy or a retained allocation. - auto const is_transient = [&](size_t stripe_id, size_t strm_id) { - return segmentation.stripes[stripe_id].size > 1 and stripe_strm_has_slack[stripe_id][strm_id]; - }; - - // Each non-empty region starts on a 256-byte boundary to match the alignment that RMM - // guaranteed for the original per-region device_uvector allocations -- the ORC encoder kernels - // and downstream compressors rely on natural alignment (uncomp_block_align is only 1 for - // compression == NONE, so we cannot rely on the per-rg alignment fix-up alone). - constexpr size_t region_alignment = rmm::CUDA_ALLOCATION_ALIGNMENT; - std::vector> stripe_strm_offsets(segmentation.num_stripes(), - std::vector(num_streams, 0)); - std::vector> must_gather(segmentation.num_stripes(), - std::vector(num_streams, false)); + transient_regions transient{segmentation.num_stripes(), num_streams}; + std::vector region_offsets(segmentation.num_stripes() * num_streams, 0); size_t persistent_arena_size = 0; size_t transient_arena_size = 0; for (size_t s = 0; s < segmentation.num_stripes(); ++s) { for (size_t strm_id = 0; strm_id < num_streams; ++strm_id) { - auto const sz = stripe_strm_sizes[s][strm_id]; - if (sz == 0) { continue; } - must_gather[s][strm_id] = is_transient(s, strm_id); - auto& arena_size = must_gather[s][strm_id] ? transient_arena_size : persistent_arena_size; - arena_size = util::round_up_unsafe(arena_size, region_alignment); - stripe_strm_offsets[s][strm_id] = arena_size; - arena_size += sz; + auto const idx = region_idx(s, strm_id); + if (region_sizes[idx] == 0) { continue; } + if (segmentation.stripes[s].size > 1 and region_has_slack[idx]) { + transient.insert(s, strm_id); + } + auto& arena_size = + transient.contains(s, strm_id) ? transient_arena_size : persistent_arena_size; + arena_size = util::round_up_unsafe(arena_size, region_alignment); + region_offsets[idx] = arena_size; + arena_size += region_sizes[idx]; } } @@ -1055,11 +1053,11 @@ encoded_data encode_columns(orc_table_view const& orc_table, segmentation.num_stripes(), std::vector>(num_streams)); for (size_t s = 0; s < segmentation.num_stripes(); ++s) { for (size_t strm_id = 0; strm_id < num_streams; ++strm_id) { - auto const sz = stripe_strm_sizes[s][strm_id]; - if (sz == 0) { continue; } - auto& arena = must_gather[s][strm_id] ? transient_buffer : persistent_buffer; + auto const idx = region_idx(s, strm_id); + if (region_sizes[idx] == 0) { continue; } + auto& arena = transient.contains(s, strm_id) ? transient_buffer : persistent_buffer; encoded_views[s][strm_id] = - device_span{arena.data() + stripe_strm_offsets[s][strm_id], sz}; + device_span{arena.data() + region_offsets[idx], region_sizes[idx]}; } } @@ -1126,12 +1124,12 @@ encoded_data encode_columns(orc_table_view const& orc_table, } chunk_streams.device_to_host(stream); - return {std::move(persistent_buffer), - std::move(transient_buffer), - rmm::device_uvector{0, stream}, // gathered_buffer (filled by gather_stripes) - std::move(encoded_views), - std::move(must_gather), - std::move(chunk_streams)}; + return {encoded_data{std::move(persistent_buffer), + std::move(transient_buffer), + rmm::device_uvector{0, stream}, // filled by gather_stripes + std::move(encoded_views), + std::move(chunk_streams)}, + std::move(transient)}; } // TODO: remove StripeInformation from this function and return strm_desc instead @@ -1141,6 +1139,7 @@ encoded_data encode_columns(orc_table_view const& orc_table, * * @param[in] num_index_streams Total number of index streams * @param[in] segmentation stripe and rowgroup ranges + * @param[in] transient Regions held in `enc_data->transient_buffer`, which must all be gathered * @param[in,out] enc_data ORC per-chunk streams of encoded data * @param[in,out] strm_desc List of stream descriptors [stripe][data_stream] * @param[in] stream CUDA stream used for device memory operations and kernel launches @@ -1148,6 +1147,7 @@ encoded_data encode_columns(orc_table_view const& orc_table, */ std::vector gather_stripes(size_t num_index_streams, file_segmentation const& segmentation, + transient_regions const& transient, encoded_data* enc_data, hostdevice_2dvector* strm_desc, rmm::cuda_stream_view stream) @@ -1155,6 +1155,9 @@ std::vector gather_stripes(size_t num_index_streams, if (segmentation.num_stripes() == 0) { return {}; } auto const num_streams_in_data = enc_data->data[0].size(); + auto const region_idx = [num_streams_in_data](size_t stripe_id, size_t strm_id) { + return stripe_id * num_streams_in_data + strm_id; + }; // Pass 1: compute per-(stripe, stream) actual sizes and decide which need a tight // gathered copy. Streams of single-rowgroup stripes are read in place, to avoid the @@ -1164,12 +1167,8 @@ std::vector gather_stripes(size_t num_index_streams, size_t actual_size; bool gathered; }; - std::vector> gather_meta( - segmentation.num_stripes(), - std::vector(num_streams_in_data, gather_info{0, false})); - - // Releasing `transient_buffer` below is only valid if every region it holds was gathered. - bool all_transient_gathered = true; + std::vector gather_meta(segmentation.num_stripes() * num_streams_in_data, + gather_info{0, false}); for (auto const& stripe : segmentation.stripes) { for (size_t col_idx = 0; col_idx < enc_data->streams.size().first; col_idx++) { @@ -1191,45 +1190,43 @@ std::vector gather_stripes(size_t num_index_streams, // Compact when the chunks are not already contiguous, i.e. when the encoder wrote less // than the region was sized for. Regions in `transient_buffer` are compacted regardless, // so that arena can be released below. - bool const gathered = (stripe.size > 1 and (enc_data->must_gather[stripe.id][stream_id] or + bool const gathered = (stripe.size > 1 and (transient.contains(stripe.id, stream_id) or allocated_stripe_size > actual_stripe_size)); - gather_meta[stripe.id][stream_id] = {actual_stripe_size, gathered}; + gather_meta[region_idx(stripe.id, stream_id)] = {actual_stripe_size, gathered}; } } } - // Verify the invariant the release relies on. Uses the membership recorded by `encode_columns`, - // because the loop above only visits stream types below CI_INDEX. - for (size_t s = 0; s < segmentation.num_stripes() and all_transient_gathered; ++s) { + // Every transient region is gathered by construction, since `encode_columns` only makes a + // region transient when its stripe spans several rowgroups, which is what the predicate above + // tests. Check it anyway, because the release below is a use-after-free if it ever stops + // holding, and the loop above cannot see stream types at or above CI_INDEX. + for (size_t s = 0; s < segmentation.num_stripes(); ++s) { for (size_t strm_id = 0; strm_id < num_streams_in_data; ++strm_id) { - if (enc_data->must_gather[s][strm_id] and not gather_meta[s][strm_id].gathered) { - all_transient_gathered = false; - break; - } + CUDF_EXPECTS( + not transient.contains(s, strm_id) or gather_meta[region_idx(s, strm_id)].gathered, + "Internal ORC writer error: transient encoded region was not gathered"); } } - // Lay out gather destinations within a single arena. As with the encoded - // arena above, each region starts on a 256-byte boundary to match the - // alignment RMM provided for the original per-region device_uvectors. - constexpr size_t region_alignment = rmm::CUDA_ALLOCATION_ALIGNMENT; - std::vector> gather_offsets(segmentation.num_stripes(), - std::vector(num_streams_in_data, 0)); + // Lay out gather destinations within a single arena, with the same alignment rule as the + // encoded arenas. + std::vector gather_offsets(segmentation.num_stripes() * num_streams_in_data, 0); size_t gather_total = 0; for (size_t s = 0; s < segmentation.num_stripes(); ++s) { for (size_t strm_id = 0; strm_id < num_streams_in_data; ++strm_id) { - if (!gather_meta[s][strm_id].gathered) { continue; } - gather_total = util::round_up_unsafe(gather_total, region_alignment); - gather_offsets[s][strm_id] = gather_total; - gather_total += gather_meta[s][strm_id].actual_size; + auto const idx = region_idx(s, strm_id); + if (!gather_meta[idx].gathered) { continue; } + gather_total = util::round_up_unsafe(gather_total, region_alignment); + gather_offsets[idx] = gather_total; + gather_total += gather_meta[idx].actual_size; } } rmm::device_uvector gather_buffer(gather_total, stream); // Pass 2: build strm_desc entries and record gather destination spans. - std::vector>> gather_views( - segmentation.num_stripes(), - std::vector>(num_streams_in_data, device_span{})); + std::vector> gather_views(segmentation.num_stripes() * num_streams_in_data, + device_span{}); std::vector stripes(segmentation.num_stripes()); for (auto const& stripe : segmentation.stripes) { for (size_t col_idx = 0; col_idx < enc_data->streams.size().first; col_idx++) { @@ -1238,11 +1235,15 @@ std::vector gather_stripes(size_t num_index_streams, auto const stream_id = col_streams[0].ids[k]; if (stream_id == -1) { continue; } - auto const& meta = gather_meta[stripe.id][stream_id]; + auto const idx = region_idx(stripe.id, stream_id); + auto const& meta = gather_meta[idx]; uint8_t* dst_ptr = nullptr; if (meta.gathered) { - dst_ptr = gather_buffer.data() + gather_offsets[stripe.id][stream_id]; - gather_views[stripe.id][stream_id] = device_span{dst_ptr, meta.actual_size}; + // Non-null even when the region is empty, unlike the empty device_uvector this replaced. + // `init_batched_memcpy_kernel` repoints the per-rowgroup data_ptrs at this arena, which + // is what lets `transient_buffer` be released without leaving them dangling. + dst_ptr = gather_buffer.data() + gather_offsets[idx]; + gather_views[idx] = device_span{dst_ptr, meta.actual_size}; } auto* ss = &(*strm_desc)[stripe.id][stream_id - num_index_streams]; @@ -1271,19 +1272,16 @@ std::vector gather_stripes(size_t num_index_streams, // spans, so consumers that read enc_data->data observe the post-gather state. for (size_t stripe_id = 0; stripe_id < enc_data->data.size(); ++stripe_id) { for (size_t stream_id = 0; stream_id < num_streams_in_data; ++stream_id) { - if (gather_meta[stripe_id][stream_id].gathered) { - enc_data->data[stripe_id][stream_id] = gather_views[stripe_id][stream_id]; - } + auto const idx = region_idx(stripe_id, stream_id); + if (gather_meta[idx].gathered) { enc_data->data[stripe_id][stream_id] = gather_views[idx]; } } } // Hold the gathered arena for lifetime management, and release the arena it copied from. // `persistent_buffer` stays alive: some of its regions are read in place, via per-rowgroup // data_ptrs that still point into it. - enc_data->gathered_buffer = std::move(gather_buffer); - if (all_transient_gathered) { - enc_data->transient_buffer = rmm::device_uvector{0, stream}; - } + enc_data->gathered_buffer = std::move(gather_buffer); + enc_data->transient_buffer = rmm::device_uvector{0, stream}; return stripes; } @@ -2517,13 +2515,14 @@ auto convert_table_to_orc_data(table_view const& input, auto const block_align = compress_required_chunk_alignment(compression); - auto streams = create_streams(orc_table.columns, + auto streams = create_streams(orc_table.columns, segmentation, decimal_column_sizes(dec_chunk_sizes.rg_sizes), enable_dictionary, compression, write_mode); - auto enc_data = encode_columns( + + auto [enc_data, transient] = encode_columns( orc_table, std::move(dec_chunk_sizes), segmentation, streams, block_align, stream); stripe_dicts.on_encode_complete(stream); @@ -2535,7 +2534,8 @@ auto convert_table_to_orc_data(table_view const& input, auto const num_data_streams = streams.size() - num_index_streams; hostdevice_2dvector strm_descs( segmentation.num_stripes(), num_data_streams, stream); - auto stripes = gather_stripes(num_index_streams, segmentation, &enc_data, &strm_descs, stream); + auto stripes = + gather_stripes(num_index_streams, segmentation, transient, &enc_data, &strm_descs, stream); if (num_rows == 0) { return std::tuple{std::move(enc_data), @@ -2601,13 +2601,11 @@ auto convert_table_to_orc_data(table_view const& input, comp_results, stream); - // deallocate encoded data as it is not needed anymore. Free the remaining arenas and clear - // the spans that referenced them. + // deallocate encoded data as it is not needed anymore. Frees the arenas that outlived the + // gather and clears the spans that referenced them. enc_data.persistent_buffer = rmm::device_uvector{0, stream}; - enc_data.transient_buffer = rmm::device_uvector{0, stream}; enc_data.gathered_buffer = rmm::device_uvector{0, stream}; enc_data.data.clear(); - enc_data.must_gather.clear(); strm_descs.device_to_host_async(stream); comp_results.device_to_host(stream); diff --git a/cpp/src/io/orc/writer_impl.hpp b/cpp/src/io/orc/writer_impl.hpp index d67c516f1074..d79d36e23518 100644 --- a/cpp/src/io/orc/writer_impl.hpp +++ b/cpp/src/io/orc/writer_impl.hpp @@ -104,16 +104,48 @@ struct file_segmentation { * non-owning views in `data`. The encoder output is split across two arenas by * whether `gather_stripes` is certain to copy the region into `gathered_buffer`, * so that `transient_buffer` can be released as soon as gathering completes. + * `encode_columns` predicts that split and `gather_stripes` refines it by + * measuring what the encoder wrote, so the two arenas are not simply a function + * of the stripe layout. */ struct encoded_data { rmm::device_uvector persistent_buffer; // regions that may be read in place rmm::device_uvector transient_buffer; // regions always copied out by the gather rmm::device_uvector gathered_buffer; // arena for gather_stripes output std::vector>> data; // [stripe][strm_id] views - std::vector> must_gather; // entry lives in `transient_buffer` hostdevice_2dvector streams; // streams of encoded data, per chunk }; +/** + * @brief Set of encoded regions placed in `encoded_data::transient_buffer`. + * + * `gather_stripes` has to copy every region in the set into + * `encoded_data::gathered_buffer` before it can release that arena. + */ +class transient_regions { + public: + transient_regions(size_t num_stripes, size_t num_streams) + : _num_streams{num_streams}, _flags(num_stripes * num_streams, false) + { + } + + void insert(size_t stripe_id, size_t strm_id) { _flags[index(stripe_id, strm_id)] = true; } + + [[nodiscard]] bool contains(size_t stripe_id, size_t strm_id) const + { + return _flags[index(stripe_id, strm_id)]; + } + + private: + [[nodiscard]] size_t index(size_t stripe_id, size_t strm_id) const + { + return stripe_id * _num_streams + strm_id; + } + + size_t _num_streams; + std::vector _flags; +}; + /** * @brief Dictionary data for string columns and their device views, per column. */ diff --git a/cpp/tests/io/orc_test.cpp b/cpp/tests/io/orc_test.cpp index 6159facde017..a3713a425b0a 100644 --- a/cpp/tests/io/orc_test.cpp +++ b/cpp/tests/io/orc_test.cpp @@ -1874,6 +1874,67 @@ TEST_F(OrcWriterTest, EmptyRowGroup) CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result.tbl->view()); } +// Stripes spanning several rowgroups, with compression disabled so that the writer does not pad +// the per-rowgroup streams to a block alignment. Both are needed to cover the case where the +// encoded output happens to be exactly as large as the writer sized it for. +TEST_F(OrcWriterTest, MultiRowgroupStripeUncompressed) +{ + constexpr auto num_rows = 5 * 10'000 + 7; // several rowgroups, the last one partial + + auto const ints = random_values(num_rows); + int32_col int_col(ints.begin(), ints.end()); + auto str_elements = cudf::detail::make_counting_transform_iterator( + 0, [](auto i) { return "v" + std::to_string(i % 137); }); + str_col string_col(str_elements, str_elements + num_rows); + table_view expected({int_col, string_col}); + + auto filepath = temp_env->get_temp_filepath("MultiRowgroupStripeUncompressed.orc"); + cudf::io::orc_writer_options out_opts = + cudf::io::orc_writer_options::builder(cudf::io::sink_info{filepath}, expected) + .compression(cudf::io::compression_type::NONE); + cudf::io::write_orc(out_opts); + + // The rowgroups have to end up in one stripe for this to test what it means to. + EXPECT_EQ(read_orc_metadata(cudf::io::source_info{filepath}).num_stripes(), 1); + + cudf::io::orc_reader_options in_opts = + cudf::io::orc_reader_options::builder(cudf::io::source_info{filepath}); + auto result = cudf::io::read_orc(in_opts); + CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result.tbl->view()); +} + +// Null-free floating point columns are written to the data stream unencoded, so the writer hands +// the encoder a region it does not write to. Spans several rowgroups so that those regions are +// laid out with the rest of the encoded data rather than on their own. +TEST_F(OrcWriterTest, MultiRowgroupPassthroughFloats) +{ + constexpr auto num_rows = 3 * 10'000 + 1; + + auto const floats = random_values(num_rows); + auto const doubles = random_values(num_rows); + float32_col float_col(floats.begin(), floats.end()); + float64_col double_col(doubles.begin(), doubles.end()); + // A nullable column alongside them, so the stripe also holds encoded (non-pass-through) streams. + auto const ints = random_values(num_rows); + auto mask = cudf::detail::make_counting_transform_iterator(0, [](auto i) { return i % 3 != 0; }); + int32_col int_col(ints.begin(), ints.end(), mask); + table_view expected({float_col, double_col, int_col}); + + for (auto const compression : + {cudf::io::compression_type::NONE, cudf::io::compression_type::SNAPPY}) { + auto filepath = temp_env->get_temp_filepath("MultiRowgroupPassthroughFloats.orc"); + cudf::io::orc_writer_options out_opts = + cudf::io::orc_writer_options::builder(cudf::io::sink_info{filepath}, expected) + .compression(compression); + cudf::io::write_orc(out_opts); + + cudf::io::orc_reader_options in_opts = + cudf::io::orc_reader_options::builder(cudf::io::source_info{filepath}); + auto result = cudf::io::read_orc(in_opts); + CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result.tbl->view()); + } +} + TEST_F(OrcWriterTest, NoNullsAsNonNullable) { auto valids = cudf::test::iterators::no_nulls(); From e0f9dd91027c17eea21b6641030e6835c092f101 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Mon, 27 Jul 2026 23:34:26 +0000 Subject: [PATCH 05/20] Say why the arena regions are aligned, and check it The alignment is not there because the encoder needs naturally aligned streams; `encode_columns` already realigns every per-rowgroup pointer by absolute address, so the encoded arenas would tolerate any base. It is the gather that needs it: compaction drops that per-rowgroup alignment, leaving the region base as the pointer the compressor receives, so it has to meet the codec's requirement. Check that rather than assuming 256 bytes is always enough. --- cpp/src/io/orc/writer_impl.cu | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/cpp/src/io/orc/writer_impl.cu b/cpp/src/io/orc/writer_impl.cu index ed5f96cf4c02..2c2bd822c7c0 100644 --- a/cpp/src/io/orc/writer_impl.cu +++ b/cpp/src/io/orc/writer_impl.cu @@ -63,10 +63,12 @@ namespace cudf::io::orc::detail { -// Alignment of every non-empty region within the encoded and gathered arenas. Matches what RMM -// guaranteed for the per-region device_uvector allocations the arenas replaced -- the ORC encoder -// kernels and downstream compressors rely on natural alignment, and uncomp_block_align is 1 for -// some codecs, so the per-rowgroup alignment fix-up alone is not enough. +// Alignment of every non-empty region within the encoded and gathered arenas. Compression is what +// requires it: `gather_stripes` compacts the per-rowgroup chunks without re-applying the codec +// alignment that `encode_columns` gave them, so a gathered region's base is the pointer the +// compressor receives, and it has to satisfy `compress_required_chunk_alignment` (checked in +// `encode_columns`). RMM's allocation alignment is what the per-region allocations these arenas +// replaced provided, so using it keeps every downstream access at least as aligned as before. constexpr size_t region_alignment = rmm::CUDA_ALLOCATION_ALIGNMENT; template @@ -1025,6 +1027,9 @@ std::pair encode_columns(orc_table_view const& // A region is certain to be compacted when its stripe spans several rowgroups (so its chunks // are laid out with gaps) and its size is an upper bound (so a gap is non-empty). Guessing // wrong either way is safe, and only costs a copy or a retained allocation. + CUDF_EXPECTS(region_alignment >= uncomp_block_align, + "Internal ORC writer error: arena regions are not aligned enough for the codec"); + transient_regions transient{segmentation.num_stripes(), num_streams}; std::vector region_offsets(segmentation.num_stripes() * num_streams, 0); size_t persistent_arena_size = 0; From fc05e8ecf4cde901d14890a5b8e81e941f3784d8 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 28 Jul 2026 00:08:18 +0000 Subject: [PATCH 06/20] Cover the exactly-sized stream types in the uncompressed stripe test The string column in the test is dictionary encoded, whose per-stripe stream is a single region no matter how many rowgroups it spans. Add the types whose per-rowgroup sizes are known exactly and chained end to end -- direct encoded strings and decimals -- since those are what actually exercise reading a multi-rowgroup stripe in place. --- cpp/tests/io/orc_test.cpp | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/cpp/tests/io/orc_test.cpp b/cpp/tests/io/orc_test.cpp index a3713a425b0a..fa337e74befe 100644 --- a/cpp/tests/io/orc_test.cpp +++ b/cpp/tests/io/orc_test.cpp @@ -1876,17 +1876,26 @@ TEST_F(OrcWriterTest, EmptyRowGroup) // Stripes spanning several rowgroups, with compression disabled so that the writer does not pad // the per-rowgroup streams to a block alignment. Both are needed to cover the case where the -// encoded output happens to be exactly as large as the writer sized it for. +// encoded output is exactly as large as the writer sized it for, which is when the writer reads +// the encoded chunks of a multi-rowgroup stripe in place instead of compacting them first. The +// column types are the ones whose stream sizes the writer knows exactly up front: dictionary and +// direct encoded strings, and decimals. TEST_F(OrcWriterTest, MultiRowgroupStripeUncompressed) { constexpr auto num_rows = 5 * 10'000 + 7; // several rowgroups, the last one partial auto const ints = random_values(num_rows); int32_col int_col(ints.begin(), ints.end()); - auto str_elements = cudf::detail::make_counting_transform_iterator( + auto dict_elements = cudf::detail::make_counting_transform_iterator( 0, [](auto i) { return "v" + std::to_string(i % 137); }); - str_col string_col(str_elements, str_elements + num_rows); - table_view expected({int_col, string_col}); + str_col dict_string_col(dict_elements, dict_elements + num_rows); + // Distinct values of varying length, so this column is written without a dictionary. + auto direct_elements = cudf::detail::make_counting_transform_iterator( + 0, [](auto i) { return std::to_string(i) + std::string(i % 19, 'x'); }); + str_col direct_string_col(direct_elements, direct_elements + num_rows); + auto const decimals = random_values(num_rows); + dec64_col decimal_col(decimals.begin(), decimals.end(), numeric::scale_type{-2}); + table_view expected({int_col, dict_string_col, direct_string_col, decimal_col}); auto filepath = temp_env->get_temp_filepath("MultiRowgroupStripeUncompressed.orc"); cudf::io::orc_writer_options out_opts = From d905aef834bd8f7987e078c155879ff8e949bb28 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 28 Jul 2026 00:31:25 +0000 Subject: [PATCH 07/20] Sweep compression in the decimal writer test instead of a new test The decimal parameterization already covers stripes of one and of several rowgroups, and decimal data streams are the exactly sized kind, so adding a compression axis to it covers writing a multi-rowgroup stripe straight from the encoder output. Drops the standalone test that was doing the same thing. --- cpp/tests/io/orc_test.cpp | 56 +++++++++------------------------------ 1 file changed, 13 insertions(+), 43 deletions(-) diff --git a/cpp/tests/io/orc_test.cpp b/cpp/tests/io/orc_test.cpp index fa337e74befe..037f9ec879c7 100644 --- a/cpp/tests/io/orc_test.cpp +++ b/cpp/tests/io/orc_test.cpp @@ -1235,12 +1235,13 @@ TEST_F(OrcReaderTest, MultipleInputs) CUDF_TEST_EXPECT_TABLES_EQUAL(*result.tbl, *full_table); } -struct OrcWriterTestDecimal : public OrcWriterTest, - public ::testing::WithParamInterface> {}; +struct OrcWriterTestDecimal + : public OrcWriterTest, + public ::testing::WithParamInterface> {}; TEST_P(OrcWriterTestDecimal, Decimal64) { - auto const [num_rows, scale] = GetParam(); + auto const [num_rows, scale, compression] = GetParam(); // Using int16_t because scale causes values to overflow if they already require 32 bits auto const vals = random_values(num_rows); @@ -1250,7 +1251,8 @@ TEST_P(OrcWriterTestDecimal, Decimal64) auto filepath = temp_env->get_temp_filepath("Decimal64.orc"); cudf::io::orc_writer_options out_opts = - cudf::io::orc_writer_options::builder(cudf::io::sink_info{filepath}, tbl); + cudf::io::orc_writer_options::builder(cudf::io::sink_info{filepath}, tbl) + .compression(compression); cudf::io::write_orc(out_opts); @@ -1261,10 +1263,16 @@ TEST_P(OrcWriterTestDecimal, Decimal64) CUDF_TEST_EXPECT_COLUMNS_EQUAL(tbl.column(0), result.tbl->view().column(0)); } +// The row counts span stripes of one and of several rowgroups. Decimal data stream sizes are known +// exactly up front, and without compression the writer does not pad them to a block alignment, so +// the encoded chunks of a multi-rowgroup stripe come out contiguous and are written from the +// encoder output as-is rather than being compacted first. INSTANTIATE_TEST_CASE_P(OrcWriterTest, OrcWriterTestDecimal, ::testing::Combine(::testing::Values(1, 10000, 10001, 34567), - ::testing::Values(-2, 0, 2))); + ::testing::Values(-2, 0, 2), + ::testing::Values(cudf::io::compression_type::AUTO, + cudf::io::compression_type::NONE))); TEST_F(OrcWriterTest, Decimal32) { @@ -1874,44 +1882,6 @@ TEST_F(OrcWriterTest, EmptyRowGroup) CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result.tbl->view()); } -// Stripes spanning several rowgroups, with compression disabled so that the writer does not pad -// the per-rowgroup streams to a block alignment. Both are needed to cover the case where the -// encoded output is exactly as large as the writer sized it for, which is when the writer reads -// the encoded chunks of a multi-rowgroup stripe in place instead of compacting them first. The -// column types are the ones whose stream sizes the writer knows exactly up front: dictionary and -// direct encoded strings, and decimals. -TEST_F(OrcWriterTest, MultiRowgroupStripeUncompressed) -{ - constexpr auto num_rows = 5 * 10'000 + 7; // several rowgroups, the last one partial - - auto const ints = random_values(num_rows); - int32_col int_col(ints.begin(), ints.end()); - auto dict_elements = cudf::detail::make_counting_transform_iterator( - 0, [](auto i) { return "v" + std::to_string(i % 137); }); - str_col dict_string_col(dict_elements, dict_elements + num_rows); - // Distinct values of varying length, so this column is written without a dictionary. - auto direct_elements = cudf::detail::make_counting_transform_iterator( - 0, [](auto i) { return std::to_string(i) + std::string(i % 19, 'x'); }); - str_col direct_string_col(direct_elements, direct_elements + num_rows); - auto const decimals = random_values(num_rows); - dec64_col decimal_col(decimals.begin(), decimals.end(), numeric::scale_type{-2}); - table_view expected({int_col, dict_string_col, direct_string_col, decimal_col}); - - auto filepath = temp_env->get_temp_filepath("MultiRowgroupStripeUncompressed.orc"); - cudf::io::orc_writer_options out_opts = - cudf::io::orc_writer_options::builder(cudf::io::sink_info{filepath}, expected) - .compression(cudf::io::compression_type::NONE); - cudf::io::write_orc(out_opts); - - // The rowgroups have to end up in one stripe for this to test what it means to. - EXPECT_EQ(read_orc_metadata(cudf::io::source_info{filepath}).num_stripes(), 1); - - cudf::io::orc_reader_options in_opts = - cudf::io::orc_reader_options::builder(cudf::io::source_info{filepath}); - auto result = cudf::io::read_orc(in_opts); - CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result.tbl->view()); -} - // Null-free floating point columns are written to the data stream unencoded, so the writer hands // the encoder a region it does not write to. Spans several rowgroups so that those regions are // laid out with the rest of the encoded data rather than on their own. From 9e7c5a0f144872cda187308e4cfbad80b145f8d2 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 28 Jul 2026 18:51:31 +0000 Subject: [PATCH 08/20] Drop the pass-through float test as redundant OrcChunkedReaderInputLimitTest already writes null-free doubles with compression disabled at 20k-row stripes, which is the same multi-rowgroup pass-through layout at a much larger scale. --- cpp/tests/io/orc_test.cpp | 32 -------------------------------- 1 file changed, 32 deletions(-) diff --git a/cpp/tests/io/orc_test.cpp b/cpp/tests/io/orc_test.cpp index 037f9ec879c7..9eb500f17a11 100644 --- a/cpp/tests/io/orc_test.cpp +++ b/cpp/tests/io/orc_test.cpp @@ -1882,38 +1882,6 @@ TEST_F(OrcWriterTest, EmptyRowGroup) CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result.tbl->view()); } -// Null-free floating point columns are written to the data stream unencoded, so the writer hands -// the encoder a region it does not write to. Spans several rowgroups so that those regions are -// laid out with the rest of the encoded data rather than on their own. -TEST_F(OrcWriterTest, MultiRowgroupPassthroughFloats) -{ - constexpr auto num_rows = 3 * 10'000 + 1; - - auto const floats = random_values(num_rows); - auto const doubles = random_values(num_rows); - float32_col float_col(floats.begin(), floats.end()); - float64_col double_col(doubles.begin(), doubles.end()); - // A nullable column alongside them, so the stripe also holds encoded (non-pass-through) streams. - auto const ints = random_values(num_rows); - auto mask = cudf::detail::make_counting_transform_iterator(0, [](auto i) { return i % 3 != 0; }); - int32_col int_col(ints.begin(), ints.end(), mask); - table_view expected({float_col, double_col, int_col}); - - for (auto const compression : - {cudf::io::compression_type::NONE, cudf::io::compression_type::SNAPPY}) { - auto filepath = temp_env->get_temp_filepath("MultiRowgroupPassthroughFloats.orc"); - cudf::io::orc_writer_options out_opts = - cudf::io::orc_writer_options::builder(cudf::io::sink_info{filepath}, expected) - .compression(compression); - cudf::io::write_orc(out_opts); - - cudf::io::orc_reader_options in_opts = - cudf::io::orc_reader_options::builder(cudf::io::source_info{filepath}); - auto result = cudf::io::read_orc(in_opts); - CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result.tbl->view()); - } -} - TEST_F(OrcWriterTest, NoNullsAsNonNullable) { auto valids = cudf::test::iterators::no_nulls(); From d519992dc3216f00187d008e732c7f8b931ec538 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 28 Jul 2026 19:13:32 +0000 Subject: [PATCH 09/20] Call the arena storage of a stripe stream an extent "Region" was a second name for something ORC already names. An extent is the aligned byte range within an arena that holds one (stripe, stream) pair, which is what replaces the allocation each pair used to get. --- cpp/src/io/orc/writer_impl.cu | 88 +++++++++++++++++----------------- cpp/src/io/orc/writer_impl.hpp | 28 +++++------ 2 files changed, 58 insertions(+), 58 deletions(-) diff --git a/cpp/src/io/orc/writer_impl.cu b/cpp/src/io/orc/writer_impl.cu index 2c2bd822c7c0..3bcf03bae534 100644 --- a/cpp/src/io/orc/writer_impl.cu +++ b/cpp/src/io/orc/writer_impl.cu @@ -63,13 +63,13 @@ namespace cudf::io::orc::detail { -// Alignment of every non-empty region within the encoded and gathered arenas. Compression is what +// Alignment of every non-empty extent within the encoded and gathered arenas. Compression is what // requires it: `gather_stripes` compacts the per-rowgroup chunks without re-applying the codec -// alignment that `encode_columns` gave them, so a gathered region's base is the pointer the +// alignment that `encode_columns` gave them, so a gathered extent's base is the pointer the // compressor receives, and it has to satisfy `compress_required_chunk_alignment` (checked in -// `encode_columns`). RMM's allocation alignment is what the per-region allocations these arenas +// `encode_columns`). RMM's allocation alignment is what the per-stream allocations these arenas // replaced provided, so using it keeps every downstream access at least as aligned as before. -constexpr size_t region_alignment = rmm::CUDA_ALLOCATION_ALIGNMENT; +constexpr size_t extent_alignment = rmm::CUDA_ALLOCATION_ALIGNMENT; template [[nodiscard]] CUDF_HOST_DEVICE constexpr int varint_size(T val) @@ -858,7 +858,7 @@ struct segmented_valid_cnt_input { std::vector indices; }; -std::pair encode_columns(orc_table_view const& orc_table, +std::pair encode_columns(orc_table_view const& orc_table, encoder_decimal_info&& dec_chunk_sizes, file_segmentation const& segmentation, orc_streams const& streams, @@ -951,13 +951,13 @@ std::pair encode_columns(orc_table_view const& // Pass 1: compute per-rowgroup stream lengths and per-(stripe, strm_id) sizes, along with // whether a size is a strict upper bound on what the encoder will write. Offsets into the - // arenas are assigned below. Per-region data is indexed `stripe_id * num_streams + strm_id`. + // arenas are assigned below. Per-extent data is indexed `stripe_id * num_streams + strm_id`. auto const num_streams = streams.size(); - auto const region_idx = [num_streams](size_t stripe_id, size_t strm_id) { + auto const extent_idx = [num_streams](size_t stripe_id, size_t strm_id) { return stripe_id * num_streams + strm_id; }; - std::vector region_sizes(segmentation.num_stripes() * num_streams, 0); - std::vector region_has_slack(segmentation.num_stripes() * num_streams, false); + std::vector extent_sizes(segmentation.num_stripes() * num_streams, 0); + std::vector extent_has_slack(segmentation.num_stripes() * num_streams, false); for (auto const& stripe : segmentation.stripes) { for (size_t col_idx = 0; col_idx < num_columns; col_idx++) { for (int strm_type = 0; strm_type < CI_NUM_STREAMS; ++strm_type) { @@ -973,7 +973,7 @@ std::pair encode_columns(orc_table_view const& if (strm_id < 0) { continue; } size_t stripe_size = 0; - // Alignment padding is slack in itself, in every rowgroup of the region. + // Alignment padding is slack in itself, in every rowgroup of the extent. bool has_slack = uncomp_block_align > 1; std::for_each(stripe.cbegin(), stripe.cend(), [&](auto rg_idx) { #if defined(__GNUC__) && (__GNUC__ >= 14) @@ -1016,53 +1016,53 @@ std::pair encode_columns(orc_table_view const& stripe_size += strm.lengths[strm_type] + uncomp_block_align - 1; }); - region_sizes[region_idx(stripe.id, strm_id)] = stripe_size; - region_has_slack[region_idx(stripe.id, strm_id)] = has_slack; + extent_sizes[extent_idx(stripe.id, strm_id)] = stripe_size; + extent_has_slack[extent_idx(stripe.id, strm_id)] = has_slack; } } } - // Regions that `gather_stripes` is certain to compact go into their own arena, so it can be - // freed as soon as gathering completes rather than staying pinned by the regions read in place. - // A region is certain to be compacted when its stripe spans several rowgroups (so its chunks + // Extents that `gather_stripes` is certain to compact go into their own arena, so it can be + // freed as soon as gathering completes rather than staying pinned by the extents read in place. + // An extent is certain to be compacted when its stripe spans several rowgroups (so its chunks // are laid out with gaps) and its size is an upper bound (so a gap is non-empty). Guessing // wrong either way is safe, and only costs a copy or a retained allocation. - CUDF_EXPECTS(region_alignment >= uncomp_block_align, - "Internal ORC writer error: arena regions are not aligned enough for the codec"); + CUDF_EXPECTS(extent_alignment >= uncomp_block_align, + "Internal ORC writer error: arena extents are not aligned enough for the codec"); - transient_regions transient{segmentation.num_stripes(), num_streams}; - std::vector region_offsets(segmentation.num_stripes() * num_streams, 0); + transient_extents transient{segmentation.num_stripes(), num_streams}; + std::vector extent_offsets(segmentation.num_stripes() * num_streams, 0); size_t persistent_arena_size = 0; size_t transient_arena_size = 0; for (size_t s = 0; s < segmentation.num_stripes(); ++s) { for (size_t strm_id = 0; strm_id < num_streams; ++strm_id) { - auto const idx = region_idx(s, strm_id); - if (region_sizes[idx] == 0) { continue; } - if (segmentation.stripes[s].size > 1 and region_has_slack[idx]) { + auto const idx = extent_idx(s, strm_id); + if (extent_sizes[idx] == 0) { continue; } + if (segmentation.stripes[s].size > 1 and extent_has_slack[idx]) { transient.insert(s, strm_id); } auto& arena_size = transient.contains(s, strm_id) ? transient_arena_size : persistent_arena_size; - arena_size = util::round_up_unsafe(arena_size, region_alignment); - region_offsets[idx] = arena_size; - arena_size += region_sizes[idx]; + arena_size = util::round_up_unsafe(arena_size, extent_alignment); + extent_offsets[idx] = arena_size; + arena_size += extent_sizes[idx]; } } rmm::device_uvector persistent_buffer(persistent_arena_size, stream); rmm::device_uvector transient_buffer(transient_arena_size, stream); - // Zero-size regions keep a null pointer, matching the empty per-region device_uvector they + // Zero-size extents keep a null pointer, matching the empty per-stream device_uvector they // replaced; a null `data_ptrs` entry selects the pass-through path in the encoder kernels. std::vector>> encoded_views( segmentation.num_stripes(), std::vector>(num_streams)); for (size_t s = 0; s < segmentation.num_stripes(); ++s) { for (size_t strm_id = 0; strm_id < num_streams; ++strm_id) { - auto const idx = region_idx(s, strm_id); - if (region_sizes[idx] == 0) { continue; } + auto const idx = extent_idx(s, strm_id); + if (extent_sizes[idx] == 0) { continue; } auto& arena = transient.contains(s, strm_id) ? transient_buffer : persistent_buffer; encoded_views[s][strm_id] = - device_span{arena.data() + region_offsets[idx], region_sizes[idx]}; + device_span{arena.data() + extent_offsets[idx], extent_sizes[idx]}; } } @@ -1144,7 +1144,7 @@ std::pair encode_columns(orc_table_view const& * * @param[in] num_index_streams Total number of index streams * @param[in] segmentation stripe and rowgroup ranges - * @param[in] transient Regions held in `enc_data->transient_buffer`, which must all be gathered + * @param[in] transient Extents held in `enc_data->transient_buffer`, which must all be gathered * @param[in,out] enc_data ORC per-chunk streams of encoded data * @param[in,out] strm_desc List of stream descriptors [stripe][data_stream] * @param[in] stream CUDA stream used for device memory operations and kernel launches @@ -1152,7 +1152,7 @@ std::pair encode_columns(orc_table_view const& */ std::vector gather_stripes(size_t num_index_streams, file_segmentation const& segmentation, - transient_regions const& transient, + transient_extents const& transient, encoded_data* enc_data, hostdevice_2dvector* strm_desc, rmm::cuda_stream_view stream) @@ -1160,7 +1160,7 @@ std::vector gather_stripes(size_t num_index_streams, if (segmentation.num_stripes() == 0) { return {}; } auto const num_streams_in_data = enc_data->data[0].size(); - auto const region_idx = [num_streams_in_data](size_t stripe_id, size_t strm_id) { + auto const extent_idx = [num_streams_in_data](size_t stripe_id, size_t strm_id) { return stripe_id * num_streams_in_data + strm_id; }; @@ -1193,24 +1193,24 @@ std::vector gather_stripes(size_t num_index_streams, "Internal ORC writer error: insufficient allocation size for encoded data"); // Compact when the chunks are not already contiguous, i.e. when the encoder wrote less - // than the region was sized for. Regions in `transient_buffer` are compacted regardless, + // than the extent was sized for. Extents in `transient_buffer` are compacted regardless, // so that arena can be released below. bool const gathered = (stripe.size > 1 and (transient.contains(stripe.id, stream_id) or allocated_stripe_size > actual_stripe_size)); - gather_meta[region_idx(stripe.id, stream_id)] = {actual_stripe_size, gathered}; + gather_meta[extent_idx(stripe.id, stream_id)] = {actual_stripe_size, gathered}; } } } - // Every transient region is gathered by construction, since `encode_columns` only makes a - // region transient when its stripe spans several rowgroups, which is what the predicate above + // Every transient extent is gathered by construction, since `encode_columns` only makes an + // extent transient when its stripe spans several rowgroups, which is what the predicate above // tests. Check it anyway, because the release below is a use-after-free if it ever stops // holding, and the loop above cannot see stream types at or above CI_INDEX. for (size_t s = 0; s < segmentation.num_stripes(); ++s) { for (size_t strm_id = 0; strm_id < num_streams_in_data; ++strm_id) { CUDF_EXPECTS( - not transient.contains(s, strm_id) or gather_meta[region_idx(s, strm_id)].gathered, - "Internal ORC writer error: transient encoded region was not gathered"); + not transient.contains(s, strm_id) or gather_meta[extent_idx(s, strm_id)].gathered, + "Internal ORC writer error: transient encoded extent was not gathered"); } } @@ -1220,9 +1220,9 @@ std::vector gather_stripes(size_t num_index_streams, size_t gather_total = 0; for (size_t s = 0; s < segmentation.num_stripes(); ++s) { for (size_t strm_id = 0; strm_id < num_streams_in_data; ++strm_id) { - auto const idx = region_idx(s, strm_id); + auto const idx = extent_idx(s, strm_id); if (!gather_meta[idx].gathered) { continue; } - gather_total = util::round_up_unsafe(gather_total, region_alignment); + gather_total = util::round_up_unsafe(gather_total, extent_alignment); gather_offsets[idx] = gather_total; gather_total += gather_meta[idx].actual_size; } @@ -1240,11 +1240,11 @@ std::vector gather_stripes(size_t num_index_streams, auto const stream_id = col_streams[0].ids[k]; if (stream_id == -1) { continue; } - auto const idx = region_idx(stripe.id, stream_id); + auto const idx = extent_idx(stripe.id, stream_id); auto const& meta = gather_meta[idx]; uint8_t* dst_ptr = nullptr; if (meta.gathered) { - // Non-null even when the region is empty, unlike the empty device_uvector this replaced. + // Non-null even when the extent is empty, unlike the empty device_uvector this replaced. // `init_batched_memcpy_kernel` repoints the per-rowgroup data_ptrs at this arena, which // is what lets `transient_buffer` be released without leaving them dangling. dst_ptr = gather_buffer.data() + gather_offsets[idx]; @@ -1277,13 +1277,13 @@ std::vector gather_stripes(size_t num_index_streams, // spans, so consumers that read enc_data->data observe the post-gather state. for (size_t stripe_id = 0; stripe_id < enc_data->data.size(); ++stripe_id) { for (size_t stream_id = 0; stream_id < num_streams_in_data; ++stream_id) { - auto const idx = region_idx(stripe_id, stream_id); + auto const idx = extent_idx(stripe_id, stream_id); if (gather_meta[idx].gathered) { enc_data->data[stripe_id][stream_id] = gather_views[idx]; } } } // Hold the gathered arena for lifetime management, and release the arena it copied from. - // `persistent_buffer` stays alive: some of its regions are read in place, via per-rowgroup + // `persistent_buffer` stays alive: some of its extents are read in place, via per-rowgroup // data_ptrs that still point into it. enc_data->gathered_buffer = std::move(gather_buffer); enc_data->transient_buffer = rmm::device_uvector{0, stream}; diff --git a/cpp/src/io/orc/writer_impl.hpp b/cpp/src/io/orc/writer_impl.hpp index d79d36e23518..34fcd5ff320c 100644 --- a/cpp/src/io/orc/writer_impl.hpp +++ b/cpp/src/io/orc/writer_impl.hpp @@ -99,32 +99,32 @@ struct file_segmentation { /** * @brief ORC per-chunk streams of encoded data. * - * The encoded buffers for every (stripe, stream) pair are packed into arena - * allocations rather than one `device_uvector` per pair, and exposed as - * non-owning views in `data`. The encoder output is split across two arenas by - * whether `gather_stripes` is certain to copy the region into `gathered_buffer`, - * so that `transient_buffer` can be released as soon as gathering completes. - * `encode_columns` predicts that split and `gather_stripes` refines it by - * measuring what the encoder wrote, so the two arenas are not simply a function - * of the stripe layout. + * Every (stripe, stream) pair is stored in an extent -- an aligned byte range + * within one of the arenas below -- instead of getting a `device_uvector` of its + * own, and is exposed as a non-owning view in `data`. The encoder writes into + * two of the arenas, split by whether `gather_stripes` is certain to copy the + * extent into `gathered_buffer`, so that `transient_buffer` can be released as + * soon as gathering completes. `encode_columns` predicts that split and + * `gather_stripes` refines it by measuring what the encoder wrote, so the two + * are not simply a function of the stripe layout. */ struct encoded_data { - rmm::device_uvector persistent_buffer; // regions that may be read in place - rmm::device_uvector transient_buffer; // regions always copied out by the gather + rmm::device_uvector persistent_buffer; // extents that may be read in place + rmm::device_uvector transient_buffer; // extents always copied out by the gather rmm::device_uvector gathered_buffer; // arena for gather_stripes output std::vector>> data; // [stripe][strm_id] views hostdevice_2dvector streams; // streams of encoded data, per chunk }; /** - * @brief Set of encoded regions placed in `encoded_data::transient_buffer`. + * @brief Set of encoded extents placed in `encoded_data::transient_buffer`. * - * `gather_stripes` has to copy every region in the set into + * `gather_stripes` has to copy every extent in the set into * `encoded_data::gathered_buffer` before it can release that arena. */ -class transient_regions { +class transient_extents { public: - transient_regions(size_t num_stripes, size_t num_streams) + transient_extents(size_t num_stripes, size_t num_streams) : _num_streams{num_streams}, _flags(num_stripes * num_streams, false) { } From be3ca744468367a4d7a7227e5790e9f4223ceed8 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 28 Jul 2026 20:35:08 +0000 Subject: [PATCH 10/20] trim comments --- cpp/src/io/orc/writer_impl.cu | 26 +++++++------------------- 1 file changed, 7 insertions(+), 19 deletions(-) diff --git a/cpp/src/io/orc/writer_impl.cu b/cpp/src/io/orc/writer_impl.cu index 3bcf03bae534..b9a702318a90 100644 --- a/cpp/src/io/orc/writer_impl.cu +++ b/cpp/src/io/orc/writer_impl.cu @@ -63,12 +63,7 @@ namespace cudf::io::orc::detail { -// Alignment of every non-empty extent within the encoded and gathered arenas. Compression is what -// requires it: `gather_stripes` compacts the per-rowgroup chunks without re-applying the codec -// alignment that `encode_columns` gave them, so a gathered extent's base is the pointer the -// compressor receives, and it has to satisfy `compress_required_chunk_alignment` (checked in -// `encode_columns`). RMM's allocation alignment is what the per-stream allocations these arenas -// replaced provided, so using it keeps every downstream access at least as aligned as before. +// Alignment of every non-empty extent within the encoded and gathered arenas. constexpr size_t extent_alignment = rmm::CUDA_ALLOCATION_ALIGNMENT; template @@ -949,9 +944,8 @@ std::pair encode_columns(orc_table_view const& hostdevice_2dvector chunk_streams( num_columns, segmentation.num_rowgroups(), stream); - // Pass 1: compute per-rowgroup stream lengths and per-(stripe, strm_id) sizes, along with - // whether a size is a strict upper bound on what the encoder will write. Offsets into the - // arenas are assigned below. Per-extent data is indexed `stripe_id * num_streams + strm_id`. + // Compute per-rowgroup stream lengths and per-(stripe, strm_id) extent sizes, along with + // whether an extent size may exceed what the encoder ends up writing. auto const num_streams = streams.size(); auto const extent_idx = [num_streams](size_t stripe_id, size_t strm_id) { return stripe_id * num_streams + strm_id; @@ -1022,14 +1016,11 @@ std::pair encode_columns(orc_table_view const& } } - // Extents that `gather_stripes` is certain to compact go into their own arena, so it can be - // freed as soon as gathering completes rather than staying pinned by the extents read in place. - // An extent is certain to be compacted when its stripe spans several rowgroups (so its chunks - // are laid out with gaps) and its size is an upper bound (so a gap is non-empty). Guessing - // wrong either way is safe, and only costs a copy or a retained allocation. CUDF_EXPECTS(extent_alignment >= uncomp_block_align, "Internal ORC writer error: arena extents are not aligned enough for the codec"); + // Extents that `gather_stripes` is certain to compact go into `transient_buffer`, so that arena + // can be freed as soon as gathering completes. transient_extents transient{segmentation.num_stripes(), num_streams}; std::vector extent_offsets(segmentation.num_stripes() * num_streams, 0); size_t persistent_arena_size = 0; @@ -1052,8 +1043,7 @@ std::pair encode_columns(orc_table_view const& rmm::device_uvector persistent_buffer(persistent_arena_size, stream); rmm::device_uvector transient_buffer(transient_arena_size, stream); - // Zero-size extents keep a null pointer, matching the empty per-stream device_uvector they - // replaced; a null `data_ptrs` entry selects the pass-through path in the encoder kernels. + // Zero-size extents keep a null pointer, matching the empty device_uvector they replaced. std::vector>> encoded_views( segmentation.num_stripes(), std::vector>(num_streams)); for (size_t s = 0; s < segmentation.num_stripes(); ++s) { @@ -1066,9 +1056,7 @@ std::pair encode_columns(orc_table_view const& } } - // Pass 2: write per-chunk data_ptrs and apply alignment fix-up. The lengths - // computed in pass 1 (now stored in `chunk_streams[col_idx][rg_idx].lengths`) - // are read here as-is. + // Write per-chunk data_ptrs and apply alignment fix-up. for (auto const& stripe : segmentation.stripes) { for (size_t col_idx = 0; col_idx < num_columns; col_idx++) { for (int strm_type = 0; strm_type < CI_NUM_STREAMS; ++strm_type) { From 8bf6fb3034ae5190e79af2339e3f1e6f760f74f8 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 28 Jul 2026 21:06:21 +0000 Subject: [PATCH 11/20] Replace transient_extents class with a flat vector viewed as host_2dspan The bespoke class was a 2D array of bools with a hand-rolled index; cudf already has host_2dspan for passing [row][column] host data between the writer's helpers. --- cpp/src/io/orc/writer_impl.cu | 46 ++++++++++++++++++---------------- cpp/src/io/orc/writer_impl.hpp | 41 +++--------------------------- 2 files changed, 28 insertions(+), 59 deletions(-) diff --git a/cpp/src/io/orc/writer_impl.cu b/cpp/src/io/orc/writer_impl.cu index b9a702318a90..9c58d752c503 100644 --- a/cpp/src/io/orc/writer_impl.cu +++ b/cpp/src/io/orc/writer_impl.cu @@ -853,12 +853,14 @@ struct segmented_valid_cnt_input { std::vector indices; }; -std::pair encode_columns(orc_table_view const& orc_table, - encoder_decimal_info&& dec_chunk_sizes, - file_segmentation const& segmentation, - orc_streams const& streams, - uint32_t uncomp_block_align, - rmm::cuda_stream_view stream) +// Returns the encoded data, along with a [stripe][strm_id] flag for each extent placed in +// `encoded_data::transient_buffer`, flattened with `streams.size()` elements per row. +std::pair> encode_columns(orc_table_view const& orc_table, + encoder_decimal_info&& dec_chunk_sizes, + file_segmentation const& segmentation, + orc_streams const& streams, + uint32_t uncomp_block_align, + rmm::cuda_stream_view stream) { auto const num_columns = orc_table.num_columns(); hostdevice_2dvector chunks(num_columns, segmentation.num_rowgroups(), stream); @@ -1021,7 +1023,7 @@ std::pair encode_columns(orc_table_view const& // Extents that `gather_stripes` is certain to compact go into `transient_buffer`, so that arena // can be freed as soon as gathering completes. - transient_extents transient{segmentation.num_stripes(), num_streams}; + std::vector extent_is_transient(segmentation.num_stripes() * num_streams, 0); std::vector extent_offsets(segmentation.num_stripes() * num_streams, 0); size_t persistent_arena_size = 0; size_t transient_arena_size = 0; @@ -1029,11 +1031,8 @@ std::pair encode_columns(orc_table_view const& for (size_t strm_id = 0; strm_id < num_streams; ++strm_id) { auto const idx = extent_idx(s, strm_id); if (extent_sizes[idx] == 0) { continue; } - if (segmentation.stripes[s].size > 1 and extent_has_slack[idx]) { - transient.insert(s, strm_id); - } - auto& arena_size = - transient.contains(s, strm_id) ? transient_arena_size : persistent_arena_size; + extent_is_transient[idx] = segmentation.stripes[s].size > 1 and extent_has_slack[idx]; + auto& arena_size = extent_is_transient[idx] ? transient_arena_size : persistent_arena_size; arena_size = util::round_up_unsafe(arena_size, extent_alignment); extent_offsets[idx] = arena_size; arena_size += extent_sizes[idx]; @@ -1050,7 +1049,7 @@ std::pair encode_columns(orc_table_view const& for (size_t strm_id = 0; strm_id < num_streams; ++strm_id) { auto const idx = extent_idx(s, strm_id); if (extent_sizes[idx] == 0) { continue; } - auto& arena = transient.contains(s, strm_id) ? transient_buffer : persistent_buffer; + auto& arena = extent_is_transient[idx] ? transient_buffer : persistent_buffer; encoded_views[s][strm_id] = device_span{arena.data() + extent_offsets[idx], extent_sizes[idx]}; } @@ -1122,7 +1121,7 @@ std::pair encode_columns(orc_table_view const& rmm::device_uvector{0, stream}, // filled by gather_stripes std::move(encoded_views), std::move(chunk_streams)}, - std::move(transient)}; + std::move(extent_is_transient)}; } // TODO: remove StripeInformation from this function and return strm_desc instead @@ -1132,7 +1131,8 @@ std::pair encode_columns(orc_table_view const& * * @param[in] num_index_streams Total number of index streams * @param[in] segmentation stripe and rowgroup ranges - * @param[in] transient Extents held in `enc_data->transient_buffer`, which must all be gathered + * @param[in] extent_is_transient Marks extents held in `enc_data->transient_buffer`, which must all + * be gathered * @param[in,out] enc_data ORC per-chunk streams of encoded data * @param[in,out] strm_desc List of stream descriptors [stripe][data_stream] * @param[in] stream CUDA stream used for device memory operations and kernel launches @@ -1140,7 +1140,7 @@ std::pair encode_columns(orc_table_view const& */ std::vector gather_stripes(size_t num_index_streams, file_segmentation const& segmentation, - transient_extents const& transient, + host_2dspan extent_is_transient, encoded_data* enc_data, hostdevice_2dvector* strm_desc, rmm::cuda_stream_view stream) @@ -1183,7 +1183,7 @@ std::vector gather_stripes(size_t num_index_streams, // Compact when the chunks are not already contiguous, i.e. when the encoder wrote less // than the extent was sized for. Extents in `transient_buffer` are compacted regardless, // so that arena can be released below. - bool const gathered = (stripe.size > 1 and (transient.contains(stripe.id, stream_id) or + bool const gathered = (stripe.size > 1 and (extent_is_transient[stripe.id][stream_id] or allocated_stripe_size > actual_stripe_size)); gather_meta[extent_idx(stripe.id, stream_id)] = {actual_stripe_size, gathered}; } @@ -1197,7 +1197,7 @@ std::vector gather_stripes(size_t num_index_streams, for (size_t s = 0; s < segmentation.num_stripes(); ++s) { for (size_t strm_id = 0; strm_id < num_streams_in_data; ++strm_id) { CUDF_EXPECTS( - not transient.contains(s, strm_id) or gather_meta[extent_idx(s, strm_id)].gathered, + not extent_is_transient[s][strm_id] or gather_meta[extent_idx(s, strm_id)].gathered, "Internal ORC writer error: transient encoded extent was not gathered"); } } @@ -2515,7 +2515,7 @@ auto convert_table_to_orc_data(table_view const& input, compression, write_mode); - auto [enc_data, transient] = encode_columns( + auto [enc_data, extent_is_transient] = encode_columns( orc_table, std::move(dec_chunk_sizes), segmentation, streams, block_align, stream); stripe_dicts.on_encode_complete(stream); @@ -2527,8 +2527,12 @@ auto convert_table_to_orc_data(table_view const& input, auto const num_data_streams = streams.size() - num_index_streams; hostdevice_2dvector strm_descs( segmentation.num_stripes(), num_data_streams, stream); - auto stripes = - gather_stripes(num_index_streams, segmentation, transient, &enc_data, &strm_descs, stream); + auto stripes = gather_stripes(num_index_streams, + segmentation, + host_2dspan{extent_is_transient, streams.size()}, + &enc_data, + &strm_descs, + stream); if (num_rows == 0) { return std::tuple{std::move(enc_data), diff --git a/cpp/src/io/orc/writer_impl.hpp b/cpp/src/io/orc/writer_impl.hpp index 34fcd5ff320c..361787fd39f4 100644 --- a/cpp/src/io/orc/writer_impl.hpp +++ b/cpp/src/io/orc/writer_impl.hpp @@ -99,14 +99,9 @@ struct file_segmentation { /** * @brief ORC per-chunk streams of encoded data. * - * Every (stripe, stream) pair is stored in an extent -- an aligned byte range - * within one of the arenas below -- instead of getting a `device_uvector` of its - * own, and is exposed as a non-owning view in `data`. The encoder writes into - * two of the arenas, split by whether `gather_stripes` is certain to copy the - * extent into `gathered_buffer`, so that `transient_buffer` can be released as - * soon as gathering completes. `encode_columns` predicts that split and - * `gather_stripes` refines it by measuring what the encoder wrote, so the two - * are not simply a function of the stripe layout. + * The encoded bytes of each (stripe, stream) pair occupy an aligned byte range (extent) within one + * of the arenas below. Streams with size that is not known in advance are written into the + * transient arena, which is freed as soon as gathering completes. */ struct encoded_data { rmm::device_uvector persistent_buffer; // extents that may be read in place @@ -116,36 +111,6 @@ struct encoded_data { hostdevice_2dvector streams; // streams of encoded data, per chunk }; -/** - * @brief Set of encoded extents placed in `encoded_data::transient_buffer`. - * - * `gather_stripes` has to copy every extent in the set into - * `encoded_data::gathered_buffer` before it can release that arena. - */ -class transient_extents { - public: - transient_extents(size_t num_stripes, size_t num_streams) - : _num_streams{num_streams}, _flags(num_stripes * num_streams, false) - { - } - - void insert(size_t stripe_id, size_t strm_id) { _flags[index(stripe_id, strm_id)] = true; } - - [[nodiscard]] bool contains(size_t stripe_id, size_t strm_id) const - { - return _flags[index(stripe_id, strm_id)]; - } - - private: - [[nodiscard]] size_t index(size_t stripe_id, size_t strm_id) const - { - return stripe_id * _num_streams + strm_id; - } - - size_t _num_streams; - std::vector _flags; -}; - /** * @brief Dictionary data for string columns and their device views, per column. */ From 6c57e108ae35f5fbaab8ec9cfbe887e6964558ae Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 28 Jul 2026 22:20:05 +0000 Subject: [PATCH 12/20] misc cleanup --- cpp/src/io/orc/writer_impl.cu | 32 ++++++++++++++------------------ cpp/tests/io/orc_test.cpp | 8 ++++---- 2 files changed, 18 insertions(+), 22 deletions(-) diff --git a/cpp/src/io/orc/writer_impl.cu b/cpp/src/io/orc/writer_impl.cu index 9c58d752c503..e1a234181bd0 100644 --- a/cpp/src/io/orc/writer_impl.cu +++ b/cpp/src/io/orc/writer_impl.cu @@ -853,8 +853,8 @@ struct segmented_valid_cnt_input { std::vector indices; }; -// Returns the encoded data, along with a [stripe][strm_id] flag for each extent placed in -// `encoded_data::transient_buffer`, flattened with `streams.size()` elements per row. +// Returns the encoded data, along with a [stripe][strm_id] flag per extent, set when the extent is +// placed in `encoded_data::transient_buffer`. Flattened with `streams.size()` elements per row. std::pair> encode_columns(orc_table_view const& orc_table, encoder_decimal_info&& dec_chunk_sizes, file_segmentation const& segmentation, @@ -862,6 +862,9 @@ std::pair> encode_columns(orc_table_view cons uint32_t uncomp_block_align, rmm::cuda_stream_view stream) { + CUDF_EXPECTS(extent_alignment >= uncomp_block_align, + "Internal ORC writer error: arena extents are not aligned enough for the codec"); + auto const num_columns = orc_table.num_columns(); hostdevice_2dvector chunks(num_columns, segmentation.num_rowgroups(), stream); @@ -969,7 +972,7 @@ std::pair> encode_columns(orc_table_view cons if (strm_id < 0) { continue; } size_t stripe_size = 0; - // Alignment padding is slack in itself, in every rowgroup of the extent. + // Alignment padding leaves a gap after every rowgroup's chunk, so it is slack in itself. bool has_slack = uncomp_block_align > 1; std::for_each(stripe.cbegin(), stripe.cend(), [&](auto rg_idx) { #if defined(__GNUC__) && (__GNUC__ >= 14) @@ -990,8 +993,8 @@ std::pair> encode_columns(orc_table_view cons (strm_type == CI_DICTIONARY) ? stripe_dict.char_count : (((stripe_dict.entry_count + 0x1ff) >> 9) * (512 * 4 + 2)); - // `char_count` is exact; the RLE-encoded lengths are a worst case. - has_slack |= strm_type != CI_DICTIONARY; + // Only the size of RLE-encoded lengths is an estimate + has_slack |= (strm_type != CI_DICTIONARY); } else { strm.lengths[strm_type] = 0; } @@ -1018,9 +1021,6 @@ std::pair> encode_columns(orc_table_view cons } } - CUDF_EXPECTS(extent_alignment >= uncomp_block_align, - "Internal ORC writer error: arena extents are not aligned enough for the codec"); - // Extents that `gather_stripes` is certain to compact go into `transient_buffer`, so that arena // can be freed as soon as gathering completes. std::vector extent_is_transient(segmentation.num_stripes() * num_streams, 0); @@ -1042,12 +1042,12 @@ std::pair> encode_columns(orc_table_view cons rmm::device_uvector persistent_buffer(persistent_arena_size, stream); rmm::device_uvector transient_buffer(transient_arena_size, stream); - // Zero-size extents keep a null pointer, matching the empty device_uvector they replaced. std::vector>> encoded_views( segmentation.num_stripes(), std::vector>(num_streams)); for (size_t s = 0; s < segmentation.num_stripes(); ++s) { for (size_t strm_id = 0; strm_id < num_streams; ++strm_id) { auto const idx = extent_idx(s, strm_id); + // Zero-size extents keep a null pointer, matching the empty device_uvector they replaced. if (extent_sizes[idx] == 0) { continue; } auto& arena = extent_is_transient[idx] ? transient_buffer : persistent_buffer; encoded_views[s][strm_id] = @@ -1055,7 +1055,7 @@ std::pair> encode_columns(orc_table_view cons } } - // Write per-chunk data_ptrs and apply alignment fix-up. + // Point each rowgroup's chunk at its place within the extent, rounded up to the codec alignment. for (auto const& stripe : segmentation.stripes) { for (size_t col_idx = 0; col_idx < num_columns; col_idx++) { for (int strm_type = 0; strm_type < CI_NUM_STREAMS; ++strm_type) { @@ -1152,16 +1152,12 @@ std::vector gather_stripes(size_t num_index_streams, return stripe_id * num_streams_in_data + strm_id; }; - // Pass 1: compute per-(stripe, stream) actual sizes and decide which need a tight - // gathered copy. Streams of single-rowgroup stripes are read in place, to avoid the - // overhead of the additional copy. When there are multiple rowgroups, the chunks are - // copied anyway to make them contiguous. + // Compute per-(stripe, stream) actual sizes and decide which need a gathered copy. struct gather_info { - size_t actual_size; - bool gathered; + size_t actual_size{0}; + bool gathered{false}; }; - std::vector gather_meta(segmentation.num_stripes() * num_streams_in_data, - gather_info{0, false}); + std::vector gather_meta(segmentation.num_stripes() * num_streams_in_data); for (auto const& stripe : segmentation.stripes) { for (size_t col_idx = 0; col_idx < enc_data->streams.size().first; col_idx++) { diff --git a/cpp/tests/io/orc_test.cpp b/cpp/tests/io/orc_test.cpp index 9eb500f17a11..89ddc93e0bdb 100644 --- a/cpp/tests/io/orc_test.cpp +++ b/cpp/tests/io/orc_test.cpp @@ -1263,10 +1263,10 @@ TEST_P(OrcWriterTestDecimal, Decimal64) CUDF_TEST_EXPECT_COLUMNS_EQUAL(tbl.column(0), result.tbl->view().column(0)); } -// The row counts span stripes of one and of several rowgroups. Decimal data stream sizes are known -// exactly up front, and without compression the writer does not pad them to a block alignment, so -// the encoded chunks of a multi-rowgroup stripe come out contiguous and are written from the -// encoder output as-is rather than being compacted first. +// The cases with more than 10000 rows and no compression test the writer's non-compaction path, +// where encoded streams are written straight from the encoder output, because decimal data stream +// sizes are known exactly up front and uncompressed streams get no alignment padding, which leaves +// the chunks of a multi-rowgroup stripe already contiguous. INSTANTIATE_TEST_CASE_P(OrcWriterTest, OrcWriterTestDecimal, ::testing::Combine(::testing::Values(1, 10000, 10001, 34567), From 18193b4b10d474766dcce4ced08d9eae70ce1be3 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 28 Jul 2026 23:01:08 +0000 Subject: [PATCH 13/20] remove useless check --- cpp/src/io/orc/writer_impl.cu | 17 ++--------------- 1 file changed, 2 insertions(+), 15 deletions(-) diff --git a/cpp/src/io/orc/writer_impl.cu b/cpp/src/io/orc/writer_impl.cu index e1a234181bd0..4359e29e95af 100644 --- a/cpp/src/io/orc/writer_impl.cu +++ b/cpp/src/io/orc/writer_impl.cu @@ -1186,18 +1186,6 @@ std::vector gather_stripes(size_t num_index_streams, } } - // Every transient extent is gathered by construction, since `encode_columns` only makes an - // extent transient when its stripe spans several rowgroups, which is what the predicate above - // tests. Check it anyway, because the release below is a use-after-free if it ever stops - // holding, and the loop above cannot see stream types at or above CI_INDEX. - for (size_t s = 0; s < segmentation.num_stripes(); ++s) { - for (size_t strm_id = 0; strm_id < num_streams_in_data; ++strm_id) { - CUDF_EXPECTS( - not extent_is_transient[s][strm_id] or gather_meta[extent_idx(s, strm_id)].gathered, - "Internal ORC writer error: transient encoded extent was not gathered"); - } - } - // Lay out gather destinations within a single arena, with the same alignment rule as the // encoded arenas. std::vector gather_offsets(segmentation.num_stripes() * num_streams_in_data, 0); @@ -1266,9 +1254,8 @@ std::vector gather_stripes(size_t num_index_streams, } } - // Hold the gathered arena for lifetime management, and release the arena it copied from. - // `persistent_buffer` stays alive: some of its extents are read in place, via per-rowgroup - // data_ptrs that still point into it. + // Hold the gathered arena for lifetime management, and release the arena it copied from. Every + // transient extent was gathered above, so nothing points into `transient_buffer` any more. enc_data->gathered_buffer = std::move(gather_buffer); enc_data->transient_buffer = rmm::device_uvector{0, stream}; From 82e424e7a7469f89ae71cfb2ef4f3d1b22de84c3 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 28 Jul 2026 23:15:00 +0000 Subject: [PATCH 14/20] comment cleanup --- cpp/src/io/orc/writer_impl.cu | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/cpp/src/io/orc/writer_impl.cu b/cpp/src/io/orc/writer_impl.cu index 4359e29e95af..cf5ac3043396 100644 --- a/cpp/src/io/orc/writer_impl.cu +++ b/cpp/src/io/orc/writer_impl.cu @@ -1186,8 +1186,7 @@ std::vector gather_stripes(size_t num_index_streams, } } - // Lay out gather destinations within a single arena, with the same alignment rule as the - // encoded arenas. + // Lay out gather destinations in a single arena, with the same alignment as the encoded arenas. std::vector gather_offsets(segmentation.num_stripes() * num_streams_in_data, 0); size_t gather_total = 0; for (size_t s = 0; s < segmentation.num_stripes(); ++s) { @@ -1201,7 +1200,7 @@ std::vector gather_stripes(size_t num_index_streams, } rmm::device_uvector gather_buffer(gather_total, stream); - // Pass 2: build strm_desc entries and record gather destination spans. + // Build strm_desc entries and record gather destination spans. std::vector> gather_views(segmentation.num_stripes() * num_streams_in_data, device_span{}); std::vector stripes(segmentation.num_stripes()); @@ -1254,8 +1253,7 @@ std::vector gather_stripes(size_t num_index_streams, } } - // Hold the gathered arena for lifetime management, and release the arena it copied from. Every - // transient extent was gathered above, so nothing points into `transient_buffer` any more. + // Hold the gathered arena for lifetime management, and release the arena it copied from. enc_data->gathered_buffer = std::move(gather_buffer); enc_data->transient_buffer = rmm::device_uvector{0, stream}; From ff139af78b70a242197a8e3ca25b537919776864 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Thu, 30 Jul 2026 01:02:11 +0000 Subject: [PATCH 15/20] Require the extent alignment to be a multiple of the codec alignment Extent bases are offset from the arena base by multiples of extent_alignment, so the codec's chunk alignment has to divide it. The previous >= check only implied that because every alignment in play today is a power of two. --- cpp/src/io/orc/writer_impl.cu | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/cpp/src/io/orc/writer_impl.cu b/cpp/src/io/orc/writer_impl.cu index 9181f887a98d..be094c4e4456 100644 --- a/cpp/src/io/orc/writer_impl.cu +++ b/cpp/src/io/orc/writer_impl.cu @@ -862,8 +862,9 @@ std::pair> encode_columns(orc_table_view cons uint32_t uncomp_block_align, rmm::cuda_stream_view stream) { - CUDF_EXPECTS(extent_alignment >= uncomp_block_align, - "Internal ORC writer error: arena extents are not aligned enough for the codec"); + CUDF_EXPECTS(uncomp_block_align > 0 and extent_alignment % uncomp_block_align == 0, + "Internal ORC writer error: extent alignment is not a multiple of the codec's chunk " + "alignment"); auto const num_columns = orc_table.num_columns(); hostdevice_2dvector chunks(num_columns, segmentation.num_rowgroups(), stream); From 27724c05fac7a145f44b6e1e6b13fd4b090378a9 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Thu, 6 Aug 2026 23:32:55 +0000 Subject: [PATCH 16/20] Describe each stream extent with a struct instead of parallel vectors The size, slack, arena placement and offset of a (stripe, stream) extent were four vectors keyed by the same index, one of them a vector, filled across three loops. Fold them into one extent_info per extent, and hand that to gather_stripes in place of the separate transient-flag array. --- cpp/src/io/orc/writer_impl.cu | 70 +++++++++++++++++++---------------- 1 file changed, 38 insertions(+), 32 deletions(-) diff --git a/cpp/src/io/orc/writer_impl.cu b/cpp/src/io/orc/writer_impl.cu index be094c4e4456..f91316d5fd3a 100644 --- a/cpp/src/io/orc/writer_impl.cu +++ b/cpp/src/io/orc/writer_impl.cu @@ -853,14 +853,23 @@ struct segmented_valid_cnt_input { std::vector indices; }; -// Returns the encoded data, along with a [stripe][strm_id] flag per extent, set when the extent is -// placed in `encoded_data::transient_buffer`. Flattened with `streams.size()` elements per row. -std::pair> encode_columns(orc_table_view const& orc_table, - encoder_decimal_info&& dec_chunk_sizes, - file_segmentation const& segmentation, - orc_streams const& streams, - uint32_t uncomp_block_align, - rmm::cuda_stream_view stream) +// Storage of one (stripe, stream) pair within an encoded arena. +struct extent_info { + size_t size{0}; // upper bound on what the encoder writes + size_t offset{0}; // byte offset within the arena + bool has_slack{false}; + bool is_transient{false}; // placed in `encoded_data::transient_buffer` +}; + +// Returns the encoded data, along with a [stripe][strm_id] description of every extent, flattened +// with `streams.size()` elements per row. +std::pair> encode_columns( + orc_table_view const& orc_table, + encoder_decimal_info&& dec_chunk_sizes, + file_segmentation const& segmentation, + orc_streams const& streams, + uint32_t uncomp_block_align, + rmm::cuda_stream_view stream) { CUDF_EXPECTS(uncomp_block_align > 0 and extent_alignment % uncomp_block_align == 0, "Internal ORC writer error: extent alignment is not a multiple of the codec's chunk " @@ -956,8 +965,7 @@ std::pair> encode_columns(orc_table_view cons auto const extent_idx = [num_streams](size_t stripe_id, size_t strm_id) { return stripe_id * num_streams + strm_id; }; - std::vector extent_sizes(segmentation.num_stripes() * num_streams, 0); - std::vector extent_has_slack(segmentation.num_stripes() * num_streams, false); + std::vector extents(segmentation.num_stripes() * num_streams); for (auto const& stripe : segmentation.stripes) { for (size_t col_idx = 0; col_idx < num_columns; col_idx++) { for (int strm_type = 0; strm_type < CI_NUM_STREAMS; ++strm_type) { @@ -1016,27 +1024,26 @@ std::pair> encode_columns(orc_table_view cons stripe_size += strm.lengths[strm_type] + uncomp_block_align - 1; }); - extent_sizes[extent_idx(stripe.id, strm_id)] = stripe_size; - extent_has_slack[extent_idx(stripe.id, strm_id)] = has_slack; + auto& extent = extents[extent_idx(stripe.id, strm_id)]; + extent.size = stripe_size; + extent.has_slack = has_slack; } } } // Extents that `gather_stripes` is certain to compact go into `transient_buffer`, so that arena // can be freed as soon as gathering completes. - std::vector extent_is_transient(segmentation.num_stripes() * num_streams, 0); - std::vector extent_offsets(segmentation.num_stripes() * num_streams, 0); size_t persistent_arena_size = 0; size_t transient_arena_size = 0; for (size_t s = 0; s < segmentation.num_stripes(); ++s) { for (size_t strm_id = 0; strm_id < num_streams; ++strm_id) { - auto const idx = extent_idx(s, strm_id); - if (extent_sizes[idx] == 0) { continue; } - extent_is_transient[idx] = segmentation.stripes[s].size > 1 and extent_has_slack[idx]; - auto& arena_size = extent_is_transient[idx] ? transient_arena_size : persistent_arena_size; + auto& extent = extents[extent_idx(s, strm_id)]; + if (extent.size == 0) { continue; } + extent.is_transient = segmentation.stripes[s].size > 1 and extent.has_slack; + auto& arena_size = extent.is_transient ? transient_arena_size : persistent_arena_size; arena_size = util::round_up_unsafe(arena_size, extent_alignment); - extent_offsets[idx] = arena_size; - arena_size += extent_sizes[idx]; + extent.offset = arena_size; + arena_size += extent.size; } } @@ -1047,12 +1054,11 @@ std::pair> encode_columns(orc_table_view cons segmentation.num_stripes(), std::vector>(num_streams)); for (size_t s = 0; s < segmentation.num_stripes(); ++s) { for (size_t strm_id = 0; strm_id < num_streams; ++strm_id) { - auto const idx = extent_idx(s, strm_id); + auto const& extent = extents[extent_idx(s, strm_id)]; // Zero-size extents keep a null pointer, matching the empty device_uvector they replaced. - if (extent_sizes[idx] == 0) { continue; } - auto& arena = extent_is_transient[idx] ? transient_buffer : persistent_buffer; - encoded_views[s][strm_id] = - device_span{arena.data() + extent_offsets[idx], extent_sizes[idx]}; + if (extent.size == 0) { continue; } + auto& arena = extent.is_transient ? transient_buffer : persistent_buffer; + encoded_views[s][strm_id] = device_span{arena.data() + extent.offset, extent.size}; } } @@ -1122,7 +1128,7 @@ std::pair> encode_columns(orc_table_view cons rmm::device_uvector{0, stream}, // filled by gather_stripes std::move(encoded_views), std::move(chunk_streams)}, - std::move(extent_is_transient)}; + std::move(extents)}; } // TODO: remove StripeInformation from this function and return strm_desc instead @@ -1132,8 +1138,8 @@ std::pair> encode_columns(orc_table_view cons * * @param[in] num_index_streams Total number of index streams * @param[in] segmentation stripe and rowgroup ranges - * @param[in] extent_is_transient Marks extents held in `enc_data->transient_buffer`, which must all - * be gathered + * @param[in] extents Extent descriptions [stripe][data_stream]; extents marked transient are held + * in `enc_data->transient_buffer` and must all be gathered * @param[in,out] enc_data ORC per-chunk streams of encoded data * @param[in,out] strm_desc List of stream descriptors [stripe][data_stream] * @param[in] stream CUDA stream used for device memory operations and kernel launches @@ -1141,7 +1147,7 @@ std::pair> encode_columns(orc_table_view cons */ std::vector gather_stripes(size_t num_index_streams, file_segmentation const& segmentation, - host_2dspan extent_is_transient, + host_2dspan extents, encoded_data* enc_data, hostdevice_2dvector* strm_desc, rmm::cuda_stream_view stream) @@ -1180,7 +1186,7 @@ std::vector gather_stripes(size_t num_index_streams, // Compact when the chunks are not already contiguous, i.e. when the encoder wrote less // than the extent was sized for. Extents in `transient_buffer` are compacted regardless, // so that arena can be released below. - bool const gathered = (stripe.size > 1 and (extent_is_transient[stripe.id][stream_id] or + bool const gathered = (stripe.size > 1 and (extents[stripe.id][stream_id].is_transient or allocated_stripe_size > actual_stripe_size)); gather_meta[extent_idx(stripe.id, stream_id)] = {actual_stripe_size, gathered}; } @@ -2497,7 +2503,7 @@ auto convert_table_to_orc_data(table_view const& input, compression, write_mode); - auto [enc_data, extent_is_transient] = encode_columns( + auto [enc_data, extents] = encode_columns( orc_table, std::move(dec_chunk_sizes), segmentation, streams, block_align, stream); stripe_dicts.on_encode_complete(stream); @@ -2511,7 +2517,7 @@ auto convert_table_to_orc_data(table_view const& input, segmentation.num_stripes(), num_data_streams, stream); auto stripes = gather_stripes(num_index_streams, segmentation, - host_2dspan{extent_is_transient, streams.size()}, + host_2dspan{extents, streams.size()}, &enc_data, &strm_descs, stream); From 5cc2caa2097bef8ce8f4127e7446232116972a27 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Thu, 6 Aug 2026 23:40:30 +0000 Subject: [PATCH 17/20] dedup --- cpp/src/io/orc/writer_impl.cu | 62 +++++++++++++++++------------------ 1 file changed, 30 insertions(+), 32 deletions(-) diff --git a/cpp/src/io/orc/writer_impl.cu b/cpp/src/io/orc/writer_impl.cu index f91316d5fd3a..4439e0274812 100644 --- a/cpp/src/io/orc/writer_impl.cu +++ b/cpp/src/io/orc/writer_impl.cu @@ -962,10 +962,8 @@ std::pair> encode_columns( // Compute per-rowgroup stream lengths and per-(stripe, strm_id) extent sizes, along with // whether an extent size may exceed what the encoder ends up writing. auto const num_streams = streams.size(); - auto const extent_idx = [num_streams](size_t stripe_id, size_t strm_id) { - return stripe_id * num_streams + strm_id; - }; - std::vector extents(segmentation.num_stripes() * num_streams); + std::vector extent_storage(segmentation.num_stripes() * num_streams); + auto const extents = host_2dspan{extent_storage, num_streams}; for (auto const& stripe : segmentation.stripes) { for (size_t col_idx = 0; col_idx < num_columns; col_idx++) { for (int strm_type = 0; strm_type < CI_NUM_STREAMS; ++strm_type) { @@ -1024,7 +1022,7 @@ std::pair> encode_columns( stripe_size += strm.lengths[strm_type] + uncomp_block_align - 1; }); - auto& extent = extents[extent_idx(stripe.id, strm_id)]; + auto& extent = extents[stripe.id][strm_id]; extent.size = stripe_size; extent.has_slack = has_slack; } @@ -1037,7 +1035,7 @@ std::pair> encode_columns( size_t transient_arena_size = 0; for (size_t s = 0; s < segmentation.num_stripes(); ++s) { for (size_t strm_id = 0; strm_id < num_streams; ++strm_id) { - auto& extent = extents[extent_idx(s, strm_id)]; + auto& extent = extents[s][strm_id]; if (extent.size == 0) { continue; } extent.is_transient = segmentation.stripes[s].size > 1 and extent.has_slack; auto& arena_size = extent.is_transient ? transient_arena_size : persistent_arena_size; @@ -1054,7 +1052,7 @@ std::pair> encode_columns( segmentation.num_stripes(), std::vector>(num_streams)); for (size_t s = 0; s < segmentation.num_stripes(); ++s) { for (size_t strm_id = 0; strm_id < num_streams; ++strm_id) { - auto const& extent = extents[extent_idx(s, strm_id)]; + auto const extent = extents[s][strm_id]; // Zero-size extents keep a null pointer, matching the empty device_uvector they replaced. if (extent.size == 0) { continue; } auto& arena = extent.is_transient ? transient_buffer : persistent_buffer; @@ -1128,7 +1126,7 @@ std::pair> encode_columns( rmm::device_uvector{0, stream}, // filled by gather_stripes std::move(encoded_views), std::move(chunk_streams)}, - std::move(extents)}; + std::move(extent_storage)}; } // TODO: remove StripeInformation from this function and return strm_desc instead @@ -1155,17 +1153,18 @@ std::vector gather_stripes(size_t num_index_streams, if (segmentation.num_stripes() == 0) { return {}; } auto const num_streams_in_data = enc_data->data[0].size(); - auto const extent_idx = [num_streams_in_data](size_t stripe_id, size_t strm_id) { - return stripe_id * num_streams_in_data + strm_id; - }; - // Compute per-(stripe, stream) actual sizes and decide which need a gathered copy. - struct gather_info { - size_t actual_size{0}; + // Compaction destination of one (stripe, stream) pair within the gather arena. + struct gather_extent { + size_t size{0}; // what the encoder actually wrote + size_t offset{0}; // byte offset within the arena bool gathered{false}; + device_span view{}; }; - std::vector gather_meta(segmentation.num_stripes() * num_streams_in_data); + std::vector gather_storage(segmentation.num_stripes() * num_streams_in_data); + auto const gather_extents = host_2dspan{gather_storage, num_streams_in_data}; + // Compute per-(stripe, stream) actual sizes and decide which need a gathered copy. for (auto const& stripe : segmentation.stripes) { for (size_t col_idx = 0; col_idx < enc_data->streams.size().first; col_idx++) { auto const& col_streams = (enc_data->streams)[col_idx]; @@ -1188,28 +1187,28 @@ std::vector gather_stripes(size_t num_index_streams, // so that arena can be released below. bool const gathered = (stripe.size > 1 and (extents[stripe.id][stream_id].is_transient or allocated_stripe_size > actual_stripe_size)); - gather_meta[extent_idx(stripe.id, stream_id)] = {actual_stripe_size, gathered}; + + auto& extent = gather_extents[stripe.id][stream_id]; + extent.size = actual_stripe_size; + extent.gathered = gathered; } } } // Lay out gather destinations in a single arena, with the same alignment as the encoded arenas. - std::vector gather_offsets(segmentation.num_stripes() * num_streams_in_data, 0); size_t gather_total = 0; for (size_t s = 0; s < segmentation.num_stripes(); ++s) { for (size_t strm_id = 0; strm_id < num_streams_in_data; ++strm_id) { - auto const idx = extent_idx(s, strm_id); - if (!gather_meta[idx].gathered) { continue; } - gather_total = util::round_up_unsafe(gather_total, extent_alignment); - gather_offsets[idx] = gather_total; - gather_total += gather_meta[idx].actual_size; + auto& extent = gather_extents[s][strm_id]; + if (!extent.gathered) { continue; } + gather_total = util::round_up_unsafe(gather_total, extent_alignment); + extent.offset = gather_total; + gather_total += extent.size; } } rmm::device_uvector gather_buffer(gather_total, stream); // Build strm_desc entries and record gather destination spans. - std::vector> gather_views(segmentation.num_stripes() * num_streams_in_data, - device_span{}); std::vector stripes(segmentation.num_stripes()); for (auto const& stripe : segmentation.stripes) { for (size_t col_idx = 0; col_idx < enc_data->streams.size().first; col_idx++) { @@ -1218,20 +1217,19 @@ std::vector gather_stripes(size_t num_index_streams, auto const stream_id = col_streams[0].ids[k]; if (stream_id == -1) { continue; } - auto const idx = extent_idx(stripe.id, stream_id); - auto const& meta = gather_meta[idx]; + auto& extent = gather_extents[stripe.id][stream_id]; uint8_t* dst_ptr = nullptr; - if (meta.gathered) { + if (extent.gathered) { // Non-null even when the extent is empty, unlike the empty device_uvector this replaced. // `init_batched_memcpy_kernel` repoints the per-rowgroup data_ptrs at this arena, which // is what lets `transient_buffer` be released without leaving them dangling. - dst_ptr = gather_buffer.data() + gather_offsets[idx]; - gather_views[idx] = device_span{dst_ptr, meta.actual_size}; + dst_ptr = gather_buffer.data() + extent.offset; + extent.view = device_span{dst_ptr, extent.size}; } auto* ss = &(*strm_desc)[stripe.id][stream_id - num_index_streams]; ss->data_ptr = dst_ptr; // null when not gathered; init_batched_memcpy_kernel skips - ss->stream_size = meta.actual_size; + ss->stream_size = extent.size; ss->first_chunk_id = stripe.first; ss->num_chunks = stripe.size; ss->column_id = col_idx; @@ -1255,8 +1253,8 @@ std::vector gather_stripes(size_t num_index_streams, // spans, so consumers that read enc_data->data observe the post-gather state. for (size_t stripe_id = 0; stripe_id < enc_data->data.size(); ++stripe_id) { for (size_t stream_id = 0; stream_id < num_streams_in_data; ++stream_id) { - auto const idx = extent_idx(stripe_id, stream_id); - if (gather_meta[idx].gathered) { enc_data->data[stripe_id][stream_id] = gather_views[idx]; } + auto const extent = gather_extents[stripe_id][stream_id]; + if (extent.gathered) { enc_data->data[stripe_id][stream_id] = extent.view; } } } From 1151b232be3af95c084639cb4cc215291b08629c Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Fri, 7 Aug 2026 03:24:00 +0000 Subject: [PATCH 18/20] bug fix --- cpp/src/io/orc/writer_impl.cu | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/cpp/src/io/orc/writer_impl.cu b/cpp/src/io/orc/writer_impl.cu index 4439e0274812..250c934baa7e 100644 --- a/cpp/src/io/orc/writer_impl.cu +++ b/cpp/src/io/orc/writer_impl.cu @@ -1007,7 +1007,10 @@ std::pair> encode_columns( } } else if (strm_type == CI_DATA && ck.type_kind == TypeKind::STRING && ck.encoding_kind == DIRECT_V2) { - strm.lengths[strm_type] = std::max(column.rowgroup_char_count(rg_idx), 1); + auto const char_count = column.rowgroup_char_count(rg_idx); + strm.lengths[strm_type] = std::max(char_count, 1); + // The `max` reserves a byte the encoder does not write + has_slack |= (char_count == 0); } else if (strm_type == CI_DATA && streams[strm_id].length == 0 && (ck.type_kind == DOUBLE || ck.type_kind == FLOAT)) { // Pass-through. The encoder reports this length back unchanged, so it is exact. From d3905e5d020afc7e20851463d15e3fb0b8f2c2bf Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Thu, 6 Aug 2026 20:44:24 -0700 Subject: [PATCH 19/20] zero size fix Co-authored-by: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> --- cpp/src/io/orc/writer_impl.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/io/orc/writer_impl.cu b/cpp/src/io/orc/writer_impl.cu index 250c934baa7e..22e0294d129b 100644 --- a/cpp/src/io/orc/writer_impl.cu +++ b/cpp/src/io/orc/writer_impl.cu @@ -1209,7 +1209,7 @@ std::vector gather_stripes(size_t num_index_streams, gather_total += extent.size; } } - rmm::device_uvector gather_buffer(gather_total, stream); + rmm::device_uvector gather_buffer(std::max(gather_total, 1), stream); // Build strm_desc entries and record gather destination spans. std::vector stripes(segmentation.num_stripes()); From 0afac63bae418920eb0b154fcc58414c34a4c4ee Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Mon, 10 Aug 2026 17:33:56 +0000 Subject: [PATCH 20/20] docs --- cpp/src/io/orc/writer_impl.cu | 26 +++++++++++++++++++------- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/cpp/src/io/orc/writer_impl.cu b/cpp/src/io/orc/writer_impl.cu index 22e0294d129b..122e2f98e24d 100644 --- a/cpp/src/io/orc/writer_impl.cu +++ b/cpp/src/io/orc/writer_impl.cu @@ -853,16 +853,28 @@ struct segmented_valid_cnt_input { std::vector indices; }; -// Storage of one (stripe, stream) pair within an encoded arena. +/** + * @brief Storage of one (stripe, stream) pair within an encoded arena. + */ struct extent_info { - size_t size{0}; // upper bound on what the encoder writes - size_t offset{0}; // byte offset within the arena - bool has_slack{false}; - bool is_transient{false}; // placed in `encoded_data::transient_buffer` + size_t size{0}; ///< upper bound on what the encoder writes + size_t offset{0}; ///< byte offset within the arena + bool has_slack{false}; ///< whether `size` is a strict upper bound rather than exact + bool is_transient{false}; ///< placed in `encoded_data::transient_buffer` }; -// Returns the encoded data, along with a [stripe][strm_id] description of every extent, flattened -// with `streams.size()` elements per row. +/** + * @brief Encodes the columns' data into the ORC stream layout. + * + * @param[in] orc_table Table to be written, with ORC-related information + * @param[in] dec_chunk_sizes Sizes of encoded decimal elements + * @param[in] segmentation stripe and rowgroup ranges + * @param[in] streams List of stream descriptors + * @param[in] uncomp_block_align Required alignment of the codec's chunks + * @param[in] stream CUDA stream used for device memory operations and kernel launches + * @return The encoded data, along with a [stripe][strm_id] description of every extent, flattened + * with `streams.size()` elements per row + */ std::pair> encode_columns( orc_table_view const& orc_table, encoder_decimal_info&& dec_chunk_sizes,