Skip to content

Commit

Permalink
core: coerce result to dtype in Block.eval if dtype explicitly given
Browse files Browse the repository at this point in the history
  • Loading branch information
holymonson committed Jul 21, 2018
1 parent 322dbf4 commit ad96a28
Show file tree
Hide file tree
Showing 3 changed files with 24 additions and 7 deletions.
1 change: 1 addition & 0 deletions doc/source/whatsnew/v0.24.0.txt
Original file line number Diff line number Diff line change
Expand Up @@ -436,6 +436,7 @@ Datetimelike

- Fixed bug where two :class:`DateOffset` objects with different ``normalize`` attributes could evaluate as equal (:issue:`21404`)
- Fixed bug where :meth:`Timestamp.resolution` incorrectly returned 1-microsecond ``timedelta`` instead of 1-nanosecond :class:`Timedelta` (:issue:`21336`,:issue:`21365`)
- Fixed bug where :class:`DataFrame` with ``dtype='datetime64[ns]'`` operating with :class:`DateOffset` could cast to ``dtype='object'`` (:issue:`21610`)

Timedelta
^^^^^^^^^
Expand Down
25 changes: 18 additions & 7 deletions pandas/core/internals.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@
is_float_dtype, is_numeric_dtype,
is_numeric_v_string_like, is_extension_type,
is_extension_array_dtype,
is_list_like,
is_list_like, is_offsetlike,
is_re,
is_re_compilable,
is_scalar,
Expand Down Expand Up @@ -710,6 +710,10 @@ def _try_cast_result(self, result, dtype=None):
""" try to cast the result to our original type, we may have
roundtripped thru object in the mean-time
"""
# coerce to `dtype` if explicitly given
if isinstance(dtype, np.dtype):
return result.astype(dtype)

if dtype is None:
dtype = self.dtype

Expand Down Expand Up @@ -1315,7 +1319,8 @@ def shift(self, periods, axis=0, mgr=None):

return [self.make_block(new_values)]

def eval(self, func, other, errors='raise', try_cast=False, mgr=None):
def eval(self, func, other, errors='raise', try_cast=False, mgr=None,
dtype=None):
"""
evaluate the block; return result block from the result
Expand All @@ -1327,7 +1332,8 @@ def eval(self, func, other, errors='raise', try_cast=False, mgr=None):
- ``raise`` : allow exceptions to be raised
- ``ignore`` : suppress exceptions. On error return original object
try_cast : try casting the results to the input type
try_cast : try casting the results to ``dtype`` or the input type
dtype: if not ``None`` and ``try_cast=Ture``, coerce the result to it
Returns
-------
Expand Down Expand Up @@ -1362,10 +1368,15 @@ def eval(self, func, other, errors='raise', try_cast=False, mgr=None):
values, values_mask, other, other_mask = self._try_coerce_args(
transf(values), other)
except TypeError:
# `Timestamp` operate with `DateOffset` must cast to `object`,
# then cast back to `datetime64[ns]` after operation.
if (dtype is None and is_datetime64_dtype(self)
and is_offsetlike(other)):
dtype = self.dtype

block = self.coerce_to_target_dtype(orig_other)
return block.eval(func, orig_other,
errors=errors,
try_cast=try_cast, mgr=mgr)
return block.eval(func, orig_other, errors=errors,
try_cast=try_cast, mgr=mgr, dtype=dtype)

# get the result, may need to transpose the other
def get_result(other):
Expand Down Expand Up @@ -1449,7 +1460,7 @@ def handle_error():

# try to cast if requested
if try_cast:
result = self._try_cast_result(result)
result = self._try_cast_result(result, dtype=dtype)

result = _block_shape(result, ndim=self.ndim)
return [self.make_block(result)]
Expand Down
5 changes: 5 additions & 0 deletions pandas/tests/frame/test_arithmetic.py
Original file line number Diff line number Diff line change
Expand Up @@ -211,6 +211,11 @@ def test_df_sub_datetime64_not_ns(self):
pd.Timedelta(days=2)])
tm.assert_frame_equal(res, expected)

def test_timestamp_df_add_dateoffset(self):
expected = pd.DataFrame([pd.Timestamp('2019')])
result = pd.DataFrame([pd.Timestamp('2018')]) + pd.DateOffset(years=1)
tm.assert_frame_equal(result, expected)

@pytest.mark.parametrize('data', [
[1, 2, 3],
[1.1, 2.2, 3.3],
Expand Down

0 comments on commit ad96a28

Please sign in to comment.