| 1230 | } |
| 1231 | |
| 1232 | Status DoWrite(const void* data, int64_t nbytes, |
| 1233 | std::shared_ptr<Buffer> owned_buffer = nullptr) { |
| 1234 | if (closed_) { |
| 1235 | return Status::Invalid("Operation on closed stream"); |
| 1236 | } |
| 1237 | |
| 1238 | const auto* data_ptr = reinterpret_cast<const int8_t*>(data); |
| 1239 | auto advance_ptr = [this, &data_ptr, &nbytes](const int64_t offset) { |
| 1240 | data_ptr += offset; |
| 1241 | nbytes -= offset; |
| 1242 | pos_ += offset; |
| 1243 | content_length_ += offset; |
| 1244 | }; |
| 1245 | |
| 1246 | // Handle case where we have some bytes buffered from prior calls. |
| 1247 | if (current_block_size_ > 0) { |
| 1248 | // Try to fill current buffer |
| 1249 | const int64_t to_copy = |
| 1250 | std::min(nbytes, kBlockUploadSizeBytes - current_block_size_); |
| 1251 | RETURN_NOT_OK(current_block_->Write(data_ptr, to_copy)); |
| 1252 | current_block_size_ += to_copy; |
| 1253 | advance_ptr(to_copy); |
| 1254 | |
| 1255 | // If buffer isn't full, break |
| 1256 | if (current_block_size_ < kBlockUploadSizeBytes) { |
| 1257 | return Status::OK(); |
| 1258 | } |
| 1259 | |
| 1260 | // Upload current buffer |
| 1261 | RETURN_NOT_OK(AppendCurrentBlock()); |
| 1262 | } |
| 1263 | |
| 1264 | // We can upload chunks without copying them into a buffer |
| 1265 | while (nbytes >= kBlockUploadSizeBytes) { |
| 1266 | const auto upload_size = std::min(nbytes, kMaxBlockSizeBytes); |
| 1267 | RETURN_NOT_OK(AppendBlock(data_ptr, upload_size)); |
| 1268 | advance_ptr(upload_size); |
| 1269 | } |
| 1270 | |
| 1271 | // Buffer remaining bytes |
| 1272 | if (nbytes > 0) { |
| 1273 | current_block_size_ = nbytes; |
| 1274 | |
| 1275 | if (current_block_ == nullptr) { |
| 1276 | ARROW_ASSIGN_OR_RAISE( |
| 1277 | current_block_, |
| 1278 | io::BufferOutputStream::Create(kBlockUploadSizeBytes, io_context_.pool())); |
| 1279 | } else { |
| 1280 | // Re-use the allocation from before. |
| 1281 | RETURN_NOT_OK(current_block_->Reset(kBlockUploadSizeBytes, io_context_.pool())); |
| 1282 | } |
| 1283 | |
| 1284 | RETURN_NOT_OK(current_block_->Write(data_ptr, current_block_size_)); |
| 1285 | pos_ += current_block_size_; |
| 1286 | content_length_ += current_block_size_; |
| 1287 | } |
| 1288 | |
| 1289 | return Status::OK(); |