1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements. See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership. The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License. You may obtain a copy of the License at
8//
9// http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied. See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18#include "arrow/buffer.h"
19
20#include <cstdint>
21#include <utility>
22
23#include "arrow/memory_pool.h"
24#include "arrow/status.h"
25#include "arrow/util/bit-util.h"
26#include "arrow/util/logging.h"
27
28namespace arrow {
29
30Status Buffer::Copy(const int64_t start, const int64_t nbytes, MemoryPool* pool,
31 std::shared_ptr<Buffer>* out) const {
32 // Sanity checks
33 DCHECK_LT(start, size_);
34 DCHECK_LE(nbytes, size_ - start);
35
36 std::shared_ptr<ResizableBuffer> new_buffer;
37 RETURN_NOT_OK(AllocateResizableBuffer(pool, nbytes, &new_buffer));
38
39 std::memcpy(new_buffer->mutable_data(), data() + start, static_cast<size_t>(nbytes));
40
41 *out = new_buffer;
42 return Status::OK();
43}
44
45Status Buffer::Copy(const int64_t start, const int64_t nbytes,
46 std::shared_ptr<Buffer>* out) const {
47 return Copy(start, nbytes, default_memory_pool(), out);
48}
49
50bool Buffer::Equals(const Buffer& other, const int64_t nbytes) const {
51 return this == &other || (size_ >= nbytes && other.size_ >= nbytes &&
52 (data_ == other.data_ ||
53 !memcmp(data_, other.data_, static_cast<size_t>(nbytes))));
54}
55
56bool Buffer::Equals(const Buffer& other) const {
57 return this == &other || (size_ == other.size_ &&
58 (data_ == other.data_ ||
59 !memcmp(data_, other.data_, static_cast<size_t>(size_))));
60}
61
62Status Buffer::FromString(const std::string& data, MemoryPool* pool,
63 std::shared_ptr<Buffer>* out) {
64 auto size = static_cast<int64_t>(data.size());
65 RETURN_NOT_OK(AllocateBuffer(pool, size, out));
66 std::copy(data.c_str(), data.c_str() + size, (*out)->mutable_data());
67 return Status::OK();
68}
69
70Status Buffer::FromString(const std::string& data, std::shared_ptr<Buffer>* out) {
71 return FromString(data, default_memory_pool(), out);
72}
73
74class StlStringBuffer : public Buffer {
75 public:
76 explicit StlStringBuffer(std::string&& data) : Buffer(nullptr, 0), input_(data) {
77 data_ = reinterpret_cast<const uint8_t*>(input_.c_str());
78 size_ = static_cast<int64_t>(input_.size());
79 capacity_ = size_;
80 }
81
82 private:
83 std::string input_;
84};
85
86std::shared_ptr<Buffer> Buffer::FromString(std::string&& data) {
87 return std::make_shared<StlStringBuffer>(std::move(data));
88}
89
90std::string Buffer::ToString() const {
91 return std::string(reinterpret_cast<const char*>(data_), static_cast<size_t>(size_));
92}
93
94void Buffer::CheckMutable() const { DCHECK(is_mutable()) << "buffer not mutable"; }
95
96/// A Buffer whose lifetime is tied to a particular MemoryPool
97class PoolBuffer : public ResizableBuffer {
98 public:
99 explicit PoolBuffer(MemoryPool* pool) : ResizableBuffer(nullptr, 0) {
100 if (pool == nullptr) {
101 pool = default_memory_pool();
102 }
103 pool_ = pool;
104 }
105
106 ~PoolBuffer() override {
107 if (mutable_data_ != nullptr) {
108 pool_->Free(mutable_data_, capacity_);
109 }
110 }
111
112 Status Reserve(const int64_t capacity) override {
113 if (!mutable_data_ || capacity > capacity_) {
114 uint8_t* new_data;
115 int64_t new_capacity = BitUtil::RoundUpToMultipleOf64(capacity);
116 if (mutable_data_) {
117 RETURN_NOT_OK(pool_->Reallocate(capacity_, new_capacity, &mutable_data_));
118 } else {
119 RETURN_NOT_OK(pool_->Allocate(new_capacity, &new_data));
120 mutable_data_ = new_data;
121 }
122 data_ = mutable_data_;
123 capacity_ = new_capacity;
124 }
125 return Status::OK();
126 }
127
128 Status Resize(const int64_t new_size, bool shrink_to_fit = true) override {
129 if (mutable_data_ && shrink_to_fit && new_size <= size_) {
130 // Buffer is non-null and is not growing, so shrink to the requested size without
131 // excess space.
132 int64_t new_capacity = BitUtil::RoundUpToMultipleOf64(new_size);
133 if (capacity_ != new_capacity) {
134 // Buffer hasn't got yet the requested size.
135 RETURN_NOT_OK(pool_->Reallocate(capacity_, new_capacity, &mutable_data_));
136 data_ = mutable_data_;
137 capacity_ = new_capacity;
138 }
139 } else {
140 RETURN_NOT_OK(Reserve(new_size));
141 }
142 size_ = new_size;
143
144 return Status::OK();
145 }
146
147 private:
148 MemoryPool* pool_;
149};
150
151std::shared_ptr<Buffer> SliceMutableBuffer(const std::shared_ptr<Buffer>& buffer,
152 const int64_t offset, const int64_t length) {
153 return std::make_shared<MutableBuffer>(buffer, offset, length);
154}
155
156MutableBuffer::MutableBuffer(const std::shared_ptr<Buffer>& parent, const int64_t offset,
157 const int64_t size)
158 : MutableBuffer(parent->mutable_data() + offset, size) {
159 DCHECK(parent->is_mutable()) << "Must pass mutable buffer";
160 parent_ = parent;
161}
162
163namespace {
164// A utility that does most of the work of the `AllocateBuffer` and
165// `AllocateResizableBuffer` methods. The argument `buffer` should be a smart pointer to a
166// PoolBuffer.
167template <typename PoolBufferPtr, typename BufferPtr>
168inline Status ResizePoolBuffer(PoolBufferPtr&& buffer, const int64_t size,
169 BufferPtr* out) {
170 RETURN_NOT_OK(buffer->Resize(size));
171 buffer->ZeroPadding();
172 *out = std::move(buffer);
173 return Status::OK();
174}
175} // namespace
176
177Status AllocateBuffer(MemoryPool* pool, const int64_t size,
178 std::shared_ptr<Buffer>* out) {
179 return ResizePoolBuffer(std::make_shared<PoolBuffer>(pool), size, out);
180}
181
182Status AllocateBuffer(MemoryPool* pool, const int64_t size,
183 std::unique_ptr<Buffer>* out) {
184 return ResizePoolBuffer(std::unique_ptr<PoolBuffer>(new PoolBuffer(pool)), size, out);
185}
186
187Status AllocateBuffer(const int64_t size, std::shared_ptr<Buffer>* out) {
188 return AllocateBuffer(default_memory_pool(), size, out);
189}
190
191Status AllocateBuffer(const int64_t size, std::unique_ptr<Buffer>* out) {
192 return AllocateBuffer(default_memory_pool(), size, out);
193}
194
195Status AllocateResizableBuffer(MemoryPool* pool, const int64_t size,
196 std::shared_ptr<ResizableBuffer>* out) {
197 return ResizePoolBuffer(std::make_shared<PoolBuffer>(pool), size, out);
198}
199
200Status AllocateResizableBuffer(MemoryPool* pool, const int64_t size,
201 std::unique_ptr<ResizableBuffer>* out) {
202 return ResizePoolBuffer(std::unique_ptr<PoolBuffer>(new PoolBuffer(pool)), size, out);
203}
204
205Status AllocateResizableBuffer(const int64_t size,
206 std::shared_ptr<ResizableBuffer>* out) {
207 return AllocateResizableBuffer(default_memory_pool(), size, out);
208}
209
210Status AllocateResizableBuffer(const int64_t size,
211 std::unique_ptr<ResizableBuffer>* out) {
212 return AllocateResizableBuffer(default_memory_pool(), size, out);
213}
214
215Status AllocateBitmap(MemoryPool* pool, int64_t length, std::shared_ptr<Buffer>* out) {
216 return AllocateBuffer(pool, BitUtil::BytesForBits(length), out);
217}
218
219Status AllocateEmptyBitmap(MemoryPool* pool, int64_t length,
220 std::shared_ptr<Buffer>* out) {
221 RETURN_NOT_OK(AllocateBitmap(pool, length, out));
222 memset((*out)->mutable_data(), 0, static_cast<size_t>((*out)->size()));
223 return Status::OK();
224}
225
226Status AllocateEmptyBitmap(int64_t length, std::shared_ptr<Buffer>* out) {
227 return AllocateEmptyBitmap(default_memory_pool(), length, out);
228}
229
230} // namespace arrow
231