Add a MemoryPool abstract interface, change builder instances to request memory from pool via Buffer subclass

Author: wesm

cpp/src/arrow/array-test.cc

@@ -28,6 +28,7 @@
 #include "arrow/types/integer.h"
 #include "arrow/types/primitive.h"
 #include "arrow/util/buffer.h"
+#include "arrow/util/memory-pool.h"
 
 using std::string;
 using std::vector;
@@ -41,8 +42,10 @@ static TypePtr int32_nn = TypePtr(new Int32Type(false));
 class TestArray : public ::testing::Test {
  public:
   void SetUp() {
-    auto data = std::make_shared<OwnedMutableBuffer>();
-    auto nulls = std::make_shared<OwnedMutableBuffer>();
+    pool_ = GetDefaultMemoryPool();
+
+    auto data = std::make_shared<PoolBuffer>(pool_);
+    auto nulls = std::make_shared<PoolBuffer>(pool_);
 
     ASSERT_OK(data->Resize(400));
     ASSERT_OK(nulls->Resize(128));
@@ -51,6 +54,7 @@ class TestArray : public ::testing::Test {
   }
 
  protected:
+  MemoryPool* pool_;
   std::unique_ptr<Int32Array> arr_;
 };
 

cpp/src/arrow/builder.cc

@@ -30,7 +30,7 @@ Status ArrayBuilder::Init(int64_t capacity) {
 
   if (nullable_) {
     int64_t to_alloc = util::ceil_byte(capacity) / 8;
-    nulls_ = std::make_shared<OwnedMutableBuffer>();
+    nulls_ = std::make_shared<PoolBuffer>(pool_);
     RETURN_NOT_OK(nulls_->Resize(to_alloc));
     null_bits_ = nulls_->mutable_data();
     memset(null_bits_, 0, to_alloc);

cpp/src/arrow/builder.h

@@ -30,18 +30,20 @@
 namespace arrow {
 
 class Array;
+class MemoryPool;
 
 static constexpr int64_t MIN_BUILDER_CAPACITY = 1 << 8;
 
 // Base class for all data array builders
 class ArrayBuilder {
  public:
-  explicit ArrayBuilder(const TypePtr& type)
-      : type_(type),
-        nullable_(type_->nullable),
-        nulls_(nullptr), null_bits_(nullptr),
-        length_(0),
-        capacity_(0) {}
+  explicit ArrayBuilder(MemoryPool* pool, const TypePtr& type) :
+      pool_(pool),
+      type_(type),
+      nullable_(type_->nullable),
+      nulls_(nullptr), null_bits_(nullptr),
+      length_(0),
+      capacity_(0) {}
 
   virtual ~ArrayBuilder() {}
 
@@ -71,18 +73,20 @@ class ArrayBuilder {
   // this function responsibly.
   Status Advance(int64_t elements);
 
-  const std::shared_ptr<OwnedMutableBuffer>& nulls() const { return nulls_;}
+  const std::shared_ptr<PoolBuffer>& nulls() const { return nulls_;}
 
   // Creates new array object to hold the contents of the builder and transfers
   // ownership of the data
   virtual Status ToArray(Array** out) = 0;
 
  protected:
+  MemoryPool* pool_;
+
   TypePtr type_;
   bool nullable_;
 
   // If the type is not nullable, then null_ is nullptr after initialization
-  std::shared_ptr<OwnedMutableBuffer> nulls_;
+  std::shared_ptr<PoolBuffer> nulls_;
   uint8_t* null_bits_;
 
   // Array length, so far. Also, the index of the next element to be added

cpp/src/arrow/types/construct.cc

@@ -32,12 +32,13 @@ class ArrayBuilder;
 // Initially looked at doing this with vtables, but shared pointers makes it
 // difficult
 
-#define BUILDER_CASE(ENUM, BuilderType)                         \
-    case TypeEnum::ENUM:                                        \
-      *out = static_cast<ArrayBuilder*>(new BuilderType(type)); \
+#define BUILDER_CASE(ENUM, BuilderType)                                 \
+    case TypeEnum::ENUM:                                                \
+      *out = static_cast<ArrayBuilder*>(new BuilderType(pool, type));   \
       return Status::OK();
 
-Status make_builder(const TypePtr& type, ArrayBuilder** out) {
+Status make_builder(MemoryPool* pool, const TypePtr& type,
+    ArrayBuilder** out) {
   switch (type->type) {
     BUILDER_CASE(UINT8, UInt8Builder);
     BUILDER_CASE(INT8, Int8Builder);
@@ -59,10 +60,10 @@ Status make_builder(const TypePtr& type, ArrayBuilder** out) {
       {
         ListType* list_type = static_cast<ListType*>(type.get());
         ArrayBuilder* value_builder;
-        RETURN_NOT_OK(make_builder(list_type->value_type, &value_builder));
+        RETURN_NOT_OK(make_builder(pool, list_type->value_type, &value_builder));
 
         // The ListBuilder takes ownership of the value_builder
-        ListBuilder* builder = new ListBuilder(type, value_builder);
+        ListBuilder* builder = new ListBuilder(pool, type, value_builder);
         *out = static_cast<ArrayBuilder*>(builder);
         return Status::OK();
       }

cpp/src/arrow/types/construct.h

@@ -23,9 +23,11 @@
 namespace arrow {
 
 class ArrayBuilder;
+class MemoryPool;
 class Status;
 
-Status make_builder(const TypePtr& type, ArrayBuilder** out);
+Status make_builder(MemoryPool* pool, const TypePtr& type,
+    ArrayBuilder** out);
 
 } // namespace arrow
 

cpp/src/arrow/types/list-test.cc

@@ -76,7 +76,7 @@ class TestListBuilder : public TestBuilder {
     type_ = TypePtr(new ListType(value_type_));
 
     ArrayBuilder* tmp;
-    ASSERT_OK(make_builder(type_, &tmp));
+    ASSERT_OK(make_builder(pool_, type_, &tmp));
     builder_.reset(static_cast<ListBuilder*>(tmp));
   }
 

cpp/src/arrow/types/list.h

@@ -34,6 +34,8 @@
 
 namespace arrow {
 
+class MemoryPool;
+
 struct ListType : public DataType {
   // List can contain any other logical value type
   TypePtr value_type;
@@ -100,8 +102,9 @@ class ListArray : public Array {
 // have been appended to the child array)
 class ListBuilder : public Int32Builder {
  public:
-  ListBuilder(const TypePtr& type, ArrayBuilder* value_builder)
-      : Int32Builder(type) {
+  ListBuilder(MemoryPool* pool, const TypePtr& type,
+      ArrayBuilder* value_builder)
+      : Int32Builder(pool, type) {
     value_builder_.reset(value_builder);
   }
 

cpp/src/arrow/types/primitive-test.cc

@@ -104,10 +104,10 @@ class TestPrimitiveBuilder : public TestBuilder {
     type_nn_ = Attrs::type(false);
 
     ArrayBuilder* tmp;
-    ASSERT_OK(make_builder(type_, &tmp));
+    ASSERT_OK(make_builder(pool_, type_, &tmp));
     builder_.reset(static_cast<BuilderType*>(tmp));
 
-    ASSERT_OK(make_builder(type_nn_, &tmp));
+    ASSERT_OK(make_builder(pool_, type_nn_, &tmp));
     builder_nn_.reset(static_cast<BuilderType*>(tmp));
   }
 

cpp/src/arrow/types/primitive.h

@@ -31,6 +31,8 @@
 
 namespace arrow {
 
+class MemoryPool;
+
 template <typename Derived>
 struct PrimitiveType : public DataType {
   explicit PrimitiveType(bool nullable = true)
@@ -113,8 +115,9 @@ class PrimitiveBuilder : public ArrayBuilder {
  public:
   typedef typename Type::c_type T;
 
-  explicit PrimitiveBuilder(const TypePtr& type)
-      : ArrayBuilder(type), values_(nullptr) {
+  explicit PrimitiveBuilder(MemoryPool* pool, const TypePtr& type) :
+      ArrayBuilder(pool, type),
+      values_(nullptr) {
     elsize_ = sizeof(T);
   }
 
@@ -139,7 +142,7 @@ class PrimitiveBuilder : public ArrayBuilder {
   Status Init(int64_t capacity) {
     RETURN_NOT_OK(ArrayBuilder::Init(capacity));
 
-    values_ = std::make_shared<OwnedMutableBuffer>();
+    values_ = std::make_shared<PoolBuffer>(pool_);
     return values_->Resize(capacity * elsize_);
   }
 
@@ -231,7 +234,7 @@ class PrimitiveBuilder : public ArrayBuilder {
   }
 
  protected:
-  std::shared_ptr<OwnedMutableBuffer> values_;
+  std::shared_ptr<PoolBuffer> values_;
   int64_t elsize_;
 };
 

cpp/src/arrow/types/string-test.cc

@@ -175,7 +175,7 @@ class TestStringBuilder : public TestBuilder {
     type_ = TypePtr(new StringType());
 
     ArrayBuilder* tmp;
-    ASSERT_OK(make_builder(type_, &tmp));
+    ASSERT_OK(make_builder(pool_, type_, &tmp));
     builder_.reset(static_cast<StringBuilder*>(tmp));
   }