REF: add indices prop to sparse (pandas-dev#43827)

Author: mzeitlin11

pandas/_libs/sparse.pyi

@@ -17,6 +17,8 @@ class SparseIndex:
     def ngaps(self) -> int: ...
     @property
     def nbytes(self) -> int: ...
+    @property
+    def indices(self) -> npt.NDArray[np.int32]: ...
     def equals(self, other) -> bool: ...
     def lookup(self, index: int) -> np.int32: ...
     def lookup_array(self, indexer: npt.NDArray[np.int32]) -> npt.NDArray[np.int32]: ...

pandas/_libs/sparse.pyx

@@ -72,7 +72,7 @@ cdef class IntIndex(SparseIndex):
     def nbytes(self) -> int:
         return self.indices.nbytes
 
-    def check_integrity(self):
+    cdef check_integrity(self):
         """
         Checks the following:
 
@@ -118,7 +118,7 @@ cdef class IntIndex(SparseIndex):
     def ngaps(self) -> int:
         return self.length - self.npoints
 
-    def to_int_index(self):
+    cpdef to_int_index(self):
         return self
 
     def to_block_index(self):
@@ -327,7 +327,7 @@ cdef class BlockIndex(SparseIndex):
     def ngaps(self) -> int:
         return self.length - self.npoints
 
-    cpdef check_integrity(self):
+    cdef check_integrity(self):
         """
         Check:
         - Locations are in ascending order
@@ -375,7 +375,7 @@ cdef class BlockIndex(SparseIndex):
     def to_block_index(self):
         return self
 
-    def to_int_index(self):
+    cpdef to_int_index(self):
         cdef:
             int32_t i = 0, j, b
             int32_t offset
@@ -392,6 +392,10 @@ cdef class BlockIndex(SparseIndex):
 
         return IntIndex(self.length, indices)
 
+    @property
+    def indices(self):
+        return self.to_int_index().indices
+
     cpdef BlockIndex intersect(self, SparseIndex other):
         """
         Intersect two BlockIndex objects

pandas/core/arrays/sparse/accessor.py

@@ -344,7 +344,7 @@ def to_coo(self):
             if sp_arr.fill_value != 0:
                 raise ValueError("fill value must be 0 when converting to COO matrix")
 
-            row = sp_arr.sp_index.to_int_index().indices
+            row = sp_arr.sp_index.indices
             cols.append(np.repeat(col, len(row)))
             rows.append(row)
             data.append(sp_arr.sp_values.astype(dtype, copy=False))

pandas/core/arrays/sparse/array.py

@@ -574,7 +574,7 @@ def __array__(self, dtype: NpDtype | None = None) -> np.ndarray:
                 dtype = object
 
         out = np.full(self.shape, fill_value, dtype=dtype)
-        out[self.sp_index.to_int_index().indices] = self.sp_values
+        out[self.sp_index.indices] = self.sp_values
         return out
 
     def __setitem__(self, key, value):
@@ -796,7 +796,7 @@ def _first_fill_value_loc(self):
         if len(self) == 0 or self.sp_index.npoints == len(self):
             return -1
 
-        indices = self.sp_index.to_int_index().indices
+        indices = self.sp_index.indices
         if not len(indices) or indices[0] > 0:
             return 0
 
@@ -903,7 +903,7 @@ def __getitem__(
                 if end < 0:
                     end += len(self)
 
-                indices = self.sp_index.to_int_index().indices
+                indices = self.sp_index.indices
                 keep_inds = np.flatnonzero((indices >= start) & (indices < end))
                 sp_vals = self.sp_values[keep_inds]
 
@@ -1111,7 +1111,7 @@ def _concat_same_type(
             indices = []
 
             for arr in to_concat:
-                int_idx = arr.sp_index.to_int_index().indices.copy()
+                int_idx = arr.sp_index.indices.copy()
                 int_idx += length  # TODO: wraparound
                 length += arr.sp_index.length
 
@@ -1324,9 +1324,9 @@ def __setstate__(self, state):
 
     def nonzero(self):
         if self.fill_value == 0:
-            return (self.sp_index.to_int_index().indices,)
+            return (self.sp_index.indices,)
         else:
-            return (self.sp_index.to_int_index().indices[self.sp_values != 0],)
+            return (self.sp_index.indices[self.sp_values != 0],)
 
     # ------------------------------------------------------------------------
     # Reductions