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

REF: implement find_result_type #45304

Merged
merged 5 commits into from
Jan 12, 2022
Merged
Show file tree
Hide file tree
Changes from all 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
52 changes: 52 additions & 0 deletions pandas/core/dtypes/cast.py
Original file line number Diff line number Diff line change
Expand Up @@ -1418,6 +1418,58 @@ def ensure_nanosecond_dtype(dtype: DtypeObj) -> DtypeObj:
return dtype


# TODO: other value-dependent functions to standardize here include
# dtypes.concat.cast_to_common_type and Index._find_common_type_compat
def find_result_type(left: ArrayLike, right: Any) -> DtypeObj:
"""
Find the type/dtype for a the result of an operation between these objects.

This is similar to find_common_type, but looks at the objects instead
of just their dtypes. This can be useful in particular when one of the
objects does not have a `dtype`.

Parameters
----------
left : np.ndarray or ExtensionArray
right : Any

Returns
-------
np.dtype or ExtensionDtype

See also
--------
find_common_type
numpy.result_type
"""
new_dtype: DtypeObj

if left.dtype.kind in ["i", "u", "c"] and (
lib.is_integer(right) or lib.is_float(right)
):
# e.g. with int8 dtype and right=512, we want to end up with
# np.int16, whereas infer_dtype_from(512) gives np.int64,
# which will make us upcast too far.
if lib.is_float(right) and right.is_integer() and left.dtype.kind != "f":
right = int(right)

# Argument 1 to "result_type" has incompatible type "Union[ExtensionArray,
# ndarray[Any, Any]]"; expected "Union[Union[_SupportsArray[dtype[Any]],
# _NestedSequence[_SupportsArray[dtype[Any]]], bool, int, float, complex,
# str, bytes, _NestedSequence[Union[bool, int, float, complex, str, bytes]]],
# Union[dtype[Any], None, Type[Any], _SupportsDType[dtype[Any]], str,
# Union[Tuple[Any, int], Tuple[Any, Union[SupportsIndex,
# Sequence[SupportsIndex]]], List[Any], _DTypeDict, Tuple[Any, Any]]]]"
new_dtype = np.result_type(left, right) # type:ignore[arg-type]

else:
dtype, _ = infer_dtype_from(right, pandas_dtype=True)

new_dtype = find_common_type([left.dtype, dtype])

return new_dtype


@overload
def find_common_type(types: list[np.dtype]) -> np.dtype:
...
Expand Down
7 changes: 2 additions & 5 deletions pandas/core/internals/blocks.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@
from pandas.core.dtypes.astype import astype_array_safe
from pandas.core.dtypes.cast import (
can_hold_element,
find_common_type,
find_result_type,
infer_dtype_from,
maybe_downcast_numeric,
maybe_downcast_to_dtype,
Expand Down Expand Up @@ -1031,10 +1031,7 @@ def coerce_to_target_dtype(self, other) -> Block:
we can also safely try to coerce to the same dtype
and will receive the same block
"""
# if we cannot then coerce to object
dtype, _ = infer_dtype_from(other, pandas_dtype=True)

new_dtype = find_common_type([self.dtype, dtype])
new_dtype = find_result_type(self.values, other)

return self.astype(new_dtype, copy=False)

Expand Down
14 changes: 8 additions & 6 deletions pandas/tests/frame/indexing/test_where.py
Original file line number Diff line number Diff line change
Expand Up @@ -753,13 +753,15 @@ def test_where_try_cast_deprecated(frame_or_series):
obj.where(mask, -1, try_cast=False)


@pytest.mark.xfail(
reason="After fixing a bug in can_hold_element, we don't go through "
"the deprecated path, and also up-cast to int64 instead of int32 "
"(for now)."
)
def test_where_int_downcasting_deprecated(using_array_manager):
def test_where_int_downcasting_deprecated(using_array_manager, request):
# GH#44597
if not using_array_manager:
mark = pytest.mark.xfail(
reason="After fixing a bug in can_hold_element, we don't go through "
"the deprecated path, and also up-cast both columns to int32 "
"instead of just 1."
)
request.node.add_marker(mark)
arr = np.arange(6).astype(np.int16).reshape(3, 2)
df = DataFrame(arr)

Expand Down
8 changes: 4 additions & 4 deletions pandas/tests/indexing/test_loc.py
Original file line number Diff line number Diff line change
Expand Up @@ -2734,14 +2734,14 @@ def test_loc_getitem_nullable_index_with_duplicates():
tm.assert_series_equal(res, expected)


def test_loc_setitem_uint8_upcast():
@pytest.mark.parametrize("value", [300, np.uint16(300), np.int16(300)])
def test_loc_setitem_uint8_upcast(value):
# GH#26049

df = DataFrame([1, 2, 3, 4], columns=["col1"], dtype="uint8")
df.loc[2, "col1"] = 300 # value that can't be held in uint8
df.loc[2, "col1"] = value # value that can't be held in uint8

# TODO: would be better to get uint16?
expected = DataFrame([1, 2, 300, 4], columns=["col1"], dtype="int64")
expected = DataFrame([1, 2, 300, 4], columns=["col1"], dtype="uint16")
tm.assert_frame_equal(df, expected)


Expand Down
33 changes: 0 additions & 33 deletions pandas/tests/series/indexing/test_setitem.py
Original file line number Diff line number Diff line change
Expand Up @@ -1077,27 +1077,6 @@ def key(self):
def expected(self):
return Series([1, 512, 3], dtype=np.int16)

def test_int_key(self, obj, key, expected, val, indexer_sli, is_inplace, request):
if not isinstance(val, np.int16):
# with python int we end up with int64
mark = pytest.mark.xfail
request.node.add_marker(mark)
super().test_int_key(obj, key, expected, val, indexer_sli, is_inplace)

def test_mask_key(self, obj, key, expected, val, indexer_sli, request):
if not isinstance(val, np.int16):
# with python int we end up with int64
mark = pytest.mark.xfail
request.node.add_marker(mark)
super().test_mask_key(obj, key, expected, val, indexer_sli)

def test_series_where(self, obj, key, expected, val, is_inplace, request):
if not isinstance(val, np.int16):
# with python int we end up with int64
mark = pytest.mark.xfail
request.node.add_marker(mark)
super().test_series_where(obj, key, expected, val, is_inplace)


@pytest.mark.parametrize("val", [2 ** 33 + 1.0, 2 ** 33 + 1.1, 2 ** 62])
class TestSmallIntegerSetitemUpcast(SetitemCastingEquivalents):
Expand All @@ -1118,18 +1097,6 @@ def expected(self, val):
dtype = "i8"
return Series([val, 2, 3], dtype=dtype)

def test_int_key(self, obj, key, expected, val, indexer_sli, is_inplace, request):
if val % 1 == 0 and isinstance(val, float):
mark = pytest.mark.xfail
request.node.add_marker(mark)
super().test_int_key(obj, key, expected, val, indexer_sli, is_inplace)

def test_mask_key(self, obj, key, expected, val, indexer_sli, request):
if val % 1 == 0 and isinstance(val, float):
mark = pytest.mark.xfail
request.node.add_marker(mark)
super().test_mask_key(obj, key, expected, val, indexer_sli)


class CoercionTest(SetitemCastingEquivalents):
# Tests ported from tests.indexing.test_coercion
Expand Down