Permit buffers (write padding) in a non-multiple of 64 in an IPC context, to allow zero-copy writing of NumPy arrays

Change-Id: I4747841f805e08f943e626ec3c77ac4b4fb5a0f7

cpp/src/arrow/ipc/adapter.cc

@@ -162,17 +162,14 @@ class RecordBatchWriter {
     for (size_t i = 0; i < buffers_.size(); ++i) {
       const Buffer* buffer = buffers_[i].get();
       int64_t size = 0;
+      int64_t padding = 0;
 
       // The buffer might be null if we are handling zero row lengths.
       if (buffer) {
-        DCHECK(buffer->capacity() >= buffer->size());
-
-        // We use capacity here, because size might not reflect the padding
-        // requirements of buffers but capacity always should.
-        size = buffer->capacity();
-        // check that padding is appropriate
-        RETURN_NOT_OK(CheckMultipleOf64(size));
+        size = buffer->size();
+        padding = util::RoundUpToMultipleOf64(size) - size;
       }
+
       // TODO(wesm): We currently have no notion of shared memory page id's,
       // but we've included it in the metadata IDL for when we have it in the
       // future. Use page=0 for now
@@ -181,12 +178,17 @@ class RecordBatchWriter {
       // are using from any OS-level shared memory. The thought is that systems
       // may (in the future) associate integer page id's with physical memory
       // pages (according to whatever is the desired shared memory mechanism)
-      buffer_meta_.push_back(flatbuf::Buffer(0, position, size));
+      buffer_meta_.push_back(flatbuf::Buffer(0, position, size + padding));
 
       if (size > 0) {
         RETURN_NOT_OK(dst->Write(buffer->data(), size));
         position += size;
       }
+
+      if (padding > 0) {
+        RETURN_NOT_OK(dst->Write(kPaddingBytes, padding));
+        position += padding;
+      }
     }
 
     *body_end_offset = position;

cpp/src/arrow/ipc/util.h

@@ -29,7 +29,11 @@ namespace ipc {
 
 // Align on 8-byte boundaries
 static constexpr int kArrowAlignment = 8;
-static constexpr uint8_t kPaddingBytes[kArrowAlignment] = {0};
+
+// Buffers are padded to 64-byte boundaries (for SIMD)
+static constexpr int kArrowBufferAlignment = 64;
+
+static constexpr uint8_t kPaddingBytes[kArrowBufferAlignment] = {0};
 
 static inline int64_t PaddedLength(int64_t nbytes, int64_t alignment = kArrowAlignment) {
   return ((nbytes + alignment - 1) / alignment) * alignment;

cpp/src/arrow/util/bit-util.h

@@ -19,6 +19,7 @@
 #define ARROW_UTIL_BIT_UTIL_H
 
 #include <cstdint>
+#include <limits>
 #include <memory>
 #include <vector>
 
@@ -77,6 +78,18 @@ static inline bool is_multiple_of_64(int64_t n) {
   return (n & 63) == 0;
 }
 
+inline int64_t RoundUpToMultipleOf64(int64_t num) {
+  // TODO(wesm): is this definitely needed?
+  // DCHECK_GE(num, 0);
+  constexpr int64_t round_to = 64;
+  constexpr int64_t force_carry_addend = round_to - 1;
+  constexpr int64_t truncate_bitmask = ~(round_to - 1);
+  constexpr int64_t max_roundable_num = std::numeric_limits<int64_t>::max() - round_to;
+  if (num <= max_roundable_num) { return (num + force_carry_addend) & truncate_bitmask; }
+  // handle overflow case.  This should result in a malloc error upstream
+  return num;
+}
+
 void bytes_to_bits(const std::vector<uint8_t>& bytes, uint8_t* bits);
 ARROW_EXPORT Status bytes_to_bits(const std::vector<uint8_t>&, std::shared_ptr<Buffer>*);
 

cpp/src/arrow/util/buffer.cc

@@ -20,25 +20,13 @@
 #include <cstdint>
 #include <limits>
 
+#include "arrow/util/bit-util.h"
 #include "arrow/util/logging.h"
 #include "arrow/util/memory-pool.h"
 #include "arrow/util/status.h"
 
 namespace arrow {
 
-namespace {
-int64_t RoundUpToMultipleOf64(int64_t num) {
-  DCHECK_GE(num, 0);
-  constexpr int64_t round_to = 64;
-  constexpr int64_t force_carry_addend = round_to - 1;
-  constexpr int64_t truncate_bitmask = ~(round_to - 1);
-  constexpr int64_t max_roundable_num = std::numeric_limits<int64_t>::max() - round_to;
-  if (num <= max_roundable_num) { return (num + force_carry_addend) & truncate_bitmask; }
-  // handle overflow case.  This should result in a malloc error upstream
-  return num;
-}
-}  // namespace
-
 Buffer::Buffer(const std::shared_ptr<Buffer>& parent, int64_t offset, int64_t size) {
   data_ = parent->data() + offset;
   size_ = size;
@@ -64,7 +52,7 @@ PoolBuffer::~PoolBuffer() {
 Status PoolBuffer::Reserve(int64_t new_capacity) {
   if (!mutable_data_ || new_capacity > capacity_) {
     uint8_t* new_data;
-    new_capacity = RoundUpToMultipleOf64(new_capacity);
+    new_capacity = util::RoundUpToMultipleOf64(new_capacity);
     if (mutable_data_) {
       RETURN_NOT_OK(pool_->Allocate(new_capacity, &new_data));
       memcpy(new_data, mutable_data_, size_);

cpp/src/arrow/util/buffer.h

@@ -23,7 +23,6 @@
 #include <cstring>
 #include <memory>
 
-#include "arrow/util/bit-util.h"
 #include "arrow/util/macros.h"
 #include "arrow/util/status.h"
 #include "arrow/util/visibility.h"

python/pyarrow/array.pyx

@@ -243,12 +243,14 @@ def from_pandas_series(object series, object mask=None, timestamps_to_ms=False):
         series_values = series_values.astype('datetime64[ms]')
 
     if mask is None:
-        check_status(pyarrow.PandasToArrow(pyarrow.GetMemoryPool(),
-                                           series_values, &out))
+        with nogil:
+            check_status(pyarrow.PandasToArrow(pyarrow.GetMemoryPool(),
+                                               series_values, &out))
     else:
         mask = series_as_ndarray(mask)
-        check_status(pyarrow.PandasMaskedToArrow(
-            pyarrow.GetMemoryPool(), series_values, mask, &out))
+        with nogil:
+            check_status(pyarrow.PandasMaskedToArrow(
+                pyarrow.GetMemoryPool(), series_values, mask, &out))
 
     return box_arrow_array(out)
 

python/pyarrow/tests/test_ipc.py

@@ -26,6 +26,8 @@ import pyarrow.ipc as ipc
 
 
 def test_ipc_simple_roundtrip():
+    # Also tests writing zero-copy NumPy array with additional padding
+
     nrows = 5
     df = pd.DataFrame({
         'one': np.random.randn(nrows),
@@ -55,9 +57,11 @@ def test_ipc_simple_roundtrip():
 
     for i in range(num_batches):
         # it works. Must convert back to DataFrame
-        reader.get_record_batch(i)
+        batch = reader.get_record_batch(i)
 
 
+# XXX: For benchmarking
+
 def big_batch():
     df = pd.DataFrame(
         np.random.randn(2**4, 2**20).T,
@@ -85,6 +89,3 @@ def read_file(source):
     reader = ipc.ArrowFileReader(source)
     return [reader.get_record_batch(i)
             for i in range(reader.num_record_batches)]
-
-batch = big_batch()
-mem = write_to_memory(batch)

python/src/pyarrow/adapters/pandas.cc

@@ -602,6 +602,8 @@ class ArrowDeserializer {
   }
 
   Status AllocateOutput(int type) {
+    PyAcquireGIL lock;
+
     npy_intp dims[1] = {col_->length()};
     out_ = reinterpret_cast<PyArrayObject*>(PyArray_SimpleNew(1, dims, type));
 
@@ -616,6 +618,8 @@ class ArrowDeserializer {
   }
 
   Status OutputFromData(int type, void* data) {
+    PyAcquireGIL lock;
+
     // Zero-Copy. We can pass the data pointer directly to NumPy.
     Py_INCREF(py_ref_);
     OwnedRef py_ref(py_ref_);
@@ -706,6 +710,8 @@ class ArrowDeserializer {
   inline typename std::enable_if<
     arrow_traits<T2>::is_boolean, Status>::type
   ConvertValues(const std::shared_ptr<Array>& arr) {
+    PyAcquireGIL lock;
+
     arrow::BooleanArray* bool_arr = static_cast<arrow::BooleanArray*>(arr.get());
 
     if (arr->null_count() > 0) {
@@ -743,6 +749,8 @@ class ArrowDeserializer {
   inline typename std::enable_if<
     T2 == arrow::Type::STRING, Status>::type
   ConvertValues(const std::shared_ptr<Array>& arr) {
+    PyAcquireGIL lock;
+
     RETURN_NOT_OK(AllocateOutput(NPY_OBJECT));
 
     PyObject** out_values = reinterpret_cast<PyObject**>(PyArray_DATA(out_));

python/src/pyarrow/common.h

@@ -120,7 +120,7 @@ class PYARROW_EXPORT NumPyBuffer : public arrow::Buffer {
 
     data_ = reinterpret_cast<const uint8_t*>(PyArray_DATA(arr_));
     size_ = PyArray_SIZE(arr_);
-    capacity_ = size_;
+    capacity_ = size_ * PyArray_DESCR(arr_)->elsize;
   }
 
   virtual ~NumPyBuffer() {
@@ -140,6 +140,22 @@ class PYARROW_EXPORT PyBytesBuffer : public arrow::Buffer {
   PyObject* obj_;
 };
 
+
+class PyAcquireGIL {
+ public:
+  PyAcquireGIL() {
+    state_ = PyGILState_Ensure();
+  }
+
+  ~PyAcquireGIL() {
+    PyGILState_Release(state_);
+  }
+
+ private:
+  PyGILState_STATE state_;
+  DISALLOW_COPY_AND_ASSIGN(PyAcquireGIL);
+};
+
 } // namespace pyarrow
 
 #endif // PYARROW_COMMON_H