ENH/API: DataFrame.stack() support for level=None, sequentially=True/False, and NaN level values.

Author: seth-p

pandas/core/frame.py

@@ -3718,7 +3718,7 @@ def pivot(self, index=None, columns=None, values=None):
  from pandas.core.reshape import pivot
  return pivot(self, index=index, columns=columns, values=values)
 
- def stack(self, level=-1, dropna=True):
+ def stack(self, level=-1, dropna=True, sequentially=True):
  """
  Pivot a level of the (possibly hierarchical) column labels, returning a
  DataFrame (or Series in the case of an object with a single level of
@@ -3728,11 +3728,15 @@ def stack(self, level=-1, dropna=True):
 
  Parameters
  ----------
- level : int, string, or list of these, default last level
- Level(s) to stack, can pass level name
+ level : int, string, list of these, or None; default -1 (last level)
+ Level(s) to stack, can pass level name(s).
+ None specifies all column levels, i.e. list(range(columns.nlevels)).
  dropna : boolean, default True
  Whether to drop rows in the resulting Frame/Series with no valid
  values
+ sequentially : boolean, default True
+ When level is a list (or None), whether the multiple column levels
+ should be stacked sequentially (if True) or simultaneously (if False).
 
  Examples
  ----------
@@ -3751,14 +3755,20 @@ def stack(self, level=-1, dropna=True):
  -------
  stacked : DataFrame or Series
  """
- from pandas.core.reshape import stack, stack_multiple
+ from pandas.core.reshape import stack_levels_sequentially, stack_multi_levels_simultaneously
 
- if isinstance(level, (tuple, list)):
- return stack_multiple(self, level, dropna=dropna)
+ level_nums = self.columns._get_level_numbers(level, allow_mixed_names_and_numbers=False)
+ if level_nums == []:
+ if dropna:
+ return self.dropna(axis=0, how='all')
+ else:
+ return self
+ elif (not sequentially) and isinstance(self.columns, MultiIndex):
+ return stack_multi_levels_simultaneously(self, level_nums, dropna=dropna)
  else:
- return stack(self, level, dropna=dropna)
+ return stack_levels_sequentially(self, level_nums, dropna=dropna)
 
- def unstack(self, level=-1):
+ def unstack(self, level=-1, dropna=False, sequentially=False):
  """
  Pivot a level of the (necessarily hierarchical) index labels, returning
  a DataFrame having a new level of column labels whose inner-most level
@@ -3769,8 +3779,15 @@ def unstack(self, level=-1):
 
  Parameters
  ----------
- level : int, string, or list of these, default -1 (last level)
- Level(s) of index to unstack, can pass level name
+ level : int, string, list of these, or None; default -1 (last level)
+ Level(s) of index to unstack, can pass level name(s).
+ None specifies all index levels, i.e. list(range(index.nlevels)).
+ dropna : boolean, default False
+ Whether to drop columns in the resulting Frame/Series with no valid
+ values
+ sequentially : boolean, default True
+ When level is a list (or None), whether the multiple index levels
+ should be stacked sequentially (if True) or simultaneously (if False).
 
  See also
  --------
@@ -3812,7 +3829,44 @@ def unstack(self, level=-1):
  unstacked : DataFrame or Series
  """
  from pandas.core.reshape import unstack
- return unstack(self, level)
+
+ level_nums = self.index._get_level_numbers(level, allow_mixed_names_and_numbers=False)
+ if level_nums == []:
+ if dropna:
+ return self.dropna(axis=1, how='all')
+ else:
+ return self
+ if sequentially and isinstance(level_nums, list) and (len(level_nums) > 1):
+ result = self
+ # Adjust level_nums to account for the fact that levels move "up"
+ # as a result of stacking of earlier levels.
+ adjusted_level_nums = [x - sum((y < x) for y in level_nums[:i])
+ for i, x in enumerate(level_nums)]
+ for level_num in adjusted_level_nums:
+ result = unstack(result, level_num)
+ else:
+ result = unstack(self, level_nums)
+
+ if isinstance(result, DataFrame):
+ # fix dtypes, if necessary
+ desired_dtypes = self.dtypes.values.repeat(len(result.columns) // len(self.columns))
+ result_dtypes = result.dtypes.values
+ for i, c in enumerate(result.columns):
+ if result_dtypes[i] != desired_dtypes[i]:
+ if result_dtypes[i] == np.object:
+ # use default Series constructor to set type
+ result[c] = Series(result[c].values.tolist(), index=result.index)
+ else:
+ # try to convert type directly
+ result[c] = result[c].astype(desired_dtypes[i], raise_on_error=False)
+ # drop empty columns, if necessary
+ if dropna:
+ result = result.dropna(axis=1, how='all')
+ else:
+ if dropna:
+ result = result.dropna()
+
+ return result
 
  #----------------------------------------------------------------------
  # Time series-related

pandas/core/index.py

@@ -1033,7 +1033,7 @@ def _validate_index_level(self, level):
  verification must be done like in MultiIndex.
 
  """
- if isinstance(level, int):
+ if com.is_integer(level):
  if level < 0 and level != -1:
  raise IndexError("Too many levels: Index has only 1 level,"
  " %d is not a valid level number" % (level,))
@@ -1045,10 +1045,44 @@ def _validate_index_level(self, level):
  raise KeyError('Level %s must be same as name (%s)'
  % (level, self.name))
 
- def _get_level_number(self, level):
+ def _get_level_number(self, level, ignore_names=False):
+ """
+ Returns level number corresponding to level.
+ If level is a level name and ignore_names is False,
+ the level number corresponding to such level name is returned.
+ Otherwise level must be a number.
+ If level is a positive number, it is returned.
+ If level is a negative number, its sum with self.nlevels is returned.
+ """
+ if ignore_names and (not com.is_integer(level)):
+ raise KeyError('Level %s not found' % str(level))
  self._validate_index_level(level)
  return 0
 
+ def _get_level_numbers(self, levels, allow_mixed_names_and_numbers=False):
+ """
+ Returns level numbers corresponding to levels.
+ If levels is None, a list of all level numbers is returned.
+ If levels is a single number or level name,
+ then a single number is returned (using _get_level_number()).
+ If levels is a list of numbers or level names,
+ then a list of numbers is returned (each using _get_level_number()).
+ If allow_mixed_names_and_numbers is False, then levels must be
+ either all level numbers or all level names.
+ """
+ if levels is None:
+ return list(range(self.nlevels))
+ elif isinstance(levels, (list, tuple, set)):
+ if (not allow_mixed_names_and_numbers) and (not all(lev in self.names for lev in levels)):
+ if all(isinstance(lev, int) for lev in levels):
+ return type(levels)(self._get_level_number(level, ignore_names=True) for level in levels)
+ else:
+ raise ValueError("level should contain all level names or all level numbers, "
+ "not a mixture of the two.")
+ return type(levels)(self._get_level_number(level) for level in levels)
+ else:
+ return self._get_level_number(levels)
+
  @cache_readonly
  def inferred_type(self):
  """ return a string of the type inferred from the values """
@@ -4294,28 +4328,38 @@ def _from_elements(values, labels=None, levels=None, names=None,
  sortorder=None):
  return MultiIndex(levels, labels, names, sortorder=sortorder)
 
- def _get_level_number(self, level):
- try:
+ def _get_level_number(self, level, ignore_names=False):
+ """
+ Returns level number corresponding to level.
+ If level is a level name and ignore_names is False,
+ the level number corresponding to such level name is returned.
+ Otherwise level must be a number.
+ If level is a positive number, it is returned.
+ If level is a negative number, its sum with self.nlevels is returned.
+ """
+ if not ignore_names:
  count = self.names.count(level)
  if count > 1:
  raise ValueError('The name %s occurs multiple times, use a '
  'level number' % level)
- level = self.names.index(level)
- except ValueError:
- if not isinstance(level, int):
- raise KeyError('Level %s not found' % str(level))
- elif level < 0:
- level += self.nlevels
- if level < 0:
- orig_level = level - self.nlevels
- raise IndexError(
- 'Too many levels: Index has only %d levels, '
- '%d is not a valid level number' % (self.nlevels, orig_level)
- )
- # Note: levels are zero-based
- elif level >= self.nlevels:
- raise IndexError('Too many levels: Index has only %d levels, '
- 'not %d' % (self.nlevels, level + 1))
+ try:
+ return self.names.index(level)
+ except ValueError:
+ pass
+ if not com.is_integer(level):
+ raise KeyError('Level %s not found' % str(level))
+ elif level < 0:
+ level += self.nlevels
+ if level < 0:
+ orig_level = level - self.nlevels
+ raise IndexError(
+ 'Too many levels: Index has only %d levels, '
+ '%d is not a valid level number' % (self.nlevels, orig_level)
+ )
+ # Note: levels are zero-based
+ elif level >= self.nlevels:
+ raise IndexError('Too many levels: Index has only %d levels, '
+ 'not %d' % (self.nlevels, level + 1))
  return level
 
  _tuples = None
@@ -4891,14 +4935,16 @@ def _drop_from_level(self, labels, level):
 
  return self[mask]
 
- def droplevel(self, level=0):
+ def droplevel(self, level=0, ignore_names=False):
  """
  Return Index with requested level removed. If MultiIndex has only 2
  levels, the result will be of Index type not MultiIndex.
 
  Parameters
  ----------
  level : int/level name or list thereof
+ ignore_names : boolean, default True
+ If True, level must be an int or list thereof
 
  Notes
  -----
@@ -4916,7 +4962,7 @@ def droplevel(self, level=0):
  new_labels = list(self.labels)
  new_names = list(self.names)
 
- levnums = sorted(self._get_level_number(lev) for lev in levels)[::-1]
+ levnums = sorted((self._get_level_number(lev, ignore_names) for lev in levels), reverse=True)
 
  for i in levnums:
  new_levels.pop(i)
@@ -4929,6 +4975,9 @@ def droplevel(self, level=0):
  mask = new_labels[0] == -1
  result = new_levels[0].take(new_labels[0])
  if mask.any():
+ if result.is_integer():
+ # cannot store NaNs in an integer index, so promote to Float64Index
+ result = Float64Index(result.values, name=result.name)
  result = result.putmask(mask, np.nan)
 
  result.name = new_names[0]
@@ -5539,7 +5588,7 @@ def convert_indexer(start, stop, step, indexer=indexer, labels=labels):
 
  else:
 
- loc = level_index.get_loc(key)
+ loc = -1 if com.is_float(key) and np.isnan(key) else level_index.get_loc(key)
  if level > 0 or self.lexsort_depth == 0:
  return np.array(labels == loc,dtype=bool)
  else:
@@ -6050,7 +6099,7 @@ def _trim_front(strings):
 
 
 def _sanitize_and_check(indexes):
- kinds = list(set([type(index) for index in indexes]))
+ kinds = list(set(type(index) for index in indexes))
 
  if list in kinds:
  if len(kinds) > 1:
@@ -6071,11 +6120,11 @@ def _get_consensus_names(indexes):
 
  # find the non-none names, need to tupleify to make
  # the set hashable, then reverse on return
- consensus_names = set([
+ consensus_names = set(
  tuple(i.names) for i in indexes if all(n is not None for n in i.names)
- ])
+ )
  if len(consensus_names) == 1:
- return list(list(consensus_names)[0])
+ return list(consensus_names.pop())
  return [None] * indexes[0].nlevels