Support UInt8 type in ClickHouse

trinodb · Mar 16, 2022 · 0777c46 · 0777c46
1 parent 6d1d817
commit 0777c46
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 0 deletions.
diff --git a/docs/src/main/sphinx/connector/clickhouse.rst b/docs/src/main/sphinx/connector/clickhouse.rst
@@ -167,6 +167,7 @@ ClickHouse        Trino             Notes
 ``Int16``         ``SMALLINT``      ``SMALLINT`` and ``INT2`` are aliases of ``Int16``
 ``Int32``         ``INTEGER``       ``INT``, ``INT4`` and ``INTEGER`` are aliases of ``Int32``
 ``Int64``         ``BIGINT``        ``BIGINT`` is an alias of ``Int64``
+``UInt8``         ``SMALLINT``
 ``UInt16``        ``INTEGER``
 ``UInt32``        ``BIGINT``
 ``UInt64``        ``DECIMAL(20,0)``

diff --git a/plugin/trino-clickhouse/src/main/java/io/trino/plugin/clickhouse/ClickHouseClient.java b/plugin/trino-clickhouse/src/main/java/io/trino/plugin/clickhouse/ClickHouseClient.java
@@ -21,6 +21,7 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.net.InetAddresses;
+import com.google.common.primitives.Shorts;
 import io.airlift.slice.Slice;
 import io.trino.plugin.base.aggregation.AggregateFunctionRewriter;
 import io.trino.plugin.base.aggregation.AggregateFunctionRule;
@@ -166,6 +167,9 @@ public class ClickHouseClient
 {
     private static final Splitter TABLE_PROPERTY_SPLITTER = Splitter.on(',').omitEmptyStrings().trimResults();
 
+    private static final long UINT8_MIN_VALUE = 0L;
+    private static final long UINT8_MAX_VALUE = 255L;
+
     private static final long UINT16_MIN_VALUE = 0L;
     private static final long UINT16_MAX_VALUE = 65535L;
 
@@ -487,6 +491,8 @@ public Optional<ColumnMapping> toColumnMapping(ConnectorSession session, Connect
         ClickHouseColumn column = ClickHouseColumn.of("", jdbcTypeName);
         ClickHouseDataType columnDataType = column.getDataType();
         switch (columnDataType) {
+            case UInt8:
+                return Optional.of(ColumnMapping.longMapping(SMALLINT, ResultSet::getShort, uInt8WriteFunction()));
             case UInt16:
                 return Optional.of(ColumnMapping.longMapping(INTEGER, ResultSet::getInt, uInt16WriteFunction()));
             case UInt32:
@@ -660,6 +666,17 @@ private Optional<String> formatProperty(List<String> prop)
         return Optional.of("(" + String.join(",", prop) + ")");
     }
 
+    private static LongWriteFunction uInt8WriteFunction()
+    {
+        return (statement, index, value) -> {
+            // ClickHouse stores incorrect results when the values are out of supported range.
+            if (value < UINT8_MIN_VALUE || value > UINT8_MAX_VALUE) {
+                throw new TrinoException(INVALID_ARGUMENTS, format("Value must be between %s and %s in ClickHouse: %s", UINT8_MIN_VALUE, UINT8_MAX_VALUE, value));
+            }
+            statement.setShort(index, Shorts.checkedCast(value));
+        };
+    }
+
     private static LongWriteFunction uInt16WriteFunction()
     {
         return (statement, index, value) -> {

diff --git a/.../trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseTypeMapping.java b/.../trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseTypeMapping.java
@@ -230,6 +230,42 @@ public void testUnsupportedBigint()
                 .execute(getQueryRunner(), clickhouseCreateAndInsert("tpch.test_unsupported_bigint"));
     }
 
+    @Test
+    public void testUint8()
+    {
+        SqlDataTypeTest.create()
+                .addRoundTrip("UInt8", "0", SMALLINT, "SMALLINT '0'") // min value in ClickHouse
+                .addRoundTrip("UInt8", "255", SMALLINT, "SMALLINT '255'") // max value in ClickHouse
+                .addRoundTrip("Nullable(UInt8)", "NULL", SMALLINT, "CAST(null AS SMALLINT)")
+                .execute(getQueryRunner(), clickhouseCreateAndInsert("tpch.test_uint8"));
+
+        SqlDataTypeTest.create()
+                .addRoundTrip("UInt8", "0", SMALLINT, "SMALLINT '0'") // min value in ClickHouse
+                .addRoundTrip("UInt8", "255", SMALLINT, "SMALLINT '255'") // max value in ClickHouse
+                .addRoundTrip("Nullable(UInt8)", "NULL", SMALLINT, "CAST(null AS SMALLINT)")
+                .execute(getQueryRunner(), clickhouseCreateTrinoInsert("tpch.test_uint8"));
+    }
+
+    @Test
+    public void testUnsupportedUint8()
+    {
+        // ClickHouse stores incorrect results when the values are out of supported range. This test should be fixed when ClickHouse changes the behavior.
+        SqlDataTypeTest.create()
+                .addRoundTrip("UInt8", "-1", SMALLINT, "SMALLINT '255'")
+                .addRoundTrip("UInt8", "256", SMALLINT, "SMALLINT '0'")
+                .execute(getQueryRunner(), clickhouseCreateAndInsert("tpch.test_unsupported_uint8"));
+
+        // Prevent writing incorrect results in the connector
+        try (TestTable table = new TestTable(clickhouseServer::execute, "tpch.test_unsupported_uint8", "(value UInt8) ENGINE=Log")) {
+            assertQueryFails(
+                    format("INSERT INTO %s VALUES (-1)", table.getName()),
+                    "Value must be between 0 and 255 in ClickHouse: -1");
+            assertQueryFails(
+                    format("INSERT INTO %s VALUES (256)", table.getName()),
+                    "Value must be between 0 and 255 in ClickHouse: 256");
+        }
+    }
+
     @Test
     public void testUint16()
     {