@@ -26,6 +26,7 @@ from numpy cimport (
26
26
int16_t,
27
27
int32_t,
28
28
int64_t,
29
+ intp_t,
29
30
ndarray,
30
31
uint8_t,
31
32
uint16_t,
@@ -1105,14 +1106,13 @@ def rank_2d(
1105
1106
Py_ssize_t infs
1106
1107
ndarray[float64_t, ndim= 2 ] ranks
1107
1108
ndarray[rank_t, ndim= 2 ] values
1108
- ndarray[int64_t , ndim= 2 ] argsorted
1109
+ ndarray[intp_t , ndim= 2 ] argsort_indexer
1109
1110
ndarray[uint8_t, ndim= 2 ] mask
1110
1111
rank_t val, nan_value
1111
1112
float64_t count, sum_ranks = 0.0
1112
1113
int tiebreak = 0
1113
1114
int64_t idx
1114
1115
bint check_mask, condition, keep_na
1115
- const int64_t[:] labels
1116
1116
1117
1117
tiebreak = tiebreakers[ties_method]
1118
1118
@@ -1158,40 +1158,19 @@ def rank_2d(
1158
1158
1159
1159
n, k = (< object > values).shape
1160
1160
ranks = np.empty((n, k), dtype = ' f8' )
1161
- # For compatibility when calling rank_1d
1162
- labels = np.zeros(k, dtype = np.int64)
1163
1161
1164
- if rank_t is object :
1165
- try :
1166
- _as = values.argsort(1 )
1167
- except TypeError :
1168
- values = in_arr
1169
- for i in range (len (values)):
1170
- ranks[i] = rank_1d(
1171
- in_arr[i],
1172
- labels = labels,
1173
- ties_method = ties_method,
1174
- ascending = ascending,
1175
- pct = pct
1176
- )
1177
- if axis == 0 :
1178
- return ranks.T
1179
- else :
1180
- return ranks
1162
+ if tiebreak == TIEBREAK_FIRST:
1163
+ # need to use a stable sort here
1164
+ argsort_indexer = values.argsort(axis = 1 , kind = ' mergesort' )
1165
+ if not ascending:
1166
+ tiebreak = TIEBREAK_FIRST_DESCENDING
1181
1167
else :
1182
- if tiebreak == TIEBREAK_FIRST:
1183
- # need to use a stable sort here
1184
- _as = values.argsort(axis = 1 , kind = ' mergesort' )
1185
- if not ascending:
1186
- tiebreak = TIEBREAK_FIRST_DESCENDING
1187
- else :
1188
- _as = values.argsort(1 )
1168
+ argsort_indexer = values.argsort(1 )
1189
1169
1190
1170
if not ascending:
1191
- _as = _as [:, ::- 1 ]
1171
+ argsort_indexer = argsort_indexer [:, ::- 1 ]
1192
1172
1193
- values = _take_2d(values, _as)
1194
- argsorted = _as.astype(' i8' )
1173
+ values = _take_2d(values, argsort_indexer)
1195
1174
1196
1175
for i in range (n):
1197
1176
dups = sum_ranks = infs = 0
@@ -1200,7 +1179,7 @@ def rank_2d(
1200
1179
count = 0.0
1201
1180
for j in range (k):
1202
1181
val = values[i, j]
1203
- idx = argsorted [i, j]
1182
+ idx = argsort_indexer [i, j]
1204
1183
if keep_na and check_mask and mask[i, idx]:
1205
1184
ranks[i, idx] = NaN
1206
1185
infs += 1
@@ -1215,38 +1194,38 @@ def rank_2d(
1215
1194
condition = (
1216
1195
j == k - 1 or
1217
1196
are_diff(values[i, j + 1 ], val) or
1218
- (keep_na and check_mask and mask[i, argsorted [i, j + 1 ]])
1197
+ (keep_na and check_mask and mask[i, argsort_indexer [i, j + 1 ]])
1219
1198
)
1220
1199
else :
1221
1200
condition = (
1222
1201
j == k - 1 or
1223
1202
values[i, j + 1 ] != val or
1224
- (keep_na and check_mask and mask[i, argsorted [i, j + 1 ]])
1203
+ (keep_na and check_mask and mask[i, argsort_indexer [i, j + 1 ]])
1225
1204
)
1226
1205
1227
1206
if condition:
1228
1207
if tiebreak == TIEBREAK_AVERAGE:
1229
1208
for z in range (j - dups + 1 , j + 1 ):
1230
- ranks[i, argsorted [i, z]] = sum_ranks / dups
1209
+ ranks[i, argsort_indexer [i, z]] = sum_ranks / dups
1231
1210
elif tiebreak == TIEBREAK_MIN:
1232
1211
for z in range (j - dups + 1 , j + 1 ):
1233
- ranks[i, argsorted [i, z]] = j - dups + 2
1212
+ ranks[i, argsort_indexer [i, z]] = j - dups + 2
1234
1213
elif tiebreak == TIEBREAK_MAX:
1235
1214
for z in range (j - dups + 1 , j + 1 ):
1236
- ranks[i, argsorted [i, z]] = j + 1
1215
+ ranks[i, argsort_indexer [i, z]] = j + 1
1237
1216
elif tiebreak == TIEBREAK_FIRST:
1238
1217
if rank_t is object :
1239
1218
raise ValueError (' first not supported for non-numeric data' )
1240
1219
else :
1241
1220
for z in range (j - dups + 1 , j + 1 ):
1242
- ranks[i, argsorted [i, z]] = z + 1
1221
+ ranks[i, argsort_indexer [i, z]] = z + 1
1243
1222
elif tiebreak == TIEBREAK_FIRST_DESCENDING:
1244
1223
for z in range (j - dups + 1 , j + 1 ):
1245
- ranks[i, argsorted [i, z]] = 2 * j - z - dups + 2
1224
+ ranks[i, argsort_indexer [i, z]] = 2 * j - z - dups + 2
1246
1225
elif tiebreak == TIEBREAK_DENSE:
1247
1226
total_tie_count += 1
1248
1227
for z in range (j - dups + 1 , j + 1 ):
1249
- ranks[i, argsorted [i, z]] = total_tie_count
1228
+ ranks[i, argsort_indexer [i, z]] = total_tie_count
1250
1229
sum_ranks = dups = 0
1251
1230
if pct:
1252
1231
if tiebreak == TIEBREAK_DENSE:
0 commit comments