BUG: groupby agg over PeriodIndex unboxes to ordinals #2569

Author: changhiskhan

pandas/core/algorithms.py

@@ -115,13 +115,14 @@ def factorize(values, sort=False, order=None, na_sentinel=-1):
     Returns
     -------
     """
-    values = np.asarray(values)
-    is_datetime = com.is_datetime64_dtype(values)
-    (hash_klass, vec_klass), values = _get_data_algo(values, _hashtables)
+    from pandas.tseries.period import PeriodIndex
+    vals = np.asarray(values)
+    is_datetime = com.is_datetime64_dtype(vals)
+    (hash_klass, vec_klass), vals = _get_data_algo(vals, _hashtables)
 
-    table = hash_klass(len(values))
+    table = hash_klass(len(vals))
     uniques = vec_klass()
-    labels = table.get_labels(values, uniques, 0, na_sentinel)
+    labels = table.get_labels(vals, uniques, 0, na_sentinel)
 
     labels = com._ensure_platform_int(labels)
 
@@ -140,6 +141,8 @@ def factorize(values, sort=False, order=None, na_sentinel=-1):
 
     if is_datetime:
         uniques = uniques.view('M8[ns]')
+    if isinstance(values, PeriodIndex):
+        uniques = PeriodIndex(ordinal=uniques, freq=values.freq)
 
     return labels, uniques
 

pandas/core/index.py

@@ -86,6 +86,7 @@ class Index(np.ndarray):
     _engine_type = _index.ObjectEngine
 
     def __new__(cls, data, dtype=None, copy=False, name=None):
+        from pandas.tseries.period import PeriodIndex
         if isinstance(data, np.ndarray):
             if issubclass(data.dtype.type, np.datetime64):
                 from pandas.tseries.index import DatetimeIndex
@@ -100,6 +101,8 @@ def __new__(cls, data, dtype=None, copy=False, name=None):
                     data = np.array(data, dtype=dtype, copy=copy)
                 except TypeError:
                     pass
+            elif isinstance(data, PeriodIndex):
+                return PeriodIndex(data, copy=copy, name=name)
 
             if issubclass(data.dtype.type, np.integer):
                 return Int64Index(data, copy=copy, name=name)
@@ -122,9 +125,8 @@ def __new__(cls, data, dtype=None, copy=False, name=None):
                     from pandas.tseries.index import DatetimeIndex
                     return DatetimeIndex(subarr, copy=copy, name=name)
 
-                if lib.is_period_array(subarr):
-                    from pandas.tseries.period import PeriodIndex
-                    return PeriodIndex(subarr, name=name)
+            if lib.is_period_array(subarr):
+                return PeriodIndex(subarr, name=name)
 
         subarr = subarr.view(cls)
         subarr.name = name

pandas/tests/test_groupby.py

@@ -224,6 +224,13 @@ def test_agg_datetimes_mixed(self):
 
         assert(len(gb1) == len(gb2))
 
+    def test_agg_period_index(self):
+        from pandas import period_range, PeriodIndex
+        prng = period_range('2012-1-1', freq='M', periods=3)
+        df = DataFrame(np.random.randn(3, 2), index=prng)
+        rs = df.groupby(level=0).sum()
+        self.assert_(isinstance(rs.index, PeriodIndex))
+
     def test_agg_must_agg(self):
         grouped = self.df.groupby('A')['C']
         self.assertRaises(Exception, grouped.agg, lambda x: x.describe())

pandas/tests/test_index.py

@@ -86,6 +86,13 @@ def test_constructor_corner(self):
         # corner case
         self.assertRaises(Exception, Index, 0)
 
+    def test_index_ctor_infer_periodindex(self):
+        from pandas import period_range, PeriodIndex
+        xp = period_range('2012-1-1', freq='M', periods=3)
+        rs = Index(xp)
+        assert_array_equal(rs, xp)
+        self.assert_(isinstance(rs, PeriodIndex))
+
     def test_copy(self):
         i = Index([], name='Foo')
         i_copy = i.copy()