clickhouse-driver 0.2.9__cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl

clickhouse_driver/columns/util.py

@@ -0,0 +1,61 @@
+
+def get_inner_spec(column_name, spec):
+    brackets = 0
+    offset = len(column_name)
+    i = offset
+    for i, ch in enumerate(spec[offset:], offset):
+        if ch == '(':
+            brackets += 1
+
+        elif ch == ')':
+            brackets -= 1
+
+        if brackets == 0:
+            break
+
+    return spec[offset + 1:i]
+
+
+def get_inner_columns(spec):
+    brackets = 0
+    column_begin = 0
+
+    columns = []
+    for i, x in enumerate(spec + ','):
+        if x == ',':
+            if brackets == 0:
+                columns.append(spec[column_begin:i])
+                column_begin = i + 1
+        elif x == '(':
+            brackets += 1
+        elif x == ')':
+            brackets -= 1
+        elif x == ' ':
+            if brackets == 0:
+                column_begin = i + 1
+    return columns
+
+
+def get_inner_columns_with_types(spec):
+    spec = spec.strip()
+    brackets = 0
+    prev_comma = 0
+    prev_space = 0
+
+    columns = []
+    for i, x in enumerate(spec.strip() + ','):
+        if x == ',':
+            if brackets == 0:
+                columns.append((
+                    spec[prev_comma:prev_space].strip(),
+                    spec[prev_space:i]
+                ))
+                prev_comma = i + 1
+        elif x == '(':
+            brackets += 1
+        elif x == ')':
+            brackets -= 1
+        elif x == ' ':
+            if brackets == 0:
+                prev_space = i + 1
+    return columns

clickhouse_driver/columns/uuidcolumn.py

@@ -0,0 +1,64 @@
+from uuid import UUID
+
+from .base import FormatColumn
+from .. import errors
+from ..writer import MAX_UINT64
+
+
+class UUIDColumn(FormatColumn):
+    ch_type = 'UUID'
+    py_types = (str, UUID)
+    format = 'Q'
+
+    # UUID is stored by two uint64 numbers.
+    def write_items(self, items, buf):
+        n_items = len(items)
+
+        uint_64_pairs = [None] * 2 * n_items
+        for i, x in enumerate(items):
+            i2 = 2 * i
+            uint_64_pairs[i2] = (x >> 64) & MAX_UINT64
+            uint_64_pairs[i2 + 1] = x & MAX_UINT64
+
+        s = self.make_struct(2 * n_items)
+        buf.write(s.pack(*uint_64_pairs))
+
+    def read_items(self, n_items, buf):
+        # TODO: cythonize
+        s = self.make_struct(2 * n_items)
+        items = s.unpack(buf.read(s.size))
+
+        uint_128_items = [None] * n_items
+        for i in range(n_items):
+            i2 = 2 * i
+            uint_128_items[i] = (items[i2] << 64) + items[i2 + 1]
+
+        return tuple(uint_128_items)
+
+    def after_read_items(self, items, nulls_map=None):
+        if nulls_map is None:
+            return tuple(UUID(int=item) for item in items)
+        else:
+            return tuple(
+                (None if is_null else UUID(int=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
+
+            try:
+                if not isinstance(item, UUID):
+                    item = UUID(int=item) if item is null_value else UUID(item)
+
+            except ValueError:
+                raise errors.CannotParseUuidError(
+                    "Cannot parse uuid '{}'".format(item)
+                )
+
+            items[i] = item.int

clickhouse_driver/compression/__init__.py

@@ -0,0 +1,28 @@
+import importlib
+
+from .. import errors
+from ..protocol import CompressionMethodByte
+
+
+def get_compressor_cls(alg):
+    try:
+        module = importlib.import_module('.' + alg, __name__)
+        return module.Compressor
+
+    except ImportError:
+        raise errors.UnknownCompressionMethod(
+            "Unknown compression method: '{}'".format(alg)
+        )
+
+
+def get_decompressor_cls(method_type):
+    if method_type == CompressionMethodByte.LZ4:
+        module = importlib.import_module('.lz4', __name__)
+
+    elif method_type == CompressionMethodByte.ZSTD:
+        module = importlib.import_module('.zstd', __name__)
+
+    else:
+        raise errors.UnknownCompressionMethod()
+
+    return module.Decompressor

clickhouse_driver/compression/base.py

@@ -0,0 +1,87 @@
+from io import BytesIO
+
+from ..reader import read_binary_uint32
+from ..writer import write_binary_uint8, write_binary_uint32
+from .. import errors
+
+try:
+    from clickhouse_cityhash.cityhash import CityHash128
+except ImportError:
+    raise RuntimeError(
+        'Package clickhouse-cityhash is required to use compression'
+    )
+
+
+class BaseCompressor(object):
+    """
+    Partial file-like object with write method.
+    """
+    method = None
+    method_byte = None
+
+    def __init__(self):
+        self.data = BytesIO()
+
+        super(BaseCompressor, self).__init__()
+
+    def get_value(self):
+        value = self.data.getvalue()
+        self.data.seek(0)
+        self.data.truncate()
+        return value
+
+    def write(self, p_str):
+        self.data.write(p_str)
+
+    def compress_data(self, data):
+        raise NotImplementedError
+
+    def get_compressed_data(self, extra_header_size):
+        rv = BytesIO()
+
+        data = self.get_value()
+        compressed = self.compress_data(data)
+
+        header_size = extra_header_size + 4 + 4  # sizes
+
+        write_binary_uint32(header_size + len(compressed), rv)
+        write_binary_uint32(len(data), rv)
+        rv.write(compressed)
+
+        return rv.getvalue()
+
+
+class BaseDecompressor(object):
+    method = None
+    method_byte = None
+
+    def __init__(self, real_stream):
+        self.stream = real_stream
+        super(BaseDecompressor, self).__init__()
+
+    def decompress_data(self, data, uncompressed_size):
+        raise NotImplementedError
+
+    def check_hash(self, compressed_data, compressed_hash):
+        if CityHash128(compressed_data) != compressed_hash:
+            raise errors.ChecksumDoesntMatchError()
+
+    def get_decompressed_data(self, method_byte, compressed_hash,
+                              extra_header_size):
+        size_with_header = read_binary_uint32(self.stream)
+        compressed_size = size_with_header - extra_header_size - 4
+
+        compressed = BytesIO(self.stream.read(compressed_size))
+
+        block_check = BytesIO()
+        write_binary_uint8(method_byte, block_check)
+        write_binary_uint32(size_with_header, block_check)
+        block_check.write(compressed.getvalue())
+
+        self.check_hash(block_check.getvalue(), compressed_hash)
+
+        uncompressed_size = read_binary_uint32(compressed)
+
+        compressed = compressed.read(compressed_size - 4)
+
+        return self.decompress_data(compressed, uncompressed_size)

clickhouse_driver/compression/lz4.py

@@ -0,0 +1,21 @@
+from lz4 import block
+
+from .base import BaseCompressor, BaseDecompressor
+from ..protocol import CompressionMethod, CompressionMethodByte
+
+
+class Compressor(BaseCompressor):
+    method = CompressionMethod.LZ4
+    method_byte = CompressionMethodByte.LZ4
+    mode = 'default'
+
+    def compress_data(self, data):
+        return block.compress(data, store_size=False, mode=self.mode)
+
+
+class Decompressor(BaseDecompressor):
+    method = CompressionMethod.LZ4
+    method_byte = CompressionMethodByte.LZ4
+
+    def decompress_data(self, data, uncompressed_size):
+        return block.decompress(data, uncompressed_size=uncompressed_size)

clickhouse_driver/compression/lz4hc.py

@@ -0,0 +1,9 @@
+from .lz4 import Compressor as BaseCompressor, Decompressor as BaseDecompressor
+
+
+class Compressor(BaseCompressor):
+    mode = 'high_compression'
+
+
+class Decompressor(BaseDecompressor):
+    pass

clickhouse_driver/compression/zstd.py

@@ -0,0 +1,20 @@
+import zstd
+
+from .base import BaseCompressor, BaseDecompressor
+from ..protocol import CompressionMethod, CompressionMethodByte
+
+
+class Compressor(BaseCompressor):
+    method = CompressionMethod.ZSTD
+    method_byte = CompressionMethodByte.ZSTD
+
+    def compress_data(self, data):
+        return zstd.compress(data)
+
+
+class Decompressor(BaseDecompressor):
+    method = CompressionMethod.ZSTD
+    method_byte = CompressionMethodByte.ZSTD
+
+    def decompress_data(self, data, uncompressed_size):
+        return zstd.decompress(data)