Repository: pandas-dev/pandas
Pull request: BUG: Series.where casting dt64 to int64
- [x] closes #37682
- [x] tests added / passed
- [x] passes `black pandas`
- [x] passes `git diff upstream/master -u -- "*.py" | flake8 --diff`
- [x] whatsnew entry

This sits on top of #38021

Problem statement:
BUG: Index.where vs Series.where mismatch
If in merge._maybe_add_join_keys we replace a call to Index.where with a call to Series.where, we get a test failure because of the following mismatch:

```
dr = pd.date_range("2001-01-01", periods=3)
lvals = pd.DatetimeIndex([dr[0], dr[1], pd.NaT])._data
rvals = pd.Categorical([dr[0], pd.NaT, dr[2]])

mask = np.array([False, False, True])

>>> pd.Series(lvals).where(~mask, rvals)._values
array([Timestamp('2001-01-01 00:00:00'), Timestamp('2001-01-02 00:00:00'),
       978480000000000000], dtype=object)

>>> pd.Index(lvals).where(~mask, rvals)._values
<DatetimeArray>
['2001-01-01 00:00:00', '2001-01-02 00:00:00', '2001-01-03 00:00:00']
Length: 3, dtype: datetime64[ns]
```

The Index version looks better behaved.  These should share an implementation.

Hints from issue discussion:
Not closed by #38021, will be closed by #38073.

Commit message: Merge branch 'master' of https://github.com/pandas-dev/pandas into bug-series-where

Patch to review:
diff --git a/doc/source/whatsnew/v1.3.0.rst b/doc/source/whatsnew/v1.3.0.rst
index 64b9a11b1980d..edfe0d10ee1b7 100644
--- a/doc/source/whatsnew/v1.3.0.rst
+++ b/doc/source/whatsnew/v1.3.0.rst
@@ -180,6 +180,7 @@ Datetimelike
 ^^^^^^^^^^^^
 - Bug in :class:`DataFrame` and :class:`Series` constructors sometimes dropping nanoseconds from :class:`Timestamp` (resp. :class:`Timedelta`) ``data``, with ``dtype=datetime64[ns]`` (resp. ``timedelta64[ns]``) (:issue:`38032`)
 - Bug in :meth:`DataFrame.first` and :meth:`Series.first` returning two months for offset one month when first day is last calendar day (:issue:`29623`)
+- Bug in :meth:`Series.where` incorrectly casting ``datetime64`` values to ``int64`` (:issue:`37682`)
 -
 
 Timedelta
diff --git a/pandas/core/arrays/numpy_.py b/pandas/core/arrays/numpy_.py
index 50d12703c3a30..ae131d8a51ba1 100644
--- a/pandas/core/arrays/numpy_.py
+++ b/pandas/core/arrays/numpy_.py
@@ -161,7 +161,8 @@ def __init__(self, values: Union[np.ndarray, "PandasArray"], copy: bool = False)
                 f"'values' must be a NumPy array, not {type(values).__name__}"
             )
 
-        if values.ndim != 1:
+        if values.ndim == 0:
+            # Technically we support 2, but do not advertise that fact.
             raise ValueError("PandasArray must be 1-dimensional.")
 
         if copy:
diff --git a/pandas/core/internals/blocks.py b/pandas/core/internals/blocks.py
index 8752224356f61..cf5aeb0152e5e 100644
--- a/pandas/core/internals/blocks.py
+++ b/pandas/core/internals/blocks.py
@@ -1389,6 +1389,25 @@ def shift(self, periods: int, axis: int = 0, fill_value=None):
 
         return [self.make_block(new_values)]
 
+    def _maybe_reshape_where_args(self, values, other, cond, axis):
+        transpose = self.ndim == 2
+
+        cond = _extract_bool_array(cond)
+
+        # If the default broadcasting would go in the wrong direction, then
+        # explicitly reshape other instead
+        if getattr(other, "ndim", 0) >= 1:
+            if values.ndim - 1 == other.ndim and axis == 1:
+                other = other.reshape(tuple(other.shape + (1,)))
+            elif transpose and values.ndim == self.ndim - 1:
+                # TODO(EA2D): not neceesssary with 2D EAs
+                cond = cond.T
+
+        if not hasattr(cond, "shape"):
+            raise ValueError("where must have a condition that is ndarray like")
+
+        return other, cond
+
     def where(
         self, other, cond, errors="raise", try_cast: bool = False, axis: int = 0
     ) -> List["Block"]:
@@ -1411,7 +1430,6 @@ def where(
         """
         import pandas.core.computation.expressions as expressions
 
-        cond = _extract_bool_array(cond)
         assert not isinstance(other, (ABCIndex, ABCSeries, ABCDataFrame))
 
         assert errors in ["raise", "ignore"]
@@ -1422,17 +1440,7 @@ def where(
         if transpose:
             values = values.T
 
-        # If the default broadcasting would go in the wrong direction, then
-        # explicitly reshape other instead
-        if getattr(other, "ndim", 0) >= 1:
-            if values.ndim - 1 == other.ndim and axis == 1:
-                other = other.reshape(tuple(other.shape + (1,)))
-            elif transpose and values.ndim == self.ndim - 1:
-                # TODO(EA2D): not neceesssary with 2D EAs
-                cond = cond.T
-
-        if not hasattr(cond, "shape"):
-            raise ValueError("where must have a condition that is ndarray like")
+        other, cond = self._maybe_reshape_where_args(values, other, cond, axis)
 
         if cond.ravel("K").all():
             result = values
@@ -2189,6 +2197,26 @@ def to_native_types(self, na_rep="NaT", **kwargs):
         result = arr._format_native_types(na_rep=na_rep, **kwargs)
         return self.make_block(result)
 
+    def where(
+        self, other, cond, errors="raise", try_cast: bool = False, axis: int = 0
+    ) -> List["Block"]:
+        # TODO(EA2D): reshape unnecessary with 2D EAs
+        arr = self.array_values().reshape(self.shape)
+
+        other, cond = self._maybe_reshape_where_args(arr, other, cond, axis)
+
+        try:
+            res_values = arr.T.where(cond, other).T
+        except (ValueError, TypeError):
+            return super().where(
+                other, cond, errors=errors, try_cast=try_cast, axis=axis
+            )
+
+        # TODO(EA2D): reshape not needed with 2D EAs
+        res_values = res_values.reshape(self.values.shape)
+        nb = self.make_block_same_class(res_values)
+        return [nb]
+
     def _can_hold_element(self, element: Any) -> bool:
         arr = self.array_values()
 
@@ -2257,6 +2285,7 @@ class DatetimeTZBlock(ExtensionBlock, DatetimeBlock):
     fillna = DatetimeBlock.fillna  # i.e. Block.fillna
     fill_value = DatetimeBlock.fill_value
     _can_hold_na = DatetimeBlock._can_hold_na
+    where = DatetimeBlock.where
 
     array_values = ExtensionBlock.array_values
 
diff --git a/pandas/tests/arrays/test_array.py b/pandas/tests/arrays/test_array.py
index 779cb7a2350ee..f14d5349dcea3 100644
--- a/pandas/tests/arrays/test_array.py
+++ b/pandas/tests/arrays/test_array.py
@@ -278,7 +278,7 @@ def test_array_inference_fails(data):
     tm.assert_extension_array_equal(result, expected)
 
 
-@pytest.mark.parametrize("data", [np.array([[1, 2], [3, 4]]), [[1, 2], [3, 4]]])
+@pytest.mark.parametrize("data", [np.array(0)])
 def test_nd_raises(data):
     with pytest.raises(ValueError, match="PandasArray must be 1-dimensional"):
         pd.array(data, dtype="int64")
diff --git a/pandas/tests/series/indexing/test_where.py b/pandas/tests/series/indexing/test_where.py
index 27bbb47e1d0d1..e30daf43df187 100644
--- a/pandas/tests/series/indexing/test_where.py
+++ b/pandas/tests/series/indexing/test_where.py
@@ -464,3 +464,34 @@ def test_where_categorical(klass):
     df = klass(["A", "A", "B", "B", "C"], dtype="category")
     res = df.where(df != "C")
     tm.assert_equal(exp, res)
+
+
+@pytest.mark.parametrize("tz", [None, "US/Pacific"])
+def test_where_datetimelike_categorical(tz):
+    # GH#37682
+    dr = pd.date_range("2001-01-01", periods=3, tz=tz)._with_freq(None)
+    lvals = pd.DatetimeIndex([dr[0], dr[1], pd.NaT])
+    rvals = pd.Categorical([dr[0], pd.NaT, dr[2]])
+
+    mask = np.array([True, True, False])
+
+    # DatetimeIndex.where
+    res = lvals.where(mask, rvals)
+    tm.assert_index_equal(res, dr)
+
+    # DatetimeArray.where
+    res = lvals._data.where(mask, rvals)
+    tm.assert_datetime_array_equal(res, dr._data)
+
+    # Series.where
+    res = Series(lvals).where(mask, rvals)
+    tm.assert_series_equal(res, Series(dr))
+
+    # DataFrame.where
+    if tz is None:
+        res = pd.DataFrame(lvals).where(mask[:, None], pd.DataFrame(rvals))
+    else:
+        with pytest.xfail(reason="frame._values loses tz"):
+            res = pd.DataFrame(lvals).where(mask[:, None], pd.DataFrame(rvals))
+
+    tm.assert_frame_equal(res, pd.DataFrame(dr))


Code review comments:
diff_hunk: @@ -464,3 +464,34 @@ def test_where_categorical(klass):
     df = klass(["A", "A", "B", "B", "C"], dtype="category")
     res = df.where(df != "C")
     tm.assert_equal(exp, res)
+
+
+@pytest.mark.parametrize("tz", [None, "US/Pacific"])
+def test_where_datetimelike_categorical(tz):
path: pandas/tests/series/indexing/test_where.py
text: @user1:
can you use index_or_series fixture?

@author:
thatd be awkward since this also does the DatetimeArray and DataFrame cases