Partial review feedback implementation.

cpp/src/arrow/python/numpy_convert.cc

@@ -281,18 +281,19 @@ Status SparseTensorCOOToNdarray(const std::shared_ptr<SparseTensorCOO>& sparse_t
   PyAcquireGIL lock;
 
   int type_num_data;
-  int type_num_coords;
+  RETURN_NOT_OK(GetNumPyType(*sparse_tensor->type(), &type_num_data));
+  PyArray_Descr* dtype_data = PyArray_DescrNewFromType(type_num_data);
+  RETURN_IF_PYERROR();
+
   const auto& sparse_index = arrow::internal::checked_cast<const SparseCOOIndex&>(
       *sparse_tensor->sparse_index());
   const std::shared_ptr<arrow::NumericTensor<arrow::Int64Type>> sparse_index_coords =
       sparse_index.indices();
 
-  RETURN_NOT_OK(GetNumPyType(*sparse_tensor->type(), &type_num_data));
-  PyArray_Descr* dtype_data = PyArray_DescrNewFromType(type_num_data);
-  RETURN_NOT_OK(GetNumPyType(*sparse_index_coords->type(), &type_num_coords));
-  PyArray_Descr* dtype_coords = PyArray_DescrNewFromType(type_num_coords);
   RETURN_IF_PYERROR();
 
+  std::vector<npy_intp> npy_shape_data({sparse_index.non_zero_length(), 1});
+
   const int ndim_coords = sparse_tensor->ndim();
   std::vector<npy_intp> npy_shape_coords(ndim_coords);
 
@@ -300,17 +301,8 @@ Status SparseTensorCOOToNdarray(const std::shared_ptr<SparseTensorCOO>& sparse_t
     npy_shape_coords[i] = sparse_index_coords->shape()[i];
   }
 
-  std::vector<npy_intp> npy_shape_data({sparse_index.non_zero_length(), 1});
-
-  const void* immutable_data = nullptr;
-  if (sparse_tensor->data()) {
-    immutable_data = sparse_tensor->data()->data();
-  }
-
-  const void* immutable_coords = nullptr;
-  if (sparse_index_coords->data()) {
-    immutable_coords = sparse_index_coords->data()->data();
-  }
+  const void* immutable_data = sparse_tensor->data()->data();
+  const void* immutable_coords = sparse_index_coords->data()->data();
 
   // Remove const =(
   void* mutable_data = const_cast<void*>(immutable_data);
@@ -321,20 +313,21 @@ Status SparseTensorCOOToNdarray(const std::shared_ptr<SparseTensorCOO>& sparse_t
     array_flags |= NPY_ARRAY_WRITEABLE;
   }
 
-  PyObject* result_data =
-      PyArray_NewFromDescr(&PyArray_Type, dtype_data, 1, npy_shape_data.data(), nullptr,
-                           mutable_data, array_flags, nullptr);
-  PyObject* result_coords =
-      PyArray_NewFromDescr(&PyArray_Type, dtype_coords, 2, npy_shape_coords.data(),
-                           nullptr, mutable_coords, array_flags, nullptr);
+  int type_num_coords;
+  RETURN_NOT_OK(GetNumPyType(*sparse_index_coords->type(), &type_num_coords));
+  PyArray_Descr* dtype_coords = PyArray_DescrNewFromType(type_num_coords);
 
+  PyObject* result_data =
+      PyArray_NewFromDescr(&PyArray_Type, dtype_data, 2, npy_shape_data.data(), nullptr,
+                           mutable_data, array_flags, nullptr);
+  PyObject* result_coords = PyArray_NewFromDescr(&PyArray_Type, dtype_coords, ndim_coords,
+                                                 npy_shape_coords.data(), nullptr,
+                                                 mutable_coords, array_flags, nullptr);
   RETURN_IF_PYERROR()
 
-  if (base == Py_None || base == nullptr) {
-    base = py::wrap_sparse_tensor_coo(sparse_tensor);
-  } else {
-    Py_XINCREF(base);
-  }
+  Py_XINCREF(base);
+  Py_XINCREF(base);
+
   PyArray_SetBaseObject(reinterpret_cast<PyArrayObject*>(result_data), base);
   PyArray_SetBaseObject(reinterpret_cast<PyArrayObject*>(result_coords), base);
 
@@ -376,20 +369,9 @@ Status SparseTensorCSRToNdarray(const std::shared_ptr<SparseTensorCSR>& sparse_t
     npy_shape_indices[i] = sparse_index_indices->shape()[i];
   }
 
-  const void* immutable_data = nullptr;
-  if (sparse_tensor->data()) {
-    immutable_data = sparse_tensor->data()->data();
-  }
-
-  const void* immutable_indptr = nullptr;
-  if (sparse_index_indptr->data()) {
-    immutable_indptr = sparse_index_indptr->data()->data();
-  }
-
-  const void* immutable_indices = nullptr;
-  if (sparse_index_indices->data()) {
-    immutable_indices = sparse_index_indices->data()->data();
-  }
+  const void* immutable_data = sparse_tensor->data()->data();
+  const void* immutable_indptr = sparse_index_indptr->data()->data();
+  const void* immutable_indices = sparse_index_indices->data()->data();
 
   // Remove const =(
   void* mutable_data = const_cast<void*>(immutable_data);
@@ -410,21 +392,20 @@ Status SparseTensorCSRToNdarray(const std::shared_ptr<SparseTensorCSR>& sparse_t
   PyArray_Descr* dtype_indices = PyArray_DescrNewFromType(type_num_indices);
 
   PyObject* result_data =
-      PyArray_NewFromDescr(&PyArray_Type, dtype_data, 1, npy_shape_data.data(), nullptr,
+      PyArray_NewFromDescr(&PyArray_Type, dtype_data, 2, npy_shape_data.data(), nullptr,
                            mutable_data, array_flags, nullptr);
-  PyObject* result_indptr =
-      PyArray_NewFromDescr(&PyArray_Type, dtype_indptr, 1, npy_shape_indptr.data(),
-                           nullptr, mutable_indptr, array_flags, nullptr);
-  PyObject* result_indices =
-      PyArray_NewFromDescr(&PyArray_Type, dtype_indices, 1, npy_shape_indices.data(),
-                           nullptr, mutable_indices, array_flags, nullptr);
+  PyObject* result_indptr = PyArray_NewFromDescr(&PyArray_Type, dtype_indptr, ndim_indptr,
+                                                 npy_shape_indptr.data(), nullptr,
+                                                 mutable_indptr, array_flags, nullptr);
+  PyObject* result_indices = PyArray_NewFromDescr(
+      &PyArray_Type, dtype_indices, ndim_indices, npy_shape_indices.data(), nullptr,
+      mutable_indices, array_flags, nullptr);
   RETURN_IF_PYERROR()
 
-  if (base == Py_None || base == nullptr) {
-    base = py::wrap_sparse_tensor_csr(sparse_tensor);
-  } else {
-    Py_XINCREF(base);
-  }
+  Py_XINCREF(base);
+  Py_XINCREF(base);
+  Py_XINCREF(base);
+
   PyArray_SetBaseObject(reinterpret_cast<PyArrayObject*>(result_data), base);
   PyArray_SetBaseObject(reinterpret_cast<PyArrayObject*>(result_indptr), base);
   PyArray_SetBaseObject(reinterpret_cast<PyArrayObject*>(result_indices), base);
@@ -435,19 +416,17 @@ Status SparseTensorCSRToNdarray(const std::shared_ptr<SparseTensorCSR>& sparse_t
   return Status::OK();
 }
 
-Status NdarrayToSparseTensorCOO(MemoryPool* pool, PyObject* ao,
-                                std::shared_ptr<SparseTensorCOO>* out) {
+Status DenseNdarrayToSparseTensorCOO(MemoryPool* pool, PyObject* ao,
+                                     const std::vector<std::string>& dim_names,
+                                     std::shared_ptr<SparseTensorCOO>* out) {
   if (!PyArray_Check(ao)) {
     return Status::TypeError("Did not pass ndarray object");
   }
 
   PyArrayObject* ndarray = reinterpret_cast<PyArrayObject*>(ao);
 
-  // TODO(wesm): What do we want to do with non-contiguous memory and negative strides?
-
   int ndim = PyArray_NDIM(ndarray);
 
-  // This is also holding the GIL, so don't already draw it.
   std::shared_ptr<Buffer> data = std::make_shared<NumPyBuffer>(ao);
   std::vector<int64_t> shape(ndim);
   std::vector<int64_t> strides(ndim);
@@ -467,25 +446,23 @@ Status NdarrayToSparseTensorCOO(MemoryPool* pool, PyObject* ao,
     std::shared_ptr<DataType> type;
     RETURN_NOT_OK(
         GetTensorType(reinterpret_cast<PyObject*>(PyArray_DESCR(ndarray)), &type));
-    Tensor tensor(type, data, shape, strides);
+    Tensor tensor(type, data, shape, strides, dim_names);
     *out = std::make_shared<SparseTensorCOO>(tensor);
     return Status::OK();
   }
 }
 
-Status NdarrayToSparseTensorCSR(MemoryPool* pool, PyObject* ao,
-                                std::shared_ptr<SparseTensorCSR>* out) {
+Status DenseNdarrayToSparseTensorCSR(MemoryPool* pool, PyObject* ao,
+                                     const std::vector<std::string>& dim_names,
+                                     std::shared_ptr<SparseTensorCSR>* out) {
   if (!PyArray_Check(ao)) {
     return Status::TypeError("Did not pass ndarray object");
   }
 
   PyArrayObject* ndarray = reinterpret_cast<PyArrayObject*>(ao);
 
-  // TODO(wesm): What do we want to do with non-contiguous memory and negative strides?
-
   int ndim = PyArray_NDIM(ndarray);
 
-  // This is also holding the GIL, so don't already draw it.
   std::shared_ptr<Buffer> data = std::make_shared<NumPyBuffer>(ao);
   std::vector<int64_t> shape(ndim);
   std::vector<int64_t> strides(ndim);
@@ -505,7 +482,7 @@ Status NdarrayToSparseTensorCSR(MemoryPool* pool, PyObject* ao,
     std::shared_ptr<DataType> type;
     RETURN_NOT_OK(
         GetTensorType(reinterpret_cast<PyObject*>(PyArray_DESCR(ndarray)), &type));
-    Tensor tensor(type, data, shape, strides);
+    Tensor tensor(type, data, shape, strides, dim_names);
     *out = std::make_shared<SparseTensorCSR>(tensor);
     return Status::OK();
   }
@@ -529,22 +506,34 @@ Status NdarraysToSparseTensorCOO(MemoryPool* pool, PyObject* data_ao, PyObject*
 
   int coords_ndim = PyArray_NDIM(ndarray_coords);
 
-  std::shared_ptr<DataType> type;
+  std::shared_ptr<DataType> type_data;
+  std::shared_ptr<DataType> type_coords;
+  RETURN_NOT_OK(GetTensorType(reinterpret_cast<PyObject*>(PyArray_DESCR(ndarray_data)),
+                              &type_data));
+  RETURN_NOT_OK(GetTensorType(reinterpret_cast<PyObject*>(PyArray_DESCR(ndarray_coords)),
+                              &type_coords));
+  ARROW_CHECK_EQ(type_coords->id(), Type::INT64);
+
+  const int64_t i64_size = sizeof(int64_t);
   npy_intp* coords_array_shape = PyArray_SHAPE(ndarray_coords);
   std::vector<int64_t> coords_shape(coords_ndim);
+  std::vector<int64_t> coords_strides(coords_ndim);
 
   for (int i = 0; i < coords_ndim; ++i) {
     coords_shape[i] = coords_array_shape[i];
+    if (i == 0) {
+      coords_strides[i] = i64_size;
+    } else {
+      coords_strides[i] = coords_strides[i - 1] * coords_array_shape[i - 1];
+    }
   }
 
   std::shared_ptr<NumericTensor<Int64Type>> coords =
-      std::make_shared<NumericTensor<Int64Type>>(coords_buffer, coords_shape);
-
-  RETURN_NOT_OK(
-      GetTensorType(reinterpret_cast<PyObject*>(PyArray_DESCR(ndarray_data)), &type));
+      std::make_shared<NumericTensor<Int64Type>>(coords_buffer, coords_shape,
+                                                 coords_strides);
 
   std::shared_ptr<SparseCOOIndex> sparse_index = std::make_shared<SparseCOOIndex>(coords);
-  *out = std::make_shared<SparseTensorImpl<SparseCOOIndex>>(sparse_index, type, data,
+  *out = std::make_shared<SparseTensorImpl<SparseCOOIndex>>(sparse_index, type_data, data,
                                                             shape, dim_names);
   return Status::OK();
 }
@@ -564,7 +553,6 @@ Status NdarraysToSparseTensorCSR(MemoryPool* pool, PyObject* data_ao, PyObject*
   PyArrayObject* ndarray_indptr = reinterpret_cast<PyArrayObject*>(indptr_ao);
   PyArrayObject* ndarray_indices = reinterpret_cast<PyArrayObject*>(indices_ao);
 
-  // This is also holding the GIL, so don't already draw it.
   std::shared_ptr<Buffer> data = std::make_shared<NumPyBuffer>(data_ao);
   std::shared_ptr<Buffer> indptr_buffer = std::make_shared<NumPyBuffer>(indptr_ao);
   std::shared_ptr<Buffer> indices_buffer = std::make_shared<NumPyBuffer>(indices_ao);
@@ -572,7 +560,9 @@ Status NdarraysToSparseTensorCSR(MemoryPool* pool, PyObject* data_ao, PyObject*
   int indptr_ndim = PyArray_NDIM(ndarray_indptr);
   int indices_ndim = PyArray_NDIM(ndarray_indices);
 
-  std::shared_ptr<DataType> type;
+  std::shared_ptr<DataType> type_data;
+  std::shared_ptr<DataType> type_indptr;
+  std::shared_ptr<DataType> type_indices;
   npy_intp* indptr_array_shape = PyArray_SHAPE(ndarray_indptr);
   npy_intp* indices_array_shape = PyArray_SHAPE(ndarray_indices);
   std::vector<int64_t> indptr_shape(indptr_ndim);
@@ -585,17 +575,23 @@ Status NdarraysToSparseTensorCSR(MemoryPool* pool, PyObject* data_ao, PyObject*
     indices_shape[i] = indices_array_shape[i];
   }
 
+  RETURN_NOT_OK(GetTensorType(reinterpret_cast<PyObject*>(PyArray_DESCR(ndarray_data)),
+                              &type_data));
+  RETURN_NOT_OK(GetTensorType(reinterpret_cast<PyObject*>(PyArray_DESCR(ndarray_indptr)),
+                              &type_indptr));
+  RETURN_NOT_OK(GetTensorType(reinterpret_cast<PyObject*>(PyArray_DESCR(ndarray_indices)),
+                              &type_indices));
+  ARROW_CHECK_EQ(type_indptr->id(), Type::INT64);
+  ARROW_CHECK_EQ(type_indices->id(), Type::INT64);
+
   std::shared_ptr<SparseCSRIndex::IndexTensor> indptr =
       std::make_shared<SparseCSRIndex::IndexTensor>(indptr_buffer, indptr_shape);
   std::shared_ptr<SparseCSRIndex::IndexTensor> indices =
       std::make_shared<SparseCSRIndex::IndexTensor>(indices_buffer, indices_shape);
 
-  RETURN_NOT_OK(
-      GetTensorType(reinterpret_cast<PyObject*>(PyArray_DESCR(ndarray_data)), &type));
-
   std::shared_ptr<SparseCSRIndex> sparse_index =
       std::make_shared<SparseCSRIndex>(indptr, indices);
-  *out = std::make_shared<SparseTensorImpl<SparseCSRIndex>>(sparse_index, type, data,
+  *out = std::make_shared<SparseTensorImpl<SparseCSRIndex>>(sparse_index, type_data, data,
                                                             shape, dim_names);
   return Status::OK();
 }

cpp/src/arrow/python/numpy_convert.h

@@ -70,23 +70,20 @@ ARROW_PYTHON_EXPORT Status NdarrayToTensor(MemoryPool* pool, PyObject* ao,
 ARROW_PYTHON_EXPORT Status TensorToNdarray(const std::shared_ptr<Tensor>& tensor,
                                            PyObject* base, PyObject** out);
 
-ARROW_PYTHON_EXPORT Status SparseTensorCSRToNdarray(
-    const std::shared_ptr<SparseTensorCSR>& sparse_tensor, PyObject* base,
-    PyObject** out_data, PyObject** out_indptr, PyObject** out_indices);
-
 ARROW_PYTHON_EXPORT Status
 SparseTensorCOOToNdarray(const std::shared_ptr<SparseTensorCOO>& sparse_tensor,
                          PyObject* base, PyObject** out_data, PyObject** out_coords);
 
-ARROW_PYTHON_EXPORT Status NdarrayToSparseTensorCSR(
-    MemoryPool* pool, PyObject* ao, std::shared_ptr<SparseTensorCSR>* out);
+ARROW_PYTHON_EXPORT Status SparseTensorCSRToNdarray(
+    const std::shared_ptr<SparseTensorCSR>& sparse_tensor, PyObject* base,
+    PyObject** out_data, PyObject** out_indptr, PyObject** out_indices);
 
-ARROW_PYTHON_EXPORT Status NdarrayToSparseTensorCOO(
-    MemoryPool* pool, PyObject* ao, std::shared_ptr<SparseTensorCOO>* out);
+ARROW_PYTHON_EXPORT Status DenseNdarrayToSparseTensorCOO(
+    MemoryPool* pool, PyObject* ao, const std::vector<std::string>& dim_names,
+    std::shared_ptr<SparseTensorCOO>* out);
 
-ARROW_PYTHON_EXPORT Status NdarraysToSparseTensorCSR(
-    MemoryPool* pool, PyObject* data_ao, PyObject* indptr_ao, PyObject* indices_ao,
-    const std::vector<int64_t>& shape, const std::vector<std::string>& dim_names,
+ARROW_PYTHON_EXPORT Status DenseNdarrayToSparseTensorCSR(
+    MemoryPool* pool, PyObject* ao, const std::vector<std::string>& dim_names,
     std::shared_ptr<SparseTensorCSR>* out);
 
 ARROW_PYTHON_EXPORT Status NdarraysToSparseTensorCOO(
@@ -94,14 +91,19 @@ ARROW_PYTHON_EXPORT Status NdarraysToSparseTensorCOO(
     const std::vector<int64_t>& shape, const std::vector<std::string>& dim_names,
     std::shared_ptr<SparseTensorCOO>* out);
 
-ARROW_PYTHON_EXPORT Status
-TensorToSparseTensorCSR(const std::shared_ptr<Tensor>& tensor,
-                        std::shared_ptr<SparseTensorCSR>* csparse_tensor);
+ARROW_PYTHON_EXPORT Status NdarraysToSparseTensorCSR(
+    MemoryPool* pool, PyObject* data_ao, PyObject* indptr_ao, PyObject* indices_ao,
+    const std::vector<int64_t>& shape, const std::vector<std::string>& dim_names,
+    std::shared_ptr<SparseTensorCSR>* out);
 
 ARROW_PYTHON_EXPORT Status
 TensorToSparseTensorCOO(const std::shared_ptr<Tensor>& tensor,
                         std::shared_ptr<SparseTensorCOO>* csparse_tensor);
 
+ARROW_PYTHON_EXPORT Status
+TensorToSparseTensorCSR(const std::shared_ptr<Tensor>& tensor,
+                        std::shared_ptr<SparseTensorCSR>* csparse_tensor);
+
 }  // namespace py
 }  // namespace arrow
 

cpp/src/arrow/python/pyarrow_lib.h

@@ -61,18 +61,6 @@ __PYX_EXTERN_C std::shared_ptr< arrow::Table>  __pyx_f_7pyarrow_3lib_pyarrow_unw
 __PYX_EXTERN_C std::shared_ptr< arrow::Tensor>  __pyx_f_7pyarrow_3lib_pyarrow_unwrap_tensor(PyObject *);
 __PYX_EXTERN_C std::shared_ptr< arrow::SparseTensorCOO>  __pyx_f_7pyarrow_3lib_pyarrow_unwrap_sparse_tensor_coo(PyObject *);
 __PYX_EXTERN_C std::shared_ptr< arrow::SparseTensorCSR>  __pyx_f_7pyarrow_3lib_pyarrow_unwrap_sparse_tensor_csr(PyObject *);
-__PYX_EXTERN_C int pyarrow_is_buffer(PyObject *);
-__PYX_EXTERN_C int pyarrow_is_data_type(PyObject *);
-__PYX_EXTERN_C int pyarrow_is_field(PyObject *);
-__PYX_EXTERN_C int pyarrow_is_schema(PyObject *);
-__PYX_EXTERN_C int pyarrow_is_array(PyObject *);
-__PYX_EXTERN_C PyObject *pyarrow_wrap_chunked_array(std::shared_ptr< arrow::ChunkedArray>  const &);
-__PYX_EXTERN_C int pyarrow_is_tensor(PyObject *);
-__PYX_EXTERN_C int pyarrow_is_sparse_tensor_coo(PyObject *);
-__PYX_EXTERN_C int pyarrow_is_sparse_tensor_csr(PyObject *);
-__PYX_EXTERN_C int pyarrow_is_column(PyObject *);
-__PYX_EXTERN_C int pyarrow_is_table(PyObject *);
-__PYX_EXTERN_C int pyarrow_is_batch(PyObject *);
 
 #endif /* !__PYX_HAVE_API__pyarrow__lib */
 

python/pyarrow/__init__.pxd

@@ -22,7 +22,7 @@ from pyarrow.includes.libarrow cimport (CArray, CBuffer, CColumn, CDataType,
                                         CField, CRecordBatch, CSchema,
                                         CTable, CTensor,
                                         CSparseTensorCSR, CSparseTensorCOO)
-
+from pyarrow.compat import frombytes
 
 cdef extern from "arrow/python/pyarrow.h" namespace "arrow::py":
     cdef int import_pyarrow() except -1

python/pyarrow/array.pxi

@@ -891,15 +891,27 @@ type: {0.type}
 shape: {0.shape}""".format(self)
 
     @staticmethod
-    def from_dense_numpy(obj):
+    def from_dense_numpy(obj, dim_names=None):
+        """
+        Convert numpy.ndarray to arrow::SparseTensorCOO
+        """
         cdef shared_ptr[CSparseTensorCOO] csparse_tensor
-        with nogil:
-            check_status(NdarrayToSparseTensorCOO(c_default_memory_pool(), obj,
-                                                  &csparse_tensor))
+        cdef vector[c_string] c_dim_names
+
+        if dim_names is not None:
+            for x in dim_names:
+                c_dim_names.push_back(tobytes(x))
+
+        check_status(DenseNdarrayToSparseTensorCOO(c_default_memory_pool(),
+                                                   obj, c_dim_names,
+                                                   &csparse_tensor))
         return pyarrow_wrap_sparse_tensor_coo(csparse_tensor)
 
     @staticmethod
     def from_numpy(data, coords, shape, dim_names=None):
+        """
+        Create arrow::SparseTensorCOO from numpy.ndarrays
+        """
         cdef shared_ptr[CSparseTensorCOO] csparse_tensor
         cdef vector[int64_t] c_shape
         cdef vector[c_string] c_dim_names
@@ -909,18 +921,23 @@ shape: {0.shape}""".format(self)
         if dim_names is not None:
             for x in dim_names:
                 c_dim_names.push_back(tobytes(x))
-        with nogil:
-            check_status(NdarraysToSparseTensorCOO(c_default_memory_pool(),
-                         data, coords, c_shape, c_dim_names, &csparse_tensor))
+
+        if coords.dtype != 'i8':
+            coords = coords.astype('i8')
+
+        check_status(NdarraysToSparseTensorCOO(c_default_memory_pool(),
+                     data, coords, c_shape, c_dim_names, &csparse_tensor))
         return pyarrow_wrap_sparse_tensor_coo(csparse_tensor)
 
     @staticmethod
     def from_tensor(obj):
+        """
+        Convert arrow::Tensor to arrow::SparseTensorCOO
+        """
         cdef shared_ptr[CSparseTensorCOO] csparse_tensor
         cdef shared_ptr[CTensor] ctensor = pyarrow_unwrap_tensor(obj)
 
-        with nogil:
-            check_status(TensorToSparseTensorCOO(ctensor, &csparse_tensor))
+        check_status(TensorToSparseTensorCOO(ctensor, &csparse_tensor))
 
         return pyarrow_wrap_sparse_tensor_coo(csparse_tensor)
 
@@ -931,9 +948,8 @@ shape: {0.shape}""".format(self)
         cdef PyObject* out_data
         cdef PyObject* out_coords
 
-        with nogil:
-            check_status(SparseTensorCOOToNdarray(self.sp_sparse_tensor, self,
-                                                  &out_data, &out_coords))
+        check_status(SparseTensorCOOToNdarray(self.sp_sparse_tensor, self,
+                                              &out_data, &out_coords))
         return PyObject_to_object(out_data), PyObject_to_object(out_coords)
 
     def equals(self, SparseTensorCOO other):
@@ -966,11 +982,11 @@ shape: {0.shape}""".format(self)
         return self.stp.size()
 
     def dim_name(self, i):
-        return self.stp.dim_name(i)
+        return frombytes(self.stp.dim_name(i))
 
     @property
     def dim_names(self):
-        return self.stp.dim_names()
+        return [frombytes(x) for x in tuple(self.stp.dim_names())]
 
     @property
     def non_zero_length(self):
@@ -998,15 +1014,27 @@ type: {0.type}
 shape: {0.shape}""".format(self)
 
     @staticmethod
-    def from_dense_numpy(obj):
+    def from_dense_numpy(obj, dim_names=None):
+        """
+        Convert numpy.ndarray to arrow::SparseTensorCSR
+        """
         cdef shared_ptr[CSparseTensorCSR] csparse_tensor
-        with nogil:
-            check_status(NdarrayToSparseTensorCSR(c_default_memory_pool(), obj,
-                                                  &csparse_tensor))
+        cdef vector[c_string] c_dim_names
+
+        if dim_names is not None:
+            for x in dim_names:
+                c_dim_names.push_back(tobytes(x))
+
+        check_status(DenseNdarrayToSparseTensorCSR(c_default_memory_pool(),
+                                                   obj, c_dim_names,
+                                                   &csparse_tensor))
         return pyarrow_wrap_sparse_tensor_csr(csparse_tensor)
 
     @staticmethod
     def from_numpy(data, indptr, indices, shape, dim_names=None):
+        """
+        Create arrow::SparseTensorCSR from numpy.ndarrays
+        """
         cdef shared_ptr[CSparseTensorCSR] csparse_tensor
         cdef vector[int64_t] c_shape
         cdef vector[c_string] c_dim_names
@@ -1016,19 +1044,26 @@ shape: {0.shape}""".format(self)
         if dim_names is not None:
             for x in dim_names:
                 c_dim_names.push_back(tobytes(x))
-        with nogil:
-            check_status(NdarraysToSparseTensorCSR(c_default_memory_pool(),
-                         data, indptr, indices, c_shape, c_dim_names,
-                         &csparse_tensor))
+
+        if indptr.dtype != 'i8':
+            indptr = indptr.astype('i8')
+        if indices.dtype != 'i8':
+            indices = indices.astype('i8')
+
+        check_status(NdarraysToSparseTensorCSR(c_default_memory_pool(),
+                     data, indptr, indices, c_shape, c_dim_names,
+                     &csparse_tensor))
         return pyarrow_wrap_sparse_tensor_csr(csparse_tensor)
 
     @staticmethod
     def from_tensor(obj):
+        """
+        Convert arrow::Tensor to arrow::SparseTensorCSR
+        """
         cdef shared_ptr[CSparseTensorCSR] csparse_tensor
         cdef shared_ptr[CTensor] ctensor = pyarrow_unwrap_tensor(obj)
 
-        with nogil:
-            check_status(TensorToSparseTensorCSR(ctensor, &csparse_tensor))
+        check_status(TensorToSparseTensorCSR(ctensor, &csparse_tensor))
 
         return pyarrow_wrap_sparse_tensor_csr(csparse_tensor)
 
@@ -1040,9 +1075,8 @@ shape: {0.shape}""".format(self)
         cdef PyObject* out_indptr
         cdef PyObject* out_indices
 
-        with nogil:
-            check_status(SparseTensorCSRToNdarray(self.sp_sparse_tensor, self,
-                         &out_data, &out_indptr, &out_indices))
+        check_status(SparseTensorCSRToNdarray(self.sp_sparse_tensor, self,
+                     &out_data, &out_indptr, &out_indices))
         return (PyObject_to_object(out_data), PyObject_to_object(out_indptr),
                 PyObject_to_object(out_indices))
 
@@ -1076,11 +1110,11 @@ shape: {0.shape}""".format(self)
         return self.stp.size()
 
     def dim_name(self, i):
-        return self.stp.dim_name(i)
+        return frombytes(self.stp.dim_name(i))
 
     @property
     def dim_names(self):
-        return self.stp.dim_names()
+        return [frombytes(x) for x in tuple(self.stp.dim_names())]
 
     @property
     def non_zero_length(self):

python/pyarrow/includes/libarrow.pxd

@@ -1259,11 +1259,13 @@ cdef extern from "arrow/python/api.h" namespace "arrow::py" nogil:
                                       const vector[c_string]& dim_names,
                                       shared_ptr[CSparseTensorCSR]* out)
 
-    CStatus NdarrayToSparseTensorCOO(CMemoryPool* pool, object ao,
-                                     shared_ptr[CSparseTensorCOO]* out)
+    CStatus DenseNdarrayToSparseTensorCOO(CMemoryPool* pool, object ao,
+                                          const vector[c_string]& dim_names,
+                                          shared_ptr[CSparseTensorCOO]* out)
 
-    CStatus NdarrayToSparseTensorCSR(CMemoryPool* pool, object ao,
-                                     shared_ptr[CSparseTensorCSR]* out)
+    CStatus DenseNdarrayToSparseTensorCSR(CMemoryPool* pool, object ao,
+                                          const vector[c_string]& dim_names,
+                                          shared_ptr[CSparseTensorCSR]* out)
 
     CStatus TensorToSparseTensorCOO(shared_ptr[CTensor],
                                     shared_ptr[CSparseTensorCOO]* out)

python/pyarrow/tests/test_sparse_tensor.py

@@ -34,29 +34,52 @@ def test_sparse_tensor_attrs(sparse_tensor_type):
         [0, 0, 0, 0, 0],
         [0, 3, 0, 0, 0],
     ])
-    sparse_tensor = sparse_tensor_type.from_dense_numpy(data)
+    dim_names = ['x', 'y']
+    sparse_tensor = sparse_tensor_type.from_dense_numpy(data, dim_names)
 
     assert sparse_tensor.ndim == 2
     assert sparse_tensor.size == 25
     assert sparse_tensor.shape == data.shape
     assert sparse_tensor.is_mutable
-    assert sparse_tensor.dim_name(0) == b''
-    assert sparse_tensor.dim_names == []
+    assert sparse_tensor.dim_name(0) == dim_names[0]
+    assert sparse_tensor.dim_names == dim_names
     assert sparse_tensor.non_zero_length == 4
 
 
 def test_sparse_tensor_coo_base_object():
-    sparse_tensor = pa.SparseTensorCOO.from_dense_numpy(np.random.randn(10, 4))
+    data = np.array([[4], [9], [7], [5]])
+    coords = np.array([[0, 0], [1, 3], [0, 2], [1, 3]])
+    array = np.array([[4, 0, 9, 0],
+                      [0, 7, 0, 0],
+                      [0, 0, 0, 0],
+                      [0, 0, 0, 5]])
+    sparse_tensor = pa.SparseTensorCOO.from_dense_numpy(array)
     n = sys.getrefcount(sparse_tensor)
-    data, coords = sparse_tensor.to_numpy()  # noqa
-    assert sys.getrefcount(sparse_tensor) == n + 1
+    result_data, result_coords = sparse_tensor.to_numpy()
+    assert sys.getrefcount(sparse_tensor) == n + 2
+
+    sparse_tensor = None
+    assert np.array_equal(data, result_data)
+    assert np.array_equal(coords, result_coords)
 
 
 def test_sparse_tensor_csr_base_object():
-    sparse_tensor = pa.SparseTensorCSR.from_dense_numpy(np.random.randn(10, 4))
+    data = np.array([[1], [2], [3], [4], [5], [6]])
+    indptr = np.array([0, 2, 3, 6])
+    indices = np.array([0, 2, 2, 0, 1, 2])
+    array = np.array([[1, 0, 2],
+                      [0, 0, 3],
+                      [4, 5, 6]])
+
+    sparse_tensor = pa.SparseTensorCSR.from_dense_numpy(array)
     n = sys.getrefcount(sparse_tensor)
-    data, indptr, indices = sparse_tensor.to_numpy()  # noqa
-    assert sys.getrefcount(sparse_tensor) == n + 1
+    result_data, result_indptr, result_indices = sparse_tensor.to_numpy()
+    assert sys.getrefcount(sparse_tensor) == n + 3
+
+    sparse_tensor = None
+    assert np.array_equal(data, result_data)
+    assert np.array_equal(indptr, result_indptr)
+    assert np.array_equal(indices, result_indices)
 
 
 @pytest.mark.skip
@@ -102,29 +125,29 @@ def test_sparse_tensor_equals(sparse_tensor_type):
 ])
 def test_sparse_tensor_coo_from_dense(dtype_str, arrow_type):
     dtype = np.dtype(dtype_str)
-    data_vector = np.array([4, 9, 7, 5]).astype(dtype)
+    data = np.array([[4], [9], [7], [5]]).astype(dtype)
     coords = np.array([[0, 0], [1, 3], [0, 2], [1, 3]])
-    data = np.array([[4, 0, 9, 0],
-                     [0, 7, 0, 0],
-                     [0, 0, 0, 0],
-                     [0, 0, 0, 5]]).astype(dtype)
-    tensor = pa.Tensor.from_numpy(data)
+    array = np.array([[4, 0, 9, 0],
+                      [0, 7, 0, 0],
+                      [0, 0, 0, 0],
+                      [0, 0, 0, 5]]).astype(dtype)
+    tensor = pa.Tensor.from_numpy(array)
 
     # Test from numpy array
-    sparse_tensor = pa.SparseTensorCOO.from_dense_numpy(data)
+    sparse_tensor = pa.SparseTensorCOO.from_dense_numpy(array)
     repr(sparse_tensor)
     assert sparse_tensor.type == arrow_type
     result_data, result_coords = sparse_tensor.to_numpy()
-    assert (data_vector == result_data).all()
-    assert (result_coords == coords).all()
+    assert np.array_equal(data, result_data)
+    assert np.array_equal(coords, result_coords)
 
     # Test from Tensor
     sparse_tensor = pa.SparseTensorCOO.from_tensor(tensor)
     repr(sparse_tensor)
     assert sparse_tensor.type == arrow_type
     result_data, result_coords = sparse_tensor.to_numpy()
-    assert (data_vector == result_data).all()
-    assert (result_coords == coords).all()
+    assert np.array_equal(data, result_data)
+    assert np.array_equal(coords, result_coords)
 
 
 @pytest.mark.parametrize('dtype_str,arrow_type', [
@@ -146,7 +169,7 @@ def test_sparse_tensor_csr_from_dense(dtype_str, arrow_type):
                            [0, 0, 3],
                            [4, 5, 6]]).astype(dtype)
 
-    data = np.array([1, 2, 3, 4, 5, 6]).astype(dtype)
+    data = np.array([[1], [2], [3], [4], [5], [6]])
     indptr = np.array([0, 2, 3, 6])
     indices = np.array([0, 2, 2, 0, 1, 2])
     tensor = pa.Tensor.from_numpy(dense_data)
@@ -155,21 +178,20 @@ def test_sparse_tensor_csr_from_dense(dtype_str, arrow_type):
     sparse_tensor = pa.SparseTensorCSR.from_dense_numpy(dense_data)
     repr(sparse_tensor)
     result_data, result_indptr, result_indices = sparse_tensor.to_numpy()
-    assert (data == result_data).all()
-    assert (indptr == result_indptr).all()
-    assert (indices == result_indices).all()
+    assert np.array_equal(data, result_data)
+    assert np.array_equal(indptr, result_indptr)
+    assert np.array_equal(indices, result_indices)
 
     # Test from Tensor
     sparse_tensor = pa.SparseTensorCSR.from_tensor(tensor)
     repr(sparse_tensor)
     assert sparse_tensor.type == arrow_type
     result_data, result_indptr, result_indices = sparse_tensor.to_numpy()
-    assert (data == result_data).all()
-    assert (indptr == result_indptr).all()
-    assert (indices == result_indices).all()
+    assert np.array_equal(data, result_data)
+    assert np.array_equal(indptr, result_indptr)
+    assert np.array_equal(indices, result_indices)
 
 
-@pytest.mark.skip
 @pytest.mark.parametrize('dtype_str,arrow_type', [
     ('i1', pa.int8()),
     ('i2', pa.int16()),
@@ -185,8 +207,8 @@ def test_sparse_tensor_csr_from_dense(dtype_str, arrow_type):
 ])
 def test_sparse_tensor_coo_numpy_roundtrip(dtype_str, arrow_type):
     dtype = np.dtype(dtype_str)
-    data = np.array([4, 5, 7, 9]).astype(dtype)
-    coords = np.array([[0, 0], [3, 3], [1, 1], [0, 2]]).astype('i8')
+    data = np.array([[4], [9], [7], [5]]).astype(dtype)
+    coords = np.array([[0, 0], [3, 3], [1, 1], [0, 2]])
     shape = (4, 4)
     dim_names = ["x", "y"]
 
@@ -195,8 +217,9 @@ def test_sparse_tensor_coo_numpy_roundtrip(dtype_str, arrow_type):
     repr(sparse_tensor)
     assert sparse_tensor.type == arrow_type
     result_data, result_coords = sparse_tensor.to_numpy()
-    assert (data == result_data).all()
-    assert (coords == result_coords).all()
+    assert np.array_equal(data, result_data)
+    assert np.array_equal(coords, result_coords)
+    assert sparse_tensor.dim_names == dim_names
 
 
 @pytest.mark.parametrize('dtype_str,arrow_type', [
@@ -214,9 +237,9 @@ def test_sparse_tensor_coo_numpy_roundtrip(dtype_str, arrow_type):
 ])
 def test_sparse_tensor_csr_numpy_roundtrip(dtype_str, arrow_type):
     dtype = np.dtype(dtype_str)
-    data = np.array([1, 2, 3, 4, 5, 6]).astype(dtype)
-    indptr = np.array([0, 2, 3, 6]).astype('i8')
-    indices = np.array([0, 2, 2, 0, 1, 2]).astype('i8')
+    data = np.array([[1], [2], [3], [4], [5], [6]]).astype(dtype)
+    indptr = np.array([0, 2, 3, 6])
+    indices = np.array([0, 2, 2, 0, 1, 2])
     shape = (3, 3)
     dim_names = ["x", "y"]
 
@@ -225,6 +248,7 @@ def test_sparse_tensor_csr_numpy_roundtrip(dtype_str, arrow_type):
     repr(sparse_tensor)
     assert sparse_tensor.type == arrow_type
     result_data, result_indptr, result_indices = sparse_tensor.to_numpy()
-    assert (data == result_data).all()
-    assert (indptr == result_indptr).all()
-    assert (indices == result_indices).all()
+    assert np.array_equal(data, result_data)
+    assert np.array_equal(indptr, result_indptr)
+    assert np.array_equal(indices, result_indices)
+    assert sparse_tensor.dim_names == dim_names