Performance and allocations improvement in TRowSet generation

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

@@ -37,11 +37,7 @@ object RowSet {
       .buildChecked(HiveResult)
       .invokeChecked[BinaryFormatter]()
 
-  def toHiveString(
-      valueAndType: (Any, DataType),
-      nested: JBoolean = false,
-      timeFormatters: TimeFormatters,
-      binaryFormatter: BinaryFormatter): String =
+  private val toHiveStringMethod =
     DynMethods.builder("toHiveString")
       .impl( // for Spark 3.5 and before
         HiveResult.getClass,
@@ -55,5 +51,11 @@ object RowSet {
         classOf[TimeFormatters],
         classOf[BinaryFormatter])
       .buildChecked(HiveResult)
-      .invokeChecked[String](valueAndType, nested, timeFormatters, binaryFormatter)
+
+  def toHiveString(
+      valueAndType: (Any, DataType),
+      nested: JBoolean = false,
+      timeFormatters: TimeFormatters,
+      binaryFormatter: BinaryFormatter): String =
+    toHiveStringMethod.invokeChecked[String](valueAndType, nested, timeFormatters, binaryFormatter)
 }

kyuubi-common/src/main/scala/org/apache/kyuubi/engine/result/TColumnGenerator.scala

@@ -20,8 +20,6 @@ import java.lang.{Boolean => JBoolean, Byte => JByte, Double => JDouble, Float =
 import java.nio.ByteBuffer
 import java.util.{ArrayList => JArrayList, BitSet => JBitSet, List => JList}
 
-import scala.collection.JavaConverters._
-
 import org.apache.kyuubi.shaded.hive.service.rpc.thrift._
 
 trait TColumnGenerator[RowT] extends TRowSetColumnGetter[RowT] {
@@ -32,19 +30,19 @@ trait TColumnGenerator[RowT] extends TRowSetColumnGetter[RowT] {
       convertFunc: (RowT, Int) => T = null): (JList[T], ByteBuffer) = {
     val rowSize = rows.length
     val ret = new JArrayList[T](rowSize)
-    val nulls = new JBitSet()
+    val nulls = new JBitSet(rowSize)
+    val valueOf: (RowT, Int) => T =
+      if (convertFunc == null) (row, ord) => getColumnAs[T](row, ord) else convertFunc
+    val iter = rows.iterator
     var idx = 0
-    val isConvertFuncNull = convertFunc == null
-    rows.foreach { row =>
-      val value = if (isColumnNullAt(row, ordinal)) {
-        nulls.set(idx, true)
-        defaultVal
-      } else if (isConvertFuncNull) {
-        getColumnAs[T](row, ordinal)
+    while (iter.hasNext) {
+      val row = iter.next()
+      if (isColumnNullAt(row, ordinal)) {
+        nulls.set(idx)
+        ret.add(defaultVal)
       } else {
-        convertFunc(row, ordinal)
+        ret.add(valueOf(row, ordinal))
       }
-      ret.add(value)
       idx += 1
     }
     (ret, ByteBuffer.wrap(nulls.toByteArray))
@@ -85,9 +83,12 @@ trait TColumnGenerator[RowT] extends TRowSetColumnGetter[RowT] {
   }
 
   def asFloatTColumn(rows: Seq[RowT], ordinal: Int): TColumn = {
-    val (values, nulls) = getColumnToList[JFloat](rows, ordinal, 0.toFloat)
-    val doubleValues = values.asScala.map(f => JDouble.valueOf(f.toString)).asJava
-    TColumn.doubleVal(new TDoubleColumn(doubleValues, nulls))
+    val (values, nulls) = getColumnToList[JDouble](
+      rows,
+      ordinal,
+      0.toDouble,
+      (row, ord) => JDouble.valueOf(getColumnAs[JFloat](row, ord).toString))
+    TColumn.doubleVal(new TDoubleColumn(values, nulls))
   }
 
   def asDoubleTColumn(rows: Seq[RowT], ordinal: Int): TColumn = {
@@ -105,8 +106,11 @@ trait TColumnGenerator[RowT] extends TRowSetColumnGetter[RowT] {
   }
 
   def asByteArrayTColumn(rows: Seq[RowT], ordinal: Int): TColumn = {
-    val (values, nulls) = getColumnToList[Array[Byte]](rows, ordinal, defaultVal = Array[Byte]())
-    val byteBufferValues = values.asScala.map(ByteBuffer.wrap).asJava
-    TColumn.binaryVal(new TBinaryColumn(byteBufferValues, nulls))
+    val (values, nulls) = getColumnToList[ByteBuffer](
+      rows,
+      ordinal,
+      defaultVal = ByteBuffer.wrap(Array[Byte]()),
+      (row, ord) => ByteBuffer.wrap(getColumnAs[Array[Byte]](row, ord)))
+    TColumn.binaryVal(new TBinaryColumn(values, nulls))
   }
 }

kyuubi-common/src/main/scala/org/apache/kyuubi/engine/result/TRowSetGenerator.scala

@@ -18,8 +18,6 @@
 package org.apache.kyuubi.engine.result
 import java.util.{ArrayList => JArrayList}
 
-import scala.collection.JavaConverters._
-
 import org.apache.kyuubi.shaded.hive.service.rpc.thrift._
 
 trait TRowSetGenerator[SchemaT, RowT, ColumnT]
@@ -42,23 +40,24 @@ trait TRowSetGenerator[SchemaT, RowT, ColumnT]
   }
 
   def toRowBasedSet(rows: Seq[RowT], schema: SchemaT): TRowSet = {
-    val tRows = rows.map { row =>
-      var i = 0
-      val columnSize = getColumnSizeFromSchemaType(schema)
+    val columnSize = getColumnSizeFromSchemaType(schema)
+    val tRows = new JArrayList[TRow](rows.size)
+    val rowIter = rows.iterator
+    while (rowIter.hasNext) {
+      val row = rowIter.next()
       val tColumnValues = new JArrayList[TColumnValue](columnSize)
+      var i = 0
       while (i < columnSize) {
-        val columnValue = toTColumnValue(row, i, schema)
-        tColumnValues.add(columnValue)
+        tColumnValues.add(toTColumnValue(row, i, schema))
         i += 1
       }
-      new TRow(tColumnValues)
-    }.asJava
+      tRows.add(new TRow(tColumnValues))
+    }
     new TRowSet(0, tRows)
   }
 
   def toColumnBasedSet(rows: Seq[RowT], schema: SchemaT): TRowSet = {
-    val rowSize = rows.length
-    val tRowSet = new TRowSet(0, new JArrayList[TRow](rowSize))
+    val tRowSet = new TRowSet(0, new JArrayList[TRow](0))
     var i = 0
     val columnSize = getColumnSizeFromSchemaType(schema)
     val tColumns = new JArrayList[TColumn](columnSize)