add comment in StructField

apache · Dec 3, 2023 · 689e773 · 689e773
1 parent 52c45c9
commit 689e773
Showing 1 changed file with 21 additions and 19 deletions.
diff --git a/...i-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/schema/RowSetSuite.scala b/...i-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/schema/RowSetSuite.scala
@@ -80,24 +80,26 @@ class RowSetSuite extends KyuubiFunSuite {
   }
 
   val schemaStructFields: Seq[StructField] = Seq(
-    ("a", "boolean"),
-    ("b", "tinyint"),
-    ("c", "smallint"),
-    ("d", "int"),
-    ("e", "bigint"),
-    ("f", "float"),
-    ("g", "double"),
-    ("h", "string"),
-    ("i", "decimal"),
-    ("j", "date"),
-    ("k", "timestamp"),
-    ("l", "binary"),
-    ("m", "array<double>"),
-    ("n", "map<int, double>"),
-    ("o", "interval"),
-    ("p", "date"),
-    ("q", "timestamp")).map {
-    case (colName, typeName) => StructField(colName, CatalystSqlParser.parseDataType(typeName))
+    ("a", "boolean", "boolVal"),
+    ("b", "tinyint", "byteVal"),
+    ("c", "smallint", "shortVal"),
+    ("d", "int", "value"),
+    ("e", "bigint", "longVal"),
+    ("f", "float", "floatVal"),
+    ("g", "double", "doubleVal"),
+    ("h", "string", "stringVal"),
+    ("i", "decimal", "decimalVal"),
+    ("j", "date", "dateVal"),
+    ("k", "timestamp", "timestampVal"),
+    ("l", "binary", "binaryVal"),
+    ("m", "array<double>", "arrVal"),
+    ("n", "map<int, double>", "mapVal"),
+    ("o", "interval", "interval"),
+    ("p", "date", "localDate"),
+    ("q", "timestamp", "instant")).map {
+    case (colName, typeName, comment) =>
+      StructField(colName, CatalystSqlParser.parseDataType(typeName))
+        .withComment(comment)
   }
 
   val schema: StructType = {
@@ -298,7 +300,7 @@ class RowSetSuite extends KyuubiFunSuite {
 
     def singleColumn(field: StructField, index: Int, protocolVersion: TProtocolVersion): Unit = {
       benchmarkToTRowSet(
-        field.dataType.typeName,
+        field.getComment().getOrElse(field.dataType.typeName),
         allRows.map(row => Row(row.get(index))).asInstanceOf[Seq[Row]],
         StructType(Seq(field)),
         protocolVersion)