BUG: detect/raise exception if possible int64 overflow #851

Author: wesm

pandas/core/groupby.py

@@ -1468,16 +1468,27 @@ def get_group_index(label_list, shape):
         return label_list[0]
 
     n = len(label_list[0])
-    group_index = np.zeros(n, dtype=int)
+    group_index = np.zeros(n, dtype=np.int64)
     mask = np.zeros(n, dtype=bool)
-    for i in xrange(len(shape)):
-        stride = np.prod([x for x in shape[i+1:]], dtype=int)
-        group_index += com._ensure_int64(label_list[i]) * stride
-        mask |= label_list[i] < 0
+
+    if _int64_overflow_possible(shape):
+        raise Exception('Possible int64 overflow, raise exception for now')
+    else:
+        for i in xrange(len(shape)):
+            stride = np.prod([x for x in shape[i+1:]], dtype=np.int64)
+            group_index += com._ensure_int64(label_list[i]) * stride
+            mask |= label_list[i] < 0
 
     np.putmask(group_index, mask, -1)
     return group_index
 
+def _int64_overflow_possible(shape):
+    the_prod = 1L
+    for x in shape:
+        the_prod *= long(x)
+
+    return the_prod >= 2**63
+
 def decons_group_index(comp_labels, shape):
     # reconstruct labels
     label_list = []

pandas/tests/test_groupby.py

@@ -1467,6 +1467,20 @@ def test_int32_overflow(self):
         right = df.groupby(['B', 'A']).sum()
         self.assert_(len(left) == len(right))
 
+    def test_int64_overflow(self):
+        B = np.concatenate((np.arange(100000), np.arange(100000),
+                            np.arange(50000)))
+        A = np.arange(250000)
+        df = DataFrame({'A' : A, 'B' : B,
+                        'C' : A, 'D' : B,
+                        'values' : np.random.randn(250000)})
+
+        self.assertRaises(Exception, df.groupby(['A', 'B', 'C', 'D']).sum)
+
+        # left = df.groupby(['A', 'B', 'C', 'D']).sum()
+        # right = df.groupby(['D', 'C', 'B', 'A']).sum()
+        # self.assert_(len(left) == len(right))
+
 def test_decons():
     from pandas.core.groupby import decons_group_index, get_group_index