clickhouse-driver 0.2.10__cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl

clickhouse_driver/columns/numpy/intcolumn.py

@@ -0,0 +1,43 @@
+import numpy as np
+
+from .base import NumpyColumn
+
+
+class NumpyInt8Column(NumpyColumn):
+    dtype = np.dtype(np.int8)
+    ch_type = 'Int8'
+
+
+class NumpyUInt8Column(NumpyColumn):
+    dtype = np.dtype(np.uint8)
+    ch_type = 'UInt8'
+
+
+class NumpyInt16Column(NumpyColumn):
+    dtype = np.dtype(np.int16)
+    ch_type = 'Int16'
+
+
+class NumpyUInt16Column(NumpyColumn):
+    dtype = np.dtype(np.uint16)
+    ch_type = 'UInt16'
+
+
+class NumpyInt32Column(NumpyColumn):
+    dtype = np.dtype(np.int32)
+    ch_type = 'Int32'
+
+
+class NumpyUInt32Column(NumpyColumn):
+    dtype = np.dtype(np.uint32)
+    ch_type = 'UInt32'
+
+
+class NumpyInt64Column(NumpyColumn):
+    dtype = np.dtype(np.int64)
+    ch_type = 'Int64'
+
+
+class NumpyUInt64Column(NumpyColumn):
+    dtype = np.dtype(np.uint64)
+    ch_type = 'UInt64'

clickhouse_driver/columns/numpy/lowcardinalitycolumn.py

@@ -0,0 +1,96 @@
+from math import log
+
+import numpy as np
+import pandas as pd
+
+from ..lowcardinalitycolumn import LowCardinalityColumn
+from ...reader import read_binary_uint64
+from ...writer import write_binary_int64
+from .intcolumn import (
+    NumpyUInt8Column, NumpyUInt16Column, NumpyUInt32Column, NumpyUInt64Column
+)
+
+
+class NumpyLowCardinalityColumn(LowCardinalityColumn):
+    int_types = {
+        0: NumpyUInt8Column,
+        1: NumpyUInt16Column,
+        2: NumpyUInt32Column,
+        3: NumpyUInt64Column
+    }
+
+    def __init__(self, nested_column, **kwargs):
+        super(NumpyLowCardinalityColumn, self).__init__(nested_column,
+                                                        **kwargs)
+
+    def _write_data(self, items, buf):
+        # Do not write anything for empty column.
+        # May happen while writing empty arrays.
+        if not len(items):
+            return
+
+        # Replace nans with defaults if not nullabe.
+        if isinstance(items, np.ndarray) and not self.nested_column.nullable:
+            nulls = pd.isnull(items)
+            items = np.where(nulls, self.nested_column.null_value, items)
+
+        c = pd.Categorical(items)
+
+        int_type = int(log(len(c.codes), 2) / 8)
+        int_column = self.int_types[int_type](**self.init_kwargs)
+
+        serialization_type = self.serialization_type | int_type
+
+        index = c.categories
+        keys = c.codes
+
+        if self.nested_column.nullable:
+            # First element represents NULL if column is nullable.
+            index = index.insert(0, self.nested_column.null_value)
+            keys = keys + 1
+            # Prevent null map writing. Reset nested column nullable flag.
+            self.nested_column.nullable = False
+
+        write_binary_int64(serialization_type, buf)
+        write_binary_int64(len(index), buf)
+
+        self.nested_column.write_data(index.to_numpy(items.dtype), 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)
+
+        read_binary_uint64(buf)  # number of keys
+        keys = keys_column.read_data(n_items, buf)
+
+        if nullable:
+            # Shift all codes by one ("No value" code is -1 for pandas
+            # categorical) and drop corresponding first index
+            # this is analog of original operation:
+            # index = (None, ) + index[1:]
+            keys = np.array(keys, dtype='int64')  # deal with possible overflow
+            keys = keys - 1
+            index = index[1:]
+        return pd.Categorical.from_codes(keys, index)
+
+
+def create_numpy_low_cardinality_column(spec, column_by_spec_getter,
+                                        column_options):
+    inner = spec[15:-1]
+    nested = column_by_spec_getter(inner)
+    return NumpyLowCardinalityColumn(nested, **column_options)

clickhouse_driver/columns/numpy/service.py

@@ -0,0 +1,58 @@
+from ..service import aliases
+from ... import errors
+from .datecolumn import NumpyDateColumn
+from .datetimecolumn import create_numpy_datetime_column
+from .floatcolumn import NumpyFloat32Column, NumpyFloat64Column
+from .intcolumn import (
+    NumpyInt8Column, NumpyInt16Column, NumpyInt32Column, NumpyInt64Column,
+    NumpyUInt8Column, NumpyUInt16Column, NumpyUInt32Column, NumpyUInt64Column
+)
+from .boolcolumn import NumpyBoolColumn
+from .lowcardinalitycolumn import create_numpy_low_cardinality_column
+from .stringcolumn import create_string_column
+from .tuplecolumn import create_tuple_column
+from ..nullablecolumn import create_nullable_column
+
+column_by_type = {c.ch_type: c for c in [
+    NumpyDateColumn,
+    NumpyFloat32Column, NumpyFloat64Column,
+    NumpyInt8Column, NumpyInt16Column, NumpyInt32Column, NumpyInt64Column,
+    NumpyUInt8Column, NumpyUInt16Column, NumpyUInt32Column, NumpyUInt64Column,
+    NumpyBoolColumn
+]}
+
+
+def get_numpy_column_by_spec(spec, column_options):
+    def create_column_with_options(x):
+        return get_numpy_column_by_spec(x, column_options)
+
+    if spec == 'String' or spec.startswith('FixedString'):
+        return create_string_column(spec, column_options)
+
+    elif spec.startswith('DateTime'):
+        return create_numpy_datetime_column(spec, column_options)
+
+    elif spec.startswith('Tuple'):
+        return create_tuple_column(
+            spec, create_column_with_options, column_options
+        )
+
+    elif spec.startswith('Nullable'):
+        return create_nullable_column(spec, create_column_with_options)
+
+    elif spec.startswith('LowCardinality'):
+        return create_numpy_low_cardinality_column(
+            spec, create_column_with_options, column_options
+        )
+    else:
+        for alias, primitive in aliases:
+            if spec.startswith(alias):
+                return create_column_with_options(
+                    primitive + spec[len(alias):]
+                )
+
+        if spec in column_by_type:
+            cls = column_by_type[spec]
+            return cls(**column_options)
+
+        raise errors.UnknownTypeError('Unknown type {}'.format(spec))

clickhouse_driver/columns/numpy/stringcolumn.py

@@ -0,0 +1,78 @@
+import numpy as np
+
+from ... import defines
+from .base import NumpyColumn
+
+
+class NumpyStringColumn(NumpyColumn):
+    null_value = ''
+
+    default_encoding = defines.STRINGS_ENCODING
+
+    def __init__(self, encoding=default_encoding, **kwargs):
+        self.encoding = encoding
+        super(NumpyStringColumn, self).__init__(**kwargs)
+
+    def read_items(self, n_items, buf):
+        return np.array(
+            buf.read_strings(n_items, encoding=self.encoding), dtype=self.dtype
+        )
+
+    def write_items(self, items, buf):
+        return buf.write_strings(items.tolist(), encoding=self.encoding)
+
+
+class NumpyByteStringColumn(NumpyColumn):
+    null_value = b''
+
+    def read_items(self, n_items, buf):
+        return np.array(buf.read_strings(n_items), dtype=self.dtype)
+
+    def write_items(self, items, buf):
+        return buf.write_strings(items.tolist())
+
+
+class NumpyFixedString(NumpyStringColumn):
+    def __init__(self, length, **kwargs):
+        self.length = length
+        super(NumpyFixedString, self).__init__(**kwargs)
+
+    def read_items(self, n_items, buf):
+        return np.array(buf.read_fixed_strings(
+            n_items, self.length, encoding=self.encoding
+        ), dtype=self.dtype)
+
+    def write_items(self, items, buf):
+        return buf.write_fixed_strings(
+            items.tolist(), self.length, encoding=self.encoding
+        )
+
+
+class NumpyByteFixedString(NumpyByteStringColumn):
+    def __init__(self, length, **kwargs):
+        self.length = length
+        super(NumpyByteFixedString, self).__init__(**kwargs)
+
+    def read_items(self, n_items, buf):
+        return np.array(
+            buf.read_fixed_strings(n_items, self.length), dtype=self.dtype
+        )
+
+    def write_items(self, items, buf):
+        return buf.write_fixed_strings(items.tolist(), self.length)
+
+
+def create_string_column(spec, column_options):
+    client_settings = column_options['context'].client_settings
+    strings_as_bytes = client_settings['strings_as_bytes']
+    encoding = client_settings.get(
+        'strings_encoding', NumpyStringColumn.default_encoding
+    )
+
+    if spec == 'String':
+        cls = NumpyByteStringColumn if strings_as_bytes else NumpyStringColumn
+        return cls(encoding=encoding, **column_options)
+    else:
+        length = int(spec[12:-1])
+        cls = NumpyByteFixedString if strings_as_bytes else NumpyFixedString
+        return cls(length, encoding=encoding, **column_options)

clickhouse_driver/columns/numpy/tuplecolumn.py

@@ -0,0 +1,37 @@
+import numpy as np
+
+from .base import NumpyColumn
+from ..util import get_inner_spec, get_inner_columns
+
+
+class TupleColumn(NumpyColumn):
+    def __init__(self, nested_columns, **kwargs):
+        self.nested_columns = nested_columns
+        super(TupleColumn, self).__init__(**kwargs)
+
+    def write_data(self, items, buf):
+        names = items.dtype.names
+        for i, (x, name) in enumerate(zip(self.nested_columns, names)):
+            x.write_data(items[name], buf)
+
+    def write_items(self, items, buf):
+        return self.write_data(items, buf)
+
+    def read_data(self, n_items, buf):
+        data = [x.read_data(n_items, buf) for x in self.nested_columns]
+        dtype = [('f{}'.format(i), x.dtype) for i, x in enumerate(data)]
+        rv = np.empty(n_items, dtype=dtype)
+        for i, x in enumerate(data):
+            rv['f{}'.format(i)] = x
+        return rv
+
+    def read_items(self, n_items, buf):
+        return self.read_data(n_items, buf)
+
+
+def create_tuple_column(spec, column_by_spec_getter, column_options):
+    inner_spec = get_inner_spec('Tuple', spec)
+    columns = get_inner_columns(inner_spec)
+
+    return TupleColumn([column_by_spec_getter(x) for x in columns],
+                       **column_options)

clickhouse_driver/columns/service.py

@@ -0,0 +1,185 @@
+import logging
+
+from .. import errors
+from .arraycolumn import create_array_column
+from .boolcolumn import BoolColumn
+from .datecolumn import DateColumn, Date32Column
+from .datetimecolumn import create_datetime_column
+from .decimalcolumn import create_decimal_column
+from . import exceptions as column_exceptions
+from .enumcolumn import create_enum_column
+from .floatcolumn import Float32Column, Float64Column
+from .intcolumn import (
+    Int8Column, Int16Column, Int32Column, Int64Column,
+    Int128Column, UInt128Column, Int256Column, UInt256Column,
+    UInt8Column, UInt16Column, UInt32Column, UInt64Column
+)
+from .lowcardinalitycolumn import create_low_cardinality_column
+from .jsoncolumn import create_json_column
+from .mapcolumn import create_map_column
+from .nothingcolumn import NothingColumn
+from .nullcolumn import NullColumn
+from .nullablecolumn import create_nullable_column
+from .simpleaggregatefunctioncolumn import (
+    create_simple_aggregate_function_column
+)
+from .stringcolumn import create_string_column
+from .tuplecolumn import create_tuple_column
+from .nestedcolumn import create_nested_column
+from .uuidcolumn import UUIDColumn
+from .intervalcolumn import (
+    IntervalYearColumn, IntervalMonthColumn, IntervalWeekColumn,
+    IntervalDayColumn, IntervalHourColumn, IntervalMinuteColumn,
+    IntervalSecondColumn
+)
+from .ipcolumn import IPv4Column, IPv6Column
+
+
+column_by_type = {c.ch_type: c for c in [
+    DateColumn, Date32Column, Float32Column, Float64Column,
+    Int8Column, Int16Column, Int32Column, Int64Column,
+    Int128Column, UInt128Column, Int256Column, UInt256Column,
+    UInt8Column, UInt16Column, UInt32Column, UInt64Column,
+    NothingColumn, NullColumn, UUIDColumn,
+    IntervalYearColumn, IntervalMonthColumn, IntervalWeekColumn,
+    IntervalDayColumn, IntervalHourColumn, IntervalMinuteColumn,
+    IntervalSecondColumn, IPv4Column, IPv6Column, BoolColumn
+]}
+
+logger = logging.getLogger(__name__)
+
+
+aliases = [
+    # Begin Geo types
+    ('Point', 'Tuple(Float64, Float64)'),
+    ('Ring', 'Array(Point)'),
+    ('Polygon', 'Array(Ring)'),
+    ('MultiPolygon', 'Array(Polygon)')
+    # End Geo types
+]
+
+
+def get_column_by_spec(spec, column_options, use_numpy=None):
+    context = column_options['context']
+
+    if use_numpy is None:
+        use_numpy = context.client_settings['use_numpy'] if context else False
+
+    if use_numpy:
+        from .numpy.service import get_numpy_column_by_spec
+
+        try:
+            return get_numpy_column_by_spec(spec, column_options)
+        except errors.UnknownTypeError:
+            use_numpy = False
+            logger.warning('NumPy support is not implemented for %s. '
+                           'Using generic column', spec)
+
+    def create_column_with_options(x):
+        return get_column_by_spec(x, column_options, use_numpy=use_numpy)
+
+    if spec == 'String' or spec.startswith('FixedString'):
+        return create_string_column(spec, column_options)
+
+    elif spec.startswith('Enum'):
+        return create_enum_column(spec, column_options)
+
+    elif spec.startswith('DateTime'):
+        return create_datetime_column(spec, column_options)
+
+    elif spec.startswith('Decimal'):
+        return create_decimal_column(spec, column_options)
+
+    elif spec.startswith('Array'):
+        return create_array_column(
+            spec, create_column_with_options, column_options
+        )
+
+    elif spec.startswith('Tuple'):
+        return create_tuple_column(
+            spec, create_column_with_options, column_options
+        )
+
+    elif spec.startswith('Nested'):
+        return create_nested_column(
+            spec, create_column_with_options, column_options
+        )
+
+    elif spec.startswith('Nullable'):
+        return create_nullable_column(spec, create_column_with_options)
+
+    elif spec.startswith('LowCardinality'):
+        return create_low_cardinality_column(
+            spec, create_column_with_options, column_options
+        )
+
+    elif spec.startswith('SimpleAggregateFunction'):
+        return create_simple_aggregate_function_column(
+            spec, create_column_with_options
+        )
+
+    elif spec.startswith('Map'):
+        return create_map_column(
+            spec, create_column_with_options, column_options
+        )
+
+    elif spec.startswith("Object('json')"):
+        return create_json_column(
+            spec, create_column_with_options, column_options
+        )
+
+    else:
+        for alias, primitive in aliases:
+            if spec.startswith(alias):
+                return create_column_with_options(
+                    primitive + spec[len(alias):]
+                )
+
+        try:
+            cls = column_by_type[spec]
+            return cls(**column_options)
+
+        except KeyError:
+            raise errors.UnknownTypeError('Unknown type {}'.format(spec))
+
+
+def read_column(context, column_spec, n_items, buf, use_numpy=None,
+                has_custom_serialization=False):
+    column_options = {
+        'context': context,
+        'has_custom_serialization': has_custom_serialization
+    }
+    col = get_column_by_spec(column_spec, column_options, use_numpy=use_numpy)
+    col.read_state_prefix(buf)
+    return col.read_data(n_items, buf)
+
+
+def write_column(context, column_name, column_spec, items, buf,
+                 types_check=False):
+    column_options = {
+        'context': context,
+        'types_check': types_check
+    }
+    column = get_column_by_spec(column_spec, column_options)
+
+    try:
+        column.write_state_prefix(buf)
+        column.write_data(items, buf)
+
+    except column_exceptions.ColumnTypeMismatchException as e:
+        raise errors.TypeMismatchError(
+            'Type mismatch in VALUES section. '
+            'Expected {} got {}: {} for column "{}".'.format(
+                column_spec, type(e.args[0]), e.args[0], column_name
+            )
+        )
+
+    except (column_exceptions.StructPackException, OverflowError) as e:
+        error = e.args[0]
+        raise errors.TypeMismatchError(
+            'Type mismatch in VALUES section. '
+            'Repeat query with types_check=True for detailed info. '
+            'Column {}: {}'.format(
+                column_name, str(error)
+            )
+        )

clickhouse_driver/columns/simpleaggregatefunctioncolumn.py

@@ -0,0 +1,7 @@
+
+
+def create_simple_aggregate_function_column(spec, column_by_spec_getter):
+    # SimpleAggregateFunction(Func, Type) -> Type
+    inner = spec[24:-1].split(',', 1)[1].strip()
+    nested = column_by_spec_getter(inner)
+    return nested

clickhouse_driver/columns/stringcolumn.py

@@ -0,0 +1,73 @@
+
+from .. import defines
+from .base import Column
+
+
+class String(Column):
+    ch_type = 'String'
+    py_types = (str, )
+    null_value = ''
+
+    default_encoding = defines.STRINGS_ENCODING
+
+    def __init__(self, encoding=default_encoding, **kwargs):
+        self.encoding = encoding
+        super(String, self).__init__(**kwargs)
+
+    def write_items(self, items, buf):
+        buf.write_strings(items, encoding=self.encoding)
+
+    def read_items(self, n_items, buf):
+        return buf.read_strings(n_items, encoding=self.encoding)
+
+
+class ByteString(String):
+    py_types = (bytes, )
+    null_value = b''
+
+    def write_items(self, items, buf):
+        buf.write_strings(items)
+
+    def read_items(self, n_items, buf):
+        return buf.read_strings(n_items)
+
+
+class FixedString(String):
+    ch_type = 'FixedString'
+
+    def __init__(self, length, **kwargs):
+        self.length = length
+        super(FixedString, self).__init__(**kwargs)
+
+    def read_items(self, n_items, buf):
+        return buf.read_fixed_strings(
+            n_items, self.length, encoding=self.encoding
+        )
+
+    def write_items(self, items, buf):
+        buf.write_fixed_strings(items, self.length, encoding=self.encoding)
+
+
+class ByteFixedString(FixedString):
+    py_types = (bytearray, bytes)
+    null_value = b''
+
+    def read_items(self, n_items, buf):
+        return buf.read_fixed_strings(n_items, self.length)
+
+    def write_items(self, items, buf):
+        buf.write_fixed_strings(items, self.length)
+
+
+def create_string_column(spec, column_options):
+    client_settings = column_options['context'].client_settings
+    strings_as_bytes = client_settings['strings_as_bytes']
+    encoding = client_settings.get('strings_encoding', String.default_encoding)
+
+    if spec == 'String':
+        cls = ByteString if strings_as_bytes else String
+        return cls(encoding=encoding, **column_options)
+    else:
+        length = int(spec[12:-1])
+        cls = ByteFixedString if strings_as_bytes else FixedString
+        return cls(length, encoding=encoding, **column_options)

clickhouse_driver/columns/tuplecolumn.py

@@ -0,0 +1,63 @@
+
+from .base import Column
+from .util import get_inner_spec, get_inner_columns_with_types
+
+
+class TupleColumn(Column):
+    py_types = (list, tuple)
+
+    def __init__(self, names, nested_columns, **kwargs):
+        self.names = names
+        self.nested_columns = nested_columns
+        client_settings = kwargs['context'].client_settings
+        settings = kwargs['context'].settings
+        self.namedtuple_as_json = (
+            settings.get('allow_experimental_object_type', False) and
+            client_settings.get('namedtuple_as_json', True)
+        )
+
+        super(TupleColumn, self).__init__(**kwargs)
+        self.null_value = tuple(x.null_value for x in nested_columns)
+
+    def write_data(self, items, buf):
+        items = self.prepare_items(items)
+        items = list(zip(*items))
+
+        for i, x in enumerate(self.nested_columns):
+            x.write_data(list(items[i]), buf)
+
+    def write_items(self, items, buf):
+        return self.write_data(items, buf)
+
+    def read_data(self, n_items, buf):
+        rv = [x.read_data(n_items, buf) for x in self.nested_columns]
+        rv = list(zip(*rv))
+
+        if self.names[0] and self.namedtuple_as_json:
+            return [dict(zip(self.names, x)) for x in rv]
+        else:
+            return rv
+
+    def read_items(self, n_items, buf):
+        return self.read_data(n_items, buf)
+
+    def read_state_prefix(self, buf):
+        super(TupleColumn, self).read_state_prefix(buf)
+
+        for x in self.nested_columns:
+            x.read_state_prefix(buf)
+
+    def write_state_prefix(self, buf):
+        super(TupleColumn, self).write_state_prefix(buf)
+
+        for x in self.nested_columns:
+            x.write_state_prefix(buf)
+
+
+def create_tuple_column(spec, column_by_spec_getter, column_options):
+    inner_spec = get_inner_spec('Tuple', spec)
+    columns_with_types = get_inner_columns_with_types(inner_spec)
+    names, types = zip(*columns_with_types)
+
+    return TupleColumn(names, [column_by_spec_getter(x) for x in types],
+                       **column_options)