From 58e48752c0b3a8e71489b706794cf0e847cb7002 Mon Sep 17 00:00:00 2001 From: Wes McKinney Date: Mon, 17 Oct 2016 10:33:58 -0400 Subject: [PATCH] Add Slice, Copy methods to Buffer Change-Id: Ia6f232bf020a0249c19453a2c39eea51b52f86cf --- cpp/src/arrow/util/buffer-test.cc | 41 +++++++++++++++++++++++++++++++ cpp/src/arrow/util/buffer.cc | 26 ++++++++++++++++++++ cpp/src/arrow/util/buffer.h | 23 +++++++++++++---- 3 files changed, 85 insertions(+), 5 deletions(-) diff --git a/cpp/src/arrow/util/buffer-test.cc b/cpp/src/arrow/util/buffer-test.cc index cc4ec98e4fb29..095b07b7ab309 100644 --- a/cpp/src/arrow/util/buffer-test.cc +++ b/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(data_str.c_str()); + + Buffer buf(data, data_str.size()); + + std::shared_ptr 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(data_str.c_str()); + + auto buf = std::make_shared(data, data_str.size()); + + std::shared_ptr out = SliceBuffer(buf, 5, 4); + Buffer expected(data + 5, 4); + ASSERT_TRUE(out->Equals(expected)); + + ASSERT_EQ(2, buf.use_count()); +} + } // namespace arrow diff --git a/cpp/src/arrow/util/buffer.cc b/cpp/src/arrow/util/buffer.cc index 1266c4f5930c0..a230259e5930d 100644 --- a/cpp/src/arrow/util/buffer.cc +++ b/cpp/src/arrow/util/buffer.cc @@ -36,6 +36,32 @@ Buffer::Buffer(const std::shared_ptr& parent, int64_t offset, int64_t si Buffer::~Buffer() {} +Status Buffer::Copy( + int64_t start, int64_t nbytes, MemoryPool* pool, std::shared_ptr* out) const { + // Sanity checks + DCHECK_LT(start, size_); + DCHECK_LE(nbytes, size_ - start); + + auto new_buffer = std::make_shared(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* out) const { + return Copy(start, nbytes, default_memory_pool(), out); +} + +std::shared_ptr SliceBuffer( + const std::shared_ptr& buffer, int64_t offset, int64_t length) { + DCHECK_LT(offset, buffer->size()); + DCHECK_LE(length, buffer->size() - offset); + return std::make_shared(buffer, offset, length); +} + std::shared_ptr MutableBuffer::GetImmutableView() { return std::make_shared(this->get_shared_ptr(), 0, size()); } diff --git a/cpp/src/arrow/util/buffer.h b/cpp/src/arrow/util/buffer.h index bc0df86221c45..04ad6c2dffde4 100644 --- a/cpp/src/arrow/util/buffer.h +++ b/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 { 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 { std::shared_ptr 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 { (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* out) const; + + // Default memory pool + Status Copy(int64_t start, int64_t nbytes, std::shared_ptr* 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(parent_); } - const std::shared_ptr 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 { 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 SliceBuffer( + const std::shared_ptr& 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; }