| 2314 | } |
| 2315 | |
| 2316 | Result<std::shared_ptr<SparseIndex>> ReadSparseCSXIndex( |
| 2317 | const flatbuf::SparseTensor* sparse_tensor, const std::vector<int64_t>& shape, |
| 2318 | int64_t non_zero_length, io::RandomAccessFile* file) { |
| 2319 | if (shape.size() != 2) { |
| 2320 | return Status::Invalid("Invalid shape length for a sparse matrix"); |
| 2321 | } |
| 2322 | |
| 2323 | auto* sparse_index = sparse_tensor->sparseIndex_as_SparseMatrixIndexCSX(); |
| 2324 | |
| 2325 | std::shared_ptr<DataType> indptr_type, indices_type; |
| 2326 | RETURN_NOT_OK( |
| 2327 | internal::GetSparseCSXIndexMetadata(sparse_index, &indptr_type, &indices_type)); |
| 2328 | const int indptr_byte_width = indptr_type->byte_width(); |
| 2329 | |
| 2330 | auto* indptr_buffer = sparse_index->indptrBuffer(); |
| 2331 | ARROW_ASSIGN_OR_RAISE(auto indptr_data, |
| 2332 | file->ReadAt(indptr_buffer->offset(), indptr_buffer->length())); |
| 2333 | |
| 2334 | auto* indices_buffer = sparse_index->indicesBuffer(); |
| 2335 | ARROW_ASSIGN_OR_RAISE(auto indices_data, |
| 2336 | file->ReadAt(indices_buffer->offset(), indices_buffer->length())); |
| 2337 | |
| 2338 | std::vector<int64_t> indices_shape({non_zero_length}); |
| 2339 | const auto indices_minimum_bytes = indices_shape[0] * indices_type->byte_width(); |
| 2340 | if (indices_minimum_bytes > indices_buffer->length()) { |
| 2341 | return Status::Invalid("shape is inconsistent to the size of indices buffer"); |
| 2342 | } |
| 2343 | |
| 2344 | switch (sparse_index->compressedAxis()) { |
| 2345 | case flatbuf::SparseMatrixCompressedAxis::SparseMatrixCompressedAxis_Row: { |
| 2346 | std::vector<int64_t> indptr_shape({shape[0] + 1}); |
| 2347 | const int64_t indptr_minimum_bytes = indptr_shape[0] * indptr_byte_width; |
| 2348 | if (indptr_minimum_bytes > indptr_buffer->length()) { |
| 2349 | return Status::Invalid("shape is inconsistent to the size of indptr buffer"); |
| 2350 | } |
| 2351 | return std::make_shared<SparseCSRIndex>( |
| 2352 | std::make_shared<Tensor>(indptr_type, indptr_data, indptr_shape), |
| 2353 | std::make_shared<Tensor>(indices_type, indices_data, indices_shape)); |
| 2354 | } |
| 2355 | case flatbuf::SparseMatrixCompressedAxis::SparseMatrixCompressedAxis_Column: { |
| 2356 | std::vector<int64_t> indptr_shape({shape[1] + 1}); |
| 2357 | const int64_t indptr_minimum_bytes = indptr_shape[0] * indptr_byte_width; |
| 2358 | if (indptr_minimum_bytes > indptr_buffer->length()) { |
| 2359 | return Status::Invalid("shape is inconsistent to the size of indptr buffer"); |
| 2360 | } |
| 2361 | return std::make_shared<SparseCSCIndex>( |
| 2362 | std::make_shared<Tensor>(indptr_type, indptr_data, indptr_shape), |
| 2363 | std::make_shared<Tensor>(indices_type, indices_data, indices_shape)); |
| 2364 | } |
| 2365 | default: |
| 2366 | return Status::Invalid("Invalid value of SparseMatrixCompressedAxis"); |
| 2367 | } |
| 2368 | } |
| 2369 | |
| 2370 | Result<std::shared_ptr<SparseIndex>> ReadSparseCSFIndex( |
| 2371 | const flatbuf::SparseTensor* sparse_tensor, const std::vector<int64_t>& shape, |
no test coverage detected