| 2606 | namespace { |
| 2607 | |
| 2608 | Result<std::shared_ptr<SparseTensor>> ReadSparseTensor(const Buffer& metadata, |
| 2609 | io::RandomAccessFile* file) { |
| 2610 | std::shared_ptr<DataType> type; |
| 2611 | std::vector<int64_t> shape; |
| 2612 | std::vector<std::string> dim_names; |
| 2613 | int64_t non_zero_length; |
| 2614 | SparseTensorFormat::type sparse_tensor_format_id; |
| 2615 | const flatbuf::SparseTensor* sparse_tensor; |
| 2616 | const flatbuf::Buffer* buffer; |
| 2617 | |
| 2618 | RETURN_NOT_OK(ReadSparseTensorMetadata(metadata, &type, &shape, &dim_names, |
| 2619 | &non_zero_length, &sparse_tensor_format_id, |
| 2620 | &sparse_tensor, &buffer)); |
| 2621 | |
| 2622 | ARROW_ASSIGN_OR_RAISE(auto data, file->ReadAt(buffer->offset(), buffer->length())); |
| 2623 | |
| 2624 | std::shared_ptr<SparseIndex> sparse_index; |
| 2625 | switch (sparse_tensor_format_id) { |
| 2626 | case SparseTensorFormat::COO: { |
| 2627 | ARROW_ASSIGN_OR_RAISE( |
| 2628 | sparse_index, ReadSparseCOOIndex(sparse_tensor, shape, non_zero_length, file)); |
| 2629 | return MakeSparseTensorWithSparseCOOIndex( |
| 2630 | type, shape, dim_names, checked_pointer_cast<SparseCOOIndex>(sparse_index), |
| 2631 | non_zero_length, data); |
| 2632 | } |
| 2633 | case SparseTensorFormat::CSR: { |
| 2634 | ARROW_ASSIGN_OR_RAISE( |
| 2635 | sparse_index, ReadSparseCSXIndex(sparse_tensor, shape, non_zero_length, file)); |
| 2636 | return MakeSparseTensorWithSparseCSRIndex( |
| 2637 | type, shape, dim_names, checked_pointer_cast<SparseCSRIndex>(sparse_index), |
| 2638 | non_zero_length, data); |
| 2639 | } |
| 2640 | case SparseTensorFormat::CSC: { |
| 2641 | ARROW_ASSIGN_OR_RAISE( |
| 2642 | sparse_index, ReadSparseCSXIndex(sparse_tensor, shape, non_zero_length, file)); |
| 2643 | return MakeSparseTensorWithSparseCSCIndex( |
| 2644 | type, shape, dim_names, checked_pointer_cast<SparseCSCIndex>(sparse_index), |
| 2645 | non_zero_length, data); |
| 2646 | } |
| 2647 | case SparseTensorFormat::CSF: { |
| 2648 | ARROW_ASSIGN_OR_RAISE(sparse_index, ReadSparseCSFIndex(sparse_tensor, shape, file)); |
| 2649 | return MakeSparseTensorWithSparseCSFIndex( |
| 2650 | type, shape, dim_names, checked_pointer_cast<SparseCSFIndex>(sparse_index), |
| 2651 | data); |
| 2652 | } |
| 2653 | default: |
| 2654 | return Status::Invalid("Unsupported sparse index format"); |
| 2655 | } |
| 2656 | } |
| 2657 | |
| 2658 | } // namespace |
| 2659 | |