Add / test Buffer wrapper. Test that we can write an arrow file to a wrapped buffer. Resize buffer in BufferOutputStream on close

Change-Id: I94d2689172554a55259370df99a3147e9047d421

cpp/src/arrow/io/io-memory-test.cc

@@ -121,5 +121,30 @@ TEST_F(TestMemoryMappedFile, InvalidFile) {
       IOError, MemoryMappedFile::Open(non_existent_path, FileMode::READ, &result));
 }
 
+class TestBufferOutputStream : public ::testing::Test {
+ public:
+  void SetUp() {
+    buffer_.reset(new PoolBuffer(default_memory_pool()));
+    stream_.reset(new BufferOutputStream(buffer_));
+  }
+
+ protected:
+  std::shared_ptr<PoolBuffer> buffer_;
+  std::unique_ptr<OutputStream> stream_;
+};
+
+TEST_F(TestBufferOutputStream, CloseResizes) {
+  std::string data = "data123456";
+
+  const int64_t nbytes = static_cast<int64_t>(data.size());
+  const int K = 100;
+  for (int i = 0; i < K; ++i) {
+    EXPECT_OK(stream_->Write(reinterpret_cast<const uint8_t*>(data.c_str()), nbytes));
+  }
+
+  ASSERT_OK(stream_->Close());
+  ASSERT_EQ(K * nbytes, buffer_->size());
+}
+
 }  // namespace io
 }  // namespace arrow

cpp/src/arrow/io/memory.cc

@@ -212,7 +212,11 @@ BufferOutputStream::BufferOutputStream(const std::shared_ptr<ResizableBuffer>& b
       mutable_data_(buffer->mutable_data()) {}
 
 Status BufferOutputStream::Close() {
-  return Status::OK();
+  if (position_ < capacity_) {
+    return buffer_->Resize(position_);
+  } else {
+    return Status::OK();
+  }
 }
 
 Status BufferOutputStream::Tell(int64_t* position) {
@@ -228,8 +232,11 @@ Status BufferOutputStream::Write(const uint8_t* data, int64_t nbytes) {
 }
 
 Status BufferOutputStream::Reserve(int64_t nbytes) {
-  while (position_ + nbytes > capacity_) {
-    int64_t new_capacity = std::max(kBufferMinimumSize, capacity_ * 2);
+  int64_t new_capacity = capacity_;
+  while (position_ + nbytes > new_capacity) {
+    new_capacity = std::max(kBufferMinimumSize, new_capacity * 2);
+  }
+  if (new_capacity > capacity_) {
     RETURN_NOT_OK(buffer_->Resize(new_capacity));
     capacity_ = new_capacity;
   }

python/pyarrow/array.pyx

@@ -37,7 +37,7 @@ import pyarrow.schema as schema
 
 
 def total_allocated_bytes():
-    cdef MemoryPool* pool = pyarrow.GetMemoryPool()
+    cdef MemoryPool* pool = pyarrow.get_memory_pool()
 
 
 cdef class Array:
@@ -243,13 +243,13 @@ def from_pandas_series(object series, object mask=None, timestamps_to_ms=False):
 
     if mask is None:
         with nogil:
-            check_status(pyarrow.PandasToArrow(pyarrow.GetMemoryPool(),
+            check_status(pyarrow.PandasToArrow(pyarrow.get_memory_pool(),
                                                series_values, &out))
     else:
         mask = series_as_ndarray(mask)
         with nogil:
             check_status(pyarrow.PandasMaskedToArrow(
-                pyarrow.GetMemoryPool(), series_values, mask, &out))
+                pyarrow.get_memory_pool(), series_values, mask, &out))
 
     return box_arrow_array(out)
 

python/pyarrow/includes/libarrow.pxd

@@ -58,6 +58,14 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
         uint8_t* data()
         int64_t size()
 
+    cdef cppclass ResizableBuffer(CBuffer):
+        CStatus Resize(int64_t nbytes)
+        CStatus Reserve(int64_t nbytes)
+
+    cdef cppclass PoolBuffer(ResizableBuffer):
+        PoolBuffer()
+        PoolBuffer(MemoryPool*)
+
     cdef MemoryPool* default_memory_pool()
 
     cdef cppclass CListType" arrow::ListType"(CDataType):

python/pyarrow/includes/libarrow_io.pxd

@@ -18,7 +18,7 @@
 # distutils: language = c++
 
 from pyarrow.includes.common cimport *
-from pyarrow.includes.libarrow cimport MemoryPool, CBuffer
+from pyarrow.includes.libarrow cimport *
 
 cdef extern from "arrow/io/interfaces.h" namespace "arrow::io" nogil:
     enum FileMode" arrow::io::FileMode::type":
@@ -148,5 +148,4 @@ cdef extern from "arrow/io/memory.h" namespace "arrow::io" nogil:
         CBufferReader(const uint8_t* data, int64_t nbytes)
 
     cdef cppclass BufferOutputStream(OutputStream):
-        # TODO(wesm)
-        pass
+        BufferOutputStream(const shared_ptr[ResizableBuffer]& buffer)

python/pyarrow/includes/pyarrow.pxd

@@ -18,8 +18,8 @@
 # distutils: language = c++
 
 from pyarrow.includes.common cimport *
-from pyarrow.includes.libarrow cimport (CArray, CColumn, CDataType, CStatus,
-                                        Type, MemoryPool)
+from pyarrow.includes.libarrow cimport (CArray, CBuffer, CColumn,
+                                        CDataType, CStatus, Type, MemoryPool)
 
 cimport pyarrow.includes.libarrow_io as arrow_io
 
@@ -53,7 +53,12 @@ cdef extern from "pyarrow/api.h" namespace "pyarrow" nogil:
     PyStatus ArrowToPandas(const shared_ptr[CColumn]& arr, object py_ref,
                            PyObject** out)
 
-    MemoryPool* GetMemoryPool()
+    MemoryPool* get_memory_pool()
+
+
+cdef extern from "pyarrow/common.h" namespace "pyarrow" nogil:
+    cdef cppclass PyBytesBuffer(CBuffer):
+        PyBytesBuffer(object o)
 
 
 cdef extern from "pyarrow/io.h" namespace "pyarrow" nogil:

python/pyarrow/io.pyx

@@ -74,10 +74,16 @@ cdef class NativeFile:
         if not self.is_readonly:
             raise IOError("only valid on readonly files")
 
+        if not self.is_open:
+            raise IOError("file not open")
+
     def _assert_writeable(self):
         if self.is_readonly:
             raise IOError("only valid on writeonly files")
 
+        if not self.is_open:
+            raise IOError("file not open")
+
     def size(self):
         cdef int64_t size
         self._assert_readable()
@@ -169,6 +175,7 @@ cdef class BytesReader(NativeFile):
 # ----------------------------------------------------------------------
 # Arrow buffers
 
+
 cdef class Buffer:
 
     def __cinit__(self):
@@ -177,16 +184,65 @@ cdef class Buffer:
     cdef init(self, const shared_ptr[CBuffer]& buffer):
         self.buffer = buffer
 
+    def __len__(self):
+        return self.size
+
+    property size:
+
+        def __get__(self):
+            return self.buffer.get().size()
+
+    def __getitem__(self, key):
+        # TODO(wesm): buffer slicing
+        raise NotImplementedError
+
     def to_pybytes(self):
         return cp.PyBytes_FromStringAndSize(
             <const char*>self.buffer.get().data(),
             self.buffer.get().size())
 
 
+cdef shared_ptr[PoolBuffer] allocate_buffer():
+    cdef shared_ptr[PoolBuffer] result
+    result.reset(new PoolBuffer(pyarrow.get_memory_pool()))
+    return result
+
+
 cdef class InMemoryOutputStream(NativeFile):
 
+    cdef:
+        shared_ptr[PoolBuffer] buffer
+
+    def __cinit__(self):
+        self.buffer = allocate_buffer()
+        self.wr_file.reset(new BufferOutputStream(
+            <shared_ptr[ResizableBuffer]> self.buffer))
+        self.is_readonly = 0
+        self.is_open = True
+
     def get_result(self):
-        pass
+        cdef Buffer result = Buffer()
+
+        check_cstatus(self.wr_file.get().Close())
+        result.init(<shared_ptr[CBuffer]> self.buffer)
+
+        self.is_open = False
+        return result
+
+
+def buffer_from_bytes(object obj):
+    """
+    Construct an Arrow buffer from a Python bytes object
+    """
+    if not isinstance(obj, bytes):
+        raise ValueError('Must pass bytes object')
+
+    cdef shared_ptr[CBuffer] buf
+    buf.reset(new pyarrow.PyBytesBuffer(obj))
+
+    cdef Buffer result = Buffer()
+    result.init(buf)
+    return result
 
 # ----------------------------------------------------------------------
 # HDFS IO implementation

python/pyarrow/tests/test_io.py

@@ -98,3 +98,44 @@ def test_bytes_reader():
 def test_bytes_reader_non_bytes():
     with pytest.raises(ValueError):
         io.BytesReader(u('some sample data'))
+
+
+
+# ----------------------------------------------------------------------
+# Buffers
+
+
+def test_buffer_bytes():
+    val = b'some data'
+
+    buf = io.buffer_from_bytes(val)
+    assert isinstance(buf, io.Buffer)
+
+    result = buf.to_pybytes()
+
+    assert result == val
+
+
+def test_memory_output_stream():
+    # 10 bytes
+    val = b'dataabcdef'
+
+    f = io.InMemoryOutputStream()
+
+    K = 1000
+    for i in range(K):
+        f.write(val)
+
+    buf = f.get_result()
+
+    assert len(buf) == len(val) * K
+    assert buf.to_pybytes() == val * K
+
+
+def test_inmemory_write_after_closed():
+    f = io.InMemoryOutputStream()
+    f.write(b'ok')
+    f.get_result()
+
+    with pytest.raises(IOError):
+        f.write(b'not ok')

python/pyarrow/tests/test_ipc.py

@@ -21,46 +21,68 @@ import numpy as np
 import pandas as pd
 
 import pyarrow as A
+import pyarrow.io as arrow_io
 import pyarrow.ipc as ipc
 
 
-def test_ipc_file_simple_roundtrip():
+class RoundtripTest(object):
     # Also tests writing zero-copy NumPy array with additional padding
 
-    nrows = 5
-    df = pd.DataFrame({
-        'one': np.random.randn(nrows),
-        'two': ['foo', np.nan, 'bar', 'bazbaz', 'qux']})
+    def __init__(self):
+        self.sink = self._get_sink()
 
-    batch = A.RecordBatch.from_pandas(df)
+    def _get_sink(self):
+        return io.BytesIO()
 
-    sink = io.BytesIO()
-    writer = ipc.ArrowFileWriter(sink, batch.schema)
+    def _get_source(self):
+        return self.sink.getvalue()
 
-    num_batches = 5
-    frames = []
-    batches = []
-    for i in range(num_batches):
-        unique_df = df.copy()
-        unique_df['one'] = np.random.randn(nrows)
+    def run(self):
+        nrows = 5
+        df = pd.DataFrame({
+            'one': np.random.randn(nrows),
+            'two': ['foo', np.nan, 'bar', 'bazbaz', 'qux']})
 
-        batch = A.RecordBatch.from_pandas(unique_df)
-        writer.write_record_batch(batch)
-        frames.append(unique_df)
-        batches.append(batch)
+        batch = A.RecordBatch.from_pandas(df)
+        writer = ipc.ArrowFileWriter(self.sink, batch.schema)
 
-    writer.close()
+        num_batches = 5
+        frames = []
+        batches = []
+        for i in range(num_batches):
+            unique_df = df.copy()
+            unique_df['one'] = np.random.randn(nrows)
 
-    file_contents = sink.getvalue()
-    reader = ipc.ArrowFileReader(file_contents)
+            batch = A.RecordBatch.from_pandas(unique_df)
+            writer.write_record_batch(batch)
+            frames.append(unique_df)
+            batches.append(batch)
 
-    assert reader.num_record_batches == num_batches
+        writer.close()
 
-    for i in range(num_batches):
-        # it works. Must convert back to DataFrame
-        batch = reader.get_record_batch(i)
+        file_contents = self._get_source()
+        reader = ipc.ArrowFileReader(file_contents)
 
-        assert batches[i].equals(batch)
+        assert reader.num_record_batches == num_batches
+
+        for i in range(num_batches):
+            # it works. Must convert back to DataFrame
+            batch = reader.get_record_batch(i)
+            assert batches[i].equals(batch)
+
+
+class InMemoryStreamTest(RoundtripTest):
+
+    def _get_sink(self):
+        return arrow_io.InMemoryOutputStream()
+
+    def _get_source(self):
+        return self.sink.get_result()
+
+
+def test_ipc_file_simple_roundtrip():
+    helper = RoundtripTest()
+    helper.run()
 
 
 # XXX: For benchmarking

python/src/pyarrow/adapters/builtin.cc

@@ -426,7 +426,7 @@ Status ConvertPySequence(PyObject* obj, std::shared_ptr<arrow::Array>* out) {
 
   // Give the sequence converter an array builder
   std::shared_ptr<ArrayBuilder> builder;
-  RETURN_ARROW_NOT_OK(arrow::MakeBuilder(GetMemoryPool(), type, &builder));
+  RETURN_ARROW_NOT_OK(arrow::MakeBuilder(get_memory_pool(), type, &builder));
   converter->Init(builder);
 
   PY_RETURN_NOT_OK(converter->AppendData(obj));

python/src/pyarrow/common.cc

@@ -63,7 +63,7 @@ class PyArrowMemoryPool : public arrow::MemoryPool {
   int64_t bytes_allocated_;
 };
 
-arrow::MemoryPool* GetMemoryPool() {
+arrow::MemoryPool* get_memory_pool() {
   static PyArrowMemoryPool memory_pool;
   return &memory_pool;
 }

python/src/pyarrow/common.h

@@ -109,7 +109,8 @@ class PyGILGuard {
     return Status::UnknownError(message);           \
   }
 
-PYARROW_EXPORT arrow::MemoryPool* GetMemoryPool();
+// Return the common PyArrow memory pool
+PYARROW_EXPORT arrow::MemoryPool* get_memory_pool();
 
 class PYARROW_EXPORT NumPyBuffer : public arrow::Buffer {
  public: