Fix 14072 pivot_table dropna

Author: OXPHOS

pandas/_libs/hashtable_class_helper.pxi.in

@@ -330,7 +330,7 @@ cdef class {{name}}HashTable(HashTable):
  @cython.boundscheck(False)
  def get_labels(self, {{dtype}}_t[:] values, {{name}}Vector uniques,
  Py_ssize_t count_prior, Py_ssize_t na_sentinel,
- bint check_null=True):
+ bint check_null=True, bint dropna=True):
  cdef:
  Py_ssize_t i, n = len(values)
  int64_t[:] labels
@@ -642,7 +642,7 @@ cdef class StringHashTable(HashTable):
  @cython.boundscheck(False)
  def get_labels(self, ndarray[object] values, ObjectVector uniques,
  Py_ssize_t count_prior, int64_t na_sentinel,
- bint check_null=1):
+ bint check_null=1, bint dropna=True):
  cdef:
  Py_ssize_t i, n = len(values)
  int64_t[:] labels
@@ -815,7 +815,7 @@ cdef class PyObjectHashTable(HashTable):
 
  def get_labels(self, ndarray[object] values, ObjectVector uniques,
  Py_ssize_t count_prior, int64_t na_sentinel,
- bint check_null=True):
+ bint check_null=True, bint dropna=True):
  cdef:
  Py_ssize_t i, n = len(values)
  int64_t[:] labels
@@ -830,7 +830,11 @@ cdef class PyObjectHashTable(HashTable):
  val = values[i]
  hash(val)
 
- if check_null and val != val or val is None:
+ if check_null and val != val:
+ labels[i] = na_sentinel
+ continue
+
+ if dropna and val is None:
  labels[i] = na_sentinel
  continue
 
@@ -968,5 +972,5 @@ cdef class MultiIndexHashTable(HashTable):
 
  def get_labels(self, object mi, ObjectVector uniques,
  Py_ssize_t count_prior, int64_t na_sentinel,
- bint check_null=True):
+ bint check_null=True, bint dropna=True):
  raise NotImplementedError

pandas/core/algorithms.py

@@ -519,7 +519,8 @@ def sort_mixed(values):
  return ordered, _ensure_platform_int(new_labels)
 
 
-def factorize(values, sort=False, order=None, na_sentinel=-1, size_hint=None):
+def factorize(values, sort=False, order=None, na_sentinel=-1, size_hint=None,
+ dropna=False):
  """
  Encode input values as an enumerated type or categorical variable
 
@@ -552,7 +553,8 @@ def factorize(values, sort=False, order=None, na_sentinel=-1, size_hint=None):
  table = hash_klass(size_hint or len(values))
  uniques = vec_klass()
  check_nulls = not is_integer_dtype(original)
- labels = table.get_labels(values, uniques, 0, na_sentinel, check_nulls)
+ labels = table.get_labels(values, uniques, 0, na_sentinel, check_nulls,
+ dropna)
 
  labels = _ensure_platform_int(labels)
  uniques = uniques.to_array()

pandas/core/categorical.py

@@ -548,10 +548,6 @@ def _validate_categories(cls, categories, fastpath=False):
 
  if not fastpath:
 
- # Categories cannot contain NaN.
- if categories.hasnans:
- raise ValueError('Categorial categories cannot be null')
-
  # Categories must be unique.
  if not categories.is_unique:
  raise ValueError('Categorical categories must be unique')

pandas/core/generic.py

@@ -4217,7 +4217,7 @@ def clip_lower(self, threshold, axis=None):
  return self.where(subset, threshold, axis=axis)
 
  def groupby(self, by=None, axis=0, level=None, as_index=True, sort=True,
- group_keys=True, squeeze=False, **kwargs):
+ group_keys=True, squeeze=False, dropna=True, **kwargs):
  """
  Group series using mapper (dict or key function, apply given function
  to group, return result as series) or by a series of columns.
@@ -4273,7 +4273,7 @@ def groupby(self, by=None, axis=0, level=None, as_index=True, sort=True,
  axis = self._get_axis_number(axis)
  return groupby(self, by=by, axis=axis, level=level, as_index=as_index,
  sort=sort, group_keys=group_keys, squeeze=squeeze,
- **kwargs)
+ dropna=dropna, **kwargs)
 
  def asfreq(self, freq, method=None, how=None, normalize=False,
  fill_value=None):

pandas/core/groupby.py

@@ -361,7 +361,8 @@ class _GroupBy(PandasObject, SelectionMixin):
 
  def __init__(self, obj, keys=None, axis=0, level=None,
  grouper=None, exclusions=None, selection=None, as_index=True,
- sort=True, group_keys=True, squeeze=False, **kwargs):
+ sort=True, group_keys=True, squeeze=False, dropna=True,
+ **kwargs):
 
  self._selection = selection
 
@@ -388,7 +389,8 @@ def __init__(self, obj, keys=None, axis=0, level=None,
  axis=axis,
  level=level,
  sort=sort,
- mutated=self.mutated)
+ mutated=self.mutated,
+ dropna=dropna)
 
  self.obj = obj
  self.axis = obj._get_axis_number(axis)
@@ -1614,15 +1616,15 @@ def tail(self, n=5):
 
 
 @Appender(GroupBy.__doc__)
-def groupby(obj, by, **kwds):
+def groupby(obj, by, dropna=True, **kwds):
  if isinstance(obj, Series):
  klass = SeriesGroupBy
  elif isinstance(obj, DataFrame):
  klass = DataFrameGroupBy
  else: # pragma: no cover
  raise TypeError('invalid type: %s' % type(obj))
 
- return klass(obj, by, **kwds)
+ return klass(obj, by, dropna=dropna, **kwds)
 
 
 def _get_axes(group):
@@ -2339,7 +2341,7 @@ class Grouping(object):
  """
 
  def __init__(self, index, grouper=None, obj=None, name=None, level=None,
- sort=True, in_axis=False):
+ sort=True, in_axis=False, dropna=True):
 
  self.name = name
  self.level = level
@@ -2348,6 +2350,7 @@ def __init__(self, index, grouper=None, obj=None, name=None, level=None,
  self.sort = sort
  self.obj = obj
  self.in_axis = in_axis
+ self.dropna = dropna
 
  # right place for this?
  if isinstance(grouper, (Series, Index)) and name is None:
@@ -2468,7 +2471,7 @@ def group_index(self):
  def _make_labels(self):
  if self._labels is None or self._group_index is None:
  labels, uniques = algorithms.factorize(
- self.grouper, sort=self.sort)
+ self.grouper, sort=self.sort, dropna=self.dropna)
  uniques = Index(uniques, name=self.name)
  self._labels = labels
  self._group_index = uniques
@@ -2480,7 +2483,7 @@ def groups(self):
 
 
 def _get_grouper(obj, key=None, axis=0, level=None, sort=True,
- mutated=False):
+ mutated=False, dropna=True):
  """
  create and return a BaseGrouper, which is an internal
  mapping of how to create the grouper indexers.
@@ -2633,7 +2636,8 @@ def is_in_obj(gpr):
  name=name,
  level=level,
  sort=sort,
- in_axis=in_axis) \
+ in_axis=in_axis,
+ dropna=dropna) \
  if not isinstance(gpr, Grouping) else gpr
 
  groupings.append(ping)

pandas/core/reshape/pivot.py

@@ -132,7 +132,7 @@ def pivot_table(data, values=None, index=None, columns=None, aggfunc='mean',
  pass
  values = list(values)
 
- grouped = data.groupby(keys)
+ grouped = data.groupby(keys, dropna=dropna)
  agged = grouped.agg(aggfunc)
 
  table = agged
@@ -159,15 +159,15 @@ def pivot_table(data, values=None, index=None, columns=None, aggfunc='mean',
  if isinstance(table, DataFrame):
  table = table.sort_index(axis=1)
 
- if fill_value is not None:
- table = table.fillna(value=fill_value, downcast='infer')
-
  if margins:
  if dropna:
  data = data[data.notnull().all(axis=1)]
  table = _add_margins(table, data, values, rows=index,
  cols=columns, aggfunc=aggfunc,
- margins_name=margins_name)
+ margins_name=margins_name, dropna=dropna)
+
+ if fill_value is not None:
+ table = table.fillna(value=fill_value, downcast='infer')
 
  # discard the top level
  if values_passed and not values_multi and not table.empty and \
@@ -188,7 +188,7 @@ def pivot_table(data, values=None, index=None, columns=None, aggfunc='mean',
 
 
 def _add_margins(table, data, values, rows, cols, aggfunc,
- margins_name='All'):
+ margins_name='All', dropna=True):
  if not isinstance(margins_name, compat.string_types):
  raise ValueError('margins_name argument must be a string')
 
@@ -219,7 +219,8 @@ def _add_margins(table, data, values, rows, cols, aggfunc,
  marginal_result_set = _generate_marginal_results(table, data, values,
  rows, cols, aggfunc,
  grand_margin,
- margins_name)
+ margins_name,
+ dropna=dropna)
  if not isinstance(marginal_result_set, tuple):
  return marginal_result_set
  result, margin_keys, row_margin = marginal_result_set
@@ -277,8 +278,7 @@ def _compute_grand_margin(data, values, aggfunc,
 
 
 def _generate_marginal_results(table, data, values, rows, cols, aggfunc,
- grand_margin,
- margins_name='All'):
+ grand_margin, margins_name='All', dropna=True):
  if len(cols) > 0:
  # need to "interleave" the margins
  table_pieces = []
@@ -288,7 +288,8 @@ def _all_key(key):
  return (key, margins_name) + ('',) * (len(cols) - 1)
 
  if len(rows) > 0:
- margin = data[rows + values].groupby(rows).agg(aggfunc)
+ margin = data[rows +
+ values].groupby(rows, dropna=dropna).agg(aggfunc)
  cat_axis = 1
 
  for key, piece in table.groupby(level=0, axis=cat_axis):
@@ -325,7 +326,8 @@ def _all_key(key):
  margin_keys = table.columns
 
  if len(cols) > 0:
- row_margin = data[cols + values].groupby(cols).agg(aggfunc)
+ row_margin = data[cols +
+ values].groupby(cols, dropna=dropna).agg(aggfunc)
  row_margin = row_margin.stack()
 
  # slight hack

pandas/tests/reshape/test_pivot.py

@@ -90,6 +90,39 @@ def test_pivot_table_dropna(self):
  tm.assert_index_equal(pv_col.columns, m)
  tm.assert_index_equal(pv_ind.index, m)
 
+ def test_pivot_table_dropna_margins(self):
+ # GH 14072
+ df = DataFrame([
+ [1, 'a', 'A'],
+ [1, 'b', 'B'],
+ [1, 'c', None]],
+ columns=['x', 'y', 'z'])
+
+ result_false = df.pivot_table(values='x', index='y', columns='z',
+ aggfunc='sum', fill_value=0,
+ margins=True, dropna=False)
+ expected_index = Series(['a', 'b', 'c', 'All'], name='y')
+ expected_columns = Series([None, 'A', 'B', 'All'], name='z')
+ expected_false = DataFrame([[0, 1, 0, 1],
+ [0, 0, 1, 1],
+ [1, 0, 0, 1],
+ [1, 1, 1, 3]],
+ index=expected_index,
+ columns=expected_columns)
+ tm.assert_frame_equal(expected_false, result_false)
+
+ result_true = df.pivot_table(values='x', index='y', columns='z',
+ aggfunc='sum', fill_value=0,
+ margins=True, dropna=True)
+ expected_index = Series(['a', 'b', 'All'], name='y')
+ expected_columns = Series(['A', 'B', 'All'], name='z')
+ expected_true = DataFrame([[1, 0, 1],
+ [0, 1, 1],
+ [1, 1, 2]],
+ index=expected_index,
+ columns=expected_columns)
+ tm.assert_frame_equal(expected_true, result_true)
+
  def test_pivot_table_dropna_categoricals(self):
  # GH 15193
  categories = ['a', 'b', 'c', 'd']