feat(clickhouse): add enum, ipaddr, json, lowcardinality to type parser

ibis-project · Aug 19, 2022 · 8f0287f · 8f0287f
1 parent 5ea894a
commit 8f0287f
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 0 deletions.
diff --git a/ibis/backends/clickhouse/datatypes.py b/ibis/backends/clickhouse/datatypes.py
@@ -151,6 +151,32 @@ def struct():
             nullable=False,
         )
 
+    @p.generate
+    def enum_value():
+        yield dt.SPACES
+        key = yield dt.RAW_STRING
+        yield dt.spaceless_string('=')
+        value = yield p.digit.at_least(1).concat()
+        return (key, int(value))
+
+    @p.generate
+    def lowcardinality():
+        yield dt.spaceless_string('LowCardinality')
+        yield dt.LPAREN
+        r = yield ty
+        yield dt.RPAREN
+        return r
+
+    @p.generate
+    def enum():
+        yield dt.spaceless_string('enum')
+        enumsz = yield p.digit.at_least(1).concat()
+        enumsz = int(enumsz)
+        yield dt.LPAREN
+        yield enum_value.sep_by(dt.COMMA).map(dict)  # ignore values
+        yield dt.RPAREN
+        return dt.String(nullable=False)
+
     ty = (
         nullable
         | nested
@@ -160,6 +186,11 @@ def struct():
         | array
         | map
         | struct
+        | enum
+        | lowcardinality
+        | dt.spaceless_string("IPv4", "IPv6").result(dt.inet(nullable=False))
+        | dt.spaceless_string("Object('json')").result(dt.json(nullable=False))
+        | dt.spaceless_string("JSON").result(dt.json(nullable=False))
     )
     return ty.parse(text)
 

diff --git a/ibis/backends/clickhouse/tests/test_types.py b/ibis/backends/clickhouse/tests/test_types.py
@@ -30,6 +30,15 @@ def test_columns_types_with_additional_argument(con):
 @pytest.mark.parametrize(
     ('ch_type', 'ibis_type'),
     [
+        (
+            "Enum8('' = 0, 'CDMA' = 1, 'GSM' = 2, 'LTE' = 3, 'NR' = 4)",
+            dt.String(nullable=False),
+        ),
+        ('IPv4', dt.inet(nullable=False)),
+        ('IPv6', dt.inet(nullable=False)),
+        ('JSON', dt.json(nullable=False)),
+        ("Object('json')", dt.json(nullable=False)),
+        ('LowCardinality(String)', dt.String(nullable=False)),
         ('Array(Int8)', dt.Array(dt.Int8(nullable=False), nullable=False)),
         ('Array(Int16)', dt.Array(dt.Int16(nullable=False), nullable=False)),
         ('Array(Int32)', dt.Array(dt.Int32(nullable=False), nullable=False)),