109 lines
3.1 KiB
Python
109 lines
3.1 KiB
Python
from os import getenv
|
|
from datetime import date, timedelta
|
|
|
|
from .base import FormatColumn
|
|
|
|
|
|
epoch_start = date(1970, 1, 1)
|
|
epoch_end = date(2149, 6, 6)
|
|
|
|
epoch_start_date32 = date(1900, 1, 1)
|
|
epoch_end_date32 = date(2299, 12, 31)
|
|
|
|
|
|
class LazyLUT(dict):
|
|
def __init__(self, *args, _factory, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self._default_factory = _factory
|
|
|
|
def __missing__(self, key):
|
|
return self.setdefault(key, self._default_factory(key))
|
|
|
|
|
|
def make_date_lut_range(date_start, date_end):
|
|
return range(
|
|
(date_start - epoch_start).days,
|
|
(date_end - epoch_start).days + 1,
|
|
)
|
|
|
|
|
|
enable_lazy_date_lut = getenv('CLICKHOUSE_DRIVER_LASY_DATE_LUT', False)
|
|
if enable_lazy_date_lut:
|
|
try:
|
|
start, end = enable_lazy_date_lut.split(':')
|
|
start_date = date.fromisoformat(start)
|
|
end_date = date.fromisoformat(end)
|
|
|
|
date_range = make_date_lut_range(start_date, end_date)
|
|
except ValueError:
|
|
date_range = ()
|
|
|
|
# Since we initialize lazy lut with some initially warmed values,
|
|
# we use iterator and not dict comprehension for memory & time optimization
|
|
_date_lut = LazyLUT(
|
|
((x, epoch_start + timedelta(days=x)) for x in date_range),
|
|
_factory=lambda x: epoch_start + timedelta(days=x),
|
|
)
|
|
_date_lut_reverse = LazyLUT(
|
|
((value, key) for key, value in _date_lut.items()),
|
|
_factory=lambda x: (x - epoch_start).days,
|
|
)
|
|
else:
|
|
# If lazy lut is not enabled, we fallback to static dict initialization
|
|
# In both cases, we use same lut for both data types,
|
|
# since one encompasses the other and we can avoid duplicating overlap
|
|
date_range = make_date_lut_range(epoch_start_date32, epoch_end_date32)
|
|
_date_lut = {x: epoch_start + timedelta(days=x) for x in date_range}
|
|
_date_lut_reverse = {value: key for key, value in _date_lut.items()}
|
|
|
|
|
|
class DateColumn(FormatColumn):
|
|
ch_type = 'Date'
|
|
py_types = (date, )
|
|
format = 'H'
|
|
|
|
min_value = epoch_start
|
|
max_value = epoch_end
|
|
|
|
date_lut = _date_lut
|
|
date_lut_reverse = _date_lut_reverse
|
|
|
|
def before_write_items(self, items, nulls_map=None):
|
|
null_value = self.null_value
|
|
|
|
date_lut_reverse = self.date_lut_reverse
|
|
min_value = self.min_value
|
|
max_value = self.max_value
|
|
|
|
for i, item in enumerate(items):
|
|
if nulls_map and nulls_map[i]:
|
|
items[i] = null_value
|
|
continue
|
|
|
|
if item is not date:
|
|
item = date(item.year, item.month, item.day)
|
|
|
|
if min_value <= item <= max_value:
|
|
items[i] = date_lut_reverse[item]
|
|
else:
|
|
items[i] = 0
|
|
|
|
def after_read_items(self, items, nulls_map=None):
|
|
date_lut = self.date_lut
|
|
|
|
if nulls_map is None:
|
|
return tuple(date_lut[item] for item in items)
|
|
else:
|
|
return tuple(
|
|
(None if is_null else date_lut[items[i]])
|
|
for i, is_null in enumerate(nulls_map)
|
|
)
|
|
|
|
|
|
class Date32Column(DateColumn):
|
|
ch_type = 'Date32'
|
|
format = 'i'
|
|
|
|
min_value = epoch_start_date32
|
|
max_value = epoch_end_date32
|