Consolidated group_any_all Cython func

Author: WillAyd

pandas/_libs/groupby.pyx

@@ -312,58 +312,23 @@ def group_fillna_indexer(ndarray[int64_t] out, ndarray[int64_t] labels,
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
-def group_any(ndarray[uint8_t] out,
-              ndarray[int64_t] labels,
-              ndarray[uint8_t] values,
-              ndarray[uint8_t] mask,
-              bint skipna):
-    """Aggregated boolean values to show if any group element is truthful
+def group_any_all(ndarray[uint8_t] out,
+                  ndarray[int64_t] labels,
+                  ndarray[uint8_t] values,
+                  ndarray[uint8_t] mask,
+                  object val_test,
+                  bint skipna):
+    """Aggregated boolean values to show truthfulness of group elements
 
     Parameters
     ----------
     out : array of values which this method will write its results to
-    labels : array containing unique label for each group, with its ordering
-        matching up to the corresponding record in `values`
-    values : array containing the truth value of each element
-    mask : array indicating whether a value is na or not
-    skipna : boolean
-        Flag to ignore nan values during truth testing
-
-    Notes
-    -----
-    This method modifies the `out` parameter rather than returning an object.
-    The returned values will either be 0 or 1 (False or True, respectively).
-    """
-    cdef:
-        Py_ssize_t i, N=len(labels)
-        int64_t lab
-
-    with nogil:
-        for i in range(N):
-            lab = labels[i]
-            if lab < 0 or (skipna and mask[i]):
-                continue
-
-            if values[i]:
-                out[lab] = 1
-
-
-@cython.boundscheck(False)
-@cython.wraparound(False)
-def group_all(ndarray[uint8_t] out,
-              ndarray[int64_t] labels,
-              ndarray[uint8_t] values,
-              ndarray[uint8_t] mask,
-              bint skipna):
-    """Aggregated boolean values to show if all group elements are truthful
-
-    Parameters
-    ----------
-    out : array of values which this method will write its results to
-    labels : array containing unique label for each group, with its ordering
-        matching up to the corresponding record in `values`
+    labels : array containing unique label for each group, with its
+        ordering matching up to the corresponding record in `values`
     values : array containing the truth value of each element
     mask : array indicating whether a value is na or not
+    val_test : str {'any', 'all'}
+        String object dictating whether to use any or all truth testing
     skipna : boolean
         Flag to ignore nan values during truth testing
 
@@ -374,23 +339,31 @@ def group_all(ndarray[uint8_t] out,
     """
     cdef:
         Py_ssize_t i, N=len(labels)
-        int64_t lab
-        ndarray[int64_t] bool_mask
-        ndarray[uint8_t] isna_mask
+        int64_t lab, flag_val
+
+    if val_test == 'all':
+        # Because the 'all' value of an empty iterable in Python is True we can
+        # start with an array full of ones and set to zero when a False value
+        # is encountered
+        flag_val = 0
+    elif val_test == 'any':
+        # Because the 'any' value of an empty iterable in Python is False we
+        # can start with an array full of zeros and set to one only if any
+        # value encountered is True
+        flag_val = 1
+    else:
+        raise ValueError("'bool_func' must be either 'any' or 'all'!")
 
-    # Because the 'all' value of an empty iterable in Python is True we can
-    # start with an array full of ones and set to zero when a False value is
-    # encountered
-    out.fill(1)
+    out.fill(1 - flag_val)
 
     with nogil:
         for i in range(N):
             lab = labels[i]
             if lab < 0 or (skipna and mask[i]):
                 continue
 
-            if not values[i]:
-                out[lab] = 0
+            if values[i] == flag_val:
+                out[lab] = flag_val
 
 
 # generated from template

pandas/core/groupby.py

@@ -1219,7 +1219,7 @@ class GroupBy(_GroupBy):
     """
     _apply_whitelist = _common_apply_whitelist
 
-    def _bool_agg(self, how, skipna):
+    def _bool_agg(self, val_test, skipna):
         """Shared func to call any / all Cython GroupBy implementations"""
 
         def objs_to_bool(vals):
@@ -1233,14 +1233,14 @@ def objs_to_bool(vals):
         def result_to_bool(result):
             return result.astype(np.bool, copy=False)
 
-        return self._get_cythonized_result(how, self.grouper,
+        return self._get_cythonized_result('group_any_all', self.grouper,
                                            aggregate=True,
                                            cython_dtype=np.uint8,
                                            needs_values=True,
                                            needs_mask=True,
                                            pre_processing=objs_to_bool,
                                            post_processing=result_to_bool,
-                                           skipna=skipna)
+                                           val_test=val_test, skipna=skipna)
 
     @Substitution(name='groupby')
     @Appender(_doc_template)
@@ -1252,7 +1252,7 @@ def any(self, skipna=True):
         skipna : bool, default True
             Flag to ignore nan values during truth testing
         """
-        return self._bool_agg('group_any', skipna)
+        return self._bool_agg('any', skipna)
 
     @Substitution(name='groupby')
     @Appender(_doc_template)
@@ -1264,7 +1264,7 @@ def all(self, skipna=True):
         skipna : bool, default True
             Flag to ignore nan values during truth testing
         """
-        return self._bool_agg('group_all', skipna)
+        return self._bool_agg('all', skipna)
 
     @Substitution(name='groupby')
     @Appender(_doc_template)