| 2312 | } |
| 2313 | |
| 2314 | Result<std::shared_ptr<SparseIndex>> ReadSparseCOOIndex( |
| 2315 | const flatbuf::SparseTensor* sparse_tensor, const std::vector<int64_t>& shape, |
| 2316 | int64_t non_zero_length, io::RandomAccessFile* file) { |
| 2317 | auto* sparse_index = sparse_tensor->sparseIndex_as_SparseTensorIndexCOO(); |
| 2318 | const auto ndim = static_cast<int64_t>(shape.size()); |
| 2319 | |
| 2320 | std::shared_ptr<DataType> indices_type; |
| 2321 | RETURN_NOT_OK(internal::GetSparseCOOIndexMetadata(sparse_index, &indices_type)); |
| 2322 | const int64_t indices_elsize = indices_type->byte_width(); |
| 2323 | |
| 2324 | auto* indices_buffer = sparse_index->indicesBuffer(); |
| 2325 | ARROW_ASSIGN_OR_RAISE(auto indices_data, |
| 2326 | file->ReadAt(indices_buffer->offset(), indices_buffer->length(), |
| 2327 | /*allow_short_read=*/false)); |
| 2328 | std::vector<int64_t> indices_shape({non_zero_length, ndim}); |
| 2329 | int64_t indices_minimum_bytes; |
| 2330 | if (MultiplyWithOverflow(non_zero_length, ndim, &indices_minimum_bytes) || |
| 2331 | MultiplyWithOverflow(indices_minimum_bytes, indices_elsize, |
| 2332 | &indices_minimum_bytes) || |
| 2333 | indices_minimum_bytes > indices_buffer->length()) { |
| 2334 | return Status::Invalid("shape is inconsistent to the size of indices buffer"); |
| 2335 | } |
| 2336 | auto* indices_strides = sparse_index->indicesStrides(); |
| 2337 | std::vector<int64_t> strides(2); |
| 2338 | if (indices_strides && indices_strides->size() > 0) { |
| 2339 | if (indices_strides->size() != 2) { |
| 2340 | return Status::Invalid("Wrong size for indicesStrides in SparseCOOIndex"); |
| 2341 | } |
| 2342 | strides[0] = indices_strides->Get(0); |
| 2343 | strides[1] = indices_strides->Get(1); |
| 2344 | } else { |
| 2345 | // Row-major by default |
| 2346 | strides[0] = indices_elsize * ndim; |
| 2347 | strides[1] = indices_elsize; |
| 2348 | } |
| 2349 | return SparseCOOIndex::Make( |
| 2350 | std::make_shared<Tensor>(indices_type, indices_data, indices_shape, strides), |
| 2351 | sparse_index->isCanonical()); |
| 2352 | } |
| 2353 | |
| 2354 | Result<std::shared_ptr<SparseIndex>> ReadSparseCSXIndex( |
| 2355 | const flatbuf::SparseTensor* sparse_tensor, const std::vector<int64_t>& shape, |
no test coverage detected