Speed up ORC reader null decoding for narrow tables

cpp/benchmarks/io/orc/orc_reader_input.cpp

@@ -13,6 +13,8 @@
 
 #include <nvbench/nvbench.cuh>
 
+#include <cstddef>
+
 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 <bool is_chunked_read>
 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<double>(data_size) / time, "bytes_per_second");
+  state.add_element_count(static_cast<double>(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,18 +83,22 @@ void orc_read_common(cudf::size_type num_rows_to_read,
 template <data_type DataType>
 void BM_orc_read_data(nvbench::state& state, nvbench::type_list<nvbench::enum_type<DataType>>)
 {
+  constexpr std::size_t single_column_data_size = 64 << 20;
+
   auto const d_type                 = get_type_or_group(static_cast<int32_t>(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<cudf::size_type>(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");
   cuio_source_sink_pair source_sink(source_type);
 
   auto const num_rows_written = [&]() {
     auto const tbl = create_random_table(
-      cycle_dtypes(d_type, num_cols),
-      table_size_bytes{data_size},
+      cycle_dtypes(d_type, num_cols_to_read),
+      table_size_bytes{bytes},
       data_profile_builder().cardinality(cardinality).avg_run_length(run_length));
     auto const view = tbl->view();
 
@@ -103,7 +111,7 @@ void BM_orc_read_data(nvbench::state& state, nvbench::type_list<nvbench::enum_ty
     return view.num_rows();
   }();
 
-  orc_read_common<false>(num_rows_written, source_sink, state);
+  orc_read_common<false>(num_rows_written, num_cols_to_read, bytes, source_sink, state);
 }
 
 template <bool chunked_read>
@@ -147,7 +155,7 @@ void orc_read_io_compression(nvbench::state& state)
     return view.num_rows();
   }();
 
-  orc_read_common<chunked_read>(num_rows_written, source_sink, state);
+  orc_read_common<chunked_read>(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});
 

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,13 +297,17 @@ 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,
                                           dictionary_entry* global_dictionary,
                                           size_type num_columns,
                                           size_type num_stripes,
                                           int64_t first_row,
+                                          device_2dspan<row_group> row_groups,
+                                          size_t level,
                                           cuda::stream_ref stream);
 
 /**

cpp/src/io/orc/reader_impl_decode.cu

@@ -389,15 +389,22 @@ 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;
     });
   });
 
   // Allocate global dictionary for deserializing
   rmm::device_uvector<dictionary_entry> 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