PERF: Sparse Series to scipy COO sparse matrix

asv_bench/benchmarks/sparse.py

@@ -67,16 +67,28 @@ def time_sparse_series_from_coo(self):
 
 
 class ToCoo:
- def setup(self):
+ params = [True, False]
+ param_names = ["sort_labels"]
+
+ def setup(self, sort_labels):
  s = Series([np.nan] * 10000)
  s[0] = 3.0
  s[100] = -1.0
  s[999] = 12.1
- s.index = MultiIndex.from_product([range(10)] * 4)
- self.ss = s.astype("Sparse")
 
- def time_sparse_series_to_coo(self):
- self.ss.sparse.to_coo(row_levels=[0, 1], column_levels=[2, 3], sort_labels=True)
+ s_mult_lvl = s.set_axis(MultiIndex.from_product([range(10)] * 4))
+ self.ss_mult_lvl = s_mult_lvl.astype("Sparse")
+
+ s_two_lvl = s.set_axis(MultiIndex.from_product([range(100)] * 2))
+ self.ss_two_lvl = s_two_lvl.astype("Sparse")
+
+ def time_sparse_series_to_coo(self, sort_labels):
+ self.ss_mult_lvl.sparse.to_coo(
+ row_levels=[0, 1], column_levels=[2, 3], sort_labels=sort_labels
+ )
+
+ def time_sparse_series_to_coo_single_level(self, sort_labels):
+ self.ss_two_lvl.sparse.to_coo(sort_labels=sort_labels)
 
 
 class Arithmetic:

pandas/core/arrays/sparse/scipy_sparse.py

@@ -3,10 +3,9 @@
 
 Currently only includes to_coo helpers.
 """
-from pandas.core.indexes.api import (
- Index,
- MultiIndex,
-)
+import numpy as np
+
+from pandas.core.indexes.api import MultiIndex
 from pandas.core.series import Series
 
 
@@ -19,6 +18,34 @@ def _check_is_partition(parts, whole):
  raise ValueError("Is not a partition because union is not the whole.")
 
 
+def _levels_to_axis(levels_codes, levels_labels, valid_ilocs, sort_labels=False):
+ if sort_labels and levels_codes.shape[0] == 1:
+ ax_coords = levels_codes[0][valid_ilocs]
+ ax_labels = levels_labels[0].tolist()
+
+ else:
+ # Why return_index anyway : https://github.com/numpy/numpy/issues/16923
+ ucodes, ucodes_idx, ucodes_inv = np.unique(
+ levels_codes.T, axis=0, return_index=True, return_inverse=True
+ )
+
+ if sort_labels:
+ ax_coords = ucodes_inv[valid_ilocs]
+
+ else:
+ og_order = np.argsort(ucodes_idx)
+ ucodes = ucodes[og_order, :]
+ ax_coords = og_order.argsort()[ucodes_inv[valid_ilocs]]
+
+ ax_labels = list(
+ zip(
+ *(tuple(lbls[ucodes[:, lvl]]) for lvl, lbls in enumerate(levels_labels))
+ )
+ )
+
+ return ax_coords, ax_labels
+
+
 def _to_ijv(ss, row_levels=(0,), column_levels=(1,), sort_labels=False):
  """
  For arbitrary (MultiIndexed) sparse Series return
@@ -27,65 +54,26 @@ def _to_ijv(ss, row_levels=(0,), column_levels=(1,), sort_labels=False):
  """
  # index and column levels must be a partition of the index
  _check_is_partition([row_levels, column_levels], range(ss.index.nlevels))
-
  # from the sparse Series: get the labels and data for non-null entries
  values = ss.array._valid_sp_values
 
- nonnull_labels = ss.dropna()
-
- def get_indexers(levels):
- """Return sparse coords and dense labels for subset levels"""
- # TODO: how to do this better? cleanly slice nonnull_labels given the
- # coord
- values_ilabels = [tuple(x[i] for i in levels) for x in nonnull_labels.index]
- if len(levels) == 1:
- values_ilabels = [x[0] for x in values_ilabels]
-
- # # performance issues with groupby ###################################
- # TODO: these two lines can replace the code below but
- # groupby is too slow (in some cases at least)
- # labels_to_i = ss.groupby(level=levels, sort=sort_labels).first()
- # labels_to_i[:] = np.arange(labels_to_i.shape[0])
-
- def _get_label_to_i_dict(labels, sort_labels=False):
- """
- Return dict of unique labels to number.
- Optionally sort by label.
- """
- labels = Index(map(tuple, labels)).unique().tolist() # squish
- if sort_labels:
- labels = sorted(labels)
- return {k: i for i, k in enumerate(labels)}
-
- def _get_index_subset_to_coord_dict(index, subset, sort_labels=False):
- ilabels = list(zip(*(index._get_level_values(i) for i in subset)))
- labels_to_i = _get_label_to_i_dict(ilabels, sort_labels=sort_labels)
- labels_to_i = Series(labels_to_i)
- if len(subset) > 1:
- labels_to_i.index = MultiIndex.from_tuples(labels_to_i.index)
- labels_to_i.index.names = [index.names[i] for i in subset]
- else:
- labels_to_i.index = Index(x[0] for x in labels_to_i.index)
- labels_to_i.index.name = index.names[subset[0]]
-
- labels_to_i.name = "value"
- return labels_to_i
-
- labels_to_i = _get_index_subset_to_coord_dict(
- ss.index, levels, sort_labels=sort_labels
- )
- # #####################################################################
- # #####################################################################
+ codes = ss.index.codes
+ labels = ss.index.levels
+ valid_ilocs = np.where(ss.notnull())[0]
 
- i_coord = labels_to_i[values_ilabels].tolist()
- i_labels = labels_to_i.index.tolist()
-
- return i_coord, i_labels
+ row_labels = [labels[lvl] for lvl in row_levels]
+ row_codes = np.asarray([codes[lvl] for lvl in row_levels])
+ i_coords, i_labels = _levels_to_axis(
+ row_codes, row_labels, valid_ilocs, sort_labels=sort_labels
+ )
 
- i_coord, i_labels = get_indexers(row_levels)
- j_coord, j_labels = get_indexers(column_levels)
+ col_labels = [labels[lvl] for lvl in column_levels]
+ col_codes = np.asarray([codes[lvl] for lvl in column_levels])
+ j_coords, j_labels = _levels_to_axis(
+ col_codes, col_labels, valid_ilocs, sort_labels=sort_labels
+ )
 
- return values, i_coord, j_coord, i_labels, j_labels
+ return values, i_coords, j_coords, i_labels, j_labels
 
 
 def sparse_series_to_coo(ss, row_levels=(0,), column_levels=(1,), sort_labels=False):
@@ -97,7 +85,7 @@ def sparse_series_to_coo(ss, row_levels=(0,), column_levels=(1,), sort_labels=Fa
  import scipy.sparse
 
  if ss.index.nlevels < 2:
- raise ValueError("to_coo requires MultiIndex with nlevels > 2")
+ raise ValueError("to_coo requires MultiIndex with nlevels >= 2.")
  if not ss.index.is_unique:
  raise ValueError(
  "Duplicate index entries are not allowed in to_coo transformation."

pandas/tests/arrays/sparse/test_array.py

@@ -1198,13 +1198,42 @@ def test_from_coo(self):
  def test_to_coo(self):
  import scipy.sparse
 
- ser = pd.Series(
- [1, 2, 3],
- index=pd.MultiIndex.from_product([[0], [1, 2, 3]], names=["a", "b"]),
- dtype="Sparse[int]",
+ s = pd.Series([np.nan] * 6)
+ s[2] = 1
+ s[5] = 3
+ s.index = pd.MultiIndex.from_tuples(
+ [
+ ("b", 2, "z", 1),
+ ("a", 2, "z", 2),
+ ("a", 2, "z", 1),
+ ("a", 2, "x", 2),
+ ("b", 1, "z", 1),
+ ("a", 1, "z", 0),
+ ]
+ )
+ ss = s.astype("Sparse")
+
+ expected_A = np.zeros((4, 4))
+ expected_A[1, 0] = 1
+ expected_A[3, 3] = 3
+ A, rows, cols = ss.sparse.to_coo(
+ row_levels=(0, 1), column_levels=(2, 3), sort_labels=False
+ )
+ assert isinstance(A, scipy.sparse.coo.coo_matrix)
+ assert np.all(A.toarray() == expected_A)
+ assert rows == [("b", 2), ("a", 2), ("b", 1), ("a", 1)]
+ assert cols == [("z", 1), ("z", 2), ("x", 2), ("z", 0)]
+
+ expected_A = np.zeros((4, 4))
+ expected_A[1, 2] = 1
+ expected_A[0, 1] = 3
+ A, rows, cols = ss.sparse.to_coo(
+ row_levels=(0, 1), column_levels=(2, 3), sort_labels=True
  )
- A, _, _ = ser.sparse.to_coo()
  assert isinstance(A, scipy.sparse.coo.coo_matrix)
+ assert np.all(A.toarray() == expected_A)
+ assert rows == [("a", 1), ("a", 2), ("b", 1), ("b", 2)]
+ assert cols == [("x", 2), ("z", 0), ("z", 1), ("z", 2)]
 
  def test_non_sparse_raises(self):
  ser = pd.Series([1, 2, 3])