@@ -100,13 +100,14 @@ def take_nd(
100
100
return arr .take (indexer , fill_value = fill_value , allow_fill = allow_fill )
101
101
102
102
arr = np .asarray (arr )
103
- return _take_nd_ndarray (arr , indexer , axis , fill_value , allow_fill )
103
+ return _take_nd_ndarray (arr , indexer , axis , None , fill_value , allow_fill )
104
104
105
105
106
106
def _take_nd_ndarray (
107
107
arr : np .ndarray ,
108
108
indexer ,
109
109
axis : int ,
110
+ out : np .ndarray | None ,
110
111
fill_value ,
111
112
allow_fill : bool ,
112
113
) -> np .ndarray :
@@ -117,11 +118,8 @@ def _take_nd_ndarray(
117
118
else :
118
119
indexer = ensure_platform_int (indexer )
119
120
120
- if not allow_fill :
121
- return arr .take (indexer , axis = axis )
122
-
123
- dtype , fill_value , mask_info = _take_preprocess_indexer_and_fill_value (
124
- arr , indexer , fill_value
121
+ indexer , dtype , fill_value , mask_info = _take_preprocess_indexer_and_fill_value (
122
+ arr , indexer , out , fill_value , allow_fill
125
123
)
126
124
127
125
flip_order = False
@@ -131,20 +129,23 @@ def _take_nd_ndarray(
131
129
if flip_order :
132
130
arr = arr .T
133
131
axis = arr .ndim - axis - 1
132
+ if out is not None :
133
+ out = out .T
134
134
135
135
# at this point, it's guaranteed that dtype can hold both the arr values
136
136
# and the fill_value
137
- out_shape_ = list (arr .shape )
138
- out_shape_ [axis ] = len (indexer )
139
- out_shape = tuple (out_shape_ )
140
- if arr .flags .f_contiguous and axis == arr .ndim - 1 :
141
- # minor tweak that can make an order-of-magnitude difference
142
- # for dataframes initialized directly from 2-d ndarrays
143
- # (s.t. df.values is c-contiguous and df._mgr.blocks[0] is its
144
- # f-contiguous transpose)
145
- out = np .empty (out_shape , dtype = dtype , order = "F" )
146
- else :
147
- out = np .empty (out_shape , dtype = dtype )
137
+ if out is None :
138
+ out_shape_ = list (arr .shape )
139
+ out_shape_ [axis ] = len (indexer )
140
+ out_shape = tuple (out_shape_ )
141
+ if arr .flags .f_contiguous and axis == arr .ndim - 1 :
142
+ # minor tweak that can make an order-of-magnitude difference
143
+ # for dataframes initialized directly from 2-d ndarrays
144
+ # (s.t. df.values is c-contiguous and df._mgr.blocks[0] is its
145
+ # f-contiguous transpose)
146
+ out = np .empty (out_shape , dtype = dtype , order = "F" )
147
+ else :
148
+ out = np .empty (out_shape , dtype = dtype )
148
149
149
150
func = _get_take_nd_function (
150
151
arr .ndim , arr .dtype , out .dtype , axis = axis , mask_info = mask_info
@@ -188,8 +189,8 @@ def take_1d(
188
189
if not allow_fill :
189
190
return arr .take (indexer )
190
191
191
- dtype , fill_value , mask_info = _take_preprocess_indexer_and_fill_value (
192
- arr , indexer , fill_value
192
+ indexer , dtype , fill_value , mask_info = _take_preprocess_indexer_and_fill_value (
193
+ arr , indexer , None , fill_value , True
193
194
)
194
195
195
196
# at this point, it's guaranteed that dtype can hold both the arr values
@@ -515,22 +516,32 @@ def _take_2d_multi_object(
515
516
def _take_preprocess_indexer_and_fill_value (
516
517
arr : np .ndarray ,
517
518
indexer : np .ndarray ,
519
+ out : np .ndarray | None ,
518
520
fill_value ,
521
+ allow_fill : bool ,
519
522
):
520
523
mask_info = None
521
524
522
- # check for promotion based on types only (do this first because
523
- # it's faster than computing a mask)
524
- dtype , fill_value = maybe_promote (arr .dtype , fill_value )
525
- if dtype != arr .dtype :
526
- # check if promotion is actually required based on indexer
527
- mask = indexer == - 1
528
- needs_masking = mask .any ()
529
- mask_info = mask , needs_masking
530
- if not needs_masking :
531
- # if not, then depromote, set fill_value to dummy
532
- # (it won't be used but we don't want the cython code
533
- # to crash when trying to cast it to dtype)
534
- dtype , fill_value = arr .dtype , arr .dtype .type ()
535
-
536
- return dtype , fill_value , mask_info
525
+ if not allow_fill :
526
+ dtype , fill_value = arr .dtype , arr .dtype .type ()
527
+ mask_info = None , False
528
+ else :
529
+ # check for promotion based on types only (do this first because
530
+ # it's faster than computing a mask)
531
+ dtype , fill_value = maybe_promote (arr .dtype , fill_value )
532
+ if dtype != arr .dtype :
533
+ # check if promotion is actually required based on indexer
534
+ mask = indexer == - 1
535
+ needs_masking = mask .any ()
536
+ mask_info = mask , needs_masking
537
+ if needs_masking :
538
+ if out is not None and out .dtype != dtype :
539
+ raise TypeError ("Incompatible type for fill_value" )
540
+ else :
541
+ # if not, then depromote, set fill_value to dummy
542
+ # (it won't be used but we don't want the cython code
543
+ # to crash when trying to cast it to dtype)
544
+ dtype , fill_value = arr .dtype , arr .dtype .type ()
545
+
546
+ indexer = ensure_platform_int (indexer )
547
+ return indexer , dtype , fill_value , mask_info
0 commit comments