diff --git a/fastparquet/dataframe.py b/fastparquet/dataframe.py index 690ea0ca..5c980481 100644 --- a/fastparquet/dataframe.py +++ b/fastparquet/dataframe.py @@ -5,7 +5,8 @@ from pandas import ( Categorical, DataFrame, Series, CategoricalIndex, RangeIndex, Index, MultiIndex, - DatetimeIndex, CategoricalDtype + DatetimeIndex, CategoricalDtype, + DatetimeTZDtype ) from pandas.core.arrays.masked import BaseMaskedDtype import warnings @@ -195,9 +196,10 @@ def set_cats(values, i=i, col=col, **kwargs): values = Categorical.from_codes(codes=code, dtype=bvalues.dtype) - elif getattr(bvalues.dtype, 'tz', None): - values = np.zeros(shape=shape, dtype='M8[ns]') - values = type(bvalues)(values, dtype=bvalues.dtype) + elif isinstance(bvalues.dtype, DatetimeTZDtype): + dt = "M8[ns]" if PANDAS_VERSION.major < 2 else f'M8[{bvalues.dtype.unit}]' + values = np.zeros(shape=shape, dtype=dt) + values = type(bvalues)._from_sequence(values, copy=False, dtype=bvalues.dtype) else: if not isinstance(bvalues, np.ndarray): # e.g. DatetimeLikeBlock backed by DatetimeArray/TimedeltaArray