[REVIEW] Expose the HOST_BUFFER parquet sink to python

CHANGELOG.md

@@ -141,6 +141,7 @@
 - PR #5045 Remove call to `unique()` in concat when `axis=1`
 - PR #5085 Print more precise numerical strings in unit tests
 - PR #5028 Add Docker 19 support to local gpuci build
+- PR #5094 Expose the HOST_BUFFER parquet sink to python
 
 ## Bug Fixes
 

python/cudf/cudf/_lib/io/utils.pyx

@@ -5,6 +5,7 @@ from cpython.memoryview cimport PyMemoryView_FromMemory
 from libcpp.memory cimport unique_ptr
 from libcpp.pair cimport pair
 from libcpp.string cimport string
+from libcpp.vector cimport vector
 from cudf._lib.cpp.io.types cimport source_info, sink_info, data_sink, io_type
 
 import errno
@@ -21,6 +22,9 @@ cdef source_info make_source_info(src) except*:
             buf = src
         else:
             empty_buffer = True
+    elif isinstance(src, HostBuffer):
+        return source_info((<HostBuffer>src).buf.data(),
+                           (<HostBuffer>src).buf.size())
     elif isinstance(src, io.BytesIO):
         buf = src.getbuffer()
     # Otherwise src is expected to be a numeric fd, string path, or PathLike.
@@ -42,6 +46,8 @@ cdef source_info make_source_info(src) except*:
 
 # Converts the Python sink input to libcudf++ IO sink_info.
 cdef sink_info make_sink_info(src, unique_ptr[data_sink] * data) except*:
+    if isinstance(src, HostBuffer):
+        return sink_info(&(<HostBuffer>src).buf)
     if isinstance(src, io.IOBase):
         data.reset(new iobase_data_sink(src))
         return sink_info(data.get())
@@ -50,7 +56,6 @@ cdef sink_info make_sink_info(src, unique_ptr[data_sink] * data) except*:
     else:
         raise TypeError("Unrecognized input type: {}".format(type(src)))
 
-
 # Adapts a python io.IOBase object as a libcudf++ IO data_sink. This lets you
 # write from cudf to any python file-like object (File/BytesIO/SocketIO etc)
 cdef cppclass iobase_data_sink(data_sink):
@@ -67,3 +72,61 @@ cdef cppclass iobase_data_sink(data_sink):
 
     size_t bytes_written() with gil:
         return buf.tell()
+
+
+cdef class HostBuffer:
+    """ HostBuffer lets you spill cudf DataFrames from device memory to host
+    memory. Once in host memory the dataframe can either be re-loaded back
+    into gpu memory, or spilled to disk. This is designed to reduce the amount
+    of unnecessary host memory copies.
+
+    Examples
+    --------
+    .. code-block:: python
+      import shutil
+      import cudf
+
+      # read cudf DataFrame into buffer on host
+      df = cudf.DataFrame({'a': [1, 1, 2], 'b': [1, 2, 3]})
+      buffer = cudf.io.HostBuffer()
+      df.to_parquet(buffer)
+
+      # Copy HostBuffer back to DataFrame
+      cudf.read_parquet(df)
+
+      # Write HostBuffer to disk
+      shutil.copyfileobj(buffer, open("output.parquet", "wb"))
+    """
+    cdef vector[char] buf
+    cdef size_t pos
+
+    def __cinit__(self, int initial_capacity=0):
+        self.pos = 0
+        if initial_capacity:
+            self.buf.reserve(initial_capacity)
+
+    def __len__(self):
+        return self.buf.size()
+
+    def seek(self, size_t pos):
+        self.pos = pos
+
+    def tell(self):
+        return self.pos
+
+    def readall(self):
+        return self.read(-1)
+
+    def read(self, int n=-1):
+        if self.pos >= self.buf.size():
+            return b""
+
+        cdef size_t count = n
+        if ((n < 0) or (n > self.buf.size() - self.pos)):
+            count = self.buf.size() - self.pos
+
+        cdef size_t start = self.pos
+        self.pos += count
+
+        return PyMemoryView_FromMemory(self.buf.data() + start, count,
+                                       PyBUF_READ)

python/cudf/cudf/io/__init__.py

@@ -1,5 +1,6 @@
 # Copyright (c) 2018, NVIDIA CORPORATION.
 
+import cudf._lib.io.utils
 from cudf.io.avro import read_avro
 from cudf.io.csv import read_csv, to_csv
 from cudf.io.dlpack import from_dlpack
@@ -13,3 +14,5 @@
     read_parquet_metadata,
     write_to_dataset,
 )
+
+HostBuffer = cudf._lib.io.utils.HostBuffer

python/cudf/cudf/tests/test_parquet.py

@@ -609,6 +609,18 @@ def test_parquet_writer_bytes_io(simple_gdf):
     assert_eq(cudf.read_parquet(output), cudf.concat([simple_gdf, simple_gdf]))
 
 
+def test_parquet_writer_host_buffer(tmpdir, simple_gdf):
+    buffer = cudf.io.HostBuffer()
+    simple_gdf.to_parquet(buffer)
+
+    assert_eq(cudf.read_parquet(buffer), simple_gdf)
+
+    gdf_fname = tmpdir.join("gdf.parquet")
+    with open(gdf_fname, "wb") as o:
+        o.write(buffer.read())
+    assert_eq(cudf.read_parquet(gdf_fname), simple_gdf)
+
+
 @pytest.mark.parametrize("cols", [["b"], ["c", "b"]])
 def test_parquet_write_partitioned(tmpdir_factory, cols):
     # Checks that write_to_dataset is wrapping to_parquet