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

clickhouse_driver/columns/numpy/intcolumn.py

@@ -1,43 +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'
+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

@@ -1,83 +1,96 @@
-from math import log
-
-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):
-        # TODO: nullable support
-
-        # Do not write anything for empty column.
-        # May happen while writing empty arrays.
-        if not len(items):
-            return
-
-        c = pd.Categorical(items)
-
-        int_type = int(log(len(c.codes), 2) / 8)
-        int_column = self.int_types[int_type]()
-
-        serialization_type = self.serialization_type | int_type
-
-        index = c.categories
-        keys = c.codes
-
-        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_data(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]()
-
-        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 = keys - 1
-            index = index[1:]
-        return pd.Categorical.from_codes(keys, index)
-
-
-def create_numpy_low_cardinality_column(spec, column_by_spec_getter):
-    inner = spec[15:-1]
-    nested = column_by_spec_getter(inner)
-    return NumpyLowCardinalityColumn(nested)
+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

@@ -1,80 +1,58 @@
-from ... import errors
-from ..arraycolumn import create_array_column
-from .datecolumn import NumpyDateColumn
-from .datetimecolumn import create_numpy_datetime_column
-from ..decimalcolumn import create_decimal_column
-from ..enumcolumn import create_enum_column
-from .floatcolumn import NumpyFloat32Column, NumpyFloat64Column
-from .intcolumn import (
-    NumpyInt8Column, NumpyInt16Column, NumpyInt32Column, NumpyInt64Column,
-    NumpyUInt8Column, NumpyUInt16Column, NumpyUInt32Column, NumpyUInt64Column
-)
-from .lowcardinalitycolumn import create_numpy_low_cardinality_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 ..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 [
-    NumpyDateColumn,
-    NumpyFloat32Column, NumpyFloat64Column,
-    NumpyInt8Column, NumpyInt16Column, NumpyInt32Column, NumpyInt64Column,
-    NumpyUInt8Column, NumpyUInt16Column, NumpyUInt32Column, NumpyUInt64Column,
-    NothingColumn, NullColumn, UUIDColumn,
-    IntervalYearColumn, IntervalMonthColumn, IntervalWeekColumn,
-    IntervalDayColumn, IntervalHourColumn, IntervalMinuteColumn,
-    IntervalSecondColumn, IPv4Column, IPv6Column
-]}
-
-
-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('Enum'):
-        return create_enum_column(spec, column_options)
-
-    elif spec.startswith('DateTime'):
-        return create_numpy_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)
-
-    elif spec.startswith('Tuple'):
-        return create_tuple_column(spec, create_column_with_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)
-
-    elif spec.startswith('SimpleAggregateFunction'):
-        return create_simple_aggregate_function_column(
-            spec, create_column_with_options)
-
-    else:
-        try:
-            cls = column_by_type[spec]
-            return cls(**column_options)
-
-        except KeyError as e:
-            raise errors.UnknownTypeError('Unknown type {}'.format(e.args[0]))
+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

@@ -1,76 +1,78 @@
-import numpy as np
-
-from ... import defines
-from .base import NumpyColumn
-
-
-class NumpyStringColumn(NumpyColumn):
-    dtype = np.dtype('object')
-
-    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):
-    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)
+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)