Expose pack/unpack API to Python

cpp/include/cudf/copying.hpp

@@ -591,7 +591,7 @@ packed_columns pack(cudf::table_view const& input,
  * guaranteeing that that all of the columns in the table point into `contiguous_buffer`.
  *
  * @param input View of the table to pack
- * @param contgiuous_buffer A contiguous buffer of device memory which contains the data referenced
+ * @param contiguous_buffer A contiguous buffer of device memory which contains the data referenced
  * by the columns in `table`
  * @param buffer_size The size of `contiguous_buffer`.
  * @return Vector of bytes representing the metadata used to `unpack` a packed_columns struct.

python/cudf/cudf/_lib/copying.pxd

@@ -0,0 +1,24 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.
+
+from cudf._lib.table cimport Table
+
+from cudf._lib.cpp.copying cimport packed_columns
+
+cdef class PackedColumns:
+    cdef packed_columns c_obj
+    cdef object column_names
+    cdef object column_dtypes
+    cdef object index_names
+
+    @staticmethod
+    cdef PackedColumns from_py_table(Table input_table, keep_index=*)
+
+    cdef Table unpack(self)
+
+    cdef const void* c_metadata_ptr(self) except *
+
+    cdef size_t c_metadata_size(self) except *
+
+    cdef void* c_gpu_data_ptr(self) except *
+
+    cdef size_t c_gpu_data_size(self) except *

python/cudf/cudf/_lib/copying.pyx

@@ -6,7 +6,9 @@ from libcpp cimport bool
 from libcpp.memory cimport make_unique, unique_ptr, shared_ptr, make_shared
 from libcpp.vector cimport vector
 from libcpp.utility cimport move
-from libc.stdint cimport int32_t, int64_t
+from libc.stdint cimport int32_t, int64_t, uint8_t, uintptr_t
+
+from rmm._lib.device_buffer cimport DeviceBuffer, device_buffer
 
 from cudf._lib.column cimport Column
 from cudf._lib.scalar import as_device_scalar
@@ -735,3 +737,203 @@ def segmented_gather(Column source_column, Column gather_map):
 
     result = Column.from_unique_ptr(move(c_result))
     return result
+
+
+cdef class PackedColumns:
+    """
+    A packed representation of a ``cudf.DataFrame``, with all columns residing
+    in a single GPU memory buffer.
+    """
+
+    def __reduce__(self):
+        return self.deserialize, self.serialize()
+
+    @property
+    def __cuda_array_interface__(self):
+        cdef dict intf = {
+            "data": (self.gpu_data_ptr, False),
+            "shape": (self.gpu_data_size,),
+            "strides": None,
+            "typestr": "|u1",
+            "version": 0
+        }
+        return intf
+
+    @property
+    def metadata_ptr(self):
+        return int(<uintptr_t>self.c_metadata_ptr())
+
+    @property
+    def metadata_size(self):
+        return int(self.c_metadata_size())
+
+    @property
+    def gpu_data_ptr(self):
+        return int(<uintptr_t>self.c_gpu_data_ptr())
+
+    @property
+    def gpu_data_size(self):
+        return int(self.c_gpu_data_size())
+
+    def serialize(self):
+        header = {}
+
+        header["column-names"] = self.column_names
+        # we probably want to split dtypes into their own header/frames
+        header["column-dtypes"] = self.column_dtypes
+        header["index-names"] = self.index_names
+        header["gpu-data-ptr"] = self.gpu_data_ptr
+        header["gpu-data-size"] = self.gpu_data_size
+        header["metadata"] = list(
+            <uint8_t[:self.metadata_size]>self.c_metadata_ptr()
+        )
+
+        return header, []
+
+    @classmethod
+    def deserialize(cls, header, frames):
+        cdef PackedColumns p = PackedColumns.__new__(PackedColumns)
+
+        dbuf = DeviceBuffer(
+            ptr=header["gpu-data-ptr"],
+            size=header["gpu-data-size"]
+        )
+
+        cdef cpp_copying.packed_columns data_
+        data_.metadata_ = move(
+            make_unique[cpp_copying.metadata](
+                move(<vector[uint8_t]>header["metadata"])
+            )
+        )
+        data_.gpu_data = move(dbuf.c_obj)
+
+        p.c_obj = move(data_)
+        p.column_names = header["column-names"]
+        p.column_dtypes = header["column-dtypes"]
+        p.index_names = header["index-names"]
+
+        return p
+
+    @staticmethod
+    cdef PackedColumns from_py_table(Table input_table, keep_index=True):
+        """
+        Construct a PackedColumns object from a ``cudf.DataFrame``.
+        """
+        from cudf.core import RangeIndex
+
+        cdef PackedColumns p = PackedColumns.__new__(PackedColumns)
+
+        if keep_index and not input_table.index.equals(
+            RangeIndex(start=0, stop=len(input_table), step=1)
+        ):
+            input_table_view = input_table.view()
+            p.index_names = input_table._index_names
+        else:
+            input_table_view = input_table.data_view()
+
+        p.c_obj = move(cpp_copying.pack(input_table_view))
+        p.column_names = input_table._column_names
+        p.column_dtypes = [c.dtype for c in input_table._columns]
+
+        return p
+
+    cdef Table unpack(self):
+        output_table = Table.from_table_view(
+            cpp_copying.unpack(self.c_obj),
+            self,
+            self.column_names,
+            self.index_names
+        )
+
+        # we need to do something with dtypes here
+
+        return output_table
+
+    cdef const void* c_metadata_ptr(self) except *:
+        return self.c_obj.metadata_.get()[0].data()
+
+    cdef size_t c_metadata_size(self) except *:
+        return self.c_obj.metadata_.get()[0].size()
+
+    cdef void* c_gpu_data_ptr(self) except *:
+        return self.c_obj.gpu_data.get()[0].data()
+
+    cdef size_t c_gpu_data_size(self) except *:
+        return self.c_obj.gpu_data.get()[0].size()
+
+
+def pack(Table input_table, keep_index=True):
+    """
+    Pack the columns of a ``cudf.DataFrame`` into a single GPU memory buffer.
+
+    Parameters
+    ----------
+    input_table : DataFrame
+        DataFrame to return packed representation of.
+
+    keep_index : boolean, default True
+        Pack the DataFrame's index columns along with the data columns.
+        This step will implicitly be skipped if the index columns can be
+        reconstructed on the fly at unpacking time.
+
+    Returns
+    -------
+    PackedColumns
+        A packed representation of ``input_table``.
+
+    Examples
+    --------
+    Build a dataframe and pack it:
+
+    >>> import cudf
+    >>> from cudf._lib.copying import pack
+    >>> df = cudf.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
+    >>> df
+        a   b
+    0   1   4
+    1   2   5
+    2   3   6
+    >>> packed = pack(df)
+    >>> packed
+    <cudf._lib.copying.PackedColumns at 0x7f34ff6a4b80>
+    """
+    return PackedColumns.from_py_table(input_table, keep_index)
+
+
+def unpack(PackedColumns packed):
+    """
+    Unpack the results of packing a ``cudf.DataFrame``, returning a new
+    ``DataFrame`` in the process.
+
+    Parameters
+    ----------
+    packed : PackedColumns
+        The results of calling ``pack()`` on a ``cudf.DataFrame``.
+
+    Returns
+    -------
+    cudf.DataFrame
+        A copy of the original DataFrame which was packed.
+
+    Examples
+    --------
+    Build a dataframe, pack it, then unpack it:
+
+    >>> import cudf
+    >>> from cudf._lib.copying import pack, unpack
+    >>> df = cudf.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
+    >>> df
+        a   b
+    0   1   4
+    1   2   5
+    2   3   6
+    >>> packed = pack(df)
+    >>> unpacked = unpack(packed)
+    >>> unpacked
+        a   b
+    0   1   4
+    1   2   5
+    2   3   6
+    """
+
+    return packed.unpack()

python/cudf/cudf/_lib/cpp/copying.pxd

@@ -3,7 +3,7 @@
 from rmm._lib.device_buffer cimport device_buffer
 
 from libcpp cimport bool
-from libc.stdint cimport int32_t, int64_t
+from libc.stdint cimport int32_t, int64_t, uint8_t
 from libcpp.memory cimport unique_ptr
 from libcpp.vector cimport vector
 
@@ -20,6 +20,12 @@ from cudf._lib.cpp.types cimport size_type
 
 ctypedef const scalar constscalar
 
+cdef extern from "cudf/copying.hpp" namespace "cudf::packed_columns" nogil:
+    cdef struct metadata:
+        metadata(vector[uint8_t]&& v)
+        const uint8_t* data () except +
+        size_type size () except +
+
 cdef extern from "cudf/copying.hpp" namespace "cudf" nogil:
     ctypedef enum out_of_bounds_policy:
         NULLIFY 'cudf::out_of_bounds_policy::NULLIFY'
@@ -111,6 +117,10 @@ cdef extern from "cudf/copying.hpp" namespace "cudf" nogil:
         vector[size_type] splits
     ) except +
 
+    cdef struct packed_columns:
+        unique_ptr[metadata] metadata_
+        unique_ptr[device_buffer] gpu_data
+
     cdef struct contiguous_split_result:
         table_view table
         vector[device_buffer] all_data
@@ -120,6 +130,10 @@ cdef extern from "cudf/copying.hpp" namespace "cudf" nogil:
         vector[size_type] splits
     ) except +
 
+    cdef packed_columns pack (const table_view& input) except +
+
+    cdef table_view unpack (const packed_columns& input) except +
+
     cdef unique_ptr[column] copy_if_else (
         column_view lhs,
         column_view rhs,