feat: experimental zero-copy row views for writes of complex-typed data [experiment]

.github/workflows/pr_build_linux.yml

@@ -303,6 +303,7 @@ jobs:
           - name: "scans"
             value: |
               org.apache.comet.parquet.CometParquetWriterSuite
+              org.apache.comet.exec.CometWriteRowViewSuite
               org.apache.comet.parquet.ParquetReadV1Suite
               org.apache.comet.parquet.ParquetReadV2Suite
               org.apache.comet.parquet.ParquetReadFromFakeHadoopFsSuite

.github/workflows/pr_build_macos.yml

@@ -119,6 +119,7 @@ jobs:
           - name: "scans"
             value: |
               org.apache.comet.parquet.CometParquetWriterSuite
+              org.apache.comet.exec.CometWriteRowViewSuite
               org.apache.comet.parquet.ParquetReadV1Suite
               org.apache.comet.parquet.ParquetReadV2Suite
               org.apache.comet.parquet.ParquetReadFromFakeHadoopFsSuite

spark/src/main/scala/org/apache/comet/CometConf.scala

@@ -289,6 +289,19 @@ object CometConf extends ShimCometConf {
       .booleanConf
       .createWithDefault(false)
 
+  val COMET_WRITE_ROW_VIEW_ENABLED: ConfigEntry[Boolean] =
+    conf(s"$COMET_EXEC_CONFIG_PREFIX.write.rowView.enabled")
+      .category(CATEGORY_EXEC)
+      .doc(
+        "Whether to feed Spark's file write path zero-copy row views over Arrow batches " +
+          "instead of materializing an UnsafeRow per row. Nothing on the write path requires " +
+          "an UnsafeRow, so the copy is undone immediately by the writer. Only applies to " +
+          "unpartitioned, unbucketed writes of a schema containing a struct, array or map, " +
+          "where the writer is guaranteed not to retain rows and the saving is large enough " +
+          "to be worth measuring. This feature is experimental.")
+      .booleanConf
+      .createWithDefault(false)
+
   val COMET_EXEC_SORT_MERGE_JOIN_WITH_JOIN_FILTER_ENABLED: ConfigEntry[Boolean] =
     conf("spark.comet.exec.sortMergeJoinWithJoinFilter.enabled")
       .category(CATEGORY_ENABLE_EXEC)

spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala

@@ -22,12 +22,15 @@ package org.apache.comet.rules
 import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.util.sideBySide
-import org.apache.spark.sql.comet.{CometCollectLimitExec, CometColumnarToRowExec, CometMapInBatchExec, CometNativeColumnarToRowExec, CometNativeWriteExec, CometPlan, CometSparkToColumnarExec}
+import org.apache.spark.sql.comet.{CometCollectLimitExec, CometColumnarToRowExec, CometColumnarToRowViewExec, CometMapInBatchExec, CometNativeColumnarToRowExec, CometNativeWriteExec, CometPlan, CometSparkToColumnarExec}
 import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometShuffleExchangeExec}
 import org.apache.spark.sql.comet.shims.{MapInBatchInfo, ShimCometMapInBatch}
 import org.apache.spark.sql.execution.{ColumnarToRowExec, RowToColumnarExec, SparkPlan}
 import org.apache.spark.sql.execution.adaptive.QueryStageExec
+import org.apache.spark.sql.execution.command.DataWritingCommandExec
+import org.apache.spark.sql.execution.datasources.{InsertIntoHadoopFsRelationCommand, WriteFilesExec}
 import org.apache.spark.sql.execution.exchange.ReusedExchangeExec
+import org.apache.spark.sql.types.{ArrayType, DataType, MapType, StructType}
 
 import org.apache.comet.CometConf
 import org.apache.comet.CometSparkSessionExtensions.withInfo
@@ -91,6 +94,29 @@ case class EliminateRedundantTransitions(session: SparkSession)
       // Write should be final operation in the plan
       case ColumnarToRowExec(nativeWrite: CometNativeWriteExec) =>
         nativeWrite
+
+      // Spark's file writers consume `InternalRow` and never need an `UnsafeRow`, so the
+      // materializing transition below a write can be swapped for a zero-copy row view over the
+      // Arrow batch. `transformUp` has already rewritten the child into one of the Comet
+      // transitions by the time these arms are visited.
+      //
+      // `plannedWrite` (the default since Spark 3.4) puts the transition under `WriteFilesExec`;
+      // with it disabled the write command executes its child directly. Both are handled, and both
+      // are gated on the write being unpartitioned and unbucketed (`rowViewSafeForWrite`) and on
+      // the schema containing a complex type (`rowView`).
+      case w: WriteFilesExec
+          if rowViewSafeForWrite(
+            w.partitionColumns.nonEmpty,
+            w.bucketSpec.isDefined,
+            w.fileFormat.getClass.getName) =>
+        rowView(w.child).map(c => w.withNewChildren(Seq(c))).getOrElse(w)
+      case d @ DataWritingCommandExec(cmd: InsertIntoHadoopFsRelationCommand, _)
+          if rowViewSafeForWrite(
+            cmd.partitionColumns.nonEmpty || cmd.staticPartitions.nonEmpty,
+            cmd.bucketSpec.isDefined,
+            cmd.fileFormat.getClass.getName) =>
+        rowView(d.child).map(c => d.withNewChildren(Seq(c))).getOrElse(d)
+
       case c @ ColumnarToRowExec(child) if hasCometNativeChild(child) =>
         val op = createColumnarToRowExec(child)
         if (c.logicalLink.isEmpty) {
@@ -169,6 +195,64 @@ case class EliminateRedundantTransitions(session: SparkSession)
     }
   }
 
+  /**
+   * Whether a write can consume the reused, mutable rows produced by
+   * [[CometColumnarToRowViewExec]] rather than materialized `UnsafeRow`s.
+   *
+   * The row view is only correct for a consumer that finishes with a row before pulling the next
+   * one. That holds for `SingleDirectoryDataWriter`, which is what `FileFormatWriter` picks when
+   * there are no partition and no bucket columns; it writes the row straight through to the
+   * `OutputWriter`, and `BasicWriteTaskStatsTracker.newRow` ignores the row entirely. The
+   * partitioned and bucketed writers do not qualify:
+   *
+   *   - `FileFormatWriter` requires an ordering on the partition/bucket columns, so a `SortExec`
+   *     sits between this transition and the writer, and `UnsafeExternalSorter` needs
+   *     `UnsafeRow`.
+   *   - `DynamicPartitionDataConcurrentWriter` spills through `UnsafeKVExternalSorter.insertKV`,
+   *     which is typed on `UnsafeRow`.
+   *
+   * The format check keeps this to Spark's own `FileFormat` implementations. Their
+   * `OutputWriter`s encode each row on the spot (Parquet through `ParquetWriteSupport`, ORC
+   * through `OrcSerializer` into a `VectorizedRowBatch`, the text formats directly), whereas a
+   * third-party format is free to buffer the `InternalRow` it is handed.
+   */
+  private def rowViewSafeForWrite(
+      partitioned: Boolean,
+      bucketed: Boolean,
+      fileFormat: String): Boolean =
+    CometConf.COMET_WRITE_ROW_VIEW_ENABLED.get() &&
+      !partitioned && !bucketed &&
+      fileFormat.startsWith("org.apache.spark.sql.execution.datasources.")
+
+  /**
+   * Rewrites a Comet columnar-to-row transition into the zero-copy row view. Returns `None` for
+   * anything else, which leaves the plan untouched - notably for the `WriteFilesExec` under a
+   * `DataWritingCommandExec`, and for a write whose input was never columnar to begin with.
+   *
+   * Also declines a schema of nothing but flat types. There the `UnsafeProjection` this replaces
+   * is a generated fixed-width copy that measures at 0-2% of a Parquet write, which does not pay
+   * for the reused-mutable-row hazard. The saving only becomes real once a struct, array or map
+   * is in play, because then the projection has to build nested `UnsafeRow` / `UnsafeArrayData`
+   * with offset-and-length bookkeeping that `ParquetWriteSupport` immediately walks back out. See
+   * `CometParquetWriteBenchmark` for the measurements behind this cut-off.
+   */
+  private def rowView(plan: SparkPlan): Option[SparkPlan] = plan match {
+    case CometColumnarToRowExec(child) if hasComplexType(child.schema) =>
+      Some(CometColumnarToRowViewExec(child))
+    case CometNativeColumnarToRowExec(child) if hasComplexType(child.schema) =>
+      Some(CometColumnarToRowViewExec(child))
+    case _ => None
+  }
+
+  /** Whether the schema has a struct, array or map anywhere in it. */
+  private def hasComplexType(schema: StructType): Boolean = {
+    def isComplex(dataType: DataType): Boolean = dataType match {
+      case _: StructType | _: ArrayType | _: MapType => true
+      case _ => false
+    }
+    schema.fields.exists(f => isComplex(f.dataType))
+  }
+
   /**
    * If the given plan is a Comet ColumnarToRow transition, returns the columnar child the Python
    * UDF operator can consume directly. By the time this rule runs the earlier

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

@@ -0,0 +1,92 @@
+/*
+ * 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.spark.sql.comet
+
+import scala.jdk.CollectionConverters._
+
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, SortOrder}
+import org.apache.spark.sql.catalyst.plans.physical.Partitioning
+import org.apache.spark.sql.execution.{ColumnarToRowTransition, SparkPlan}
+import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
+import org.apache.spark.util.Utils
+
+/**
+ * A columnar-to-row transition that hands the consumer `ColumnarBatch.rowIterator()` directly
+ * instead of materializing an `UnsafeRow` per row.
+ *
+ * `ColumnarBatch.rowIterator()` returns a single mutable `ColumnarBatchRow` that is advanced over
+ * the batch, so each row is a zero-copy view over the underlying Arrow buffers. That makes this
+ * operator roughly free, but it is only correct for a consumer that fully consumes a row before
+ * requesting the next one and never retains a reference to it.
+ *
+ * Spark's file write path is such a consumer: `OutputWriter.write`, `FileFormatDataWriter.write`
+ * and `WriteTaskStatsTracker.newRow` all take a plain `InternalRow`, and `ParquetWriteSupport`
+ * reads fields through `SpecializedGetters` and encodes them immediately. Nothing there requires
+ * an `UnsafeRow`, so the `UnsafeProjection` performed by [[CometColumnarToRowExec]] is a copy
+ * that the writer only undoes again.
+ *
+ * This is deliberately NOT a `CodegenSupport` node: whole-stage codegen would generate an
+ * `UnsafeRowWriter` loop and reintroduce exactly the copy this operator exists to avoid.
+ *
+ * Only [[org.apache.comet.rules.EliminateRedundantTransitions]] introduces this operator, and
+ * only where it has proven the consumer is a non-retaining one. Do not use it as a general
+ * replacement for [[CometColumnarToRowExec]].
+ *
+ * @param child
+ *   The child plan that produces columnar batches
+ */
+case class CometColumnarToRowViewExec(child: SparkPlan)
+    extends ColumnarToRowTransition
+    with CometPlan {
+
+  // supportsColumnar requires to be only called on driver side, see also SPARK-37779.
+  assert(Utils.isInRunningSparkTask || child.supportsColumnar)
+
+  override def output: Seq[Attribute] = child.output
+
+  override def outputPartitioning: Partitioning = child.outputPartitioning
+
+  override def outputOrdering: Seq[SortOrder] = child.outputOrdering
+
+  override def nodeName: String = "CometColumnarToRowView"
+
+  override lazy val metrics: Map[String, SQLMetric] = Map(
+    "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows"),
+    "numInputBatches" -> SQLMetrics.createMetric(sparkContext, "number of input batches"))
+
+  override def doExecute(): RDD[InternalRow] = {
+    val numOutputRows = longMetric("numOutputRows")
+    val numInputBatches = longMetric("numInputBatches")
+    child.executeColumnar().mapPartitionsInternal { batches =>
+      batches.flatMap { batch =>
+        numInputBatches += 1
+        numOutputRows += batch.numRows()
+        // `flatMap` does not advance to the next batch until this row iterator is exhausted, so
+        // the Arrow buffers the returned rows point at stay live for as long as the rows are used.
+        batch.rowIterator().asScala
+      }
+    }
+  }
+
+  override def withNewChildInternal(newChild: SparkPlan): SparkPlan =
+    copy(child = newChild)
+}