From 10e3dec478eba0c790afa8a9a223c67941adf2dc Mon Sep 17 00:00:00 2001 From: senmiaoliu Date: Sun, 7 Apr 2024 21:05:35 +0800 Subject: [PATCH] [KYUUBI #5374][FOLLOWUP] Fix JDBC ClickHouse TRowSet Generator MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # :mag: Description ## Issue References ๐Ÿ”— ## Describe Your Solution ๐Ÿ”ง Some data type like `UInt8` in ck can not cast to `short`, we should fix it ## Types of changes :bookmark: - [x] Bugfix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) ## Test Plan ๐Ÿงช #### Behavior Without This Pull Request :coffin: #### Behavior With This Pull Request :tada: #### Related Unit Tests --- # Checklist ๐Ÿ“ - [ ] This patch was not authored or co-authored using [Generative Tooling](https://www.apache.org/legal/generative-tooling.html) **Be nice. Be informative.** Closes #6270 from lsm1/branch-fix-jdbc-ck-short. Closes #5374 b5dac0f59 [senmiaoliu] ck fix RowSetGenerator Authored-by: senmiaoliu Signed-off-by: Cheng Pan --- .../ClickHouseTRowSetGenerator.scala | 45 +++++++++++++++++++ .../clickhouse/ClickHouseStatementSuite.scala | 22 +++++++-- .../engine/result/TColumnGenerator.scala | 21 ++++++--- .../engine/result/TColumnValueGenerator.scala | 33 +++++++++++--- 4 files changed, 106 insertions(+), 15 deletions(-) diff --git a/externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/clickhouse/ClickHouseTRowSetGenerator.scala b/externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/clickhouse/ClickHouseTRowSetGenerator.scala index fd0e24e8f..ce1761fe9 100644 --- a/externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/clickhouse/ClickHouseTRowSetGenerator.scala +++ b/externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/clickhouse/ClickHouseTRowSetGenerator.scala @@ -16,6 +16,7 @@ */ package org.apache.kyuubi.engine.jdbc.clickhouse +import java.lang.{Long => JLong, Short => JShort} import java.sql.Types.{ARRAY, OTHER} import org.apache.kyuubi.engine.jdbc.schema.DefaultJdbcTRowSetGenerator @@ -26,6 +27,39 @@ class ClickHouseTRowSetGenerator extends DefaultJdbcTRowSetGenerator { super.asByteTColumn(rows, ordinal) } + override def toSmallIntTColumn(rows: Seq[Seq[_]], ordinal: Int): TColumn = { + val colHead = if (rows.isEmpty) None else rows.head(ordinal) + colHead match { + case _: JShort => super.toSmallIntTColumn(rows, ordinal) + case _ => super.asShortTColumn( + rows, + ordinal, + convertFunc = (row, ordinal) => JShort.valueOf(row(ordinal).toString)) + } + } + + override def toIntegerTColumn(rows: Seq[Seq[_]], ordinal: Int): TColumn = { + val colHead = if (rows.isEmpty) None else rows.head(ordinal) + colHead match { + case _: Integer => super.toIntegerTColumn(rows, ordinal) + case _ => super.asIntegerTColumn( + rows, + ordinal, + convertFunc = (row, ordinal) => Integer.valueOf(row(ordinal).toString)) + } + } + + override def toBigIntTColumn(rows: Seq[Seq[_]], ordinal: Int): TColumn = { + val colHead = if (rows.isEmpty) None else rows.head(ordinal) + colHead match { + case _: JLong => super.toBigIntTColumn(rows, ordinal) + case _ => super.asLongTColumn( + rows, + ordinal, + convertFunc = (row, ordinal) => JLong.valueOf(row(ordinal).toString)) + } + } + override def toVarcharTColumn(rows: Seq[Seq[_]], ordinal: Int): TColumn = { val colHead = if (rows.isEmpty) None else rows.head(ordinal) colHead match { @@ -39,6 +73,17 @@ class ClickHouseTRowSetGenerator extends DefaultJdbcTRowSetGenerator { super.asByteTColumnValue(row, ordinal) } + override def toSmallIntTColumnValue(row: Seq[_], ordinal: Int): TColumnValue = { + asShortTColumnValue(row, ordinal, rawValue => JShort.valueOf(rawValue.toString)) + } + + override def toIntegerTColumnValue(row: Seq[_], ordinal: Int): TColumnValue = + asIntegerTColumnValue(row, ordinal, rawValue => Integer.valueOf(rawValue.toString)) + + override def toBigIntTColumnValue(row: Seq[_], ordinal: Int): TColumnValue = { + asLongTColumnValue(row, ordinal, rawValue => JLong.valueOf(rawValue.toString)) + } + override def toHiveString(data: Any, sqlType: Int): String = (data, sqlType) match { case (array: Array[_], ARRAY) => arrayToString(array) diff --git a/externals/kyuubi-jdbc-engine/src/test/scala/org/apache/kyuubi/engine/jdbc/clickhouse/ClickHouseStatementSuite.scala b/externals/kyuubi-jdbc-engine/src/test/scala/org/apache/kyuubi/engine/jdbc/clickhouse/ClickHouseStatementSuite.scala index 303614910..6051893e5 100644 --- a/externals/kyuubi-jdbc-engine/src/test/scala/org/apache/kyuubi/engine/jdbc/clickhouse/ClickHouseStatementSuite.scala +++ b/externals/kyuubi-jdbc-engine/src/test/scala/org/apache/kyuubi/engine/jdbc/clickhouse/ClickHouseStatementSuite.scala @@ -64,17 +64,25 @@ class ClickHouseStatementSuite extends WithClickHouseEngine with HiveJDBCTestHel | boolean_col boolean, | double_col double, | float_col float, - | x UUID) + | x UUID, + | ui8 UInt8, + | ui16 UInt16, + | ui32 UInt32, + | ui64 UInt64, + | ui128 UInt128, + | ui256 UInt256, + | i128 Int128, + | i256 Int256) | ENGINE=File(TabSeparated) |""".stripMargin) statement.execute( """ insert into db1.type_test | (id, tiny_col, smallint_col, int_col, bigint_col, largeint_col, decimal_col, | date_col, datetime_col, char_col, varchar_col, string_col, boolean_col, - | double_col, float_col, x) + | double_col, float_col, x, ui8, ui16, ui32, ui64, ui128, ui256, i128, i256) | VALUES (1, 2, 3, 4, 5, 6, 7.7, | '2022-05-08', '2022-05-08 17:47:45', 'a', 'Hello', 'Hello, Kyuubi', true, - | 8.8, 9.9, generateUUIDv4()) + | 8.8, 9.9, generateUUIDv4(), 8, 16, 32, 64, 128, 256, -128, -256) |""".stripMargin) val resultSet1 = statement.executeQuery("select * from db1.type_test") while (resultSet1.next()) { @@ -94,6 +102,14 @@ class ClickHouseStatementSuite extends WithClickHouseEngine with HiveJDBCTestHel assert(resultSet1.getObject(14) == 8.8) assert(resultSet1.getObject(15) == 9.9) assert(resultSet1.getString(16).length == 36) + assert(resultSet1.getObject(17) == 8) + assert(resultSet1.getObject(18) == 16) + assert(resultSet1.getObject(19) == 32) + assert(resultSet1.getObject(20) == "64") + assert(resultSet1.getObject(21) == "128") + assert(resultSet1.getObject(22) == "256") + assert(resultSet1.getObject(23) == "-128") + assert(resultSet1.getObject(24) == "-256") } } } diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/engine/result/TColumnGenerator.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/engine/result/TColumnGenerator.scala index e2c8f1ea6..7d437fc61 100644 --- a/kyuubi-common/src/main/scala/org/apache/kyuubi/engine/result/TColumnGenerator.scala +++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/engine/result/TColumnGenerator.scala @@ -62,18 +62,27 @@ trait TColumnGenerator[RowT] extends TRowSetColumnGetter[RowT] { TColumn.byteVal(new TByteColumn(values, nulls)) } - def asShortTColumn(rows: Seq[RowT], ordinal: Int): TColumn = { - val (values, nulls) = getColumnToList[JShort](rows, ordinal, 0.toShort) + def asShortTColumn( + rows: Seq[RowT], + ordinal: Int, + convertFunc: (RowT, Int) => JShort = null): TColumn = { + val (values, nulls) = getColumnToList[JShort](rows, ordinal, 0.toShort, convertFunc) TColumn.i16Val(new TI16Column(values, nulls)) } - def asIntegerTColumn(rows: Seq[RowT], ordinal: Int): TColumn = { - val (values, nulls) = getColumnToList[Integer](rows, ordinal, 0) + def asIntegerTColumn( + rows: Seq[RowT], + ordinal: Int, + convertFunc: (RowT, Int) => Integer = null): TColumn = { + val (values, nulls) = getColumnToList[Integer](rows, ordinal, 0, convertFunc) TColumn.i32Val(new TI32Column(values, nulls)) } - def asLongTColumn(rows: Seq[RowT], ordinal: Int): TColumn = { - val (values, nulls) = getColumnToList[JLong](rows, ordinal, 0.toLong) + def asLongTColumn( + rows: Seq[RowT], + ordinal: Int, + convertFunc: (RowT, Int) => JLong = null): TColumn = { + val (values, nulls) = getColumnToList[JLong](rows, ordinal, 0.toLong, convertFunc) TColumn.i64Val(new TI64Column(values, nulls)) } diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/engine/result/TColumnValueGenerator.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/engine/result/TColumnValueGenerator.scala index 0ff3a250d..032a28a18 100644 --- a/kyuubi-common/src/main/scala/org/apache/kyuubi/engine/result/TColumnValueGenerator.scala +++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/engine/result/TColumnValueGenerator.scala @@ -39,26 +39,47 @@ trait TColumnValueGenerator[RowT] extends TRowSetColumnGetter[RowT] { TColumnValue.byteVal(tValue) } - def asShortTColumnValue(row: RowT, ordinal: Int): TColumnValue = { + def asShortTColumnValue( + row: RowT, + ordinal: Int, + convertFunc: Any => JShort = null): TColumnValue = { val tValue = new TI16Value if (!isColumnNullAt(row, ordinal)) { - tValue.setValue(getColumnAs[JShort](row, ordinal)) + val short = getColumnAs[JShort](row, ordinal) match { + case sObj: JShort => sObj + case obj if convertFunc != null => convertFunc(obj) + } + tValue.setValue(short) } TColumnValue.i16Val(tValue) } - def asIntegerTColumnValue(row: RowT, ordinal: Int): TColumnValue = { + def asIntegerTColumnValue( + row: RowT, + ordinal: Int, + convertFunc: Any => Integer = null): TColumnValue = { val tValue = new TI32Value if (!isColumnNullAt(row, ordinal)) { - tValue.setValue(getColumnAs[Integer](row, ordinal)) + val integer = getColumnAs[Integer](row, ordinal) match { + case iObj: Integer => iObj + case obj if convertFunc != null => convertFunc(obj) + } + tValue.setValue(integer) } TColumnValue.i32Val(tValue) } - def asLongTColumnValue(row: RowT, ordinal: Int): TColumnValue = { + def asLongTColumnValue( + row: RowT, + ordinal: Int, + convertFunc: Any => JLong = null): TColumnValue = { val tValue = new TI64Value if (!isColumnNullAt(row, ordinal)) { - tValue.setValue(getColumnAs[JLong](row, ordinal)) + val long = getColumnAs[JLong](row, ordinal) match { + case lObj: JLong => lObj + case obj if convertFunc != null => convertFunc(obj) + } + tValue.setValue(long) } TColumnValue.i64Val(tValue) }