| 68 | : tensor_(tensor), index_value_type_(index_value_type), pool_(pool) {} |
| 69 | |
| 70 | Status Convert() { |
| 71 | RETURN_NOT_OK(::arrow::internal::CheckSparseIndexMaximumValue(index_value_type_, |
| 72 | tensor_.shape())); |
| 73 | |
| 74 | const int index_elsize = index_value_type_->byte_width(); |
| 75 | const int value_elsize = tensor_.type()->byte_width(); |
| 76 | |
| 77 | const int64_t ndim = tensor_.ndim(); |
| 78 | // Axis order as ascending order of dimension size is a good heuristic but is not |
| 79 | // necessarily optimal. |
| 80 | std::vector<int64_t> axis_order = internal::ArgSort(tensor_.shape()); |
| 81 | ARROW_ASSIGN_OR_RAISE(int64_t nonzero_count, tensor_.CountNonZero()); |
| 82 | |
| 83 | ARROW_ASSIGN_OR_RAISE(auto values_buffer, |
| 84 | AllocateBuffer(value_elsize * nonzero_count, pool_)); |
| 85 | auto* values = values_buffer->mutable_data(); |
| 86 | |
| 87 | std::vector<int64_t> counts(ndim, 0); |
| 88 | std::vector<int64_t> coord(ndim, 0); |
| 89 | std::vector<int64_t> previous_coord(ndim, -1); |
| 90 | std::vector<BufferBuilder> indptr_buffer_builders(ndim - 1); |
| 91 | std::vector<BufferBuilder> indices_buffer_builders(ndim); |
| 92 | |
| 93 | const auto* tensor_data = tensor_.raw_data(); |
| 94 | uint8_t index_buffer[sizeof(int64_t)]; |
| 95 | |
| 96 | if (ndim <= 1) { |
| 97 | return Status::NotImplemented("TODO for ndim <= 1"); |
| 98 | } else { |
| 99 | const auto& shape = tensor_.shape(); |
| 100 | for (int64_t n = tensor_.size(); n > 0; n--) { |
| 101 | const auto offset = tensor_.CalculateValueOffset(coord); |
| 102 | const auto xp = tensor_data + offset; |
| 103 | |
| 104 | if (std::any_of(xp, xp + value_elsize, IsNonZero)) { |
| 105 | bool tree_split = false; |
| 106 | |
| 107 | std::copy_n(xp, value_elsize, values); |
| 108 | values += value_elsize; |
| 109 | |
| 110 | for (int64_t i = 0; i < ndim; ++i) { |
| 111 | int64_t dimension = axis_order[i]; |
| 112 | |
| 113 | tree_split = tree_split || (coord[dimension] != previous_coord[dimension]); |
| 114 | if (tree_split) { |
| 115 | if (i < ndim - 1) { |
| 116 | AssignIndex(index_buffer, counts[i + 1], index_elsize); |
| 117 | RETURN_NOT_OK( |
| 118 | indptr_buffer_builders[i].Append(index_buffer, index_elsize)); |
| 119 | } |
| 120 | |
| 121 | AssignIndex(index_buffer, coord[dimension], index_elsize); |
| 122 | RETURN_NOT_OK( |
| 123 | indices_buffer_builders[i].Append(index_buffer, index_elsize)); |
| 124 | |
| 125 | ++counts[i]; |
| 126 | } |
| 127 | } |
no test coverage detected