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 7 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
19 changes: 19 additions & 0 deletions pandas/core/tools/numeric.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,9 @@
ensure_object,
is_datetime_or_timedelta_dtype,
is_decimal,
is_extension_array_dtype,
is_float_dtype,
is_integer_dtype,
is_number,
is_numeric_dtype,
is_scalar,
Expand All @@ -15,6 +18,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 @@ -118,6 +122,7 @@ def to_numeric(arg, errors="raise", downcast=None):
is_series = False
is_index = False
is_scalars = False
is_numeric_extension_dtype = False

if isinstance(arg, ABCSeries):
is_series = True
Expand All @@ -142,6 +147,10 @@ def to_numeric(arg, errors="raise", downcast=None):
else:
values = arg

if is_extension_array_dtype(arg) and isinstance(values, NumericArray):
Copy link
Contributor

Choose a reason for hiding this comment

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

this really should be part of the above logic. also mask needs to be defined for all cases (can be default to None)

Copy link
Member Author

Choose a reason for hiding this comment

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

this really should be part of the above logic.

The reason for having it down here is to handle EA dtype Series and array in a single place (and Index when #34159/#37869 go through)

also mask needs to be defined for all cases (can be default to None)

done

is_numeric_extension_dtype = True
mask, values = values._mask, values._data
Copy link
Contributor

Choose a reason for hiding this comment

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

you likely need to just take the masked values only for the following block of code

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


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

if is_numeric_extension_dtype:
Copy link
Contributor

Choose a reason for hiding this comment

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

L194 might need to handle the mask

Copy link
Member Author

Choose a reason for hiding this comment

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

Do you mean?

            float_32_ind = typecodes.index(float_32_char)

I feel like there's a testcase I haven't looked at if yes (the ones I have work as is)

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

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

values = FloatingArray(values, mask)

if is_series:
return arg._constructor(values, index=arg.index, name=arg.name)
elif is_index:
Expand Down
25 changes: 25 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,28 @@ 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, pd.NA], "Int64", "integer", "Int8"),
([450, 300], "Int64", "integer", "Int16"),
([1, 1], "Int64", "signed", "Int8"),
([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, pd.NA], "Int64", "unsigned", "UInt8"),
([450, -300], "Int64", "unsigned", "Int64"),
([-1, -1], "Int32", "unsigned", "Int32"),
([1, 1], "Float64", "float", "Float32"),
([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)