ARROW-2998: [C++] Add unique_ptr versions of Allocate[Resizable]Buffer

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<ResizableBuffer> 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<int64_t>(expected.size()), buf->size());
 }
 
-TEST(TestBuffer, Resize) {
-  std::shared_ptr<ResizableBuffer> 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<ResizableBuffer> buf;
-  ASSERT_OK(AllocateResizableBuffer(0, &buf));
-
-  ASSERT_EQ(0, buf->size());
-  ASSERT_OK(buf->TypedResize<double>(100));
-  ASSERT_EQ(800, buf->size());
-  ASSERT_OK(buf->TypedResize<double>(200));
-  ASSERT_EQ(1600, buf->size());
-
-  ASSERT_OK(buf->TypedResize<double>(50, true));
-  ASSERT_EQ(400, buf->size());
-  ASSERT_EQ(448, buf->capacity());
-
-  ASSERT_OK(buf->TypedResize<double>(100));
-  ASSERT_EQ(832, buf->capacity());
-  ASSERT_OK(buf->TypedResize<double>(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<ResizableBuffer> buf;
-  ASSERT_OK(AllocateResizableBuffer(0, &buf));
-  ASSERT_OK(buf->Resize(100));
-  int64_t to_alloc = std::numeric_limits<int64_t>::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 <typename T>
+class TypedTestBuffer : public ::testing::Test {};
+
+using BufferPtrs =
+    ::testing::Types<std::shared_ptr<ResizableBuffer>, std::unique_ptr<ResizableBuffer>>;
+
+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<double>(100));
+  ASSERT_EQ(800, buf->size());
+  ASSERT_OK(buf->template TypedResize<double>(200));
+  ASSERT_EQ(1600, buf->size());
+
+  ASSERT_OK(buf->template TypedResize<double>(50, true));
+  ASSERT_EQ(400, buf->size());
+  ASSERT_EQ(448, buf->capacity());
+
+  ASSERT_OK(buf->template TypedResize<double>(100));
+  ASSERT_EQ(832, buf->capacity());
+  ASSERT_OK(buf->template TypedResize<double>(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<int64_t>::max();
+  ASSERT_RAISES(OutOfMemory, buf->Resize(to_alloc));
+#endif
+}
+
 }  // namespace arrow

cpp/src/arrow/buffer.cc

@@ -18,6 +18,7 @@
 #include "arrow/buffer.h"
 
 #include <cstdint>
+#include <utility>
 
 #include "arrow/memory_pool.h"
 #include "arrow/status.h"
@@ -146,33 +147,58 @@ MutableBuffer::MutableBuffer(const std::shared_ptr<Buffer>& parent, const int64_
   parent_ = parent;
 }
 
-Status AllocateBuffer(MemoryPool* pool, const int64_t size,
-                      std::shared_ptr<Buffer>* out) {
-  auto buffer = std::make_shared<PoolBuffer>(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 <typename PoolBufferPtr, typename BufferPtr>
+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<Buffer>* out) {
+  return ResizePoolBuffer(std::make_shared<PoolBuffer>(pool), size, out);
+}
+
+Status AllocateBuffer(MemoryPool* pool, const int64_t size,
+                      std::unique_ptr<Buffer>* out) {
+  return ResizePoolBuffer(std::unique_ptr<PoolBuffer>(new PoolBuffer(pool)), size, out);
+}
 
 Status AllocateBuffer(const int64_t size, std::shared_ptr<Buffer>* out) {
   return AllocateBuffer(default_memory_pool(), size, out);
 }
 
+Status AllocateBuffer(const int64_t size, std::unique_ptr<Buffer>* out) {
+  return AllocateBuffer(default_memory_pool(), size, out);
+}
+
 Status AllocateResizableBuffer(MemoryPool* pool, const int64_t size,
                                std::shared_ptr<ResizableBuffer>* out) {
-  auto buffer = std::make_shared<PoolBuffer>(pool);
-  RETURN_NOT_OK(buffer->Resize(size));
-  buffer->ZeroPadding();
-  *out = buffer;
-  return Status::OK();
+  return ResizePoolBuffer(std::make_shared<PoolBuffer>(pool), size, out);
+}
+
+Status AllocateResizableBuffer(MemoryPool* pool, const int64_t size,
+                               std::unique_ptr<ResizableBuffer>* out) {
+  return ResizePoolBuffer(std::unique_ptr<PoolBuffer>(new PoolBuffer(pool)), size, out);
 }
 
 Status AllocateResizableBuffer(const int64_t size,
                                std::shared_ptr<ResizableBuffer>* out) {
   return AllocateResizableBuffer(default_memory_pool(), size, out);
 }
 
+Status AllocateResizableBuffer(const int64_t size,
+                               std::unique_ptr<ResizableBuffer>* out) {
+  return AllocateResizableBuffer(default_memory_pool(), size, out);
+}
+
 Status AllocateEmptyBitmap(MemoryPool* pool, int64_t length,
                            std::shared_ptr<Buffer>* out) {
   RETURN_NOT_OK(AllocateBuffer(pool, BitUtil::BytesForBits(length), out));

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<Buffer>* 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<Buffer>* 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<Buff
 ARROW_EXPORT
 Status AllocateBuffer(const int64_t size, std::shared_ptr<Buffer>* 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<Buffer>* 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<ResizableBuffer>* 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<ResizableBuffer>* 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<ResizableBuffer>* 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<ResizableBuffer>* out);
+
 /// \brief Allocate a zero-initialized bitmap buffer from a memory pool
 ///
 /// \param[in] pool memory pool to allocate memory from