Repository: pandas-dev/pandas

Issue:
BUG: fillna() for tz-aware timestamps causes an exception
Suppose I have a DataFrame of timestamps:

```python
df = pd.DataFrame({'A': [pd.Timestamp('2012-11-11 00:00:00'),
                         pd.Timestamp('2012-11-11 00:00:00')]})
```

I can attempt to fill any nulls:

```python
In [23]: df.fillna(method='pad')
Out[23]:
           A
0 2012-11-11
1 2012-11-11
```

Now suppose my timestsamps are tz-aware:

```python
df = pd.DataFrame({'A': [pd.Timestamp('2012-11-11 00:00:00+00:00'),
                         pd.Timestamp('2012-11-11 00:00:00+00:00')]})
```

This causes an error:

```python
In [25]: df.fillna(method='pad')
...
ValueError: total size of new array must be unchanged
```

This does not appear to be fixed by #14960.



Discussion:
hmm, does look buggy. These should be coerced to/from i8, filled, then re-localized.
Also, this works as expected:

```
pd.DataFrame({'test': [pd.Timestamp('2012-01-01 13:00:00'),
                       pd.Timestamp('2012-01-01 13:00:00')]}) == -1
```

However, this raises an exception:

```
pd.DataFrame({'test': [pd.Timestamp('2012-01-01 13:00:00+00:00'),
                       pd.Timestamp('2012-01-01 13:00:00+00:00')]}) == -1
```

----

Furthermore, this works as expected:

```
pd.DataFrame({'time': [pd.Timestamp('2012-01-01 13:00:00+00:00')],
              'A': [3]}).groupby('A', as_index=False).head(1)
```

However, this loses the timezone:

```
pd.DataFrame({'time': [pd.Timestamp('2012-01-01 13:00:00+00:00')],
              'A': [3]}).groupby('A', as_index=False).first()
```

```
pd.DataFrame({'test': [pd.Timestamp('2012-01-01 13:00:00+00:00'),
                       pd.Timestamp('2012-01-01 13:00:00+00:00')]}) == -1
```

is the same exact bug (the issue is it can't reshape properly), it should have a ``try/except`` around that ``.reshape``. same issue as in #15869 
can you make a new issue for: 

```
pd.DataFrame({'time': [pd.Timestamp('2012-01-01 13:00:00+00:00')],
              'A': [3]}).groupby('A', as_index=False).first()
```

I *thought* we had an issue for this. But can't seem to find it.
I have split-out that issue as #15884.
This is something that fixes the `fillna()` issue:

```
diff --git a/pandas/core/internals.py b/pandas/core/internals.py
index 8db801f..5736188 100644
--- a/pandas/core/internals.py
+++ b/pandas/core/internals.py
@@ -2475,7 +2475,7 @@ class DatetimeTZBlock(NonConsolidatableMixIn, DatetimeBlock):
         if isinstance(result, np.ndarray):
             # allow passing of > 1dim if its trivial
             if result.ndim > 1:
-                result = result.reshape(len(result))
+                result = result.reshape(np.prod(result.shape))
             result = self.values._shallow_copy(result)
 
         return result
```

I.e., use the total size of the result instead of just the length of one dimension. Here's an example of it at work:

```
In [12]: df = pd.DataFrame({'A': [pd.Timestamp('2012-11-11 00:00:00+01:00'),
    ...:                          pd.NaT]})
    ...: df.fillna(method='pad')
    ...: 
Out[12]: 
                          A
0 2012-11-11 00:00:00+01:00
1 2012-11-11 00:00:00+01:00
```

Does this make sense as a fix, or is there more to it than that?
in other cases like this, we have a ``try/except`` around reshape (in internals). but this is ok too. does it pass everything?

Resolving patch:
diff --git a/pandas/core/internals.py b/pandas/core/internals.py
--- a/pandas/core/internals.py
+++ b/pandas/core/internals.py
@@ -2475,7 +2475,7 @@ def _try_coerce_result(self, result):
         if isinstance(result, np.ndarray):
             # allow passing of > 1dim if its trivial
             if result.ndim > 1:
-                result = result.reshape(len(result))
+                result = result.reshape(np.prod(result.shape))
             result = self.values._shallow_copy(result)
 
         return result


Test patch:
diff --git a/pandas/tests/frame/test_missing.py b/pandas/tests/frame/test_missing.py
--- a/pandas/tests/frame/test_missing.py
+++ b/pandas/tests/frame/test_missing.py
@@ -257,6 +257,20 @@ def test_fillna(self):
         result = df.fillna(value={'Date': df['Date2']})
         assert_frame_equal(result, expected)
 
+        # with timezone
+        # GH 15855
+        df = pd.DataFrame({'A': [pd.Timestamp('2012-11-11 00:00:00+01:00'),
+                                 pd.NaT]})
+        exp = pd.DataFrame({'A': [pd.Timestamp('2012-11-11 00:00:00+01:00'),
+                                  pd.Timestamp('2012-11-11 00:00:00+01:00')]})
+        assert_frame_equal(df.fillna(method='pad'), exp)
+
+        df = pd.DataFrame({'A': [pd.NaT,
+                                 pd.Timestamp('2012-11-11 00:00:00+01:00')]})
+        exp = pd.DataFrame({'A': [pd.Timestamp('2012-11-11 00:00:00+01:00'),
+                                  pd.Timestamp('2012-11-11 00:00:00+01:00')]})
+        assert_frame_equal(df.fillna(method='bfill'), exp)
+
     def test_fillna_downcast(self):
         # GH 15277
         # infer int64 from float64
diff --git a/pandas/tests/series/test_missing.py b/pandas/tests/series/test_missing.py
--- a/pandas/tests/series/test_missing.py
+++ b/pandas/tests/series/test_missing.py
@@ -258,6 +258,18 @@ def test_datetime64_tz_fillna(self):
             self.assert_series_equal(expected, result)
             self.assert_series_equal(pd.isnull(s), null_loc)
 
+        # with timezone
+        # GH 15855
+        df = pd.Series([pd.Timestamp('2012-11-11 00:00:00+01:00'), pd.NaT])
+        exp = pd.Series([pd.Timestamp('2012-11-11 00:00:00+01:00'),
+                         pd.Timestamp('2012-11-11 00:00:00+01:00')])
+        assert_series_equal(df.fillna(method='pad'), exp)
+
+        df = pd.Series([pd.NaT, pd.Timestamp('2012-11-11 00:00:00+01:00')])
+        exp = pd.Series([pd.Timestamp('2012-11-11 00:00:00+01:00'),
+                         pd.Timestamp('2012-11-11 00:00:00+01:00')])
+        assert_series_equal(df.fillna(method='bfill'), exp)
+
     def test_datetime64tz_fillna_round_issue(self):
         # GH 14872