| 2352 | } |
| 2353 | |
| 2354 | Result<std::shared_ptr<SparseIndex>> ReadSparseCSXIndex( |
| 2355 | const flatbuf::SparseTensor* sparse_tensor, const std::vector<int64_t>& shape, |
| 2356 | int64_t non_zero_length, io::RandomAccessFile* file) { |
| 2357 | if (shape.size() != 2) { |
| 2358 | return Status::Invalid("Invalid shape length for a sparse matrix"); |
| 2359 | } |
| 2360 | |
| 2361 | auto* sparse_index = sparse_tensor->sparseIndex_as_SparseMatrixIndexCSX(); |
| 2362 | |
| 2363 | std::shared_ptr<DataType> indptr_type, indices_type; |
| 2364 | RETURN_NOT_OK( |
| 2365 | internal::GetSparseCSXIndexMetadata(sparse_index, &indptr_type, &indices_type)); |
| 2366 | const int indptr_byte_width = indptr_type->byte_width(); |
| 2367 | |
| 2368 | auto* indptr_buffer = sparse_index->indptrBuffer(); |
| 2369 | ARROW_ASSIGN_OR_RAISE(auto indptr_data, |
| 2370 | file->ReadAt(indptr_buffer->offset(), indptr_buffer->length(), |
| 2371 | /*allow_short_read=*/false)); |
| 2372 | |
| 2373 | auto* indices_buffer = sparse_index->indicesBuffer(); |
| 2374 | ARROW_ASSIGN_OR_RAISE(auto indices_data, |
| 2375 | file->ReadAt(indices_buffer->offset(), indices_buffer->length(), |
| 2376 | /*allow_short_read=*/false)); |
| 2377 | |
| 2378 | std::vector<int64_t> indices_shape({non_zero_length}); |
| 2379 | const auto indices_minimum_bytes = |
| 2380 | MultiplyWithOverflow<int64_t>({indices_shape[0], indices_type->byte_width()}); |
| 2381 | if (!indices_minimum_bytes.has_value() || |
| 2382 | indices_minimum_bytes.value() > indices_buffer->length()) { |
| 2383 | return Status::Invalid("shape is inconsistent to the size of indices buffer"); |
| 2384 | } |
| 2385 | |
| 2386 | switch (sparse_index->compressedAxis()) { |
| 2387 | case flatbuf::SparseMatrixCompressedAxis::SparseMatrixCompressedAxis_Row: { |
| 2388 | const auto indptr_length = AddWithOverflow<int64_t>({shape[0], 1}); |
| 2389 | if (!indptr_length.has_value()) { |
| 2390 | return Status::Invalid("shape is inconsistent to the size of indptr buffer"); |
| 2391 | } |
| 2392 | std::vector<int64_t> indptr_shape({indptr_length.value()}); |
| 2393 | const auto indptr_minimum_bytes = |
| 2394 | MultiplyWithOverflow<int64_t>({indptr_shape[0], indptr_byte_width}); |
| 2395 | if (!indptr_minimum_bytes.has_value() || |
| 2396 | indptr_minimum_bytes.value() > indptr_buffer->length()) { |
| 2397 | return Status::Invalid("shape is inconsistent to the size of indptr buffer"); |
| 2398 | } |
| 2399 | return std::make_shared<SparseCSRIndex>( |
| 2400 | std::make_shared<Tensor>(indptr_type, indptr_data, indptr_shape), |
| 2401 | std::make_shared<Tensor>(indices_type, indices_data, indices_shape)); |
| 2402 | } |
| 2403 | case flatbuf::SparseMatrixCompressedAxis::SparseMatrixCompressedAxis_Column: { |
| 2404 | const auto indptr_length = AddWithOverflow<int64_t>({shape[1], 1}); |
| 2405 | if (!indptr_length.has_value()) { |
| 2406 | return Status::Invalid("shape is inconsistent to the size of indptr buffer"); |
| 2407 | } |
| 2408 | std::vector<int64_t> indptr_shape({indptr_length.value()}); |
| 2409 | const auto indptr_minimum_bytes = |
| 2410 | MultiplyWithOverflow<int64_t>({indptr_shape[0], indptr_byte_width}); |
| 2411 | if (!indptr_minimum_bytes.has_value() || |
no test coverage detected