[KYUUBI#2405] Support Flink StringData Data Type

externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/schema/RowSet.scala

@@ -29,6 +29,7 @@ import scala.collection.mutable.ListBuffer
 import scala.language.implicitConversions
 
 import org.apache.flink.table.catalog.Column
+import org.apache.flink.table.data.StringData
 import org.apache.flink.table.types.logical._
 import org.apache.flink.types.Row
 import org.apache.hive.service.rpc.thrift._
@@ -136,18 +137,21 @@ object RowSet {
           tDoubleValue.setValue(row.getField(ordinal).asInstanceOf[Double])
         }
         TColumnValue.doubleVal(tDoubleValue)
-      case _: VarCharType =>
+      case t @ (_: VarCharType | _: CharType) =>
         val tStringValue = new TStringValue
-        if (row.getField(ordinal) != null) {
-          val stringValue = row.getField(ordinal).asInstanceOf[String]
-          tStringValue.setValue(stringValue)
-        }
-        TColumnValue.stringVal(tStringValue)
-      case _: CharType =>
-        val tStringValue = new TStringValue
-        if (row.getField(ordinal) != null) {
-          val stringValue = row.getField(ordinal).asInstanceOf[String]
-          tStringValue.setValue(stringValue)
+        val fieldValue = row.getField(ordinal)
+        fieldValue match {
+          case value: String =>
+            tStringValue.setValue(value)
+          case _: StringData =>
+            val stringValue = fieldValue.asInstanceOf[StringData]
+            tStringValue.setValue(stringValue.toString)
+          case null =>
+            tStringValue.setValue(null)
+          case other =>
+            throw new IllegalArgumentException(
+              s"Unsupported conversion class ${other.getClass} " +
+                s"for type ${t.getClass}.")
         }
         TColumnValue.stringVal(tStringValue)
       case t =>
@@ -165,6 +169,14 @@ object RowSet {
 
   private def toTColumn(rows: Seq[Row], ordinal: Int, logicalType: LogicalType): TColumn = {
     val nulls = new java.util.BitSet()
+    // for each column, determine the conversion class by sampling the first row
+    // if there's zero row, set empty as the column
+    var sampleField: Object = null
+    var i = -1
+    while (sampleField == null && i + 1 < rows.length) {
+      i += 1
+      sampleField = rows(i).getField(ordinal)
+    }
     logicalType match {
       case _: BooleanType =>
         val values = getOrSetAsNull[lang.Boolean](rows, ordinal, nulls, true)
@@ -188,11 +200,22 @@ object RowSet {
       case _: DoubleType =>
         val values = getOrSetAsNull[lang.Double](rows, ordinal, nulls, 0.0)
         TColumn.doubleVal(new TDoubleColumn(values, nulls))
-      case _: VarCharType =>
-        val values = getOrSetAsNull[String](rows, ordinal, nulls, "")
-        TColumn.stringVal(new TStringColumn(values, nulls))
-      case _: CharType =>
-        val values = getOrSetAsNull[String](rows, ordinal, nulls, "")
+      case t @ (_: VarCharType | _: CharType) =>
+        val values: util.List[String] = new util.ArrayList[String](0)
+        sampleField match {
+          case _: String =>
+            values.addAll(getOrSetAsNull[String](rows, ordinal, nulls, ""))
+          case _: StringData =>
+            val stringDataValues =
+              getOrSetAsNull[StringData](rows, ordinal, nulls, StringData.fromString(""))
+            stringDataValues.forEach(e => values.add(e.toString))
+          case null =>
+            values.addAll(getOrSetAsNull[String](rows, ordinal, nulls, ""))
+          case other =>
+            throw new IllegalArgumentException(
+              s"Unsupported conversion class ${other.getClass} " +
+                s"for type ${t.getClass}.")
+        }
         TColumn.stringVal(new TStringColumn(values, nulls))
       case _ =>
         var i = 0

externals/kyuubi-flink-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/result/ResultSetSuite.scala

@@ -0,0 +1,52 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kyuubi.engine.flink.result
+
+import org.apache.flink.table.api.{DataTypes, ResultKind}
+import org.apache.flink.table.catalog.Column
+import org.apache.flink.table.data.StringData
+import org.apache.flink.types.Row
+
+import org.apache.kyuubi.KyuubiFunSuite
+import org.apache.kyuubi.engine.flink.schema.RowSet
+
+class ResultSetSuite extends KyuubiFunSuite {
+
+  test("StringData type conversion") {
+    val strings = List[String]("apache", "kyuubi", null)
+
+    val rowsOld: Array[Row] = strings.map(s => Row.of(s)).toArray
+    val resultSetOld = ResultSet.builder
+      .resultKind(ResultKind.SUCCESS_WITH_CONTENT)
+      .columns(Column.physical("str1", DataTypes.STRING))
+      .data(rowsOld)
+      .build
+
+    val rowsNew: Array[Row] = strings.map(s => Row.of(StringData.fromString(s))).toArray
+    val resultSetNew = ResultSet.builder
+      .resultKind(ResultKind.SUCCESS_WITH_CONTENT)
+      .columns(Column.physical("str1", DataTypes.STRING))
+      .data(rowsNew)
+      .build
+
+    assert(RowSet.toRowBaseSet(rowsNew, resultSetNew)
+      === RowSet.toRowBaseSet(rowsOld, resultSetOld))
+    assert(RowSet.toColumnBasedSet(rowsNew, resultSetNew)
+      === RowSet.toColumnBasedSet(rowsOld, resultSetOld))
+  }
+}