Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

ENH: support downcasting of nullable EAs in pd.to_numeric #38746

Merged
merged 20 commits into from
Dec 30, 2020
Merged
Show file tree
Hide file tree
Changes from 11 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions doc/source/whatsnew/v1.3.0.rst
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ Other enhancements
- Improved consistency of error message when passing an invalid ``win_type`` argument in :class:`Window` (:issue:`15969`)
- :func:`pandas.read_sql_query` now accepts a ``dtype`` argument to cast the columnar data from the SQL database based on user input (:issue:`10285`)
- Improved integer type mapping from pandas to SQLAlchemy when using :meth:`DataFrame.to_sql` (:issue:`35076`)
- :func:`to_numeric` now supports downcasting of nullable ``ExtensionDtype`` objects (:issue:`33013`)

.. ---------------------------------------------------------------------------

Expand Down
22 changes: 22 additions & 0 deletions pandas/core/tools/numeric.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@
ensure_object,
is_datetime_or_timedelta_dtype,
is_decimal,
is_float_dtype,
is_integer_dtype,
is_number,
is_numeric_dtype,
is_scalar,
Expand All @@ -15,6 +17,7 @@
from pandas.core.dtypes.generic import ABCIndex, ABCSeries

import pandas as pd
from pandas.core.arrays.numeric import NumericArray


def to_numeric(arg, errors="raise", downcast=None):
Expand Down Expand Up @@ -142,6 +145,12 @@ def to_numeric(arg, errors="raise", downcast=None):
else:
values = arg

if isinstance(values, NumericArray):
jreback marked this conversation as resolved.
Show resolved Hide resolved
mask = values._mask
values = values._data[~mask]
else:
mask = None

values_dtype = getattr(values, "dtype", None)
if is_numeric_dtype(values_dtype):
pass
Expand Down Expand Up @@ -188,6 +197,19 @@ def to_numeric(arg, errors="raise", downcast=None):
if values.dtype == dtype:
break

if mask is not None:
data = np.zeros(mask.shape, dtype=values.dtype)
data[~mask] = values

if is_integer_dtype(data):
from pandas.core.arrays import IntegerArray

values = IntegerArray(data, mask)
elif is_float_dtype(data):
from pandas.core.arrays import FloatingArray

values = FloatingArray(data, mask)

if is_series:
return arg._constructor(values, index=arg.index, name=arg.name)
elif is_index:
Expand Down
31 changes: 31 additions & 0 deletions pandas/tests/tools/test_to_numeric.py
Original file line number Diff line number Diff line change
Expand Up @@ -725,3 +725,34 @@ def test_to_numeric_from_nullable_string(values, expected):
s = Series(values, dtype="string")
result = to_numeric(s)
tm.assert_series_equal(result, expected)


@pytest.mark.parametrize(
"data, input_dtype, downcast, expected_dtype",
(
([1, 1], "Int64", "integer", "Int8"),
([1.0, 1.0], "Float64", "integer", "Int8"),
arw2019 marked this conversation as resolved.
Show resolved Hide resolved
([1.0, 1.1], "Float64", "integer", "Float64"),
([1, pd.NA], "Int64", "integer", "Int8"),
([450, 300], "Int64", "integer", "Int16"),
([1, 1], "Int64", "signed", "Int8"),
([1.0, 1.0], "Float32", "signed", "Int8"),
([1.0, 1.1], "Float64", "signed", "Float64"),
([1, pd.NA], "Int64", "signed", "Int8"),
jreback marked this conversation as resolved.
Show resolved Hide resolved
([450, -300], "Int64", "signed", "Int16"),
([1, 1], "Int64", "unsigned", "UInt8"),
([1.0, 1.0], "Float32", "unsigned", "UInt8"),
([1.0, 1.1], "Float64", "unsigned", "Float64"),
([1, pd.NA], "Int64", "unsigned", "UInt8"),
([450, -300], "Int64", "unsigned", "Int64"),
([-1, -1], "Int32", "unsigned", "Int32"),
([1, 1], "Float64", "float", "Float32"),
([1, 1.1], "Float64", "float", "Float32"),
([1, 1], "Float64", "integer", "Int8"),
arw2019 marked this conversation as resolved.
Show resolved Hide resolved
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: maybe move this one up to the other "integer" downcast cases

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done

),
)
def test_downcast_nullable_numeric(data, input_dtype, downcast, expected_dtype):
arr = pd.array(data, dtype=input_dtype)
result = pd.to_numeric(arr, downcast=downcast)
expected = pd.array(data, dtype=expected_dtype)
tm.assert_extension_array_equal(result, expected)