MCPcopy Create free account
hub / github.com/apache/arrow / ReadSparseTensorPayload

Function ReadSparseTensorPayload

cpp/src/arrow/ipc/reader.cc:2514–2602  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

2512}
2513
2514Result<std::shared_ptr<SparseTensor>> ReadSparseTensorPayload(const IpcPayload& payload) {
2515 std::shared_ptr<DataType> type;
2516 std::vector<int64_t> shape;
2517 std::vector<std::string> dim_names;
2518 int64_t non_zero_length;
2519 SparseTensorFormat::type sparse_tensor_format_id;
2520 const flatbuf::SparseTensor* sparse_tensor;
2521 const flatbuf::Buffer* buffer;
2522
2523 RETURN_NOT_OK(ReadSparseTensorMetadata(*payload.metadata, &type, &shape, &dim_names,
2524 &non_zero_length, &sparse_tensor_format_id,
2525 &sparse_tensor, &buffer));
2526
2527 RETURN_NOT_OK(CheckSparseTensorBodyBufferCount(payload, sparse_tensor_format_id,
2528 static_cast<size_t>(shape.size())));
2529
2530 switch (sparse_tensor_format_id) {
2531 case SparseTensorFormat::COO: {
2532 std::shared_ptr<SparseCOOIndex> sparse_index;
2533 std::shared_ptr<DataType> indices_type;
2534 RETURN_NOT_OK(internal::GetSparseCOOIndexMetadata(
2535 sparse_tensor->sparseIndex_as_SparseTensorIndexCOO(), &indices_type));
2536 ARROW_ASSIGN_OR_RAISE(sparse_index,
2537 SparseCOOIndex::Make(indices_type, shape, non_zero_length,
2538 payload.body_buffers[0]));
2539 return MakeSparseTensorWithSparseCOOIndex(type, shape, dim_names, sparse_index,
2540 non_zero_length, payload.body_buffers[1]);
2541 }
2542 case SparseTensorFormat::CSR: {
2543 std::shared_ptr<SparseCSRIndex> sparse_index;
2544 std::shared_ptr<DataType> indptr_type;
2545 std::shared_ptr<DataType> indices_type;
2546 RETURN_NOT_OK(internal::GetSparseCSXIndexMetadata(
2547 sparse_tensor->sparseIndex_as_SparseMatrixIndexCSX(), &indptr_type,
2548 &indices_type));
2549 ARROW_CHECK_EQ(indptr_type, indices_type);
2550 ARROW_ASSIGN_OR_RAISE(
2551 sparse_index,
2552 SparseCSRIndex::Make(indices_type, shape, non_zero_length,
2553 payload.body_buffers[0], payload.body_buffers[1]));
2554 return MakeSparseTensorWithSparseCSRIndex(type, shape, dim_names, sparse_index,
2555 non_zero_length, payload.body_buffers[2]);
2556 }
2557 case SparseTensorFormat::CSC: {
2558 std::shared_ptr<SparseCSCIndex> sparse_index;
2559 std::shared_ptr<DataType> indptr_type;
2560 std::shared_ptr<DataType> indices_type;
2561 RETURN_NOT_OK(internal::GetSparseCSXIndexMetadata(
2562 sparse_tensor->sparseIndex_as_SparseMatrixIndexCSX(), &indptr_type,
2563 &indices_type));
2564 ARROW_CHECK_EQ(indptr_type, indices_type);
2565 ARROW_ASSIGN_OR_RAISE(
2566 sparse_index,
2567 SparseCSCIndex::Make(indices_type, shape, non_zero_length,
2568 payload.body_buffers[0], payload.body_buffers[1]));
2569 return MakeSparseTensorWithSparseCSCIndex(type, shape, dim_names, sparse_index,
2570 non_zero_length, payload.body_buffers[2]);
2571 }

Callers

nothing calls this directly

Calls 13

ReadSparseTensorMetadataFunction · 0.85
ARROW_ASSIGN_OR_RAISEFunction · 0.70
MakeFunction · 0.50
InvalidFunction · 0.50

Tested by

no test coverage detected