KLUDGE: check for iNaT in integer data prior to accumulate/transform in groupby

xref pandas-dev#15053

Author: jreback

pandas/core/groupby.py

@@ -24,6 +24,7 @@
                                  is_bool_dtype,
                                  is_scalar,
                                  is_list_like,
+                                 needs_i8_conversion,
                                  _ensure_float64,
                                  _ensure_platform_int,
                                  _ensure_int64,
@@ -1844,15 +1845,21 @@ def _cython_operation(self, kind, values, how, axis):
                                           "supported for the 'how' argument")
             out_shape = (self.ngroups,) + values.shape[1:]
 
+        is_datetimelike = needs_i8_conversion(values.dtype)
         is_numeric = is_numeric_dtype(values.dtype)
 
-        if is_datetime_or_timedelta_dtype(values.dtype):
+        if is_datetimelike:
             values = values.view('int64')
             is_numeric = True
         elif is_bool_dtype(values.dtype):
             values = _ensure_float64(values)
         elif is_integer_dtype(values):
-            values = values.astype('int64', copy=False)
+            # we use iNaT for the missing value on ints
+            # so pre-convert to guard this condition
+            if (values == tslib.iNaT).any():
+                values = _ensure_float64(values)
+            else:
+                values = values.astype('int64', copy=False)
         elif is_numeric and not is_complex_dtype(values):
             values = _ensure_float64(values)
         else:
@@ -1881,20 +1888,17 @@ def _cython_operation(self, kind, values, how, axis):
                                  fill_value=np.nan)
             counts = np.zeros(self.ngroups, dtype=np.int64)
             result = self._aggregate(
-                result, counts, values, labels, func, is_numeric)
+                result, counts, values, labels, func,
+                is_numeric, is_datetimelike)
         elif kind == 'transform':
             result = _maybe_fill(np.empty_like(values, dtype=out_dtype),
                                  fill_value=np.nan)
 
             # temporary storange for running-total type tranforms
             accum = np.empty(out_shape, dtype=out_dtype)
             result = self._transform(
-                result, accum, values, labels, func, is_numeric)
-
-        if is_integer_dtype(result):
-            if len(result[result == tslib.iNaT]) > 0:
-                result = result.astype('float64')
-                result[result == tslib.iNaT] = np.nan
+                result, accum, values, labels, func,
+                is_numeric, is_datetimelike)
 
         if kind == 'aggregate' and \
            self._filter_empty_groups and not counts.all():
@@ -1929,8 +1933,19 @@ def aggregate(self, values, how, axis=0):
     def transform(self, values, how, axis=0):
         return self._cython_operation('transform', values, how, axis)
 
+    def _maybe_mask_missing(self, result, is_datetimelike):
+        # we use iNaT as a marker for missing values
+        # but we *only* care for non-datetimelikes
+        if is_integer_dtype(result) and not is_datetimelike:
+            mask = result == tslib.iNaT
+            if mask.any():
+                result = result.astype('float64')
+                result[mask] = np.nan
+        return result
+
     def _aggregate(self, result, counts, values, comp_ids, agg_func,
-                   is_numeric):
+                   is_numeric, is_datetimelike):
+
         if values.ndim > 3:
             # punting for now
             raise NotImplementedError("number of dimensions is currently "
@@ -1943,11 +1958,12 @@ def _aggregate(self, result, counts, values, comp_ids, agg_func,
         else:
             agg_func(result, counts, values, comp_ids)
 
-        return result
+        return self._maybe_mask_missing(result, is_datetimelike)
 
     def _transform(self, result, accum, values, comp_ids, transform_func,
-                   is_numeric):
+                   is_numeric, is_datetimelike):
         comp_ids, _, ngroups = self.group_info
+
         if values.ndim > 3:
             # punting for now
             raise NotImplementedError("number of dimensions is currently "
@@ -1961,7 +1977,7 @@ def _transform(self, result, accum, values, comp_ids, transform_func,
         else:
             transform_func(result, values, comp_ids, accum)
 
-        return result
+        return self._maybe_mask_missing(result, is_datetimelike)
 
     def agg_series(self, obj, func):
         try: