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]