implement mutable arrow python buffers

cpp/src/plasma/plasma/plasma.pyx

@@ -25,8 +25,8 @@ from libcpp.string cimport string as c_string
 from libcpp.vector cimport vector as c_vector
 from libc.stdint cimport int64_t, uint8_t, uintptr_t
 
-from pyarrow.lib cimport Buffer, NativeFile, check_status
-from pyarrow.includes.libarrow cimport MutableBuffer, CBuffer, CFixedSizeBufferWrite, CStatus
+from pyarrow.lib cimport Buffer, MutableBuffer, NativeFile, check_status
+from pyarrow.includes.libarrow cimport CMutableBuffer, CBuffer, CFixedSizeBufferWrite, CStatus
 
 cdef class FixedSizeBufferOutputStream(NativeFile):
 
@@ -140,6 +140,24 @@ cdef class PlasmaBuffer(Buffer):
     """
     self.client.release(self.object_id)
 
+cdef class MutablePlasmaBuffer(MutableBuffer):
+
+  cdef:
+    ObjectID object_id
+    PlasmaClient client
+
+  def __cinit__(self, ObjectID object_id, PlasmaClient client):
+    """Initialize a PlasmaBuffer."""
+    self.object_id = object_id
+    self.client = client
+
+  def __dealloc__(self):
+    """Notify Plasma that the object is no longer needed.
+
+    If the plasma client has been shut down, then don't do anything.
+    """
+    self.client.release(self.object_id)
+
 cdef class PlasmaClient:
   """The PlasmaClient is used to interface with a plasma store and manager.
 
@@ -165,10 +183,17 @@ cdef class PlasmaClient:
     check_status(self.client.get().Get(ids.data(), ids.size(), timeout_ms, result[0].data()))
 
   cdef _make_plasma_buffer(self, ObjectID object_id, uint8_t* data, int64_t size):
-    cdef shared_ptr[MutableBuffer] buffer
-    buffer.reset(new MutableBuffer(data, size))
+    cdef shared_ptr[CBuffer] buffer
+    buffer.reset(new CBuffer(data, size))
     result = PlasmaBuffer(object_id, self)
-    result.init(<shared_ptr[CBuffer]>(buffer))
+    result.init(buffer)
+    return result
+
+  cdef _make_mutable_plasma_buffer(self, ObjectID object_id, uint8_t* data, int64_t size):
+    cdef shared_ptr[CMutableBuffer] buffer
+    buffer.reset(new CMutableBuffer(data, size))
+    result = MutablePlasmaBuffer(object_id, self)
+    result.init_mutable(buffer)
     return result
 
   def connect(self, store_socket_name, manager_socket_name, release_delay):
@@ -205,7 +230,7 @@ cdef class PlasmaClient:
     """
     cdef uint8_t* data
     check_status(self.client.get().Create(object_id.data, data_size, <uint8_t*>(metadata.data()), metadata.size(), &data))
-    return self._make_plasma_buffer(object_id, data, data_size)
+    return self._make_mutable_plasma_buffer(object_id, data, data_size)
 
   def get(self, object_ids, timeout_ms=-1):
     """Returns data buffer from the PlasmaStore based on object ID.

python/pyarrow/includes/libarrow.pxd

@@ -153,8 +153,9 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
         int64_t size()
         shared_ptr[CBuffer] parent()
 
-    cdef cppclass MutableBuffer(CBuffer):
-        MutableBuffer(const uint8_t* data, int64_t size)
+    cdef cppclass CMutableBuffer" arrow::MutableBuffer"(CBuffer):
+        CMutableBuffer(const uint8_t* data, int64_t size)
+        uint8_t* mutable_data()
 
     cdef cppclass ResizableBuffer(CBuffer):
         CStatus Resize(int64_t nbytes)

python/pyarrow/io.pxi

@@ -541,6 +541,30 @@ cdef class Buffer:
         return self.size
 
 
+cdef class MutableBuffer(Buffer):
+
+    def __cinit__(self):
+        pass
+
+    cdef void init_mutable(self, const shared_ptr[CMutableBuffer]& buffer):
+        Buffer.init(self, <shared_ptr[CBuffer]>(buffer))
+        self.mutable_buffer = buffer
+
+    def __getbuffer__(self, cp.Py_buffer* buffer, int flags):
+
+        buffer.buf = <char *>self.mutable_buffer.get().mutable_data()
+        buffer.format = 'b'
+        buffer.internal = NULL
+        buffer.itemsize = 1
+        buffer.len = self.size
+        buffer.ndim = 1
+        buffer.obj = self
+        buffer.readonly = 0
+        buffer.shape = self.shape
+        buffer.strides = self.strides
+        buffer.suboffsets = NULL
+
+
 cdef shared_ptr[PoolBuffer] allocate_buffer(CMemoryPool* pool):
     cdef shared_ptr[PoolBuffer] result
     result.reset(new PoolBuffer(pool))

python/pyarrow/lib.pxd

@@ -313,6 +313,13 @@ cdef class Buffer:
     cdef void init(self, const shared_ptr[CBuffer]& buffer)
 
 
+cdef class MutableBuffer(Buffer):
+    cdef:
+        shared_ptr[CMutableBuffer] mutable_buffer
+
+    cdef void init_mutable(self, const shared_ptr[CMutableBuffer]& buffer)
+
+
 cdef class NativeFile:
     cdef:
         shared_ptr[RandomAccessFile] rd_file