perf: Mark more operators as FFI safe to avoid deep copies (#3765)

* mark more operators as ffi safe

* convert CometLocalTableScanExec to extend CometSink

* scalastyle

Author: andygrove

spark/src/main/scala/org/apache/comet/serde/operator/CometSink.scala

@@ -41,7 +41,7 @@ import org.apache.comet.serde.QueryPlanSerde.{serializeDataType, supportedDataTy
 abstract class CometSink[T <: SparkPlan] extends CometOperatorSerde[T] {
 
   /** Whether the data produced by the Comet operator is FFI safe */
-  def isFfiSafe: Boolean = false
+  def isFfiSafe: Boolean = true
 
   override def enabledConfig: Option[ConfigEntry[Boolean]] = None
 
@@ -90,8 +90,6 @@ abstract class CometSink[T <: SparkPlan] extends CometOperatorSerde[T] {
 
 object CometExchangeSink extends CometSink[SparkPlan] {
 
-  override def isFfiSafe: Boolean = true
-
   override def convert(
       op: SparkPlan,
       builder: Operator.Builder,

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

@@ -279,13 +279,6 @@ case class CometBroadcastExchangeExec(
 
 object CometBroadcastExchangeExec extends CometSink[BroadcastExchangeExec] {
 
-  /**
-   * Exchange data is FFI safe because there is no use of mutable buffers involved.
-   *
-   * Source of broadcast exchange batches is ArrowStreamReader.
-   */
-  override def isFfiSafe: Boolean = true
-
   override def enabledConfig: Option[ConfigEntry[Boolean]] = Some(
     CometConf.COMET_EXEC_BROADCAST_EXCHANGE_ENABLED)
 

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

@@ -19,8 +19,6 @@
 
 package org.apache.spark.sql.comet
 
-import scala.jdk.CollectionConverters._
-
 import org.apache.spark.TaskContext
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
@@ -34,9 +32,8 @@ import org.apache.spark.sql.vectorized.ColumnarBatch
 import com.google.common.base.Objects
 
 import org.apache.comet.{CometConf, ConfigEntry}
-import org.apache.comet.serde.{CometOperatorSerde, OperatorOuterClass}
 import org.apache.comet.serde.OperatorOuterClass.Operator
-import org.apache.comet.serde.QueryPlanSerde.serializeDataType
+import org.apache.comet.serde.operator.CometSink
 
 case class CometLocalTableScanExec(
     originalPlan: LocalTableScanExec,
@@ -106,24 +103,14 @@ case class CometLocalTableScanExec(
   override def hashCode(): Int = Objects.hashCode(originalPlan, originalPlan.schema, output)
 }
 
-object CometLocalTableScanExec extends CometOperatorSerde[LocalTableScanExec] {
+object CometLocalTableScanExec extends CometSink[LocalTableScanExec] {
+
+  // uses CometArrowConverters, which re-uses arrays
+  override def isFfiSafe: Boolean = false
 
   override def enabledConfig: Option[ConfigEntry[Boolean]] = Some(
     CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED)
 
-  override def convert(
-      op: LocalTableScanExec,
-      builder: Operator.Builder,
-      childOp: Operator*): Option[Operator] = {
-    val scanTypes = op.output.flatten(attr => serializeDataType(attr.dataType))
-    val scanBuilder = OperatorOuterClass.Scan
-      .newBuilder()
-      .setSource(op.getClass.getSimpleName)
-      .addAllFields(scanTypes.asJava)
-      .setArrowFfiSafe(false)
-    Some(builder.setScan(scanBuilder).build())
-  }
-
   override def createExec(nativeOp: Operator, op: LocalTableScanExec): CometNativeExec = {
     CometScanWrapper(nativeOp, CometLocalTableScanExec(op, op.rows, op.output))
   }

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

@@ -139,6 +139,10 @@ case class CometSparkToColumnarExec(child: SparkPlan)
 }
 
 object CometSparkToColumnarExec extends CometSink[SparkPlan] with DataTypeSupport {
+
+  // uses CometArrowConverters, which re-uses arrays
+  override def isFfiSafe: Boolean = false
+
   override def createExec(
       nativeOp: OperatorOuterClass.Operator,
       op: SparkPlan): CometNativeExec = {