]>
Commit | Line | Data |
---|---|---|
1d09f67e TL |
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 <algorithm> | |
21 | #include <cstdint> | |
22 | #include <utility> | |
23 | ||
24 | #include "arrow/result.h" | |
25 | #include "arrow/status.h" | |
26 | #include "arrow/util/bit_util.h" | |
27 | #include "arrow/util/int_util_internal.h" | |
28 | #include "arrow/util/logging.h" | |
29 | #include "arrow/util/string.h" | |
30 | ||
31 | namespace arrow { | |
32 | ||
33 | Result<std::shared_ptr<Buffer>> Buffer::CopySlice(const int64_t start, | |
34 | const int64_t nbytes, | |
35 | MemoryPool* pool) const { | |
36 | // Sanity checks | |
37 | ARROW_CHECK_LE(start, size_); | |
38 | ARROW_CHECK_LE(nbytes, size_ - start); | |
39 | DCHECK_GE(nbytes, 0); | |
40 | ||
41 | ARROW_ASSIGN_OR_RAISE(auto new_buffer, AllocateResizableBuffer(nbytes, pool)); | |
42 | std::memcpy(new_buffer->mutable_data(), data() + start, static_cast<size_t>(nbytes)); | |
43 | return std::move(new_buffer); | |
44 | } | |
45 | ||
46 | namespace { | |
47 | ||
48 | Status CheckBufferSlice(const Buffer& buffer, int64_t offset, int64_t length) { | |
49 | return internal::CheckSliceParams(buffer.size(), offset, length, "buffer"); | |
50 | } | |
51 | ||
52 | Status CheckBufferSlice(const Buffer& buffer, int64_t offset) { | |
53 | if (ARROW_PREDICT_FALSE(offset < 0)) { | |
54 | // Avoid UBSAN in subtraction below | |
55 | return Status::Invalid("Negative buffer slice offset"); | |
56 | } | |
57 | return CheckBufferSlice(buffer, offset, buffer.size() - offset); | |
58 | } | |
59 | ||
60 | } // namespace | |
61 | ||
62 | Result<std::shared_ptr<Buffer>> SliceBufferSafe(const std::shared_ptr<Buffer>& buffer, | |
63 | int64_t offset) { | |
64 | RETURN_NOT_OK(CheckBufferSlice(*buffer, offset)); | |
65 | return SliceBuffer(buffer, offset); | |
66 | } | |
67 | ||
68 | Result<std::shared_ptr<Buffer>> SliceBufferSafe(const std::shared_ptr<Buffer>& buffer, | |
69 | int64_t offset, int64_t length) { | |
70 | RETURN_NOT_OK(CheckBufferSlice(*buffer, offset, length)); | |
71 | return SliceBuffer(buffer, offset, length); | |
72 | } | |
73 | ||
74 | Result<std::shared_ptr<Buffer>> SliceMutableBufferSafe( | |
75 | const std::shared_ptr<Buffer>& buffer, int64_t offset) { | |
76 | RETURN_NOT_OK(CheckBufferSlice(*buffer, offset)); | |
77 | return SliceMutableBuffer(buffer, offset); | |
78 | } | |
79 | ||
80 | Result<std::shared_ptr<Buffer>> SliceMutableBufferSafe( | |
81 | const std::shared_ptr<Buffer>& buffer, int64_t offset, int64_t length) { | |
82 | RETURN_NOT_OK(CheckBufferSlice(*buffer, offset, length)); | |
83 | return SliceMutableBuffer(buffer, offset, length); | |
84 | } | |
85 | ||
86 | std::string Buffer::ToHexString() { | |
87 | return HexEncode(data(), static_cast<size_t>(size())); | |
88 | } | |
89 | ||
90 | bool Buffer::Equals(const Buffer& other, const int64_t nbytes) const { | |
91 | return this == &other || (size_ >= nbytes && other.size_ >= nbytes && | |
92 | (data_ == other.data_ || | |
93 | !memcmp(data_, other.data_, static_cast<size_t>(nbytes)))); | |
94 | } | |
95 | ||
96 | bool Buffer::Equals(const Buffer& other) const { | |
97 | return this == &other || (size_ == other.size_ && | |
98 | (data_ == other.data_ || | |
99 | !memcmp(data_, other.data_, static_cast<size_t>(size_)))); | |
100 | } | |
101 | ||
102 | std::string Buffer::ToString() const { | |
103 | return std::string(reinterpret_cast<const char*>(data_), static_cast<size_t>(size_)); | |
104 | } | |
105 | ||
106 | void Buffer::CheckMutable() const { DCHECK(is_mutable()) << "buffer not mutable"; } | |
107 | ||
108 | void Buffer::CheckCPU() const { | |
109 | DCHECK(is_cpu()) << "not a CPU buffer (device: " << device()->ToString() << ")"; | |
110 | } | |
111 | ||
112 | Result<std::shared_ptr<io::RandomAccessFile>> Buffer::GetReader( | |
113 | std::shared_ptr<Buffer> buf) { | |
114 | return buf->memory_manager_->GetBufferReader(buf); | |
115 | } | |
116 | ||
117 | Result<std::shared_ptr<io::OutputStream>> Buffer::GetWriter(std::shared_ptr<Buffer> buf) { | |
118 | if (!buf->is_mutable()) { | |
119 | return Status::Invalid("Expected mutable buffer"); | |
120 | } | |
121 | return buf->memory_manager_->GetBufferWriter(buf); | |
122 | } | |
123 | ||
124 | Result<std::shared_ptr<Buffer>> Buffer::Copy(std::shared_ptr<Buffer> source, | |
125 | const std::shared_ptr<MemoryManager>& to) { | |
126 | return MemoryManager::CopyBuffer(source, to); | |
127 | } | |
128 | ||
129 | Result<std::shared_ptr<Buffer>> Buffer::View(std::shared_ptr<Buffer> source, | |
130 | const std::shared_ptr<MemoryManager>& to) { | |
131 | return MemoryManager::ViewBuffer(source, to); | |
132 | } | |
133 | ||
134 | Result<std::shared_ptr<Buffer>> Buffer::ViewOrCopy( | |
135 | std::shared_ptr<Buffer> source, const std::shared_ptr<MemoryManager>& to) { | |
136 | auto maybe_buffer = MemoryManager::ViewBuffer(source, to); | |
137 | if (maybe_buffer.ok()) { | |
138 | return maybe_buffer; | |
139 | } | |
140 | return MemoryManager::CopyBuffer(source, to); | |
141 | } | |
142 | ||
143 | class StlStringBuffer : public Buffer { | |
144 | public: | |
145 | explicit StlStringBuffer(std::string data) | |
146 | : Buffer(nullptr, 0), input_(std::move(data)) { | |
147 | data_ = reinterpret_cast<const uint8_t*>(input_.c_str()); | |
148 | size_ = static_cast<int64_t>(input_.size()); | |
149 | capacity_ = size_; | |
150 | } | |
151 | ||
152 | private: | |
153 | std::string input_; | |
154 | }; | |
155 | ||
156 | std::shared_ptr<Buffer> Buffer::FromString(std::string data) { | |
157 | return std::make_shared<StlStringBuffer>(std::move(data)); | |
158 | } | |
159 | ||
160 | std::shared_ptr<Buffer> SliceMutableBuffer(const std::shared_ptr<Buffer>& buffer, | |
161 | const int64_t offset, const int64_t length) { | |
162 | return std::make_shared<MutableBuffer>(buffer, offset, length); | |
163 | } | |
164 | ||
165 | MutableBuffer::MutableBuffer(const std::shared_ptr<Buffer>& parent, const int64_t offset, | |
166 | const int64_t size) | |
167 | : MutableBuffer(reinterpret_cast<uint8_t*>(parent->mutable_address()) + offset, | |
168 | size) { | |
169 | DCHECK(parent->is_mutable()) << "Must pass mutable buffer"; | |
170 | parent_ = parent; | |
171 | } | |
172 | ||
173 | Result<std::shared_ptr<Buffer>> AllocateBitmap(int64_t length, MemoryPool* pool) { | |
174 | ARROW_ASSIGN_OR_RAISE(auto buf, AllocateBuffer(BitUtil::BytesForBits(length), pool)); | |
175 | // Zero out any trailing bits | |
176 | if (buf->size() > 0) { | |
177 | buf->mutable_data()[buf->size() - 1] = 0; | |
178 | } | |
179 | return std::move(buf); | |
180 | } | |
181 | ||
182 | Result<std::shared_ptr<Buffer>> AllocateEmptyBitmap(int64_t length, MemoryPool* pool) { | |
183 | ARROW_ASSIGN_OR_RAISE(auto buf, AllocateBuffer(BitUtil::BytesForBits(length), pool)); | |
184 | memset(buf->mutable_data(), 0, static_cast<size_t>(buf->size())); | |
185 | return std::move(buf); | |
186 | } | |
187 | ||
188 | Status AllocateEmptyBitmap(int64_t length, std::shared_ptr<Buffer>* out) { | |
189 | return AllocateEmptyBitmap(length).Value(out); | |
190 | } | |
191 | ||
192 | Result<std::shared_ptr<Buffer>> ConcatenateBuffers( | |
193 | const std::vector<std::shared_ptr<Buffer>>& buffers, MemoryPool* pool) { | |
194 | int64_t out_length = 0; | |
195 | for (const auto& buffer : buffers) { | |
196 | out_length += buffer->size(); | |
197 | } | |
198 | ARROW_ASSIGN_OR_RAISE(auto out, AllocateBuffer(out_length, pool)); | |
199 | auto out_data = out->mutable_data(); | |
200 | for (const auto& buffer : buffers) { | |
201 | std::memcpy(out_data, buffer->data(), buffer->size()); | |
202 | out_data += buffer->size(); | |
203 | } | |
204 | return std::move(out); | |
205 | } | |
206 | ||
207 | } // namespace arrow |