pivot_table very slow on Categorical data; how about an observed keyword argument? #24923

asv_bench/benchmarks/reshape.py

@@ -127,6 +127,10 @@ def setup(self):
  'value1': np.random.randn(N),
  'value2': np.random.randn(N),
  'value3': np.random.randn(N)})
+ self.df2 = DataFrame({'col1': list('abcde'), 'col2': list('fghij'),
+ 'col3': [1, 2, 3, 4, 5]})
+ self.df2.col1 = self.df2.col1.astype('category')
+ self.df2.col2 = self.df2.col2.astype('category')
 
  def time_pivot_table(self):
  self.df.pivot_table(index='key1', columns=['key2', 'key3'])
@@ -139,6 +143,14 @@ def time_pivot_table_margins(self):
  self.df.pivot_table(index='key1', columns=['key2', 'key3'],
  margins=True)
 
+ def time_pivot_table_categorical(self):
+ self.df2.pivot_table(index='col1', values='col3', columns='col2',
+ aggfunc=np.sum, fill_value=0)
+
+ def time_pivot_table_categorical_observed(self):
+ self.df2.pivot_table(index='col1', values='col3', columns='col2',
+ aggfunc=np.sum, fill_value=0, observed=True)
+
 
 class Crosstab:
 

doc/source/whatsnew/v0.25.0.rst

@@ -28,6 +28,7 @@ Other Enhancements
 - Indexing of ``DataFrame`` and ``Series`` now accepts zerodim ``np.ndarray`` (:issue:`24919`)
 - :meth:`Timestamp.replace` now supports the ``fold`` argument to disambiguate DST transition times (:issue:`25017`)
 - :meth:`DataFrame.at_time` and :meth:`Series.at_time` now support :meth:`datetime.time` objects with timezones (:issue:`24043`)
+- :meth:`DataFrame.pivot_table` now accepts an ``observed`` parameter which is passed to underlying calls to :meth:`DataFrame.groupby` to speed up grouping categorical data. (:issue:`24923`)
 - ``Series.str`` has gained :meth:`Series.str.casefold` method to removes all case distinctions present in a string (:issue:`25405`)
 - :meth:`DataFrame.set_index` now works for instances of ``abc.Iterator``, provided their output is of the same length as the calling frame (:issue:`22484`, :issue:`24984`)
 - :meth:`DatetimeIndex.union` now supports the ``sort`` argument. The behaviour of the sort parameter matches that of :meth:`Index.union` (:issue:`24994`)

pandas/core/frame.py

@@ -5701,6 +5701,12 @@ def pivot(self, index=None, columns=None, values=None):
  margins_name : string, default 'All'
  Name of the row / column that will contain the totals
  when margins is True.
+ observed : boolean, default False
+ This only applies if any of the groupers are Categoricals.
+ If True: only show observed values for categorical groupers.
+ If False: show all values for categorical groupers.
+
+ .. versionchanged :: 0.25.0
 
  Returns
  -------
@@ -5791,12 +5797,12 @@ def pivot(self, index=None, columns=None, values=None):
  @Appender(_shared_docs['pivot_table'])
  def pivot_table(self, values=None, index=None, columns=None,
  aggfunc='mean', fill_value=None, margins=False,
- dropna=True, margins_name='All'):
+ dropna=True, margins_name='All', observed=False):
  from pandas.core.reshape.pivot import pivot_table
  return pivot_table(self, values=values, index=index, columns=columns,
  aggfunc=aggfunc, fill_value=fill_value,
  margins=margins, dropna=dropna,
- margins_name=margins_name)
+ margins_name=margins_name, observed=observed)
 
  def stack(self, level=-1, dropna=True):
  """

pandas/core/reshape/pivot.py

@@ -22,7 +22,7 @@
 @Appender(_shared_docs['pivot_table'], indents=1)
 def pivot_table(data, values=None, index=None, columns=None, aggfunc='mean',
  fill_value=None, margins=False, dropna=True,
- margins_name='All'):
+ margins_name='All', observed=False):
  index = _convert_by(index)
  columns = _convert_by(columns)
 
@@ -34,7 +34,8 @@ def pivot_table(data, values=None, index=None, columns=None, aggfunc='mean',
  columns=columns,
  fill_value=fill_value, aggfunc=func,
  margins=margins, dropna=dropna,
- margins_name=margins_name)
+ margins_name=margins_name,
+ observed=observed)
  pieces.append(table)
  keys.append(getattr(func, '__name__', func))
 
@@ -77,7 +78,7 @@ def pivot_table(data, values=None, index=None, columns=None, aggfunc='mean',
  pass
  values = list(values)
 
- grouped = data.groupby(keys, observed=False)
+ grouped = data.groupby(keys, observed=observed)
  agged = grouped.agg(aggfunc)
  if dropna and isinstance(agged, ABCDataFrame) and len(agged.columns):
  agged = agged.dropna(how='all')

pandas/tests/reshape/test_pivot.py

@@ -37,18 +37,18 @@ def setup_method(self, method):
  'E': np.random.randn(11),
  'F': np.random.randn(11)})
 
- def test_pivot_table(self):
+ def test_pivot_table(self, observed):
  index = ['A', 'B']
  columns = 'C'
  table = pivot_table(self.data, values='D',
- index=index, columns=columns)
+ index=index, columns=columns, observed=observed)
 
  table2 = self.data.pivot_table(
- values='D', index=index, columns=columns)
+ values='D', index=index, columns=columns, observed=observed)
  tm.assert_frame_equal(table, table2)
 
  # this works
- pivot_table(self.data, values='D', index=index)
+ pivot_table(self.data, values='D', index=index, observed=observed)
 
  if len(index) > 1:
  assert table.index.names == tuple(index)
@@ -64,6 +64,28 @@ def test_pivot_table(self):
  index + [columns])['D'].agg(np.mean).unstack()
  tm.assert_frame_equal(table, expected)
 
+ def test_pivot_table_categorical_observed_equal(self, observed):
+ # issue #24923
+ df = pd.DataFrame({'col1': list('abcde'),
+ 'col2': list('fghij'),
+ 'col3': [1, 2, 3, 4, 5]})
+
+ expected = df.pivot_table(index='col1', values='col3',
+ columns='col2', aggfunc=np.sum,
+ fill_value=0)
+
+ expected.index = expected.index.astype('category')
+ expected.columns = expected.columns.astype('category')
+
+ df.col1 = df.col1.astype('category')
+ df.col2 = df.col2.astype('category')
+
+ result = df.pivot_table(index='col1', values='col3',
+ columns='col2', aggfunc=np.sum,
+ fill_value=0, observed=observed)
+
+ tm.assert_frame_equal(result, expected)
+
  def test_pivot_table_nocols(self):
  df = DataFrame({'rows': ['a', 'b', 'c'],
  'cols': ['x', 'y', 'z'],