BUG: Groupby.agg with reduction function with tz aware data

doc/source/whatsnew/v0.25.0.rst

@@ -172,8 +172,8 @@ Plotting
 Groupby/Resample/Rolling
 ^^^^^^^^^^^^^^^^^^^^^^^^
 
--
--
+- Bug in :func:`pandas.core.groupby.GroupBy.agg` when applying a aggregation function to timezone aware data (:issue:`23683`)
+- Bug in :func:`pandas.core.groupby.GroupBy.first` and :func:`pandas.core.groupby.GroupBy.last` where timezone information would be dropped (:issue:`21603`)
 -
 
 

pandas/_libs/reduction.pyx

@@ -342,7 +342,9 @@ cdef class SeriesGrouper:
  index = None
  else:
  values = dummy.values
- if dummy.dtype != self.arr.dtype:
+ # GH 23683: datetimetz types are equivalent to datetime types here
+ if (dummy.dtype != self.arr.dtype
+ and values.dtype != self.arr.dtype):
  raise ValueError('Dummy array must be same dtype')
  if not values.flags.contiguous:
  values = values.copy()

pandas/core/groupby/groupby.py

@@ -26,7 +26,9 @@ class providing the base-class of operations.
 
 from pandas.core.dtypes.cast import maybe_downcast_to_dtype
 from pandas.core.dtypes.common import (
- ensure_float, is_extension_array_dtype, is_numeric_dtype, is_scalar)
+ ensure_float, is_datetime64tz_dtype, is_extension_array_dtype,
+ is_numeric_dtype, is_scalar)
+from pandas.core.dtypes.dtypes import DatetimeTZDtype
 from pandas.core.dtypes.missing import isna, notna
 
 import pandas.core.algorithms as algorithms
@@ -768,7 +770,17 @@ def _try_cast(self, result, obj, numeric_only=False):
  # The function can return something of any type, so check
  # if the type is compatible with the calling EA.
  try:
+ tz_dtype = None
+ if is_datetime64tz_dtype(dtype):
+ # GH 23683
+ # Prior results were generated in UTC. Ensure
+ # We localize to UTC first before converting to
+ # the target timezone
+ tz_dtype = dtype
+ dtype = DatetimeTZDtype(tz='UTC')
  result = obj._values._from_sequence(result, dtype=dtype)
+ if tz_dtype is not None:
+ result = result.astype(tz_dtype)
  except Exception:
  # https://github.com/pandas-dev/pandas/issues/22850
  # pandas has no control over what 3rd-party ExtensionArrays

pandas/tests/groupby/aggregate/test_other.py

@@ -512,3 +512,18 @@ def test_agg_list_like_func():
  expected = pd.DataFrame({'A': [str(x) for x in range(3)],
  'B': [[str(x)] for x in range(3)]})
  tm.assert_frame_equal(result, expected)
+
+
+def test_agg_lambda_with_timezone():
+ # GH 23683
+ df = pd.DataFrame({
+ 'tag': [1, 1],
+ 'date': [
+ pd.Timestamp('2018-01-01', tz='UTC'),
+ pd.Timestamp('2018-01-02', tz='UTC')]
+ })
+ result = df.groupby('tag').agg({'date': lambda e: e.head(1)})
+ expected = pd.DataFrame([pd.Timestamp('2018-01-01', tz='UTC')],
+ index=pd.Index([1], name='tag'),
+ columns=['date'])
+ tm.assert_frame_equal(result, expected)

pandas/tests/groupby/test_nth.py

@@ -278,6 +278,26 @@ def test_first_last_tz(data, expected_first, expected_last):
  assert_frame_equal(result, expected[['id', 'time']])
 
 
+@pytest.mark.parametrize('method, ts, alpha', [
+ ['first', Timestamp('2013-01-01', tz='US/Eastern'), 'a'],
+ ['last', Timestamp('2013-01-02', tz='US/Eastern'), 'b']
+])
+def test_first_last_tz_multi_column(method, ts, alpha):
+ # GH 21603
+ df = pd.DataFrame({'group': [1, 1, 2],
+ 'category_string': pd.Series(list('abc')).astype(
+ 'category'),
+ 'datetimetz': pd.date_range('20130101', periods=3,
+ tz='US/Eastern')})
+ result = getattr(df.groupby('group'), method)()
+ expepcted = pd.DataFrame({'category_string': [alpha, 'c'],
+ 'datetimetz': [ts,
+ Timestamp('2013-01-03',
+ tz='US/Eastern')]},
+ index=pd.Index([1, 2], name='group'))
+ assert_frame_equal(result, expepcted)
+
+
 def test_nth_multi_index_as_expected():
  # PR 9090, related to issue 8979
  # test nth on MultiIndex