clickhouse-driver 0.2.1__cp39-cp39-win_amd64.whl → 0.2.8__cp39-cp39-win_amd64.whl

clickhouse_driver/columns/base.py

@@ -1,147 +1,221 @@
-from struct import Struct, error as struct_error
-
-from . import exceptions
-
-
-class Column(object):
-    ch_type = None
-    py_types = None
-
-    check_item = None
-    after_read_items = None
-    before_write_items = None
-
-    types_check_enabled = False
-
-    null_value = 0
-
-    def __init__(self, types_check=False, **kwargs):
-        self.nullable = False
-        self.types_check_enabled = types_check
-        super(Column, self).__init__()
-
-    def make_null_struct(self, n_items):
-        return Struct('<{}B'.format(n_items))
-
-    def _read_nulls_map(self, n_items, buf):
-        s = self.make_null_struct(n_items)
-        return s.unpack(buf.read(s.size))
-
-    def _write_nulls_map(self, items, buf):
-        s = self.make_null_struct(len(items))
-        items = [x is None for x in items]
-        buf.write(s.pack(*items))
-
-    def check_item_type(self, value):
-        if not isinstance(value, self.py_types):
-            raise exceptions.ColumnTypeMismatchException(value)
-
-    def prepare_items(self, items):
-        nullable = self.nullable
-        null_value = self.null_value
-
-        check_item = self.check_item
-        if self.types_check_enabled:
-            check_item_type = self.check_item_type
-        else:
-            check_item_type = False
-
-        if (not self.nullable and not check_item_type and
-                not check_item and not self.before_write_items):
-            return items
-
-        nulls_map = [False] * len(items) if self.nullable else None
-        for i, x in enumerate(items):
-            if x is None and nullable:
-                nulls_map[i] = True
-                x = null_value
-
-            else:
-                if check_item_type:
-                    check_item_type(x)
-
-                if check_item:
-                    check_item(x)
-
-            items[i] = x
-
-        if self.before_write_items:
-            self.before_write_items(items, nulls_map=nulls_map)
-
-        return items
-
-    def write_data(self, items, buf):
-        if self.nullable:
-            self._write_nulls_map(items, buf)
-
-        self._write_data(items, buf)
-
-    def _write_data(self, items, buf):
-        prepared = self.prepare_items(items)
-        self.write_items(prepared, buf)
-
-    def write_items(self, items, buf):
-        raise NotImplementedError
-
-    def read_data(self, n_items, buf):
-        if self.nullable:
-            nulls_map = self._read_nulls_map(n_items, buf)
-        else:
-            nulls_map = None
-
-        return self._read_data(n_items, buf, nulls_map=nulls_map)
-
-    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:
-            return tuple(
-                (None if is_null else items[i])
-                for i, is_null in enumerate(nulls_map)
-            )
-        return items
-
-    def read_items(self, n_items, buf):
-        raise NotImplementedError
-
-    def read_state_prefix(self, buf):
-        pass
-
-    def write_state_prefix(self, buf):
-        pass
-
-
-class FormatColumn(Column):
-    """
-    Uses struct.pack for bulk items writing.
-    """
-
-    format = None
-
-    def make_struct(self, n_items):
-        return Struct('<{}{}'.format(n_items, self.format))
-
-    def write_items(self, items, buf):
-        s = self.make_struct(len(items))
-        try:
-            buf.write(s.pack(*items))
-
-        except struct_error as e:
-            raise exceptions.StructPackException(e)
-
-    def read_items(self, n_items, buf):
-        s = self.make_struct(n_items)
-        return s.unpack(buf.read(s.size))
-
-
-# How to write new column?
-# - Check ClickHouse documentation for column
-# - Wireshark and tcpdump are your friends.
-# - Use `clickhouse-client --compression 0` to see what's going on data
-#   transmission.
-# - Check for similar existing columns and tests.
-# - Use `FormatColumn` for columns that use "simple" types under the hood.
-# - Some columns have before_write and after_read hooks.
-#   Use them to convert items in column into "simple" types.
+from struct import Struct, error as struct_error
+
+from . import exceptions
+from ..varint import read_varint
+
+
+class CommonSerialization(object):
+    def __init__(self, column):
+        self.column = column
+        super(CommonSerialization, self).__init__()
+
+    def read_sparse(self, n_items, buf):
+        return n_items
+
+    def apply_sparse(self, items):
+        return items
+
+
+class SparseSerialization(CommonSerialization):
+
+    def __init__(self, *args, **kwargs):
+        self.sparse_indexes = []
+        self.items_total = None
+        super(SparseSerialization, self).__init__(*args, **kwargs)
+
+    def read_sparse(self, n_items, buf):
+        sparse_indexes = []
+        items_total = 0
+        non_default_items = 0
+
+        END_OF_GRANULE_FLAG = 1 << 62
+        end_of_granule = False
+
+        while not end_of_granule:
+            group_size = read_varint(buf)
+            end_of_granule = group_size & END_OF_GRANULE_FLAG
+            group_size &= ~END_OF_GRANULE_FLAG
+
+            items_total += group_size + 1
+            if not end_of_granule:
+                non_default_items += 1
+                sparse_indexes.append(items_total)
+
+        self.sparse_indexes = sparse_indexes
+        self.items_total = items_total
+
+        return non_default_items
+
+    def apply_sparse(self, items):
+        default = self.column.null_value
+        if self.column.after_read_items:
+            default = self.column.after_read_items([default])[0]
+
+        rv = [default] * (self.items_total - 1)
+        for item_number, i in enumerate(self.sparse_indexes):
+            rv[i - 1] = items[item_number]
+
+        return rv
+
+
+class Column(object):
+    ch_type = None
+    py_types = None
+
+    check_item = None
+    after_read_items = None
+    before_write_items = None
+
+    types_check_enabled = False
+
+    null_value = 0
+
+    def __init__(self, types_check=False, has_custom_serialization=False,
+                 **kwargs):
+        self.nullable = False
+        self.types_check_enabled = types_check
+        self.has_custom_serialization = has_custom_serialization
+        self.serialization = CommonSerialization(self)
+        self.input_null_as_default = False
+
+        self.context = kwargs['context']
+        self.input_null_as_default = self.context.client_settings \
+            .get('input_format_null_as_default', False)
+
+        super(Column, self).__init__()
+
+    def make_null_struct(self, n_items):
+        return Struct('<{}B'.format(n_items))
+
+    def _read_nulls_map(self, n_items, buf):
+        s = self.make_null_struct(n_items)
+        return s.unpack(buf.read(s.size))
+
+    def _write_nulls_map(self, items, buf):
+        s = self.make_null_struct(len(items))
+        items = [x is None for x in items]
+        buf.write(s.pack(*items))
+
+    def check_item_type(self, value):
+        if not isinstance(value, self.py_types):
+            raise exceptions.ColumnTypeMismatchException(value)
+
+    def prepare_items(self, items):
+        nullable = self.nullable
+        null_value = self.null_value
+        null_as_default = self.input_null_as_default
+
+        check_item = self.check_item
+        if self.types_check_enabled:
+            check_item_type = self.check_item_type
+        else:
+            check_item_type = False
+
+        if (not (self.nullable or null_as_default) and not check_item_type and
+                not check_item and not self.before_write_items):
+            return items
+
+        nulls_map = [False] * len(items) if self.nullable else None
+        for i, x in enumerate(items):
+            if x is None:
+                if nullable:
+                    nulls_map[i] = True
+                    x = null_value
+                elif null_as_default:
+                    x = null_value
+
+            else:
+                if check_item_type:
+                    check_item_type(x)
+
+                if check_item:
+                    check_item(x)
+
+            items[i] = x
+
+        if self.before_write_items:
+            self.before_write_items(items, nulls_map=nulls_map)
+
+        return items
+
+    def write_data(self, items, buf):
+        if self.nullable:
+            self._write_nulls_map(items, buf)
+
+        self._write_data(items, buf)
+
+    def _write_data(self, items, buf):
+        prepared = self.prepare_items(items)
+        self.write_items(prepared, buf)
+
+    def write_items(self, items, buf):
+        raise NotImplementedError
+
+    def read_data(self, n_items, buf):
+        n_items = self.serialization.read_sparse(n_items, buf)
+
+        if self.nullable:
+            nulls_map = self._read_nulls_map(n_items, buf)
+        else:
+            nulls_map = None
+
+        items = self._read_data(n_items, buf, nulls_map=nulls_map)
+        return self.serialization.apply_sparse(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:
+            return tuple(
+                (None if is_null else items[i])
+                for i, is_null in enumerate(nulls_map)
+            )
+        return items
+
+    def read_items(self, n_items, buf):
+        raise NotImplementedError
+
+    def read_state_prefix(self, buf):
+        if self.has_custom_serialization:
+            use_custom_serialization = read_varint(buf)
+            if use_custom_serialization:
+                self.serialization = SparseSerialization(self)
+
+    def write_state_prefix(self, buf):
+        pass
+
+
+class FormatColumn(Column):
+    """
+    Uses struct.pack for bulk items writing.
+    """
+
+    format = None
+
+    def make_struct(self, n_items):
+        return Struct('<{}{}'.format(n_items, self.format))
+
+    def write_items(self, items, buf):
+        s = self.make_struct(len(items))
+        try:
+            buf.write(s.pack(*items))
+
+        except struct_error as e:
+            raise exceptions.StructPackException(e)
+
+    def read_items(self, n_items, buf):
+        s = self.make_struct(n_items)
+        return s.unpack(buf.read(s.size))
+
+
+# How to write new column?
+# - Check ClickHouse documentation for column
+# - Wireshark and tcpdump are your friends.
+# - Use `clickhouse-client --compression 0` to see what's going on data
+#   transmission.
+# - Check for similar existing columns and tests.
+# - Use `FormatColumn` for columns that use "simple" types under the hood.
+# - Some columns have before_write and after_read hooks.
+#   Use them to convert items in column into "simple" types.

clickhouse_driver/columns/boolcolumn.py

@@ -0,0 +1,7 @@
+from .base import FormatColumn
+
+
+class BoolColumn(FormatColumn):
+    ch_type = 'Bool'
+    py_types = (bool, )
+    format = '?'

clickhouse_driver/columns/datecolumn.py

@@ -1,49 +1,108 @@
-from datetime import date, timedelta
-
-from .base import FormatColumn
-
-
-epoch_start = date(1970, 1, 1)
-
-
-class DateColumn(FormatColumn):
-    ch_type = 'Date'
-    py_types = (date, )
-    format = 'H'
-
-    epoch_start = epoch_start
-    epoch_end = date(2105, 12, 31)
-
-    date_lut = {x: epoch_start + timedelta(x) for x in range(65535)}
-    date_lut_reverse = {value: key for key, value in date_lut.items()}
-
-    def before_write_items(self, items, nulls_map=None):
-        null_value = self.null_value
-
-        date_lut_reverse = self.date_lut_reverse
-        epoch_start = self.epoch_start
-        epoch_end = self.epoch_end
-
-        for i, item in enumerate(items):
-            if nulls_map and nulls_map[i]:
-                items[i] = null_value
-                continue
-
-            if type(item) != date:
-                item = date(item.year, item.month, item.day)
-
-            if item > epoch_end or item < epoch_start:
-                items[i] = 0
-            else:
-                items[i] = date_lut_reverse[item]
-
-    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)
-            )
+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