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 | |
28 | namespace arrow { |
29 | |
30 | Status 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 | |
45 | Status 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 | |
50 | bool 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 | |
56 | bool 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 | |
62 | Status 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 | |
70 | Status Buffer::FromString(const std::string& data, std::shared_ptr<Buffer>* out) { |
71 | return FromString(data, default_memory_pool(), out); |
72 | } |
73 | |
74 | class 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 | |
86 | std::shared_ptr<Buffer> Buffer::FromString(std::string&& data) { |
87 | return std::make_shared<StlStringBuffer>(std::move(data)); |
88 | } |
89 | |
90 | std::string Buffer::ToString() const { |
91 | return std::string(reinterpret_cast<const char*>(data_), static_cast<size_t>(size_)); |
92 | } |
93 | |
94 | void Buffer::CheckMutable() const { DCHECK(is_mutable()) << "buffer not mutable" ; } |
95 | |
96 | /// A Buffer whose lifetime is tied to a particular MemoryPool |
97 | class 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 | |
151 | std::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 | |
156 | MutableBuffer::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 | |
163 | namespace { |
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. |
167 | template <typename PoolBufferPtr, typename BufferPtr> |
168 | inline 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 | |
177 | Status 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 | |
182 | Status 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 | |
187 | Status AllocateBuffer(const int64_t size, std::shared_ptr<Buffer>* out) { |
188 | return AllocateBuffer(default_memory_pool(), size, out); |
189 | } |
190 | |
191 | Status AllocateBuffer(const int64_t size, std::unique_ptr<Buffer>* out) { |
192 | return AllocateBuffer(default_memory_pool(), size, out); |
193 | } |
194 | |
195 | Status 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 | |
200 | Status 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 | |
205 | Status AllocateResizableBuffer(const int64_t size, |
206 | std::shared_ptr<ResizableBuffer>* out) { |
207 | return AllocateResizableBuffer(default_memory_pool(), size, out); |
208 | } |
209 | |
210 | Status AllocateResizableBuffer(const int64_t size, |
211 | std::unique_ptr<ResizableBuffer>* out) { |
212 | return AllocateResizableBuffer(default_memory_pool(), size, out); |
213 | } |
214 | |
215 | Status AllocateBitmap(MemoryPool* pool, int64_t length, std::shared_ptr<Buffer>* out) { |
216 | return AllocateBuffer(pool, BitUtil::BytesForBits(length), out); |
217 | } |
218 | |
219 | Status 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 | |
226 | Status AllocateEmptyBitmap(int64_t length, std::shared_ptr<Buffer>* out) { |
227 | return AllocateEmptyBitmap(default_memory_pool(), length, out); |
228 | } |
229 | |
230 | } // namespace arrow |
231 | |