@@ -981,40 +981,43 @@ def putmask(self, mask, new) -> list[Block]:
981
981
new = self .fill_value
982
982
983
983
new = self ._standardize_fill_value (new )
984
+ new = extract_array (new , extract_numpy = True )
984
985
985
- if self ._can_hold_element (new ):
986
- putmask_without_repeat (values .T , mask , new )
986
+ if noop :
987
987
return [self ]
988
988
989
- elif noop :
989
+ try :
990
+ casted = np_can_hold_element (values .dtype , new )
991
+ putmask_without_repeat (values .T , mask , casted )
990
992
return [self ]
993
+ except LossySetitemError :
991
994
992
- elif self .ndim == 1 or self .shape [0 ] == 1 :
993
- # no need to split columns
995
+ if self .ndim == 1 or self .shape [0 ] == 1 :
996
+ # no need to split columns
994
997
995
- if not is_list_like (new ):
996
- # using just new[indexer] can't save us the need to cast
997
- return self .coerce_to_target_dtype (new ).putmask (mask , new )
998
- else :
999
- indexer = mask .nonzero ()[0 ]
1000
- nb = self .setitem (indexer , new [indexer ])
1001
- return [nb ]
998
+ if not is_list_like (new ):
999
+ # using just new[indexer] can't save us the need to cast
1000
+ return self .coerce_to_target_dtype (new ).putmask (mask , new )
1001
+ else :
1002
+ indexer = mask .nonzero ()[0 ]
1003
+ nb = self .setitem (indexer , new [indexer ])
1004
+ return [nb ]
1002
1005
1003
- else :
1004
- is_array = isinstance (new , np .ndarray )
1006
+ else :
1007
+ is_array = isinstance (new , np .ndarray )
1005
1008
1006
- res_blocks = []
1007
- nbs = self ._split ()
1008
- for i , nb in enumerate (nbs ):
1009
- n = new
1010
- if is_array :
1011
- # we have a different value per-column
1012
- n = new [:, i : i + 1 ]
1009
+ res_blocks = []
1010
+ nbs = self ._split ()
1011
+ for i , nb in enumerate (nbs ):
1012
+ n = new
1013
+ if is_array :
1014
+ # we have a different value per-column
1015
+ n = new [:, i : i + 1 ]
1013
1016
1014
- submask = orig_mask [:, i : i + 1 ]
1015
- rbs = nb .putmask (submask , n )
1016
- res_blocks .extend (rbs )
1017
- return res_blocks
1017
+ submask = orig_mask [:, i : i + 1 ]
1018
+ rbs = nb .putmask (submask , n )
1019
+ res_blocks .extend (rbs )
1020
+ return res_blocks
1018
1021
1019
1022
def where (self , other , cond ) -> list [Block ]:
1020
1023
"""
@@ -1057,9 +1060,32 @@ def where(self, other, cond) -> list[Block]:
1057
1060
casted = np_can_hold_element (values .dtype , other )
1058
1061
except (ValueError , TypeError , LossySetitemError ):
1059
1062
# we cannot coerce, return a compat dtype
1060
- block = self .coerce_to_target_dtype (other )
1061
- blocks = block .where (orig_other , cond )
1062
- return self ._maybe_downcast (blocks , "infer" )
1063
+
1064
+ if self .ndim == 1 or self .shape [0 ] == 1 :
1065
+ # no need to split columns
1066
+
1067
+ block = self .coerce_to_target_dtype (other )
1068
+ blocks = block .where (orig_other , cond )
1069
+ return self ._maybe_downcast (blocks , "infer" )
1070
+
1071
+ else :
1072
+ # since _maybe_downcast would split blocks anyway, we
1073
+ # can avoid some potential upcast/downcast by splitting
1074
+ # on the front end.
1075
+ is_array = isinstance (other , (np .ndarray , ExtensionArray ))
1076
+
1077
+ res_blocks = []
1078
+ nbs = self ._split ()
1079
+ for i , nb in enumerate (nbs ):
1080
+ oth = other
1081
+ if is_array :
1082
+ # we have a different value per-column
1083
+ oth = other [:, i : i + 1 ]
1084
+
1085
+ submask = cond [:, i : i + 1 ]
1086
+ rbs = nb .where (oth , submask )
1087
+ res_blocks .extend (rbs )
1088
+ return res_blocks
1063
1089
1064
1090
else :
1065
1091
other = casted
0 commit comments