| 262 | |
| 263 | template <typename T> |
| 264 | inline int BitReader::GetBatch(int num_bits, T* v, int batch_size) { |
| 265 | constexpr uint64_t kBitsPerByte = 8; |
| 266 | |
| 267 | ARROW_DCHECK(buffer_ != NULLPTR); |
| 268 | ARROW_DCHECK_LE(num_bits, static_cast<int>(sizeof(T) * 8)) << "num_bits: " << num_bits; |
| 269 | |
| 270 | const int64_t needed_bits = num_bits * static_cast<int64_t>(batch_size); |
| 271 | const int64_t remaining_bits = |
| 272 | static_cast<int64_t>(max_bytes_ - byte_offset_) * kBitsPerByte - bit_offset_; |
| 273 | if (remaining_bits < needed_bits) { |
| 274 | batch_size = static_cast<int>(remaining_bits / num_bits); |
| 275 | } |
| 276 | |
| 277 | const ::arrow::internal::UnpackOptions opts{ |
| 278 | .batch_size = batch_size, |
| 279 | .bit_width = num_bits, |
| 280 | .bit_offset = bit_offset_, |
| 281 | .max_read_bytes = max_bytes_ - byte_offset_, |
| 282 | }; |
| 283 | |
| 284 | if constexpr (std::is_same_v<T, bool>) { |
| 285 | ::arrow::internal::unpack(buffer_ + byte_offset_, v, opts); |
| 286 | |
| 287 | } else { |
| 288 | ::arrow::internal::unpack(buffer_ + byte_offset_, |
| 289 | reinterpret_cast<std::make_unsigned_t<T>*>(v), opts); |
| 290 | } |
| 291 | |
| 292 | Advance(batch_size * num_bits); |
| 293 | |
| 294 | return batch_size; |
| 295 | } |
| 296 | |
| 297 | template <typename T> |
| 298 | inline bool BitReader::GetAligned(int num_bytes, T* v) { |