Set base object in to_numpy methods.

cpp/src/arrow/python/numpy_convert.cc

@@ -315,17 +315,25 @@ Status SparseTensorCOOToNdarray(const std::shared_ptr<SparseTensorCOO>& sparse_t
     array_flags |= NPY_ARRAY_WRITEABLE;
   }
 
-  PyObject* data =
+  PyObject* result_data =
       PyArray_NewFromDescr(&PyArray_Type, dtype_data, 1, npy_shape_data.data(), nullptr,
                            mutable_data, array_flags, nullptr);
-  PyObject* coords =
+  PyObject* result_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;
+  if (base == Py_None || base == nullptr) {
+    base = py::wrap_sparse_tensor_coo(sparse_tensor);
+  } else {
+    Py_XINCREF(base);
+  }
+  PyArray_SetBaseObject(reinterpret_cast<PyArrayObject*>(result_data), base);
+  PyArray_SetBaseObject(reinterpret_cast<PyArrayObject*>(result_coords), base);
+
+  *out_data = result_data;
+  *out_coords = result_coords;
   return Status::OK();
 }
 
@@ -394,6 +402,15 @@ Status SparseTensorCSRToNdarray(const std::shared_ptr<SparseTensorCSR>& sparse_t
                            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);
+  }
+  PyArray_SetBaseObject(reinterpret_cast<PyArrayObject*>(result_data), base);
+  PyArray_SetBaseObject(reinterpret_cast<PyArrayObject*>(result_indptr), base);
+  PyArray_SetBaseObject(reinterpret_cast<PyArrayObject*>(result_indices), base);
+
   *out_data = result_data;
   *out_indptr = result_indptr;
   *out_indices = result_indices;

python/pyarrow/tests/test_sparse_tensor.py

@@ -16,6 +16,7 @@
 # under the License.
 
 import pytest
+import sys
 
 import numpy as np
 import pyarrow as pa
@@ -44,6 +45,20 @@ def test_sparse_tensor_attrs(sparse_tensor_type):
     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))
+    n = sys.getrefcount(sparse_tensor)
+    data, coords = sparse_tensor.to_numpy()  # noqa
+    assert sys.getrefcount(sparse_tensor) == n + 1
+
+
+def test_sparse_tensor_csr_base_object():
+    sparse_tensor = pa.SparseTensorCSR.from_dense_numpy(np.random.randn(10, 4))
+    n = sys.getrefcount(sparse_tensor)
+    data, indptr, indices = sparse_tensor.to_numpy()  # noqa
+    assert sys.getrefcount(sparse_tensor) == n + 1
+
+
 @pytest.mark.parametrize('dtype_str,arrow_type', [
     ('i1', pa.int8()),
     ('i2', pa.int16()),
@@ -126,7 +141,6 @@ def test_sparse_tensor_csr_from_dense(dtype_str, arrow_type):
     assert (indices == result_indices).all()
 
 
-@pytest.mark.skip
 @pytest.mark.parametrize('dtype_str,arrow_type', [
     ('i1', pa.int8()),
     ('i2', pa.int16()),
@@ -142,8 +156,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)
-    coords = np.array([[0, 0], [3, 3], [1, 1], [0, 2]]).astype('i8')
     data = np.array([4, 5, 7, 9]).astype(dtype)
+    coords = np.array([[0, 0], [3, 3], [1, 1], [0, 2]]).astype('i8')
     shape = (4, 4)
     dim_names = ["x", "y"]