TST: test MultiIndexPyIntEngine

Author: toobaz

pandas/tests/test_multilevel.py

@@ -1590,6 +1590,38 @@ def test_unstack_group_index_overflow(self):
         result = s.unstack(4)
         assert result.shape == (500, 2)
 
+    def test_pyint_engine(self):
+        # GH 18519 : when combinations of codes cannot be represented in 64
+        # bits, the index underlying the MultiIndex engine works with Python
+        # integers, rather than uint64.
+        N = 5
+        keys = [tuple(l) for l in [[0] * 10 * N,
+                                   [1] * 10 * N,
+                                   [2] * 10 * N,
+                                   [np.nan] * N + [2] * 9 * N,
+                                   [0] * N + [2] * 9 * N,
+                                   [np.nan] * N + [2] * 8 * N + [0] * N]]
+        # Each level contains 4 elements (including NaN), so it is represented
+        # in 2 bits, for a total of 2*N*10 = 100 > 64 bits. If we were using a
+        # 64 bit engine and truncating the first levels, the fourth and fifth
+        # keys would collide; if truncating the last levels, the fifth and
+        # sixth; if rotating bits rather than shifting, the third and fifth.
+
+        for idx in range(len(keys)):
+            index = MultiIndex.from_tuples(keys)
+            assert index.get_loc(keys[idx]) == idx
+
+            expected = np.arange(idx + 1, dtype='int64')
+            result = index.get_indexer([keys[i] for i in expected])
+            tm.assert_numpy_array_equal(result, expected)
+
+        # With missing key:
+        idces = range(len(keys))
+        expected = np.array([-1] + list(idces), dtype='int64')
+        missing = tuple([0, 1] * 5 * N)
+        result = index.get_indexer([missing] + [keys[i] for i in idces])
+        tm.assert_numpy_array_equal(result, expected)
+
     def test_getitem_lowerdim_corner(self):
         pytest.raises(KeyError, self.frame.loc.__getitem__,
                       (('bar', 'three'), 'B'))