Skip to content

REF: add indices prop to sparse #43827

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Oct 3, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions pandas/_libs/sparse.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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]: ...
Expand Down
12 changes: 8 additions & 4 deletions pandas/_libs/sparse.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -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:

Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down
2 changes: 1 addition & 1 deletion pandas/core/arrays/sparse/accessor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
12 changes: 6 additions & 6 deletions pandas/core/arrays/sparse/array.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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]

Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down