def backfill_1d(values, limit=None, mask=None, dtype=None):
if dtype is None:
dtype = values.dtype
_method = None
if is_float_dtype(values):
name = 'backfill_inplace_{name}'.format(name=dtype.name)
_method = getattr(algos, name, None)
elif is_datetime64_dtype(dtype) or is_datetime64tz_dtype(dtype):
_method = _backfill_1d_datetime
elif is_integer_dtype(values):
values = ensure_float64(values)
_method = algos.backfill_inplace_float64
elif values.dtype == np.object_:
_method = algos.backfill_inplace_object
elif is_timedelta64_dtype(values):
# NaTs are treated identically to datetime64, so we can dispatch
# to that implementation
_method = _backfill_1d_datetime
if _method is None:
raise ValueError('Invalid dtype for backfill_1d [{name}]'
.format(name=dtype.name))
if mask is None:
mask = isna(values)
mask = mask.view(np.uint8)
_method(values, mask, limit=limit)
return values
def backfill_2d(values, limit=None, mask=None, dtype=None):
if dtype is None:
dtype = values.dtype
_method = None
if is_float_dtype(values):
name = 'backfill_2d_inplace_{name}'.format(name=dtype.name)
_method = getattr(algos, name, None)
elif is_datetime64_dtype(dtype) or is_datetime64tz_dtype(dtype):
_method = _backfill_2d_datetime
elif is_integer_dtype(values):
values = ensure_float64(values)
_method = algos.backfill_2d_inplace_float64
elif values.dtype == np.object_:
_method = algos.backfill_2d_inplace_object
if _method is None:
raise ValueError('Invalid dtype for backfill_2d [{name}]'
.format(name=dtype.name))
if mask is None:
mask = isna(values)
mask = mask.view(np.uint8)
if np.all(values.shape):
_method(values, mask, limit=limit)
else:
# for test coverage
pass
return values
def _cast_values_for_fillna(values, dtype):
"""
Cast values to a dtype that algos.pad and algos.backfill can handle.
"""
# TODO: for int-dtypes we make a copy, but for everything else this
# alters the values in-place. Is this intentional?
if (is_datetime64_dtype(dtype) or is_datetime64tz_dtype(dtype)
or is_timedelta64_dtype(dtype)):
values = values.view(np.int64)
elif is_integer_dtype(values):
# NB: this check needs to come after the datetime64 check above
values = ensure_float64(values)
return values
def _cast_values_for_fillna(values, dtype):
"""
Cast values to a dtype that algos.pad and algos.backfill can handle.
"""
# TODO: for int-dtypes we make a copy, but for everything else this
# alters the values in-place. Is this intentional?
if (is_datetime64_dtype(dtype) or is_datetime64tz_dtype(dtype) or
is_timedelta64_dtype(dtype)):
values = values.view(np.int64)
elif is_integer_dtype(values):
# NB: this check needs to come after the datetime64 check above
values = ensure_float64(values)
return values
def _cast_values_for_fillna(values, dtype: DtypeObj, has_mask: bool):
"""
Cast values to a dtype that algos.pad and algos.backfill can handle.
"""
# TODO: for int-dtypes we make a copy, but for everything else this
# alters the values in-place. Is this intentional?
if needs_i8_conversion(dtype):
values = values.view(np.int64)
elif is_integer_dtype(values) and not has_mask:
# NB: this check needs to come after the datetime64 check above
# has_mask check to avoid casting i8 values that have already
# been cast from PeriodDtype
values = ensure_float64(values)
return values
def pad_1d(values, limit=None, mask=None, dtype=None):
if dtype is None:
dtype = values.dtype
_method = None
if is_float_dtype(values):
name = 'pad_inplace_{name}'.format(name=dtype.name)
_method = getattr(algos, name, None)
elif is_datetime64_dtype(dtype) or is_datetime64tz_dtype(dtype):
_method = _pad_1d_datetime
elif is_integer_dtype(values):
values = ensure_float64(values)
_method = algos.pad_inplace_float64
elif values.dtype == np.object_:
_method = algos.pad_inplace_object
if _method is None:
raise ValueError(
'Invalid dtype for pad_1d [{name}]'.format(name=dtype.name))
if mask is None:
mask = isna(values)
mask = mask.view(np.uint8)
_method(values, mask, limit=limit)
return values
def pad_1d(values, limit=None, mask=None, dtype=None):
if dtype is None:
dtype = values.dtype
_method = None
if is_float_dtype(values):
name = 'pad_inplace_{name}'.format(name=dtype.name)
_method = getattr(algos, name, None)
elif is_datetime64_dtype(dtype) or is_datetime64tz_dtype(dtype):
_method = _pad_1d_datetime
elif is_integer_dtype(values):
values = ensure_float64(values)
_method = algos.pad_inplace_float64
elif values.dtype == np.object_:
_method = algos.pad_inplace_object
if _method is None:
raise ValueError('Invalid dtype for pad_1d [{name}]'
.format(name=dtype.name))
if mask is None:
mask = isna(values)
mask = mask.view(np.uint8)
_method(values, mask, limit=limit)
return values
def _cython_operation(self, kind, values, how, axis, min_count=-1,
**kwargs):
assert kind in ['transform', 'aggregate']
# can we do this operation with our cython functions
# if not raise NotImplementedError
# we raise NotImplemented if this is an invalid operation
# entirely, e.g. adding datetimes
# categoricals are only 1d, so we
# are not setup for dim transforming
if is_categorical_dtype(values):
raise NotImplementedError(
"categoricals are not support in cython ops ATM")
elif is_datetime64_any_dtype(values):
if how in ['add', 'prod', 'cumsum', 'cumprod']:
raise NotImplementedError(
"datetime64 type does not support {} "
"operations".format(how))
elif is_timedelta64_dtype(values):
if how in ['prod', 'cumprod']:
raise NotImplementedError(
"timedelta64 type does not support {} "
"operations".format(how))
arity = self._cython_arity.get(how, 1)
vdim = values.ndim
swapped = False
if vdim == 1:
values = values[:, None]
out_shape = (self.ngroups, arity)
else:
if axis > 0:
swapped = True
values = values.swapaxes(0, axis)
if arity > 1:
raise NotImplementedError("arity of more than 1 is not "
"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_datetimelike:
values = values.view('int64')
is_numeric = True
elif is_bool_dtype(values.dtype):
values = ensure_float64(values)
elif is_integer_dtype(values):
# we use iNaT for the missing value on ints
# so pre-convert to guard this condition
if (values == iNaT).any():
values = ensure_float64(values)
else:
values = ensure_int64_or_float64(values)
elif is_numeric and not is_complex_dtype(values):
values = ensure_float64(values)
else:
values = values.astype(object)
try:
func = self._get_cython_function(
kind, how, values, is_numeric)
except NotImplementedError:
if is_numeric:
values = ensure_float64(values)
func = self._get_cython_function(
kind, how, values, is_numeric)
else:
raise
if how == 'rank':
out_dtype = 'float'
else:
if is_numeric:
out_dtype = '{kind}{itemsize}'.format(
kind=values.dtype.kind, itemsize=values.dtype.itemsize)
else:
out_dtype = 'object'
labels, _, _ = self.group_info
if kind == 'aggregate':
result = _maybe_fill(np.empty(out_shape, dtype=out_dtype),
fill_value=np.nan)
counts = np.zeros(self.ngroups, dtype=np.int64)
result = self._aggregate(
result, counts, values, labels, func, is_numeric,
is_datetimelike, min_count)
elif kind == 'transform':
result = _maybe_fill(np.empty_like(values, dtype=out_dtype),
fill_value=np.nan)
# TODO: min_count
result = self._transform(
result, values, labels, func, is_numeric, is_datetimelike,
**kwargs)
if is_integer_dtype(result) and not is_datetimelike:
mask = result == iNaT
if mask.any():
result = result.astype('float64')
result[mask] = np.nan
if (kind == 'aggregate' and
self._filter_empty_groups and not counts.all()):
if result.ndim == 2:
try:
result = lib.row_bool_subset(
result, (counts > 0).view(np.uint8))
except ValueError:
result = lib.row_bool_subset_object(
ensure_object(result),
(counts > 0).view(np.uint8))
else:
result = result[counts > 0]
if vdim == 1 and arity == 1:
result = result[:, 0]
if how in self._name_functions:
# TODO
names = self._name_functions[how]()
else:
names = None
if swapped:
result = result.swapaxes(0, axis)
return result, names
def _cython_operation(
self, kind: str, values, how: str, axis: int, min_count: int = -1, **kwargs
) -> ArrayLike:
"""
Returns the values of a cython operation.
"""
orig_values = values
assert kind in ["transform", "aggregate"]
if values.ndim > 2:
raise NotImplementedError("number of dimensions is currently limited to 2")
elif values.ndim == 2:
# Note: it is *not* the case that axis is always 0 for 1-dim values,
# as we can have 1D ExtensionArrays that we need to treat as 2D
assert axis == 1, axis
dtype = values.dtype
is_numeric = is_numeric_dtype(dtype)
cy_op = WrappedCythonOp(kind=kind, how=how)
# can we do this operation with our cython functions
# if not raise NotImplementedError
cy_op.disallow_invalid_ops(dtype, is_numeric)
if is_extension_array_dtype(dtype):
return self._ea_wrap_cython_operation(
kind, values, how, axis, min_count, **kwargs
)
elif values.ndim == 1:
# expand to 2d, dispatch, then squeeze if appropriate
values2d = values[None, :]
res = self._cython_operation(
kind=kind,
values=values2d,
how=how,
axis=1,
min_count=min_count,
**kwargs,
)
if res.shape[0] == 1:
return res[0]
# otherwise we have OHLC
return res.T
is_datetimelike = needs_i8_conversion(dtype)
if is_datetimelike:
values = values.view("int64")
is_numeric = True
elif is_bool_dtype(dtype):
values = ensure_int_or_float(values)
elif is_integer_dtype(dtype):
# we use iNaT for the missing value on ints
# so pre-convert to guard this condition
if (values == iNaT).any():
values = ensure_float64(values)
else:
values = ensure_int_or_float(values)
elif is_numeric:
if not is_complex_dtype(dtype):
values = ensure_float64(values)
ngroups = self.ngroups
comp_ids, _, _ = self.group_info
assert axis == 1
values = values.T
out_shape = cy_op.get_output_shape(ngroups, values)
func, values = cy_op.get_cython_func_and_vals(values, is_numeric)
out_dtype = cy_op.get_out_dtype(values.dtype)
result = maybe_fill(np.empty(out_shape, dtype=out_dtype))
if kind == "aggregate":
counts = np.zeros(ngroups, dtype=np.int64)
func(result, counts, values, comp_ids, min_count)
elif kind == "transform":
# TODO: min_count
func(result, values, comp_ids, ngroups, is_datetimelike, **kwargs)
if is_integer_dtype(result.dtype) and not is_datetimelike:
mask = result == iNaT
if mask.any():
result = result.astype("float64")
result[mask] = np.nan
if kind == "aggregate" and self._filter_empty_groups and not counts.all():
assert result.ndim != 2
result = result[counts > 0]
result = result.T
if how not in base.cython_cast_blocklist:
# e.g. if we are int64 and need to restore to datetime64/timedelta64
# "rank" is the only member of cython_cast_blocklist we get here
dtype = maybe_cast_result_dtype(orig_values.dtype, how)
op_result = maybe_downcast_to_dtype(result, dtype)
else:
op_result = result
return op_result
def _cython_operation(
self, kind: str, values, how: str, axis, min_count: int = -1, **kwargs
) -> Tuple[np.ndarray, Optional[List[str]]]:
"""
Returns the values of a cython operation as a Tuple of [data, names].
Names is only useful when dealing with 2D results, like ohlc
(see self._name_functions).
"""
assert kind in ["transform", "aggregate"]
orig_values = values
if values.ndim > 2:
raise NotImplementedError("number of dimensions is currently limited to 2")
elif values.ndim == 2:
# Note: it is *not* the case that axis is always 0 for 1-dim values,
# as we can have 1D ExtensionArrays that we need to treat as 2D
assert axis == 1, axis
# can we do this operation with our cython functions
# if not raise NotImplementedError
# we raise NotImplemented if this is an invalid operation
# entirely, e.g. adding datetimes
# categoricals are only 1d, so we
# are not setup for dim transforming
if is_categorical_dtype(values) or is_sparse(values):
raise NotImplementedError(f"{values.dtype} dtype not supported")
elif is_datetime64_any_dtype(values):
if how in ["add", "prod", "cumsum", "cumprod"]:
raise NotImplementedError(
f"datetime64 type does not support {how} operations"
)
elif is_timedelta64_dtype(values):
if how in ["prod", "cumprod"]:
raise NotImplementedError(
f"timedelta64 type does not support {how} operations"
)
if is_datetime64tz_dtype(values.dtype):
# Cast to naive; we'll cast back at the end of the function
# TODO: possible need to reshape? kludge can be avoided when
# 2D EA is allowed.
values = values.view("M8[ns]")
is_datetimelike = needs_i8_conversion(values.dtype)
is_numeric = is_numeric_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):
# we use iNaT for the missing value on ints
# so pre-convert to guard this condition
if (values == iNaT).any():
values = ensure_float64(values)
else:
values = ensure_int_or_float(values)
elif is_numeric and not is_complex_dtype(values):
values = ensure_float64(values)
else:
values = values.astype(object)
arity = self._cython_arity.get(how, 1)
vdim = values.ndim
swapped = False
if vdim == 1:
values = values[:, None]
out_shape = (self.ngroups, arity)
else:
if axis > 0:
swapped = True
assert axis == 1, axis
values = values.T
if arity > 1:
raise NotImplementedError(
"arity of more than 1 is not supported for the 'how' argument"
)
out_shape = (self.ngroups,) + values.shape[1:]
func, values = self._get_cython_func_and_vals(kind, how, values, is_numeric)
if how == "rank":
out_dtype = "float"
else:
if is_numeric:
out_dtype = f"{values.dtype.kind}{values.dtype.itemsize}"
else:
out_dtype = "object"
codes, _, _ = self.group_info
if kind == "aggregate":
result = _maybe_fill(
np.empty(out_shape, dtype=out_dtype), fill_value=np.nan
)
counts = np.zeros(self.ngroups, dtype=np.int64)
result = self._aggregate(result, counts, values, codes, func, min_count)
elif kind == "transform":
result = _maybe_fill(
np.empty_like(values, dtype=out_dtype), fill_value=np.nan
)
# TODO: min_count
result = self._transform(
result, values, codes, func, is_datetimelike, **kwargs
)
if is_integer_dtype(result) and not is_datetimelike:
mask = result == iNaT
if mask.any():
result = result.astype("float64")
result[mask] = np.nan
elif (
how == "add"
and is_integer_dtype(orig_values.dtype)
and is_extension_array_dtype(orig_values.dtype)
):
# We need this to ensure that Series[Int64Dtype].resample().sum()
# remains int64 dtype.
# Two options for avoiding this special case
# 1. mask-aware ops and avoid casting to float with NaN above
# 2. specify the result dtype when calling this method
result = result.astype("int64")
if kind == "aggregate" and self._filter_empty_groups and not counts.all():
assert result.ndim != 2
result = result[counts > 0]
if vdim == 1 and arity == 1:
result = result[:, 0]
names: Optional[List[str]] = self._name_functions.get(how, None)
if swapped:
result = result.swapaxes(0, axis)
if is_datetime64tz_dtype(orig_values.dtype) or is_period_dtype(
orig_values.dtype
):
# We need to use the constructors directly for these dtypes
# since numpy won't recognize them
# https://github.com/pandas-dev/pandas/issues/31471
result = type(orig_values)(result.astype(np.int64), dtype=orig_values.dtype)
elif is_datetimelike and kind == "aggregate":
result = result.astype(orig_values.dtype)
return result, names
def _cython_operation(self, kind, values, how, axis, min_count=-1,
**kwargs):
assert kind in ['transform', 'aggregate']
# can we do this operation with our cython functions
# if not raise NotImplementedError
# we raise NotImplemented if this is an invalid operation
# entirely, e.g. adding datetimes
# categoricals are only 1d, so we
# are not setup for dim transforming
if is_categorical_dtype(values):
raise NotImplementedError(
"categoricals are not support in cython ops ATM")
elif is_datetime64_any_dtype(values):
if how in ['add', 'prod', 'cumsum', 'cumprod']:
raise NotImplementedError(
"datetime64 type does not support {} "
"operations".format(how))
elif is_timedelta64_dtype(values):
if how in ['prod', 'cumprod']:
raise NotImplementedError(
"timedelta64 type does not support {} "
"operations".format(how))
arity = self._cython_arity.get(how, 1)
vdim = values.ndim
swapped = False
if vdim == 1:
values = values[:, None]
out_shape = (self.ngroups, arity)
else:
if axis > 0:
swapped = True
values = values.swapaxes(0, axis)
if arity > 1:
raise NotImplementedError("arity of more than 1 is not "
"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_datetimelike:
values = values.view('int64')
is_numeric = True
elif is_bool_dtype(values.dtype):
values = ensure_float64(values)
elif is_integer_dtype(values):
# we use iNaT for the missing value on ints
# so pre-convert to guard this condition
if (values == iNaT).any():
values = ensure_float64(values)
else:
values = ensure_int64_or_float64(values)
elif is_numeric and not is_complex_dtype(values):
values = ensure_float64(values)
else:
values = values.astype(object)
try:
func = self._get_cython_function(
kind, how, values, is_numeric)
except NotImplementedError:
if is_numeric:
values = ensure_float64(values)
func = self._get_cython_function(
kind, how, values, is_numeric)
else:
raise
if how == 'rank':
out_dtype = 'float'
else:
if is_numeric:
out_dtype = '{kind}{itemsize}'.format(
kind=values.dtype.kind, itemsize=values.dtype.itemsize)
else:
out_dtype = 'object'
labels, _, _ = self.group_info
if kind == 'aggregate':
result = _maybe_fill(np.empty(out_shape, dtype=out_dtype),
fill_value=np.nan)
counts = np.zeros(self.ngroups, dtype=np.int64)
result = self._aggregate(
result, counts, values, labels, func, is_numeric,
is_datetimelike, min_count)
elif kind == 'transform':
result = _maybe_fill(np.empty_like(values, dtype=out_dtype),
fill_value=np.nan)
# TODO: min_count
result = self._transform(
result, values, labels, func, is_numeric, is_datetimelike,
**kwargs)
if is_integer_dtype(result) and not is_datetimelike:
mask = result == iNaT
if mask.any():
result = result.astype('float64')
result[mask] = np.nan
if (kind == 'aggregate' and
self._filter_empty_groups and not counts.all()):
if result.ndim == 2:
try:
result = lib.row_bool_subset(
result, (counts > 0).view(np.uint8))
except ValueError:
result = lib.row_bool_subset_object(
ensure_object(result),
(counts > 0).view(np.uint8))
else:
result = result[counts > 0]
if vdim == 1 and arity == 1:
result = result[:, 0]
if how in self._name_functions:
# TODO
names = self._name_functions[how]()
else:
names = None
if swapped:
result = result.swapaxes(0, axis)
return result, names
def _cython_operation(self,
kind: str,
values,
how,
axis,
min_count=-1,
**kwargs):
assert kind in ["transform", "aggregate"]
orig_values = values
# can we do this operation with our cython functions
# if not raise NotImplementedError
# we raise NotImplemented if this is an invalid operation
# entirely, e.g. adding datetimes
# categoricals are only 1d, so we
# are not setup for dim transforming
if is_categorical_dtype(values) or is_sparse(values):
raise NotImplementedError(
"{dtype} dtype not supported".format(dtype=values.dtype))
elif is_datetime64_any_dtype(values):
if how in ["add", "prod", "cumsum", "cumprod"]:
raise NotImplementedError(
"datetime64 type does not support {how} operations".format(
how=how))
elif is_timedelta64_dtype(values):
if how in ["prod", "cumprod"]:
raise NotImplementedError(
"timedelta64 type does not support {how} operations".
format(how=how))
if is_datetime64tz_dtype(values.dtype):
# Cast to naive; we'll cast back at the end of the function
# TODO: possible need to reshape? kludge can be avoided when
# 2D EA is allowed.
values = values.view("M8[ns]")
is_datetimelike = needs_i8_conversion(values.dtype)
is_numeric = is_numeric_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):
# we use iNaT for the missing value on ints
# so pre-convert to guard this condition
if (values == iNaT).any():
values = ensure_float64(values)
else:
values = ensure_int_or_float(values)
elif is_numeric and not is_complex_dtype(values):
values = ensure_float64(values)
else:
values = values.astype(object)
arity = self._cython_arity.get(how, 1)
vdim = values.ndim
swapped = False
if vdim == 1:
values = values[:, None]
out_shape = (self.ngroups, arity)
else:
if axis > 0:
swapped = True
assert axis == 1, axis
values = values.T
if arity > 1:
raise NotImplementedError(
"arity of more than 1 is not supported for the 'how' argument"
)
out_shape = (self.ngroups, ) + values.shape[1:]
try:
func = self._get_cython_function(kind, how, values, is_numeric)
except NotImplementedError:
if is_numeric:
try:
values = ensure_float64(values)
except TypeError:
if lib.infer_dtype(values, skipna=False) == "complex":
values = values.astype(complex)
else:
raise
func = self._get_cython_function(kind, how, values, is_numeric)
else:
raise
if how == "rank":
out_dtype = "float"
else:
if is_numeric:
out_dtype = "{kind}{itemsize}".format(
kind=values.dtype.kind, itemsize=values.dtype.itemsize)
else:
out_dtype = "object"
labels, _, _ = self.group_info
if kind == "aggregate":
result = _maybe_fill(np.empty(out_shape, dtype=out_dtype),
fill_value=np.nan)
counts = np.zeros(self.ngroups, dtype=np.int64)
result = self._aggregate(result, counts, values, labels, func,
is_datetimelike, min_count)
elif kind == "transform":
result = _maybe_fill(np.empty_like(values, dtype=out_dtype),
fill_value=np.nan)
# TODO: min_count
result = self._transform(result, values, labels, func,
is_datetimelike, **kwargs)
if is_integer_dtype(result) and not is_datetimelike:
mask = result == iNaT
if mask.any():
result = result.astype("float64")
result[mask] = np.nan
if kind == "aggregate" and self._filter_empty_groups and not counts.all(
):
assert result.ndim != 2
result = result[counts > 0]
if vdim == 1 and arity == 1:
result = result[:, 0]
if how in self._name_functions:
names = self._name_functions[how]() # type: Optional[List[str]]
else:
names = None
if swapped:
result = result.swapaxes(0, axis)
if is_datetime64tz_dtype(orig_values.dtype):
result = type(orig_values)(result.astype(np.int64),
dtype=orig_values.dtype)
elif is_datetimelike and kind == "aggregate":
result = result.astype(orig_values.dtype)
return result, names
def _cython_operation(self,
kind: str,
values,
how: str,
axis: int,
min_count: int = -1,
**kwargs) -> Tuple[np.ndarray, Optional[List[str]]]:
"""
Returns the values of a cython operation as a Tuple of [data, names].
Names is only useful when dealing with 2D results, like ohlc
(see self._name_functions).
"""
orig_values = values
assert kind in ["transform", "aggregate"]
if values.ndim > 2:
raise NotImplementedError(
"number of dimensions is currently limited to 2")
elif values.ndim == 2:
# Note: it is *not* the case that axis is always 0 for 1-dim values,
# as we can have 1D ExtensionArrays that we need to treat as 2D
assert axis == 1, axis
# can we do this operation with our cython functions
# if not raise NotImplementedError
self._disallow_invalid_ops(values, how)
if is_extension_array_dtype(values.dtype):
return self._ea_wrap_cython_operation(kind, values, how, axis,
min_count, **kwargs)
is_datetimelike = needs_i8_conversion(values.dtype)
is_numeric = is_numeric_dtype(values.dtype)
if is_datetimelike:
values = values.view("int64")
is_numeric = True
elif is_bool_dtype(values.dtype):
values = ensure_int_or_float(values)
elif is_integer_dtype(values):
# we use iNaT for the missing value on ints
# so pre-convert to guard this condition
if (values == iNaT).any():
values = ensure_float64(values)
else:
values = ensure_int_or_float(values)
elif is_numeric and not is_complex_dtype(values):
values = ensure_float64(ensure_float(values))
else:
values = values.astype(object)
arity = self._cython_arity.get(how, 1)
vdim = values.ndim
swapped = False
if vdim == 1:
values = values[:, None]
out_shape = (self.ngroups, arity)
else:
if axis > 0:
swapped = True
assert axis == 1, axis
values = values.T
if arity > 1:
raise NotImplementedError(
"arity of more than 1 is not supported for the 'how' argument"
)
out_shape = (self.ngroups, ) + values.shape[1:]
func, values = self._get_cython_func_and_vals(kind, how, values,
is_numeric)
if how == "rank":
out_dtype = "float"
else:
if is_numeric:
out_dtype = f"{values.dtype.kind}{values.dtype.itemsize}"
else:
out_dtype = "object"
codes, _, _ = self.group_info
if kind == "aggregate":
result = maybe_fill(np.empty(out_shape, dtype=out_dtype),
fill_value=np.nan)
counts = np.zeros(self.ngroups, dtype=np.int64)
result = self._aggregate(result, counts, values, codes, func,
min_count)
elif kind == "transform":
result = maybe_fill(np.empty_like(values, dtype=out_dtype),
fill_value=np.nan)
# TODO: min_count
result = self._transform(result, values, codes, func,
is_datetimelike, **kwargs)
if is_integer_dtype(result) and not is_datetimelike:
mask = result == iNaT
if mask.any():
result = result.astype("float64")
result[mask] = np.nan
if kind == "aggregate" and self._filter_empty_groups and not counts.all(
):
assert result.ndim != 2
result = result[counts > 0]
if vdim == 1 and arity == 1:
result = result[:, 0]
names: Optional[List[str]] = self._name_functions.get(how, None)
if swapped:
result = result.swapaxes(0, axis)
if is_datetimelike and kind == "aggregate":
result = result.astype(orig_values.dtype)
return result, names
def _call_cython_op(
self,
values: np.ndarray, # np.ndarray[ndim=2]
*,
min_count: int,
ngroups: int,
comp_ids: np.ndarray,
mask: np.ndarray | None,
**kwargs,
) -> np.ndarray: # np.ndarray[ndim=2]
orig_values = values
dtype = values.dtype
is_numeric = is_numeric_dtype(dtype)
is_datetimelike = needs_i8_conversion(dtype)
if is_datetimelike:
values = values.view("int64")
is_numeric = True
elif is_bool_dtype(dtype):
values = values.astype("int64")
elif is_integer_dtype(dtype):
# e.g. uint8 -> uint64, int16 -> int64
dtype_str = dtype.kind + "8"
values = values.astype(dtype_str, copy=False)
elif is_numeric:
if not is_complex_dtype(dtype):
values = ensure_float64(values)
values = values.T
if mask is not None:
mask = mask.reshape(values.shape, order="C")
out_shape = self.get_output_shape(ngroups, values)
func, values = self.get_cython_func_and_vals(values, is_numeric)
out_dtype = self.get_out_dtype(values.dtype)
result = maybe_fill(np.empty(out_shape, dtype=out_dtype))
if self.kind == "aggregate":
counts = np.zeros(ngroups, dtype=np.int64)
if self.how in ["min", "max"]:
func(
result,
counts,
values,
comp_ids,
min_count,
is_datetimelike=is_datetimelike,
)
else:
func(result, counts, values, comp_ids, min_count)
else:
# TODO: min_count
if self.uses_mask():
func(
result,
values,
comp_ids,
ngroups,
is_datetimelike,
mask=mask,
**kwargs,
)
else:
func(result, values, comp_ids, ngroups, is_datetimelike,
**kwargs)
if self.kind == "aggregate":
# i.e. counts is defined. Locations where count<min_count
# need to have the result set to np.nan, which may require casting,
# see GH#40767
if is_integer_dtype(result.dtype) and not is_datetimelike:
cutoff = max(1, min_count)
empty_groups = counts < cutoff
if empty_groups.any():
# Note: this conversion could be lossy, see GH#40767
result = result.astype("float64")
result[empty_groups] = np.nan
result = result.T
if self.how not in self.cast_blocklist:
# e.g. if we are int64 and need to restore to datetime64/timedelta64
# "rank" is the only member of cast_blocklist we get here
res_dtype = self.get_result_dtype(orig_values.dtype)
# error: Argument 2 to "maybe_downcast_to_dtype" has incompatible type
# "Union[dtype[Any], ExtensionDtype]"; expected "Union[str, dtype[Any]]"
op_result = maybe_downcast_to_dtype(
result,
res_dtype # type: ignore[arg-type]
)
else:
op_result = result
# error: Incompatible return value type (got "Union[ExtensionArray, ndarray]",
# expected "ndarray")
return op_result # type: ignore[return-value]
def _cython_operation(self,
kind: str,
values,
how: str,
axis: int,
min_count: int = -1,
**kwargs) -> ArrayLike:
"""
Returns the values of a cython operation.
"""
orig_values = values
assert kind in ["transform", "aggregate"]
if values.ndim > 2:
raise NotImplementedError(
"number of dimensions is currently limited to 2")
elif values.ndim == 2:
# Note: it is *not* the case that axis is always 0 for 1-dim values,
# as we can have 1D ExtensionArrays that we need to treat as 2D
assert axis == 1, axis
dtype = values.dtype
is_numeric = is_numeric_dtype(dtype)
# can we do this operation with our cython functions
# if not raise NotImplementedError
self._disallow_invalid_ops(dtype, how, is_numeric)
if is_extension_array_dtype(dtype):
return self._ea_wrap_cython_operation(kind, values, how, axis,
min_count, **kwargs)
is_datetimelike = needs_i8_conversion(dtype)
if is_datetimelike:
values = values.view("int64")
is_numeric = True
elif is_bool_dtype(dtype):
values = ensure_int_or_float(values)
elif is_integer_dtype(dtype):
# we use iNaT for the missing value on ints
# so pre-convert to guard this condition
if (values == iNaT).any():
values = ensure_float64(values)
else:
values = ensure_int_or_float(values)
elif is_numeric and not is_complex_dtype(dtype):
values = ensure_float64(values)
else:
values = values.astype(object)
arity = self._cython_arity.get(how, 1)
vdim = values.ndim
swapped = False
if vdim == 1:
values = values[:, None]
out_shape = (self.ngroups, arity)
else:
if axis > 0:
swapped = True
assert axis == 1, axis
values = values.T
if arity > 1:
raise NotImplementedError(
"arity of more than 1 is not supported for the 'how' argument"
)
out_shape = (self.ngroups, ) + values.shape[1:]
func, values = self._get_cython_func_and_vals(kind, how, values,
is_numeric)
if how == "rank":
out_dtype = "float"
else:
if is_numeric:
out_dtype = f"{values.dtype.kind}{values.dtype.itemsize}"
else:
out_dtype = "object"
codes, _, _ = self.group_info
if kind == "aggregate":
result = maybe_fill(np.empty(out_shape, dtype=out_dtype))
counts = np.zeros(self.ngroups, dtype=np.int64)
result = self._aggregate(result, counts, values, codes, func,
min_count)
elif kind == "transform":
result = maybe_fill(np.empty(values.shape, dtype=out_dtype))
# TODO: min_count
result = self._transform(result, values, codes, func,
is_datetimelike, **kwargs)
if is_integer_dtype(result.dtype) and not is_datetimelike:
mask = result == iNaT
if mask.any():
result = result.astype("float64")
result[mask] = np.nan
if kind == "aggregate" and self._filter_empty_groups and not counts.all(
):
assert result.ndim != 2
result = result[counts > 0]
if vdim == 1 and arity == 1:
result = result[:, 0]
if swapped:
result = result.swapaxes(0, axis)
if how not in base.cython_cast_blocklist:
# e.g. if we are int64 and need to restore to datetime64/timedelta64
# "rank" is the only member of cython_cast_blocklist we get here
dtype = maybe_cast_result_dtype(orig_values.dtype, how)
op_result = maybe_downcast_to_dtype(result, dtype)
else:
op_result = result
return op_result
def _call_cython_op(
self,
values: np.ndarray, # np.ndarray[ndim=2]
*,
min_count: int,
ngroups: int,
comp_ids: np.ndarray,
mask: np.ndarray | None,
result_mask: np.ndarray | None,
**kwargs,
) -> np.ndarray: # np.ndarray[ndim=2]
orig_values = values
dtype = values.dtype
is_numeric = is_numeric_dtype(dtype)
is_datetimelike = needs_i8_conversion(dtype)
if is_datetimelike:
values = values.view("int64")
is_numeric = True
elif is_bool_dtype(dtype):
values = values.astype("int64")
elif is_integer_dtype(dtype):
# GH#43329 If the dtype is explicitly of type uint64 the type is not
# changed to prevent overflow.
if dtype != np.uint64:
values = values.astype(np.int64, copy=False)
elif is_numeric:
if not is_complex_dtype(dtype):
values = ensure_float64(values)
values = values.T
if mask is not None:
mask = mask.T
if result_mask is not None:
result_mask = result_mask.T
out_shape = self._get_output_shape(ngroups, values)
func, values = self.get_cython_func_and_vals(values, is_numeric)
out_dtype = self.get_out_dtype(values.dtype)
result = maybe_fill(np.empty(out_shape, dtype=out_dtype))
if self.kind == "aggregate":
counts = np.zeros(ngroups, dtype=np.int64)
if self.how in ["min", "max", "mean"]:
func(
result,
counts,
values,
comp_ids,
min_count,
mask=mask,
result_mask=result_mask,
is_datetimelike=is_datetimelike,
)
elif self.how in ["add"]:
# We support datetimelike
func(
result,
counts,
values,
comp_ids,
min_count,
datetimelike=is_datetimelike,
)
else:
func(result, counts, values, comp_ids, min_count)
else:
# TODO: min_count
if self.uses_mask():
func(
result,
values,
comp_ids,
ngroups,
is_datetimelike,
mask=mask,
**kwargs,
)
else:
func(result, values, comp_ids, ngroups, is_datetimelike, **kwargs)
if self.kind == "aggregate":
# i.e. counts is defined. Locations where count<min_count
# need to have the result set to np.nan, which may require casting,
# see GH#40767
if is_integer_dtype(result.dtype) and not is_datetimelike:
cutoff = max(1, min_count)
empty_groups = counts < cutoff
if empty_groups.any():
# Note: this conversion could be lossy, see GH#40767
result = result.astype("float64")
result[empty_groups] = np.nan
result = result.T
if self.how not in self.cast_blocklist:
# e.g. if we are int64 and need to restore to datetime64/timedelta64
# "rank" is the only member of cast_blocklist we get here
res_dtype = self._get_result_dtype(orig_values.dtype)
op_result = maybe_downcast_to_dtype(result, res_dtype)
else:
op_result = result
# error: Incompatible return value type (got "Union[ExtensionArray, ndarray]",
# expected "ndarray")
return op_result # type: ignore[return-value]
def _cython_operation(
self,
kind: str,
values,
how: str,
axis: int,
min_count: int = -1,
mask: np.ndarray | None = None,
**kwargs,
) -> ArrayLike:
"""
Returns the values of a cython operation.
"""
orig_values = values
assert kind in ["transform", "aggregate"]
if values.ndim > 2:
raise NotImplementedError("number of dimensions is currently limited to 2")
elif values.ndim == 2:
# Note: it is *not* the case that axis is always 0 for 1-dim values,
# as we can have 1D ExtensionArrays that we need to treat as 2D
assert axis == 1, axis
dtype = values.dtype
is_numeric = is_numeric_dtype(dtype)
cy_op = WrappedCythonOp(kind=kind, how=how)
# can we do this operation with our cython functions
# if not raise NotImplementedError
cy_op.disallow_invalid_ops(dtype, is_numeric)
func_uses_mask = cy_op.uses_mask()
if is_extension_array_dtype(dtype):
if isinstance(values, BaseMaskedArray) and func_uses_mask:
return self._masked_ea_wrap_cython_operation(
kind, values, how, axis, min_count, **kwargs
)
else:
return self._ea_wrap_cython_operation(
kind, values, how, axis, min_count, **kwargs
)
elif values.ndim == 1:
# expand to 2d, dispatch, then squeeze if appropriate
values2d = values[None, :]
res = self._cython_operation(
kind=kind,
values=values2d,
how=how,
axis=1,
min_count=min_count,
mask=mask,
**kwargs,
)
if res.shape[0] == 1:
return res[0]
# otherwise we have OHLC
return res.T
is_datetimelike = needs_i8_conversion(dtype)
if is_datetimelike:
values = values.view("int64")
is_numeric = True
elif is_bool_dtype(dtype):
values = values.astype("int64")
elif is_integer_dtype(dtype):
# e.g. uint8 -> uint64, int16 -> int64
dtype = dtype.kind + "8"
values = values.astype(dtype, copy=False)
elif is_numeric:
if not is_complex_dtype(dtype):
values = ensure_float64(values)
ngroups = self.ngroups
comp_ids, _, _ = self.group_info
assert axis == 1
values = values.T
if mask is not None:
mask = mask.reshape(values.shape, order="C")
out_shape = cy_op.get_output_shape(ngroups, values)
func, values = cy_op.get_cython_func_and_vals(values, is_numeric)
out_dtype = cy_op.get_out_dtype(values.dtype)
result = maybe_fill(np.empty(out_shape, dtype=out_dtype))
if kind == "aggregate":
counts = np.zeros(ngroups, dtype=np.int64)
if how in ["min", "max"]:
func(
result,
counts,
values,
comp_ids,
min_count,
is_datetimelike=is_datetimelike,
)
else:
func(result, counts, values, comp_ids, min_count)
elif kind == "transform":
# TODO: min_count
if func_uses_mask:
func(
result,
values,
comp_ids,
ngroups,
is_datetimelike,
mask=mask,
**kwargs,
)
else:
func(result, values, comp_ids, ngroups, is_datetimelike, **kwargs)
if kind == "aggregate":
# i.e. counts is defined. Locations where count<min_count
# need to have the result set to np.nan, which may require casting,
# see GH#40767
if is_integer_dtype(result.dtype) and not is_datetimelike:
cutoff = max(1, min_count)
empty_groups = counts < cutoff
if empty_groups.any():
# Note: this conversion could be lossy, see GH#40767
result = result.astype("float64")
result[empty_groups] = np.nan
if self._filter_empty_groups and not counts.all():
assert result.ndim != 2
result = result[counts > 0]
result = result.T
if how not in cy_op.cast_blocklist:
# e.g. if we are int64 and need to restore to datetime64/timedelta64
# "rank" is the only member of cast_blocklist we get here
dtype = maybe_cast_result_dtype(orig_values.dtype, how)
op_result = maybe_downcast_to_dtype(result, dtype)
else:
op_result = result
return op_result
