ENH: Implement export of datetime64[ns, tz] dtypes with a fixed HDF5 store #11411

Author: jreback

doc/source/whatsnew/v0.17.1.txt

@@ -66,6 +66,7 @@ Enhancements
      pd.Index([1, np.nan, 3]).fillna(2)
 
 - ``pivot_table`` now has a ``margins_name`` argument so you can use something other than the default of 'All' (:issue:`3335`)
+- Implement export of ``datetime64[ns, tz]`` dtypes with a fixed HDF5 store (:issue:`11411`)
 
 .. _whatsnew_0171.api:
 
@@ -159,4 +160,3 @@ Bug Fixes
 - Bug in ``DataFrame.join()`` with ``how='right'`` producing a ``TypeError`` (:issue:`11519`)
 - Bug in ``Series.quantile`` with empty list results has ``Index`` with ``object`` dtype (:issue:`11588`)
 - Bug in ``pd.merge`` results in empty ``Int64Index`` rather than ``Index(dtype=object)`` when the merge result is empty (:issue:`11588`)
-

pandas/io/pytables.py

@@ -47,6 +47,7 @@
 # versioning attribute
 _version = '0.15.2'
 
+### encoding ###
 # PY3 encoding if we don't specify
 _default_encoding = 'UTF-8'
 
@@ -64,22 +65,8 @@ def _ensure_encoding(encoding):
             encoding = _default_encoding
     return encoding
 
-def _set_tz(values, tz, preserve_UTC=False):
-    """ set the timezone if values are an Index """
-    if tz is not None and isinstance(values, Index):
-        tz = _ensure_decoded(tz)
-        if values.tz is None:
-            values = values.tz_localize('UTC').tz_convert(tz)
-        if preserve_UTC:
-            if tslib.get_timezone(tz) == 'UTC':
-                values = list(values)
-
-    return values
-
-
 Term = Expr
 
-
 def _ensure_term(where, scope_level):
     """
     ensure that the where is a Term or a list of Term
@@ -1947,14 +1934,11 @@ def set_atom_datetime64tz(self, block, info, values=None):
         if values is None:
             values = block.values
 
-        # convert this column to datetime64[ns] utc, and save the tz
-        values = values.tz_convert('UTC').values.view('i8').reshape(block.shape)
+        # convert this column to i8 in UTC, and save the tz
+        values = values.asi8.reshape(block.shape)
 
         # store a converted timezone
-        zone = tslib.get_timezone(block.values.tz)
-        if zone is None:
-            zone = tslib.tot_seconds(block.values.tz.utcoffset())
-        self.tz = zone
+        self.tz = _get_tz(block.values.tz)
         self.update_info(info)
 
         self.kind = 'datetime64'
@@ -2015,18 +1999,9 @@ def convert(self, values, nan_rep, encoding):
 
             # reverse converts
             if dtype == u('datetime64'):
-                # recreate the timezone
-                if self.tz is not None:
-
-                    # data should be 2-dim here
-                    # we stored as utc, so just set the tz
 
-                    index = DatetimeIndex(
-                        self.data.ravel(), tz='UTC').tz_convert(tslib.maybe_get_tz(self.tz))
-                    self.data = index
-
-                else:
-                    self.data = np.asarray(self.data, dtype='M8[ns]')
+                # recreate with tz if indicated
+                self.data = _set_tz(self.data, self.tz, coerce=True)
 
             elif dtype == u('timedelta64'):
                 self.data = np.asarray(self.data, dtype='m8[ns]')
@@ -2347,7 +2322,10 @@ def read_array(self, key):
                 ret = data
 
             if dtype == u('datetime64'):
-                ret = np.asarray(ret, dtype='M8[ns]')
+
+                # reconstruct a timezone if indicated
+                ret = _set_tz(ret, getattr(attrs, 'tz', None), coerce=True)
+
             elif dtype == u('timedelta64'):
                 ret = np.asarray(ret, dtype='m8[ns]')
 
@@ -2397,10 +2375,7 @@ def write_index(self, key, index):
                 node._v_attrs.freq = index.freq
 
             if hasattr(index, 'tz') and index.tz is not None:
-                zone = tslib.get_timezone(index.tz)
-                if zone is None:
-                    zone = tslib.tot_seconds(index.tz.utcoffset())
-                node._v_attrs.tz = zone
+                node._v_attrs.tz = _get_tz(index.tz)
 
     def write_block_index(self, key, index):
         self.write_array('%s_blocs' % key, index.blocs)
@@ -2574,11 +2549,20 @@ def write_array(self, key, value, items=None):
             if empty_array:
                 self.write_array_empty(key, value)
             else:
-                if value.dtype.type == np.datetime64:
+                if com.is_datetime64_dtype(value.dtype):
                     self._handle.create_array(self.group, key, value.view('i8'))
                     getattr(
                         self.group, key)._v_attrs.value_type = 'datetime64'
-                elif value.dtype.type == np.timedelta64:
+                elif com.is_datetime64tz_dtype(value.dtype):
+                    # store as UTC
+                    # with a zone
+                    self._handle.create_array(self.group, key,
+                                              value.asi8)
+
+                    node = getattr(self.group, key)
+                    node._v_attrs.tz = _get_tz(value.tz)
+                    node._v_attrs.value_type = 'datetime64'
+                elif com.is_timedelta64_dtype(value.dtype):
                     self._handle.create_array(self.group, key, value.view('i8'))
                     getattr(
                         self.group, key)._v_attrs.value_type = 'timedelta64'
@@ -4248,6 +4232,40 @@ def _get_info(info, name):
         idx = info[name] = dict()
     return idx
 
+### tz to/from coercion ###
+def _get_tz(tz):
+    """ for a tz-aware type, return an encoded zone """
+    zone = tslib.get_timezone(tz)
+    if zone is None:
+        zone = tslib.tot_seconds(tz.utcoffset())
+    return zone
+
+def _set_tz(values, tz, preserve_UTC=False, coerce=False):
+    """
+    coerce the values to a DatetimeIndex if tz is set
+    preserve the input shape if possible
+
+    Parameters
+    ----------
+    values : ndarray
+    tz : string/pickled tz object
+    preserve_UTC : boolean,
+        preserve the UTC of the result
+    coerce : if we do not have a passed timezone, coerce to M8[ns] ndarray
+    """
+    if tz is not None:
+        values = values.ravel()
+        tz = tslib.get_timezone(_ensure_decoded(tz))
+        values = DatetimeIndex(values)
+        if values.tz is None:
+            values = values.tz_localize('UTC').tz_convert(tz)
+        if preserve_UTC:
+            if tz == 'UTC':
+                values = list(values)
+    elif coerce:
+        values = np.asarray(values, dtype='M8[ns]')
+
+    return values
 
 def _convert_index(index, encoding=None, format_type=None):
     index_name = getattr(index, 'name', None)

pandas/io/tests/test_pytables.py

@@ -4909,15 +4909,26 @@ def test_tseries_select_index_column(self):
             result = store.select_column('frame', 'index')
             self.assertEqual(rng.tz, result.dt.tz)
 
-    def test_timezones(self):
-        rng = date_range('1/1/2000', '1/30/2000', tz='US/Eastern')
-        frame = DataFrame(np.random.randn(len(rng), 4), index=rng)
-
+    def test_timezones_fixed(self):
         with ensure_clean_store(self.path) as store:
-            store['frame'] = frame
-            recons = store['frame']
-            self.assertTrue(recons.index.equals(rng))
-            self.assertEqual(rng.tz, recons.index.tz)
+
+            # index
+            rng = date_range('1/1/2000', '1/30/2000', tz='US/Eastern')
+            df = DataFrame(np.random.randn(len(rng), 4), index=rng)
+            store['df'] = df
+            result = store['df']
+            assert_frame_equal(result, df)
+
+            # as data
+            # GH11411
+            _maybe_remove(store, 'df')
+            df = DataFrame({'A' : rng,
+                            'B' : rng.tz_convert('UTC').tz_localize(None),
+                            'C' : rng.tz_convert('CET'),
+                            'D' : range(len(rng))}, index=rng)
+            store['df'] = df
+            result = store['df']
+            assert_frame_equal(result, df)
 
     def test_fixed_offset_tz(self):
         rng = date_range('1/1/2000 00:00:00-07:00', '1/30/2000 00:00:00-07:00')