Cython wrapper for SparseTensor.

cpp/src/arrow/python/numpy_convert.cc

@@ -25,6 +25,7 @@
 #include <vector>
 
 #include "arrow/buffer.h"
+#include "arrow/sparse_tensor.h"
 #include "arrow/tensor.h"
 #include "arrow/type.h"
 
@@ -274,5 +275,313 @@ Status TensorToNdarray(const std::shared_ptr<Tensor>& tensor, PyObject* base,
   return Status::OK();
 }
 
+Status SparseTensorCOOToNdarray(const std::shared_ptr<SparseTensorCOO>& sparse_tensor,
+                                PyObject* base, PyObject** out_data,
+                                PyObject** out_coords) {
+  PyAcquireGIL lock;
+
+  int type_num_data;
+  int type_num_coords;
+  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});
+  std::vector<npy_intp> npy_shape_coords({sparse_index_coords->shape()[0], 2});
+
+  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();
+  }
+
+  // Remove const =(
+  void* mutable_data = const_cast<void*>(immutable_data);
+  void* mutable_coords = const_cast<void*>(immutable_coords);
+
+  int array_flags = 0;
+  if (sparse_tensor->is_mutable()) {
+    array_flags |= NPY_ARRAY_WRITEABLE;
+  }
+
+  PyObject* data =
+      PyArray_NewFromDescr(&PyArray_Type, dtype_data, 1, npy_shape_data.data(), nullptr,
+                           mutable_data, array_flags, nullptr);
+  PyObject* coords =
+      PyArray_NewFromDescr(&PyArray_Type, dtype_coords, 2, npy_shape_coords.data(),
+                           nullptr, mutable_coords, array_flags, nullptr);
+
+  RETURN_IF_PYERROR()
+
+  *out_data = data;
+  *out_coords = coords;
+  return Status::OK();
+}
+
+Status SparseTensorCSRToNdarray(const std::shared_ptr<SparseTensorCSR>& sparse_tensor,
+                                PyObject* base, PyObject** out_data,
+                                PyObject** out_indptr, PyObject** out_indices) {
+  PyAcquireGIL lock;
+
+  int type_num_data;
+  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 SparseCSRIndex&>(
+      *sparse_tensor->sparse_index());
+  const std::shared_ptr<arrow::NumericTensor<arrow::Int64Type>> sparse_index_indptr =
+      sparse_index.indptr();
+  const std::shared_ptr<arrow::NumericTensor<arrow::Int64Type>> sparse_index_indices =
+      sparse_index.indices();
+
+  std::vector<npy_intp> npy_shape_data({sparse_index.non_zero_length(), 1});
+  std::vector<npy_intp> npy_shape_indptr({sparse_index_indptr->shape()[0], 1});
+  std::vector<npy_intp> npy_shape_indices({sparse_index_indices->shape()[0], 1});
+
+  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();
+  }
+
+  // Remove const =(
+  void* mutable_data = const_cast<void*>(immutable_data);
+  void* mutable_indptr = const_cast<void*>(immutable_indptr);
+  void* mutable_indices = const_cast<void*>(immutable_indices);
+
+  int array_flags = 0;
+  if (sparse_tensor->is_mutable()) {
+    array_flags |= NPY_ARRAY_WRITEABLE;
+  }
+
+  int type_num_indptr;
+  RETURN_NOT_OK(GetNumPyType(*sparse_index_indptr->type(), &type_num_indptr));
+  PyArray_Descr* dtype_indptr = PyArray_DescrNewFromType(type_num_indptr);
+
+  int type_num_indices;
+  RETURN_NOT_OK(GetNumPyType(*sparse_index_indptr->type(), &type_num_indices));
+  PyArray_Descr* dtype_indices = PyArray_DescrNewFromType(type_num_indices);
+
+  PyObject* result_data =
+      PyArray_NewFromDescr(&PyArray_Type, dtype_data, 1, 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);
+  RETURN_IF_PYERROR()
+
+  *out_data = result_data;
+  *out_indptr = result_indptr;
+  *out_indices = result_indices;
+  return Status::OK();
+}
+
+Status NdarrayToSparseTensorCOO(MemoryPool* pool, PyObject* ao,
+                                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);
+
+  {
+    PyAcquireGIL lock;
+    npy_intp* array_strides = PyArray_STRIDES(ndarray);
+    npy_intp* array_shape = PyArray_SHAPE(ndarray);
+    for (int i = 0; i < ndim; ++i) {
+      if (array_strides[i] < 0) {
+        return Status::Invalid("Negative ndarray strides not supported");
+      }
+      shape[i] = array_shape[i];
+      strides[i] = array_strides[i];
+    }
+
+    std::shared_ptr<DataType> type;
+    RETURN_NOT_OK(
+        GetTensorType(reinterpret_cast<PyObject*>(PyArray_DESCR(ndarray)), &type));
+    Tensor tensor(type, data, shape, strides);
+    *out = std::make_shared<SparseTensorCOO>(tensor);
+    return Status::OK();
+  }
+}
+
+Status NdarrayToSparseTensorCSR(MemoryPool* pool, PyObject* ao,
+                                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);
+
+  {
+    PyAcquireGIL lock;
+    npy_intp* array_strides = PyArray_STRIDES(ndarray);
+    npy_intp* array_shape = PyArray_SHAPE(ndarray);
+    for (int i = 0; i < ndim; ++i) {
+      if (array_strides[i] < 0) {
+        return Status::Invalid("Negative ndarray strides not supported");
+      }
+      shape[i] = array_shape[i];
+      strides[i] = array_strides[i];
+    }
+
+    std::shared_ptr<DataType> type;
+    RETURN_NOT_OK(
+        GetTensorType(reinterpret_cast<PyObject*>(PyArray_DESCR(ndarray)), &type));
+    Tensor tensor(type, data, shape, strides);
+    *out = std::make_shared<SparseTensorCSR>(tensor);
+    return Status::OK();
+  }
+}
+
+Status NdarraysToSparseTensorCOO(MemoryPool* pool, PyObject* data_ao, PyObject* coords_ao,
+                                 const std::vector<int64_t>& shape,
+                                 const std::vector<std::string>& dim_names,
+                                 std::shared_ptr<SparseTensorCOO>* out) {
+  if (!PyArray_Check(data_ao) || !PyArray_Check(coords_ao)) {
+    return Status::TypeError("Did not pass ndarray object");
+  }
+
+  PyAcquireGIL lock;
+
+  PyArrayObject* ndarray_data = reinterpret_cast<PyArrayObject*>(data_ao);
+  PyArrayObject* ndarray_coords = reinterpret_cast<PyArrayObject*>(coords_ao);
+
+  std::shared_ptr<Buffer> data = std::make_shared<NumPyBuffer>(data_ao);
+  std::shared_ptr<Buffer> coords_buffer = std::make_shared<NumPyBuffer>(coords_ao);
+
+  int coords_ndim = PyArray_NDIM(ndarray_coords);
+
+  std::shared_ptr<DataType> type;
+  npy_intp* coords_array_shape = PyArray_SHAPE(ndarray_coords);
+  std::vector<int64_t> coords_shape(coords_ndim);
+
+  for (int i = 0; i < coords_ndim; ++i) {
+    coords_shape[i] = coords_array_shape[i];
+  }
+
+  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::shared_ptr<SparseCOOIndex> sparse_index = std::make_shared<SparseCOOIndex>(coords);
+  *out = std::make_shared<SparseTensorImpl<SparseCOOIndex>>(sparse_index, type, data,
+                                                            shape, dim_names);
+  return Status::OK();
+}
+
+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) {
+  if (!PyArray_Check(data_ao) || !PyArray_Check(indptr_ao) ||
+      !PyArray_Check(indices_ao)) {
+    return Status::TypeError("Did not pass ndarray object");
+  }
+
+  PyAcquireGIL lock;
+
+  PyArrayObject* ndarray_data = reinterpret_cast<PyArrayObject*>(data_ao);
+  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);
+
+  int indptr_ndim = PyArray_NDIM(ndarray_indptr);
+  int indices_ndim = PyArray_NDIM(ndarray_indices);
+
+  std::shared_ptr<DataType> type;
+  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);
+  std::vector<int64_t> indices_shape(indices_ndim);
+
+  for (int i = 0; i < indptr_ndim; ++i) {
+    indptr_shape[i] = indptr_array_shape[i];
+  }
+  for (int i = 0; i < indices_ndim; ++i) {
+    indices_shape[i] = indices_array_shape[i];
+  }
+
+  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,
+                                                            shape, dim_names);
+  return Status::OK();
+}
+
+Status TensorToSparseTensorCOO(const std::shared_ptr<Tensor>& tensor,
+                               std::shared_ptr<SparseTensorCOO>* out) {
+  {
+    PyAcquireGIL lock;
+    *out = std::make_shared<SparseTensorCOO>(*tensor);
+    return Status::OK();
+  }
+}
+
+Status TensorToSparseTensorCSR(const std::shared_ptr<Tensor>& tensor,
+                               std::shared_ptr<SparseTensorCSR>* out) {
+  {
+    PyAcquireGIL lock;
+    *out = std::make_shared<SparseTensorCSR>(*tensor);
+    return Status::OK();
+  }
+}
+
 }  // namespace py
 }  // namespace arrow

cpp/src/arrow/python/numpy_convert.h

@@ -25,9 +25,11 @@
 
 #include <memory>
 #include <string>
+#include <vector>
 
 #include "arrow/buffer.h"
 #include "arrow/python/visibility.h"
+#include "arrow/sparse_tensor.h"
 
 namespace arrow {
 
@@ -68,6 +70,38 @@ 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 NdarrayToSparseTensorCOO(
+    MemoryPool* pool, PyObject* ao, 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,
+    std::shared_ptr<SparseTensorCSR>* out);
+
+ARROW_PYTHON_EXPORT Status NdarraysToSparseTensorCOO(
+    MemoryPool* pool, PyObject* data_ao, PyObject* coords_ao,
+    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
+TensorToSparseTensorCOO(const std::shared_ptr<Tensor>& tensor,
+                        std::shared_ptr<SparseTensorCOO>* csparse_tensor);
+
 }  // namespace py
 }  // namespace arrow
 

cpp/src/arrow/python/pyarrow.cc

@@ -123,6 +123,44 @@ PyObject* wrap_tensor(const std::shared_ptr<Tensor>& tensor) {
   return ::pyarrow_wrap_tensor(tensor);
 }
 
+bool is_sparse_tensor_csr(PyObject* sparse_tensor) {
+  return ::pyarrow_is_sparse_tensor_csr(sparse_tensor) != 0;
+}
+
+Status unwrap_sparse_tensor_csr(PyObject* sparse_tensor,
+                                std::shared_ptr<SparseTensorCSR>* out) {
+  *out = ::pyarrow_unwrap_sparse_tensor_csr(sparse_tensor);
+  if (*out) {
+    return Status::OK();
+  } else {
+    return Status::Invalid(
+        "Could not unwrap SparseTensorCSR from the passed Python object.");
+  }
+}
+
+PyObject* wrap_sparse_tensor_csr(const std::shared_ptr<SparseTensorCSR>& sparse_tensor) {
+  return ::pyarrow_wrap_sparse_tensor_csr(sparse_tensor);
+}
+
+bool is_sparse_tensor_coo(PyObject* sparse_tensor) {
+  return ::pyarrow_is_sparse_tensor_coo(sparse_tensor) != 0;
+}
+
+Status unwrap_sparse_tensor_coo(PyObject* sparse_tensor,
+                                std::shared_ptr<SparseTensorCOO>* out) {
+  *out = ::pyarrow_unwrap_sparse_tensor_coo(sparse_tensor);
+  if (*out) {
+    return Status::OK();
+  } else {
+    return Status::Invalid(
+        "Could not unwrap SparseTensorCOO from the passed Python object.");
+  }
+}
+
+PyObject* wrap_sparse_tensor_coo(const std::shared_ptr<SparseTensorCOO>& sparse_tensor) {
+  return ::pyarrow_wrap_sparse_tensor_coo(sparse_tensor);
+}
+
 bool is_column(PyObject* column) { return ::pyarrow_is_column(column) != 0; }
 
 Status unwrap_column(PyObject* column, std::shared_ptr<Column>* out) {

cpp/src/arrow/python/pyarrow.h

@@ -24,6 +24,8 @@
 
 #include "arrow/python/visibility.h"
 
+#include "arrow/sparse_tensor.h"
+
 namespace arrow {
 
 class Array;
@@ -67,6 +69,18 @@ ARROW_PYTHON_EXPORT bool is_tensor(PyObject* tensor);
 ARROW_PYTHON_EXPORT Status unwrap_tensor(PyObject* tensor, std::shared_ptr<Tensor>* out);
 ARROW_PYTHON_EXPORT PyObject* wrap_tensor(const std::shared_ptr<Tensor>& tensor);
 
+ARROW_PYTHON_EXPORT bool is_sparse_tensor_coo(PyObject* sparse_tensor);
+ARROW_PYTHON_EXPORT Status
+unwrap_sparse_tensor_coo(PyObject* sparse_tensor, std::shared_ptr<SparseTensorCOO>* out);
+ARROW_PYTHON_EXPORT PyObject* wrap_sparse_tensor_coo(
+    const std::shared_ptr<SparseTensorCOO>& sparse_tensor);
+
+ARROW_PYTHON_EXPORT bool is_sparse_tensor_csr(PyObject* sparse_tensor);
+ARROW_PYTHON_EXPORT Status
+unwrap_sparse_tensor_csr(PyObject* sparse_tensor, std::shared_ptr<SparseTensorCSR>* out);
+ARROW_PYTHON_EXPORT PyObject* wrap_sparse_tensor_csr(
+    const std::shared_ptr<SparseTensorCSR>& sparse_tensor);
+
 ARROW_PYTHON_EXPORT bool is_column(PyObject* column);
 ARROW_PYTHON_EXPORT Status unwrap_column(PyObject* column, std::shared_ptr<Column>* out);
 ARROW_PYTHON_EXPORT PyObject* wrap_column(const std::shared_ptr<Column>& column);

cpp/src/arrow/python/pyarrow_api.h

@@ -50,6 +50,10 @@ static PyObject *(*__pyx_api_f_7pyarrow_3lib_pyarrow_wrap_table)(std::shared_ptr
 #define pyarrow_wrap_table __pyx_api_f_7pyarrow_3lib_pyarrow_wrap_table
 static PyObject *(*__pyx_api_f_7pyarrow_3lib_pyarrow_wrap_tensor)(std::shared_ptr< arrow::Tensor>  const &) = 0;
 #define pyarrow_wrap_tensor __pyx_api_f_7pyarrow_3lib_pyarrow_wrap_tensor
+static PyObject *(*__pyx_api_f_7pyarrow_3lib_pyarrow_wrap_sparse_tensor_csr)(std::shared_ptr< arrow::SparseTensorCSR>  const &) = 0;
+#define pyarrow_wrap_sparse_tensor_csr __pyx_api_f_7pyarrow_3lib_pyarrow_wrap_sparse_tensor_csr
+static PyObject *(*__pyx_api_f_7pyarrow_3lib_pyarrow_wrap_sparse_tensor_coo)(std::shared_ptr< arrow::SparseTensorCOO>  const &) = 0;
+#define pyarrow_wrap_sparse_tensor_coo __pyx_api_f_7pyarrow_3lib_pyarrow_wrap_sparse_tensor_coo
 static std::shared_ptr< arrow::Array>  (*__pyx_api_f_7pyarrow_3lib_pyarrow_unwrap_array)(PyObject *) = 0;
 #define pyarrow_unwrap_array __pyx_api_f_7pyarrow_3lib_pyarrow_unwrap_array
 static std::shared_ptr< arrow::RecordBatch>  (*__pyx_api_f_7pyarrow_3lib_pyarrow_unwrap_batch)(PyObject *) = 0;
@@ -68,6 +72,10 @@ static std::shared_ptr< arrow::Table>  (*__pyx_api_f_7pyarrow_3lib_pyarrow_unwra
 #define pyarrow_unwrap_table __pyx_api_f_7pyarrow_3lib_pyarrow_unwrap_table
 static std::shared_ptr< arrow::Tensor>  (*__pyx_api_f_7pyarrow_3lib_pyarrow_unwrap_tensor)(PyObject *) = 0;
 #define pyarrow_unwrap_tensor __pyx_api_f_7pyarrow_3lib_pyarrow_unwrap_tensor
+static std::shared_ptr< arrow::SparseTensorCSR>  (*__pyx_api_f_7pyarrow_3lib_pyarrow_unwrap_sparse_tensor_csr)(PyObject *) = 0;
+#define pyarrow_unwrap_sparse_tensor_csr __pyx_api_f_7pyarrow_3lib_pyarrow_unwrap_sparse_tensor_csr
+static std::shared_ptr< arrow::SparseTensorCOO>  (*__pyx_api_f_7pyarrow_3lib_pyarrow_unwrap_sparse_tensor_coo)(PyObject *) = 0;
+#define pyarrow_unwrap_sparse_tensor_coo __pyx_api_f_7pyarrow_3lib_pyarrow_unwrap_sparse_tensor_coo
 static int (*__pyx_api_f_7pyarrow_3lib_pyarrow_internal_check_status)(arrow::Status const &) = 0;
 #define pyarrow_internal_check_status __pyx_api_f_7pyarrow_3lib_pyarrow_internal_check_status
 static int (*__pyx_api_f_7pyarrow_3lib_pyarrow_is_buffer)(PyObject *) = 0;
@@ -84,6 +92,10 @@ static PyObject *(*__pyx_api_f_7pyarrow_3lib_pyarrow_wrap_scalar)(std::shared_pt
 #define pyarrow_wrap_scalar __pyx_api_f_7pyarrow_3lib_pyarrow_wrap_scalar
 static int (*__pyx_api_f_7pyarrow_3lib_pyarrow_is_tensor)(PyObject *) = 0;
 #define pyarrow_is_tensor __pyx_api_f_7pyarrow_3lib_pyarrow_is_tensor
+static int (*__pyx_api_f_7pyarrow_3lib_pyarrow_is_sparse_tensor_csr)(PyObject *) = 0;
+#define pyarrow_is_sparse_tensor_csr __pyx_api_f_7pyarrow_3lib_pyarrow_is_sparse_tensor_csr
+static int (*__pyx_api_f_7pyarrow_3lib_pyarrow_is_sparse_tensor_coo)(PyObject *) = 0;
+#define pyarrow_is_sparse_tensor_coo __pyx_api_f_7pyarrow_3lib_pyarrow_is_sparse_tensor_coo
 static int (*__pyx_api_f_7pyarrow_3lib_pyarrow_is_column)(PyObject *) = 0;
 #define pyarrow_is_column __pyx_api_f_7pyarrow_3lib_pyarrow_is_column
 static int (*__pyx_api_f_7pyarrow_3lib_pyarrow_is_table)(PyObject *) = 0;
@@ -167,6 +179,8 @@ static int import_pyarrow__lib(void) {
   if (__Pyx_ImportFunction(module, "pyarrow_wrap_schema", (void (**)(void))&__pyx_api_f_7pyarrow_3lib_pyarrow_wrap_schema, "PyObject *(std::shared_ptr< arrow::Schema>  const &)") < 0) goto bad;
   if (__Pyx_ImportFunction(module, "pyarrow_wrap_table", (void (**)(void))&__pyx_api_f_7pyarrow_3lib_pyarrow_wrap_table, "PyObject *(std::shared_ptr< arrow::Table>  const &)") < 0) goto bad;
   if (__Pyx_ImportFunction(module, "pyarrow_wrap_tensor", (void (**)(void))&__pyx_api_f_7pyarrow_3lib_pyarrow_wrap_tensor, "PyObject *(std::shared_ptr< arrow::Tensor>  const &)") < 0) goto bad;
+  if (__Pyx_ImportFunction(module, "pyarrow_wrap_sparse_tensor_csr", (void (**)(void))&__pyx_api_f_7pyarrow_3lib_pyarrow_wrap_sparse_tensor_csr, "PyObject *(std::shared_ptr< arrow::SparseTensorCSR>  const &)") < 0) goto bad;
+  if (__Pyx_ImportFunction(module, "pyarrow_wrap_sparse_tensor_coo", (void (**)(void))&__pyx_api_f_7pyarrow_3lib_pyarrow_wrap_sparse_tensor_csr, "PyObject *(std::shared_ptr< arrow::SparseTensorCOO>  const &)") < 0) goto bad;
   if (__Pyx_ImportFunction(module, "pyarrow_unwrap_array", (void (**)(void))&__pyx_api_f_7pyarrow_3lib_pyarrow_unwrap_array, "std::shared_ptr< arrow::Array>  (PyObject *)") < 0) goto bad;
   if (__Pyx_ImportFunction(module, "pyarrow_unwrap_batch", (void (**)(void))&__pyx_api_f_7pyarrow_3lib_pyarrow_unwrap_batch, "std::shared_ptr< arrow::RecordBatch>  (PyObject *)") < 0) goto bad;
   if (__Pyx_ImportFunction(module, "pyarrow_unwrap_buffer", (void (**)(void))&__pyx_api_f_7pyarrow_3lib_pyarrow_unwrap_buffer, "std::shared_ptr< arrow::Buffer>  (PyObject *)") < 0) goto bad;
@@ -176,6 +190,8 @@ static int import_pyarrow__lib(void) {
   if (__Pyx_ImportFunction(module, "pyarrow_unwrap_schema", (void (**)(void))&__pyx_api_f_7pyarrow_3lib_pyarrow_unwrap_schema, "std::shared_ptr< arrow::Schema>  (PyObject *)") < 0) goto bad;
   if (__Pyx_ImportFunction(module, "pyarrow_unwrap_table", (void (**)(void))&__pyx_api_f_7pyarrow_3lib_pyarrow_unwrap_table, "std::shared_ptr< arrow::Table>  (PyObject *)") < 0) goto bad;
   if (__Pyx_ImportFunction(module, "pyarrow_unwrap_tensor", (void (**)(void))&__pyx_api_f_7pyarrow_3lib_pyarrow_unwrap_tensor, "std::shared_ptr< arrow::Tensor>  (PyObject *)") < 0) goto bad;
+  if (__Pyx_ImportFunction(module, "pyarrow_unwrap_sparse_tensor_csr", (void (**)(void))&__pyx_api_f_7pyarrow_3lib_pyarrow_unwrap_sparse_tensor_csr, "std::shared_ptr< arrow::SparseTensorCSR>  (PyObject *)") < 0) goto bad;
+  if (__Pyx_ImportFunction(module, "pyarrow_unwrap_sparse_tensor_coo", (void (**)(void))&__pyx_api_f_7pyarrow_3lib_pyarrow_unwrap_sparse_tensor_coo, "std::shared_ptr< arrow::SparseTensorCOO>  (PyObject *)") < 0) goto bad;
   if (__Pyx_ImportFunction(module, "pyarrow_internal_check_status", (void (**)(void))&__pyx_api_f_7pyarrow_3lib_pyarrow_internal_check_status, "int (arrow::Status const &)") < 0) goto bad;
   if (__Pyx_ImportFunction(module, "pyarrow_is_buffer", (void (**)(void))&__pyx_api_f_7pyarrow_3lib_pyarrow_is_buffer, "int (PyObject *)") < 0) goto bad;
   if (__Pyx_ImportFunction(module, "pyarrow_is_data_type", (void (**)(void))&__pyx_api_f_7pyarrow_3lib_pyarrow_is_data_type, "int (PyObject *)") < 0) goto bad;
@@ -184,6 +200,8 @@ static int import_pyarrow__lib(void) {
   if (__Pyx_ImportFunction(module, "pyarrow_is_array", (void (**)(void))&__pyx_api_f_7pyarrow_3lib_pyarrow_is_array, "int (PyObject *)") < 0) goto bad;
   if (__Pyx_ImportFunction(module, "pyarrow_wrap_scalar", (void (**)(void))&__pyx_api_f_7pyarrow_3lib_pyarrow_wrap_scalar, "PyObject *(std::shared_ptr< arrow::Scalar>  const &)") < 0) goto bad;
   if (__Pyx_ImportFunction(module, "pyarrow_is_tensor", (void (**)(void))&__pyx_api_f_7pyarrow_3lib_pyarrow_is_tensor, "int (PyObject *)") < 0) goto bad;
+  if (__Pyx_ImportFunction(module, "pyarrow_is_sparse_tensor_csr", (void (**)(void))&__pyx_api_f_7pyarrow_3lib_pyarrow_is_sparse_tensor_csr, "int (PyObject *)") < 0) goto bad;
+  if (__Pyx_ImportFunction(module, "pyarrow_is_sparse_tensor_coo", (void (**)(void))&__pyx_api_f_7pyarrow_3lib_pyarrow_is_sparse_tensor_coo, "int (PyObject *)") < 0) goto bad;
   if (__Pyx_ImportFunction(module, "pyarrow_is_column", (void (**)(void))&__pyx_api_f_7pyarrow_3lib_pyarrow_is_column, "int (PyObject *)") < 0) goto bad;
   if (__Pyx_ImportFunction(module, "pyarrow_is_table", (void (**)(void))&__pyx_api_f_7pyarrow_3lib_pyarrow_is_table, "int (PyObject *)") < 0) goto bad;
   if (__Pyx_ImportFunction(module, "pyarrow_is_batch", (void (**)(void))&__pyx_api_f_7pyarrow_3lib_pyarrow_is_batch, "int (PyObject *)") < 0) goto bad;

cpp/src/arrow/python/pyarrow_lib.h

@@ -48,6 +48,8 @@ __PYX_EXTERN_C PyObject *__pyx_f_7pyarrow_3lib_pyarrow_wrap_resizable_buffer(std
 __PYX_EXTERN_C PyObject *__pyx_f_7pyarrow_3lib_pyarrow_wrap_schema(std::shared_ptr< arrow::Schema>  const &);
 __PYX_EXTERN_C PyObject *__pyx_f_7pyarrow_3lib_pyarrow_wrap_table(std::shared_ptr< arrow::Table>  const &);
 __PYX_EXTERN_C PyObject *__pyx_f_7pyarrow_3lib_pyarrow_wrap_tensor(std::shared_ptr< arrow::Tensor>  const &);
+__PYX_EXTERN_C PyObject *__pyx_f_7pyarrow_3lib_pyarrow_wrap_sparse_tensor_coo(std::shared_ptr< arrow::SparseTensorCOO>  const &);
+__PYX_EXTERN_C PyObject *__pyx_f_7pyarrow_3lib_pyarrow_wrap_sparse_tensor_csr(std::shared_ptr< arrow::SparseTensorCSR>  const &);
 __PYX_EXTERN_C std::shared_ptr< arrow::Array>  __pyx_f_7pyarrow_3lib_pyarrow_unwrap_array(PyObject *);
 __PYX_EXTERN_C std::shared_ptr< arrow::RecordBatch>  __pyx_f_7pyarrow_3lib_pyarrow_unwrap_batch(PyObject *);
 __PYX_EXTERN_C std::shared_ptr< arrow::Buffer>  __pyx_f_7pyarrow_3lib_pyarrow_unwrap_buffer(PyObject *);
@@ -57,6 +59,20 @@ __PYX_EXTERN_C std::shared_ptr< arrow::Field>  __pyx_f_7pyarrow_3lib_pyarrow_unw
 __PYX_EXTERN_C std::shared_ptr< arrow::Schema>  __pyx_f_7pyarrow_3lib_pyarrow_unwrap_schema(PyObject *);
 __PYX_EXTERN_C std::shared_ptr< arrow::Table>  __pyx_f_7pyarrow_3lib_pyarrow_unwrap_table(PyObject *);
 __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 */
 

docs/source/python/extending.rst

@@ -116,6 +116,16 @@ C++ objects.
    Return whether *obj* wraps an Arrow C++ :class:`Tensor` pointer;
    in other words, whether *obj* is a :py:class:`pyarrow.Tensor` instance.
 
+.. function:: bool is_sparse_tensor_coo(PyObject* obj)
+
+   Return whether *obj* wraps an Arrow C++ :class:`SparseTensorCOO` pointer;
+   in other words, whether *obj* is a :py:class:`pyarrow.SparseTensorCOO` instance.
+
+.. function:: bool is_sparse_tensor_csr(PyObject* obj)
+
+   Return whether *obj* wraps an Arrow C++ :class:`SparseTensorCSR` pointer;
+   in other words, whether *obj* is a :py:class:`pyarrow.SparseTensorCSR` instance.
+
 The following functions expect a pyarrow object, unwrap the underlying
 Arrow C++ API pointer, and put it in the *out* parameter.  The returned
 :class:`Status` object must be inspected first to know whether any error
@@ -157,6 +167,14 @@ occurred.  If successful, *out* is guaranteed to be non-NULL.
 
    Unwrap the Arrow C++ :class:`Tensor` pointer from *obj* and put it in *out*.
 
+.. function:: Status unwrap_sparse_tensor_coo(PyObject* obj, std::shared_ptr<SparseTensorCOO>* out)
+
+   Unwrap the Arrow C++ :class:`SparseTensorCOO` pointer from *obj* and put it in *out*.
+
+.. function:: Status unwrap_sparse_tensor_csr(PyObject* obj, std::shared_ptr<SparseTensorCSR>* out)
+
+   Unwrap the Arrow C++ :class:`SparseTensorCSR` pointer from *obj* and put it in *out*.
+
 The following functions take an Arrow C++ API pointer and wrap it in a
 pyarray object of the corresponding type.  A new reference is returned.
 On error, NULL is returned and a Python exception is set.
@@ -197,6 +215,14 @@ On error, NULL is returned and a Python exception is set.
 
    Wrap the Arrow C++ *tensor* in a :py:class:`pyarrow.Tensor` instance.
 
+.. function:: PyObject* wrap_sparse_tensor_coo(const std::shared_ptr<SparseTensorCOO>& sparse_tensor)
+
+   Wrap the Arrow C++ *COO sparse tensor* in a :py:class:`pyarrow.SparseTensorCOO` instance.
+
+.. function:: PyObject* wrap_sparse_tensor_csr(const std::shared_ptr<SparseTensorCSR>& sparse_tensor)
+
+   Wrap the Arrow C++ *CSR sparse tensor* in a :py:class:`pyarrow.SparseTensorCSR` instance.
+
 
 Cython API
 ----------
@@ -257,6 +283,14 @@ an exception) if the input is not of the right type.
 
    Unwrap the Arrow C++ :cpp:class:`Tensor` pointer from *obj*.
 
+.. function:: pyarrow_unwrap_sparse_tensor_coo(obj) -> shared_ptr[CSparseTensorCOO]
+
+   Unwrap the Arrow C++ :cpp:class:`SparseTensorCOO` pointer from *obj*.
+
+.. function:: pyarrow_unwrap_sparse_tensor_csr(obj) -> shared_ptr[CSparseTensorCSR]
+
+   Unwrap the Arrow C++ :cpp:class:`SparseTensorCSR` pointer from *obj*.
+
 The following functions take a Arrow C++ API pointer and wrap it in a
 pyarray object of the corresponding type.  An exception is raised on error.
 
@@ -300,6 +334,14 @@ pyarray object of the corresponding type.  An exception is raised on error.
 
    Wrap the Arrow C++ *tensor* in a Python :class:`pyarrow.Tensor` instance.
 
+.. function:: pyarrow_wrap_sparse_tensor_coo(sp_array: const shared_ptr[CSparseTensorCOO]& sparse_tensor) -> object
+
+   Wrap the Arrow C++ *COO sparse tensor* in a Python :class:`pyarrow.SparseTensorCOO` instance.
+
+.. function:: pyarrow_wrap_sparse_tensor_csr(sp_array: const shared_ptr[CSparseTensorCSR]& sparse_tensor) -> object
+
+   Wrap the Arrow C++ *CSR sparse tensor* in a Python :class:`pyarrow.SparseTensorCSR` instance.
+
 Example
 ~~~~~~~
 

python/pyarrow/__init__.pxd

@@ -20,7 +20,8 @@ from __future__ import absolute_import
 from libcpp.memory cimport shared_ptr
 from pyarrow.includes.libarrow cimport (CArray, CBuffer, CColumn, CDataType,
                                         CField, CRecordBatch, CSchema,
-                                        CTable, CTensor)
+                                        CTable, CTensor,
+                                        CSparseTensorCSR, CSparseTensorCOO)
 
 
 cdef extern from "arrow/python/pyarrow.h" namespace "arrow::py":
@@ -31,6 +32,10 @@ cdef extern from "arrow/python/pyarrow.h" namespace "arrow::py":
     cdef object wrap_schema(const shared_ptr[CSchema]& schema)
     cdef object wrap_array(const shared_ptr[CArray]& sp_array)
     cdef object wrap_tensor(const shared_ptr[CTensor]& sp_tensor)
+    cdef object wrap_sparse_tensor_coo(
+        const shared_ptr[CSparseTensorCOO]& sp_sparse_tensor)
+    cdef object wrap_sparse_tensor_csr(
+        const shared_ptr[CSparseTensorCSR]& sp_sparse_tensor)
     cdef object wrap_column(const shared_ptr[CColumn]& ccolumn)
     cdef object wrap_table(const shared_ptr[CTable]& ctable)
     cdef object wrap_batch(const shared_ptr[CRecordBatch]& cbatch)

python/pyarrow/__init__.py

@@ -66,6 +66,7 @@ from pyarrow.lib import (null, bool_,
                          schema,
                          Array, Tensor,
                          array, chunked_array, column, table,
+                         SparseTensorCSR, SparseTensorCOO,
                          infer_type, from_numpy_dtype,
                          NullArray,
                          NumericArray, IntegerArray, FloatingPointArray,

python/pyarrow/array.pxi

@@ -870,6 +870,199 @@ cdef class Array(_PandasConvertible):
         return res
 
 
+cdef class SparseTensorCOO:
+    """
+    A sparse COO tensor.
+    """
+
+    def __init__(self):
+        raise TypeError("Do not call SparseTensorCOO's constructor directly, "
+                        "use one of the `pyarrow.SparseTensorCOO.from_*` "
+                        "functions instead.")
+
+    cdef void init(self, const shared_ptr[CSparseTensorCOO]& sp_sparse_tensor):
+        self.sp_sparse_tensor = sp_sparse_tensor
+        self.stp = sp_sparse_tensor.get()
+        self.type = pyarrow_wrap_data_type(self.stp.type())
+
+    def __repr__(self):
+        return """<pyarrow.SparseTensorCOO>
+type: {0.type}
+shape: {0.shape}""".format(self)
+
+    @staticmethod
+    def from_dense_numpy(obj):
+        cdef shared_ptr[CSparseTensorCOO] csparse_tensor
+        with nogil:
+            check_status(NdarrayToSparseTensorCOO(c_default_memory_pool(), obj,
+                                                  &csparse_tensor))
+        return pyarrow_wrap_sparse_tensor_coo(csparse_tensor)
+
+    @staticmethod
+    def from_numpy(data, coords, shape, dim_names=None):
+        cdef shared_ptr[CSparseTensorCOO] csparse_tensor
+        cdef vector[int64_t] c_shape
+        cdef vector[c_string] c_dim_names
+
+        for x in shape:
+            c_shape.push_back(x)
+        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))
+        return pyarrow_wrap_sparse_tensor_coo(csparse_tensor)
+
+    @staticmethod
+    def from_tensor(obj):
+        cdef shared_ptr[CSparseTensorCOO] csparse_tensor
+        cdef shared_ptr[CTensor] ctensor = pyarrow_unwrap_tensor(obj)
+
+        with nogil:
+            check_status(TensorToSparseTensorCOO(ctensor, &csparse_tensor))
+
+        return pyarrow_wrap_sparse_tensor_coo(csparse_tensor)
+
+    def to_numpy(self):
+        """
+        Convert arrow::SparseTensorCOO to numpy.ndarrays with zero copy
+        """
+        cdef PyObject* out_data
+        cdef PyObject* out_coords
+
+        with nogil:
+            check_status(SparseTensorCOOToNdarray(self.sp_sparse_tensor, self,
+                                                  &out_data, &out_coords))
+        return PyObject_to_object(out_data), PyObject_to_object(out_coords)
+
+    @property
+    def is_mutable(self):
+        return self.stp.is_mutable()
+
+    @property
+    def ndim(self):
+        return self.stp.ndim()
+
+    @property
+    def shape(self):
+        # Cython knows how to convert a vector[T] to a Python list
+        return tuple(self.stp.shape())
+
+    @property
+    def size(self):
+        return self.stp.size()
+
+    def dim_name(self, i):
+        return self.stp.dim_name(i)
+
+    @property
+    def dim_names(self):
+        return self.stp.dim_names()
+
+    @property
+    def non_zero_length(self):
+        return self.stp.non_zero_length()
+
+
+cdef class SparseTensorCSR:
+    """
+    A sparse CSR tensor.
+    """
+
+    def __init__(self):
+        raise TypeError("Do not call SparseTensorCSR's constructor directly, "
+                        "use one of the `pyarrow.SparseTensorCSR.from_*` "
+                        "functions instead.")
+
+    cdef void init(self, const shared_ptr[CSparseTensorCSR]& sp_sparse_tensor):
+        self.sp_sparse_tensor = sp_sparse_tensor
+        self.stp = sp_sparse_tensor.get()
+        self.type = pyarrow_wrap_data_type(self.stp.type())
+
+    def __repr__(self):
+        return """<pyarrow.SparseTensorCSR>
+type: {0.type}
+shape: {0.shape}""".format(self)
+
+    @staticmethod
+    def from_dense_numpy(obj):
+        cdef shared_ptr[CSparseTensorCSR] csparse_tensor
+        with nogil:
+            check_status(NdarrayToSparseTensorCSR(c_default_memory_pool(), obj,
+                                                  &csparse_tensor))
+        return pyarrow_wrap_sparse_tensor_csr(csparse_tensor)
+
+    @staticmethod
+    def from_numpy(data, indptr, indices, shape, dim_names=None):
+        cdef shared_ptr[CSparseTensorCSR] csparse_tensor
+        cdef vector[int64_t] c_shape
+        cdef vector[c_string] c_dim_names
+
+        for x in shape:
+            c_shape.push_back(x)
+        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))
+        return pyarrow_wrap_sparse_tensor_csr(csparse_tensor)
+
+    @staticmethod
+    def from_tensor(obj):
+        cdef shared_ptr[CSparseTensorCSR] csparse_tensor
+        cdef shared_ptr[CTensor] ctensor = pyarrow_unwrap_tensor(obj)
+
+        with nogil:
+            check_status(TensorToSparseTensorCSR(ctensor, &csparse_tensor))
+
+        return pyarrow_wrap_sparse_tensor_csr(csparse_tensor)
+
+    def to_numpy(self):
+        """
+        Convert arrow::SparseTensorCSR to numpy.ndarrays with zero copy
+        """
+        cdef PyObject* out_data
+        cdef PyObject* out_indptr
+        cdef PyObject* out_indices
+
+        with nogil:
+            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))
+
+    @property
+    def is_mutable(self):
+        return self.stp.is_mutable()
+
+    @property
+    def ndim(self):
+        return self.stp.ndim()
+
+    @property
+    def shape(self):
+        # Cython knows how to convert a vector[T] to a Python list
+        return tuple(self.stp.shape())
+
+    @property
+    def size(self):
+        return self.stp.size()
+
+    def dim_name(self, i):
+        return self.stp.dim_name(i)
+
+    @property
+    def dim_names(self):
+        return self.stp.dim_names()
+
+    @property
+    def non_zero_length(self):
+        return self.stp.non_zero_length()
+
+
 cdef class Tensor:
     """
     A n-dimensional array a.k.a Tensor.

python/pyarrow/includes/libarrow.pxd

@@ -600,6 +600,38 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
         Type type_id()
         c_bool Equals(const CTensor& other)
 
+    cdef cppclass CSparseTensorCOO" arrow::SparseTensorCOO":
+        shared_ptr[CDataType] type()
+        shared_ptr[CBuffer] data()
+
+        const vector[int64_t]& shape()
+        int64_t size()
+        int64_t non_zero_length()
+
+        int ndim()
+        const vector[c_string]& dim_names()
+        const c_string& dim_name(int i)
+
+        c_bool is_mutable()
+        Type type_id()
+        c_bool Equals(const CSparseTensorCOO& other)
+
+    cdef cppclass CSparseTensorCSR" arrow::SparseTensorCSR":
+        shared_ptr[CDataType] type()
+        shared_ptr[CBuffer] data()
+
+        const vector[int64_t]& shape()
+        int64_t size()
+        int64_t non_zero_length()
+
+        int ndim()
+        const vector[c_string]& dim_names()
+        const c_string& dim_name(int i)
+
+        c_bool is_mutable()
+        Type type_id()
+        c_bool Equals(const CSparseTensorCSR& other)
+
     cdef cppclass CScalar" arrow::Scalar":
         shared_ptr[CDataType] type
 
@@ -1207,6 +1239,38 @@ cdef extern from "arrow/python/api.h" namespace "arrow::py" nogil:
     CStatus TensorToNdarray(const shared_ptr[CTensor]& tensor, object base,
                             PyObject** out)
 
+    CStatus SparseTensorCOOToNdarray(
+        const shared_ptr[CSparseTensorCOO]& sparse_tensor, object base,
+        PyObject** out_data, PyObject** out_coords)
+
+    CStatus SparseTensorCSRToNdarray(
+        const shared_ptr[CSparseTensorCSR]& sparse_tensor, object base,
+        PyObject** out_data, PyObject** out_indptr, PyObject** out_indices)
+
+    CStatus NdarraysToSparseTensorCOO(CMemoryPool* pool, object data_ao,
+                                      object coords_ao,
+                                      const vector[int64_t]& shape,
+                                      const vector[c_string]& dim_names,
+                                      shared_ptr[CSparseTensorCOO]* out)
+
+    CStatus NdarraysToSparseTensorCSR(CMemoryPool* pool, object data_ao,
+                                      object indptr_ao, object indices_ao,
+                                      const vector[int64_t]& shape,
+                                      const vector[c_string]& dim_names,
+                                      shared_ptr[CSparseTensorCSR]* out)
+
+    CStatus NdarrayToSparseTensorCOO(CMemoryPool* pool, object ao,
+                                     shared_ptr[CSparseTensorCOO]* out)
+
+    CStatus NdarrayToSparseTensorCSR(CMemoryPool* pool, object ao,
+                                     shared_ptr[CSparseTensorCSR]* out)
+
+    CStatus TensorToSparseTensorCOO(shared_ptr[CTensor],
+                                    shared_ptr[CSparseTensorCOO]* out)
+
+    CStatus TensorToSparseTensorCSR(shared_ptr[CTensor],
+                                    shared_ptr[CSparseTensorCSR]* out)
+
     CStatus ConvertArrayToPandas(const PandasOptions& options,
                                  const shared_ptr[CArray]& arr,
                                  object py_ref, PyObject** out)

python/pyarrow/lib.pxd

@@ -231,6 +231,28 @@ cdef class Tensor:
     cdef void init(self, const shared_ptr[CTensor]& sp_tensor)
 
 
+cdef class SparseTensorCSR:
+    cdef:
+        shared_ptr[CSparseTensorCSR] sp_sparse_tensor
+        CSparseTensorCSR* stp
+
+    cdef readonly:
+        DataType type
+
+    cdef void init(self, const shared_ptr[CSparseTensorCSR]& sp_sparse_tensor)
+
+
+cdef class SparseTensorCOO:
+    cdef:
+        shared_ptr[CSparseTensorCOO] sp_sparse_tensor
+        CSparseTensorCOO* stp
+
+    cdef readonly:
+        DataType type
+
+    cdef void init(self, const shared_ptr[CSparseTensorCOO]& sp_sparse_tensor)
+
+
 cdef class NullArray(Array):
     pass
 
@@ -452,6 +474,10 @@ cdef public object pyarrow_wrap_resizable_buffer(
 cdef public object pyarrow_wrap_schema(const shared_ptr[CSchema]& type)
 cdef public object pyarrow_wrap_table(const shared_ptr[CTable]& ctable)
 cdef public object pyarrow_wrap_tensor(const shared_ptr[CTensor]& sp_tensor)
+cdef public object pyarrow_wrap_sparse_tensor_coo(
+    const shared_ptr[CSparseTensorCOO]& sp_sparse_tensor)
+cdef public object pyarrow_wrap_sparse_tensor_csr(
+    const shared_ptr[CSparseTensorCSR]& sp_sparse_tensor)
 
 cdef public shared_ptr[CArray] pyarrow_unwrap_array(object array)
 cdef public shared_ptr[CRecordBatch] pyarrow_unwrap_batch(object batch)
@@ -462,3 +488,7 @@ cdef public shared_ptr[CField] pyarrow_unwrap_field(object field)
 cdef public shared_ptr[CSchema] pyarrow_unwrap_schema(object schema)
 cdef public shared_ptr[CTable] pyarrow_unwrap_table(object table)
 cdef public shared_ptr[CTensor] pyarrow_unwrap_tensor(object tensor)
+cdef public shared_ptr[CSparseTensorCOO] pyarrow_unwrap_sparse_tensor_coo(
+    object sparse_tensor)
+cdef public shared_ptr[CSparseTensorCSR] pyarrow_unwrap_sparse_tensor_csr(
+    object sparse_tensor)

python/pyarrow/public-api.pxi

@@ -18,7 +18,8 @@
 from libcpp.memory cimport shared_ptr
 from pyarrow.includes.libarrow cimport (CArray, CColumn, CDataType, CField,
                                         CRecordBatch, CSchema,
-                                        CTable, CTensor)
+                                        CTable, CTensor,
+                                        CSparseTensorCSR, CSparseTensorCOO)
 
 # You cannot assign something to a dereferenced pointer in Cython thus these
 # methods don't use Status to indicate a successful operation.
@@ -225,6 +226,7 @@ cdef api object pyarrow_wrap_scalar(const shared_ptr[CScalar]& sp_scalar):
     scalar.init(sp_scalar)
     return scalar
 
+
 cdef api bint pyarrow_is_tensor(object tensor):
     return isinstance(tensor, Tensor)
 
@@ -248,6 +250,52 @@ cdef api object pyarrow_wrap_tensor(
     return tensor
 
 
+cdef api bint pyarrow_is_sparse_tensor_coo(object sparse_tensor):
+    return isinstance(sparse_tensor, SparseTensorCOO)
+
+cdef api shared_ptr[CSparseTensorCOO] pyarrow_unwrap_sparse_tensor_coo(
+        object sparse_tensor):
+    cdef SparseTensorCOO sten
+    if pyarrow_is_sparse_tensor_coo(sparse_tensor):
+        sten = <SparseTensorCOO>(sparse_tensor)
+        return sten.sp_sparse_tensor
+
+    return shared_ptr[CSparseTensorCOO]()
+
+cdef api object pyarrow_wrap_sparse_tensor_coo(
+        const shared_ptr[CSparseTensorCOO]& sp_sparse_tensor):
+    if sp_sparse_tensor.get() == NULL:
+        raise ValueError('SparseTensorCOO was NULL')
+
+    cdef SparseTensorCOO sparse_tensor = SparseTensorCOO.__new__(
+        SparseTensorCOO)
+    sparse_tensor.init(sp_sparse_tensor)
+    return sparse_tensor
+
+
+cdef api bint pyarrow_is_sparse_tensor_csr(object sparse_tensor):
+    return isinstance(sparse_tensor, SparseTensorCSR)
+
+cdef api shared_ptr[CSparseTensorCSR] pyarrow_unwrap_sparse_tensor_csr(
+        object sparse_tensor):
+    cdef SparseTensorCSR sten
+    if pyarrow_is_sparse_tensor_csr(sparse_tensor):
+        sten = <SparseTensorCSR>(sparse_tensor)
+        return sten.sp_sparse_tensor
+
+    return shared_ptr[CSparseTensorCSR]()
+
+cdef api object pyarrow_wrap_sparse_tensor_csr(
+        const shared_ptr[CSparseTensorCSR]& sp_sparse_tensor):
+    if sp_sparse_tensor.get() == NULL:
+        raise ValueError('SparseTensorCSR was NULL')
+
+    cdef SparseTensorCSR sparse_tensor = SparseTensorCSR.__new__(
+        SparseTensorCSR)
+    sparse_tensor.init(sp_sparse_tensor)
+    return sparse_tensor
+
+
 cdef api bint pyarrow_is_column(object column):
     return isinstance(column, Column)
 

python/pyarrow/tests/test_sparse_tensor.py

@@ -0,0 +1,187 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pytest
+
+import numpy as np
+import pyarrow as pa
+
+
+@pytest.mark.parametrize('sparse_tensor_type', [
+    pa.SparseTensorCSR,
+    pa.SparseTensorCOO,
+])
+def test_sparse_tensor_attrs(sparse_tensor_type):
+    data = np.array([
+        [0, 1, 0, 0, 1],
+        [0, 0, 0, 0, 0],
+        [0, 0, 0, 1, 0],
+        [0, 0, 0, 0, 0],
+        [0, 3, 0, 0, 0],
+    ])
+    sparse_tensor = sparse_tensor_type.from_dense_numpy(data)
+
+    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.non_zero_length == 4
+
+
+@pytest.mark.parametrize('dtype_str,arrow_type', [
+    ('i1', pa.int8()),
+    ('i2', pa.int16()),
+    ('i4', pa.int32()),
+    ('i8', pa.int64()),
+    ('u1', pa.uint8()),
+    ('u2', pa.uint16()),
+    ('u4', pa.uint32()),
+    ('u8', pa.uint64()),
+    ('f2', pa.float16()),
+    ('f4', pa.float32()),
+    ('f8', pa.float64())
+])
+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)
+    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)
+
+    # Test from numpy array
+    sparse_tensor = pa.SparseTensorCOO.from_dense_numpy(data)
+    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()
+
+    # 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()
+
+
+@pytest.mark.parametrize('dtype_str,arrow_type', [
+    ('i1', pa.int8()),
+    ('i2', pa.int16()),
+    ('i4', pa.int32()),
+    ('i8', pa.int64()),
+    ('u1', pa.uint8()),
+    ('u2', pa.uint16()),
+    ('u4', pa.uint32()),
+    ('u8', pa.uint64()),
+    ('f2', pa.float16()),
+    ('f4', pa.float32()),
+    ('f8', pa.float64())
+])
+def test_sparse_tensor_csr_from_dense(dtype_str, arrow_type):
+    dtype = np.dtype(dtype_str)
+    dense_data = np.array([[1, 0, 2],
+                           [0, 0, 3],
+                           [4, 5, 6]]).astype(dtype)
+
+    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])
+    tensor = pa.Tensor.from_numpy(dense_data)
+
+    # Test from numpy array
+    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()
+
+    # 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()
+
+
+@pytest.mark.skip
+@pytest.mark.parametrize('dtype_str,arrow_type', [
+    ('i1', pa.int8()),
+    ('i2', pa.int16()),
+    ('i4', pa.int32()),
+    ('i8', pa.int64()),
+    ('u1', pa.uint8()),
+    ('u2', pa.uint16()),
+    ('u4', pa.uint32()),
+    ('u8', pa.uint64()),
+    ('f2', pa.float16()),
+    ('f4', pa.float32()),
+    ('f8', pa.float64())
+])
+def test_sparse_tensor_coo_numpy_roundtrip(dtype_str, arrow_type):
+    dtype = np.dtype(dtype_str)
+    coords = np.array([[0, 0], [3, 3], [1, 1], [0, 2]]).astype('i8')
+    data = np.array([4, 5, 7, 9]).astype(dtype)
+    shape = (4, 4)
+    dim_names = ["x", "y"]
+
+    sparse_tensor = pa.SparseTensorCOO.from_numpy(data, coords, shape,
+                                                  dim_names)
+    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()
+
+
+@pytest.mark.parametrize('dtype_str,arrow_type', [
+    ('i1', pa.int8()),
+    ('i2', pa.int16()),
+    ('i4', pa.int32()),
+    ('i8', pa.int64()),
+    ('u1', pa.uint8()),
+    ('u2', pa.uint16()),
+    ('u4', pa.uint32()),
+    ('u8', pa.uint64()),
+    ('f2', pa.float16()),
+    ('f4', pa.float32()),
+    ('f8', pa.float64())
+])
+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')
+    shape = (3, 3)
+    dim_names = ["x", "y"]
+
+    sparse_tensor = pa.SparseTensorCSR.from_numpy(data, indptr, indices,
+                                                  shape, dim_names)
+    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()