diff --git a/cpp/benchmarks/io/orc/orc_reader_input.cpp b/cpp/benchmarks/io/orc/orc_reader_input.cpp index 1b40a4f4a67f..5fc39903c7d7 100644 --- a/cpp/benchmarks/io/orc/orc_reader_input.cpp +++ b/cpp/benchmarks/io/orc/orc_reader_input.cpp @@ -13,6 +13,8 @@ #include +#include + namespace { // Size of the data in the benchmark dataframe; chosen to be low enough to allow benchmarks to @@ -23,6 +25,8 @@ constexpr std::size_t Mbytes = 1024 * 1024; template void orc_read_common(cudf::size_type num_rows_to_read, + cudf::size_type num_cols_to_read, + std::size_t throughput_bytes, cuio_source_sink_pair& source_sink, nvbench::state& state) { @@ -62,13 +66,13 @@ void orc_read_common(cudf::size_type num_rows_to_read, auto const result = cudf::io::read_orc(read_opts); timer.stop(); - CUDF_EXPECTS(result.tbl->num_columns() == num_cols, "Unexpected number of columns"); + CUDF_EXPECTS(result.tbl->num_columns() == num_cols_to_read, "Unexpected number of columns"); CUDF_EXPECTS(result.tbl->num_rows() == num_rows_to_read, "Unexpected number of rows"); }); } auto const time = state.get_summary("nv/cold/time/gpu/mean").get_float64("value"); - state.add_element_count(static_cast(data_size) / time, "bytes_per_second"); + state.add_element_count(static_cast(throughput_bytes) / time, "bytes_per_second"); state.add_buffer_size( mem_stats_logger.peak_memory_usage(), "peak_memory_usage", "peak_memory_usage"); state.add_buffer_size(source_sink.size(), "encoded_file_size", "encoded_file_size"); @@ -79,9 +83,13 @@ void orc_read_common(cudf::size_type num_rows_to_read, template void BM_orc_read_data(nvbench::state& state, nvbench::type_list>) { + constexpr std::size_t single_column_data_size = 64 << 20; + auto const d_type = get_type_or_group(static_cast(DataType)); cudf::size_type const cardinality = state.get_int64("cardinality"); cudf::size_type const run_length = state.get_int64("run_length"); + auto const num_cols_to_read = static_cast(state.get_int64("num_cols")); + auto const bytes = num_cols_to_read == 1 ? single_column_data_size : data_size; auto const source_type = retrieve_io_type_enum(state.get_string("io_type")); auto const stripe_size_bytes = state.get_int64("stripe_size_bytes"); auto const stripe_size_rows = state.get_int64("stripe_size_rows"); @@ -89,8 +97,8 @@ void BM_orc_read_data(nvbench::state& state, nvbench::type_listview(); @@ -103,7 +111,7 @@ void BM_orc_read_data(nvbench::state& state, nvbench::type_list(num_rows_written, source_sink, state); + orc_read_common(num_rows_written, num_cols_to_read, bytes, source_sink, state); } template @@ -147,7 +155,7 @@ void orc_read_io_compression(nvbench::state& state) return view.num_rows(); }(); - orc_read_common(num_rows_written, source_sink, state); + orc_read_common(num_rows_written, num_cols, data_size, source_sink, state); } void BM_orc_read_io_compression(nvbench::state& state) @@ -175,6 +183,7 @@ NVBENCH_BENCH_TYPES(BM_orc_read_data, NVBENCH_TYPE_AXES(d_type_list)) .set_min_samples(4) .add_int64_axis("cardinality", {0, 1000}) .add_int64_axis("run_length", {1, 32}) + .add_int64_axis("num_cols", {1, num_cols}) .add_int64_axis("stripe_size_bytes", {0}) .add_int64_axis("stripe_size_rows", {0}); diff --git a/cpp/src/io/orc/orc_gpu.hpp b/cpp/src/io/orc/orc_gpu.hpp index ab2475fc28b5..ad847b0f36ce 100644 --- a/cpp/src/io/orc/orc_gpu.hpp +++ b/cpp/src/io/orc/orc_gpu.hpp @@ -79,6 +79,12 @@ enum stream_index_type { CI_NUM_STREAMS }; +// Streams that carry per-row-group positions in the row index, in the order they appear there. +constexpr int num_indexed_streams = CI_PRESENT + 1; +static_assert(CI_DATA == 0 && CI_DATA2 == 1 && CI_PRESENT == 2, + "Row index parsing and the per-stream arrays sized by num_indexed_streams assume " + "these are the first three stream types"); + /** * @brief Struct to describe a single entry in the global dictionary */ @@ -120,9 +126,11 @@ struct column_desc { * @brief Struct to describe a groups of row belonging to a column stripe */ struct row_group { - uint32_t chunk_id; // Column chunk this entry belongs to - int64_t strm_offset[2]; // Index offset for CI_DATA and CI_DATA2 streams - uint16_t run_pos[2]; // Run position for CI_DATA and CI_DATA2 + uint32_t chunk_id; // Column chunk this entry belongs to + // Index offset for the CI_DATA, CI_DATA2 and CI_PRESENT streams + int64_t strm_offset[num_indexed_streams]; + // Run position for the same streams; a bit position for CI_PRESENT and for BOOLEAN CI_DATA + uint16_t run_pos[num_indexed_streams]; uint32_t num_rows; // number of rows in rowgroup int64_t start_row; // starting row of the rowgroup uint32_t num_child_rows; // number of rows of children in rowgroup in case of list type @@ -289,6 +297,8 @@ void parse_row_group_index(row_group* row_groups, * @param[in] num_columns Number of columns * @param[in] num_stripes Number of stripes * @param[in] first_row Crop all rows below first_row + * @param[in] row_groups Row group descriptors [rowgroup][column], empty if the row index is unused + * @param[in] level Nesting level being decoded * @param[in] stream CUDA stream used for device memory operations and kernel launches */ void decode_nulls_and_string_dictionaries(column_desc* chunks, @@ -296,6 +306,8 @@ void decode_nulls_and_string_dictionaries(column_desc* chunks, size_type num_columns, size_type num_stripes, int64_t first_row, + device_2dspan row_groups, + size_t level, cuda::stream_ref stream); /** diff --git a/cpp/src/io/orc/reader_impl_decode.cu b/cpp/src/io/orc/reader_impl_decode.cu index d6181389e1cd..8701ee22b5c6 100644 --- a/cpp/src/io/orc/reader_impl_decode.cu +++ b/cpp/src/io/orc/reader_impl_decode.cu @@ -389,6 +389,7 @@ void decode_stream_data(int64_t num_dicts, auto& chunk = chunks[stripe_idx][col_idx]; chunk.column_data_base = out_buffers[col_idx].data(); chunk.valid_map_base = out_buffers[col_idx].null_mask(); + chunk.null_count = 0; }); }); @@ -396,8 +397,14 @@ void decode_stream_data(int64_t num_dicts, rmm::device_uvector global_dict(num_dicts, stream); chunks.host_to_device_async(stream); - decode_nulls_and_string_dictionaries( - chunks.base_device_ptr(), global_dict.data(), num_columns, num_stripes, skip_rows, stream); + decode_nulls_and_string_dictionaries(chunks.base_device_ptr(), + global_dict.data(), + num_columns, + num_stripes, + skip_rows, + row_groups, + level, + stream); if (level > 0) { // Update nullmasks for children if parent was a struct and had null mask diff --git a/cpp/src/io/orc/stripe_data.cu b/cpp/src/io/orc/stripe_data.cu index 92670a0465c7..6283603c7e75 100644 --- a/cpp/src/io/orc/stripe_data.cu +++ b/cpp/src/io/orc/stripe_data.cu @@ -13,6 +13,8 @@ #include #include +#include + namespace cudf::io::orc::detail { using cudf::io::detail::string_index_pair; @@ -27,6 +29,8 @@ constexpr int num_warps = 32; constexpr int block_size = 32 * num_warps; // Add some margin to look ahead to future rows in case there are many zeroes constexpr int row_decoder_buffer_size = block_size + 128; +// Longest byte RLE run, measured in the rows it covers in a PRESENT stream +constexpr uint32_t max_byte_rle_run_bits = 130 * 8; inline __device__ uint8_t is_rlev1(uint8_t encoding_mode) { return encoding_mode < DIRECT_V2; } inline __device__ uint8_t is_dictionary(uint8_t encoding_mode) { return encoding_mode & 1; } @@ -1262,6 +1266,9 @@ static __device__ int decode_decimals(orc_bytestream_s* bs, * @param[in] num_stripes Number of stripes * @param[in] max_num_rows Maximum number of rows to load * @param[in] first_row Crop all rows below first_row + * @param[in] row_groups Row group descriptors [rowgroup][column], empty if the row index is unused + * @param[in] decode_nulls_by_rowgroup Whether the null decode grid is sized one block per row group + * rather than one per (column, stripe) */ // blockDim {block_size,1,1} template @@ -1270,7 +1277,9 @@ CUDF_KERNEL void __launch_bounds__(block_size) dictionary_entry* global_dictionary, size_type num_columns, size_type num_stripes, - int64_t first_row) + int64_t first_row, + device_2dspan row_groups, + bool decode_nulls_by_rowgroup) { __shared__ __align__(16) orcdec_state_s state_g; using warp_reduce = cub::WarpReduce; @@ -1280,31 +1289,34 @@ CUDF_KERNEL void __launch_bounds__(block_size) typename block_reduce::TempStorage bk_storage; } temp_storage; - orcdec_state_s* const s = &state_g; - // Need the modulo because we have twice as many threads as columns*stripes - uint32_t const column = blockIdx.x / num_stripes; - uint32_t const stripe = blockIdx.x % num_stripes; - uint32_t const chunk_id = stripe * num_columns + column; - int t = threadIdx.x; + orcdec_state_s* const s = &state_g; + int t = threadIdx.x; + auto const num_rowgroups = static_cast(row_groups.size().first); + bool const is_nulldec = (blockIdx.y == 0); + // The null decode is spread over row groups when the host sized the grid for it + bool const by_rowgroup = is_nulldec && decode_nulls_by_rowgroup; + + uint32_t column, stripe, chunk_id; + size_type rowgroup_idx = 0; + if (by_rowgroup) { + column = blockIdx.x / num_rowgroups; + rowgroup_idx = blockIdx.x % num_rowgroups; + chunk_id = row_groups[rowgroup_idx][column].chunk_id; + stripe = chunk_id / num_columns; + } else { + if (blockIdx.x >= static_cast(num_columns) * num_stripes) { return; } + // Need the modulo because we have twice as many threads as columns*stripes + column = blockIdx.x / num_stripes; + stripe = blockIdx.x % num_stripes; + chunk_id = stripe * num_columns + column; + } if (t == 0) s->chunk = chunks[chunk_id]; __syncthreads(); size_t const max_num_rows = s->chunk.column_num_rows - s->chunk.parent_validity_info.null_count; - bool const is_nulldec = (blockIdx.y == 0); if (is_nulldec) { uint32_t null_count = 0; - // Decode NULLs - if (t == 0) { - s->chunk.skip_count = 0; - s->top.nulls_desc_row = 0; - bytestream_init(&s->bs, s->chunk.streams[CI_PRESENT], s->chunk.strm_len[CI_PRESENT]); - } - __syncthreads(); - if (s->chunk.strm_len[CI_PRESENT] == 0) { - // No present stream: all rows are valid - s->vals.u32[t] = ~0; - } auto const prev_parent_null_count = (s->chunk.parent_null_count_prefix_sums != nullptr && stripe > 0) ? s->chunk.parent_null_count_prefix_sums[stripe - 1] @@ -1314,20 +1326,58 @@ CUDF_KERNEL void __launch_bounds__(block_size) ? s->chunk.parent_null_count_prefix_sums[stripe] - prev_parent_null_count : 0; auto const num_elems = s->chunk.num_rows - parent_null_count; - while (s->top.nulls_desc_row < num_elems) { + + // Row range of this chunk that this block is responsible for + int64_t begin_row = 0; + int64_t end_row = num_elems; + // Bits of the seeked-to RLE run that belong to earlier row groups and must be discarded + uint32_t bit_skip = 0; + if (by_rowgroup) { + auto const& rg = row_groups[rowgroup_idx][column]; + // At level 0 a row group's `start_row` is its offset within the chunk, which is also what + // the PRESENT stream position below is relative to. + begin_row = rg.start_row; + end_row = min(begin_row + static_cast(rg.num_rows), num_elems); + if (begin_row >= end_row) { return; } + // Clamping keeps a corrupt index from underflowing the batch size below. + bit_skip = min(static_cast(rg.run_pos[CI_PRESENT]), max_byte_rle_run_bits); + } + + // Decode NULLs + if (t == 0) { + s->chunk.skip_count = 0; + s->top.nulls_desc_row = begin_row; + auto const strm_ofs = by_rowgroup + ? min(row_groups[rowgroup_idx][column].strm_offset[CI_PRESENT], + s->chunk.strm_len[CI_PRESENT]) + : 0; + bytestream_init(&s->bs, + s->chunk.streams[CI_PRESENT] + strm_ofs, + static_cast(s->chunk.strm_len[CI_PRESENT] - strm_ofs)); + } + __syncthreads(); + if (s->chunk.strm_len[CI_PRESENT] == 0) { + // No present stream: all rows are valid + s->vals.u32[t] = ~0; + bit_skip = 0; + } + while (s->top.nulls_desc_row < end_row) { + // The skipped bits share the decode buffer with the rows we want, so leave room for them auto const nrows_max = - static_cast(min(num_elems - s->top.nulls_desc_row, blockDim.x * 32ul)); + static_cast(min(end_row - s->top.nulls_desc_row, blockDim.x * 32ul - bit_skip)); bytestream_fill(&s->bs, t); __syncthreads(); uint32_t nrows; if (s->chunk.strm_len[CI_PRESENT] > 0) { - uint32_t nbytes = byte_rle(&s->bs, &s->u.rle8, s->vals.u8, (nrows_max + 7) >> 3, t); - nrows = min(nrows_max, nbytes * 8u); + uint32_t nbytes = + byte_rle(&s->bs, &s->u.rle8, s->vals.u8, (bit_skip + nrows_max + 7) >> 3, t); + nrows = min(nrows_max, nbytes * 8u - min(bit_skip, nbytes * 8u)); if (!nrows) { // Error: mark all remaining rows as null - nrows = nrows_max; + nrows = nrows_max; + bit_skip = 0; if (t * 32 < nrows) { s->vals.u32[t] = 0; } } } else { @@ -1338,12 +1388,15 @@ CUDF_KERNEL void __launch_bounds__(block_size) auto const row_in = s->chunk.start_row + s->top.nulls_desc_row - prev_parent_null_count; if (row_in + nrows > first_row && row_in < first_row + max_num_rows && s->chunk.valid_map_base != nullptr) { - int64_t dst_row = row_in - first_row; - int64_t dst_pos = max(dst_row, (int64_t)0); - uint32_t startbit = -static_cast(min(dst_row, (int64_t)0)); - uint32_t nbits = nrows - min(startbit, nrows); - uint32_t* valid = s->chunk.valid_map_base + (dst_pos >> 5); - uint32_t bitpos = static_cast(dst_pos) & 0x1f; + int64_t dst_row = row_in - first_row; + int64_t dst_pos = max(dst_row, int64_t{0}); + // Leading rows of this batch that fall below `first_row` and so are not stored. Unlike + // `bit_skip`, these are counted in `nrows`, so they also shorten the run of stored bits. + auto const rows_below_first = static_cast(-min(dst_row, int64_t{0})); + uint32_t startbit = bit_skip + rows_below_first; + uint32_t nbits = nrows - min(rows_below_first, nrows); + uint32_t* valid = s->chunk.valid_map_base + (dst_pos >> 5); + uint32_t bitpos = static_cast(dst_pos) & 0x1f; if ((size_t)(dst_pos + nbits) > max_num_rows) { nbits = static_cast(max_num_rows - min((size_t)dst_pos, max_num_rows)); } @@ -1393,14 +1446,25 @@ CUDF_KERNEL void __launch_bounds__(block_size) } __syncthreads(); if (t == 0) { s->top.nulls_desc_row += nrows; } + // Only the first batch starts part-way into an RLE run + bit_skip = 0; __syncthreads(); } __syncthreads(); // Sum up the valid counts and infer null_count null_count = block_reduce(temp_storage.bk_storage).Sum(null_count); if (t == 0) { - chunks[chunk_id].null_count = parent_null_count + null_count; - chunks[chunk_id].skip_count = s->chunk.skip_count; + if (by_rowgroup) { + // Every row group of the chunk contributes, and `parent_null_count` is zero on this path. + cuda::atomic_ref ref{chunks[chunk_id].null_count}; + ref.fetch_add(static_cast(null_count), cuda::std::memory_order_relaxed); + // `skip_count` arrives holding the index-stream bitmap that `parse_row_group_index_kernel` + // consumed, so it has to be overwritten even though `first_row` is zero here. + if (rowgroup_idx == s->chunk.rowgroup_id) { chunks[chunk_id].skip_count = 0; } + } else { + chunks[chunk_id].null_count = parent_null_count + null_count; + chunks[chunk_id].skip_count = s->chunk.skip_count; + } } } else { // Decode string dictionary @@ -2057,6 +2121,8 @@ CUDF_KERNEL void __launch_bounds__(block_size) * @param[in] num_columns Number of columns * @param[in] num_stripes Number of stripes * @param[in] first_row Crop all rows below first_row + * @param[in] row_groups Row group descriptors [rowgroup][column], empty if the row index is unused + * @param[in] level Nesting level being decoded * @param[in] stream CUDA stream used for device memory operations and kernel launches */ void __host__ decode_nulls_and_string_dictionaries(column_desc* chunks, @@ -2064,13 +2130,32 @@ void __host__ decode_nulls_and_string_dictionaries(column_desc* chunks, size_type num_columns, size_type num_stripes, int64_t first_row, + device_2dspan row_groups, + size_t level, cuda::stream_ref stream) { - dim3 dim_grid(num_columns * num_stripes, 2); + // A row index lets the null decode use one block per row group rather than one per stripe. Its + // PRESENT positions only line up with the output rows when nothing is skipped. + auto const num_rowgroups = static_cast(row_groups.size().first); + auto const rowgroup_blocks = static_cast(num_columns) * num_rowgroups; + auto const stripe_blocks = static_cast(num_columns) * num_stripes; + constexpr auto max_blocks = std::numeric_limits::max(); + bool const decode_nulls_by_rowgroup = + level == 0 && num_rowgroups > 0 && first_row == 0 && rowgroup_blocks <= max_blocks; + + // The dictionary half of the grid stays per (column, stripe) and ignores the extra blocks. + auto const nulldec_blocks = decode_nulls_by_rowgroup ? rowgroup_blocks : stripe_blocks; + CUDF_EXPECTS(nulldec_blocks <= max_blocks, "Too many stripes to decode in a single pass"); + dim3 dim_grid(static_cast(nulldec_blocks), 2); decode_nulls_and_string_dictionaries_kernel - <<>>( - chunks, global_dictionary, num_columns, num_stripes, first_row); + <<>>(chunks, + global_dictionary, + num_columns, + num_stripes, + first_row, + row_groups, + decode_nulls_by_rowgroup); CUDF_CUDA_TRY(cudaGetLastError()); } diff --git a/cpp/src/io/orc/stripe_init.cu b/cpp/src/io/orc/stripe_init.cu index 15cd3241dd3f..c37af9c4aef0 100644 --- a/cpp/src/io/orc/stripe_init.cu +++ b/cpp/src/io/orc/stripe_init.cu @@ -196,6 +196,14 @@ CUDF_KERNEL void __launch_bounds__(128, 8) } } +// Positions a row index entry records for each stream, in the order they are parsed +enum row_index_pos_e { + RI_BLOCK = 0, // Offset of the compression block holding the row group, zero if uncompressed + RI_OFFSET, // Offset of the row group within the decompressed stream + RI_RUN, // Position within the run found at that offset + RI_NUM_POS +}; + /** * @brief Shared mem state for parse_row_group_index_kernel */ @@ -204,11 +212,10 @@ struct rowindex_state_s { uint32_t rowgroup_start{}; uint32_t rowgroup_end{}; int is_compressed{}; - uint32_t row_index_entry[3] - [CI_PRESENT]{}; // NOTE: Assumes CI_PRESENT follows CI_DATA and CI_DATA2 - compressed_stream_info strm_info[2]{}; + uint32_t row_index_entry[RI_NUM_POS][num_indexed_streams]{}; + compressed_stream_info strm_info[num_indexed_streams]{}; row_group rowgroups[128]{}; - uint32_t compressed_offset[128][2]{}; + uint32_t compressed_offset[128][num_indexed_streams]{}; }; enum row_entry_state_e { @@ -232,7 +239,10 @@ static auto __device__ index_order_from_index_types(uint32_t index_types_bitmap) { constexpr cuda::std::array full_order = {CI_PRESENT, CI_DATA, CI_DATA2}; + // `copy_if` only fills an entry per stream the column actually indexes; the rest keep this + // marker so the parser can tell them apart from a real stream id. cuda::std::array partial_order; + partial_order.fill(CI_NUM_STREAMS); thrust::copy_if(thrust::seq, full_order.cbegin(), full_order.cend(), @@ -305,19 +315,20 @@ static uint32_t __device__ protobuf_parse_row_index_entry(rowindex_state_s* s, } break; case STORE_INDEX0: - // Start of a new entry; determine the stream index types - ci_id = stream_order[idx_id++]; + // Start of a new entry; determine the stream index type. Positions past the streams this + // column indexes read as CI_NUM_STREAMS and are skipped by the checks below. + ci_id = (idx_id < stream_order.size()) ? stream_order[idx_id++] : CI_NUM_STREAMS; if (s->is_compressed) { - if (ci_id < CI_PRESENT) s->row_index_entry[0][ci_id] = v; + if (ci_id < num_indexed_streams) s->row_index_entry[RI_BLOCK][ci_id] = v; if (cur >= start + pos_end) return length; state = STORE_INDEX1; break; } else { - if (ci_id < CI_PRESENT) s->row_index_entry[0][ci_id] = 0; + if (ci_id < num_indexed_streams) s->row_index_entry[RI_BLOCK][ci_id] = 0; // Fall through to STORE_INDEX1 for uncompressed (always block0) } case STORE_INDEX1: - if (ci_id < CI_PRESENT) s->row_index_entry[1][ci_id] = v; + if (ci_id < num_indexed_streams) s->row_index_entry[RI_OFFSET][ci_id] = v; if (cur >= start + pos_end) return length; state = (ci_id == CI_DATA && s->chunk.encoding_kind != DICTIONARY && s->chunk.encoding_kind != DICTIONARY_V2 && @@ -329,9 +340,11 @@ static uint32_t __device__ protobuf_parse_row_index_entry(rowindex_state_s* s, : STORE_INDEX2; break; case STORE_INDEX2: - if (ci_id < CI_PRESENT) { - // Boolean columns have an extra byte to indicate the position of the bit within the byte - s->row_index_entry[2][ci_id] = (s->chunk.type_kind == BOOLEAN) ? (v << 3) + *cur : v; + if (ci_id < num_indexed_streams) { + // Bit-packed streams have an extra byte to indicate the position of the bit within the + // byte; the PRESENT stream is always bit-packed, and so is the data of a BOOLEAN column + auto const is_bit_packed = (ci_id == CI_PRESENT) || (s->chunk.type_kind == BOOLEAN); + s->row_index_entry[RI_RUN][ci_id] = is_bit_packed ? (v << 3) + *cur : v; } if (ci_id == CI_PRESENT || s->chunk.type_kind == BOOLEAN) cur++; if (cur >= start + pos_end) return length; @@ -353,20 +366,19 @@ static __device__ void read_row_group_index_entries(rowindex_state_s* s, int num uint8_t const* index_data = s->chunk.streams[CI_INDEX]; int index_data_len = s->chunk.strm_len[CI_INDEX]; for (int i = 0; i < num_rowgroups; i++) { - s->row_index_entry[0][0] = 0; - s->row_index_entry[0][1] = 0; - s->row_index_entry[1][0] = 0; - s->row_index_entry[1][1] = 0; - s->row_index_entry[2][0] = 0; - s->row_index_entry[2][1] = 0; + for (int j = 0; j < num_indexed_streams; j++) { + s->row_index_entry[RI_BLOCK][j] = 0; + s->row_index_entry[RI_OFFSET][j] = 0; + s->row_index_entry[RI_RUN][j] = 0; + } if (index_data_len > 0) { int len = protobuf_parse_row_index_entry(s, index_data, index_data + index_data_len); index_data += len; index_data_len = max(index_data_len - len, 0); - for (int j = 0; j < 2; j++) { - s->rowgroups[i].strm_offset[j] = s->row_index_entry[1][j]; - s->rowgroups[i].run_pos[j] = s->row_index_entry[2][j]; - s->compressed_offset[i][j] = s->row_index_entry[0][j]; + for (int j = 0; j < num_indexed_streams; j++) { + s->rowgroups[i].strm_offset[j] = s->row_index_entry[RI_OFFSET][j]; + s->rowgroups[i].run_pos[j] = s->row_index_entry[RI_RUN][j]; + s->compressed_offset[i][j] = s->row_index_entry[RI_BLOCK][j]; } } } @@ -378,7 +390,7 @@ static __device__ void read_row_group_index_entries(rowindex_state_s* s, int num * @brief Translate block+offset compressed position into an uncompressed offset * * @param[in,out] s row group index state - * @param[in] ci_id index to convert (CI_DATA or CI_DATA2) + * @param[in] ci_id index to convert (one of the `num_indexed_streams` streams) * @param[in] num_rowgroups Number of index entries * @param[in] t thread id */ @@ -453,8 +465,9 @@ CUDF_KERNEL void __launch_bounds__(128, 8) if (t == 0) { s->chunk = chunks[chunk_id]; if (strm_info) { - if (s->chunk.strm_len[0] > 0) s->strm_info[0] = strm_info[s->chunk.strm_id[0]]; - if (s->chunk.strm_len[1] > 0) s->strm_info[1] = strm_info[s->chunk.strm_id[1]]; + for (int i = 0; i < num_indexed_streams; i++) { + if (s->chunk.strm_len[i] > 0) s->strm_info[i] = strm_info[s->chunk.strm_id[i]]; + } } uint32_t rowgroups_in_chunk = s->chunk.num_rowgroups; @@ -472,11 +485,10 @@ CUDF_KERNEL void __launch_bounds__(128, 8) __syncthreads(); if (s->is_compressed) { // Convert the block + blk_offset pair into a raw offset into the decompressed stream - if (s->chunk.strm_len[CI_DATA] > 0) { - map_row_index_to_uncompressed(s, CI_DATA, num_rowgroups, t); - } - if (s->chunk.strm_len[CI_DATA2] > 0) { - map_row_index_to_uncompressed(s, CI_DATA2, num_rowgroups, t); + for (int ci_id = 0; ci_id < num_indexed_streams; ci_id++) { + if (s->chunk.strm_len[ci_id] > 0) { + map_row_index_to_uncompressed(s, ci_id, num_rowgroups, t); + } } __syncthreads(); } diff --git a/cpp/tests/io/orc_test.cpp b/cpp/tests/io/orc_test.cpp index 4f957d547e9b..6aee4006b943 100644 --- a/cpp/tests/io/orc_test.cpp +++ b/cpp/tests/io/orc_test.cpp @@ -2038,6 +2038,45 @@ TEST_F(OrcWriterTest, EmptyRowGroup) CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result.tbl->view()); } +TEST_F(OrcReaderTest, NullDecodeSpanningRowGroups) +{ + // The reader only decodes nulls one row group per block when the row index is in use, which needs + // more rows than the 10000-row index stride. Reading the same file with the index disabled forces + // the whole-stripe decode instead, giving a direct comparison between the two paths. + constexpr cudf::size_type num_rows = 75'000; + + // Mix long runs with scattered nulls so both RLE run kinds appear in the PRESENT stream and row + // group boundaries land inside runs rather than neatly on them. + auto const valids = cudf::detail::make_counting_transform_iterator(0, [](auto i) { + if (i < 12'345) { return true; } + if (i < 12'400) { return false; } + return (i % 7) != 0; + }); + + auto const ints = cudf::detail::make_counting_transform_iterator(0, [](auto i) { return i; }); + int32_col int_column{ints, ints + num_rows, valids}; + + std::vector strings(num_rows); + std::generate(strings.begin(), strings.end(), [i = 0]() mutable { + return "value_" + std::to_string(i++ % 1000); + }); + str_col string_column{strings.begin(), strings.end(), valids}; + + table_view expected({int_column, string_column}); + + auto filepath = temp_env->get_temp_filepath("OrcNullDecodeRowGroups.orc"); + cudf::io::write_orc( + cudf::io::orc_writer_options::builder(cudf::io::sink_info{filepath}, expected).build()); + + auto const indexed = + cudf::io::read_orc(cudf::io::orc_reader_options::builder(cudf::io::source_info{filepath})); + auto const unindexed = cudf::io::read_orc( + cudf::io::orc_reader_options::builder(cudf::io::source_info{filepath}).use_index(false)); + + CUDF_TEST_EXPECT_TABLES_EQUAL(expected, indexed.tbl->view()); + CUDF_TEST_EXPECT_TABLES_EQUAL(indexed.tbl->view(), unindexed.tbl->view()); +} + TEST_F(OrcWriterTest, NoNullsAsNonNullable) { auto valids = cudf::test::iterators::no_nulls(); diff --git a/python/cudf/cudf/tests/input_output/test_orc.py b/python/cudf/cudf/tests/input_output/test_orc.py index 192a51dd6e7a..bde12e9b1cd8 100644 --- a/python/cudf/cudf/tests/input_output/test_orc.py +++ b/python/cudf/cudf/tests/input_output/test_orc.py @@ -149,6 +149,24 @@ def test_orc_reader_trailing_nulls(datadir): assert_eq(expect, got, check_categorical=True) +@pytest.mark.parametrize( + "orc_file", + [ + "TestOrcFile.nulls-at-end-snappy.orc", + "TestOrcFile.boolean_corruption_PR_6636.orc", + "TestOrcFile.boolean_corruption_PR_6702.orc", + ], +) +def test_orc_reader_null_decode_mid_run_positions(datadir, orc_file): + path = datadir / orc_file + + indexed = cudf.read_orc(path) + unindexed = cudf.read_orc(path, use_index=False) + + assert_eq(pd.read_orc(path), indexed) + assert_eq(unindexed, indexed) + + @pytest.mark.parametrize( "inputfile", ["TestOrcFile.testDate1900.orc", "TestOrcFile.testDate2038.orc"],