Fix incorrect slicing of GDS read/write calls

cpp/src/io/utilities/file_io_utilities.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021, NVIDIA CORPORATION.
+ * Copyright (c) 2021-2022, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -194,20 +194,13 @@ template <typename DataT,
 std::vector<std::future<ResultT>> make_sliced_tasks(
   F function, DataT* ptr, size_t offset, size_t size, cudf::detail::thread_pool& pool)
 {
+  constexpr size_t default_max_slice_size = 4 * 1024 * 1024;
+  static auto const max_slice_size = getenv_or("LIBCUDF_CUFILE_SLICE_SIZE", default_max_slice_size);
+  auto const slices                = make_file_io_slices(size, max_slice_size);
   std::vector<std::future<ResultT>> slice_tasks;
-  constexpr size_t default_max_slice_bytes = 4 * 1024 * 1024;
-  static auto const max_slice_bytes =
-    getenv_or("LIBCUDF_CUFILE_SLICE_SIZE", default_max_slice_bytes);
-  size_t const n_slices = util::div_rounding_up_safe(size, max_slice_bytes);
-  size_t slice_offset   = 0;
-  for (size_t t = 0; t < n_slices; ++t) {
-    DataT* ptr_slice = ptr + slice_offset;
-
-    size_t const slice_size = (t == n_slices - 1) ? size % max_slice_bytes : max_slice_bytes;
-    slice_tasks.push_back(pool.submit(function, ptr_slice, slice_size, offset + slice_offset));
-
-    slice_offset += slice_size;
-  }
+  std::transform(slices.cbegin(), slices.cend(), std::back_inserter(slice_tasks), [&](auto& slice) {
+    return pool.submit(function, ptr + slice.offset, slice.size, offset + slice.offset);
+  });
   return slice_tasks;
 }
 
@@ -318,6 +311,21 @@ std::unique_ptr<cufile_output_impl> make_cufile_output(std::string const& filepa
   return nullptr;
 }
 
+std::vector<file_io_slice> make_file_io_slices(size_t size, size_t max_slice_size)
+{
+  max_slice_size      = std::max(1024ul, max_slice_size);
+  auto const n_slices = util::div_rounding_up_safe(size, max_slice_size);
+  std::vector<file_io_slice> slices;
+  slices.reserve(n_slices);
+  std::generate_n(std::back_inserter(slices), n_slices, [&, idx = 0]() mutable {
+    auto const slice_offset = idx++ * max_slice_size;
+    auto const slice_size   = std::min(size - slice_offset, max_slice_size);
+    return file_io_slice{slice_offset, slice_size};
+  });
+
+  return slices;
+}
+
 }  // namespace detail
 }  // namespace io
 }  // namespace cudf

cpp/src/io/utilities/file_io_utilities.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021, NVIDIA CORPORATION.
+ * Copyright (c) 2021-2022, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -291,6 +291,21 @@ std::unique_ptr<cufile_input_impl> make_cufile_input(std::string const& filepath
  */
 std::unique_ptr<cufile_output_impl> make_cufile_output(std::string const& filepath);
 
+/**
+ * @brief Byte range to be read/written in a single operation.
+ */
+struct file_io_slice {
+  size_t offset;
+  size_t size;
+};
+
+/**
+ * @brief Split the total number of bytes to read/write into slices to enable parallel IO.
+ *
+ * If `max_slice_size` is below 1024, 1024 will be used instead to prevent potential misuse.
+ */
+std::vector<file_io_slice> make_file_io_slices(size_t size, size_t max_slice_size);
+
 }  // namespace detail
 }  // namespace io
 }  // namespace cudf

cpp/tests/CMakeLists.txt

@@ -199,6 +199,7 @@ ConfigureTest(
 ConfigureTest(DECOMPRESSION_TEST io/comp/decomp_test.cpp)
 
 ConfigureTest(CSV_TEST io/csv_test.cpp)
+ConfigureTest(FILE_IO_TEST io/file_io_test.cpp)
 ConfigureTest(ORC_TEST io/orc_test.cpp)
 ConfigureTest(PARQUET_TEST io/parquet_test.cpp)
 ConfigureTest(JSON_TEST io/json_test.cpp)

cpp/tests/io/file_io_test.cpp

@@ -0,0 +1,46 @@
+/*
+ * Copyright (c) 2022, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <cudf_test/base_fixture.hpp>
+#include <cudf_test/cudf_gtest.hpp>
+
+#include <src/io/utilities/file_io_utilities.hpp>
+
+#include <type_traits>
+
+// Base test fixture for tests
+struct CuFileIOTest : public cudf::test::BaseFixture {
+};
+
+TEST_F(CuFileIOTest, SliceSize)
+{
+  std::vector<std::pair<size_t, size_t>> test_cases{
+    {1 << 20, 1 << 18}, {1 << 18, 1 << 20}, {1 << 20, 3333}, {0, 1 << 18}, {0, 0}, {1 << 20, 0}};
+  for (auto const& test_case : test_cases) {
+    auto const slices = cudf::io::detail::make_file_io_slices(test_case.first, test_case.second);
+    if (slices.empty()) {
+      ASSERT_EQ(test_case.first, 0);
+    } else {
+      ASSERT_EQ(slices.front().offset, 0);
+      ASSERT_EQ(slices.back().offset + slices.back().size, test_case.first);
+      for (auto i = 1u; i < slices.size(); ++i) {
+        ASSERT_EQ(slices[i].offset, slices[i - 1].offset + slices[i - 1].size);
+      }
+    }
+  }
+}
+
+CUDF_TEST_PROGRAM_MAIN()