Map type

.travis.yml

@@ -1,4 +1,7 @@
 env:
+  - VERSION=21.3.10.1
+  - VERSION=21.2.10.48
+  - VERSION=21.1.9.41
   - VERSION=20.11.2.1
   - VERSION=20.10.2.20
   - VERSION=20.9.3.45
@@ -16,10 +19,10 @@ env:
   - VERSION=19.3.3
   - VERSION=18.12.17
   - VERSION=18.10.3
-  - VERSION=18.6.0
-  - VERSION=18.5.1
-  - VERSION=18.4.0
-  - VERSION=18.1.0
+#  - VERSION=18.6.0
+#  - VERSION=18.5.1
+#  - VERSION=18.4.0
+#  - VERSION=18.1.0
 #  - VERSION=1.1.54394
 #  - VERSION=1.1.54390
 #  - VERSION=1.1.54388

README.rst

@@ -50,6 +50,7 @@ Features
   * SimpleAggregateFunction(F, T)
   * Tuple(T1, T2, ...)
   * Nested
+  * Map(key, value)
 
 - Query progress information.
 

clickhouse_driver/columns/mapcolumn.py

@@ -0,0 +1,58 @@
+from .base import Column
+from .intcolumn import UInt64Column
+from ..util.helpers import pairwise
+
+
+class MapColumn(Column):
+    py_types = (dict, )
+
+    def __init__(self, key_column, value_column, **kwargs):
+        self.offset_column = UInt64Column()
+        self.key_column = key_column
+        self.value_column = value_column
+        super(MapColumn, self).__init__(**kwargs)
+
+    def read_state_prefix(self, buf):
+        self.key_column.read_state_prefix(buf)
+        self.value_column.read_state_prefix(buf)
+
+    def write_state_prefix(self, buf):
+        self.key_column.write_state_prefix(buf)
+        self.value_column.write_state_prefix(buf)
+
+    def read_items(self, n_items, buf):
+        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):
+    key, value = spec[4:-1].split(',')
+    key_column = column_by_spec_getter(key.strip())
+    value_column = column_by_spec_getter(value.strip())
+
+    return MapColumn(key_column, value_column)

clickhouse_driver/columns/service.py

@@ -12,6 +12,7 @@
     UInt8Column, UInt16Column, UInt32Column, UInt64Column
 )
 from .lowcardinalitycolumn import create_low_cardinality_column
+from .mapcolumn import create_map_column
 from .nothingcolumn import NothingColumn
 from .nullcolumn import NullColumn
 from .nullablecolumn import create_nullable_column
@@ -80,6 +81,9 @@ def create_column_with_options(x):
         return create_simple_aggregate_function_column(
             spec, create_column_with_options)
 
+    elif spec.startswith('Map'):
+        return create_map_column(spec, create_column_with_options)
+
     else:
         try:
             cls = column_by_type[spec]

docs/types.rst

@@ -327,3 +327,12 @@ Inserting data into nested column with ``clickhouse-driver``:
       client.execute('INSERT INTO test_nested VALUES', [
           (['a', 'b', 'c'], [100, 200, 300]),
       ])
+
+Map(key, value)
+------------------
+
+*New in version 0.2.1.*
+
+INSERT types: :class:`dict`.
+
+SELECT type: :class:`dict`.

tests/columns/test_map.py

@@ -0,0 +1,102 @@
+from tests.testcase import BaseTestCase
+
+
+class MapTestCase(BaseTestCase):
+    required_server_version = (21, 1, 2)
+    # TODO: detect stable version
+    stable_support_version = (22, 14, 9)
+
+    def client_kwargs(self, version):
+        if version < self.stable_support_version:
+            return {'settings': {'allow_experimental_map_type': True}}
+
+    def cli_client_kwargs(self):
+        if self.stable_support_version > self.server_version:
+            return {'allow_experimental_map_type': 1}
+
+    def _sorted_dicts(self, text):
+        items = [
+            ('{' + ','.join(sorted(x.strip('{}\n').split(','))) + '}')
+            for x in text.strip('\n').split('\n')
+        ]
+        return '\n'.join(items) + '\n'
+
+    def test_simple(self):
+        with self.create_table('a Map(String, UInt64)'):
+            data = [
+                ({},),
+                ({'key1': 1}, ),
+                ({'key1': 2, 'key2': 20}, ),
+                ({'key1': 3, 'key2': 30, 'key3': 50}, )
+            ]
+            self.client.execute('INSERT INTO test (a) VALUES', data)
+            query = 'SELECT * FROM test'
+            inserted = self.emit_cli(query)
+            self.assertEqual(
+                self._sorted_dicts(inserted),
+                "{}\n"
+                "{'key1':1}\n"
+                "{'key1':2,'key2':20}\n"
+                "{'key1':3,'key2':30,'key3':50}\n"
+            )
+            inserted = self.client.execute(query)
+            self.assertEqual(inserted, data)
+
+    def test_nullable(self):
+        with self.create_table('a Map(Nullable(String), Nullable(UInt64))'):
+            data = [
+                ({},),
+                ({None: None},),
+                ({'key1': 1}, )
+            ]
+            self.client.execute('INSERT INTO test (a) VALUES', data)
+            query = 'SELECT * FROM test'
+            inserted = self.emit_cli(query)
+            self.assertEqual(
+                inserted,
+                "{}\n"
+                "{NULL:NULL}\n"
+                "{'key1':1}\n"
+            )
+            inserted = self.client.execute(query)
+            self.assertEqual(inserted, data)
+
+    def test_low_cardinality(self):
+        columns = 'a Map(LowCardinality(String), LowCardinality(UInt64))'
+        with self.create_table(columns):
+            data = [
+                ({'key1': 1}, ),
+                ({'key1': 1}, ),
+                ({'key1': 1}, )
+            ]
+            self.client.execute('INSERT INTO test (a) VALUES', data)
+            query = 'SELECT * FROM test'
+            inserted = self.emit_cli(query)
+            self.assertEqual(
+                inserted,
+                "{'key1':1}\n"
+                "{'key1':1}\n"
+                "{'key1':1}\n"
+            )
+            inserted = self.client.execute(query)
+            self.assertEqual(inserted, data)
+
+    def test_array(self):
+        columns = 'a Map(String, Array(UInt64))'
+        with self.create_table(columns):
+            data = [
+                ({'key1': []}, ),
+                ({'key2': [1, 2, 3]}, ),
+                ({'key3': [1, 1, 1, 1]}, )
+            ]
+            self.client.execute('INSERT INTO test (a) VALUES', data)
+            query = 'SELECT * FROM test'
+            inserted = self.emit_cli(query)
+            self.assertEqual(
+                inserted,
+                "{'key1':[]}\n"
+                "{'key2':[1,2,3]}\n"
+                "{'key3':[1,1,1,1]}\n"
+            )
+            inserted = self.client.execute(query)
+            self.assertEqual(inserted, data)