BUG: unstacking int64 overflow with many levels. re #2616

Author: wesm

RELEASE.rst

@@ -68,6 +68,7 @@ pandas 0.10.1
   - Fix bug and possible segfault when grouping by hierarchical level that
     contains NA values (GH2616_)
   - Ensure that MultiIndex tuples can be constructed with NAs (seen in #2616)
+  - Fix int64 overflow issue when unstacking MultiIndex with many levels (#2616)
 
 **API Changes**
 

pandas/core/reshape.py

@@ -83,8 +83,11 @@ def _make_sorted_values_labels(self):
         to_sort = labs[:v] + labs[v + 1:] + [labs[v]]
         sizes = [len(x) for x in levs[:v] + levs[v + 1:] + [levs[v]]]
 
-        group_index = get_group_index(to_sort, sizes)
-        comp_index, obs_ids = _compress_group_index(group_index)
+        comp_index, obs_ids = get_compressed_ids(to_sort, sizes)
+
+        # group_index = get_group_index(to_sort, sizes)
+        # comp_index, obs_ids = _compress_group_index(group_index)
+
         ngroups = len(obs_ids)
 
         indexer = algos.groupsort_indexer(comp_index, ngroups)[0]
@@ -97,10 +100,10 @@ def _make_selectors(self):
         new_levels = self.new_index_levels
 
         # make the mask
-        group_index = get_group_index(self.sorted_labels[:-1],
-                                      [len(x) for x in new_levels])
+        remaining_labels = self.sorted_labels[:-1]
+        level_sizes = [len(x) for x in new_levels]
 
-        comp_index, obs_ids = _compress_group_index(group_index)
+        comp_index, obs_ids = get_compressed_ids(remaining_labels, level_sizes)
         ngroups = len(obs_ids)
 
         comp_index = _ensure_platform_int(comp_index)
@@ -391,6 +394,36 @@ def _unstack_frame(obj, level):
                                value_columns=obj.columns)
         return unstacker.get_result()
 
+def get_compressed_ids(labels, sizes):
+    # no overflow
+    if _long_prod(sizes) < 2**63:
+        group_index = get_group_index(labels, sizes)
+        comp_index, obs_ids = _compress_group_index(group_index)
+    else:
+        n = len(labels[0])
+        mask = np.zeros(n, dtype=bool)
+        for v in labels:
+            mask |= v < 0
+
+        while _long_prod(sizes) >= 2**63:
+            i = len(sizes)
+            while _long_prod(sizes[:i]) >= 2**63:
+                i -= 1
+
+            rem_index, rem_ids = get_compressed_ids(labels[:i],
+                                                    sizes[:i])
+            sizes = [len(rem_ids)] + sizes[i:]
+            labels = [rem_index] + labels[i:]
+
+        return get_compressed_ids(labels, sizes)
+
+    return comp_index, obs_ids
+
+def _long_prod(vals):
+    result = 1L
+    for x in vals:
+        result *= x
+    return result
 
 def stack(frame, level=-1, dropna=True):
     """

pandas/tests/test_multilevel.py

@@ -1267,6 +1267,39 @@ def test_unstack_preserve_types(self):
         self.assert_(unstacked['E', 1].dtype == np.object_)
         self.assert_(unstacked['F', 1].dtype == np.float64)
 
+    def test_unstack_group_index_overflow(self):
+        labels = np.tile(np.arange(500), 2)
+        level = np.arange(500)
+
+        index = MultiIndex(levels=[level] * 8 + [[0, 1]],
+                           labels=[labels] * 8 + [np.arange(2).repeat(500)])
+
+        s = Series(np.arange(1000), index=index)
+        result = s.unstack()
+        self.assertEqual(result.shape, (500, 2))
+
+        # test roundtrip
+        stacked = result.stack()
+        assert_series_equal(s.astype(np.float64),
+                            stacked.reindex(s.index))
+
+        # put it at beginning
+        index = MultiIndex(levels=[[0, 1]] + [level] * 8,
+                           labels=[np.arange(2).repeat(500)] + [labels] * 8)
+
+        s = Series(np.arange(1000), index=index)
+        result = s.unstack(0)
+        self.assertEqual(result.shape, (500, 2))
+
+        # put it in middle
+        index = MultiIndex(levels=[level] * 4 + [[0, 1]] + [level] * 4,
+                           labels=([labels] * 4 + [np.arange(2).repeat(500)]
+                                   + [labels] * 4))
+
+        s = Series(np.arange(1000), index=index)
+        result = s.unstack(4)
+        self.assertEqual(result.shape, (500, 2))
+
     def test_getitem_lowerdim_corner(self):
         self.assertRaises(KeyError, self.frame.ix.__getitem__,
                           (('bar', 'three'), 'B'))