feat: enable native Iceberg writes after Spark required-field validation

docs/source/contributor-guide/iceberg-writes.md

@@ -121,8 +121,39 @@ conversion by calling `convertToComet`. The converted node keeps the `IcebergWri
 `originalPlan`, so AQE re-plans the write from that node.
 
 `CometIcebergNativeWrite.requiresNativeChildren` is `true`. The native writer consumes Arrow
-batches from its child over FFI, so the conversion is declined unless the child is already a Comet
-native operator. This is why a write fed by a `LocalTableScanExec` (`INSERT ... VALUES`, a local
+batches from its child over FFI. If Spark has inserted an `AssertNotNull` projection for required
+fields, `CometExecRule` keeps that projection in Spark. When row conversion is safe, it wraps the
+writer's input in `CometSparkToColumnarExec`. Any Spark sort or limit above the projection stays
+below that conversion, and Spark evaluates the assertions and their child expressions.
+
+For partitioned clustered writers, `isRowInputEvaluationSupported` also requires Spark to
+materialize each writer partition before conversion. It accepts `SortExec`, optionally wrapped
+only in `LocalLimitExec` or `GlobalLimitExec` (and planning placeholders). Unsorted or otherwise
+deferred row input keeps the JVM writer: if one row violates partition order and a later row
+violates a required-field assertion, batch read-ahead must not report the assertion first. This
+condition does not apply to unpartitioned or fanout writers.
+
+This new row conversion also keeps writes with `fixed(N)` fields, including nested fields, on the
+JVM writer even when sorted. `SparkWrite.writeSchema` retains the fixed width that Spark's
+`BinaryType` cannot enforce. A row with the wrong byte length must still fail in the JVM writer
+before a later assertion or Arrow conversion error replaces it. These two restrictions preserve
+the observable failure behavior of the newly enabled row conversion, including the exception type
+and meaningful message. They do not promise identical errors for every existing native write:
+inputs that already supply Arrow batches remain eligible for fixed fields and can report an Arrow
+cast error instead of iceberg-java's Parquet error.
+
+The conversion accepts the recursive types supported by `RowArrowReader` and requires proof that
+string outputs, including nested strings, contain valid UTF-8. `isRowConversionSupported` follows
+native inputs through preserving Spark operators and projections, including attributes, struct
+fields and construction, `If` branches, assertions, and validated string literals. It also accepts
+`MapKeys` over proven map input and `MapFromArrays` when both arrays preserve valid strings; Spark
+uses these to rebuild maps whose struct values need field alignment. Unproven row
+sources or other string-producing expressions, including string-bearing collection transforms,
+keep the JVM writer: importing invalid UTF-8 into Arrow could change the stored bytes. The writer's
+table and schema eligibility checks still apply.
+
+Other inputs must already be Comet native operators. This is why a write fed directly by a
+`LocalTableScanExec` (`INSERT ... VALUES`, a local
 DataFrame) stays on the JVM writer unless `spark.comet.exec.localTableScan.enabled` is also set.
 
 `CometIcebergWriteExec` is row-based (it emits the commit message row) over a columnar child, so

docs/source/user-guide/latest/iceberg.md

@@ -458,7 +458,7 @@ per-task Parquet write is delegated to [iceberg-rust](https://github.com/apache/
 The native writer must produce the same outcome as iceberg-java — the same Parquet features,
 statistics, and manifest metadata — so a write is only eligible when every table property it
 depends on is one the native path reproduces exactly, and additionally only when the plan
-feeding the write is fully Comet-native. For a partitioned table that plan includes the hash
+feeding the write supplies Arrow batches to Comet. For a partitioned table that plan includes the hash
 distribution and local sort Iceberg requests on its partition transforms; those stay native
 because the transforms themselves have native implementations (see
 [Iceberg system functions](#iceberg-system-functions)). Ineligible writes run through
@@ -467,7 +467,34 @@ output. A write that runs natively shows `CometIcebergWrite` under `IcebergCommi
 plan; an ineligible write keeps `IcebergWrite`.
 
 The native writer reads its input as Arrow batches from a Comet operator, so the write's input
-must itself run in Comet. A write whose input is a local relation, such as `INSERT ... VALUES` or
+must supply Arrow data. When Spark inserts `AssertNotNull` checks for nullable input written to
+required fields, those checks remain in Spark. Eligible output is converted to Arrow at the writer
+boundary, allowing native writes without a shuffle.
+
+For a partitioned clustered writer, this new conversion also requires a Spark sort to fully consume
+each writer partition before writing; only local or global limits may wrap that sort. Unsorted or
+otherwise deferred row input keeps the iceberg-java writer. For example, a row that violates
+partition order must fail before a later row's NOT NULL assertion; filling an Arrow batch first
+could reverse those errors. Unpartitioned and fanout writers do not require this additional
+condition.
+
+Writes with `fixed(N)` fields, including nested fields, keep iceberg-java when this row conversion
+would be needed, even with a sort. A row with the wrong byte length must still report the Java
+writer's validation failure before a later assertion or Arrow conversion error. These restrictions
+preserve the exception type and meaningful message for this newly enabled row conversion. They do
+not guarantee identical errors for every existing native write: writes that already supply Arrow
+batches retain their eligibility and can report an Arrow cast error instead of iceberg-java's
+Parquet error.
+
+This conversion requires proof that string values, including nested strings, contain valid UTF-8.
+For example, it accepts native input with recognized Spark projections that preserve those strings
+or use validated literals. Extracting existing map keys and rebuilding a map from arrays whose
+strings are already proven valid also qualify. Unproven sources or string-producing expressions,
+such as Spark text input, Spark binary-to-string casts, or unproven string-bearing collection
+transforms, keep the iceberg-java
+writer to preserve the original bytes. The same table and schema eligibility checks apply.
+
+A write fed directly by a local relation, such as `INSERT ... VALUES` or
 `df.writeTo(...).append()` on a DataFrame built from local data, is fed by Spark's
 `LocalTableScanExec`, which Comet only converts when `spark.comet.exec.localTableScan.enabled=true`
 (off by default). Without that setting such writes keep `IcebergWrite` and run through

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

@@ -27,6 +27,7 @@ import scala.jdk.CollectionConverters._
 import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, Expression, SortOrder}
 import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, Final, Partial, PartialMerge}
+import org.apache.spark.sql.catalyst.expressions.objects.AssertNotNull
 import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide}
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.trees.TreeNodeTag
@@ -614,7 +615,35 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false)
       // A shuffle directly under the native write stays in the child stage and is not wrapped
       // again.
       case op: IcebergWriteExec if CometConf.COMET_ICEBERG_NATIVE_WRITE_ENABLED.get(op.conf) =>
-        convertToComet(op, CometIcebergNativeWrite).getOrElse(op)
+        // Spark inserts AssertNotNull projections for nullable input to required Iceberg fields.
+        // Keep their row-wise evaluation in Spark, including any sort or limit above them, and
+        // convert the writer's input to Arrow only after Spark has evaluated those operators.
+        val needsRowConversion = !op.child.isInstanceOf[CometNativeExec] &&
+          op.child.exists {
+            case project: ProjectExec =>
+              project.projectList.exists(_.exists(_.isInstanceOf[AssertNotNull]))
+            case _ => false
+          }
+        val rowConversionReasons = ListBuffer.empty[String]
+        val prepared =
+          if (needsRowConversion &&
+            CometSparkToColumnarExec.isRowConversionSupported(op.child, rowConversionReasons) &&
+            CometIcebergNativeWrite.isRowInputEvaluationSupported(op, rowConversionReasons)) {
+            val rows = ApplyColumnarRulesAndInsertTransitions(Seq.empty, outputsColumnar = false)
+              .apply(op.child)
+            convertToComet(rows, CometSparkToColumnarExec)
+              .map(child => op.withNewChildren(Seq(child)))
+              .getOrElse(op)
+          } else {
+            rowConversionReasons.foreach(reason => withFallbackReason(op, reason))
+            op
+          }
+        convertToComet(prepared, CometIcebergNativeWrite).getOrElse {
+          prepared
+            .getTagValue(CometExplainInfo.FALLBACK_REASONS)
+            .foreach(reasons => withFallbackReasons(op, reasons))
+          op
+        }
 
       // For AQE broadcast stage on a Comet broadcast exchange
       case s @ BroadcastQueryStageExec(_, _: CometBroadcastExchangeExec, _) =>

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

@@ -22,18 +22,20 @@ package org.apache.comet.serde.operator
 import java.lang.reflect.Modifier
 import java.util.Locale
 
+import scala.collection.mutable.ListBuffer
 import scala.jdk.CollectionConverters._
 import scala.util.control.NonFatal
 
 import org.apache.hadoop.conf.Configuration
-import org.apache.spark.sql.comet.{CometIcebergWriteExec, CometNativeExec, IcebergWriteExec}
+import org.apache.spark.sql.comet.{CometIcebergWriteExec, CometNativeExec, CometScanWrapper, CometSinkPlaceHolder, IcebergWriteExec}
+import org.apache.spark.sql.execution.{GlobalLimitExec, LocalLimitExec, SortExec, SparkPlan}
 
 import org.apache.comet.{CometConf, ConfigEntry}
 import org.apache.comet.CometSparkSessionExtensions.withFallbackReason
 import org.apache.comet.iceberg.{IcebergReflection, IcebergStorageSchemes, PositionDeltaWrite, ReplaceDataWrite}
 import org.apache.comet.objectstore.NativeConfig
 import org.apache.comet.serde.{CometOperatorSerde, Compatible, OperatorOuterClass, SupportLevel, Unsupported}
-import org.apache.comet.serde.OperatorOuterClass.Operator
+import org.apache.comet.serde.OperatorOuterClass.{IcebergWriterMode, Operator}
 import org.apache.comet.serde.QueryPlanSerde.exprToProto
 
 object CometIcebergNativeWrite extends CometOperatorSerde[IcebergWriteExec] {
@@ -297,6 +299,64 @@ object CometIcebergNativeWrite extends CometOperatorSerde[IcebergWriteExec] {
         Unsupported(Some(s"Iceberg native write detection failed: ${e.getMessage}"))
     }
 
+  /**
+   * Row-to-Arrow conversion reads ahead before the writer consumes the batch. Keep fixed fields
+   * on the JVM because Spark's BinaryType does not enforce their length. For a clustered writer,
+   * materialize the writer partition so a later Spark expression failure cannot overtake an
+   * earlier partition-order error.
+   */
+  def isRowInputEvaluationSupported(
+      op: IcebergWriteExec,
+      fallbackReasons: ListBuffer[String]): Boolean = {
+    val rowInputConstraints =
+      try {
+        for {
+          sparkWrite <- IcebergReflection.getOuterSparkWrite(op.batchWrite)
+          table <- IcebergReflection.getTableFromSparkWrite(sparkWrite)
+          specId <- IcebergReflection.getOutputSpecIdFromSparkWrite(sparkWrite)
+          spec <- IcebergReflection.getPartitionSpecById(table, specId)
+          fanout <- IcebergReflection.getUseFanoutWriterFromSparkWrite(sparkWrite)
+          writeSchema <- IcebergReflection.getWriteSchemaFromSparkWrite(sparkWrite)
+        } yield {
+          val writerMode = IcebergWriteProtoTranslation.resolveWriterMode(
+            specIsUnpartitioned = isUnpartitionedSpec(spec),
+            useFanoutWriter = fanout)
+          val hasFixedFields =
+            IcebergReflection.findFieldWithTypeIds(writeSchema, Set("FIXED")).nonEmpty
+          (writerMode, hasFixedFields)
+        }
+      } catch {
+        case NonFatal(_) => None
+      }
+
+    rowInputConstraints match {
+      case None =>
+        fallbackReasons +=
+          "Iceberg writer row conversion could not resolve the effective writer constraints"
+        false
+      case Some((_, true)) =>
+        fallbackReasons += "Iceberg writer row conversion keeps fixed fields on the JVM to " +
+          "preserve fixed-width validation failures"
+        false
+      case Some((IcebergWriterMode.ICEBERG_WRITER_CLUSTERED, false))
+          if !hasMaterializedRowInput(op.child) =>
+        fallbackReasons += "Iceberg clustered writer row conversion requires materialized Spark " +
+          "input so batch read-ahead cannot change expression and partition-order error precedence"
+        false
+      case _ => true
+    }
+  }
+
+  private def hasMaterializedRowInput(plan: SparkPlan): Boolean = plan match {
+    // Planning placeholders disappear without evaluating their source.
+    case CometScanWrapper(_, source) => hasMaterializedRowInput(source)
+    case CometSinkPlaceHolder(_, _, child) => hasMaterializedRowInput(child)
+    case _: SortExec => true
+    case limit: LocalLimitExec => hasMaterializedRowInput(limit.child)
+    case limit: GlobalLimitExec => hasMaterializedRowInput(limit.child)
+    case _ => false
+  }
+
   private def checkTriggers(op: IcebergWriteExec): Option[String] = {
     op.dispatch match {
       case PositionDeltaWrite(_) =>

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

@@ -27,18 +27,21 @@ import org.apache.arrow.vector.ipc.ArrowReader
 import org.apache.spark.broadcast.Broadcast
 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.expressions.{Alias, Attribute, AttributeSet, CreateNamedStruct, Expression, GetStructField, If, Literal, MapFromArrays, MapKeys, SortOrder}
+import org.apache.spark.sql.catalyst.expressions.objects.AssertNotNull
 import org.apache.spark.sql.catalyst.plans.physical.Partitioning
 import org.apache.spark.sql.comet.execution.arrow.{CometArrowStream, CometNativeArrowSource, RowArrowReader, SparkColumnarArrowReader}
 import org.apache.spark.sql.comet.util.Utils
-import org.apache.spark.sql.execution.{RowToColumnarTransition, SparkPlan}
+import org.apache.spark.sql.execution.{CoalesceExec, CollectLimitExec, ColumnarToRowTransition, FilterExec, GlobalLimitExec, LocalLimitExec, ProjectExec, RowToColumnarTransition, SortExec, SparkPlan}
 import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
 import org.apache.spark.sql.types._
 import org.apache.spark.sql.vectorized.ColumnarBatch
+import org.apache.spark.unsafe.types.UTF8String
 
 import org.apache.comet.{CometConf, DataTypeSupport}
 import org.apache.comet.serde.OperatorOuterClass
 import org.apache.comet.serde.operator.CometSink
+import org.apache.comet.shims.CometTypeShim
 
 case class CometSparkToColumnarExec(child: SparkPlan, batchSize: Option[Int] = None)
     extends RowToColumnarTransition
@@ -131,6 +134,77 @@ case class CometSparkToColumnarExec(child: SparkPlan, batchSize: Option[Int] = N
 
 object CometSparkToColumnarExec extends CometSink[SparkPlan] with DataTypeSupport {
 
+  private object RowTypeSupport extends DataTypeSupport with CometTypeShim
+
+  // RowArrowReader recursively writes the collection types supported by DataTypeSupport.
+  def isRowSchemaSupported(schema: StructType, fallbackReasons: ListBuffer[String]): Boolean =
+    RowTypeSupport.isSchemaSupported(schema, fallbackReasons)
+
+  /**
+   * Whether the Iceberg writer may convert this Spark row subtree to Arrow. Native import
+   * normalizes invalid UTF-8, while Spark rows can retain those bytes. A new conversion is safe
+   * only when its string outputs preserve already valid strings or contain validated literals.
+   */
+  def isRowConversionSupported(plan: SparkPlan, fallbackReasons: ListBuffer[String]): Boolean = {
+    if (!isRowSchemaSupported(plan.schema, fallbackReasons)) {
+      false
+    } else if (hasValidStrings(plan)) {
+      true
+    } else {
+      fallbackReasons += "Iceberg writer row conversion cannot prove that string outputs contain " +
+        "valid UTF-8; native Arrow import could change their bytes"
+      false
+    }
+  }
+
+  private def containsString(dt: DataType): Boolean =
+    dt.existsRecursively(_.isInstanceOf[StringType])
+
+  private def hasValidStrings(plan: SparkPlan): Boolean = {
+    if (!containsString(plan.schema)) {
+      true
+    } else {
+      plan match {
+        // These planning placeholders disappear without running a native import.
+        case CometScanWrapper(_, source) => hasValidStrings(source)
+        case CometSinkPlaceHolder(_, _, child) => hasValidStrings(child)
+        case _: CometNativeExec => true
+        case project: ProjectExec =>
+          hasValidStrings(project.child) && project.projectList.forall(
+            preservesValidStrings(_, project.child.outputSet))
+        case _: SortExec | _: LocalLimitExec | _: GlobalLimitExec | _: CollectLimitExec |
+            _: FilterExec | _: CoalesceExec | _: ColumnarToRowTransition =>
+          plan.children.forall(hasValidStrings)
+        case _ => false
+      }
+    }
+  }
+
+  private def preservesValidStrings(expr: Expression, input: AttributeSet): Boolean = {
+    if (!containsString(expr.dataType)) {
+      true
+    } else {
+      expr match {
+        case attr: Attribute => input.contains(attr)
+        case alias: Alias => preservesValidStrings(alias.child, input)
+        case assertion: AssertNotNull => preservesValidStrings(assertion.child, input)
+        case field: GetStructField => preservesValidStrings(field.child, input)
+        case struct: CreateNamedStruct => struct.valExprs.forall(preservesValidStrings(_, input))
+        // Spark aligns struct-valued maps by rebuilding them with their existing keys.
+        case keys: MapKeys => preservesValidStrings(keys.child, input)
+        case map: MapFromArrays =>
+          preservesValidStrings(map.left, input) && preservesValidStrings(map.right, input)
+        case conditional: If =>
+          preservesValidStrings(conditional.trueValue, input) &&
+          preservesValidStrings(conditional.falseValue, input)
+        case Literal(null, _) => true
+        case Literal(value: UTF8String, _: StringType) => RowTypeSupport.isValidUtf8(value)
+        // String-bearing collection transforms and other producers need their own proof.
+        case _ => false
+      }
+    }
+  }
+
   override def createExec(
       nativeOp: OperatorOuterClass.Operator,
       op: SparkPlan): CometNativeExec = {