Refactor groupby group_add from tempita to fused types

Author: noamher

pandas/_libs/groupby.pyx

@@ -2,6 +2,7 @@
 
 import cython
 from cython import Py_ssize_t
+from cython cimport floating
 
 from libc.stdlib cimport malloc, free
 
@@ -382,5 +383,55 @@ def group_any_all(uint8_t[:] out,
  out[lab] = flag_val
 
 
+@cython.wraparound(False)
+@cython.boundscheck(False)
+def _group_add(floating[:, :] out,
+ int64_t[:] counts,
+ floating[:, :] values,
+ const int64_t[:] labels,
+ Py_ssize_t min_count=0):
+ """
+ Only aggregates on axis=0
+ """
+ cdef:
+ Py_ssize_t i, j, N, K, lab, ncounts = len(counts)
+ floating val, count
+ ndarray[floating, ndim=2] sumx, nobs
+
+ if not len(values) == len(labels):
+ raise AssertionError("len(index) != len(labels)")
+
+ nobs = np.zeros_like(out)
+ sumx = np.zeros_like(out)
+
+ N, K = (<object>values).shape
+
+ with nogil:
+
+ for i in range(N):
+ lab = labels[i]
+ if lab < 0:
+ continue
+
+ counts[lab] += 1
+ for j in range(K):
+ val = values[i, j]
+
+ # not nan
+ if val == val:
+ nobs[lab, j] += 1
+ sumx[lab, j] += val
+
+ for i in range(ncounts):
+ for j in range(K):
+ if nobs[i, j] < min_count:
+ out[i, j] = NAN
+ else:
+ out[i, j] = sumx[i, j]
+
+
+group_add_float32 = _group_add['float']
+group_add_float64 = _group_add['double']
+
 # generated from template
 include "groupby_helper.pxi"

pandas/_libs/groupby_helper.pxi.in

@@ -9,7 +9,7 @@ cdef extern from "numpy/npy_math.h":
 _int64_max = np.iinfo(np.int64).max
 
 # ----------------------------------------------------------------------
-# group_add, group_prod, group_var, group_mean, group_ohlc
+# group_prod, group_var, group_mean, group_ohlc
 # ----------------------------------------------------------------------
 
 {{py:
@@ -27,53 +27,6 @@ def get_dispatch(dtypes):
 {{for name, c_type in get_dispatch(dtypes)}}
 
 
-@cython.wraparound(False)
-@cython.boundscheck(False)
-def group_add_{{name}}({{c_type}}[:, :] out,
- int64_t[:] counts,
- {{c_type}}[:, :] values,
- const int64_t[:] labels,
- Py_ssize_t min_count=0):
- """
- Only aggregates on axis=0
- """
- cdef:
- Py_ssize_t i, j, N, K, lab, ncounts = len(counts)
- {{c_type}} val, count
- ndarray[{{c_type}}, ndim=2] sumx, nobs
-
- if not len(values) == len(labels):
- raise AssertionError("len(index) != len(labels)")
-
- nobs = np.zeros_like(out)
- sumx = np.zeros_like(out)
-
- N, K = (<object>values).shape
-
- with nogil:
-
- for i in range(N):
- lab = labels[i]
- if lab < 0:
- continue
-
- counts[lab] += 1
- for j in range(K):
- val = values[i, j]
-
- # not nan
- if val == val:
- nobs[lab, j] += 1
- sumx[lab, j] += val
-
- for i in range(ncounts):
- for j in range(K):
- if nobs[i, j] < min_count:
- out[i, j] = NAN
- else:
- out[i, j] = sumx[i, j]
-
-
 @cython.wraparound(False)
 @cython.boundscheck(False)
 def group_prod_{{name}}({{c_type}}[:, :] out,

pandas/core/groupby/ops.py

@@ -380,7 +380,7 @@ def get_func(fname):
  # otherwise find dtype-specific version, falling back to object
  for dt in [dtype_str, 'object']:
  f = getattr(libgroupby, "{fname}_{dtype_str}".format(
- fname=fname, dtype_str=dtype_str), None)
+ fname=fname, dtype_str=dt), None)
  if f is not None:
  return f