@@ -981,40 +981,43 @@ def putmask(self, mask, new) -> list[Block]:
981981 new = self .fill_value
982982
983983 new = self ._standardize_fill_value (new )
984+ new = extract_array (new , extract_numpy = True )
984985
985- if self ._can_hold_element (new ):
986- putmask_without_repeat (values .T , mask , new )
986+ if noop :
987987 return [self ]
988988
989- elif noop :
989+ try :
990+ casted = np_can_hold_element (values .dtype , new )
991+ putmask_without_repeat (values .T , mask , casted )
990992 return [self ]
993+ except LossySetitemError :
991994
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
994997
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 ]
10021005
1003- else :
1004- is_array = isinstance (new , np .ndarray )
1006+ else :
1007+ is_array = isinstance (new , np .ndarray )
10051008
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 ]
10131016
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
10181021
10191022 def where (self , other , cond ) -> list [Block ]:
10201023 """
@@ -1057,9 +1060,32 @@ def where(self, other, cond) -> list[Block]:
10571060 casted = np_can_hold_element (values .dtype , other )
10581061 except (ValueError , TypeError , LossySetitemError ):
10591062 # 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
10631089
10641090 else :
10651091 other = casted
0 commit comments