diff --git a/cpp/src/arrow/buffer-test.cc b/cpp/src/arrow/buffer-test.cc index 55b86e96ea4..368a9cb0138 100644 --- a/cpp/src/arrow/buffer-test.cc +++ b/cpp/src/arrow/buffer-test.cc @@ -32,19 +32,6 @@ using std::string; namespace arrow { -TEST(TestBuffer, IsMutableFlag) { - Buffer buf(nullptr, 0); - - ASSERT_FALSE(buf.is_mutable()); - - MutableBuffer mbuf(nullptr, 0); - ASSERT_TRUE(mbuf.is_mutable()); - - std::shared_ptr pool_buf; - ASSERT_OK(AllocateResizableBuffer(0, &pool_buf)); - ASSERT_TRUE(pool_buf->is_mutable()); -} - TEST(TestBuffer, FromStdString) { std::string val = "hello, world"; @@ -70,62 +57,6 @@ TEST(TestBuffer, FromStdStringWithMemory) { ASSERT_EQ(static_cast(expected.size()), buf->size()); } -TEST(TestBuffer, Resize) { - std::shared_ptr buf; - ASSERT_OK(AllocateResizableBuffer(0, &buf)); - - ASSERT_EQ(0, buf->size()); - ASSERT_OK(buf->Resize(100)); - ASSERT_EQ(100, buf->size()); - ASSERT_OK(buf->Resize(200)); - ASSERT_EQ(200, buf->size()); - - // Make it smaller, too - ASSERT_OK(buf->Resize(50, true)); - ASSERT_EQ(50, buf->size()); - // We have actually shrunken in size - // The spec requires that capacity is a multiple of 64 - ASSERT_EQ(64, buf->capacity()); - - // Resize to a larger capacity again to test shrink_to_fit = false - ASSERT_OK(buf->Resize(100)); - ASSERT_EQ(128, buf->capacity()); - ASSERT_OK(buf->Resize(50, false)); - ASSERT_EQ(128, buf->capacity()); -} - -TEST(TestBuffer, TypedResize) { - std::shared_ptr buf; - ASSERT_OK(AllocateResizableBuffer(0, &buf)); - - ASSERT_EQ(0, buf->size()); - ASSERT_OK(buf->TypedResize(100)); - ASSERT_EQ(800, buf->size()); - ASSERT_OK(buf->TypedResize(200)); - ASSERT_EQ(1600, buf->size()); - - ASSERT_OK(buf->TypedResize(50, true)); - ASSERT_EQ(400, buf->size()); - ASSERT_EQ(448, buf->capacity()); - - ASSERT_OK(buf->TypedResize(100)); - ASSERT_EQ(832, buf->capacity()); - ASSERT_OK(buf->TypedResize(50, false)); - ASSERT_EQ(832, buf->capacity()); -} - -TEST(TestBuffer, ResizeOOM) { -// This test doesn't play nice with AddressSanitizer -#ifndef ADDRESS_SANITIZER - // realloc fails, even though there may be no explicit limit - std::shared_ptr buf; - ASSERT_OK(AllocateResizableBuffer(0, &buf)); - ASSERT_OK(buf->Resize(100)); - int64_t to_alloc = std::numeric_limits::max(); - ASSERT_RAISES(OutOfMemory, buf->Resize(to_alloc)); -#endif -} - TEST(TestBuffer, EqualsWithSameContent) { MemoryPool* pool = default_memory_pool(); const int32_t bufferSize = 128 * 1024; @@ -243,4 +174,81 @@ TEST(TestBufferBuilder, ResizeReserve) { ASSERT_EQ(128, builder.capacity()); } +template +class TypedTestBuffer : public ::testing::Test {}; + +using BufferPtrs = + ::testing::Types, std::unique_ptr>; + +TYPED_TEST_CASE(TypedTestBuffer, BufferPtrs); + +TYPED_TEST(TypedTestBuffer, IsMutableFlag) { + Buffer buf(nullptr, 0); + + ASSERT_FALSE(buf.is_mutable()); + + MutableBuffer mbuf(nullptr, 0); + ASSERT_TRUE(mbuf.is_mutable()); + + TypeParam pool_buf; + ASSERT_OK(AllocateResizableBuffer(0, &pool_buf)); + ASSERT_TRUE(pool_buf->is_mutable()); +} + +TYPED_TEST(TypedTestBuffer, Resize) { + TypeParam buf; + ASSERT_OK(AllocateResizableBuffer(0, &buf)); + + ASSERT_EQ(0, buf->size()); + ASSERT_OK(buf->Resize(100)); + ASSERT_EQ(100, buf->size()); + ASSERT_OK(buf->Resize(200)); + ASSERT_EQ(200, buf->size()); + + // Make it smaller, too + ASSERT_OK(buf->Resize(50, true)); + ASSERT_EQ(50, buf->size()); + // We have actually shrunken in size + // The spec requires that capacity is a multiple of 64 + ASSERT_EQ(64, buf->capacity()); + + // Resize to a larger capacity again to test shrink_to_fit = false + ASSERT_OK(buf->Resize(100)); + ASSERT_EQ(128, buf->capacity()); + ASSERT_OK(buf->Resize(50, false)); + ASSERT_EQ(128, buf->capacity()); +} + +TYPED_TEST(TypedTestBuffer, TypedResize) { + TypeParam buf; + ASSERT_OK(AllocateResizableBuffer(0, &buf)); + + ASSERT_EQ(0, buf->size()); + ASSERT_OK(buf->template TypedResize(100)); + ASSERT_EQ(800, buf->size()); + ASSERT_OK(buf->template TypedResize(200)); + ASSERT_EQ(1600, buf->size()); + + ASSERT_OK(buf->template TypedResize(50, true)); + ASSERT_EQ(400, buf->size()); + ASSERT_EQ(448, buf->capacity()); + + ASSERT_OK(buf->template TypedResize(100)); + ASSERT_EQ(832, buf->capacity()); + ASSERT_OK(buf->template TypedResize(50, false)); + ASSERT_EQ(832, buf->capacity()); +} + +TYPED_TEST(TypedTestBuffer, ResizeOOM) { +// This test doesn't play nice with AddressSanitizer +#ifndef ADDRESS_SANITIZER + // realloc fails, even though there may be no explicit limit + TypeParam buf; + ASSERT_OK(AllocateResizableBuffer(0, &buf)); + ASSERT_OK(buf->Resize(100)); + int64_t to_alloc = std::numeric_limits::max(); + ASSERT_RAISES(OutOfMemory, buf->Resize(to_alloc)); +#endif +} + } // namespace arrow diff --git a/cpp/src/arrow/buffer.cc b/cpp/src/arrow/buffer.cc index 91b0002fcc2..388e98f56ca 100644 --- a/cpp/src/arrow/buffer.cc +++ b/cpp/src/arrow/buffer.cc @@ -18,6 +18,7 @@ #include "arrow/buffer.h" #include +#include #include "arrow/memory_pool.h" #include "arrow/status.h" @@ -146,26 +147,46 @@ MutableBuffer::MutableBuffer(const std::shared_ptr& parent, const int64_ parent_ = parent; } -Status AllocateBuffer(MemoryPool* pool, const int64_t size, - std::shared_ptr* out) { - auto buffer = std::make_shared(pool); +namespace { +// A utility that does most of the work of the `AllocateBuffer` and +// `AllocateResizableBuffer` methods. The argument `buffer` should be a smart pointer to a +// PoolBuffer. +template +inline Status ResizePoolBuffer(PoolBufferPtr&& buffer, const int64_t size, + BufferPtr* out) { RETURN_NOT_OK(buffer->Resize(size)); buffer->ZeroPadding(); - *out = buffer; + *out = std::move(buffer); return Status::OK(); } +} // namespace + +Status AllocateBuffer(MemoryPool* pool, const int64_t size, + std::shared_ptr* out) { + return ResizePoolBuffer(std::make_shared(pool), size, out); +} + +Status AllocateBuffer(MemoryPool* pool, const int64_t size, + std::unique_ptr* out) { + return ResizePoolBuffer(std::unique_ptr(new PoolBuffer(pool)), size, out); +} Status AllocateBuffer(const int64_t size, std::shared_ptr* out) { return AllocateBuffer(default_memory_pool(), size, out); } +Status AllocateBuffer(const int64_t size, std::unique_ptr* out) { + return AllocateBuffer(default_memory_pool(), size, out); +} + Status AllocateResizableBuffer(MemoryPool* pool, const int64_t size, std::shared_ptr* out) { - auto buffer = std::make_shared(pool); - RETURN_NOT_OK(buffer->Resize(size)); - buffer->ZeroPadding(); - *out = buffer; - return Status::OK(); + return ResizePoolBuffer(std::make_shared(pool), size, out); +} + +Status AllocateResizableBuffer(MemoryPool* pool, const int64_t size, + std::unique_ptr* out) { + return ResizePoolBuffer(std::unique_ptr(new PoolBuffer(pool)), size, out); } Status AllocateResizableBuffer(const int64_t size, @@ -173,6 +194,11 @@ Status AllocateResizableBuffer(const int64_t size, return AllocateResizableBuffer(default_memory_pool(), size, out); } +Status AllocateResizableBuffer(const int64_t size, + std::unique_ptr* out) { + return AllocateResizableBuffer(default_memory_pool(), size, out); +} + Status AllocateEmptyBitmap(MemoryPool* pool, int64_t length, std::shared_ptr* out) { RETURN_NOT_OK(AllocateBuffer(pool, BitUtil::BytesForBits(length), out)); diff --git a/cpp/src/arrow/buffer.h b/cpp/src/arrow/buffer.h index 486f04633b9..442c4516985 100644 --- a/cpp/src/arrow/buffer.h +++ b/cpp/src/arrow/buffer.h @@ -231,6 +231,16 @@ class ARROW_EXPORT ResizableBuffer : public MutableBuffer { ARROW_EXPORT Status AllocateBuffer(MemoryPool* pool, const int64_t size, std::shared_ptr* out); +/// \brief Allocate a fixed size mutable buffer from a memory pool, zero its padding. +/// +/// \param[in] pool a memory pool +/// \param[in] size size of buffer to allocate +/// \param[out] out the allocated buffer (contains padding) +/// +/// \return Status message +ARROW_EXPORT +Status AllocateBuffer(MemoryPool* pool, const int64_t size, std::unique_ptr* out); + /// \brief Allocate a fixed-size mutable buffer from the default memory pool /// /// \param[in] size size of buffer to allocate @@ -240,6 +250,15 @@ Status AllocateBuffer(MemoryPool* pool, const int64_t size, std::shared_ptr* out); +/// \brief Allocate a fixed-size mutable buffer from the default memory pool +/// +/// \param[in] size size of buffer to allocate +/// \param[out] out the allocated buffer (contains padding) +/// +/// \return Status message +ARROW_EXPORT +Status AllocateBuffer(const int64_t size, std::unique_ptr* out); + /// \brief Allocate a resizeable buffer from a memory pool, zero its padding. /// /// \param[in] pool a memory pool @@ -251,6 +270,17 @@ ARROW_EXPORT Status AllocateResizableBuffer(MemoryPool* pool, const int64_t size, std::shared_ptr* out); +/// \brief Allocate a resizeable buffer from a memory pool, zero its padding. +/// +/// \param[in] pool a memory pool +/// \param[in] size size of buffer to allocate +/// \param[out] out the allocated buffer +/// +/// \return Status message +ARROW_EXPORT +Status AllocateResizableBuffer(MemoryPool* pool, const int64_t size, + std::unique_ptr* out); + /// \brief Allocate a resizeable buffer from the default memory pool /// /// \param[in] size size of buffer to allocate @@ -260,6 +290,15 @@ Status AllocateResizableBuffer(MemoryPool* pool, const int64_t size, ARROW_EXPORT Status AllocateResizableBuffer(const int64_t size, std::shared_ptr* out); +/// \brief Allocate a resizeable buffer from the default memory pool +/// +/// \param[in] size size of buffer to allocate +/// \param[out] out the allocated buffer +/// +/// \return Status message +ARROW_EXPORT +Status AllocateResizableBuffer(const int64_t size, std::unique_ptr* out); + /// \brief Allocate a zero-initialized bitmap buffer from a memory pool /// /// \param[in] pool memory pool to allocate memory from