Merge pull request #3286 from jreback/GH3283

ENH/BUG: GH3283 allow pivot tables to downcast the output (e.g. float -> int) if possible

Author: jreback

RELEASE.rst

@@ -171,6 +171,8 @@ pandas 0.11.0
  - added the method ``select_column`` to select a single column from a table as a Series.
  - deprecated the ``unique`` method, can be replicated by ``select_column(key,column).unique()``
 
+ - Downcast on pivot if possible (GH3283_), adds argument ``downcast`` to ``fillna``
+
 **Bug Fixes**
 
  - Fix seg fault on empty data frame when fillna with ``pad`` or ``backfill``
@@ -376,6 +378,7 @@ pandas 0.11.0
 .. _GH3222: https://github.com/pydata/pandas/issues/3222
 .. _GH2641: https://github.com/pydata/pandas/issues/2641
 .. _GH3238: https://github.com/pydata/pandas/issues/3238
+.. _GH3283: https://github.com/pydata/pandas/issues/3283
 
 pandas 0.10.1
 =============

pandas/core/frame.py

@@ -3231,7 +3231,7 @@ def reorder_levels(self, order, axis=0):
  # Filling NA's
 
  def fillna(self, value=None, method=None, axis=0, inplace=False,
- limit=None):
+ limit=None, downcast=None):
  """
  Fill NA/NaN values using the specified method
 
@@ -3255,6 +3255,8 @@ def fillna(self, value=None, method=None, axis=0, inplace=False,
  a reference to the filled object, which is self if inplace=True
  limit : int, default None
  Maximum size gap to forward or backward fill
+ downcast : dict, default is None, a dict of item->dtype of what to 
+ downcast if possible
 
  See also
  --------
@@ -3300,7 +3302,8 @@ def fillna(self, value=None, method=None, axis=0, inplace=False,
  result[k].fillna(v, inplace=True)
  return result
  else:
- new_data = self._data.fillna(value, inplace=inplace)
+ new_data = self._data.fillna(value, inplace=inplace, 
+ downcast=downcast)
 
  if inplace:
  self._data = new_data

pandas/core/internals.py

@@ -28,6 +28,7 @@ class Block(object):
  is_bool = False
  is_object = False
  _can_hold_na = False
+ _downcast_dtype = None
 
  def __init__(self, values, items, ref_items, ndim=2):
  if issubclass(values.dtype.type, basestring):
@@ -205,7 +206,7 @@ def split_block_at(self, item):
  self.items[s:e].copy(),
  self.ref_items)
 
- def fillna(self, value, inplace=False):
+ def fillna(self, value, inplace=False, downcast=None):
  if not self._can_hold_na:
  if inplace:
  return self
@@ -216,10 +217,32 @@ def fillna(self, value, inplace=False):
  mask = com.isnull(new_values)
  np.putmask(new_values, mask, value)
 
- if inplace:
- return self
- else:
- return make_block(new_values, self.items, self.ref_items)
+ block = make_block(new_values, self.items, self.ref_items)
+ if downcast:
+ block = block.downcast()
+ return block
+
+ def downcast(self, dtypes = None):
+ """ try to downcast each item to the dict of dtypes if present """
+
+ if dtypes is None:
+ dtypes = dict()
+
+ values = self.values
+ blocks = []
+ for i, item in enumerate(self.items):
+
+ dtype = dtypes.get(item,self._downcast_dtype)
+ if dtype is None:
+ nv = _block_shape(values[i])
+ blocks.append(make_block(nv, [ item ], self.ref_items))
+ continue
+ 
+ nv = _possibly_downcast_to_dtype(values[i], np.dtype(dtype))
+ nv = _block_shape(nv)
+ blocks.append(make_block(nv, [ item ], self.ref_items))
+
+ return blocks
 
  def astype(self, dtype, copy = True, raise_on_error = True):
  """
@@ -563,6 +586,7 @@ def _try_cast_result(self, result):
  return _possibly_downcast_to_dtype(result, self.dtype)
 
 class FloatBlock(NumericBlock):
+ _downcast_dtype = 'int64'
 
  def _can_hold_element(self, element):
  if isinstance(element, np.ndarray):
@@ -974,6 +998,9 @@ def shift(self, *args, **kwargs):
  def fillna(self, *args, **kwargs):
  return self.apply('fillna', *args, **kwargs)
 
+ def downcast(self, *args, **kwargs):
+ return self.apply('downcast', *args, **kwargs)
+
  def astype(self, *args, **kwargs):
  return self.apply('astype', *args, **kwargs)
 

pandas/tools/pivot.py

@@ -111,7 +111,7 @@ def pivot_table(data, values=None, rows=None, cols=None, aggfunc='mean',
  table = table.sort_index(axis=1)
 
  if fill_value is not None:
- table = table.fillna(value=fill_value)
+ table = table.fillna(value=fill_value, downcast=True)
 
  if margins:
  table = _add_margins(table, data, values, rows=rows,

pandas/tools/tests/test_pivot.py

@@ -69,6 +69,26 @@ def test_pivot_table_multiple(self):
  expected = self.data.groupby(rows + [cols]).agg(np.mean).unstack()
  tm.assert_frame_equal(table, expected)
 
+ def test_pivot_dtypes(self):
+
+ # can convert dtypes
+ f = DataFrame({'a' : ['cat', 'bat', 'cat', 'bat'], 'v' : [1,2,3,4], 'i' : ['a','b','a','b']})
+ self.assert_(f.dtypes['v'] == 'int64')
+
+ z = pivot_table(f, values='v', rows=['a'], cols=['i'], fill_value=0, aggfunc=np.sum)
+ result = z.get_dtype_counts()
+ expected = Series(dict(int64 = 2))
+ tm.assert_series_equal(result, expected)
+
+ # cannot convert dtypes
+ f = DataFrame({'a' : ['cat', 'bat', 'cat', 'bat'], 'v' : [1.5,2.5,3.5,4.5], 'i' : ['a','b','a','b']})
+ self.assert_(f.dtypes['v'] == 'float64')
+
+ z = pivot_table(f, values='v', rows=['a'], cols=['i'], fill_value=0, aggfunc=np.mean)
+ result = z.get_dtype_counts()
+ expected = Series(dict(float64 = 2))
+ tm.assert_series_equal(result, expected)
+
  def test_pivot_multi_values(self):
  result = pivot_table(self.data, values=['D', 'E'],
  rows='A', cols=['B', 'C'], fill_value=0)