Skip to content

Commit bca157d

Browse files
improve logic in str_map
1 parent 607b95e commit bca157d

File tree

1 file changed

+22
-26
lines changed

1 file changed

+22
-26
lines changed

pandas/core/arrays/string_.py

Lines changed: 22 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -760,32 +760,28 @@ def _str_map(
760760
convert = convert and not np.all(mask)
761761

762762
if is_integer_dtype(dtype) or is_bool_dtype(dtype):
763-
# if is_integer_dtype(dtype):
764-
# na_value = np.nan
765-
# else:
766-
# na_value = False
767-
try:
768-
result = lib.map_infer_mask(
769-
arr,
770-
f,
771-
mask.view("uint8"),
772-
convert=False,
773-
na_value=na_value,
774-
dtype=np.dtype(cast(type, dtype)),
775-
)
776-
return result
777-
778-
except ValueError:
779-
result = lib.map_infer_mask(
780-
arr,
781-
f,
782-
mask.view("uint8"),
783-
convert=False,
784-
na_value=na_value,
785-
)
786-
if convert and result.dtype == object:
787-
result = lib.maybe_convert_objects(result)
788-
return result
763+
na_value_is_na = isna(na_value)
764+
if na_value_is_na:
765+
if is_integer_dtype(dtype):
766+
na_value = 0
767+
else:
768+
na_value = True
769+
770+
result = lib.map_infer_mask(
771+
arr,
772+
f,
773+
mask.view("uint8"),
774+
convert=False,
775+
na_value=na_value,
776+
dtype=np.dtype(cast(type, dtype)),
777+
)
778+
if na_value_is_na and mask.any():
779+
if is_integer_dtype(dtype):
780+
result = result.astype("float64")
781+
else:
782+
result = result.astype("object")
783+
result[mask] = np.nan
784+
return result
789785

790786
elif is_string_dtype(dtype) and not is_object_dtype(dtype):
791787
# i.e. StringDtype

0 commit comments

Comments
 (0)