def _quantile_ea_fallback(values: ExtensionArray, qs: npt.NDArray[np.float64],
interpolation: str) -> ExtensionArray:
"""
quantile compatibility for ExtensionArray subclasses that do not
implement `_from_factorized`, e.g. IntegerArray.
Notes
-----
We assume that all impacted cases are 1D-only.
"""
mask = np.atleast_2d(np.asarray(values.isna()))
npvalues = np.atleast_2d(np.asarray(values))
res = _quantile_with_mask(
npvalues,
mask=mask,
fill_value=values.dtype.na_value,
qs=qs,
interpolation=interpolation,
)
assert res.ndim == 2
assert res.shape[0] == 1
res = res[0]
try:
out = type(values)._from_sequence(res, dtype=values.dtype)
except TypeError:
# GH#42626: not able to safely cast Int64
# for floating point output
out = np.atleast_2d(np.asarray(res, dtype=np.float64))
return out
def _quantile_ea_fallback(values: ExtensionArray, qs: np.ndarray,
interpolation: str) -> ExtensionArray:
"""
quantile compatibility for ExtensionArray subclasses that do not
implement `_from_factorized`, e.g. IntegerArray.
Notes
-----
We assume that all impacted cases are 1D-only.
"""
mask = np.atleast_2d(np.asarray(values.isna()))
npvalues = np.atleast_2d(np.asarray(values))
res = _quantile_with_mask(
npvalues,
mask=mask,
fill_value=values.dtype.na_value,
qs=qs,
interpolation=interpolation,
)
assert res.ndim == 2
assert res.shape[0] == 1
res = res[0]
out = type(values)._from_sequence(res, dtype=values.dtype)
return out
def quantile_ea_compat(values: ExtensionArray, qs, interpolation: str,
axis: int) -> ExtensionArray:
"""
ExtensionArray compatibility layer for quantile_with_mask.
We pretend that an ExtensionArray with shape (N,) is actually (1, N,)
for compatibility with non-EA code.
Parameters
----------
values : ExtensionArray
qs : a scalar or list of the quantiles to be computed
interpolation: str
axis : int
Returns
-------
ExtensionArray
"""
# TODO(EA2D): make-believe not needed with 2D EAs
orig = values
# asarray needed for Sparse, see GH#24600
mask = np.asarray(values.isna())
mask = np.atleast_2d(mask)
# error: Incompatible types in assignment (expression has type "ndarray", variable
# has type "ExtensionArray")
values, fill_value = values._values_for_factorize(
) # type: ignore[assignment]
# error: No overload variant of "atleast_2d" matches argument type "ExtensionArray"
values = np.atleast_2d(values) # type: ignore[call-overload]
# error: Argument 1 to "quantile_with_mask" has incompatible type "ExtensionArray";
# expected "ndarray"
result = quantile_with_mask(
values,
mask,
fill_value,
qs,
interpolation,
axis # type: ignore[arg-type]
)
if not is_sparse(orig.dtype):
# shape[0] should be 1 as long as EAs are 1D
if result.ndim == 1:
# i.e. qs was originally a scalar
assert result.shape == (1, ), result.shape
result = type(orig)._from_factorized(result, orig)
else:
assert result.shape == (1, len(qs)), result.shape
result = type(orig)._from_factorized(result[0], orig)
# error: Incompatible return value type (got "ndarray", expected "ExtensionArray")
return result # type: ignore[return-value]
def _ea_wrap_cython_operation(
self,
values: ExtensionArray,
min_count: int,
ngroups: int,
comp_ids: np.ndarray,
**kwargs,
) -> ArrayLike:
"""
If we have an ExtensionArray, unwrap, call _cython_operation, and
re-wrap if appropriate.
"""
if isinstance(values, BaseMaskedArray) and self.uses_mask():
return self._masked_ea_wrap_cython_operation(
values,
min_count=min_count,
ngroups=ngroups,
comp_ids=comp_ids,
**kwargs,
)
elif isinstance(values, Categorical) and self.uses_mask():
assert self.how == "rank" # the only one implemented ATM
assert values.ordered # checked earlier
mask = values.isna()
npvalues = values._ndarray
res_values = self._cython_op_ndim_compat(
npvalues,
min_count=min_count,
ngroups=ngroups,
comp_ids=comp_ids,
mask=mask,
**kwargs,
)
# If we ever have more than just "rank" here, we'll need to do
# `if self.how in self.cast_blocklist` like we do for other dtypes.
return res_values
npvalues = self._ea_to_cython_values(values)
res_values = self._cython_op_ndim_compat(
npvalues,
min_count=min_count,
ngroups=ngroups,
comp_ids=comp_ids,
mask=None,
**kwargs,
)
if self.how in self.cast_blocklist:
# i.e. how in ["rank"], since other cast_blocklist methods dont go
# through cython_operation
return res_values
return self._reconstruct_ea_result(values, res_values)
def quantile_ea_compat(values: ExtensionArray, qs, interpolation: str,
axis: int) -> ExtensionArray:
"""
ExtensionArray compatibility layer for quantile_with_mask.
We pretend that an ExtensionArray with shape (N,) is actually (1, N,)
for compatibility with non-EA code.
Parameters
----------
values : ExtensionArray
qs : a scalar or list of the quantiles to be computed
interpolation: str
axis : int
Returns
-------
ExtensionArray
"""
# TODO(EA2D): make-believe not needed with 2D EAs
orig = values
# asarray needed for Sparse, see GH#24600
mask = np.asarray(values.isna())
mask = np.atleast_2d(mask)
values, fill_value = values._values_for_factorize()
values = np.atleast_2d(values)
result = quantile_with_mask(values, mask, fill_value, qs, interpolation,
axis)
if not is_sparse(orig.dtype):
# shape[0] should be 1 as long as EAs are 1D
if result.ndim == 1:
# i.e. qs was originally a scalar
assert result.shape == (1, ), result.shape
result = type(orig)._from_factorized(result, orig)
else:
assert result.shape == (1, len(qs)), result.shape
result = type(orig)._from_factorized(result[0], orig)
return result
def _quantile_ea_compat(
values: ExtensionArray, qs: np.ndarray, interpolation: str
) -> ExtensionArray:
"""
ExtensionArray compatibility layer for _quantile_with_mask.
We pretend that an ExtensionArray with shape (N,) is actually (1, N,)
for compatibility with non-EA code.
Parameters
----------
values : ExtensionArray
qs : np.ndarray[float64]
interpolation: str
Returns
-------
ExtensionArray
"""
# TODO(EA2D): make-believe not needed with 2D EAs
orig = values
# asarray needed for Sparse, see GH#24600
mask = np.asarray(values.isna())
mask = np.atleast_2d(mask)
arr, fill_value = values._values_for_factorize()
arr = np.atleast_2d(arr)
result = _quantile_with_mask(arr, mask, fill_value, qs, interpolation)
if not is_sparse(orig.dtype):
# shape[0] should be 1 as long as EAs are 1D
if orig.ndim == 2:
# i.e. DatetimeArray
result = type(orig)._from_factorized(result, orig)
else:
assert result.shape == (1, len(qs)), result.shape
result = type(orig)._from_factorized(result[0], orig)
# error: Incompatible return value type (got "ndarray", expected "ExtensionArray")
return result # type: ignore[return-value]
