@@ -117,21 +117,22 @@ def compute(self, method: str) -> Series:
117
117
# Save index and reset to default index to avoid performance impact
118
118
# from when index contains duplicates
119
119
original_index : Index = self .obj .index
120
- cur_series = self .obj .reset_index (drop = True )
120
+ default_index = self .obj .reset_index (drop = True )
121
121
122
- # slow method
123
- if n >= len (cur_series ):
122
+ # Slower method used when taking the full length of the series
123
+ # In this case, it is equivalent to a sort.
124
+ if n >= len (default_index ):
124
125
ascending = method == "nsmallest"
125
- final_series = cur_series .sort_values (
126
- ascending = ascending , kind = "stable"
127
- ). head ( n )
128
- final_series .index = original_index .take (final_series .index )
129
- return final_series
126
+ result = default_index .sort_values ( ascending = ascending , kind = "stable" ). head (
127
+ n
128
+ )
129
+ result .index = original_index .take (result .index )
130
+ return result
130
131
131
- dropped = cur_series .dropna ()
132
- nan_index = cur_series .drop (dropped .index )
132
+ # Fast method used in the general case
133
+ dropped = default_index .dropna ()
134
+ nan_index = default_index .drop (dropped .index )
133
135
134
- # fast method
135
136
new_dtype = dropped .dtype
136
137
137
138
# Similar to algorithms._ensure_data
@@ -170,7 +171,7 @@ def compute(self, method: str) -> Series:
170
171
else :
171
172
kth_val = np .nan
172
173
(ns ,) = np .nonzero (arr <= kth_val )
173
- inds = ns [arr [ns ].argsort (kind = "mergesort " )]
174
+ inds = ns [arr [ns ].argsort (kind = "stable " )]
174
175
175
176
if self .keep != "all" :
176
177
inds = inds [:n ]
@@ -185,9 +186,9 @@ def compute(self, method: str) -> Series:
185
186
# reverse indices
186
187
inds = narr - 1 - inds
187
188
188
- final_series = concat ([dropped .iloc [inds ], nan_index ]).iloc [:findex ]
189
- final_series .index = original_index .take (final_series .index )
190
- return final_series
189
+ result = concat ([dropped .iloc [inds ], nan_index ]).iloc [:findex ]
190
+ result .index = original_index .take (result .index )
191
+ return result
191
192
192
193
193
194
class SelectNFrame (SelectN [DataFrame ]):
0 commit comments