ENH: cython count_level function, cleanup and tests

Author: wesm

pandas/core/index.py

@@ -1217,6 +1217,8 @@ def sortlevel(self, level=0, ascending=True):
         -------
         sorted_index : MultiIndex
         """
+        # TODO: check if lexsorted when level=0
+
         labels = list(self.labels)
         level = self._get_level_number(level)
         primary = labels.pop(level)

pandas/core/series.py

@@ -573,36 +573,17 @@ def count(self, level=None):
 
     def _count_level(self, level):
         # TODO: GENERALIZE CODE OVERLAP WITH DATAFRAME
-        # TODO: deal with sortedness??
-        obj = self.sortlevel(level)
-        mask = notnull(obj.values)
-
-        level_index = obj.index.levels[level]
+        mask = notnull(self.values)
+        level_index = self.index.levels[level]
 
         if len(self) == 0:
             return Series(0, index=level_index)
 
-        n = len(level_index)
-        locs = obj.index.labels[level].searchsorted(np.arange(n))
-
-        # WORKAROUND: reduceat fusses about the endpoints. should file ticket?
-        start = locs.searchsorted(0, side='right') - 1
-        end = locs.searchsorted(len(mask), side='left')
-
-        result = np.zeros((n), dtype=int)
-        out = result[start:end]
-        np.add.reduceat(mask, locs[start:end], out=out)
-
-        # WORKAROUND: to see why, try this
-        # arr = np.ones((10, 4), dtype=bool)
-        # np.add.reduceat(arr, [0, 3, 3, 7, 9], axis=0)
-
-        # this stinks
-        if len(locs) > 1:
-            workaround_mask = locs[:-1] == locs[1:]
-            result[:-1][workaround_mask] = 0
-
-        return Series(result, index=level_index)
+        # call cython function
+        max_bin = len(level_index)
+        counts = lib.count_level_1d(mask.view(np.uint8),
+                                    self.index.labels[level], max_bin)
+        return Series(counts, index=level_index)
 
     def value_counts(self):
         """

pandas/src/groupby.pyx

@@ -453,6 +453,23 @@ def _bucket_locs(index, buckets, inclusive=False):
 
     return locs
 
+def count_level_1d(ndarray[uint8_t, cast=True] mask,
+                  ndarray[int32_t] labels, Py_ssize_t max_bin):
+    cdef:
+        Py_ssize_t i, n
+        ndarray[int64_t] counts
+
+    counts = np.zeros(max_bin, dtype='i8')
+
+    n = len(mask)
+
+    for i from 0 <= i < n:
+        if mask[i]:
+            counts[labels[i]] += 1
+
+    return counts
+
+
 '''
 
 def ts_upsample_mean(ndarray[object] indices,

pandas/tests/test_multilevel.py

@@ -281,6 +281,24 @@ def _check_counts(frame, axis=0):
         df = tm.makeTimeDataFrame()
         self.assertRaises(Exception, df.count, level=0)
 
+    def test_count_level_series(self):
+        index = MultiIndex(levels=[['foo', 'bar', 'baz'],
+                                   ['one', 'two', 'three', 'four']],
+                           labels=[[0, 0, 0, 2, 2],
+                                   [2, 0, 1, 1, 2]])
+
+        s = Series(np.random.randn(len(index)), index=index)
+
+        result = s.count(level=0)
+        expected = s.groupby(level=0).count()
+        assert_series_equal(result.astype('f8'),
+                            expected.reindex(result.index).fillna(0))
+
+        result = s.count(level=1)
+        expected = s.groupby(level=1).count()
+        assert_series_equal(result.astype('f8'),
+                            expected.reindex(result.index).fillna(0))
+
     def test_count_level_corner(self):
         s = self.frame['A'][:0]
         result = s.count(level=0)