Add a zero-copy reader for PyBytes

Change-Id: I131f49d73c73097c5ac250bfc0c60639358c0954

python/pyarrow/includes/libarrow_io.pxd

@@ -140,3 +140,12 @@ cdef extern from "arrow/io/hdfs.h" namespace "arrow::io" nogil:
                               int32_t buffer_size, int16_t replication,
                               int64_t default_block_size,
                               shared_ptr[HdfsOutputStream]* handle)
+
+
+cdef extern from "arrow/io/memory.h" namespace "arrow::io" nogil:
+    cdef cppclass BufferReader(ReadableFileInterface):
+        BufferReader(const uint8_t* data, int64_t nbytes)
+
+    cdef cppclass BufferOutputStream(OutputStream):
+        # TODO(wesm)
+        pass

python/pyarrow/includes/pyarrow.pxd

@@ -62,3 +62,6 @@ cdef extern from "pyarrow/io.h" namespace "pyarrow" nogil:
 
     cdef cppclass PyOutputStream(arrow_io.OutputStream):
         PyOutputStream(object fo)
+
+    cdef cppclass PyBytesReader(arrow_io.BufferReader):
+        PyBytesReader(object fo)

python/pyarrow/io.pyx

@@ -436,6 +436,20 @@ cdef class PythonFileInterface(NativeFile):
         self.is_open = True
 
 
+cdef class BytesReader(NativeFile):
+    cdef:
+        object obj
+
+    def __cinit__(self, obj):
+        if not isinstance(obj, bytes):
+            raise ValueError('Must pass bytes object')
+
+        self.obj = obj
+        self.is_readonly = 1
+        self.is_open = True
+
+        self.rd_file.reset(new pyarrow.PyBytesReader(obj))
+
 # ----------------------------------------------------------------------
 # Specialization for HDFS
 

python/pyarrow/tests/test_io.py

@@ -16,10 +16,6 @@
 # under the License.
 
 from io import BytesIO
-from os.path import join as pjoin
-import os
-import random
-
 import pytest
 
 import pyarrow.io as io
@@ -70,3 +66,28 @@ def test_python_file_read():
     assert f.read(50) == b'sample data'
 
     f.close()
+
+
+def test_bytes_reader():
+    # Like a BytesIO, but zero-copy underneath for C++ consumers
+    f = io.BytesReader(b'some sample data')
+
+    assert f.tell() == 0
+
+    assert f.read(4) == b'some'
+    assert f.tell() == 4
+
+    f.seek(0)
+    assert f.tell() == 0
+
+    f.seek(5)
+    assert f.tell() == 5
+
+    assert f.read(50) == b'sample data'
+
+    f.close()
+
+
+def test_bytes_reader_non_bytes():
+    with pytest.raises(ValueError):
+        io.BytesReader('some sample data')

python/src/pyarrow/io.cc

@@ -20,6 +20,7 @@
 #include <cstdlib>
 #include <sstream>
 
+#include <arrow/io/memory.h>
 #include <arrow/util/memory-pool.h>
 #include <arrow/util/status.h>
 
@@ -197,4 +198,18 @@ arrow::Status PyOutputStream::Write(const uint8_t* data, int64_t nbytes) {
   return file_->Write(data, nbytes);
 }
 
-} // namespace pyarrow
+// ----------------------------------------------------------------------
+// A readable file that is backed by a PyBytes
+
+PyBytesReader::PyBytesReader(PyObject* obj)
+    : arrow::io::BufferReader(reinterpret_cast<const uint8_t*>(PyBytes_AS_STRING(obj)),
+        PyBytes_GET_SIZE(obj)),
+      obj_(obj) {
+  Py_INCREF(obj_);
+}
+
+PyBytesReader::~PyBytesReader() {
+  Py_DECREF(obj_);
+}
+
+}  // namespace pyarrow

python/src/pyarrow/io.h

@@ -19,6 +19,7 @@
 #define PYARROW_IO_H
 
 #include "arrow/io/interfaces.h"
+#include "arrow/io/memory.h"
 
 #include "pyarrow/config.h"
 #include "pyarrow/visibility.h"
@@ -45,7 +46,7 @@ class PythonFile {
 };
 
 class PYARROW_EXPORT PyReadableFile : public arrow::io::ReadableFileInterface {
-public:
+ public:
   explicit PyReadableFile(PyObject* file);
   ~PyReadableFile();
 
@@ -62,12 +63,12 @@ public:
 
   bool supports_zero_copy() const override;
 
-private:
+ private:
   std::unique_ptr<PythonFile> file_;
 };
 
 class PYARROW_EXPORT PyOutputStream : public arrow::io::OutputStream {
-public:
+ public:
   explicit PyOutputStream(PyObject* file);
   ~PyOutputStream();
 
@@ -75,10 +76,20 @@ public:
   arrow::Status Tell(int64_t* position) override;
   arrow::Status Write(const uint8_t* data, int64_t nbytes) override;
 
-private:
+ private:
   std::unique_ptr<PythonFile> file_;
 };
 
+// A zero-copy reader backed by a PyBytes object
+class PYARROW_EXPORT PyBytesReader : public arrow::io::BufferReader {
+ public:
+  explicit PyBytesReader(PyObject* obj);
+  ~PyBytesReader();
+
+ private:
+  PyObject* obj_;
+};
+
 // TODO(wesm): seekable output files
 
 } // namespace pyarrow