Add Slice, Copy methods to Buffer

Change-Id: Ia6f232bf020a0249c19453a2c39eea51b52f86cf

Author: wesm

cpp/src/arrow/util/buffer-test.cc

@@ -31,6 +31,18 @@ namespace arrow {
 
 class TestBuffer : public ::testing::Test {};
 
+TEST_F(TestBuffer, IsMutableFlag) {
+  Buffer buf(nullptr, 0);
+
+  ASSERT_FALSE(buf.is_mutable());
+
+  MutableBuffer mbuf(nullptr, 0);
+  ASSERT_TRUE(mbuf.is_mutable());
+
+  PoolBuffer pbuf;
+  ASSERT_TRUE(pbuf.is_mutable());
+}
+
 TEST_F(TestBuffer, Resize) {
   PoolBuffer buf;
 
@@ -96,4 +108,33 @@ TEST_F(TestBuffer, EqualsWithSameBuffer) {
   pool->Free(rawBuffer, bufferSize);
 }
 
+TEST_F(TestBuffer, Copy) {
+  std::string data_str = "some data to copy";
+
+  auto data = reinterpret_cast<const uint8_t*>(data_str.c_str());
+
+  Buffer buf(data, data_str.size());
+
+  std::shared_ptr<Buffer> out;
+
+  ASSERT_OK(buf.Copy(5, 4, &out));
+
+  Buffer expected(data + 5, 4);
+  ASSERT_TRUE(out->Equals(expected));
+}
+
+TEST_F(TestBuffer, SliceBuffer) {
+  std::string data_str = "some data to slice";
+
+  auto data = reinterpret_cast<const uint8_t*>(data_str.c_str());
+
+  auto buf = std::make_shared<Buffer>(data, data_str.size());
+
+  std::shared_ptr<Buffer> out = SliceBuffer(buf, 5, 4);
+  Buffer expected(data + 5, 4);
+  ASSERT_TRUE(out->Equals(expected));
+
+  ASSERT_EQ(2, buf.use_count());
+}
+
 }  // namespace arrow

cpp/src/arrow/util/buffer.cc

@@ -36,6 +36,32 @@ Buffer::Buffer(const std::shared_ptr<Buffer>& parent, int64_t offset, int64_t si
 
 Buffer::~Buffer() {}
 
+Status Buffer::Copy(
+    int64_t start, int64_t nbytes, MemoryPool* pool, std::shared_ptr<Buffer>* out) const {
+  // Sanity checks
+  DCHECK_LT(start, size_);
+  DCHECK_LE(nbytes, size_ - start);
+
+  auto new_buffer = std::make_shared<PoolBuffer>(pool);
+  RETURN_NOT_OK(new_buffer->Resize(nbytes));
+
+  std::memcpy(new_buffer->mutable_data(), data() + start, nbytes);
+
+  *out = new_buffer;
+  return Status::OK();
+}
+
+Status Buffer::Copy(int64_t start, int64_t nbytes, std::shared_ptr<Buffer>* out) const {
+  return Copy(start, nbytes, default_memory_pool(), out);
+}
+
+std::shared_ptr<Buffer> SliceBuffer(
+    const std::shared_ptr<Buffer>& buffer, int64_t offset, int64_t length) {
+  DCHECK_LT(offset, buffer->size());
+  DCHECK_LE(length, buffer->size() - offset);
+  return std::make_shared<Buffer>(buffer, offset, length);
+}
+
 std::shared_ptr<Buffer> MutableBuffer::GetImmutableView() {
   return std::make_shared<Buffer>(this->get_shared_ptr(), 0, size());
 }

cpp/src/arrow/util/buffer.h

@@ -43,7 +43,8 @@ class Status;
 // The following invariant is always true: Size < Capacity
 class ARROW_EXPORT Buffer : public std::enable_shared_from_this<Buffer> {
  public:
-  Buffer(const uint8_t* data, int64_t size) : data_(data), size_(size), capacity_(size) {}
+  Buffer(const uint8_t* data, int64_t size)
+      : is_mutable_(false), data_(data), size_(size), capacity_(size) {}
   virtual ~Buffer();
 
   // An offset into data that is owned by another buffer, but we want to be
@@ -57,6 +58,8 @@ class ARROW_EXPORT Buffer : public std::enable_shared_from_this<Buffer> {
 
   std::shared_ptr<Buffer> get_shared_ptr() { return shared_from_this(); }
 
+  bool is_mutable() const { return is_mutable_; }
+
   // Return true if both buffers are the same size and contain the same bytes
   // up to the number of compared bytes
   bool Equals(const Buffer& other, int64_t nbytes) const {
@@ -71,18 +74,22 @@ class ARROW_EXPORT Buffer : public std::enable_shared_from_this<Buffer> {
                (data_ == other.data_ || !memcmp(data_, other.data_, size_)));
   }
 
+  // Copy section of buffer into a new Buffer
+  Status Copy(int64_t start, int64_t nbytes, MemoryPool* pool,
+      std::shared_ptr<Buffer>* out) const;
+
+  // Default memory pool
+  Status Copy(int64_t start, int64_t nbytes, std::shared_ptr<Buffer>* out) const;
+
   int64_t capacity() const { return capacity_; }
   const uint8_t* data() const { return data_; }
 
   int64_t size() const { return size_; }
 
-  // Returns true if this Buffer is referencing memory (possibly) owned by some
-  // other buffer
-  bool is_shared() const { return static_cast<bool>(parent_); }
-
   const std::shared_ptr<Buffer> parent() const { return parent_; }
 
  protected:
+  bool is_mutable_;
   const uint8_t* data_;
   int64_t size_;
   int64_t capacity_;
@@ -94,10 +101,16 @@ class ARROW_EXPORT Buffer : public std::enable_shared_from_this<Buffer> {
   DISALLOW_COPY_AND_ASSIGN(Buffer);
 };
 
+// Construct a view on passed buffer at the indicated offset and length. This
+// function cannot fail and does not error checking (except in debug builds)
+std::shared_ptr<Buffer> SliceBuffer(
+    const std::shared_ptr<Buffer>& buffer, int64_t offset, int64_t length);
+
 // A Buffer whose contents can be mutated. May or may not own its data.
 class ARROW_EXPORT MutableBuffer : public Buffer {
  public:
   MutableBuffer(uint8_t* data, int64_t size) : Buffer(data, size) {
+    is_mutable_ = true;
     mutable_data_ = data;
   }