Support UInt16 type in ClickHouse

trinodb · Mar 16, 2022 · 6d1d817 · 6d1d817
1 parent 51ea8dd
commit 6d1d817
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``
+``UInt16``        ``INTEGER``
 ``UInt32``        ``BIGINT``
 ``UInt64``        ``DECIMAL(20,0)``
 ``Float32``       ``REAL``          ``FLOAT`` is an alias of ``Float32``

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
@@ -153,6 +153,7 @@
 import static java.lang.Math.floorDiv;
 import static java.lang.Math.floorMod;
 import static java.lang.Math.max;
+import static java.lang.Math.toIntExact;
 import static java.lang.String.format;
 import static java.lang.String.join;
 import static java.lang.System.arraycopy;
@@ -165,6 +166,9 @@ public class ClickHouseClient
 {
     private static final Splitter TABLE_PROPERTY_SPLITTER = Splitter.on(',').omitEmptyStrings().trimResults();
 
+    private static final long UINT16_MIN_VALUE = 0L;
+    private static final long UINT16_MAX_VALUE = 65535L;
+
     private static final long UINT32_MIN_VALUE = 0L;
     private static final long UINT32_MAX_VALUE = 4294967295L;
 
@@ -483,6 +487,8 @@ public Optional<ColumnMapping> toColumnMapping(ConnectorSession session, Connect
         ClickHouseColumn column = ClickHouseColumn.of("", jdbcTypeName);
         ClickHouseDataType columnDataType = column.getDataType();
         switch (columnDataType) {
+            case UInt16:
+                return Optional.of(ColumnMapping.longMapping(INTEGER, ResultSet::getInt, uInt16WriteFunction()));
             case UInt32:
                 return Optional.of(ColumnMapping.longMapping(BIGINT, ResultSet::getLong, uInt32WriteFunction()));
             case UInt64:
@@ -654,6 +660,17 @@ private Optional<String> formatProperty(List<String> prop)
         return Optional.of("(" + String.join(",", prop) + ")");
     }
 
+    private static LongWriteFunction uInt16WriteFunction()
+    {
+        return (statement, index, value) -> {
+            // ClickHouse stores incorrect results when the values are out of supported range.
+            if (value < UINT16_MIN_VALUE || value > UINT16_MAX_VALUE) {
+                throw new TrinoException(INVALID_ARGUMENTS, format("Value must be between %s and %s in ClickHouse: %s", UINT16_MIN_VALUE, UINT16_MAX_VALUE, value));
+            }
+            statement.setInt(index, toIntExact(value));
+        };
+    }
+
     private static LongWriteFunction uInt32WriteFunction()
     {
         return (preparedStatement, parameterIndex, 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 testUint16()
+    {
+        SqlDataTypeTest.create()
+                .addRoundTrip("UInt16", "0", INTEGER, "0") // min value in ClickHouse
+                .addRoundTrip("UInt16", "65535", INTEGER, "65535") // max value in ClickHouse
+                .addRoundTrip("Nullable(UInt16)", "NULL", INTEGER, "CAST(null AS INTEGER)")
+                .execute(getQueryRunner(), clickhouseCreateAndInsert("tpch.test_uint16"));
+
+        SqlDataTypeTest.create()
+                .addRoundTrip("UInt16", "0", INTEGER, "0") // min value in ClickHouse
+                .addRoundTrip("UInt16", "65535", INTEGER, "65535") // max value in ClickHouse
+                .addRoundTrip("Nullable(UInt16)", "NULL", INTEGER, "CAST(null AS INTEGER)")
+                .execute(getQueryRunner(), clickhouseCreateTrinoInsert("tpch.test_uint16"));
+    }
+
+    @Test
+    public void testUnsupportedUint16()
+    {
+        // 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("UInt16", "-1", INTEGER, "65535")
+                .addRoundTrip("UInt16", "65536", INTEGER, "0")
+                .execute(getQueryRunner(), clickhouseCreateAndInsert("tpch.test_unsupported_uint16"));
+
+        // Prevent writing incorrect results in the connector
+        try (TestTable table = new TestTable(clickhouseServer::execute, "tpch.test_unsupported_uint16", "(value UInt16) ENGINE=Log")) {
+            assertQueryFails(
+                    format("INSERT INTO %s VALUES (-1)", table.getName()),
+                    "Value must be between 0 and 65535 in ClickHouse: -1");
+            assertQueryFails(
+                    format("INSERT INTO %s VALUES (65536)", table.getName()),
+                    "Value must be between 0 and 65535 in ClickHouse: 65536");
+        }
+    }
+
     @Test
     public void testUint32()
     {