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. """ # TODO: general case implementation overridable by EAs. 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, ) if isinstance(values, (DatetimeArray, PeriodArray, TimedeltaArray)): # All of the functions implemented here are ordinal, so we can # operate on the tz-naive equivalents npvalues = values._ndarray.view("M8[ns]") elif isinstance(values.dtype, (BooleanDtype, IntegerDtype)): # IntegerArray or BooleanArray npvalues = values.to_numpy("float64", na_value=np.nan) elif isinstance(values.dtype, FloatingDtype): # FloatingArray npvalues = values.to_numpy(values.dtype.numpy_dtype, na_value=np.nan) elif isinstance(values.dtype, StringDtype): # StringArray npvalues = values.to_numpy(object, na_value=np.nan) else: raise NotImplementedError( f"function is not implemented for this dtype: {values.dtype}" ) 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 ["rank"]: # i.e. how in WrappedCythonOp.cast_blocklist, since # other cast_blocklist methods dont go through cython_operation return res_values return self._reconstruct_ea_result(values, res_values)
def _ea_to_cython_values(self, values: ExtensionArray) -> np.ndarray: # GH#43682 if isinstance(values, (DatetimeArray, PeriodArray, TimedeltaArray)): # All of the functions implemented here are ordinal, so we can # operate on the tz-naive equivalents npvalues = values._ndarray.view("M8[ns]") elif isinstance(values.dtype, (BooleanDtype, IntegerDtype)): # IntegerArray or BooleanArray npvalues = values.to_numpy("float64", na_value=np.nan) elif isinstance(values.dtype, FloatingDtype): # FloatingArray npvalues = values.to_numpy(values.dtype.numpy_dtype, na_value=np.nan) elif isinstance(values.dtype, StringDtype): # StringArray npvalues = values.to_numpy(object, na_value=np.nan) else: raise NotImplementedError( f"function is not implemented for this dtype: {values.dtype}" ) return npvalues
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. """ # TODO: general case implementation overridable by EAs. 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, ) orig_values = values if isinstance(orig_values, (DatetimeArray, PeriodArray)): # All of the functions implemented here are ordinal, so we can # operate on the tz-naive equivalents npvalues = orig_values._ndarray.view("M8[ns]") 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 ["rank"]: # i.e. how in WrappedCythonOp.cast_blocklist, since # other cast_blocklist methods dont go through cython_operation # preserve float64 dtype return res_values res_values = res_values.view("i8") result = type(orig_values)(res_values, dtype=orig_values.dtype) return result elif isinstance(orig_values, TimedeltaArray): # We have an ExtensionArray but not ExtensionDtype res_values = self._cython_op_ndim_compat( orig_values._ndarray, min_count=min_count, ngroups=ngroups, comp_ids=comp_ids, mask=None, **kwargs, ) if self.how in ["rank"]: # i.e. how in WrappedCythonOp.cast_blocklist, since # other cast_blocklist methods dont go through cython_operation # preserve float64 dtype return res_values # otherwise res_values has the same dtype as original values return type(orig_values)(res_values) elif isinstance(values.dtype, (BooleanDtype, _IntegerDtype)): # IntegerArray or BooleanArray npvalues = values.to_numpy("float64", na_value=np.nan) 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 ["rank"]: # i.e. how in WrappedCythonOp.cast_blocklist, since # other cast_blocklist methods dont go through cython_operation return res_values dtype = self._get_result_dtype(orig_values.dtype) cls = dtype.construct_array_type() return cls._from_sequence(res_values, dtype=dtype) elif isinstance(values.dtype, FloatingDtype): # FloatingArray npvalues = values.to_numpy( values.dtype.numpy_dtype, na_value=np.nan, ) 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 ["rank"]: # i.e. how in WrappedCythonOp.cast_blocklist, since # other cast_blocklist methods dont go through cython_operation return res_values dtype = self._get_result_dtype(orig_values.dtype) cls = dtype.construct_array_type() return cls._from_sequence(res_values, dtype=dtype) raise NotImplementedError( f"function is not implemented for this dtype: {values.dtype}")
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. """ # TODO: general case implementation overridable by EAs. orig_values = values if is_datetime64tz_dtype(values.dtype) or is_period_dtype( values.dtype): # All of the functions implemented here are ordinal, so we can # operate on the tz-naive equivalents npvalues = values.view("M8[ns]") res_values = self._cython_op_ndim_compat( # error: Argument 1 to "_cython_op_ndim_compat" of # "WrappedCythonOp" has incompatible type # "Union[ExtensionArray, ndarray]"; expected "ndarray" npvalues, # type: ignore[arg-type] min_count=min_count, ngroups=ngroups, comp_ids=comp_ids, mask=None, **kwargs, ) if self.how in ["rank"]: # i.e. how in WrappedCythonOp.cast_blocklist, since # other cast_blocklist methods dont go through cython_operation # preserve float64 dtype return res_values res_values = res_values.astype("i8", copy=False) # error: Too many arguments for "ExtensionArray" result = type(orig_values)( # type: ignore[call-arg] res_values, dtype=orig_values.dtype) return result elif is_integer_dtype(values.dtype) or is_bool_dtype(values.dtype): # IntegerArray or BooleanArray npvalues = values.to_numpy("float64", na_value=np.nan) 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 ["rank"]: # i.e. how in WrappedCythonOp.cast_blocklist, since # other cast_blocklist methods dont go through cython_operation return res_values dtype = self.get_result_dtype(orig_values.dtype) # error: Item "dtype[Any]" of "Union[dtype[Any], ExtensionDtype]" # has no attribute "construct_array_type" cls = dtype.construct_array_type() # type: ignore[union-attr] return cls._from_sequence(res_values, dtype=dtype) elif is_float_dtype(values.dtype): # FloatingArray # error: "ExtensionDtype" has no attribute "numpy_dtype" npvalues = values.to_numpy( values.dtype.numpy_dtype, # type: ignore[attr-defined] na_value=np.nan, ) 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 ["rank"]: # i.e. how in WrappedCythonOp.cast_blocklist, since # other cast_blocklist methods dont go through cython_operation return res_values dtype = self.get_result_dtype(orig_values.dtype) # error: Item "dtype[Any]" of "Union[dtype[Any], ExtensionDtype]" # has no attribute "construct_array_type" cls = dtype.construct_array_type() # type: ignore[union-attr] return cls._from_sequence(res_values, dtype=dtype) raise NotImplementedError( f"function is not implemented for this dtype: {values.dtype}")