Adding support for list and struct type in ORC Reader

cpp/src/io/orc/orc.cpp

@@ -37,9 +37,7 @@ void ProtobufReader::skip_struct_field(int t)
     case PB_TYPE_FIXED64: skip_bytes(8); break;
     case PB_TYPE_FIXEDLEN: skip_bytes(get<uint32_t>()); break;
     case PB_TYPE_FIXED32: skip_bytes(4); break;
-    default:
-      // printf("invalid type (%d)\n", t);
-      break;
+    default: break;
   }
 }
 
@@ -471,20 +469,11 @@ void metadata::init_column_names() const
   auto const &types      = ff.types;
   for (int32_t col_id = 0; col_id < get_num_columns(); ++col_id) {
     std::string col_name;
-    uint32_t parent_idx = col_id;
-    uint32_t idx        = col_id;
-    do {
-      idx        = parent_idx;
-      parent_idx = (idx < types.size()) ? static_cast<uint32_t>(schema_idxs[idx].parent) : ~0;
-      if (parent_idx >= types.size()) break;
-
-      auto const field_idx =
-        (parent_idx < types.size()) ? static_cast<uint32_t>(schema_idxs[idx].field) : ~0;
-      if (field_idx < types[parent_idx].fieldNames.size()) {
-        col_name =
-          types[parent_idx].fieldNames[field_idx] + (col_name.empty() ? "" : ("." + col_name));
-      }
-    } while (parent_idx != idx);
+    uint32_t parent_idx = static_cast<uint32_t>(schema_idxs[col_id].parent);
+    uint32_t field_idx  = static_cast<uint32_t>(schema_idxs[col_id].field);
+    if (field_idx < types[parent_idx].fieldNames.size()) {
+      col_name = types[parent_idx].fieldNames[field_idx];
+    }
     // If we have no name (root column), generate a name
     column_names.push_back(col_name.empty() ? "col" + std::to_string(col_id) : col_name);
   }

cpp/src/io/orc/orc.h

@@ -537,6 +537,32 @@ class OrcDecompressor {
   std::vector<uint8_t> m_buf;
 };
 
+/**
+ * @brief Stores orc id for each column and its adjacent number of children
+ * in case of struct or number of children in case of list column.
+ * If list column has struct column, then all child columns of that struct are treated as child
+ * column of list.
+ *
+ * @code{.pseudo}
+ * Consider following data where a struct has two members and a list column
+ * {"struct": [{"a": 1, "b": 2}, {"a":3, "b":5}], "list":[[1, 2], [2, 3]]}
+ *
+ * `orc_column_meta` for struct column would be
+ * id = 0
+ * num_children = 2
+ *
+ * `orc_column_meta` for list column would be
+ * id = 3
+ * num_children = 1
+ * @endcode
+ *
+ */
+struct orc_column_meta {
+  // orc_column_meta(uint32_t _id, uint32_t _num_children) : id(_id), num_children(_num_children){};
+  uint32_t id;            // orc id for the column
+  uint32_t num_children;  // number of children at the same level of nesting in case of struct
+};
+
 /**
  * @brief A helper class for ORC file metadata. Provides some additional
  * convenience methods for initializing and accessing metadata.

cpp/src/io/orc/orc_gpu.h

@@ -95,7 +95,10 @@ struct ColumnDesc {
   uint32_t *valid_map_base;                // base pointer of valid bit map for this column
   void *column_data_base;                  // base pointer of column data
   uint32_t start_row;                      // starting row of the stripe
-  uint32_t num_rows;                       // starting row of the stripe
+  uint32_t num_rows;                       // number of rows in stripe
+  uint32_t column_num_rows;                // number of rows in whole column
+  uint32_t num_child_rows;                 // store number of child rows if it's list column
+  uint32_t num_rowgroups;                  // number of rowgroups in the chunk
   uint32_t dictionary_start;               // start position in global dictionary
   uint32_t dict_len;                       // length of local dictionary
   uint32_t null_count;                     // number of null values in this stripe's column
@@ -115,6 +118,9 @@ struct RowGroup {
   uint32_t chunk_id;        // Column chunk this entry belongs to
   uint32_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 num_rows;        // number of rows in rowgroup
+  uint32_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
 };
 
 /**
@@ -224,6 +230,9 @@ void PostDecompressionReassemble(CompressedStreamInfo *strm_info,
  * @param[in] num_columns Number of columns
  * @param[in] num_stripes Number of stripes
  * @param[in] num_rowgroups Number of row groups
+ * @param[in] rowidx_stride Row index stride
+ * @param[in] use_base_stride Whether to use base stride obtained from meta or use the computed
+ * value
  * @param[in] stream CUDA stream to use, default `rmm::cuda_stream_default`
  */
 void ParseRowGroupIndex(RowGroup *row_groups,
@@ -233,6 +242,7 @@ void ParseRowGroupIndex(RowGroup *row_groups,
                         uint32_t num_stripes,
                         uint32_t num_rowgroups,
                         uint32_t rowidx_stride,
+                        bool use_base_stride,
                         rmm::cuda_stream_view stream = rmm::cuda_stream_default);
 
 /**
@@ -242,15 +252,13 @@ void ParseRowGroupIndex(RowGroup *row_groups,
  * @param[in] global_dictionary Global dictionary device array
  * @param[in] num_columns Number of columns
  * @param[in] num_stripes Number of stripes
- * @param[in] max_rows Maximum number of rows to load
  * @param[in] first_row Crop all rows below first_row
  * @param[in] stream CUDA stream to use, default `rmm::cuda_stream_default`
  */
 void DecodeNullsAndStringDictionaries(ColumnDesc *chunks,
                                       DictionaryEntry *global_dictionary,
                                       uint32_t num_columns,
                                       uint32_t num_stripes,
-                                      size_t max_rows              = ~0,
                                       size_t first_row             = 0,
                                       rmm::cuda_stream_view stream = rmm::cuda_stream_default);
 
@@ -261,25 +269,25 @@ void DecodeNullsAndStringDictionaries(ColumnDesc *chunks,
  * @param[in] global_dictionary Global dictionary device array
  * @param[in] num_columns Number of columns
  * @param[in] num_stripes Number of stripes
- * @param[in] max_rows Maximum number of rows to load
  * @param[in] first_row Crop all rows below first_row
  * @param[in] tz_table Timezone translation table
  * @param[in] tz_len Length of timezone translation table
- * @param[in] row_groups Optional row index data
+ * @param[in] row_groups Optional row index data [rowgroup][column]
  * @param[in] num_rowgroups Number of row groups in row index data
  * @param[in] rowidx_stride Row index stride
+ * @param[in] level Current nesting level being processed
  * @param[in] stream CUDA stream to use, default `rmm::cuda_stream_default`
  */
-void DecodeOrcColumnData(ColumnDesc const *chunks,
+void DecodeOrcColumnData(ColumnDesc *chunks,
                          DictionaryEntry *global_dictionary,
+                         device_2dspan<RowGroup> row_groups,
                          uint32_t num_columns,
                          uint32_t num_stripes,
-                         size_t max_rows              = ~0,
                          size_t first_row             = 0,
                          timezone_table_view tz_table = {},
-                         const RowGroup *row_groups   = 0,
                          uint32_t num_rowgroups       = 0,
                          uint32_t rowidx_stride       = 0,
+                         size_t level                 = 0,
                          rmm::cuda_stream_view stream = rmm::cuda_stream_default);
 
 /**