Prevent overflow with skip_rows in ORC and Parquet readers

cpp/CMakeLists.txt

@@ -395,6 +395,7 @@ add_library(
   src/io/utilities/datasource.cpp
   src/io/utilities/file_io_utilities.cpp
   src/io/utilities/parsing_utils.cu
+  src/io/utilities/row_selection.cpp
   src/io/utilities/trie.cu
   src/io/utilities/type_conversion.cpp
   src/jit/cache.cpp

cpp/include/cudf/io/orc.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -58,9 +58,9 @@ class orc_reader_options {
   // List of individual stripes to read (ignored if empty)
   std::vector<std::vector<size_type>> _stripes;
   // Rows to skip from the start;
-  size_type _skip_rows = 0;
-  // Rows to read; -1 is all
-  size_type _num_rows = -1;
+  int64_t _skip_rows = 0;
+  // Rows to read; `nullopt` is all
+  std::optional<size_type> _num_rows;
 
   // Whether to use row index to speed-up reading
   bool _use_index = true;
@@ -124,14 +124,14 @@ class orc_reader_options {
    *
    * @return Number of rows to skip from the start
    */
-  size_type get_skip_rows() const { return _skip_rows; }
+  int64_t get_skip_rows() const { return _skip_rows; }
 
   /**
    * @brief Returns number of row to read.
    *
    * @return Number of row to read
    */
-  size_type get_num_rows() const { return _num_rows; }
+  std::optional<size_type> const& get_num_rows() const { return _num_rows; }
 
   /**
    * @brief Whether to use row index to speed-up reading.
@@ -178,7 +178,8 @@ class orc_reader_options {
   void set_stripes(std::vector<std::vector<size_type>> stripes)
   {
     CUDF_EXPECTS(stripes.empty() or (_skip_rows == 0), "Can't set stripes along with skip_rows");
-    CUDF_EXPECTS(stripes.empty() or (_num_rows == -1), "Can't set stripes along with num_rows");
+    CUDF_EXPECTS(stripes.empty() or not _num_rows.has_value(),
+                 "Can't set stripes along with num_rows");
     _stripes = std::move(stripes);
   }
 
@@ -187,8 +188,9 @@ class orc_reader_options {
    *
    * @param rows Number of rows
    */
-  void set_skip_rows(size_type rows)
+  void set_skip_rows(int64_t rows)
   {
+    CUDF_EXPECTS(rows >= 0, "skip_rows cannot be negative");
     CUDF_EXPECTS(rows == 0 or _stripes.empty(), "Can't set both skip_rows along with stripes");
     _skip_rows = rows;
   }
@@ -200,7 +202,8 @@ class orc_reader_options {
    */
   void set_num_rows(size_type nrows)
   {
-    CUDF_EXPECTS(nrows == -1 or _stripes.empty(), "Can't set both num_rows along with stripes");
+    CUDF_EXPECTS(nrows >= 0, "num_rows cannot be negative");
+    CUDF_EXPECTS(_stripes.empty(), "Can't set both num_rows along with stripes");
     _num_rows = nrows;
   }
 
@@ -287,7 +290,7 @@ class orc_reader_options_builder {
    * @param rows Number of rows
    * @return this for chaining
    */
-  orc_reader_options_builder& skip_rows(size_type rows)
+  orc_reader_options_builder& skip_rows(int64_t rows)
   {
     options.set_skip_rows(rows);
     return *this;

cpp/include/cudf/io/parquet.hpp

@@ -58,9 +58,9 @@ class parquet_reader_options {
   // List of individual row groups to read (ignored if empty)
   std::vector<std::vector<size_type>> _row_groups;
   // Number of rows to skip from the start
-  size_type _skip_rows = 0;
-  // Number of rows to read; -1 is all
-  size_type _num_rows = -1;
+  int64_t _skip_rows = 0;
+  // Number of rows to read; `nullopt` is all
+  std::optional<size_type> _num_rows;
 
   // Whether to store string data as categorical type
   bool _convert_strings_to_categories = false;
@@ -136,14 +136,14 @@ class parquet_reader_options {
    *
    * @return Number of rows to skip from the start
    */
-  [[nodiscard]] size_type get_skip_rows() const { return _skip_rows; }
+  [[nodiscard]] int64_t get_skip_rows() const { return _skip_rows; }
 
   /**
    * @brief Returns number of rows to read.
    *
    * @return Number of rows to read
    */
-  [[nodiscard]] size_type get_num_rows() const { return _num_rows; }
+  [[nodiscard]] std::optional<size_type> const& get_num_rows() const { return _num_rows; }
 
   /**
    * @brief Returns names of column to be read, if set.
@@ -210,7 +210,7 @@ class parquet_reader_options {
    *
    * @param val Number of rows to skip from start
    */
-  void set_skip_rows(size_type val);
+  void set_skip_rows(int64_t val);
 
   /**
    * @brief Sets number of rows to read.
@@ -314,7 +314,7 @@ class parquet_reader_options_builder {
    * @param val Number of rows to skip from start
    * @return this for chaining
    */
-  parquet_reader_options_builder& skip_rows(size_type val)
+  parquet_reader_options_builder& skip_rows(int64_t val)
   {
     options.set_skip_rows(val);
     return *this;

cpp/src/io/functions.cpp

@@ -606,27 +606,25 @@ std::unique_ptr<std::vector<uint8_t>> parquet_chunked_writer::close(
 
 void parquet_reader_options::set_row_groups(std::vector<std::vector<size_type>> row_groups)
 {
-  if ((!row_groups.empty()) and ((_skip_rows != 0) or (_num_rows != -1))) {
+  if ((!row_groups.empty()) and ((_skip_rows != 0) or _num_rows.has_value())) {
     CUDF_FAIL("row_groups can't be set along with skip_rows and num_rows");
   }
 
   _row_groups = std::move(row_groups);
 }
 
-void parquet_reader_options::set_skip_rows(size_type val)
+void parquet_reader_options::set_skip_rows(int64_t val)
 {
-  if ((val != 0) and (!_row_groups.empty())) {
-    CUDF_FAIL("skip_rows can't be set along with a non-empty row_groups");
-  }
+  CUDF_EXPECTS(val >= 0, "skip_rows cannot be negative");
+  CUDF_EXPECTS(_row_groups.empty(), "skip_rows can't be set along with a non-empty row_groups");
 
   _skip_rows = val;
 }
 
 void parquet_reader_options::set_num_rows(size_type val)
 {
-  if ((val != -1) and (!_row_groups.empty())) {
-    CUDF_FAIL("num_rows can't be set along with a non-empty row_groups");
-  }
+  CUDF_EXPECTS(val >= 0, "num_rows cannot be negative");
+  CUDF_EXPECTS(_row_groups.empty(), "num_rows can't be set along with a non-empty row_groups");
 
   _num_rows = val;
 }

cpp/src/io/orc/aggregate_orc_metadata.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2021-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -14,7 +14,9 @@
  * limitations under the License.
  */
 
-#include "aggregate_orc_metadata.hpp"
+#include <io/orc/aggregate_orc_metadata.hpp>
+
+#include <io/utilities/row_selection.hpp>
 
 #include <algorithm>
 #include <numeric>
@@ -106,10 +108,10 @@ auto metadatas_from_sources(std::vector<std::unique_ptr<datasource>> const& sour
 
 }  // namespace
 
-size_type aggregate_orc_metadata::calc_num_rows() const
+int64_t aggregate_orc_metadata::calc_num_rows() const
 {
   return std::accumulate(
-    per_file_metadata.begin(), per_file_metadata.end(), 0, [](auto const& sum, auto const& pfm) {
+    per_file_metadata.begin(), per_file_metadata.end(), 0l, [](auto const& sum, auto const& pfm) {
       return sum + pfm.get_total_rows();
     });
 }
@@ -151,22 +153,29 @@ aggregate_orc_metadata::aggregate_orc_metadata(
   }
 }
 
-std::vector<metadata::stripe_source_mapping> aggregate_orc_metadata::select_stripes(
+std::tuple<int64_t, size_type, std::vector<metadata::stripe_source_mapping>>
+aggregate_orc_metadata::select_stripes(
   std::vector<std::vector<size_type>> const& user_specified_stripes,
-  size_type& row_start,
-  size_type& row_count,
+  int64_t skip_rows_opt,
+  std::optional<size_type> const& num_rows_opt,
   rmm::cuda_stream_view stream)
 {
+  CUDF_EXPECTS(
+    (skip_rows_opt == 0 and not num_rows_opt.has_value()) or user_specified_stripes.empty(),
+    "Can't use both the row selection and the stripe selection");
+
+  auto [rows_to_skip, rows_to_read] = [&]() {
+    if (not user_specified_stripes.empty()) { return std::pair<int64_t, size_type>{0, 0}; }
+    return cudf::io::detail::skip_rows_num_rows_from_options(
+      skip_rows_opt, num_rows_opt, get_num_rows());
+  }();
+
   std::vector<metadata::stripe_source_mapping> selected_stripes_mapping;
 
   if (!user_specified_stripes.empty()) {
     CUDF_EXPECTS(user_specified_stripes.size() == per_file_metadata.size(),
                  "Must specify stripes for each source");
-    // row_start is 0 if stripes are set. If this is not true anymore, then
-    // row_start needs to be subtracted to get the correct row_count
-    CUDF_EXPECTS(row_start == 0, "Start row index should be 0");
 
-    row_count = 0;
     // Each vector entry represents a source file; each nested vector represents the
     // user_defined_stripes to get from that source file
     for (size_t src_file_idx = 0; src_file_idx < user_specified_stripes.size(); ++src_file_idx) {
@@ -181,33 +190,24 @@ std::vector<metadata::stripe_source_mapping> aggregate_orc_metadata::select_stri
           "Invalid stripe index");
         stripe_infos.push_back(
           std::pair(&per_file_metadata[src_file_idx].ff.stripes[stripe_idx], nullptr));
-        row_count += per_file_metadata[src_file_idx].ff.stripes[stripe_idx].numberOfRows;
+        rows_to_read += per_file_metadata[src_file_idx].ff.stripes[stripe_idx].numberOfRows;
       }
       selected_stripes_mapping.push_back({static_cast<int>(src_file_idx), stripe_infos});
     }
   } else {
-    row_start = std::max(row_start, 0);
-    if (row_count < 0) {
-      row_count = static_cast<size_type>(
-        std::min<int64_t>(get_num_rows(), std::numeric_limits<size_type>::max()));
-    }
-    row_count = std::min(row_count, get_num_rows() - row_start);
-    CUDF_EXPECTS(row_count >= 0, "Invalid row count");
-    CUDF_EXPECTS(row_start <= get_num_rows(), "Invalid row start");
-
     size_type count            = 0;
     size_type stripe_skip_rows = 0;
     // Iterate all source files, each source file has corelating metadata
     for (size_t src_file_idx = 0;
-         src_file_idx < per_file_metadata.size() && count < row_start + row_count;
+         src_file_idx < per_file_metadata.size() && count < rows_to_skip + rows_to_read;
          ++src_file_idx) {
       std::vector<OrcStripeInfo> stripe_infos;
 
       for (size_t stripe_idx = 0; stripe_idx < per_file_metadata[src_file_idx].ff.stripes.size() &&
-                                  count < row_start + row_count;
+                                  count < rows_to_skip + rows_to_read;
            ++stripe_idx) {
         count += per_file_metadata[src_file_idx].ff.stripes[stripe_idx].numberOfRows;
-        if (count > row_start || count == 0) {
+        if (count > rows_to_skip || count == 0) {
           stripe_infos.push_back(
             std::pair(&per_file_metadata[src_file_idx].ff.stripes[stripe_idx], nullptr));
         } else {
@@ -218,7 +218,7 @@ std::vector<metadata::stripe_source_mapping> aggregate_orc_metadata::select_stri
       selected_stripes_mapping.push_back({static_cast<int>(src_file_idx), stripe_infos});
     }
     // Need to remove skipped rows from the stripes which are not selected.
-    row_start -= stripe_skip_rows;
+    rows_to_skip -= stripe_skip_rows;
   }
 
   // Read each stripe's stripefooter metadata
@@ -246,7 +246,7 @@ std::vector<metadata::stripe_source_mapping> aggregate_orc_metadata::select_stri
     }
   }
 
-  return selected_stripes_mapping;
+  return {rows_to_skip, rows_to_read, selected_stripes_mapping};
 }
 
 column_hierarchy aggregate_orc_metadata::select_columns(

cpp/src/io/orc/aggregate_orc_metadata.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2021-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -48,7 +48,7 @@ class aggregate_orc_metadata {
   /**
    * @brief Sums up the number of rows of each source
    */
-  [[nodiscard]] size_type calc_num_rows() const;
+  [[nodiscard]] int64_t calc_num_rows() const;
 
   /**
    * @brief Number of columns in a ORC file.
@@ -62,7 +62,7 @@ class aggregate_orc_metadata {
 
  public:
   std::vector<metadata> per_file_metadata;
-  size_type const num_rows;
+  int64_t const num_rows;
   size_type const num_stripes;
   bool row_grp_idx_present{true};
 
@@ -115,10 +115,10 @@ class aggregate_orc_metadata {
    *
    * Stripes are potentially selected from multiple files.
    */
-  std::vector<metadata::stripe_source_mapping> select_stripes(
+  std::tuple<int64_t, size_type, std::vector<metadata::stripe_source_mapping>> select_stripes(
     std::vector<std::vector<size_type>> const& user_specified_stripes,
-    size_type& row_start,
-    size_type& row_count,
+    int64_t row_start,
+    std::optional<size_type> const& num_rows_opt,
     rmm::cuda_stream_view stream);
 
   /**