@@ -339,6 +339,26 @@ def _ea_wrap_cython_operation(
339
339
** kwargs ,
340
340
)
341
341
342
+ npvalues = self ._ea_to_cython_values (values )
343
+
344
+ res_values = self ._cython_op_ndim_compat (
345
+ npvalues ,
346
+ min_count = min_count ,
347
+ ngroups = ngroups ,
348
+ comp_ids = comp_ids ,
349
+ mask = None ,
350
+ ** kwargs ,
351
+ )
352
+
353
+ if self .how in ["rank" ]:
354
+ # i.e. how in WrappedCythonOp.cast_blocklist, since
355
+ # other cast_blocklist methods dont go through cython_operation
356
+ return res_values
357
+
358
+ return self ._reconstruct_ea_result (values , res_values )
359
+
360
+ def _ea_to_cython_values (self , values : ExtensionArray ):
361
+ # GH#43682
342
362
if isinstance (values , (DatetimeArray , PeriodArray , TimedeltaArray )):
343
363
# All of the functions implemented here are ordinal, so we can
344
364
# operate on the tz-naive equivalents
@@ -356,22 +376,7 @@ def _ea_wrap_cython_operation(
356
376
raise NotImplementedError (
357
377
f"function is not implemented for this dtype: { values .dtype } "
358
378
)
359
-
360
- res_values = self ._cython_op_ndim_compat (
361
- npvalues ,
362
- min_count = min_count ,
363
- ngroups = ngroups ,
364
- comp_ids = comp_ids ,
365
- mask = None ,
366
- ** kwargs ,
367
- )
368
-
369
- if self .how in ["rank" ]:
370
- # i.e. how in WrappedCythonOp.cast_blocklist, since
371
- # other cast_blocklist methods dont go through cython_operation
372
- return res_values
373
-
374
- return self ._reconstruct_ea_result (values , res_values )
379
+ return npvalues
375
380
376
381
def _reconstruct_ea_result (self , values , res_values ):
377
382
"""
@@ -387,6 +392,7 @@ def _reconstruct_ea_result(self, values, res_values):
387
392
return cls ._from_sequence (res_values , dtype = dtype )
388
393
389
394
elif needs_i8_conversion (values .dtype ):
395
+ assert res_values .dtype .kind != "f" # just to be on the safe side
390
396
i8values = res_values .view ("i8" )
391
397
return type (values )(i8values , dtype = values .dtype )
392
398
@@ -577,9 +583,12 @@ def _call_cython_op(
577
583
cutoff = max (1 , min_count )
578
584
empty_groups = counts < cutoff
579
585
if empty_groups .any ():
580
- # Note: this conversion could be lossy, see GH#40767
581
- result = result .astype ("float64" )
582
- result [empty_groups ] = np .nan
586
+ if result_mask is not None and self .uses_mask ():
587
+ assert result_mask [empty_groups ].all ()
588
+ else :
589
+ # Note: this conversion could be lossy, see GH#40767
590
+ result = result .astype ("float64" )
591
+ result [empty_groups ] = np .nan
583
592
584
593
result = result .T
585
594
0 commit comments