Use parallel memcopy from arrow (#633)

* use parallel memcopy from arrow

* fix linting

* remove memory.h

src/numbuf/python/src/pynumbuf/memory.h

@@ -1,187 +0,0 @@
-#ifndef PYNUMBUF_MEMORY_H
-#define PYNUMBUF_MEMORY_H
-
-#include <arrow/io/interfaces.h>
-
-/* C++ includes */
-#include <string>
-#include <thread>
-#include <vector>
-
-#define THREADPOOL_SIZE 8
-#define MEMCOPY_BLOCK_SIZE 64
-#define BYTES_IN_MB (1 << 20)
-
-using namespace std;
-
-namespace numbuf {
-
-class ParallelMemcopy {
- public:
-  explicit ParallelMemcopy(uint64_t block_size, bool timeit, int threadpool_size)
-      : timeit_(timeit),
-        threadpool_(threadpool_size),
-        block_size_(block_size),
-        threadpool_size_(threadpool_size) {}
-
-  void memcopy(uint8_t* dst, const uint8_t* src, uint64_t nbytes) {
-    struct timeval tv1, tv2;
-    if (timeit_) {
-      // Start the timer.
-      gettimeofday(&tv1, NULL);
-    }
-    if (nbytes >= BYTES_IN_MB) {
-      memcopy_aligned(dst, src, nbytes, block_size_);
-    } else {
-      memcpy(dst, src, nbytes);
-    }
-    if (timeit_) {
-      // Stop the timer and log the measured time.
-      gettimeofday(&tv2, NULL);
-      double elapsed =
-          ((tv2.tv_sec - tv1.tv_sec) * 1000000 + (tv2.tv_usec - tv1.tv_usec)) / 1000000.0;
-      // TODO: replace this with ARROW_LOG(ARROW_INFO) or better equivalent.
-      printf("Copied %llu bytes in time = %8.4f MBps = %8.4f\n", nbytes, elapsed,
-          nbytes / (BYTES_IN_MB * elapsed));
-    }
-  }
-
-  ~ParallelMemcopy() {
-    // Join threadpool threads just in case they are still running.
-    for (auto& t : threadpool_) {
-      if (t.joinable()) { t.join(); }
-    }
-  }
-
- private:
-  /** Controls whether the memcopy operations are timed. */
-  bool timeit_;
-  /** Specifies the desired alignment in bytes, as a power of 2. */
-  uint64_t block_size_;
-  /** Number of threads to be used for parallel memcopy operations. */
-  int threadpool_size_;
-  /** Internal threadpool to be used in the fork/join pattern. */
-  std::vector<std::thread> threadpool_;
-
-  void memcopy_aligned(
-      uint8_t* dst, const uint8_t* src, uint64_t nbytes, uint64_t block_size) {
-    uint64_t num_threads = threadpool_size_;
-    uint64_t dst_address = reinterpret_cast<uint64_t>(dst);
-    uint64_t src_address = reinterpret_cast<uint64_t>(src);
-    uint64_t left_address = (src_address + block_size - 1) & ~(block_size - 1);
-    uint64_t right_address = (src_address + nbytes) & ~(block_size - 1);
-    uint64_t num_blocks = (right_address - left_address) / block_size;
-    // Update right address
-    right_address = right_address - (num_blocks % num_threads) * block_size;
-    // Now we divide these blocks between available threads. The remainder is
-    // handled on the main thread.
-
-    uint64_t chunk_size = (right_address - left_address) / num_threads;
-    uint64_t prefix = left_address - src_address;
-    uint64_t suffix = src_address + nbytes - right_address;
-    // Now the data layout is | prefix | k * num_threads * block_size | suffix |.
-    // We have chunk_size = k * block_size, therefore the data layout is
-    // | prefix | num_threads * chunk_size | suffix |.
-    // Each thread gets a "chunk" of k blocks.
-
-    // Start all threads first and handle leftovers while threads run.
-    for (int i = 0; i < num_threads; i++) {
-      threadpool_[i] = std::thread(memcpy, dst + prefix + i * chunk_size,
-          reinterpret_cast<uint8_t*>(left_address) + i * chunk_size, chunk_size);
-    }
-
-    memcpy(dst, src, prefix);
-    memcpy(dst + prefix + num_threads * chunk_size,
-        reinterpret_cast<uint8_t*>(right_address), suffix);
-
-    for (auto& t : threadpool_) {
-      if (t.joinable()) { t.join(); }
-    }
-  }
-};
-
-class FixedBufferStream : public arrow::io::OutputStream,
-                          public arrow::io::ReadableFileInterface {
- public:
-  virtual ~FixedBufferStream() {}
-
-  explicit FixedBufferStream(uint8_t* data, int64_t nbytes)
-      : data_(data),
-        position_(0),
-        size_(nbytes),
-        memcopy_helper(MEMCOPY_BLOCK_SIZE, false, THREADPOOL_SIZE) {}
-
-  arrow::Status Read(int64_t nbytes, std::shared_ptr<arrow::Buffer>* out) override {
-    DCHECK(out);
-    if (position_ + nbytes > size_) { return arrow::Status::IOError("EOF"); }
-    *out = std::make_shared<arrow::Buffer>(data_ + position_, nbytes);
-    position_ += nbytes;
-    return arrow::Status::OK();
-  }
-
-  arrow::Status Read(int64_t nbytes, int64_t* bytes_read, uint8_t* out) override {
-    assert(0);
-    return arrow::Status::OK();
-  }
-
-  arrow::Status Seek(int64_t position) override {
-    position_ = position;
-    return arrow::Status::OK();
-  }
-
-  arrow::Status Close() override { return arrow::Status::OK(); }
-
-  arrow::Status Tell(int64_t* position) override {
-    *position = position_;
-    return arrow::Status::OK();
-  }
-
-  arrow::Status Write(const uint8_t* data, int64_t nbytes) override {
-    DCHECK(position_ >= 0 && position_ < size_);
-    DCHECK(position_ + nbytes <= size_) << "position: " << position_
-                                        << " nbytes: " << nbytes << "size: " << size_;
-    uint8_t* dst = data_ + position_;
-    memcopy_helper.memcopy(dst, data, nbytes);
-    position_ += nbytes;
-    return arrow::Status::OK();
-  }
-
-  arrow::Status GetSize(int64_t* size) override {
-    *size = size_;
-    return arrow::Status::OK();
-  }
-
-  bool supports_zero_copy() const override { return true; }
-
- private:
-  uint8_t* data_;
-  int64_t position_;
-  int64_t size_;
-  ParallelMemcopy memcopy_helper;
-};
-
-class MockBufferStream : public arrow::io::OutputStream {
- public:
-  virtual ~MockBufferStream() {}
-
-  explicit MockBufferStream() : position_(0) {}
-
-  arrow::Status Close() override { return arrow::Status::OK(); }
-
-  arrow::Status Tell(int64_t* position) override {
-    *position = position_;
-    return arrow::Status::OK();
-  }
-
-  arrow::Status Write(const uint8_t* data, int64_t nbytes) override {
-    position_ += nbytes;
-    return arrow::Status::OK();
-  }
-
- private:
-  int64_t position_;
-};
-
-}  // namespace numbuf
-
-#endif  // PYNUMBUF_MEMORY_H

src/numbuf/python/src/pynumbuf/numbuf.cc

@@ -27,12 +27,13 @@ PyObject* NumbufPlasmaObjectExistsError;
 #endif
 
 #include <arrow/api.h>
+#include <arrow/io/memory.h>
 #include <arrow/ipc/api.h>
+#include <arrow/ipc/util.h>
 #include <arrow/ipc/writer.h>
 #include <arrow/python/numpy_convert.h>
 
 #include "adapters/python.h"
-#include "memory.h"
 
 using namespace arrow;
 using namespace numbuf;
@@ -81,7 +82,7 @@ Status read_batch_and_tensors(uint8_t* data, int64_t size,
     std::vector<std::shared_ptr<Tensor>>& tensors_out) {
   std::shared_ptr<arrow::ipc::FileReader> reader;
   int64_t batch_size = *((int64_t*)data);
-  auto source = std::make_shared<FixedBufferStream>(
+  auto source = std::make_shared<arrow::io::BufferReader>(
       LENGTH_PREFIX_SIZE + data, size - LENGTH_PREFIX_SIZE);
   RETURN_NOT_OK(arrow::ipc::FileReader::Open(source, batch_size, &reader));
   RETURN_NOT_OK(reader->GetRecordBatch(0, batch_out));
@@ -151,7 +152,7 @@ static PyObject* serialize_list(PyObject* self, PyObject* args) {
     object->batch = make_batch(array);
 
     int64_t data_size, total_size;
-    auto mock = std::make_shared<MockBufferStream>();
+    auto mock = std::make_shared<arrow::ipc::MockOutputStream>();
     write_batch_and_tensors(
         mock.get(), object->batch, object->arrays, &data_size, &total_size);
 
@@ -173,9 +174,10 @@ static PyObject* write_to_buffer(PyObject* self, PyObject* args) {
   }
   if (!PyMemoryView_Check(memoryview)) { return NULL; }
   Py_buffer* buffer = PyMemoryView_GET_BUFFER(memoryview);
-  auto target = std::make_shared<FixedBufferStream>(
+  auto buf = std::make_shared<arrow::MutableBuffer>(
       LENGTH_PREFIX_SIZE + reinterpret_cast<uint8_t*>(buffer->buf),
       buffer->len - LENGTH_PREFIX_SIZE);
+  auto target = std::make_shared<arrow::io::FixedSizeBufferWriter>(buf);
   int64_t batch_size, total_size;
   ARROW_CHECK_OK(write_batch_and_tensors(
       target.get(), object->batch, object->arrays, &batch_size, &total_size));
@@ -297,7 +299,7 @@ static PyObject* store_list(PyObject* self, PyObject* args) {
   std::shared_ptr<RecordBatch> batch = make_batch(array);
 
   int64_t data_size, total_size;
-  auto mock = std::make_shared<MockBufferStream>();
+  auto mock = std::make_shared<arrow::ipc::MockOutputStream>();
   write_batch_and_tensors(mock.get(), batch, tensors, &data_size, &total_size);
 
   uint8_t* data;
@@ -321,8 +323,9 @@ static PyObject* store_list(PyObject* self, PyObject* args) {
   }
   ARROW_CHECK_OK(s);
 
-  auto target =
-      std::make_shared<FixedBufferStream>(LENGTH_PREFIX_SIZE + data, total_size);
+  auto buf =
+      std::make_shared<arrow::MutableBuffer>(LENGTH_PREFIX_SIZE + data, total_size);
+  auto target = std::make_shared<arrow::io::FixedSizeBufferWriter>(buf);
   write_batch_and_tensors(target.get(), batch, tensors, &data_size, &total_size);
   *((int64_t*)data) = data_size;
 

src/numbuf/thirdparty/download_thirdparty.sh

@@ -13,4 +13,4 @@ fi
 cd $TP_DIR/arrow
 git pull origin master
 
-git checkout 670612e6fdf699486641ed0d39d22257eb8acdb2
+git checkout 8a700ccdad745c250fe5d91a9104e7c2d6364c1b