support types with comma inside map

Author: Yibo-Chen13

proton_driver/columns/mapcolumn.py

@@ -1,7 +1,10 @@
+import re
 from .base import Column
 from .intcolumn import UInt64Column
 from ..util.helpers import pairwise
 
+comma_re = re.compile(r',(?![^()]*\))')
+
 
 class MapColumn(Column):
     py_types = (dict, )
@@ -51,7 +54,9 @@ def write_items(self, items, buf):
 
 
 def create_map_column(spec, column_by_spec_getter):
-    key, value = spec[4:-1].split(',')
+    # Match commas outside of parentheses, so we don't match the comma in
+    # Decimal types.
+    key, value = comma_re.split(spec[4:-1])
     key_column = column_by_spec_getter(key.strip())
     value_column = column_by_spec_getter(value.strip())
 

tests/columns/test_map.py

@@ -1,4 +1,5 @@
 from tests.testcase import BaseTestCase
+from decimal import Decimal
 
 
 class MapTestCase(BaseTestCase):
@@ -22,7 +23,7 @@ def _sorted_dicts(self, text):
         return '\n'.join(items) + '\n'
 
     def test_simple(self):
-        with self.create_stream('a map(string, uint64)'):
+        with self.create_stream('a map(string, int)'):
             data = [
                 ({},),
                 ({'key1': 1}, ),
@@ -100,3 +101,22 @@ def test_array(self):
             )
             inserted = self.client.execute(query)
             self.assertEqual(inserted, data)
+
+    def test_decimal(self):
+        with self.create_stream('a map(string, Decimal(9, 2))'):
+            data = [
+                ({'key1': Decimal('123.45')}, ),
+                ({'key2': Decimal('234.56')}, ),
+                ({'key3': Decimal('345.67')}, )
+            ]
+            self.client.execute('INSERT INTO test (a) VALUES', data)
+            query = 'SELECT * FROM test'
+            inserted = self.emit_cli(query)
+            self.assertEqual(
+                inserted,
+                "{'key1':123.45}\n"
+                "{'key2':234.56}\n"
+                "{'key3':345.67}\n"
+            )
+            inserted = self.client.execute(query)
+            self.assertEqual(inserted, data)