From 3cf97b7a38883fe9d4596b83d0422cc256f813e7 Mon Sep 17 00:00:00 2001 From: Martin Durant Date: Mon, 23 Oct 2023 12:47:46 -0400 Subject: [PATCH] Use dt units in empty() (tz path) (#893) * Use dt units in empty() * simplfiying suggestion * fix --- fastparquet/dataframe.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) 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