clickhouse-driver 0.2.8__cp312-cp312-macosx_11_0_arm64.whl

clickhouse_driver/columns/ipcolumn.py

@@ -0,0 +1,118 @@
+from ipaddress import IPv4Address, IPv6Address, AddressValueError
+
+from .. import errors
+from .exceptions import ColumnTypeMismatchException
+from .stringcolumn import ByteFixedString
+from .intcolumn import UInt32Column
+
+
+class IPv4Column(UInt32Column):
+    ch_type = "IPv4"
+    py_types = (str, IPv4Address, int)
+
+    def __init__(self, types_check=False, **kwargs):
+        # UIntColumn overrides before_write_item and check_item
+        # in its __init__ when types_check is True so we force
+        # __init__ without it then add the appropriate check method for IPv4
+        super(UInt32Column, self).__init__(types_check=False, **kwargs)
+
+        self.types_check_enabled = types_check
+        if types_check:
+
+            def check_item(value):
+                if isinstance(value, int) and value < 0:
+                    raise ColumnTypeMismatchException(value)
+
+                if not isinstance(value, IPv4Address):
+                    try:
+                        value = IPv4Address(value)
+                    except AddressValueError:
+                        # Cannot parse input in a valid IPv4
+                        raise ColumnTypeMismatchException(value)
+
+            self.check_item = check_item
+
+    def after_read_items(self, items, nulls_map=None):
+        if nulls_map is None:
+            return tuple(IPv4Address(item) for item in items)
+        else:
+            return tuple(
+                (None if is_null else IPv4Address(items[i]))
+                for i, is_null in enumerate(nulls_map)
+            )
+
+    def before_write_items(self, 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
+                continue
+
+            # allow Ipv4 in integer, string or IPv4Address object
+            try:
+                if isinstance(item, int):
+                    continue
+
+                if not isinstance(item, IPv4Address):
+                    item = IPv4Address(item)
+
+                items[i] = int(item)
+            except AddressValueError:
+                raise errors.CannotParseDomainError(
+                    "Cannot parse IPv4 '{}'".format(item)
+                )
+
+
+class IPv6Column(ByteFixedString):
+    ch_type = "IPv6"
+    py_types = (str, IPv6Address, bytes)
+
+    def __init__(self, types_check=False, **kwargs):
+        super(IPv6Column, self).__init__(16, types_check=types_check, **kwargs)
+
+        if types_check:
+
+            def check_item(value):
+                if isinstance(value, bytes) and len(value) != 16:
+                    raise ColumnTypeMismatchException(value)
+
+                if not isinstance(value, IPv6Address):
+                    try:
+                        value = IPv6Address(value)
+                    except AddressValueError:
+                        # Cannot parse input in a valid IPv6
+                        raise ColumnTypeMismatchException(value)
+
+            self.check_item = check_item
+
+    def after_read_items(self, items, nulls_map=None):
+        if nulls_map is None:
+            return tuple(IPv6Address(item) for item in items)
+        else:
+            return tuple(
+                (None if is_null else IPv6Address(items[i]))
+                for i, is_null in enumerate(nulls_map)
+            )
+
+    def before_write_items(self, 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
+                continue
+
+            # allow Ipv6 in bytes or python IPv6Address
+            # this is raw bytes (not encoded) in order to fit FixedString(16)
+            try:
+                if isinstance(item, bytes):
+                    continue
+
+                if not isinstance(item, IPv6Address):
+                    item = IPv6Address(item)
+                items[i] = item.packed
+            except AddressValueError:
+                raise errors.CannotParseDomainError(
+                    "Cannot parse IPv6 '{}'".format(item)
+                )

clickhouse_driver/columns/jsoncolumn.py

@@ -0,0 +1,37 @@
+from .base import Column
+from .stringcolumn import String
+from ..reader import read_binary_uint8, read_binary_str
+from ..util.compat import json
+from ..writer import write_binary_uint8
+
+
+class JsonColumn(Column):
+    py_types = (dict, )
+
+    # No NULL value actually
+    null_value = {}
+
+    def __init__(self, column_by_spec_getter, **kwargs):
+        self.column_by_spec_getter = column_by_spec_getter
+        self.string_column = String(**kwargs)
+        super(JsonColumn, self).__init__(**kwargs)
+
+    def write_state_prefix(self, buf):
+        # Read in binary format.
+        # Write in text format.
+        write_binary_uint8(1, buf)
+
+    def read_items(self, n_items, buf):
+        read_binary_uint8(buf)
+        spec = read_binary_str(buf)
+        col = self.column_by_spec_getter(spec)
+        col.read_state_prefix(buf)
+        return col.read_data(n_items, buf)
+
+    def write_items(self, items, buf):
+        items = [x if isinstance(x, str) else json.dumps(x) for x in items]
+        self.string_column.write_items(items, buf)
+
+
+def create_json_column(spec, column_by_spec_getter, column_options):
+    return JsonColumn(column_by_spec_getter, **column_options)

clickhouse_driver/columns/lowcardinalitycolumn.py

@@ -0,0 +1,142 @@
+from math import log
+
+from ..reader import read_binary_uint64
+from ..writer import write_binary_int64
+from .base import Column
+from .intcolumn import UInt8Column, UInt16Column, UInt32Column, UInt64Column
+
+
+def create_low_cardinality_column(spec, column_by_spec_getter, column_options):
+    inner = spec[15:-1]
+    nested = column_by_spec_getter(inner)
+    return LowCardinalityColumn(nested, **column_options)
+
+
+class LowCardinalityColumn(Column):
+    """
+    Stores column as index (unique elements) and keys.
+    Good for de-duplication of large values with low cardinality.
+    """
+    int_types = {
+        0: UInt8Column,
+        1: UInt16Column,
+        2: UInt32Column,
+        3: UInt64Column
+    }
+
+    # Need to read additional keys.
+    # Additional keys are stored before indexes as value N and N keys
+    # after them.
+    has_additional_keys_bit = 1 << 9
+    # Need to update dictionary.
+    # It means that previous granule has different dictionary.
+    need_update_dictionary = 1 << 10
+
+    serialization_type = has_additional_keys_bit | need_update_dictionary
+
+    def __init__(self, nested_column, **kwargs):
+        self.init_kwargs = kwargs
+        self.nested_column = nested_column
+        super(LowCardinalityColumn, self).__init__(**kwargs)
+
+    def read_state_prefix(self, buf):
+        super(LowCardinalityColumn, self).read_state_prefix(buf)
+
+        read_binary_uint64(buf)
+
+    def write_state_prefix(self, buf):
+        super(LowCardinalityColumn, self).write_state_prefix(buf)
+
+        # KeysSerializationVersion. See ClickHouse docs.
+        write_binary_int64(1, buf)
+
+    def _write_data(self, items, buf):
+        index, keys = [], []
+        key_by_index_element = {}
+        nested_is_nullable = False
+
+        if self.nested_column.nullable:
+            # First element represents NULL if column is nullable.
+            index.append(self.nested_column.null_value)
+            # Prevent null map writing. Reset nested column nullable flag.
+            self.nested_column.nullable = False
+            nested_is_nullable = True
+
+            for x in items:
+                if x is None:
+                    # Zero element for null.
+                    keys.append(0)
+
+                else:
+                    key = key_by_index_element.get(x)
+                    # Get key from index or add it to index.
+                    if key is None:
+                        key = len(key_by_index_element)
+                        key_by_index_element[x] = key
+                        index.append(x)
+
+                    keys.append(key + 1)
+        else:
+            for x in items:
+                key = key_by_index_element.get(x)
+
+                # Get key from index or add it to index.
+                if key is None:
+                    key = len(key_by_index_element)
+                    key_by_index_element[x] = len(key_by_index_element)
+                    index.append(x)
+
+                keys.append(key)
+
+        # Do not write anything for empty column.
+        # May happen while writing empty arrays.
+        if not len(index):
+            return
+
+        int_type = int(log(len(index), 2) / 8)
+        int_column = self.int_types[int_type](**self.init_kwargs)
+
+        serialization_type = self.serialization_type | int_type
+
+        write_binary_int64(serialization_type, buf)
+        write_binary_int64(len(index), buf)
+
+        if nested_is_nullable:
+            # Given we reset nested column nullable flag above,
+            # we need to write null map manually. If to invoke
+            # write_data method, it will cause an exception,
+            # because `prepare_data` may not be able to handle
+            # null value correctly.
+            self.nested_column.write_items(
+                [self.nested_column.null_value], buf)
+            # Remove null map from index, because it is already written.
+            index_to_write = index[1:]
+            self.nested_column.write_data(index_to_write, buf)
+        else:
+            self.nested_column.write_data(index, buf)
+        write_binary_int64(len(items), buf)
+        int_column.write_items(keys, buf)
+
+    def _read_data(self, n_items, buf, nulls_map=None):
+        if not n_items:
+            return tuple()
+
+        serialization_type = read_binary_uint64(buf)
+
+        # Lowest byte contains info about key type.
+        key_type = serialization_type & 0xf
+        keys_column = self.int_types[key_type](**self.init_kwargs)
+
+        nullable = self.nested_column.nullable
+        # Prevent null map reading. Reset nested column nullable flag.
+        self.nested_column.nullable = False
+
+        index_size = read_binary_uint64(buf)
+        index = self.nested_column.read_data(index_size, buf)
+        if nullable:
+            index = (None, ) + index[1:]
+
+        read_binary_uint64(buf)  # number of keys
+        keys = keys_column.read_data(n_items, buf)
+
+        return tuple(index[x] for x in keys)

clickhouse_driver/columns/mapcolumn.py

@@ -0,0 +1,73 @@
+import re
+from .base import Column
+from .intcolumn import UInt64Column
+from ..util.helpers import pairwise
+
+
+comma_re = re.compile(r',(?![^()]*\))')
+
+
+class MapColumn(Column):
+    py_types = (dict, )
+
+    null_value = {}
+
+    def __init__(self, key_column, value_column, **kwargs):
+        self.offset_column = UInt64Column(**kwargs)
+        self.key_column = key_column
+        self.value_column = value_column
+        super(MapColumn, self).__init__(**kwargs)
+
+    def read_state_prefix(self, buf):
+        super(MapColumn, self).read_state_prefix(buf)
+
+        self.key_column.read_state_prefix(buf)
+        self.value_column.read_state_prefix(buf)
+
+    def write_state_prefix(self, buf):
+        super(MapColumn, self).write_state_prefix(buf)
+
+        self.key_column.write_state_prefix(buf)
+        self.value_column.write_state_prefix(buf)
+
+    def read_items(self, n_items, buf):
+        if not n_items:
+            return [{}]
+
+        offsets = list(self.offset_column.read_items(n_items, buf))
+        last_offset = offsets[-1]
+        keys = self.key_column.read_data(last_offset, buf)
+        values = self.value_column.read_data(last_offset, buf)
+
+        offsets.insert(0, 0)
+
+        return [
+            dict(zip(keys[begin:end], values[begin:end]))
+            for begin, end in pairwise(offsets)
+        ]
+
+    def write_items(self, items, buf):
+        offsets = []
+        keys = []
+        values = []
+
+        total = 0
+        for x in items:
+            total += len(x)
+            offsets.append(total)
+            keys.extend(x.keys())
+            values.extend(x.values())
+
+        self.offset_column.write_items(offsets, buf)
+        self.key_column.write_data(keys, buf)
+        self.value_column.write_data(values, buf)
+
+
+def create_map_column(spec, column_by_spec_getter, column_options):
+    # Match commas outside of parentheses, so we don't match the comma in
+    # Decimal types.
+    key, value = comma_re.split(spec[4:-1])
+    key_column = column_by_spec_getter(key.strip())
+    value_column = column_by_spec_getter(value.strip())
+
+    return MapColumn(key_column, value_column, **column_options)

clickhouse_driver/columns/nestedcolumn.py

@@ -0,0 +1,10 @@
+
+from .arraycolumn import create_array_column
+from .util import get_inner_spec
+
+
+def create_nested_column(spec, column_by_spec_getter, column_options):
+    return create_array_column(
+        'Array(Tuple({}))'.format(get_inner_spec('Nested', spec)),
+        column_by_spec_getter, column_options
+    )

clickhouse_driver/columns/nothingcolumn.py

@@ -0,0 +1,13 @@
+from .intcolumn import FormatColumn
+
+
+class NothingColumn(FormatColumn):
+    ch_type = 'Nothing'
+    format = 'B'
+
+    @property
+    def size(self):
+        return 1
+
+    def after_read_items(self, items, nulls_map=None):
+        return (None, ) * len(items)

clickhouse_driver/columns/nullablecolumn.py

@@ -0,0 +1,7 @@
+
+
+def create_nullable_column(spec, column_by_spec_getter):
+    inner = spec[9:-1]
+    nested = column_by_spec_getter(inner)
+    nested.nullable = True
+    return nested

clickhouse_driver/columns/nullcolumn.py

@@ -0,0 +1,15 @@
+from .intcolumn import FormatColumn
+
+
+# TODO: Drop Null column support in future.
+# Compatibility with old servers.
+class NullColumn(FormatColumn):
+    ch_type = 'Null'
+    format = 'B'
+
+    @property
+    def size(self):
+        return 1
+
+    def after_read_items(self, items, nulls_map=None):
+        return (None, ) * len(items)

clickhouse_driver/columns/numpy/base.py

@@ -0,0 +1,47 @@
+import numpy as np
+import pandas as pd
+
+from ..base import Column
+
+
+class NumpyColumn(Column):
+    dtype = None
+
+    normalize_null_value = True
+
+    def read_items(self, n_items, buf):
+        data = buf.read(n_items * self.dtype.itemsize)
+        return np.frombuffer(data, self.dtype.newbyteorder('<'), n_items)
+
+    def write_items(self, items, buf):
+        buf.write(items.astype(self.dtype.newbyteorder('<')).tobytes())
+
+    def _write_nulls_map(self, items, buf):
+        s = self.make_null_struct(len(items))
+        nulls_map = self._get_nulls_map(items)
+        buf.write(s.pack(*nulls_map))
+
+    def _get_nulls_map(self, items):
+        return [bool(x) for x in pd.isnull(items)]
+
+    def _read_data(self, n_items, buf, nulls_map=None):
+        items = self.read_items(n_items, buf)
+
+        if self.after_read_items:
+            return self.after_read_items(items, nulls_map)
+        elif nulls_map is not None:
+            items = np.array(items, dtype=object)
+            np.place(items, nulls_map, None)
+
+        return items
+
+    def prepare_items(self, items):
+        nulls_map = pd.isnull(items)
+
+        # Always replace null values to null_value for proper inserts into
+        # non-nullable columns.
+        if isinstance(items, np.ndarray) and self.normalize_null_value:
+            items = np.array(items)
+            np.place(items, nulls_map, self.null_value)
+
+        return items

clickhouse_driver/columns/numpy/boolcolumn.py

@@ -0,0 +1,8 @@
+import numpy as np
+
+from .base import NumpyColumn
+
+
+class NumpyBoolColumn(NumpyColumn):
+    dtype = np.dtype(np.bool_)
+    ch_type = 'Bool'

clickhouse_driver/columns/numpy/datecolumn.py

@@ -0,0 +1,19 @@
+import numpy as np
+
+from .base import NumpyColumn
+
+
+class NumpyDateColumn(NumpyColumn):
+    dtype = np.dtype(np.uint16)
+    ch_type = 'Date'
+
+    null_value = np.datetime64(0, 'Y')
+
+    def read_items(self, n_items, buf):
+        data = super(NumpyDateColumn, self).read_items(n_items, buf)
+        return data.astype('datetime64[D]')
+
+    def write_items(self, items, buf):
+        super(NumpyDateColumn, self).write_items(
+            items.astype('datetime64[D]'), buf
+        )

clickhouse_driver/columns/numpy/datetimecolumn.py

@@ -0,0 +1,143 @@
+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).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):
+            if local_tz_name != context.server_info.timezone:
+                tz_name = context.server_info.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)

clickhouse_driver/columns/numpy/floatcolumn.py

@@ -0,0 +1,24 @@
+import numpy as np
+
+from .base import NumpyColumn
+
+# normalize_null_value = False due to float('nan')
+# With normalization pandas.isnull will threat float('nan') as NULL value.
+
+
+class NumpyFloat32Column(NumpyColumn):
+    dtype = np.dtype(np.float32)
+    ch_type = 'Float32'
+    normalize_null_value = False
+
+    def _get_nulls_map(self, items):
+        return [x is None for x in items]
+
+
+class NumpyFloat64Column(NumpyColumn):
+    dtype = np.dtype(np.float64)
+    ch_type = 'Float64'
+    normalize_null_value = False
+
+    def _get_nulls_map(self, items):
+        return [x is None for x in items]