@@ -288,20 +288,14 @@ def test_rank_tie_methods(
288
288
pytest .skip ("int64/str does not support NaN" )
289
289
290
290
ser = ser if dtype is None else ser .astype (dtype )
291
+ if dtype in ["float64[pyarrow]" , "Float64" ] and not using_nan_is_na :
292
+ # TODO: use ser.replace(np.nan, NA) once that works
293
+ ser [np .isnan (ser .to_numpy (dtype = np .float64 , na_value = np .nan ))] = NA
294
+ mask = np .isnan (exp )
295
+ exp = exp .astype (object )
296
+ exp [mask ] = NA
297
+
291
298
result = ser .rank (method = method )
292
- if dtype == "float64[pyarrow]" and not using_nan_is_na :
293
- # the NaNs are not treated as NA
294
- exp = exp .copy ()
295
- if method == "average" :
296
- exp [np .isnan (ser )] = 9.5
297
- elif method == "dense" :
298
- exp [np .isnan (ser )] = 6
299
- elif method == "max" :
300
- exp [np .isnan (ser )] = 10
301
- elif method == "min" :
302
- exp [np .isnan (ser )] = 9
303
- elif method == "first" :
304
- exp [np .isnan (ser )] = [9 , 10 ]
305
299
306
300
if dtype == "string[pyarrow]" and not using_nan_is_na :
307
301
mask = np .isnan (exp )
@@ -368,7 +362,7 @@ def test_rank_tie_methods_on_infs_nans(
368
362
order = [ranks [1 ], ranks [0 ], ranks [2 ]]
369
363
elif na_option == "bottom" :
370
364
order = [ranks [0 ], ranks [2 ], ranks [1 ]]
371
- elif dtype == "float64[pyarrow]" and not using_nan_is_na :
365
+ elif dtype in ( "float64[pyarrow]" , "Float64" ) and not using_nan_is_na :
372
366
order = [ranks [0 ], [NA ] * chunk , ranks [1 ]]
373
367
else :
374
368
order = [ranks [0 ], [np .nan ] * chunk , ranks [1 ]]
0 commit comments