Function GetCOOIndexTensorRow

cpp/src/arrow/sparse_tensor.cc:144–182

}

void GetCOOIndexTensorRow(const std::shared_ptr<Tensor>& coords, const int64_t row,
                          std::vector<int64_t>* out_index) {
  const auto& fw_index_value_type =
      internal::checked_cast<const FixedWidthType&>(*coords->type());
  const size_t indices_elsize = fw_index_value_type.bit_width() / CHAR_BIT;

  const auto& shape = coords->shape();
  const int64_t non_zero_length = shape[0];
  DCHECK(0 <= row && row < non_zero_length);

  const int64_t ndim = shape[1];
  out_index->resize(ndim);

  switch (indices_elsize) {
    case 1:  // Int8, UInt8
      for (int64_t i = 0; i < ndim; ++i) {
        (*out_index)[i] = static_cast<int64_t>(coords->Value<UInt8Type>({row, i}));
      }
      break;
    case 2:  // Int16, UInt16
      for (int64_t i = 0; i < ndim; ++i) {
        (*out_index)[i] = static_cast<int64_t>(coords->Value<UInt16Type>({row, i}));
      }
      break;
    case 4:  // Int32, UInt32
      for (int64_t i = 0; i < ndim; ++i) {
        (*out_index)[i] = static_cast<int64_t>(coords->Value<UInt32Type>({row, i}));
      }
      break;
    case 8:  // Int64
      for (int64_t i = 0; i < ndim; ++i) {
        (*out_index)[i] = coords->Value<Int64Type>({row, i});
      }
      break;
    default:
      DCHECK(false) << "Must not reach here";
      break;
  }
}

bool DetectSparseCOOIndexCanonicality(const std::shared_ptr<Tensor>& coords) {
  DCHECK_EQ(coords->ndim(), 2);