Skip to content

Commit a09552e

Browse files
committed
Save original index and remap after function completes.
1 parent bc24e84 commit a09552e

File tree

1 file changed

+14
-5
lines changed

1 file changed

+14
-5
lines changed

pandas/core/methods/selectn.py

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -111,13 +111,20 @@ def compute(self, method: str) -> Series:
111111
if n <= 0:
112112
return self.obj[[]]
113113

114-
dropped = self.obj.dropna()
115-
nan_index = self.obj.drop(dropped.index)
114+
# Save index and reset to default index to avoid performance impact
115+
# from when index contains duplicates
116+
original_index = self.obj.index
117+
cur_series = self.obj.reset_index(drop=True)
118+
119+
dropped = cur_series.dropna()
120+
nan_index = cur_series.drop(dropped.index)
116121

117122
# slow method
118-
if n >= len(self.obj):
123+
if n >= len(cur_series):
119124
ascending = method == "nsmallest"
120-
return self.obj.sort_values(ascending=ascending).head(n)
125+
final_series = cur_series.sort_values(ascending=ascending, kind="mergesort").head(n)
126+
final_series.index = original_index.take(final_series.index)
127+
return final_series
121128

122129
# fast method
123130
new_dtype = dropped.dtype
@@ -173,7 +180,9 @@ def compute(self, method: str) -> Series:
173180
# reverse indices
174181
inds = narr - 1 - inds
175182

176-
return concat([dropped.iloc[inds], nan_index]).iloc[:findex]
183+
final_series = concat([dropped.iloc[inds], nan_index]).iloc[:findex]
184+
final_series.index = original_index.take(final_series.index)
185+
return final_series
177186

178187

179188
class SelectNFrame(SelectN[DataFrame]):

0 commit comments

Comments
 (0)