| 2571 | } |
| 2572 | |
| 2573 | Result<std::shared_ptr<SparseTensor>> ReadSparseTensorPayload(const IpcPayload& payload) { |
| 2574 | std::shared_ptr<DataType> type; |
| 2575 | std::vector<int64_t> shape; |
| 2576 | std::vector<std::string> dim_names; |
| 2577 | int64_t non_zero_length; |
| 2578 | SparseTensorFormat::type sparse_tensor_format_id; |
| 2579 | const flatbuf::SparseTensor* sparse_tensor; |
| 2580 | const flatbuf::Buffer* buffer; |
| 2581 | |
| 2582 | RETURN_NOT_OK(ReadSparseTensorMetadata(*payload.metadata, &type, &shape, &dim_names, |
| 2583 | &non_zero_length, &sparse_tensor_format_id, |
| 2584 | &sparse_tensor, &buffer)); |
| 2585 | |
| 2586 | RETURN_NOT_OK(CheckSparseTensorBodyBufferCount(payload, sparse_tensor_format_id, |
| 2587 | static_cast<size_t>(shape.size()))); |
| 2588 | |
| 2589 | switch (sparse_tensor_format_id) { |
| 2590 | case SparseTensorFormat::COO: { |
| 2591 | std::shared_ptr<SparseCOOIndex> sparse_index; |
| 2592 | std::shared_ptr<DataType> indices_type; |
| 2593 | RETURN_NOT_OK(internal::GetSparseCOOIndexMetadata( |
| 2594 | sparse_tensor->sparseIndex_as_SparseTensorIndexCOO(), &indices_type)); |
| 2595 | ARROW_ASSIGN_OR_RAISE(sparse_index, |
| 2596 | SparseCOOIndex::Make(indices_type, shape, non_zero_length, |
| 2597 | payload.body_buffers[0])); |
| 2598 | return MakeSparseTensorWithSparseCOOIndex(type, shape, dim_names, sparse_index, |
| 2599 | non_zero_length, payload.body_buffers[1]); |
| 2600 | } |
| 2601 | case SparseTensorFormat::CSR: { |
| 2602 | std::shared_ptr<SparseCSRIndex> sparse_index; |
| 2603 | std::shared_ptr<DataType> indptr_type; |
| 2604 | std::shared_ptr<DataType> indices_type; |
| 2605 | RETURN_NOT_OK(internal::GetSparseCSXIndexMetadata( |
| 2606 | sparse_tensor->sparseIndex_as_SparseMatrixIndexCSX(), &indptr_type, |
| 2607 | &indices_type)); |
| 2608 | ARROW_CHECK_EQ(indptr_type, indices_type); |
| 2609 | ARROW_ASSIGN_OR_RAISE( |
| 2610 | sparse_index, |
| 2611 | SparseCSRIndex::Make(indices_type, shape, non_zero_length, |
| 2612 | payload.body_buffers[0], payload.body_buffers[1])); |
| 2613 | return MakeSparseTensorWithSparseCSRIndex(type, shape, dim_names, sparse_index, |
| 2614 | non_zero_length, payload.body_buffers[2]); |
| 2615 | } |
| 2616 | case SparseTensorFormat::CSC: { |
| 2617 | std::shared_ptr<SparseCSCIndex> sparse_index; |
| 2618 | std::shared_ptr<DataType> indptr_type; |
| 2619 | std::shared_ptr<DataType> indices_type; |
| 2620 | RETURN_NOT_OK(internal::GetSparseCSXIndexMetadata( |
| 2621 | sparse_tensor->sparseIndex_as_SparseMatrixIndexCSX(), &indptr_type, |
| 2622 | &indices_type)); |
| 2623 | ARROW_CHECK_EQ(indptr_type, indices_type); |
| 2624 | ARROW_ASSIGN_OR_RAISE( |
| 2625 | sparse_index, |
| 2626 | SparseCSCIndex::Make(indices_type, shape, non_zero_length, |
| 2627 | payload.body_buffers[0], payload.body_buffers[1])); |
| 2628 | return MakeSparseTensorWithSparseCSCIndex(type, shape, dim_names, sparse_index, |
| 2629 | non_zero_length, payload.body_buffers[2]); |
| 2630 | } |