Skip to content

Commit

Permalink
improve logic in str_map
Browse files Browse the repository at this point in the history
  • Loading branch information
jorisvandenbossche committed Apr 27, 2024
1 parent 607b95e commit bca157d
Showing 1 changed file with 22 additions and 26 deletions.
48 changes: 22 additions & 26 deletions pandas/core/arrays/string_.py
Expand Up @@ -760,32 +760,28 @@ def _str_map(
convert = convert and not np.all(mask)

if is_integer_dtype(dtype) or is_bool_dtype(dtype):
# if is_integer_dtype(dtype):
# na_value = np.nan
# else:
# na_value = False
try:
result = lib.map_infer_mask(
arr,
f,
mask.view("uint8"),
convert=False,
na_value=na_value,
dtype=np.dtype(cast(type, dtype)),
)
return result

except ValueError:
result = lib.map_infer_mask(
arr,
f,
mask.view("uint8"),
convert=False,
na_value=na_value,
)
if convert and result.dtype == object:
result = lib.maybe_convert_objects(result)
return result
na_value_is_na = isna(na_value)
if na_value_is_na:
if is_integer_dtype(dtype):
na_value = 0
else:
na_value = True

result = lib.map_infer_mask(
arr,
f,
mask.view("uint8"),
convert=False,
na_value=na_value,
dtype=np.dtype(cast(type, dtype)),
)
if na_value_is_na and mask.any():
if is_integer_dtype(dtype):
result = result.astype("float64")
else:
result = result.astype("object")
result[mask] = np.nan
return result

elif is_string_dtype(dtype) and not is_object_dtype(dtype):
# i.e. StringDtype
Expand Down

0 comments on commit bca157d

Please sign in to comment.