Function MakeSparseTensorIndexCSF

cpp/src/arrow/ipc/metadata_internal.cc:1145–1202

}

Status MakeSparseTensorIndexCSF(FBB& fbb, const SparseCSFIndex& sparse_index,
                                const std::vector<BufferMetadata>& buffers,
                                flatbuf::SparseTensorIndex* fb_sparse_index_type,
                                Offset* fb_sparse_index, size_t* num_buffers) {
  *fb_sparse_index_type =
      flatbuf::SparseTensorIndex::SparseTensorIndex_SparseTensorIndexCSF;
  const int ndim = static_cast<int>(sparse_index.axis_order().size());

  // We assume that the value type of indptr tensor is an integer.
  const auto& indptr_value_type =
      checked_cast<const IntegerType&>(*sparse_index.indptr()[0]->type());
  auto indptr_type_offset = flatbuf::CreateInt(fbb, indptr_value_type.bit_width(),
                                               indptr_value_type.is_signed());

  // We assume that the value type of indices tensor is an integer.
  const auto& indices_value_type =
      checked_cast<const IntegerType&>(*sparse_index.indices()[0]->type());
  auto indices_type_offset = flatbuf::CreateInt(fbb, indices_value_type.bit_width(),
                                                indices_value_type.is_signed());

  const int64_t indptr_elem_size = indptr_value_type.byte_width();
  const int64_t indices_elem_size = indices_value_type.byte_width();

  int64_t offset = 0;
  std::vector<flatbuf::Buffer> indptr, indices;

  for (const std::shared_ptr<arrow::Tensor>& tensor : sparse_index.indptr()) {
    const int64_t size = tensor->data()->size() / indptr_elem_size;
    const int64_t padded_size = PaddedLength(tensor->data()->size(), kArrowIpcAlignment);

    indptr.push_back({offset, size});
    offset += padded_size;
  }
  for (const std::shared_ptr<arrow::Tensor>& tensor : sparse_index.indices()) {
    const int64_t size = tensor->data()->size() / indices_elem_size;
    const int64_t padded_size = PaddedLength(tensor->data()->size(), kArrowIpcAlignment);

    indices.push_back({offset, size});
    offset += padded_size;
  }

  auto fb_indices = fbb.CreateVectorOfStructs(indices);
  auto fb_indptr = fbb.CreateVectorOfStructs(indptr);

  std::vector<int> axis_order;
  for (int i = 0; i < ndim; ++i) {
    axis_order.emplace_back(static_cast<int>(sparse_index.axis_order()[i]));
  }
  auto fb_axis_order =
      fbb.CreateVector(arrow::util::MakeNonNull(axis_order.data()), axis_order.size());

  *fb_sparse_index =
      flatbuf::CreateSparseTensorIndexCSF(fbb, indptr_type_offset, fb_indptr,
                                          indices_type_offset, fb_indices, fb_axis_order)
          .Union();
  *num_buffers = 2 * ndim - 1;
  return Status::OK();
}