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]