Function MakeTensorFromSparseCSXMatrix

cpp/src/arrow/tensor/csx_converter.cc:162–214

namespace {

Result<std::shared_ptr<Tensor>> MakeTensorFromSparseCSXMatrix(
    SparseMatrixCompressedAxis axis, MemoryPool* pool,
    const std::shared_ptr<Tensor>& indptr, const std::shared_ptr<Tensor>& indices,
    const int64_t non_zero_length, const std::shared_ptr<DataType>& value_type,
    const std::vector<int64_t>& shape, const int64_t tensor_size, const uint8_t* raw_data,
    const std::vector<std::string>& dim_names) {
  const auto* indptr_data = indptr->raw_data();
  const auto* indices_data = indices->raw_data();

  const int indptr_elsize = indptr->type()->byte_width();
  const int indices_elsize = indices->type()->byte_width();

  const auto& fw_value_type = checked_cast<const FixedWidthType&>(*value_type);
  const int value_elsize = fw_value_type.byte_width();
  ARROW_ASSIGN_OR_RAISE(auto values_buffer,
                        AllocateBuffer(value_elsize * tensor_size, pool));
  auto values = values_buffer->mutable_data();
  std::fill_n(values, value_elsize * tensor_size, 0);

  std::vector<int64_t> strides;
  RETURN_NOT_OK(ComputeRowMajorStrides(fw_value_type, shape, &strides));

  const auto nc = shape[1];

  int64_t offset = 0;
  for (int64_t i = 0; i < indptr->size() - 1; ++i) {
    const auto start =
        SparseTensorConverterMixin::GetIndexValue(indptr_data, indptr_elsize);
    const auto stop = SparseTensorConverterMixin::GetIndexValue(
        indptr_data + indptr_elsize, indptr_elsize);

    for (int64_t j = start; j < stop; ++j) {
      const auto index = SparseTensorConverterMixin::GetIndexValue(
          indices_data + j * indices_elsize, indices_elsize);
      switch (axis) {
        case SparseMatrixCompressedAxis::ROW:
          offset = (index + i * nc) * value_elsize;
          break;
        case SparseMatrixCompressedAxis::COLUMN:
          offset = (i + index * nc) * value_elsize;
          break;
      }

      std::copy_n(raw_data, value_elsize, values + offset);
      raw_data += value_elsize;
    }

    indptr_data += indptr_elsize;
  }

  return std::make_shared<Tensor>(value_type, std::move(values_buffer), shape, strides,
                                  dim_names);
}

}  // namespace