Get ready to wrap Arrow buffers in a Python object

Change-Id: I4c56812272d083722ebf5f9d89de25ad9cfade15

python/pyarrow/array.pyx

@@ -38,7 +38,6 @@ import pyarrow.schema as schema
 
 def total_allocated_bytes():
     cdef MemoryPool* pool = pyarrow.GetMemoryPool()
-    return pool.bytes_allocated()
 
 
 cdef class Array:

python/pyarrow/includes/common.pxd

@@ -47,7 +47,3 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
         c_bool IsKeyError()
         c_bool IsNotImplemented()
         c_bool IsInvalid()
-
-    cdef cppclass Buffer:
-        uint8_t* data()
-        int64_t size()

python/pyarrow/includes/libarrow.pxd

@@ -54,6 +54,10 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
     cdef cppclass MemoryPool" arrow::MemoryPool":
         int64_t bytes_allocated()
 
+    cdef cppclass CBuffer" arrow::Buffer":
+        uint8_t* data()
+        int64_t size()
+
     cdef MemoryPool* default_memory_pool()
 
     cdef cppclass CListType" arrow::ListType"(CDataType):
@@ -201,7 +205,7 @@ cdef extern from "arrow/ipc/metadata.h" namespace "arrow::ipc" nogil:
         MessageType_DICTIONARY_BATCH" arrow::ipc::Message::DICTIONARY_BATCH"
 
     cdef cppclass Message:
-        CStatus Open(const shared_ptr[Buffer]& buf,
+        CStatus Open(const shared_ptr[CBuffer]& buf,
                      shared_ptr[Message]* out)
         int64_t body_length()
         MessageType type()

python/pyarrow/includes/libarrow_io.pxd

@@ -18,7 +18,7 @@
 # distutils: language = c++
 
 from pyarrow.includes.common cimport *
-from pyarrow.includes.libarrow cimport MemoryPool
+from pyarrow.includes.libarrow cimport MemoryPool, CBuffer
 
 cdef extern from "arrow/io/interfaces.h" namespace "arrow::io" nogil:
     enum FileMode" arrow::io::FileMode::type":
@@ -36,7 +36,7 @@ cdef extern from "arrow/io/interfaces.h" namespace "arrow::io" nogil:
         FileMode mode()
 
     cdef cppclass Readable:
-        CStatus ReadB" Read"(int64_t nbytes, shared_ptr[Buffer]* out)
+        CStatus ReadB" Read"(int64_t nbytes, shared_ptr[CBuffer]* out)
         CStatus Read(int64_t nbytes, int64_t* bytes_read, uint8_t* out)
 
     cdef cppclass Seekable:
@@ -57,7 +57,7 @@ cdef extern from "arrow/io/interfaces.h" namespace "arrow::io" nogil:
         CStatus ReadAt(int64_t position, int64_t nbytes,
                        int64_t* bytes_read, uint8_t* buffer)
         CStatus ReadAt(int64_t position, int64_t nbytes,
-                       int64_t* bytes_read, shared_ptr[Buffer]* out)
+                       int64_t* bytes_read, shared_ptr[CBuffer]* out)
 
     cdef cppclass WriteableFileInterface(OutputStream, Seekable):
         CStatus WriteAt(int64_t position, const uint8_t* data,
@@ -143,8 +143,9 @@ cdef extern from "arrow/io/hdfs.h" namespace "arrow::io" nogil:
 
 
 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 CBufferReader" arrow::io::BufferReader"\
+        (ReadableFileInterface):
+        CBufferReader(const uint8_t* data, int64_t nbytes)
 
     cdef cppclass BufferOutputStream(OutputStream):
         # TODO(wesm)

python/pyarrow/includes/pyarrow.pxd

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

python/pyarrow/io.pxd

@@ -22,6 +22,11 @@ from pyarrow.includes.libarrow cimport *
 from pyarrow.includes.libarrow_io cimport (ReadableFileInterface,
                                            OutputStream)
 
+cdef class Buffer:
+    cdef:
+        shared_ptr[CBuffer] buffer
+
+    cdef init(self, const shared_ptr[CBuffer]& buffer)
 
 cdef class NativeFile:
     cdef:

python/pyarrow/io.pyx

@@ -36,6 +36,161 @@ import re
 import sys
 import threading
 
+
+cdef class NativeFile:
+
+    def __cinit__(self):
+        self.is_open = False
+        self.own_file = False
+
+    def __dealloc__(self):
+        if self.is_open and self.own_file:
+            self.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, tb):
+        self.close()
+
+    def close(self):
+        if self.is_open:
+            with nogil:
+                if self.is_readonly:
+                    check_cstatus(self.rd_file.get().Close())
+                else:
+                    check_cstatus(self.wr_file.get().Close())
+        self.is_open = False
+
+    cdef read_handle(self, shared_ptr[ReadableFileInterface]* file):
+        self._assert_readable()
+        file[0] = <shared_ptr[ReadableFileInterface]> self.rd_file
+
+    cdef write_handle(self, shared_ptr[OutputStream]* file):
+        self._assert_writeable()
+        file[0] = <shared_ptr[OutputStream]> self.wr_file
+
+    def _assert_readable(self):
+        if not self.is_readonly:
+            raise IOError("only valid on readonly files")
+
+    def _assert_writeable(self):
+        if self.is_readonly:
+            raise IOError("only valid on writeonly files")
+
+    def size(self):
+        cdef int64_t size
+        self._assert_readable()
+        with nogil:
+            check_cstatus(self.rd_file.get().GetSize(&size))
+        return size
+
+    def tell(self):
+        cdef int64_t position
+        with nogil:
+            if self.is_readonly:
+                check_cstatus(self.rd_file.get().Tell(&position))
+            else:
+                check_cstatus(self.wr_file.get().Tell(&position))
+        return position
+
+    def seek(self, int64_t position):
+        self._assert_readable()
+        with nogil:
+            check_cstatus(self.rd_file.get().Seek(position))
+
+    def write(self, data):
+        """
+        Write bytes-like (unicode, encoded to UTF-8) to file
+        """
+        self._assert_writeable()
+
+        data = tobytes(data)
+
+        cdef const uint8_t* buf = <const uint8_t*> cp.PyBytes_AS_STRING(data)
+        cdef int64_t bufsize = len(data)
+        with nogil:
+            check_cstatus(self.wr_file.get().Write(buf, bufsize))
+
+    def read(self, int nbytes):
+        cdef:
+            int64_t bytes_read = 0
+            uint8_t* buf
+            shared_ptr[CBuffer] out
+
+        self._assert_readable()
+
+        with nogil:
+            check_cstatus(self.rd_file.get()
+                          .ReadB(nbytes, &out))
+
+        result = cp.PyBytes_FromStringAndSize(
+            <const char*>out.get().data(), out.get().size())
+
+        return result
+
+
+# ----------------------------------------------------------------------
+# Python file-like objects
+
+cdef class PythonFileInterface(NativeFile):
+    cdef:
+        object handle
+
+    def __cinit__(self, handle, mode='w'):
+        self.handle = handle
+
+        if mode.startswith('w'):
+            self.wr_file.reset(new pyarrow.PyOutputStream(handle))
+            self.is_readonly = 0
+        elif mode.startswith('r'):
+            self.rd_file.reset(new pyarrow.PyReadableFile(handle))
+            self.is_readonly = 1
+        else:
+            raise ValueError('Invalid file mode: {0}'.format(mode))
+
+        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))
+
+# ----------------------------------------------------------------------
+# Arrow buffers
+
+cdef class Buffer:
+
+    def __cinit__(self):
+        pass
+
+    cdef init(self, const shared_ptr[CBuffer]& buffer):
+        self.buffer = buffer
+
+    def to_pybytes(self):
+        return cp.PyBytes_FromStringAndSize(
+            <const char*>self.buffer.get().data(),
+            self.buffer.get().size())
+
+
+cdef class InMemoryOutputStream(NativeFile):
+
+    def get_result(self):
+        pass
+
+# ----------------------------------------------------------------------
+# HDFS IO implementation
+
 _HDFS_PATH_RE = re.compile('hdfs://(.*):(\d+)(.*)')
 
 try:
@@ -323,135 +478,6 @@ cdef class HdfsClient:
         f.download(stream)
 
 
-cdef class NativeFile:
-
-    def __cinit__(self):
-        self.is_open = False
-        self.own_file = False
-
-    def __dealloc__(self):
-        if self.is_open and self.own_file:
-            self.close()
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_value, tb):
-        self.close()
-
-    def close(self):
-        if self.is_open:
-            with nogil:
-                if self.is_readonly:
-                    check_cstatus(self.rd_file.get().Close())
-                else:
-                    check_cstatus(self.wr_file.get().Close())
-        self.is_open = False
-
-    cdef read_handle(self, shared_ptr[ReadableFileInterface]* file):
-        self._assert_readable()
-        file[0] = <shared_ptr[ReadableFileInterface]> self.rd_file
-
-    cdef write_handle(self, shared_ptr[OutputStream]* file):
-        self._assert_writeable()
-        file[0] = <shared_ptr[OutputStream]> self.wr_file
-
-    def _assert_readable(self):
-        if not self.is_readonly:
-            raise IOError("only valid on readonly files")
-
-    def _assert_writeable(self):
-        if self.is_readonly:
-            raise IOError("only valid on writeonly files")
-
-    def size(self):
-        cdef int64_t size
-        self._assert_readable()
-        with nogil:
-            check_cstatus(self.rd_file.get().GetSize(&size))
-        return size
-
-    def tell(self):
-        cdef int64_t position
-        with nogil:
-            if self.is_readonly:
-                check_cstatus(self.rd_file.get().Tell(&position))
-            else:
-                check_cstatus(self.wr_file.get().Tell(&position))
-        return position
-
-    def seek(self, int64_t position):
-        self._assert_readable()
-        with nogil:
-            check_cstatus(self.rd_file.get().Seek(position))
-
-    def write(self, data):
-        """
-        Write bytes-like (unicode, encoded to UTF-8) to file
-        """
-        self._assert_writeable()
-
-        data = tobytes(data)
-
-        cdef const uint8_t* buf = <const uint8_t*> cp.PyBytes_AS_STRING(data)
-        cdef int64_t bufsize = len(data)
-        with nogil:
-            check_cstatus(self.wr_file.get().Write(buf, bufsize))
-
-    def read(self, int nbytes):
-        cdef:
-            int64_t bytes_read = 0
-            uint8_t* buf
-            shared_ptr[Buffer] out
-
-        self._assert_readable()
-
-        with nogil:
-            check_cstatus(self.rd_file.get()
-                          .ReadB(nbytes, &out))
-
-        result = cp.PyBytes_FromStringAndSize(
-            <const char*>out.get().data(), out.get().size())
-
-        return result
-
-
-# ----------------------------------------------------------------------
-# Python file-like objects
-
-cdef class PythonFileInterface(NativeFile):
-    cdef:
-        object handle
-
-    def __cinit__(self, handle, mode='w'):
-        self.handle = handle
-
-        if mode.startswith('w'):
-            self.wr_file.reset(new pyarrow.PyOutputStream(handle))
-            self.is_readonly = 0
-        elif mode.startswith('r'):
-            self.rd_file.reset(new pyarrow.PyReadableFile(handle))
-            self.is_readonly = 1
-        else:
-            raise ValueError('Invalid file mode: {0}'.format(mode))
-
-        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