Add public pyarrow.allocate_buffer API. Rename FixedSizeBufferOutputStream to FixedSizeBufferWriter for when we can support seeking at some point

Change-Id: I4fa0d8e4c6a5b6a07de993585837e29084cb8414

python/doc/source/api.rst

@@ -158,6 +158,7 @@ Input / Output and Shared Memory
 .. autosummary::
    :toctree: generated/
 
+   allocate_buffer
    Buffer
    BufferReader
    BufferOutputStream

python/doc/source/plasma.rst

@@ -235,14 +235,14 @@ API such as ``pyarrow.get_tensor_size``.
   buf = client.create(object_id, data_size)
 
 To write the Arrow ``Tensor`` object into the buffer, you can use Plasma to
-convert the ``memoryview`` buffer into a ``pyarrow.FixedSizeBufferOutputStream``
-object. A ``pyarrow.FixedSizeBufferOutputStream`` is a format suitable for Arrow's
+convert the ``memoryview`` buffer into a ``pyarrow.FixedSizeBufferWriter``
+object. A ``pyarrow.FixedSizeBufferWriter`` is a format suitable for Arrow's
 ``pyarrow.write_tensor``:
 
 .. code-block:: python
 
   # Write the tensor into the Plasma-allocated buffer
-  stream = pa.FixedSizeBufferOutputStream(buf)
+  stream = pa.FixedSizeBufferWriter(buf)
   pa.write_tensor(tensor, stream)  # Writes tensor's 552 bytes to Plasma stream
 
 To finish storing the Arrow object in Plasma, call ``seal``:
@@ -328,7 +328,7 @@ The DataFrame can now be written to the buffer as follows.
 .. code-block:: python
 
   # Write the PyArrow RecordBatch to Plasma
-  stream = pa.FixedSizeBufferOutputStream(buf)
+  stream = pa.FixedSizeBufferWriter(buf)
   stream_writer = pa.RecordBatchStreamWriter(stream, record_batch.schema)
   stream_writer.write_batch(record_batch)
   stream_writer.close()
@@ -414,7 +414,7 @@ You can test this with the following script:
   object_id = plasma.ObjectID(np.random.bytes(20))
   buf = client.create(object_id, pa.get_tensor_size(tensor))
 
-  stream = pa.FixedSizeBufferOutputStream(buf)
+  stream = pa.FixedSizeBufferWriter(buf)
   stream.set_memcopy_threads(4)
   a = time.time()
   pa.write_tensor(tensor, stream)

python/examples/plasma/sorting/sort_df.py

@@ -39,7 +39,7 @@ import multimerge
 #     num_cols = 1
 
 client = None
-object_store_size = 2 * 10 ** 9 # 2 GB
+object_store_size = 2 * 10 ** 9  # 2 GB
 num_cores = 8
 num_rows = 200000
 num_cols = 2
@@ -69,7 +69,7 @@ def put_df(df):
     buf = client.create(object_id, data_size)
 
     # Write the serialized DataFrame to the object store
-    sink = pa.FixedSizeBufferOutputStream(buf)
+    sink = pa.FixedSizeBufferWriter(buf)
     stream_writer = pa.RecordBatchStreamWriter(sink, record_batch.schema)
     stream_writer.write_batch(record_batch)
 

python/pyarrow/__init__.py

@@ -68,10 +68,10 @@ from pyarrow.lib import (null, bool_,
                          Date32Value, Date64Value, TimestampValue)
 
 from pyarrow.lib import (HdfsFile, NativeFile, PythonFile,
-                         FixedSizeBufferOutputStream,
+                         FixedSizeBufferWriter,
                          Buffer, BufferReader, BufferOutputStream,
                          OSFile, MemoryMappedFile, memory_map,
-                         frombuffer,
+                         allocate_buffer, frombuffer,
                          memory_map, create_memory_map,
                          have_libhdfs, have_libhdfs3, MockOutputStream)
 
@@ -134,6 +134,11 @@ def _plasma_store_entry_point():
 
 from pyarrow.util import _deprecate_class
 
+FixedSizeBufferOutputStream = (
+    _deprecate_class('FixedSizeBufferOutputStream',
+                     'FixedSizeBufferWriter',
+                     FixedSizeBufferWriter, '0.7.0'))
+
 # Backwards compatibility with pyarrow < 0.6.0
 HdfsClient = _deprecate_class('HdfsClient', 'pyarrow.hdfs.connect',
                               hdfs.connect, '0.6.0')

python/pyarrow/includes/libarrow.pxd

@@ -160,6 +160,12 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
         CMutableBuffer(const uint8_t* data, int64_t size)
         uint8_t* mutable_data()
 
+    CStatus AllocateBuffer(CMemoryPool* pool, const int64_t size,
+                           shared_ptr[CBuffer]* out)
+
+    CStatus AllocateResizableBuffer(CMemoryPool* pool, const int64_t size,
+                                    shared_ptr[ResizableBuffer]* out)
+
     cdef cppclass ResizableBuffer(CBuffer):
         CStatus Resize(int64_t nbytes)
         CStatus Reserve(int64_t nbytes)

python/pyarrow/io.pxi

@@ -150,7 +150,7 @@ cdef class NativeFile:
         new_position : the new absolute stream position
         """
         cdef int64_t offset
-        self._assert_readable()
+
         with nogil:
             if whence == 0:
                 offset = position
@@ -534,7 +534,7 @@ cdef class OSFile(NativeFile):
         self.wr_file = <shared_ptr[OutputStream]> handle
 
 
-cdef class FixedSizeBufferOutputStream(NativeFile):
+cdef class FixedSizeBufferWriter(NativeFile):
 
     def __cinit__(self, Buffer buffer):
         self.wr_file.reset(new CFixedSizeBufferWriter(buffer.buffer))
@@ -580,6 +580,11 @@ cdef class Buffer:
         def __get__(self):
             return self.buffer.get().size()
 
+    property is_mutable:
+
+        def __get__(self):
+            return self.buffer.get().is_mutable()
+
     property parent:
 
         def __get__(self):
@@ -638,19 +643,40 @@ cdef class Buffer:
         return self.size
 
 
-cdef shared_ptr[PoolBuffer] allocate_buffer(CMemoryPool* pool):
+cdef shared_ptr[PoolBuffer] _allocate_buffer(CMemoryPool* pool):
     cdef shared_ptr[PoolBuffer] result
     result.reset(new PoolBuffer(pool))
     return result
 
 
+def allocate_buffer(int64_t size, MemoryPool pool=None):
+    """
+    Allocate mutable fixed-size buffer
+
+    Parameters
+    ----------
+    size : int
+        Number of bytes to allocate (plus internal padding)
+    pool : MemoryPool, optional
+        Uses default memory pool if not provided
+    """
+    cdef:
+        shared_ptr[CBuffer] buffer
+        CMemoryPool* cpool = maybe_unbox_memory_pool(pool)
+
+    with nogil:
+        check_status(AllocateBuffer(cpool, size, &buffer))
+
+    return pyarrow_wrap_buffer(buffer)
+
+
 cdef class BufferOutputStream(NativeFile):
 
     cdef:
         shared_ptr[PoolBuffer] buffer
 
     def __cinit__(self, MemoryPool memory_pool=None):
-        self.buffer = allocate_buffer(maybe_unbox_memory_pool(memory_pool))
+        self.buffer = _allocate_buffer(maybe_unbox_memory_pool(memory_pool))
         self.wr_file.reset(new CBufferOutputStream(
             <shared_ptr[ResizableBuffer]> self.buffer))
         self.is_readable = 0

python/pyarrow/plasma.pyx

@@ -390,7 +390,7 @@ cdef class PlasmaClient:
                                    else ObjectID.from_random())
         serialized = pyarrow.serialize(value)
         buffer = self.create(target_id, serialized.total_bytes)
-        stream = pyarrow.FixedSizeBufferOutputStream(buffer)
+        stream = pyarrow.FixedSizeBufferWriter(buffer)
         stream.set_memcopy_threads(4)
         serialized.write_to(stream)
         self.seal(target_id)

python/pyarrow/tests/test_io.py

@@ -170,6 +170,18 @@ def test_buffer_numpy():
     assert array.base == buf
 
 
+def test_allocate_buffer():
+    buf = pa.allocate_buffer(100)
+    assert buf.size == 100
+    assert buf.is_mutable
+
+    bit = b'abcde'
+    writer = pa.FixedSizeBufferWriter(buf)
+    writer.write(bit)
+
+    assert buf.to_pybytes()[:5] == bit
+
+
 def test_buffer_memoryview_is_immutable():
     val = b'some data'
 

python/pyarrow/tests/test_plasma.py

@@ -294,7 +294,7 @@ class TestPlasmaClient(object):
         tensor = pa.Tensor.from_numpy(data)
         data_size = pa.get_tensor_size(tensor)
         buf = self.plasma_client.create(object_id, data_size)
-        stream = pa.FixedSizeBufferOutputStream(buf)
+        stream = pa.FixedSizeBufferWriter(buf)
         pa.write_tensor(tensor, stream)
         self.plasma_client.seal(object_id)
         # Read the arrow object.
@@ -320,7 +320,7 @@ class TestPlasmaClient(object):
         object_id = plasma.ObjectID(np.random.bytes(20))
 
         buf = self.plasma_client.create(object_id, data_size)
-        stream = pa.FixedSizeBufferOutputStream(buf)
+        stream = pa.FixedSizeBufferWriter(buf)
         stream_writer = pa.RecordBatchStreamWriter(stream, record_batch.schema)
         stream_writer.write_batch(record_batch)