perf: reduce GC pressure in protobuf serialization (apache#3242)

Author: andygrove

spark/src/main/scala/org/apache/comet/parquet/ParquetFilters.scala

@@ -19,7 +19,6 @@
 
 package org.apache.comet.parquet
 
-import java.io.ByteArrayOutputStream
 import java.lang.{Boolean => JBoolean, Byte => JByte, Double => JDouble, Float => JFloat, Long => JLong, Short => JShort}
 import java.math.{BigDecimal => JBigDecimal}
 import java.sql.{Date, Timestamp}
@@ -43,6 +42,8 @@ import org.apache.spark.sql.sources
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.unsafe.types.UTF8String
 
+import com.google.protobuf.CodedOutputStream
+
 import org.apache.comet.parquet.SourceFilterSerde.{createBinaryExpr, createNameExpr, createUnaryExpr, createValueExpr}
 import org.apache.comet.serde.ExprOuterClass
 import org.apache.comet.serde.QueryPlanSerde.scalarFunctionExprToProto
@@ -885,10 +886,12 @@ class ParquetFilters(
 
   def createNativeFilters(predicates: Seq[sources.Filter]): Option[Array[Byte]] = {
     predicates.reduceOption(sources.And).flatMap(createNativeFilter).map { expr =>
-      val outputStream = new ByteArrayOutputStream()
-      expr.writeTo(outputStream)
-      outputStream.close()
-      outputStream.toByteArray
+      val size = expr.getSerializedSize
+      val bytes = new Array[Byte](size)
+      val codedOutput = CodedOutputStream.newInstance(bytes)
+      expr.writeTo(codedOutput)
+      codedOutput.checkNoSpaceLeft()
+      bytes
     }
   }
 

spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala

@@ -19,8 +19,6 @@
 
 package org.apache.spark.sql.comet
 
-import java.io.ByteArrayOutputStream
-
 import scala.jdk.CollectionConverters._
 
 import org.apache.hadoop.fs.Path
@@ -34,6 +32,8 @@ import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
 import org.apache.spark.sql.vectorized.ColumnarBatch
 import org.apache.spark.util.Utils
 
+import com.google.protobuf.CodedOutputStream
+
 import org.apache.comet.CometExecIterator
 import org.apache.comet.serde.OperatorOuterClass.Operator
 
@@ -75,10 +75,12 @@ case class CometNativeWriteExec(
     sparkContext.collectionAccumulator[FileCommitProtocol.TaskCommitMessage]("taskCommitMessages")
 
   override def serializedPlanOpt: SerializedPlan = {
-    val outputStream = new ByteArrayOutputStream()
-    nativeOp.writeTo(outputStream)
-    outputStream.close()
-    SerializedPlan(Some(outputStream.toByteArray))
+    val size = nativeOp.getSerializedSize
+    val bytes = new Array[Byte](size)
+    val codedOutput = CodedOutputStream.newInstance(bytes)
+    nativeOp.writeTo(codedOutput)
+    codedOutput.checkNoSpaceLeft()
+    SerializedPlan(Some(bytes))
   }
 
   override def withNewChildInternal(newChild: SparkPlan): SparkPlan =
@@ -196,10 +198,11 @@ case class CometNativeWriteExec(
 
       val nativeMetrics = CometMetricNode.fromCometPlan(this)
 
-      val outputStream = new ByteArrayOutputStream()
-      modifiedNativeOp.writeTo(outputStream)
-      outputStream.close()
-      val planBytes = outputStream.toByteArray
+      val size = modifiedNativeOp.getSerializedSize
+      val planBytes = new Array[Byte](size)
+      val codedOutput = CodedOutputStream.newInstance(planBytes)
+      modifiedNativeOp.writeTo(codedOutput)
+      codedOutput.checkNoSpaceLeft()
 
       val execIterator = new CometExecIterator(
         CometExec.newIterId,

spark/src/main/scala/org/apache/spark/sql/comet/operators.scala

@@ -19,7 +19,6 @@
 
 package org.apache.spark.sql.comet
 
-import java.io.ByteArrayOutputStream
 import java.util.Locale
 
 import scala.collection.mutable
@@ -50,6 +49,7 @@ import org.apache.spark.util.SerializableConfiguration
 import org.apache.spark.util.io.ChunkedByteBuffer
 
 import com.google.common.base.Objects
+import com.google.protobuf.CodedOutputStream
 
 import org.apache.comet.{CometConf, CometExecIterator, CometRuntimeException, ConfigEntry}
 import org.apache.comet.CometSparkSessionExtensions.{isCometShuffleEnabled, withInfo}
@@ -139,10 +139,11 @@ object CometExec {
       partitionIdx: Int,
       broadcastedHadoopConfForEncryption: Option[Broadcast[SerializableConfiguration]],
       encryptedFilePaths: Seq[String]): CometExecIterator = {
-    val outputStream = new ByteArrayOutputStream()
-    nativePlan.writeTo(outputStream)
-    outputStream.close()
-    val bytes = outputStream.toByteArray
+    val size = nativePlan.getSerializedSize
+    val bytes = new Array[Byte](size)
+    val codedOutput = CodedOutputStream.newInstance(bytes)
+    nativePlan.writeTo(codedOutput)
+    codedOutput.checkNoSpaceLeft()
     new CometExecIterator(
       newIterId,
       inputs,
@@ -414,10 +415,12 @@ abstract class CometNativeExec extends CometExec {
   def convertBlock(): CometNativeExec = {
     def transform(arg: Any): AnyRef = arg match {
       case serializedPlan: SerializedPlan if serializedPlan.isEmpty =>
-        val out = new ByteArrayOutputStream()
-        nativeOp.writeTo(out)
-        out.close()
-        SerializedPlan(Some(out.toByteArray))
+        val size = nativeOp.getSerializedSize
+        val bytes = new Array[Byte](size)
+        val codedOutput = CodedOutputStream.newInstance(bytes)
+        nativeOp.writeTo(codedOutput)
+        codedOutput.checkNoSpaceLeft()
+        SerializedPlan(Some(bytes))
       case other: AnyRef => other
       case null => null
     }