[WIP] Serialize/deserialize with libcudf pack/unpack

CHANGELOG.md

@@ -10,6 +10,7 @@
 - PR #4789 Disallow timestamp sum and diffs via binary ops
 - PR #4815 Add JNI total memory allocated API
 - PR #4906 Add Java bindings for interleave_columns
+- PR #4941 Add `cudf::pack` and `cudf::unpack`
 - PR #4938 Add Java bindings for strip
 - PR #4923 Add Java and JNI bindings for string split
 - PR #4975 Add Java bindings for first and last aggregate expressions based on nth
@@ -120,6 +121,7 @@
 - PR #4993 Remove Java memory prediction code
 - PR #4985 Add null_count to Python Column ctors and use already computed null_count when possible
 - PR #5002 Fix Column.__reduce__ to accept `null_count`
+- PR #5025 Serialize/deserialize with libcudf pack/unpack
 
 ## Bug Fixes
 

cpp/CMakeLists.txt

@@ -568,6 +568,7 @@ add_library(cudf
             src/copying/slice.cpp
             src/copying/split.cpp
             src/copying/contiguous_split.cu
+            src/copying/pack.cpp
             src/copying/legacy/copy.cpp
             src/copying/legacy/gather.cu
             src/copying/legacy/scatter.cu

cpp/include/cudf/copying.hpp

@@ -454,7 +454,59 @@ std::vector<contiguous_split_result> contiguous_split(
   rmm::mr::device_memory_resource* mr = rmm::mr::get_default_resource());
 
 /**
- * @brief   Returns a new column, where each element is selected from either @p lhs or
+ * @brief Table data in a serialized format
+ * 
+ * Contains data from a table in two contiguous buffers: one on host, which contains table metadata
+ * and one on device which contains the table data.
+ */
+struct packed_columns {
+  packed_columns() = default;
+  packed_columns(std::unique_ptr<std::vector<uint8_t>> metadata,
+                 std::unique_ptr<rmm::device_buffer> data)
+    : metadata(std::move(metadata)), data(std::move(data)){};
+  std::unique_ptr<std::vector<uint8_t>> metadata;
+  std::unique_ptr<rmm::device_buffer> data;
+};
+
+/**
+ * @brief Deep-copy a `table_view` into a serialized contiguous memory format
+ * 
+ * The metadata from the `table_view` is copied into a host vector of bytes and the data from the
+ * `table_view` is copied into a `device_buffer`. Pass the output of this function into
+ * `cudf::experimental::unpack` to deserialize.
+ * 
+ * @param input View of the table to pack
+ * @param[in] mr Optional, The resource to use for all returned device allocations
+ * @return packed_columns A struct containing the serialized metadata and data in contiguous host
+ *         and device memory respectively
+ */
+packed_columns pack(std::vector<column_view> const& input,
+                    rmm::mr::device_memory_resource* mr = rmm::mr::get_default_resource());
+
+struct unpack_result {
+  std::vector<column_view> columns;
+  std::unique_ptr<rmm::device_buffer> all_data;
+};
+
+/**
+ * @brief Deserialize the result of `cudf::experimental::pack`
+ * 
+ * Converts the result of a serialized table into a `table_view` that points to the data stored in
+ * the contiguous device buffer `output.all_data`. The data for the table `output.all_data` is moved
+ * from the input `packed_table`'s member `table_data`. 
+ * 
+ * It is the caller's responsibility to ensure that the `table_view` in the output does not outlive 
+ * the device_buffer `all_data` in the output.
+ * 
+ * No new device memory is allocated in this function.
+ * 
+ * @param input The packed table to unpack
+ * @return contiguous_split_result The unpacked `table_view` and corresponding device data buffer
+ */
+unpack_result unpack(std::unique_ptr<packed_columns> input);
+
+/**
+ * @brief   Returns a new column, where each element is selected from either @p lhs or 
  *          @p rhs based on the value of the corresponding element in @p boolean_mask
  *
  * Selects each element i in the output column from either @p rhs or @p lhs using the following

cpp/include/cudf/detail/copy.hpp

@@ -74,6 +74,23 @@ std::vector<column_view> slice(column_view const& input,
                                std::vector<size_type> const& indices,
                                cudaStream_t stream = 0);
 
+/**
+ * @brief Information about the split for a given column. Bundled together
+ *        into a struct because tuples were getting pretty unreadable. 
+ */
+struct column_split_info {
+  size_t data_buf_size;      // size of the data (including padding)
+  size_t validity_buf_size;  // validity vector size (including padding)
+
+  size_t offsets_buf_size;  // (strings only) size of offset column (including padding)
+  size_type num_chars;      // (strings only) number of chars in the column
+  size_type chars_offset;   // (strings only) offset from head of chars data
+};
+
+unpack_result alloc_and_copy(std::vector<column_view> const& t,
+                             rmm::mr::device_memory_resource* mr,
+                             cudaStream_t stream);
+
 /**
  * @copydoc cudf::experimental::contiguous_split
  *
@@ -85,6 +102,15 @@ std::vector<contiguous_split_result> contiguous_split(
   rmm::mr::device_memory_resource* mr = rmm::mr::get_default_resource(),
   cudaStream_t stream                 = 0);
 
+/**
+ * @copydoc cudf::experimental::pack
+ *
+ * @param stream Optional CUDA stream on which to execute kernels
+ **/
+packed_columns pack(cudf::table_view const& input,
+                    cudaStream_t stream                 = 0,
+                    rmm::mr::device_memory_resource* mr = rmm::mr::get_default_resource());
+
 /**
  * @brief Creates an uninitialized new column of the specified size and same type as the `input`.
  * Supports only fixed-width types.

cpp/src/copying/contiguous_split.cu

@@ -18,6 +18,7 @@
 #include <cudf/column/column_device_view.cuh>
 #include <cudf/column/column_view.hpp>
 #include <cudf/copying.hpp>
+#include <cudf/detail/copy.hpp>
 #include <cudf/detail/null_mask.hpp>
 #include <cudf/detail/nvtx/ranges.hpp>
 #include <cudf/detail/utilities/cuda.cuh>
@@ -170,21 +171,7 @@ __launch_bounds__(block_size) __global__
 // start at that alignment.
 static constexpr size_t split_align = 64;
 
-/**
- * @brief Information about the split for a given column. Bundled together
- *        into a struct because tuples were getting pretty unreadable.
- */
-struct column_split_info {
-  size_t data_buf_size;      // size of the data (including padding)
-  size_t validity_buf_size;  // validity vector size (including padding)
-
-  size_t offsets_buf_size;  // (strings only) size of offset column (including padding)
-  size_type num_chars;      // (strings only) number of chars in the column
-  size_type chars_offset;   // (strings only) offset from head of chars data
-};
-
-/**
- * @brief Functor called by the `type_dispatcher` to incrementally compute total
+/** @brief Functor called by the `type_dispatcher` to incrementally compute total
  * memory buffer size needed to allocate a contiguous copy of all columns within
  * a source table.
  */
@@ -351,14 +338,10 @@ struct column_preprocess_info {
  * allocation.
  */
 thrust::host_vector<column_split_info> preprocess_string_column_info(
-  cudf::table_view const& t,
-  rmm::device_vector<column_split_info>& device_split_info,
-  cudaStream_t stream)
-{
-  // build a list of all the offset columns and their indices for all input string columns and put
-  // them on the gpu
+  std::vector<column_view> const& t, cudaStream_t stream) {
+  // build a list of all the offset columns and their indices for all input string columns and put them on the gpu
   thrust::host_vector<column_preprocess_info> offset_columns;
-  offset_columns.reserve(t.num_columns());  // worst case
+  offset_columns.reserve(t.size());  // worst case
 
   // collect only string columns
   size_type column_index = 0;
@@ -373,7 +356,7 @@ thrust::host_vector<column_split_info> preprocess_string_column_info(
   rmm::device_vector<column_preprocess_info> device_offset_columns = offset_columns;
 
   // compute column split information
-  rmm::device_vector<thrust::pair<size_type, size_type>> device_offsets(t.num_columns());
+  rmm::device_vector<thrust::pair<size_type, size_type>> device_offsets(t.size());
   auto* offsets_p = device_offsets.data().get();
   thrust::for_each(rmm::exec_policy(stream)->on(stream),
                    device_offset_columns.begin(),
@@ -384,7 +367,7 @@ thrust::host_vector<column_split_info> preprocess_string_column_info(
                                          cpi.offsets.head<int32_t>()[cpi.offset + cpi.size]);
                    });
   thrust::host_vector<thrust::pair<size_type, size_type>> host_offsets(device_offsets);
-  thrust::host_vector<column_split_info> split_info(t.num_columns());
+  thrust::host_vector<column_split_info> split_info(t.size());
   std::for_each(offset_columns.begin(),
                 offset_columns.end(),
                 [&split_info, &host_offsets](column_preprocess_info const& cpi) {
@@ -403,6 +386,8 @@ thrust::host_vector<column_split_info> preprocess_string_column_info(
   return split_info;
 }
 
+};  // anonymous namespace
+
 /**
  * @brief Creates a contiguous_split_result object which contains a deep-copy of the input
  * table_view into a single contiguous block of memory.
@@ -411,14 +396,11 @@ thrust::host_vector<column_split_info> preprocess_string_column_info(
  * call with the input table.  The memory referenced by the table_view and its internal column_views
  * is entirely contained in single block of memory.
  */
-contiguous_split_result alloc_and_copy(cudf::table_view const& t,
-                                       rmm::device_vector<column_split_info>& device_split_info,
-                                       rmm::mr::device_memory_resource* mr,
-                                       cudaStream_t stream)
-{
+unpack_result alloc_and_copy(std::vector<column_view> const& t,
+                             rmm::mr::device_memory_resource* mr,
+                             cudaStream_t stream) {
   // preprocess column split information for string columns.
-  thrust::host_vector<column_split_info> split_info =
-    preprocess_string_column_info(t, device_split_info, stream);
+  thrust::host_vector<column_split_info> split_info = preprocess_string_column_info(t, stream);
 
   // compute the rest of the column sizes (non-string columns, and total buffer size)
   size_t total_size      = 0;
@@ -437,7 +419,7 @@ contiguous_split_result alloc_and_copy(cudf::table_view const& t,
   // copy (this would be cleaner with a std::transform, but there's an nvcc compiler issue in the
   // way)
   std::vector<column_view> out_cols;
-  out_cols.reserve(t.num_columns());
+  out_cols.reserve(t.size());
 
   column_index = 0;
   std::for_each(
@@ -447,34 +429,26 @@ contiguous_split_result alloc_and_copy(cudf::table_view const& t,
       column_index++;
     });
 
-  return contiguous_split_result{cudf::table_view{out_cols}, std::move(device_buf)};
+  return unpack_result{out_cols, std::move(device_buf)};
 }
 
-};  // anonymous namespace
-
 std::vector<contiguous_split_result> contiguous_split(cudf::table_view const& input,
                                                       std::vector<size_type> const& splits,
                                                       rmm::mr::device_memory_resource* mr,
                                                       cudaStream_t stream)
 {
   auto subtables = cudf::experimental::split(input, splits);
 
-  // optimization : for large numbers of splits this allocation can dominate total time
-  //                spent if done inside alloc_and_copy().  so we'll allocate it once
-  //                and reuse it.
-  //
-  //                benchmark:        1 GB data, 10 columns, 256 splits.
-  //                no optimization:  106 ms (8 GB/s)
-  //                optimization:     20 ms (48 GB/s)
-  rmm::device_vector<column_split_info> device_split_info(input.num_columns());
-
   std::vector<contiguous_split_result> result;
-  std::transform(subtables.begin(),
-                 subtables.end(),
-                 std::back_inserter(result),
-                 [mr, stream, &device_split_info](table_view const& t) {
-                   return alloc_and_copy(t, device_split_info, mr, stream);
-                 });
+  std::transform(
+    subtables.begin(),
+    subtables.end(),
+    std::back_inserter(result),
+    [mr, stream](table_view const& t) {
+      std::vector<column_view> table_columns(t.begin(), t.end());
+      unpack_result result(alloc_and_copy(table_columns, mr, stream));
+      return contiguous_split_result{table_view(result.columns), std::move(result.all_data)};
+    });
 
   return result;
 }