clickhouse-driver 0.2.10__cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl

clickhouse_driver/columns/datetimecolumn.py

@@ -0,0 +1,203 @@
+from datetime import datetime
+
+from pytz import timezone as get_timezone, utc
+from ..util.compat import get_localzone_name_compat
+from .base import FormatColumn
+
+EPOCH = datetime(1970, 1, 1, tzinfo=utc)
+
+
+class DateTimeColumn(FormatColumn):
+    ch_type = 'DateTime'
+    py_types = (datetime, int)
+    format = 'I'
+
+    def __init__(self, timezone=None, offset_naive=True, **kwargs):
+        self.timezone = timezone
+        self.offset_naive = offset_naive
+        super(DateTimeColumn, self).__init__(**kwargs)
+
+    def after_read_items(self, items, nulls_map=None):
+        tz = self.timezone
+        fromts = datetime.fromtimestamp
+
+        # A bit ugly copy-paste. But it helps save time on items
+        # processing by avoiding lambda calls or if in loop.
+        if self.offset_naive:
+            if tz:
+                if nulls_map is None:
+                    return tuple(
+                        fromts(item, tz).replace(tzinfo=None)
+                        for item in items
+                    )
+                else:
+                    return tuple(
+                        (None if is_null else
+                         fromts(items[i], tz).replace(tzinfo=None))
+                        for i, is_null in enumerate(nulls_map)
+                    )
+            else:
+                if nulls_map is None:
+                    return tuple(fromts(item) for item in items)
+                else:
+                    return tuple(
+                        (None if is_null else fromts(items[i]))
+                        for i, is_null in enumerate(nulls_map)
+                    )
+
+        else:
+            if nulls_map is None:
+                return tuple(fromts(item, tz) for item in items)
+            else:
+                return tuple(
+                    (None if is_null else fromts(items[i], tz))
+                    for i, is_null in enumerate(nulls_map)
+                )
+
+    def before_write_items(self, items, nulls_map=None):
+        timezone = self.timezone
+        null_value = self.null_value
+        to_timestamp = datetime.timestamp
+
+        for i, item in enumerate(items):
+            if nulls_map and nulls_map[i]:
+                items[i] = null_value
+                continue
+
+            if isinstance(item, int):
+                # support supplying raw integers to avoid
+                # costly timezone conversions when using datetime
+                continue
+
+            if timezone:
+                # Set server's timezone for offset-naive datetime.
+                if item.tzinfo is None:
+                    item = timezone.localize(item)
+
+                item = item.astimezone(utc)
+
+            else:
+                # If datetime is offset-aware use it's timezone.
+                if item.tzinfo is not None:
+                    item = item.astimezone(utc)
+
+            items[i] = int(to_timestamp(item))
+
+
+class DateTime64Column(DateTimeColumn):
+    ch_type = 'DateTime64'
+    format = 'q'
+
+    max_scale = 6
+
+    def __init__(self, scale=0, **kwargs):
+        self.scale = scale
+        super(DateTime64Column, self).__init__(**kwargs)
+
+    def after_read_items(self, items, nulls_map=None):
+        scale = float(10 ** self.scale)
+
+        tz = self.timezone
+        fromts = datetime.fromtimestamp
+
+        # A bit ugly copy-paste. But it helps save time on items
+        # processing by avoiding lambda calls or if in loop.
+        if self.offset_naive:
+            if tz:
+                if nulls_map is None:
+                    return tuple(
+                        fromts(item / scale, tz).replace(tzinfo=None)
+                        for item in items
+                    )
+                else:
+                    return tuple(
+                        (None if is_null else
+                         fromts(items[i] / scale, tz).replace(tzinfo=None))
+                        for i, is_null in enumerate(nulls_map)
+                    )
+            else:
+                if nulls_map is None:
+                    return tuple(fromts(item / scale) for item in items)
+                else:
+                    return tuple(
+                        (None if is_null else fromts(items[i] / scale))
+                        for i, is_null in enumerate(nulls_map)
+                    )
+
+        else:
+            if nulls_map is None:
+                return tuple(fromts(item / scale, tz) for item in items)
+            else:
+                return tuple(
+                    (None if is_null else fromts(items[i] / scale, tz))
+                    for i, is_null in enumerate(nulls_map)
+                )
+
+    def before_write_items(self, items, nulls_map=None):
+        scale = 10 ** self.scale
+        frac_scale = 10 ** (self.max_scale - self.scale)
+
+        timezone = self.timezone
+        null_value = self.null_value
+        to_timestamp = datetime.timestamp
+
+        for i, item in enumerate(items):
+            if nulls_map and nulls_map[i]:
+                items[i] = null_value
+                continue
+
+            if isinstance(item, int):
+                # support supplying raw integers to avoid
+                # costly timezone conversions when using datetime
+                continue
+
+            if timezone:
+                # Set server's timezone for offset-naive datetime.
+                if item.tzinfo is None:
+                    item = timezone.localize(item)
+
+                item = item.astimezone(utc)
+
+            else:
+                # If datetime is offset-aware use it's timezone.
+                if item.tzinfo is not None:
+                    item = item.astimezone(utc)
+
+            items[i] = (
+                int(to_timestamp(item)) * scale +
+                int(item.microsecond / frac_scale)
+            )
+
+
+def create_datetime_column(spec, column_options):
+    if spec.startswith('DateTime64'):
+        cls = DateTime64Column
+        spec = spec[11:-1]
+        params = spec.split(',', 1)
+        column_options['scale'] = int(params[0])
+        if len(params) > 1:
+            spec = params[1].strip() + ')'
+    else:
+        cls = DateTimeColumn
+        spec = spec[9:]
+
+    context = column_options['context']
+
+    tz_name = timezone = None
+    offset_naive = True
+
+    # 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):
+            local_timezone = get_localzone_name_compat()
+            remote_timezone = context.server_info.get_timezone()
+            if local_timezone != remote_timezone:
+                tz_name = remote_timezone
+
+    if tz_name:
+        timezone = get_timezone(tz_name)
+
+    return cls(timezone=timezone, offset_naive=offset_naive, **column_options)

clickhouse_driver/columns/decimalcolumn.py

@@ -0,0 +1,116 @@
+from decimal import Decimal, localcontext
+
+from .base import FormatColumn
+from .exceptions import ColumnTypeMismatchException
+from .intcolumn import Int128Column, Int256Column
+
+
+class DecimalColumn(FormatColumn):
+    py_types = (Decimal, float, int)
+    max_precision = None
+
+    def __init__(self, precision, scale, types_check=False, **kwargs):
+        self.precision = precision
+        self.scale = scale
+        super(DecimalColumn, self).__init__(**kwargs)
+
+        if types_check:
+            def check_item(value):
+                parts = str(value).split('.')
+                int_part = parts[0]
+
+                if len(int_part) > precision:
+                    raise ColumnTypeMismatchException(value)
+
+            self.check_item = check_item
+
+    def after_read_items(self, items, nulls_map=None):
+        if self.scale >= 1:
+            scale = 10 ** self.scale
+
+            if nulls_map is None:
+                return tuple(Decimal(item) / scale for item in items)
+            else:
+                return tuple(
+                    (None if is_null else Decimal(items[i]) / scale)
+                    for i, is_null in enumerate(nulls_map)
+                )
+        else:
+            if nulls_map is None:
+                return tuple(Decimal(item) for item in items)
+            else:
+                return tuple(
+                    (None if is_null else Decimal(items[i]))
+                    for i, is_null in enumerate(nulls_map)
+                )
+
+    def before_write_items(self, items, nulls_map=None):
+        null_value = self.null_value
+
+        if self.scale >= 1:
+            scale = 10 ** self.scale
+
+            for i, item in enumerate(items):
+                if nulls_map and nulls_map[i]:
+                    items[i] = null_value
+                else:
+                    items[i] = int(Decimal(str(item)) * scale)
+
+        else:
+            for i, item in enumerate(items):
+                if nulls_map and nulls_map[i]:
+                    items[i] = null_value
+                else:
+                    items[i] = int(Decimal(str(item)))
+
+    # Override default precision to the maximum supported by underlying type.
+    def _write_data(self, items, buf):
+        with localcontext() as ctx:
+            ctx.prec = self.max_precision
+            super(DecimalColumn, self)._write_data(items, buf)
+
+    def _read_data(self, n_items, buf, nulls_map=None):
+        with localcontext() as ctx:
+            ctx.prec = self.max_precision
+            return super(DecimalColumn, self)._read_data(
+                n_items, buf, nulls_map=nulls_map
+            )
+
+
+class Decimal32Column(DecimalColumn):
+    format = 'i'
+    max_precision = 9
+
+
+class Decimal64Column(DecimalColumn):
+    format = 'q'
+    max_precision = 18
+
+
+class Decimal128Column(DecimalColumn, Int128Column):
+    max_precision = 38
+
+
+class Decimal256Column(DecimalColumn, Int256Column):
+    max_precision = 76
+
+
+def create_decimal_column(spec, column_options):
+    precision, scale = spec[8:-1].split(',')
+    precision, scale = int(precision), int(scale)
+
+    # Maximum precisions for underlying types are:
+    # Int32    10**9
+    # Int64   10**18
+    # Int128  10**38
+    # Int256  10**76
+    if precision <= 9:
+        cls = Decimal32Column
+    elif precision <= 18:
+        cls = Decimal64Column
+    elif precision <= 38:
+        cls = Decimal128Column
+    else:
+        cls = Decimal256Column
+
+    return cls(precision, scale, **column_options)

clickhouse_driver/columns/enumcolumn.py

@@ -0,0 +1,129 @@
+from enum import Enum
+from collections import OrderedDict
+
+from .. import errors
+from .intcolumn import IntColumn
+
+invalid_names_for_python_enum = frozenset(['mro', ''])
+
+
+class EnumColumn(IntColumn):
+    py_types = (Enum, int, str)
+
+    def __init__(self, name_by_value, value_by_name, **kwargs):
+        self.name_by_value = name_by_value
+        self.value_by_name = value_by_name
+        super(EnumColumn, self).__init__(**kwargs)
+
+    def before_write_items(self, items, nulls_map=None):
+        null_value = self.null_value
+        name_by_value = self.name_by_value
+        value_by_name = self.value_by_name
+
+        for i, item in enumerate(items):
+            if nulls_map and nulls_map[i]:
+                items[i] = null_value
+                continue
+
+            source_value = item.name if isinstance(item, Enum) else item
+
+            # Check real enum value
+            try:
+                if isinstance(source_value, str):
+                    items[i] = value_by_name[source_value]
+                else:
+                    items[i] = value_by_name[name_by_value[source_value]]
+            except (ValueError, KeyError):
+                choices = ', '.join(
+                    "'{}' = {}".format(name.replace("'", r"\'"), value)
+                    for name, value in value_by_name.items()
+                )
+                enum_str = '{}({})'.format(self.ch_type, choices)
+
+                raise errors.LogicalError(
+                    "Unknown element '{}' for type {}"
+                    .format(source_value, enum_str)
+                )
+
+    def after_read_items(self, items, nulls_map=None):
+        name_by_value = self.name_by_value
+
+        if nulls_map is None:
+            return tuple(name_by_value[item] for item in items)
+        else:
+            return tuple(
+                (None if is_null else name_by_value[items[i]])
+                for i, is_null in enumerate(nulls_map)
+            )
+
+
+class Enum8Column(EnumColumn):
+    ch_type = 'Enum8'
+    format = 'b'
+    int_size = 1
+
+
+class Enum16Column(EnumColumn):
+    ch_type = 'Enum16'
+    format = 'h'
+    int_size = 2
+
+
+def create_enum_column(spec, column_options):
+    if spec.startswith('Enum8'):
+        params = spec[6:-1]
+        cls = Enum8Column
+    else:
+        params = spec[7:-1]
+        cls = Enum16Column
+
+    name_by_value, value_by_name = _parse_options(params)
+
+    return cls(name_by_value, value_by_name, **column_options)
+
+
+def _parse_options(option_string):
+    name_by_value, value_by_name = {}, OrderedDict()
+    after_name = False
+    escaped = False
+    quote_character = None
+    name = ''
+    value = ''
+
+    for ch in option_string:
+        if escaped:
+            name += ch
+            escaped = False  # accepting escaped character
+
+        elif after_name:
+            if ch in (' ', '='):
+                pass
+            elif ch == ',':
+                value = int(value)
+                name_by_value[value] = name
+                value_by_name[name] = value
+                after_name = False
+                name = ''
+                value = ''  # reset before collecting new option
+            else:
+                value += ch
+
+        elif quote_character:
+            if ch == '\\':
+                escaped = True
+            elif ch == quote_character:
+                quote_character = None
+                after_name = True  # start collecting option value
+            else:
+                name += ch
+
+        else:
+            if ch == "'":
+                quote_character = ch
+
+    if after_name:
+        value = int(value)
+        name_by_value[value] = name
+        value_by_name[name] = value
+
+    return name_by_value, value_by_name

clickhouse_driver/columns/exceptions.py

@@ -0,0 +1,12 @@
+
+
+class ColumnException(Exception):
+    pass
+
+
+class ColumnTypeMismatchException(ColumnException):
+    pass
+
+
+class StructPackException(ColumnException):
+    pass

clickhouse_driver/columns/floatcolumn.py

@@ -0,0 +1,34 @@
+from ctypes import c_float
+
+from .base import FormatColumn
+
+
+class FloatColumn(FormatColumn):
+    py_types = (float, int)
+
+
+class Float32Column(FloatColumn):
+    ch_type = 'Float32'
+    format = 'f'
+
+    def __init__(self, types_check=False, **kwargs):
+        super(Float32Column, self).__init__(types_check=types_check, **kwargs)
+
+        if types_check:
+            # Chop only bytes that fit current type.
+            # Cast to -nan or nan if overflows.
+            def before_write_items(items, nulls_map=None):
+                null_value = self.null_value
+
+                for i, item in enumerate(items):
+                    if nulls_map and nulls_map[i]:
+                        items[i] = null_value
+                    else:
+                        items[i] = c_float(item).value
+
+            self.before_write_items = before_write_items
+
+
+class Float64Column(FloatColumn):
+    ch_type = 'Float64'
+    format = 'd'

clickhouse_driver/columns/intcolumn.py

@@ -0,0 +1,157 @@
+
+from .exceptions import ColumnTypeMismatchException
+from .base import FormatColumn
+from .largeint import (
+    int128_from_quads, int128_to_quads, uint128_from_quads, uint128_to_quads,
+    int256_from_quads, int256_to_quads, uint256_from_quads, uint256_to_quads
+)
+
+
+class IntColumn(FormatColumn):
+    py_types = (int, )
+    int_size = None
+
+    def __init__(self, types_check=False, **kwargs):
+        super(IntColumn, self).__init__(types_check=types_check, **kwargs)
+
+        if types_check:
+            self.mask = (1 << 8 * self.int_size) - 1
+
+            # Chop only bytes that fit current type.
+            # ctypes.c_intXX is slower.
+            def before_write_items(items, nulls_map=None):
+                # TODO: cythonize
+                null_value = self.null_value
+
+                for i, item in enumerate(items):
+                    if nulls_map and nulls_map[i]:
+                        items[i] = null_value
+                        continue
+
+                    if item >= 0:
+                        sign = 1
+                    else:
+                        sign = -1
+                        item = -item
+
+                    items[i] = sign * (item & self.mask)
+
+            self.before_write_items = before_write_items
+
+
+class UIntColumn(IntColumn):
+    def __init__(self, types_check=False, **kwargs):
+        super(UIntColumn, self).__init__(types_check=types_check, **kwargs)
+
+        if types_check:
+            def check_item(value):
+                if value < 0:
+                    raise ColumnTypeMismatchException(value)
+
+            self.check_item = check_item
+
+
+class Int8Column(IntColumn):
+    ch_type = 'Int8'
+    format = 'b'
+    int_size = 1
+
+
+class Int16Column(IntColumn):
+    ch_type = 'Int16'
+    format = 'h'
+    int_size = 2
+
+
+class Int32Column(IntColumn):
+    ch_type = 'Int32'
+    format = 'i'
+    int_size = 4
+
+
+class Int64Column(IntColumn):
+    ch_type = 'Int64'
+    format = 'q'
+    int_size = 8
+
+
+class UInt8Column(UIntColumn):
+    ch_type = 'UInt8'
+    format = 'B'
+    int_size = 1
+
+
+class UInt16Column(UIntColumn):
+    ch_type = 'UInt16'
+    format = 'H'
+    int_size = 2
+
+
+class UInt32Column(UIntColumn):
+    ch_type = 'UInt32'
+    format = 'I'
+    int_size = 4
+
+
+class UInt64Column(UIntColumn):
+    ch_type = 'UInt64'
+    format = 'Q'
+    int_size = 8
+
+
+class LargeIntColumn(IntColumn):
+    format = 'Q'  # We manually deal with sign in read/write.
+    factor = None
+
+    to_quads = None
+    from_quads = None
+
+    def write_items(self, items, buf):
+        n_items = len(items)
+
+        s = self.make_struct(self.factor * n_items)
+        uint_64_pairs = self.to_quads(items, n_items)
+
+        buf.write(s.pack(*uint_64_pairs))
+
+    def read_items(self, n_items, buf):
+        s = self.make_struct(self.factor * n_items)
+        items = s.unpack(buf.read(s.size))
+
+        return self.from_quads(items, n_items)
+
+
+class Int128Column(LargeIntColumn):
+    ch_type = 'Int128'
+    int_size = 16
+    factor = 2
+
+    to_quads = int128_to_quads
+    from_quads = int128_from_quads
+
+
+class UInt128Column(LargeIntColumn):
+    ch_type = 'UInt128'
+    int_size = 16
+    factor = 2
+
+    to_quads = uint128_to_quads
+    from_quads = uint128_from_quads
+
+
+class Int256Column(LargeIntColumn):
+    ch_type = 'Int256'
+    int_size = 32
+    factor = 4
+
+    to_quads = int256_to_quads
+    from_quads = int256_from_quads
+
+
+class UInt256Column(LargeIntColumn):
+    ch_type = 'UInt256'
+    int_size = 32
+    factor = 4
+
+    to_quads = uint256_to_quads
+    from_quads = uint256_from_quads

clickhouse_driver/columns/intervalcolumn.py

@@ -0,0 +1,33 @@
+from .intcolumn import Int64Column
+
+
+class IntervalColumn(Int64Column):
+    pass
+
+
+class IntervalDayColumn(IntervalColumn):
+    ch_type = 'IntervalDay'
+
+
+class IntervalWeekColumn(IntervalColumn):
+    ch_type = 'IntervalWeek'
+
+
+class IntervalMonthColumn(IntervalColumn):
+    ch_type = 'IntervalMonth'
+
+
+class IntervalYearColumn(IntervalColumn):
+    ch_type = 'IntervalYear'
+
+
+class IntervalHourColumn(IntervalColumn):
+    ch_type = 'IntervalHour'
+
+
+class IntervalMinuteColumn(IntervalColumn):
+    ch_type = 'IntervalMinute'
+
+
+class IntervalSecondColumn(IntervalColumn):
+    ch_type = 'IntervalSecond'