diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 5ddd9dae3fe..d2c840abfe8 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -434,7 +434,7 @@ if (UNIX) add_custom_target(lint ${BUILD_SUPPORT_DIR}/cpplint.py --verbose=2 --linelength=90 - --filter=-whitespace/comments,-readability/todo,-build/header_guard + --filter=-whitespace/comments,-readability/todo,-build/header_guard,-build/c++11 `find ${CMAKE_CURRENT_SOURCE_DIR}/src -name \\*.cc -or -name \\*.h`) endif (UNIX) diff --git a/cpp/src/arrow/array-test.cc b/cpp/src/arrow/array-test.cc index 5ecf91624fe..16afb9bef34 100644 --- a/cpp/src/arrow/array-test.cc +++ b/cpp/src/arrow/array-test.cc @@ -18,6 +18,7 @@ #include #include +#include #include #include #include @@ -28,6 +29,8 @@ #include "arrow/types/integer.h" #include "arrow/types/primitive.h" #include "arrow/util/buffer.h" +#include "arrow/util/memory-pool.h" +#include "arrow/util/status.h" using std::string; using std::vector; @@ -41,8 +44,10 @@ static TypePtr int32_nn = TypePtr(new Int32Type(false)); class TestArray : public ::testing::Test { public: void SetUp() { - auto data = std::make_shared(); - auto nulls = std::make_shared(); + pool_ = GetDefaultMemoryPool(); + + auto data = std::make_shared(pool_); + auto nulls = std::make_shared(pool_); ASSERT_OK(data->Resize(400)); ASSERT_OK(nulls->Resize(128)); @@ -51,6 +56,7 @@ class TestArray : public ::testing::Test { } protected: + MemoryPool* pool_; std::unique_ptr arr_; }; diff --git a/cpp/src/arrow/array.h b/cpp/src/arrow/array.h index c95450d12a4..0eaa28d528e 100644 --- a/cpp/src/arrow/array.h +++ b/cpp/src/arrow/array.h @@ -19,7 +19,6 @@ #define ARROW_ARRAY_H #include -#include #include #include "arrow/type.h" diff --git a/cpp/src/arrow/builder.cc b/cpp/src/arrow/builder.cc index 1fd74719283..cb850673150 100644 --- a/cpp/src/arrow/builder.cc +++ b/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(); + nulls_ = std::make_shared(pool_); RETURN_NOT_OK(nulls_->Resize(to_alloc)); null_bits_ = nulls_->mutable_data(); memset(null_bits_, 0, to_alloc); diff --git a/cpp/src/arrow/builder.h b/cpp/src/arrow/builder.h index b43668af77c..456bb04ae09 100644 --- a/cpp/src/arrow/builder.h +++ b/cpp/src/arrow/builder.h @@ -23,25 +23,27 @@ #include #include "arrow/type.h" -#include "arrow/util/buffer.h" #include "arrow/util/macros.h" #include "arrow/util/status.h" namespace arrow { class Array; +class MemoryPool; +class PoolBuffer; 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& nulls() const { return nulls_;} + const std::shared_ptr& 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 nulls_; + std::shared_ptr nulls_; uint8_t* null_bits_; // Array length, so far. Also, the index of the next element to be added diff --git a/cpp/src/arrow/types/construct.cc b/cpp/src/arrow/types/construct.cc index 5176cafd3ba..e1bb990063c 100644 --- a/cpp/src/arrow/types/construct.cc +++ b/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(new BuilderType(type)); \ +#define BUILDER_CASE(ENUM, BuilderType) \ + case TypeEnum::ENUM: \ + *out = static_cast(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(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(builder); return Status::OK(); } diff --git a/cpp/src/arrow/types/construct.h b/cpp/src/arrow/types/construct.h index c0bfedd27d6..b5ba436f787 100644 --- a/cpp/src/arrow/types/construct.h +++ b/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 diff --git a/cpp/src/arrow/types/list-test.cc b/cpp/src/arrow/types/list-test.cc index 47673ff898b..abfc8a31b0d 100644 --- a/cpp/src/arrow/types/list-test.cc +++ b/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(tmp)); } diff --git a/cpp/src/arrow/types/list.h b/cpp/src/arrow/types/list.h index 0f1116257c5..4ca0f13d53c 100644 --- a/cpp/src/arrow/types/list.h +++ b/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); } diff --git a/cpp/src/arrow/types/primitive-test.cc b/cpp/src/arrow/types/primitive-test.cc index 12968608094..3484294a39f 100644 --- a/cpp/src/arrow/types/primitive-test.cc +++ b/cpp/src/arrow/types/primitive-test.cc @@ -18,7 +18,6 @@ #include #include -#include #include #include #include @@ -104,10 +103,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(tmp)); - ASSERT_OK(make_builder(type_nn_, &tmp)); + ASSERT_OK(make_builder(pool_, type_nn_, &tmp)); builder_nn_.reset(static_cast(tmp)); } diff --git a/cpp/src/arrow/types/primitive.h b/cpp/src/arrow/types/primitive.h index a41911224e0..c5ae0f78a99 100644 --- a/cpp/src/arrow/types/primitive.h +++ b/cpp/src/arrow/types/primitive.h @@ -20,6 +20,7 @@ #include #include +#include #include #include "arrow/array.h" @@ -31,6 +32,8 @@ namespace arrow { +class MemoryPool; + template struct PrimitiveType : public DataType { explicit PrimitiveType(bool nullable = true) @@ -113,8 +116,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 +143,7 @@ class PrimitiveBuilder : public ArrayBuilder { Status Init(int64_t capacity) { RETURN_NOT_OK(ArrayBuilder::Init(capacity)); - values_ = std::make_shared(); + values_ = std::make_shared(pool_); return values_->Resize(capacity * elsize_); } @@ -231,7 +235,7 @@ class PrimitiveBuilder : public ArrayBuilder { } protected: - std::shared_ptr values_; + std::shared_ptr values_; int64_t elsize_; }; diff --git a/cpp/src/arrow/types/string-test.cc b/cpp/src/arrow/types/string-test.cc index 6dba3fdcbb6..a2d87ead59c 100644 --- a/cpp/src/arrow/types/string-test.cc +++ b/cpp/src/arrow/types/string-test.cc @@ -31,12 +31,9 @@ #include "arrow/types/test-common.h" #include "arrow/util/status.h" -using std::string; -using std::unique_ptr; -using std::vector; - namespace arrow { +class Buffer; TEST(TypesTest, TestCharType) { CharType t1(5); @@ -45,7 +42,7 @@ TEST(TypesTest, TestCharType) { ASSERT_TRUE(t1.nullable); ASSERT_EQ(t1.size, 5); - ASSERT_EQ(t1.ToString(), string("char(5)")); + ASSERT_EQ(t1.ToString(), std::string("char(5)")); // Test copy constructor CharType t2 = t1; @@ -63,7 +60,7 @@ TEST(TypesTest, TestVarcharType) { ASSERT_EQ(t1.size, 5); ASSERT_EQ(t1.physical_type.size, 6); - ASSERT_EQ(t1.ToString(), string("varchar(5)")); + ASSERT_EQ(t1.ToString(), std::string("varchar(5)")); // Test copy constructor VarcharType t2 = t1; @@ -78,7 +75,7 @@ TEST(TypesTest, TestStringType) { StringType str_nn(false); ASSERT_EQ(str.type, TypeEnum::STRING); - ASSERT_EQ(str.name(), string("string")); + ASSERT_EQ(str.name(), std::string("string")); ASSERT_TRUE(str.nullable); ASSERT_FALSE(str_nn.nullable); } @@ -111,11 +108,11 @@ class TestStringContainer : public ::testing::Test { } protected: - vector offsets_; - vector chars_; - vector nulls_; + std::vector offsets_; + std::vector chars_; + std::vector nulls_; - vector expected_; + std::vector expected_; std::shared_ptr value_buf_; std::shared_ptr offsets_buf_; @@ -175,7 +172,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(tmp)); } @@ -188,8 +185,8 @@ class TestStringBuilder : public TestBuilder { protected: TypePtr type_; - unique_ptr builder_; - unique_ptr result_; + std::unique_ptr builder_; + std::unique_ptr result_; }; TEST_F(TestStringBuilder, TestAttrs) { @@ -197,8 +194,8 @@ TEST_F(TestStringBuilder, TestAttrs) { } TEST_F(TestStringBuilder, TestScalarAppend) { - vector strings = {"a", "bb", "", "", "ccc"}; - vector is_null = {0, 0, 0, 1, 0}; + std::vector strings = {"a", "bb", "", "", "ccc"}; + std::vector is_null = {0, 0, 0, 1, 0}; int N = strings.size(); int reps = 1000; diff --git a/cpp/src/arrow/types/string.h b/cpp/src/arrow/types/string.h index 30d6e247db1..d0690d9a7d2 100644 --- a/cpp/src/arrow/types/string.h +++ b/cpp/src/arrow/types/string.h @@ -27,12 +27,13 @@ #include "arrow/type.h" #include "arrow/types/integer.h" #include "arrow/types/list.h" -#include "arrow/util/buffer.h" #include "arrow/util/status.h" namespace arrow { class ArrayBuilder; +class Buffer; +class MemoryPool; struct CharType : public DataType { int size; @@ -148,8 +149,9 @@ class StringArray : public ListArray { class StringBuilder : public ListBuilder { public: - explicit StringBuilder(const TypePtr& type) : - ListBuilder(type, static_cast(new UInt8Builder(value_type_))) { + explicit StringBuilder(MemoryPool* pool, const TypePtr& type) : + ListBuilder(pool, type, + static_cast(new UInt8Builder(pool, value_type_))) { byte_builder_ = static_cast(value_builder_.get()); } @@ -171,6 +173,7 @@ class StringBuilder : public ListBuilder { } protected: + std::shared_ptr list_builder_; UInt8Builder* byte_builder_; static TypePtr value_type_; diff --git a/cpp/src/arrow/types/struct.cc b/cpp/src/arrow/types/struct.cc index b7be5d8245f..a245656b516 100644 --- a/cpp/src/arrow/types/struct.cc +++ b/cpp/src/arrow/types/struct.cc @@ -17,6 +17,7 @@ #include "arrow/types/struct.h" +#include #include #include #include diff --git a/cpp/src/arrow/types/test-common.h b/cpp/src/arrow/types/test-common.h index 267e48a7f25..3ecb0dec7c0 100644 --- a/cpp/src/arrow/types/test-common.h +++ b/cpp/src/arrow/types/test-common.h @@ -25,6 +25,7 @@ #include "arrow/test-util.h" #include "arrow/type.h" +#include "arrow/util/memory-pool.h" using std::unique_ptr; @@ -33,12 +34,15 @@ namespace arrow { class TestBuilder : public ::testing::Test { public: void SetUp() { + pool_ = GetDefaultMemoryPool(); type_ = TypePtr(new UInt8Type()); type_nn_ = TypePtr(new UInt8Type(false)); - builder_.reset(new UInt8Builder(type_)); - builder_nn_.reset(new UInt8Builder(type_nn_)); + builder_.reset(new UInt8Builder(pool_, type_)); + builder_nn_.reset(new UInt8Builder(pool_, type_nn_)); } protected: + MemoryPool* pool_; + TypePtr type_; TypePtr type_nn_; unique_ptr builder_; diff --git a/cpp/src/arrow/types/union.cc b/cpp/src/arrow/types/union.cc index 54f41a7eef6..db3f81795ea 100644 --- a/cpp/src/arrow/types/union.cc +++ b/cpp/src/arrow/types/union.cc @@ -17,6 +17,7 @@ #include "arrow/types/union.h" +#include #include #include #include diff --git a/cpp/src/arrow/util/CMakeLists.txt b/cpp/src/arrow/util/CMakeLists.txt index ff8db6a0410..c53f307c9f5 100644 --- a/cpp/src/arrow/util/CMakeLists.txt +++ b/cpp/src/arrow/util/CMakeLists.txt @@ -22,6 +22,7 @@ set(UTIL_SRCS bit-util.cc buffer.cc + memory-pool.cc status.cc ) @@ -39,6 +40,7 @@ install(FILES bit-util.h buffer.h macros.h + memory-pool.h status.h DESTINATION include/arrow/util) @@ -79,3 +81,4 @@ endif() ADD_ARROW_TEST(bit-util-test) ADD_ARROW_TEST(buffer-test) +ADD_ARROW_TEST(memory-pool-test) diff --git a/cpp/src/arrow/util/bit-util.cc b/cpp/src/arrow/util/bit-util.cc index d2ddd6584a8..dbac0a42527 100644 --- a/cpp/src/arrow/util/bit-util.cc +++ b/cpp/src/arrow/util/bit-util.cc @@ -33,7 +33,7 @@ Status util::bytes_to_bits(uint8_t* bytes, int length, std::shared_ptr* out) { int bit_length = ceil_byte(length) / 8; - auto buffer = std::make_shared(); + auto buffer = std::make_shared(); RETURN_NOT_OK(buffer->Resize(bit_length)); memset(buffer->mutable_data(), 0, bit_length); bytes_to_bits(bytes, length, buffer->mutable_data()); diff --git a/cpp/src/arrow/util/bit-util.h b/cpp/src/arrow/util/bit-util.h index 61dffa30423..9ae6127c5ea 100644 --- a/cpp/src/arrow/util/bit-util.h +++ b/cpp/src/arrow/util/bit-util.h @@ -22,10 +22,9 @@ #include #include -#include "arrow/util/buffer.h" - namespace arrow { +class Buffer; class Status; namespace util { diff --git a/cpp/src/arrow/util/buffer-test.cc b/cpp/src/arrow/util/buffer-test.cc index edfd08e850b..9f1fd91432b 100644 --- a/cpp/src/arrow/util/buffer-test.cc +++ b/cpp/src/arrow/util/buffer-test.cc @@ -16,10 +16,8 @@ // under the License. #include -#include #include #include -#include #include #include "arrow/test-util.h" @@ -34,7 +32,7 @@ class TestBuffer : public ::testing::Test { }; TEST_F(TestBuffer, Resize) { - OwnedMutableBuffer buf; + PoolBuffer buf; ASSERT_EQ(0, buf.size()); ASSERT_OK(buf.Resize(100)); @@ -49,7 +47,7 @@ TEST_F(TestBuffer, Resize) { TEST_F(TestBuffer, ResizeOOM) { // realloc fails, even though there may be no explicit limit - OwnedMutableBuffer buf; + PoolBuffer buf; ASSERT_OK(buf.Resize(100)); int64_t to_alloc = std::numeric_limits::max(); ASSERT_RAISES(OutOfMemory, buf.Resize(to_alloc)); diff --git a/cpp/src/arrow/util/buffer.cc b/cpp/src/arrow/util/buffer.cc index 2fb34d59e0b..3f3807d4e20 100644 --- a/cpp/src/arrow/util/buffer.cc +++ b/cpp/src/arrow/util/buffer.cc @@ -19,6 +19,7 @@ #include +#include "arrow/util/memory-pool.h" #include "arrow/util/status.h" namespace arrow { @@ -34,19 +35,34 @@ std::shared_ptr MutableBuffer::GetImmutableView() { return std::make_shared(this->get_shared_ptr(), 0, size()); } -OwnedMutableBuffer::OwnedMutableBuffer() : - MutableBuffer(nullptr, 0) {} +PoolBuffer::PoolBuffer(MemoryPool* pool) : + ResizableBuffer(nullptr, 0) { + if (pool == nullptr) { + pool = GetDefaultMemoryPool(); + } + pool_ = pool; +} -Status OwnedMutableBuffer::Resize(int64_t new_size) { - size_ = new_size; - try { - buffer_owner_.resize(new_size); - } catch (const std::bad_alloc& e) { - return Status::OutOfMemory("resize failed"); +Status PoolBuffer::Reserve(int64_t new_capacity) { + if (!mutable_data_ || new_capacity > capacity_) { + uint8_t* new_data; + if (mutable_data_) { + RETURN_NOT_OK(pool_->Allocate(new_capacity, &new_data)); + memcpy(new_data, mutable_data_, size_); + pool_->Free(mutable_data_, capacity_); + } else { + RETURN_NOT_OK(pool_->Allocate(new_capacity, &new_data)); + } + mutable_data_ = new_data; + data_ = mutable_data_; + capacity_ = new_capacity; } - data_ = buffer_owner_.data(); - mutable_data_ = buffer_owner_.data(); + return Status::OK(); +} +Status PoolBuffer::Resize(int64_t new_size) { + RETURN_NOT_OK(Reserve(new_size)); + size_ = new_size; return Status::OK(); } diff --git a/cpp/src/arrow/util/buffer.h b/cpp/src/arrow/util/buffer.h index 3e4183936b3..8704723eb0a 100644 --- a/cpp/src/arrow/util/buffer.h +++ b/cpp/src/arrow/util/buffer.h @@ -19,15 +19,14 @@ #define ARROW_UTIL_BUFFER_H #include -#include #include #include -#include #include "arrow/util/macros.h" namespace arrow { +class MemoryPool; class Status; // ---------------------------------------------------------------------- @@ -115,17 +114,34 @@ class MutableBuffer : public Buffer { uint8_t* mutable_data_; }; -// A MutableBuffer whose memory is owned by the class instance. For example, -// for reading data out of files that you want to deallocate when this class is -// garbage-collected -class OwnedMutableBuffer : public MutableBuffer { +class ResizableBuffer : public MutableBuffer { public: - OwnedMutableBuffer(); - Status Resize(int64_t new_size); + // Change buffer reported size to indicated size, allocating memory if + // necessary + virtual Status Resize(int64_t new_size) = 0; + + // Ensure that buffer has enough memory allocated to fit the indicated + // capacity. Does not change buffer's reported size + virtual Status Reserve(int64_t new_capacity) = 0; + + protected: + ResizableBuffer(uint8_t* data, int64_t size) : + MutableBuffer(data, size), + capacity_(size) {} + + int64_t capacity_; +}; + +// A Buffer whose lifetime is tied to a particular MemoryPool +class PoolBuffer : public ResizableBuffer { + public: + explicit PoolBuffer(MemoryPool* pool = nullptr); + + virtual Status Resize(int64_t new_size); + virtual Status Reserve(int64_t new_capacity); private: - // TODO: aligned allocations - std::vector buffer_owner_; + MemoryPool* pool_; }; } // namespace arrow diff --git a/cpp/src/arrow/util/memory-pool-test.cc b/cpp/src/arrow/util/memory-pool-test.cc new file mode 100644 index 00000000000..954b5f951b5 --- /dev/null +++ b/cpp/src/arrow/util/memory-pool-test.cc @@ -0,0 +1,47 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include +#include +#include + +#include "arrow/test-util.h" +#include "arrow/util/memory-pool.h" +#include "arrow/util/status.h" + +namespace arrow { + +TEST(DefaultMemoryPool, MemoryTracking) { + MemoryPool* pool = GetDefaultMemoryPool(); + + uint8_t* data; + ASSERT_OK(pool->Allocate(100, &data)); + ASSERT_EQ(100, pool->bytes_allocated()); + + pool->Free(data, 100); + ASSERT_EQ(0, pool->bytes_allocated()); +} + +TEST(DefaultMemoryPool, OOM) { + MemoryPool* pool = GetDefaultMemoryPool(); + + uint8_t* data; + int64_t to_alloc = std::numeric_limits::max(); + ASSERT_RAISES(OutOfMemory, pool->Allocate(to_alloc, &data)); +} + +} // namespace arrow diff --git a/cpp/src/arrow/util/memory-pool.cc b/cpp/src/arrow/util/memory-pool.cc new file mode 100644 index 00000000000..5820346e5a7 --- /dev/null +++ b/cpp/src/arrow/util/memory-pool.cc @@ -0,0 +1,78 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "arrow/util/memory-pool.h" + +#include +#include +#include + +#include "arrow/util/status.h" + +namespace arrow { + +MemoryPool::~MemoryPool() {} + +class InternalMemoryPool : public MemoryPool { + public: + InternalMemoryPool() : bytes_allocated_(0) {} + virtual ~InternalMemoryPool(); + + Status Allocate(int64_t size, uint8_t** out) override; + + void Free(uint8_t* buffer, int64_t size) override; + + int64_t bytes_allocated() const override; + + private: + mutable std::mutex pool_lock_; + int64_t bytes_allocated_; +}; + +Status InternalMemoryPool::Allocate(int64_t size, uint8_t** out) { + std::lock_guard guard(pool_lock_); + *out = static_cast(std::malloc(size)); + if (*out == nullptr) { + std::stringstream ss; + ss << "malloc of size " << size << " failed"; + return Status::OutOfMemory(ss.str()); + } + + bytes_allocated_ += size; + + return Status::OK(); +} + +int64_t InternalMemoryPool::bytes_allocated() const { + std::lock_guard guard(pool_lock_); + return bytes_allocated_; +} + +void InternalMemoryPool::Free(uint8_t* buffer, int64_t size) { + std::lock_guard guard(pool_lock_); + std::free(buffer); + bytes_allocated_ -= size; +} + +InternalMemoryPool::~InternalMemoryPool() {} + +MemoryPool* GetDefaultMemoryPool() { + static InternalMemoryPool default_memory_pool; + return &default_memory_pool; +} + +} // namespace arrow diff --git a/cpp/src/arrow/util/memory-pool.h b/cpp/src/arrow/util/memory-pool.h new file mode 100644 index 00000000000..a7cb10dae17 --- /dev/null +++ b/cpp/src/arrow/util/memory-pool.h @@ -0,0 +1,41 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#ifndef ARROW_UTIL_MEMORY_POOL_H +#define ARROW_UTIL_MEMORY_POOL_H + +#include + +namespace arrow { + +class Status; + +class MemoryPool { + public: + virtual ~MemoryPool(); + + virtual Status Allocate(int64_t size, uint8_t** out) = 0; + virtual void Free(uint8_t* buffer, int64_t size) = 0; + + virtual int64_t bytes_allocated() const = 0; +}; + +MemoryPool* GetDefaultMemoryPool(); + +} // namespace arrow + +#endif // ARROW_UTIL_MEMORY_POOL_H