Allocate single compressed buffer per source in parquet reader (rapidsai#21323)

## Summary

This PR changes the parquet reader to allocate one device buffer per source file for compressed page data, instead of potentially many buffers per chunk. This reduces the number of allocations from O(num_chunks) to O(num_sources), which can reduce memory fragmentation and allocator overhead.

## Implementation

The `read_column_chunks_async` function now has three steps:

1. Calculate total compressed size per source and the offset within the buffer for each chunk
2. Allocate one buffer per source with the total size needed
3. Issue reads (still coalescing adjacent contiguous chunks) into the appropriate offsets within each source's buffer

Individual chunks still reference their data via `ColumnChunkDesc::compressed_data` pointers into these buffers, so downstream code is unchanged.

Authors:
  - Vukasin Milovanovic (https://github.com/vuule)

Approvers:
  - Bradley Dice (https://github.com/bdice)
  - Muhammad Haseeb (https://github.com/mhaseeb123)

URL: rapidsai#21323

Author: vuule

cpp/src/io/parquet/reader_impl_preprocess.cu

@@ -473,7 +473,7 @@ std::pair<bool, std::future<void>> reader_impl::read_column_chunks()
   std::vector<size_type> chunk_source_map(num_chunks);
 
   // Tracker for eventually deallocating compressed and uncompressed data
-  raw_page_data = std::vector<rmm::device_buffer>(num_chunks);
+  raw_page_data = std::vector<rmm::device_buffer>(_num_sources);
 
   // Keep track of column chunk file offsets
   std::vector<size_t> column_chunk_offsets(num_chunks);

cpp/src/io/parquet/reader_impl_preprocess_utils.cu

@@ -28,6 +28,7 @@
 #include <thrust/transform.h>
 #include <thrust/transform_scan.h>
 
+#include <algorithm>
 #include <bitset>
 #include <iostream>
 #include <numeric>
@@ -174,57 +175,73 @@ void generate_depth_remappings(
   std::vector<size_type> const& chunk_source_map,
   rmm::cuda_stream_view stream)
 {
-  // Transfer chunk data, coalescing adjacent chunks
+  // Calculate total size per source and offset for each chunk
+  std::vector<size_t> source_total_size(sources.size(), 0);
+  std::vector<size_t> chunk_buffer_offset(end_chunk - begin_chunk);
+
+  for (size_t chunk = begin_chunk; chunk < end_chunk; ++chunk) {
+    auto const source_idx                    = chunk_source_map[chunk];
+    chunk_buffer_offset[chunk - begin_chunk] = source_total_size[source_idx];
+    source_total_size[source_idx] += chunks[chunk].compressed_size;
+  }
+
+  // Allocate one buffer per source
+  std::transform(
+    source_total_size.begin(), source_total_size.end(), page_data.begin(), [&](size_t total_size) {
+      return rmm::device_buffer(
+        cudf::util::round_up_safe(total_size, cudf::io::detail::BUFFER_PADDING_MULTIPLE), stream);
+    });
+  // device_read_async is not guaranteed to follow stream-ordering (see datasource API docs).
+  stream.synchronize();
+
+  // Issue reads, coalescing adjacent chunks
   std::vector<std::future<size_t>> read_tasks;
   for (size_t chunk = begin_chunk; chunk < end_chunk;) {
-    size_t const io_offset = column_chunk_offsets[chunk];
-    size_t io_size         = chunks[chunk].compressed_size;
-    size_t next_chunk      = chunk + 1;
+    auto const source_idx    = chunk_source_map[chunk];
+    auto const io_offset     = column_chunk_offsets[chunk];
+    size_t io_size           = chunks[chunk].compressed_size;
+    size_t const first_chunk = chunk;
+    size_t next_chunk        = chunk + 1;
+
     while (next_chunk < end_chunk) {
-      size_t const next_offset = column_chunk_offsets[next_chunk];
-      if (next_offset != io_offset + io_size ||
-          chunk_source_map[chunk] != chunk_source_map[next_chunk]) {
-        break;
-      }
+      if (chunk_source_map[next_chunk] != source_idx) { break; }
+      auto const next_offset = column_chunk_offsets[next_chunk];
+      if (next_offset != io_offset + io_size) { break; }
       io_size += chunks[next_chunk].compressed_size;
       next_chunk++;
     }
+
     if (io_size != 0) {
-      auto& source = sources[chunk_source_map[chunk]];
-      // Buffer needs to be padded.
-      // Required by `gpuDecodePageData`.
-      page_data[chunk] = rmm::device_buffer(
-        cudf::util::round_up_safe(io_size, cudf::io::detail::BUFFER_PADDING_MULTIPLE), stream);
+      auto& source = sources[source_idx];
+      auto* dest   = static_cast<uint8_t*>(page_data[source_idx].data()) +
+                   chunk_buffer_offset[first_chunk - begin_chunk];
 
       if (source->is_device_read_preferred(io_size)) {
-        auto fut_read_size = source->device_read_async(
-          io_offset, io_size, static_cast<uint8_t*>(page_data[chunk].data()), stream);
-        read_tasks.emplace_back(std::move(fut_read_size));
+        auto fut = source->device_read_async(io_offset, io_size, dest, stream);
+        read_tasks.emplace_back(std::move(fut));
       } else {
-        read_tasks.emplace_back(
-          std::async(std::launch::deferred,
-                     [source = std::ref(*source),
-                      io_offset,
-                      io_size,
-                      dest = page_data[chunk].data(),
-                      stream]() {
-                       auto const read_buffer = source.get().host_read(io_offset, io_size);
-                       cudf::detail::cuda_memcpy_async(
-                         cudf::device_span<uint8_t>{static_cast<uint8_t*>(dest), io_size},
-                         cudf::host_span<uint8_t const>{read_buffer->data(), io_size},
-                         stream);
-                       return io_size;
-                     }));
+        read_tasks.emplace_back(std::async(
+          std::launch::deferred, [source = std::ref(*source), io_offset, io_size, dest, stream]() {
+            auto const read_buffer = source.get().host_read(io_offset, io_size);
+            cudf::detail::cuda_memcpy_async(
+              cudf::device_span<uint8_t>{static_cast<uint8_t*>(dest), io_size},
+              cudf::host_span<uint8_t const>{read_buffer->data(), io_size},
+              stream);
+            return io_size;
+          }));
+      }
+
+      // Set compressed_data pointers for all coalesced chunks
+      auto* ptr = static_cast<uint8_t const*>(dest);
+      for (size_t c = first_chunk; c < next_chunk; ++c) {
+        chunks[c].compressed_data = ptr;
+        ptr += chunks[c].compressed_size;
       }
-      auto d_compdata = static_cast<uint8_t const*>(page_data[chunk].data());
-      do {
-        chunks[chunk].compressed_data = d_compdata;
-        d_compdata += chunks[chunk].compressed_size;
-      } while (++chunk != next_chunk);
-    } else {
-      chunk = next_chunk;
     }
+
+    chunk = next_chunk;
   }
+
   auto sync_fn = [](decltype(read_tasks) read_tasks) {
     for (auto& task : read_tasks) {
       task.get();