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 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]
