147 lines
4.8 KiB
Python
147 lines
4.8 KiB
Python
import numpy as np
|
|
import pandas as pd
|
|
from pytz import timezone as get_timezone
|
|
|
|
from .base import NumpyColumn
|
|
from ...util.compat import get_localzone_name_compat
|
|
|
|
|
|
class NumpyDateTimeColumnBase(NumpyColumn):
|
|
datetime_dtype = None
|
|
|
|
null_value = np.datetime64(0, 'Y')
|
|
|
|
def __init__(self, timezone=None, offset_naive=True, local_timezone=None,
|
|
**kwargs):
|
|
self.timezone = timezone
|
|
self.offset_naive = offset_naive
|
|
self.local_timezone = local_timezone
|
|
super(NumpyDateTimeColumnBase, self).__init__(**kwargs)
|
|
|
|
def apply_timezones_after_read(self, dt):
|
|
timezone = self.timezone if self.timezone else self.local_timezone
|
|
|
|
if self.offset_naive and timezone.zone != 'UTC':
|
|
ts = pd.to_datetime(dt, utc=True).tz_convert(timezone)
|
|
ts = ts.tz_localize(None)
|
|
return ts.to_numpy(self.datetime_dtype)
|
|
|
|
return dt
|
|
|
|
def apply_timezones_before_write(self, items):
|
|
if isinstance(items, pd.DatetimeIndex):
|
|
ts = items
|
|
else:
|
|
timezone = self.timezone if self.timezone else self.local_timezone
|
|
ts = pd.to_datetime(items)
|
|
if not getattr(ts.dtype, 'tz', None):
|
|
ts = ts.tz_localize(timezone)
|
|
|
|
ts = ts.tz_convert('UTC')
|
|
return ts.tz_localize(None).to_numpy(self.datetime_dtype)
|
|
|
|
def is_items_integer(self, items):
|
|
return (
|
|
isinstance(items, np.ndarray) and
|
|
np.issubdtype(items.dtype, np.integer)
|
|
)
|
|
|
|
|
|
class NumpyDateTimeColumn(NumpyDateTimeColumnBase):
|
|
dtype = np.dtype(np.uint32)
|
|
datetime_dtype = 'datetime64[s]'
|
|
|
|
def write_items(self, items, buf):
|
|
# write int 'as is'.
|
|
if self.is_items_integer(items):
|
|
super(NumpyDateTimeColumn, self).write_items(items, buf)
|
|
return
|
|
|
|
items = self.apply_timezones_before_write(items)
|
|
|
|
super(NumpyDateTimeColumn, self).write_items(items, buf)
|
|
|
|
def read_items(self, n_items, buf):
|
|
items = super(NumpyDateTimeColumn, self).read_items(n_items, buf)
|
|
return self.apply_timezones_after_read(items.astype('datetime64[s]'))
|
|
|
|
|
|
class NumpyDateTime64Column(NumpyDateTimeColumnBase):
|
|
dtype = np.dtype(np.int64)
|
|
datetime_dtype = 'datetime64[ns]'
|
|
|
|
max_scale = 9
|
|
|
|
def __init__(self, scale=0, **kwargs):
|
|
self.scale = scale
|
|
super(NumpyDateTime64Column, self).__init__(**kwargs)
|
|
|
|
def read_items(self, n_items, buf):
|
|
# Clickhouse: t seconds is represented as t * 10^scale.
|
|
# datetime64[ns]: t seconds is represented as t * 10^9.
|
|
# Since 0 <= scale <= 9, multiply by the integer 10^(9 - scale).
|
|
items = super(NumpyDateTime64Column, self).read_items(n_items, buf)
|
|
|
|
tmp = np.copy(items)
|
|
tmp *= 10 ** (9 - self.scale)
|
|
dt = tmp.view(dtype='datetime64[ns]')
|
|
|
|
return self.apply_timezones_after_read(dt)
|
|
|
|
def write_items(self, items, buf):
|
|
# write int 'as is'.
|
|
if self.is_items_integer(items):
|
|
super(NumpyDateTime64Column, self).write_items(items, buf)
|
|
return
|
|
|
|
scale = 10 ** self.scale
|
|
frac_scale = 10 ** (self.max_scale - self.scale)
|
|
|
|
items = self.apply_timezones_before_write(items)
|
|
|
|
seconds = items.astype('datetime64[s]')
|
|
microseconds = (items - seconds).astype(dtype='timedelta64[ns]') \
|
|
.astype(np.uint32) // frac_scale
|
|
|
|
items = seconds.astype(self.dtype) * scale + microseconds
|
|
|
|
super(NumpyDateTime64Column, self).write_items(items, buf)
|
|
|
|
|
|
def create_numpy_datetime_column(spec, column_options):
|
|
if spec.startswith('DateTime64'):
|
|
cls = NumpyDateTime64Column
|
|
spec = spec[11:-1]
|
|
params = spec.split(',', 1)
|
|
column_options['scale'] = int(params[0])
|
|
if len(params) > 1:
|
|
spec = params[1].strip() + ')'
|
|
else:
|
|
cls = NumpyDateTimeColumn
|
|
spec = spec[9:]
|
|
|
|
context = column_options['context']
|
|
|
|
tz_name = None
|
|
offset_naive = True
|
|
|
|
# As Numpy do not use local timezone for converting timestamp to
|
|
# datetime we need always detect local timezone for manual converting.
|
|
local_tz_name = get_localzone_name_compat()
|
|
|
|
# Use column's timezone if it's specified.
|
|
if spec and spec[-1] == ')':
|
|
tz_name = spec[1:-2]
|
|
offset_naive = False
|
|
else:
|
|
if not context.settings.get('use_client_time_zone', False):
|
|
remote_timezone = context.server_info.get_timezone()
|
|
if local_tz_name != remote_timezone:
|
|
tz_name = remote_timezone
|
|
|
|
timezone = get_timezone(tz_name) if tz_name else None
|
|
local_timezone = get_timezone(local_tz_name) if local_tz_name else None
|
|
|
|
return cls(timezone=timezone, offset_naive=offset_naive,
|
|
local_timezone=local_timezone, **column_options)
|