diff --git a/docs/source/contributor-guide/iceberg-writes.md b/docs/source/contributor-guide/iceberg-writes.md index 01da96ff8ef..b238d39b39b 100644 --- a/docs/source/contributor-guide/iceberg-writes.md +++ b/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 diff --git a/docs/source/user-guide/latest/iceberg.md b/docs/source/user-guide/latest/iceberg.md index 2822847e71a..c2f13cef189 100644 --- a/docs/source/user-guide/latest/iceberg.md +++ b/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 diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index 70d108f2506..0d6c00229a1 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/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, _) => diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala index b385668291d..6702c7a7dd4 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala +++ b/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(_) => diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometSparkToColumnarExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometSparkToColumnarExec.scala index bb80b922c34..9a02511f76d 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometSparkToColumnarExec.scala +++ b/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 = { diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala index 6f752267a63..64771fad1ca 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala @@ -35,6 +35,7 @@ import scala.util.control.NonFatal import org.json4s.{DefaultFormats, Formats} import org.json4s.jackson.JsonMethods.parse +import org.apache.hadoop.conf.Configuration import org.apache.hadoop.fs.{Path => HadoopPath} import org.apache.parquet.hadoop.ParquetFileReader import org.apache.parquet.hadoop.util.HadoopInputFile @@ -45,16 +46,18 @@ import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.DataFrame import org.apache.spark.sql.Row import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.Attribute -import org.apache.spark.sql.comet.{CometIcebergWriteExec, CometMergeRowsExec, IcebergCommitExec, IcebergSchedulerTestProbe, IcebergWriteExec} +import org.apache.spark.sql.catalyst.expressions.{Attribute, Cast} +import org.apache.spark.sql.catalyst.expressions.objects.AssertNotNull +import org.apache.spark.sql.comet.{CometIcebergWriteExec, CometMergeRowsExec, CometSparkToColumnarExec, IcebergCommitExec, IcebergSchedulerTestProbe, IcebergWriteExec} +import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec import org.apache.spark.sql.connector.catalog.InMemoryTableCatalog import org.apache.spark.sql.connector.write.{BatchWrite, DataWriterFactory, PhysicalWriteInfo, Write, WriterCommitMessage} -import org.apache.spark.sql.execution.{ColumnarToRowTransition, LeafExecNode, SparkPlan} +import org.apache.spark.sql.execution.{ColumnarToRowTransition, LeafExecNode, LocalLimitExec, ProjectExec, SortExec, SparkPlan} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.streaming.Trigger -import org.apache.spark.sql.types.{DoubleType, IntegerType, StringType, StructField, StructType} +import org.apache.spark.sql.types.{BinaryType, DoubleType, IntegerType, StringType, StructField, StructType} import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark41Plus, isSpark42Plus} import org.apache.comet.iceberg.{IcebergReflection, IcebergWriteReportListener, PositionDeltaWrite} @@ -102,6 +105,702 @@ class CometIcebergWriteActionSuite } } + Seq(false, true).foreach { adaptive => + Seq( + "id INT NOT NULL, s STRUCT", + "id INT NOT NULL, s STRUCT", + "id INT, s STRUCT NOT NULL").zipWithIndex.foreach { + case (schema, index) => + test(s"native Iceberg writer: Spark validates nullable input $index, AQE=$adaptive") { + assume(icebergAvailable && isSpark35Plus) + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString) { + withIcebergCatalog { _ => + withTempPath { path => + sql("SELECT 1 AS id, named_struct('a', 2, 'b', 'value') AS s").write.parquet( + path.toString) + val input = spark.read.parquet(path.toString) + withTempView("required_write_source") { + input.createOrReplaceTempView("required_write_source") + Seq("", "PARTITIONED BY (bucket(4, id))").zipWithIndex.foreach { + case (partition, partIndex) => + val table = s"required_input_${index}_${partIndex}_$adaptive" + val reference = s"${table}_jvm" + Seq(table, reference).foreach { name => + sql(s"""CREATE TABLE cat.db.$name ($schema) USING iceberg $partition + |TBLPROPERTIES ('write.distribution-mode'='none')""".stripMargin) + } + def appendBoth(name: String): Unit = { + sql(s"INSERT INTO cat.db.$name SELECT * FROM required_write_source") + input.writeTo(s"cat.db.$name").append() + } + val snapshot = withNativeEnabled { + captureWrite(table)(appendBoth(table)) + } + assert(snapshot.snapshotDelta == 2) + assertSparkRequiredWrite(snapshot.plans) + assert( + snapshot.plans.count(p => + collectWithSubqueries(p) { case w: CometIcebergWriteExec => + w + }.nonEmpty) == 2, + snapshot.plans.mkString("\n")) + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + appendBoth(reference) + } + checkAnswer( + sql(s"SELECT * FROM cat.db.$table"), + sql(s"SELECT * FROM cat.db.$reference")) + } + } + } + } + } + } + } + + test(s"native Iceberg writer: required input keeps the default shuffle, AQE=$adaptive") { + assume(icebergAvailable && isSpark35Plus) + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString) { + withIcebergCatalog { _ => + withTempPath { path => + sql("SELECT 1 AS id UNION ALL SELECT 2").write.parquet(path.toString) + val table = s"required_shuffle_$adaptive" + sql(s"""CREATE TABLE cat.db.$table (id INT NOT NULL) + |USING iceberg PARTITIONED BY (bucket(4, id))""".stripMargin) + val snapshot = withNativeEnabled { + captureWrite(table) { + spark.read.parquet(path.toString).writeTo(s"cat.db.$table").append() + } + } + assert(snapshot.snapshotDelta == 1) + assertSparkRequiredWrite(snapshot.plans, requireBridge = false) + assert(snapshot.plans.exists(p => + collectWithSubqueries(p) { case e: ShuffleExchangeLike => e }.nonEmpty)) + checkAnswer(sql(s"SELECT * FROM cat.db.$table"), Seq(Row(1), Row(2))) + } + } + } + } + + test(s"native Iceberg writer: required input keeps writer fallback reasons, AQE=$adaptive") { + assume(icebergAvailable && isSpark35Plus) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + EXTENDED_EXPLAIN_PROVIDERS_KEY -> "org.apache.comet.ExtendedExplainInfo") { + withIcebergCatalog { _ => + withTempPath { path => + sql("SELECT 1 AS id").write.parquet(path.toString) + withTempView("required_fallback_source") { + val input = spark.read.parquet(path.toString) + assert(input.schema("id").nullable) + input.createOrReplaceTempView("required_fallback_source") + val table = s"required_fallback_$adaptive" + sql(s"""CREATE TABLE cat.db.$table (id INT NOT NULL) USING iceberg + |TBLPROPERTIES ('write.distribution-mode'='none', + |'write.parquet.row-group-check-min-record-count'='500')""".stripMargin) + val insert = s"INSERT INTO cat.db.$table SELECT * FROM required_fallback_source" + withNativeEnabled { + val reason = "write.parquet.row-group-check-min-record-count=500" + // Spark 3.5 does not invoke extended explain providers unless it is backported. + if (supportsExtendedExplainInfo(input.queryExecution)) { + val explained = sql(s"EXPLAIN EXTENDED $insert").collect().mkString("\n") + assert(explained.contains(reason), explained) + } + val snapshot = captureWrite(table)(sql(insert)) + assertExactlyOneCommit(snapshot) + val writers = snapshot.plans.flatMap(p => + collectWithSubqueries(p) { case w: IcebergWriteExec => w }) + assert(writers.nonEmpty && writers.forall(hasRequiredProjection)) + assert(snapshot.plans.forall(p => + collectWithSubqueries(p) { case w: CometIcebergWriteExec => w }.isEmpty)) + assert( + snapshot.plans.exists(p => + new ExtendedExplainInfo().generateExtendedInfo(p).contains(reason)), + snapshot.plans.mkString("\n")) + + // Removing only the declined setting must enable the same projection's bridge. + sql(s"ALTER TABLE cat.db.$table UNSET TBLPROPERTIES " + + "('write.parquet.row-group-check-min-record-count')") + val native = captureWrite(table)(sql(insert)) + assert(native.snapshotDelta == 1) + assertSparkRequiredWrite(native.plans) + } + checkAnswer(sql(s"SELECT * FROM cat.db.$table"), Seq(Row(1), Row(1))) + } + } + } + } + } + + Seq(false, true).foreach { sorted => + test( + s"native Iceberg writer: required input with array and map, sorted=$sorted, " + + s"AQE=$adaptive") { + assume(icebergAvailable && isSpark35Plus) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + CometConf.COMET_BATCH_SIZE.key -> "3") { + withIcebergCatalog { _ => + withTempPath { path => + spark + .range(16) + .coalesce(1) + .selectExpr( + "CAST(id AS INT) AS id", + "CAST(id % 4 AS STRING) AS region", + "array(CAST(id AS INT), CAST(NULL AS INT)) AS items", + "map('value', IF(id % 2 = 0, CAST(id AS INT), CAST(NULL AS INT))) AS attrs", + "array(named_struct('v', CAST(id AS INT)), " + + "CAST(NULL AS STRUCT), named_struct('v', CAST(NULL AS INT))) " + + "AS struct_items", + "array(array(CAST(id AS INT), CAST(NULL AS INT)), " + + "CAST(array() AS ARRAY), CAST(NULL AS ARRAY)) AS nested_items", + "map('present', named_struct('v', CAST(id AS INT)), " + + "'null_struct', CAST(NULL AS STRUCT), " + + "'null_field', named_struct('v', CAST(NULL AS INT))) AS struct_values", + "map('present', array(CAST(id AS INT), CAST(NULL AS INT)), " + + "'empty', CAST(array() AS ARRAY), " + + "'null_array', CAST(NULL AS ARRAY)) AS array_values") + .write + .parquet(path.toString) + val partition = if (sorted) "PARTITIONED BY (region)" else "" + val nativeTable = s"required_complex_${sorted}_$adaptive" + val reference = s"${nativeTable}_jvm" + Seq(nativeTable, reference).foreach { table => + sql(s"""CREATE TABLE cat.db.$table + |(id INT NOT NULL, region STRING, items ARRAY, attrs MAP, + |struct_items ARRAY>, nested_items ARRAY>, + |struct_values MAP>, array_values MAP>) + |USING iceberg $partition TBLPROPERTIES + |('write.distribution-mode'='none', 'write.spark.fanout.enabled'='false') + |""".stripMargin) + if (sorted) sql(s"ALTER TABLE cat.db.$table WRITE ORDERED BY id") + } + val input = spark.read.parquet(path.toString) + assert(input.schema.fields.forall(_.nullable)) + val snapshot = withNativeEnabled { + captureWrite(nativeTable) { + input + .writeTo(s"cat.db.$nativeTable") + .option("distribution-mode", "none") + .append() + } + } + assert(snapshot.snapshotDelta == 1) + assertSparkRequiredWrite(snapshot.plans) + assert(snapshot.plans.forall(p => + collectWithSubqueries(p) { case e: ShuffleExchangeLike => e }.isEmpty)) + if (sorted) { + val writers = snapshot.plans.flatMap(p => + collectWithSubqueries(p) { case w: CometIcebergWriteExec => w }) + assert(writers.forall(w => + collectWithSubqueries(w) { + case bridge: CometSparkToColumnarExec if collectWithSubqueries(bridge.child) { + case sort: SortExec if hasRequiredProjection(sort.child) => sort + }.nonEmpty => + bridge + }.nonEmpty)) + } + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + input.writeTo(s"cat.db.$reference").option("distribution-mode", "none").append() + checkAnswer(sql(s"SELECT * FROM cat.db.$nativeTable"), input) + checkAnswer( + sql(s"SELECT * FROM cat.db.$nativeTable"), + sql(s"SELECT * FROM cat.db.$reference")) + } + } + } + } + } + } + + Seq( + ( + "child NULL before division by zero", + "SELECT CAST(NULL AS BIGINT) AS den UNION ALL SELECT 0L", + "id BIGINT NOT NULL", + Seq("10L DIV den AS id")), + ( + "earlier field before later division by zero", + "SELECT CAST(NULL AS BIGINT) AS a, 0L AS den", + "a BIGINT NOT NULL, b BIGINT NOT NULL", + Seq("a", "10L DIV den AS b"))).zipWithIndex.foreach { + case ((description, source, schema, projection), index) => + test(s"native Iceberg writer: Spark preserves $description, AQE=$adaptive") { + assume(icebergAvailable && isSpark35Plus) + withSQLConf( + SQLConf.ANSI_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString) { + withIcebergCatalog { _ => + withTempPath { path => + sql(source).coalesce(1).write.parquet(path.toString) + val nativeTable = s"required_error_native_${index}_$adaptive" + val jvmTable = s"required_error_jvm_${index}_$adaptive" + Seq(nativeTable, jvmTable).foreach { table => + sql(s"""CREATE TABLE cat.db.$table ($schema) USING iceberg + |TBLPROPERTIES ('write.distribution-mode'='none')""".stripMargin) + } + val (nativeError, jvmError) = + failedWrites(nativeTable, jvmTable, checkSparkAssertions = true) { table => + withSQLConf(CometConf.COMET_ENABLED.key -> (table == nativeTable).toString) { + spark.read + .parquet(path.toString) + .selectExpr(projection: _*) + .writeTo(s"cat.db.$table") + .append() + } + } + val expected = requiredFieldError(jvmError) + val actual = requiredFieldError(nativeError) + assert(actual.getClass == expected.getClass) + assert(actual.getMessage == expected.getMessage) + Seq(nativeTable, jvmTable).foreach { table => + assert(countSnapshots(table) == 0) + assert(parquetFiles(dataDir(table)).isEmpty) + checkAnswer(sql(s"SELECT * FROM cat.db.$table"), Seq.empty[Row]) + } + } + } + } + } + } + + test(s"native Iceberg writer: clustered input error precedes a later null, AQE=$adaptive") { + assume(icebergAvailable && isSpark35Plus) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + CometConf.COMET_BATCH_SIZE.key -> "8") { + withIcebergCatalog { _ => + withTempPath { path => + val expectedRows = Seq(Row(1, 0), Row(2, 1), Row(3, 0), Row(null, 2)) + val schema = StructType( + Seq( + StructField("id", IntegerType, nullable = true), + StructField("p", IntegerType, nullable = false))) + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark + .createDataFrame(spark.sparkContext.parallelize(expectedRows, 1), schema) + .write + .parquet(path.toString) + } + val rows = spark.read.parquet(path.toString) + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + assert(rows.collect().toSeq == expectedRows) + } + val candidate = s"required_unclustered_native_$adaptive" + val reference = s"required_unclustered_jvm_$adaptive" + val control = s"required_unclustered_control_$adaptive" + Seq(candidate, reference, control).foreach { table => + sql(s"""CREATE TABLE cat.db.$table (id INT NOT NULL, p INT) + |USING iceberg PARTITIONED BY (p)""".stripMargin) + } + val successful = withNativeEnabled { + captureWrite(control)(appendUnclustered(control, rows.where("id <= 2"))) + } + assert(successful.snapshotDelta == 1) + checkAnswer(sql(s"SELECT * FROM cat.db.$control"), expectedRows.take(2)) + + // Row three revisits a closed partition before row four evaluates AssertNotNull. + // Capture failures as in failedWrites, but the safe route here is the JVM writer. + var candidateFailure: Throwable = null + val candidatePlans = withNativeEnabled { + capturePlans(spark, includeFailures = true) { + candidateFailure = intercept[Exception](appendUnclustered(candidate, rows)) + } + } + var referenceFailure: Throwable = null + var referencePlans: Seq[SparkPlan] = Seq.empty + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + referencePlans = capturePlans(spark, includeFailures = true) { + referenceFailure = intercept[Exception](appendUnclustered(reference, rows)) + } + } + val plans = candidatePlans ++ referencePlans + plans.foreach(assertColumnarContract) + assert( + plans.forall(p => + collectWithSubqueries(p) { + case s: SortExec => s + case e: ShuffleExchangeLike => e + }.isEmpty), + plans.mkString("\n")) + assert(referencePlans.exists(hasRequiredProjection), referencePlans.mkString("\n")) + Seq(candidate, reference).foreach { table => + assert(countSnapshots(table) == 0) + assert(parquetFiles(dataDir(table)).isEmpty) + checkAnswer(sql(s"SELECT * FROM cat.db.$table"), Seq.empty[Row]) + } + val expected = clusteredWriterError(referenceFailure) + assert(expected._2.startsWith( + "Incoming records violate the writer assumption that records are clustered by spec")) + val actual = withClue(s"Candidate plans:\n${candidatePlans.mkString("\n")}\n") { + clusteredWriterError(candidateFailure) + } + assert(actual == expected, s"candidate $actual != jvm $expected") + assert( + causeChain(candidateFailure)(actual._1).getClass == + causeChain(referenceFailure)(expected._1).getClass) + val writers = candidatePlans.flatMap(p => + collectWithSubqueries(p) { case w: IcebergWriteExec => w }) + assert(writers.nonEmpty, candidatePlans.mkString("\n")) + assert(writers.forall(hasRequiredProjection), candidatePlans.mkString("\n")) + assert( + candidatePlans.forall(p => + collectWithSubqueries(p) { case w: CometIcebergWriteExec => + w + }.isEmpty), + candidatePlans.mkString("\n")) + } + } + } + } + + test(s"native Iceberg writer: Spark assertion aborts after a valid prefix, AQE=$adaptive") { + assume(icebergAvailable && isSpark35Plus) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + CometConf.COMET_BATCH_SIZE.key -> "16") { + withIcebergCatalog { _ => + withTempPath { path => + // Cross the writer's 1,000-row file-size check before the assertion fails. + spark + .range(0, 2048, 1, 1) + .selectExpr("id AS a", "IF(id=2047, CAST(NULL AS BIGINT), id) AS b") + .write + .parquet(path.toString) + val tables = Seq(s"prefix_native_$adaptive", s"prefix_jvm_$adaptive") + tables.foreach { table => + sql(s"""CREATE TABLE cat.db.$table (a BIGINT NOT NULL, b BIGINT NOT NULL) + |USING iceberg TBLPROPERTIES + |('write.distribution-mode'='none', 'write.target-file-size-bytes'='1', + |'write.parquet.page-row-limit'='1000')""".stripMargin) + sql(s"INSERT INTO cat.db.$table VALUES (-1L, -1L)") + } + val snapshots = tables.map(table => table -> countSnapshots(table)).toMap + val committed = tables.map { table => + val root = dataDir(table).toPath + val files = IcebergTestFiles.regularFiles(root) + assert(files.nonEmpty, "seed write did not create a data file") + table -> files + .map(file => file -> Files.readAllBytes(root.resolve(file)).toSeq) + .toMap + }.toMap + def referencedFiles(table: String): Set[String] = + sql(s"SELECT file_path FROM cat.db.$table.files") + .collect() + .map(_.getString(0)) + .toSet + val references = tables.map(table => table -> referencedFiles(table)).toMap + val (nativeError, jvmError) = + failedWrites(tables.head, tables.last, checkSparkAssertions = true) { table => + withSQLConf(CometConf.COMET_ENABLED.key -> (table == tables.head).toString) { + val location = dataDir(table).getAbsolutePath + val existing = committed(table).keySet + spark.udf.register( + "verify_required_write_prefix", + (value: java.lang.Long) => { + if (value == null) { + val created = IcebergTestFiles.parquetFiles(new File(location).toPath) -- + existing + // Java opens the next physical file lazily; verify a completed footer. + val finalized = created.exists { file => + try { + val reader = ParquetFileReader.open( + HadoopInputFile.fromPath( + new HadoopPath(new File(location, file).toURI), + new Configuration())) + try reader.getFooter.getBlocks.asScala.exists(_.getRowCount > 0) + finally reader.close() + } catch { + case NonFatal(_) => false // The current file may still be open. + } + } + require(finalized, s"no finalized prefix before assertion: $created") + } + value + }) + spark.read + .parquet(path.toString) + .selectExpr("a", "verify_required_write_prefix(b) AS b") + .writeTo(s"cat.db.$table") + .append() + } + } + val expected = requiredFieldError(jvmError) + val actual = requiredFieldError(nativeError) + assert(actual.getClass == expected.getClass) + assert(actual.getMessage == expected.getMessage) + tables.foreach { table => + assert(countSnapshots(table) == snapshots(table)) + val root = dataDir(table).toPath + assert(IcebergTestFiles.regularFiles(root) == committed(table).keySet) + assert(referencedFiles(table) == references(table)) + committed(table).foreach { case (file, bytes) => + assert(Files.readAllBytes(root.resolve(file)).toSeq == bytes) + } + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + checkAnswer(sql(s"SELECT * FROM cat.db.$table"), Seq(Row(-1L, -1L))) + } + } + } + } + } + } + + test(s"native Iceberg writer: Spark validates only the limited input, AQE=$adaptive") { + assume(icebergAvailable && isSpark35Plus) + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString) { + withIcebergCatalog { _ => + withTempPath { path => + sql("SELECT 1 AS id UNION ALL SELECT CAST(NULL AS INT)") + .coalesce(1) + .write + .parquet(path.toString) + val tableName = s"assert_limited_input_$adaptive" + sql(s"CREATE TABLE cat.db.$tableName (id INT NOT NULL) USING iceberg") + val snapshot = withNativeEnabled { + captureWrite(tableName) { + spark.read.parquet(path.toString).limit(1).writeTo(s"cat.db.$tableName").append() + } + } + assert(snapshot.snapshotDelta == 1) + assertSparkRequiredWrite(snapshot.plans, requireBridge = false) + val writers = snapshot.plans.flatMap(p => + collectWithSubqueries(p) { case w: CometIcebergWriteExec => w }) + assert(writers.forall(w => + collectWithSubqueries(w) { + case exchange: CometShuffleExchangeExec + if exchange.nodeName == "CometColumnarExchange" && + collectWithSubqueries(exchange.child) { + case limit: LocalLimitExec if hasRequiredProjection(limit.child) => limit + }.nonEmpty => + exchange + }.nonEmpty)) + checkAnswer(sql(s"SELECT * FROM cat.db.$tableName"), Seq(Row(1))) + } + } + } + } + } + + Seq(false, true).foreach { adaptive => + test( + s"native Iceberg writer: Arrow bridge respects a single-partition limit, AQE=$adaptive") { + assume(icebergAvailable && isSpark35Plus) + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString) { + withIcebergCatalog { _ => + val table = s"required_rdd_limit_$adaptive" + sql(s"""CREATE TABLE cat.db.$table (id INT NOT NULL) USING iceberg + |TBLPROPERTIES ('write.distribution-mode'='none')""".stripMargin) + val schema = StructType(Seq(StructField("id", IntegerType, nullable = true))) + val input = spark.createDataFrame( + spark.sparkContext.parallelize(Seq(Row(1), Row(null)), 1), + schema) + val snapshot = withNativeEnabled { + captureWrite(table) { + input.coalesce(1).limit(1).writeTo(s"cat.db.$table").append() + } + } + assert(snapshot.snapshotDelta == 1) + assertSparkRequiredWrite(snapshot.plans) + assert(snapshot.plans.forall(p => + collectWithSubqueries(p) { case e: ShuffleExchangeLike => e }.isEmpty)) + val writers = snapshot.plans.flatMap(p => + collectWithSubqueries(p) { case w: CometIcebergWriteExec => w }) + assert(writers.forall(w => + collectWithSubqueries(w) { + case bridge: CometSparkToColumnarExec if collectWithSubqueries(bridge.child) { + case limit: LocalLimitExec if hasRequiredProjection(limit.child) => limit + }.nonEmpty => + bridge + }.nonEmpty)) + checkAnswer(sql(s"SELECT * FROM cat.db.$table"), Seq(Row(1))) + } + } + } + } + + Seq("text", "binary_cast", "binary_map_key", "binary_map_value").foreach { source => + test(s"native Iceberg writer: unproven $source strings preserve invalid UTF-8 on Spark") { + assume(icebergAvailable && isSpark35Plus) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_TEXT_ENABLED.key -> "false", + "spark.comet.expression.Cast.enabled" -> "false") { + withIcebergCatalog { _ => + withTempPath { path => + val table = s"required_utf8_$source" + val reference = s"${table}_jvm" + val schema = source match { + case "text" => "value STRING NOT NULL" + case "binary_cast" => "id INT NOT NULL, value STRING" + case _ => "id INT NOT NULL, value MAP" + } + Seq(table, reference).foreach { name => + sql(s"""CREATE TABLE cat.db.$name ($schema) USING iceberg + |TBLPROPERTIES ('write.distribution-mode'='none')""".stripMargin) + } + val input = if (source == "text") { + Files.write(path.toPath, Array(0xff.toByte, '\n'.toByte)) + spark.read.text(path.toString) + } else { + sql("SELECT 1 AS id, X'FF' AS bytes, array('valid') AS valid").write.parquet( + path.toString) + val value = source match { + case "binary_map_key" => + "map_from_arrays(array(CAST(bytes AS STRING)), valid)" + case "binary_map_value" => + "map_from_arrays(valid, array(CAST(bytes AS STRING)))" + case _ => "CAST(bytes AS STRING)" + } + spark.read.parquet(path.toString).selectExpr("id", s"$value AS value") + } + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + input.writeTo(s"cat.db.$reference").append() + } + val snapshot = withNativeEnabled { + captureWrite(table) { + input.writeTo(s"cat.db.$table").append() + } + } + assert(snapshot.snapshotDelta == 1) + // Read bytes through Spark too: a native reader could normalize both results. + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + val value = source match { + case "binary_map_key" => "map_keys(value)[0]" + case "binary_map_value" => "map_values(value)[0]" + case _ => "value" + } + val expected = sql(s"SELECT hex(CAST($value AS BINARY)) FROM cat.db.$reference") + val actual = sql(s"SELECT hex(CAST($value AS BINARY)) FROM cat.db.$table") + checkAnswer(expected, Seq(Row("FF"))) + checkAnswer(actual, expected) + } + val writers = snapshot.plans.flatMap(p => + collectWithSubqueries(p) { case w: IcebergWriteExec => w }) + assert(writers.nonEmpty, snapshot.plans.mkString("\n")) + assert(writers.forall(hasRequiredProjection), snapshot.plans.mkString("\n")) + assert( + snapshot.plans.forall(p => + collectWithSubqueries(p) { case w: CometIcebergWriteExec => w }.isEmpty), + snapshot.plans.mkString("\n")) + if (source != "text") { + assert( + writers.exists(w => + collectWithSubqueries(w) { + case p: ProjectExec if p.projectList.exists(_.exists { + case c: Cast => + c.child.dataType == BinaryType && c.dataType == StringType + case _ => false + }) => + p + }.nonEmpty), + snapshot.plans.mkString("\n")) + } + } + } + } + } + } + + test("native Iceberg writer: Spark validates scalar struct fields") { + assume(icebergAvailable && isSpark35Plus) + withIcebergCatalog { _ => + withTempPath { path => + sql("SELECT 1 AS id UNION ALL SELECT 2 UNION ALL SELECT 3") + .coalesce(1) + .write + .parquet(path.toString) + withTempView("assertion_scalar_source") { + spark.read.parquet(path.toString).createOrReplaceTempView("assertion_scalar_source") + sql("""CREATE TABLE cat.db.assert_scalar_struct + |(id INT, s STRUCT NOT NULL) USING iceberg + |TBLPROPERTIES ('write.distribution-mode'='none')""".stripMargin) + val snapshot = withNativeEnabled { + captureWrite("assert_scalar_struct") { + sql("""INSERT INTO cat.db.assert_scalar_struct + |SELECT id, CASE WHEN (SELECT max(id) FROM assertion_scalar_source) > 0 + |THEN named_struct('a', (SELECT max(id) FROM assertion_scalar_source)) + |ELSE CAST(NULL AS STRUCT) END AS s + |FROM assertion_scalar_source""".stripMargin) + } + } + assert(snapshot.snapshotDelta == 1) + assertSparkRequiredWrite(snapshot.plans) + checkAnswer( + sql("SELECT * FROM cat.db.assert_scalar_struct"), + Seq(Row(1, Row(3)), Row(2, Row(3)), Row(3, Row(3)))) + } + } + } + } + + test("native Iceberg writer: Spark allows a null parent of a required nested field") { + assume(icebergAvailable && isSpark35Plus) + withIcebergCatalog { _ => + withTempPath { path => + sql("""SELECT 1 AS id, named_struct('a', 2, 'b', 'value') AS s + |UNION ALL SELECT 2, CAST(NULL AS STRUCT)""".stripMargin) + .coalesce(1) + .write + .parquet(path.toString) + sql("""CREATE TABLE cat.db.assert_null_parent + |(id INT NOT NULL, s STRUCT) USING iceberg + |TBLPROPERTIES ('write.distribution-mode'='none')""".stripMargin) + val snapshot = withNativeEnabled { + captureWrite("assert_null_parent") { + spark.read.parquet(path.toString).writeTo("cat.db.assert_null_parent").append() + } + } + assert(snapshot.snapshotDelta == 1) + assertSparkRequiredWrite(snapshot.plans) + checkAnswer( + sql("SELECT * FROM cat.db.assert_null_parent"), + Seq(Row(1, Row(2, "value")), Row(2, null))) + } + } + } + + Seq( + ("id INT NOT NULL", "CAST(NULL AS INT) AS id"), + ( + "s STRUCT", + "named_struct('a', CAST(NULL AS INT), 'b', 'value') AS s"), + ( + "s STRUCT NOT NULL", + "CAST(NULL AS STRUCT) AS s")).zipWithIndex.foreach { + case ((schema, value), index) => + test(s"native Iceberg writer: Spark rejects null input without a commit $index") { + assume(icebergAvailable && isSpark35Plus) + withIcebergCatalog { _ => + withTempPath { path => + sql(s"SELECT $value").write.parquet(path.toString) + val nativeTable = s"assert_null_native_$index" + val jvmTable = s"assert_null_jvm_$index" + Seq(nativeTable, jvmTable).foreach { table => + sql(s"""CREATE TABLE cat.db.$table ($schema) USING iceberg + |TBLPROPERTIES ('write.distribution-mode'='none')""".stripMargin) + } + val (nativeError, jvmError) = + failedWrites(nativeTable, jvmTable, checkSparkAssertions = true) { table => + withSQLConf(CometConf.COMET_ENABLED.key -> (table == nativeTable).toString) { + spark.read.parquet(path.toString).writeTo(s"cat.db.$table").append() + } + } + val actual = requiredFieldError(nativeError) + val expected = requiredFieldError(jvmError) + assert(actual.getClass == expected.getClass) + assert(actual.getMessage == expected.getMessage) + Seq(nativeTable, jvmTable).foreach { table => + assert(countSnapshots(table) == 0) + assert(parquetFiles(dataDir(table)).isEmpty) + checkAnswer(sql(s"SELECT * FROM cat.db.$table"), Seq.empty[Row]) + } + } + } + } + } + // Turning off Comet, or only its native execution as an application that uses Comet just for // scans or shuffle does, keeps Spark's own write operator. Seq(CometConf.COMET_ENABLED, CometConf.COMET_EXEC_ENABLED).foreach { flag => @@ -3980,6 +4679,134 @@ class CometIcebergWriteActionSuite } } + for { + adaptive <- Seq(false, true) + laterNull <- Seq(false, true) + } { + test( + "native Iceberg writer: fixed-length validation preserves the first error, " + + s"laterNull=$laterNull, AQE=$adaptive") { + assume(icebergAvailable && isSpark35Plus) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + CometConf.COMET_BATCH_SIZE.key -> "8") { + withIcebergCatalog { _ => + withTempPath { path => + val schema = StructType( + Seq( + StructField("id", IntegerType, nullable = true), + StructField("fx", BinaryType, nullable = true))) + val invalidRows = Seq(Row(1, Array[Byte](1, 2, 3))) ++ + (if (laterNull) Seq(Row(null, Array[Byte](5, 6, 7, 8))) else Seq.empty) + val invalidPath = new File(path, "invalid").toString + val validPath = new File(path, "valid").toString + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark + .createDataFrame(spark.sparkContext.parallelize(invalidRows, 1), schema) + .write + .parquet(invalidPath) + spark + .createDataFrame( + spark.sparkContext.parallelize( + Seq(Row(1, Array[Byte](1, 2, 3, 4)), Row(2, Array[Byte](5, 6, 7, 8))), + 1), + schema) + .write + .parquet(validPath) + val expectedInput = Seq(Row(1, "010203")) ++ + (if (laterNull) Seq(Row(null, "05060708")) else Seq.empty) + assert( + spark.read.parquet(invalidPath).selectExpr("id", "hex(fx)").collect().toSeq == + expectedInput) + } + val suffix = s"${laterNull}_$adaptive" + val candidate = s"required_fixed_native_$suffix" + val reference = s"required_fixed_jvm_$suffix" + val nativeControl = s"required_fixed_control_native_$suffix" + val jvmControl = s"required_fixed_control_jvm_$suffix" + Seq(candidate, reference, nativeControl, jvmControl).foreach { table => + sql(s"""CREATE TABLE $catalog.$ns.$table (id INT NOT NULL) USING iceberg + |TBLPROPERTIES ('write.distribution-mode'='none', + |'write.parquet.dict-size-bytes'='1')""".stripMargin) + addIcebergColumn( + loadIcebergTable(spark, catalog, ns, table), + "fx", + icebergFixedType(4)) + sql(s"REFRESH TABLE $catalog.$ns.$table") + } + def assertJvmRequiredWrite(plans: Seq[SparkPlan]): Unit = { + val writers = + plans.flatMap(p => collectWithSubqueries(p) { case w: IcebergWriteExec => w }) + assert(writers.nonEmpty, plans.mkString("\n")) + assert(writers.forall(hasRequiredProjection), plans.mkString("\n")) + assert( + plans.forall(p => + collectWithSubqueries(p) { case w: CometIcebergWriteExec => w }.isEmpty), + plans.mkString("\n")) + plans.foreach(assertColumnarContract) + } + val valid = spark.read.parquet(validPath) + val successful = withNativeEnabled { + captureWrite(nativeControl) { + valid.writeTo(s"$catalog.$ns.$nativeControl").append() + } + } + assert(successful.snapshotDelta == 1) + assertJvmRequiredWrite(successful.plans) + assert( + successful.plans.forall(p => + collectWithSubqueries(p) { + case s: SortExec => s + case e: ShuffleExchangeLike => e + }.isEmpty), + successful.plans.mkString("\n")) + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + valid.writeTo(s"$catalog.$ns.$jvmControl").append() + assert(countSnapshots(jvmControl) == 1) + Seq(nativeControl, jvmControl).foreach { table => + checkAnswer( + sql(s"SELECT id, hex(fx) FROM $catalog.$ns.$table"), + Seq(Row(1, "01020304"), Row(2, "05060708"))) + } + } + + // A one-byte dictionary limit forces the first fixed(4) value through Parquet's + // plain writer, which rejects its three bytes before any later null assertion. + def appendInvalid(table: String): Unit = + spark.read.parquet(invalidPath).writeTo(s"$catalog.$ns.$table").append() + var candidateFailure: Throwable = null + val candidatePlans = withNativeEnabled { + capturePlans(spark, includeFailures = true) { + candidateFailure = intercept[Exception](appendInvalid(candidate)) + } + } + var referenceFailure: Throwable = null + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + referenceFailure = intercept[Exception](appendInvalid(reference)) + } + assertJvmRequiredWrite(candidatePlans) + Seq(candidate, reference).foreach { table => + assert(countSnapshots(table) == 0) + assert(parquetFiles(dataDir(table)).isEmpty) + checkAnswer(sql(s"SELECT * FROM $catalog.$ns.$table"), Seq.empty[Row]) + } + val expected = causeChain(referenceFailure).last + assert(expected.isInstanceOf[IllegalArgumentException], expected.toString) + assert( + expected.getMessage.contains("Fixed Binary size 3") && + expected.getMessage.contains("field type length 4"), + expected.toString) + val actual = causeChain(candidateFailure).last + assert( + actual.getClass == expected.getClass, + s"candidate root $actual != JVM root $expected") + assert(actual.getMessage == expected.getMessage) + } + } + } + } + } + test("native acceleration: fixed(N) column round-trips like the JVM writer") { assumeNativeAcceleration() withIcebergCatalog { _ => @@ -4690,8 +5517,10 @@ class CometIcebergWriteActionSuite * writer, expects both to fail, and returns the two failures. The native write has to go * through [[CometIcebergWriteExec]], so a fallback to the JVM writer cannot pass as parity. */ - private def failedWrites(nativeTable: String, jvmTable: String)( - write: String => Unit): (Throwable, Throwable) = { + private def failedWrites( + nativeTable: String, + jvmTable: String, + checkSparkAssertions: Boolean = false)(write: String => Unit): (Throwable, Throwable) = { var nativeFailure: Throwable = null val nativePlans = withNativeEnabled { capturePlans(spark, includeFailures = true) { @@ -4699,6 +5528,7 @@ class CometIcebergWriteActionSuite } } val jvmFailure = intercept[Exception](write(jvmTable)) + if (checkSparkAssertions) assertSparkRequiredWrite(nativePlans) assert( nativePlans.exists(p => collectWithSubqueries(p) { case e: CometIcebergWriteExec => e }.nonEmpty), @@ -4707,6 +5537,40 @@ class CometIcebergWriteActionSuite (nativeFailure, jvmFailure) } + private def requiredFieldError(error: Throwable): Throwable = causeChain(error) + .find { e => + e.isInstanceOf[NullPointerException] || (e match { + case s: SparkThrowable => s.getErrorClass == "NOT_NULL_ASSERT_VIOLATION" + case _ => false + }) + } + .getOrElse(fail(s"Expected Spark's null assertion failure, got $error")) + + private def hasRequiredProjection(plan: SparkPlan): Boolean = collectWithSubqueries(plan) { + case p: ProjectExec if p.projectList.exists(_.exists(_.isInstanceOf[AssertNotNull])) => p + }.nonEmpty + + private def assertSparkRequiredWrite( + plans: Seq[SparkPlan], + requireBridge: Boolean = true): Unit = { + val writers = + plans.flatMap(p => collectWithSubqueries(p) { case w: CometIcebergWriteExec => w }) + assert(writers.nonEmpty, plans.mkString("\n")) + writers.foreach { writer => + assert( + hasRequiredProjection(writer), + s"Expected Spark's required-field projection:\n$writer") + if (requireBridge) { + assert( + collectWithSubqueries(writer) { + case bridge: CometSparkToColumnarExec if hasRequiredProjection(bridge.child) => bridge + }.nonEmpty, + s"Expected Spark's required-field projection to feed Arrow into the native writer:\n$writer") + } + } + plans.foreach(assertColumnarContract) + } + private def captureWrite(tableName: String)(action: => Unit): WriteSnapshot = { val before = countSnapshots(tableName) val plans = capturePlans(spark)(action) diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergWriteBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergWriteBenchmark.scala index 3ca224ec1a4..16a29a02572 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergWriteBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergWriteBenchmark.scala @@ -19,16 +19,22 @@ package org.apache.spark.sql.benchmark +import scala.collection.mutable import scala.concurrent.duration._ import org.apache.spark.SparkConf import org.apache.spark.benchmark.Benchmark -import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.comet.CometIcebergWriteExec +import org.apache.spark.sql.{Row, SparkSession} +import org.apache.spark.sql.catalyst.expressions.Alias +import org.apache.spark.sql.catalyst.expressions.objects.AssertNotNull +import org.apache.spark.sql.catalyst.plans.QueryPlan +import org.apache.spark.sql.comet.{CometIcebergWriteExec, CometNativeScanExec, CometScanExec, CometSparkToColumnarExec, IcebergWriteExec} +import org.apache.spark.sql.execution.{ProjectExec, SparkPlan} import org.apache.spark.sql.internal.SQLConf import org.apache.comet.CometConf import org.apache.comet.CometSparkSessionExtensions +import org.apache.comet.iceberg.IcebergReflection /** * Benchmark of writes into an Iceberg table with Comet's native (iceberg-rust) writer on and off. @@ -73,6 +79,13 @@ import org.apache.comet.CometSparkSessionExtensions * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometIcebergWriteBenchmark * }}} * Results will be written to "spark/benchmarks/CometIcebergWriteBenchmark-**results.txt". + * + * Two additional required-field cases compare iceberg-java with the native writer after the same + * Spark assertion projection, unpartitioned and partitioned with a local sort. Both use Comet + * scan/exec and the split operator. Pass `required [rows]` to run only these cases (the default + * is 4194304 rows). They require a nullable Parquet source without null values, reject fallback, + * warm both writers once, then alternate them for five measured writes each. Table preparation, + * plan checks and Spark read-back checks are outside the write-and-commit timer. */ object CometIcebergWriteBenchmark extends CometBenchmarkBase { @@ -176,7 +189,8 @@ object CometIcebergWriteBenchmark extends CometBenchmarkBase { statement: String, expectedRowsAfter: Long, expectShuffle: Option[Boolean], - expectSort: Option[Boolean]) + expectSort: Option[Boolean], + requiredFields: Boolean = false) private def workloads(values: Int): Seq[Workload] = { val insert = s"INSERT INTO $targetTable SELECT * FROM parquetV1Table" @@ -285,7 +299,14 @@ object CometIcebergWriteBenchmark extends CometBenchmarkBase { } override def runCometBenchmark(mainArgs: Array[String]): Unit = { + require( + mainArgs.isEmpty || (mainArgs.head == "required" && mainArgs.length <= 2), + "Usage: CometIcebergWriteBenchmark [required [rows]]") + val requiredOnly = mainArgs.nonEmpty + val rowCount = mainArgs.lift(1).map(_.toInt).getOrElse(4 * 1024 * 1024) + require(rowCount >= sourceFiles, s"rows must be at least $sourceFiles") if (!icebergOnClasspath) { + require(!requiredOnly, "The required-field benchmark needs Iceberg on the classpath") // scalastyle:off println println("Iceberg is not on the classpath; skipping. Build with an Iceberg-enabled profile.") // scalastyle:on println @@ -295,14 +316,20 @@ object CometIcebergWriteBenchmark extends CometBenchmarkBase { withTempPath { warehouse => configureIcebergHadoopCatalog(warehouse, catalog) - runBenchmarkWithTable("Iceberg write", 4 * 1024 * 1024) { values => + runBenchmarkWithTable("Iceberg write", rowCount) { values => withTempPath { dir => withTempTable("parquetV1Table") { - saveAsParquetV1Table( - spark.sql(corpusQuery).repartition(sourceFiles).write, - dir.getCanonicalPath + "/parquetV1") try { - workloads(values).foreach(runWorkload(_, values)) + if (!requiredOnly) { + saveAsParquetV1Table( + spark.sql(corpusQuery(includeNulls = true)).repartition(sourceFiles).write, + dir.getCanonicalPath + "/parquetV1") + workloads(values).foreach(runWorkload(_, values)) + } + saveAsParquetV1Table( + spark.sql(corpusQuery(includeNulls = false)).repartition(sourceFiles).write, + dir.getCanonicalPath + "/requiredParquetV1") + runRequiredWorkloads(values) } finally { spark.sql(s"DROP TABLE IF EXISTS $targetTable") } @@ -312,6 +339,173 @@ object CometIcebergWriteBenchmark extends CometBenchmarkBase { } } + private def runRequiredWorkloads(values: Int): Unit = { + val commonConfs = Seq( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "true", + CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key -> "true", + CometConf.COMET_BATCH_SIZE.key -> "8192", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + // One file per scan partition; verify the actual count instead of inferring it from files. + SQLConf.FILES_MAX_PARTITION_BYTES.key -> "1073741824", + SQLConf.FILES_OPEN_COST_IN_BYTES.key -> "1073741824") + val writerArms = Seq(false, true).map { native => + Arm( + if (native) "native writer" else "iceberg-java writer", + Seq(CometConf.COMET_ICEBERG_NATIVE_WRITE_ENABLED.key -> native.toString), + expectComet = true, + expectNativeWrite = native) + } + withSQLConf(commonConfs: _*) { + val source = spark.table("parquetV1Table") + require(source.schema.fields.forall(_.nullable), "required source must be nullable") + require(source.rdd.getNumPartitions == sourceFiles, "unexpected source partition count") + val columns = source.columns.mkString(", ") + val digest = s"count(*), sum(CAST(xxhash64($columns) AS DECIMAL(38,0)))" + var expected = Row.empty + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + val nulls = source.columns.map(c => s"$c IS NULL").mkString(" OR ") + require( + spark.sql(s"SELECT * FROM parquetV1Table WHERE $nulls LIMIT 1").isEmpty, + "required source must contain no null values") + expected = spark.sql(s"SELECT $digest FROM parquetV1Table").head() + } + Seq(false, true).foreach { partitioned => + val workload = Workload( + s"required fields, ${if (partitioned) "partitioned and sorted" else "unpartitioned"}", + partitionSpec = if (partitioned) s"PARTITIONED BY ($partitionColumn)" else "", + properties = + Seq("'write.distribution-mode'='none'", "'write.spark.fanout.enabled'='false'"), + orderBy = if (partitioned) Some(partitionColumn) else None, + prePopulate = false, + statement = s"INSERT INTO $targetTable SELECT * FROM parquetV1Table", + expectedRowsAfter = values.toLong, + expectShuffle = Some(false), + expectSort = Some(partitioned), + requiredFields = true) + val benchmark = new Benchmark( + workload.title, + values.toLong, + warmupTime = Duration.Zero, + minTime = Duration.Zero, + output = output) + benchmark.out.println(Benchmark.getJVMOSInfo()) + benchmark.out.println(Benchmark.getProcessorName()) + benchmark.out.println(s"Workload: ${workload.title}") + benchmark.out.println( + IcebergReflection + .loadClass("org.apache.iceberg.IcebergBuild") + .getMethod("fullVersion") + .invoke(null)) + benchmark.out.println( + s"Spark=${spark.version}, Comet=${org.apache.comet.COMET_VERSION}, " + + s"revision=${org.apache.comet.COMET_REVISION}, rows=$values, partitions=$sourceFiles") + benchmark.out.println( + s"master=${spark.sparkContext.master}, " + + s"defaultParallelism=${spark.sparkContext.defaultParallelism}, " + + s"maxHeapBytes=${Runtime.getRuntime.maxMemory()}, " + + s"offHeapEnabled=${spark.conf.get("spark.memory.offHeap.enabled")}, " + + s"offHeapSize=${spark.conf.get("spark.memory.offHeap.size")}") + benchmark.out.println( + s"Settings: ${commonConfs.mkString(", ")}; " + + s"properties=${workload.properties.mkString(", ")}; compression=snappy") + val samples = writerArms.map(arm => arm.name -> mutable.ArrayBuffer.empty[Double]).toMap + var requiredProjection = Seq.empty[String] + def writeAndCheck(arm: Arm, timer: Benchmark.Timer): Unit = { + resetTable(workload) + val plans = capturePlans { + withSQLConf(arm.confs: _*) { + timer.startTiming() + spark.sql(workload.statement) + timer.stopTiming() + } + } + val projections = verifyRequiredPlan(arm, workload, plans) + if (requiredProjection.isEmpty) requiredProjection = projections + require( + projections == requiredProjection, + s"the writers must use the same Spark projection: $projections != $requiredProjection") + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + require( + spark.sql(s"SELECT $digest FROM $targetTable").head() == expected, + s"${arm.name}: Spark read-back differs from the source") + } + if (timer.iteration < 0) { + benchmark.out.println( + s"${arm.name}, ${arm.confs.mkString(", ")}\n" + + plans.mkString("\n--\n")) + } + } + // Use the same full write and read-back to warm each arm, then alternate the timed arms. + writerArms.foreach(arm => writeAndCheck(arm, new Benchmark.Timer(-1))) + (0 until 5).foreach { round => + val order = if (round % 2 == 0) writerArms else writerArms.reverse + order.foreach { arm => + val result = benchmark.measure(values.toLong, 1)(writeAndCheck(arm, _)) + samples(arm.name) += result.avgMs + benchmark.out.println(f"${arm.name}, sample=${round + 1}, " + + f"elapsed_ms=${result.avgMs}%.3f, rows_per_second=${values * 1000.0 / result.avgMs}%.3f") + } + } + val javaMean = samples(writerArms.head.name).sum / 5 + writerArms.foreach { arm => + val times = samples(arm.name) + val mean = times.sum / times.size + val stdev = math.sqrt(times.map(t => (t - mean) * (t - mean)).sum / (times.size - 1)) + benchmark.out.println(f"${arm.name}: mean_ms=$mean%.3f, stdev_ms=$stdev%.3f, " + + f"rows_per_second=${values * 1000.0 / mean}%.3f, time_relative_to_java=${mean / javaMean}%.4f") + } + } + } + } + + private def verifyRequiredPlan( + arm: Arm, + workload: Workload, + plans: Seq[SparkPlan]): Seq[String] = { + val writers = collectAcross(plans) { case w: CometIcebergWriteExec => w } + require(writers.nonEmpty == arm.expectNativeWrite, s"${arm.name}: wrong writer\n$plans") + if (!arm.expectNativeWrite) { + require( + collectAcross(plans) { case w: IcebergWriteExec => w }.nonEmpty, + s"${arm.name}: split writer is missing\n$plans") + } + val projections = collectAcross(plans) { + case p: ProjectExec if p.projectList.exists(_.exists(_.isInstanceOf[AssertNotNull])) => p + } + require(projections.nonEmpty, s"${arm.name}: Spark assertions are missing\n$plans") + writers.foreach { writer => + require( + collectAcross(Seq(writer)) { + case b: CometSparkToColumnarExec if collectAcross(Seq(b.child)) { + case p: ProjectExec + if p.projectList.exists(_.exists(_.isInstanceOf[AssertNotNull])) => + p + }.nonEmpty => + b + }.nonEmpty, + s"${arm.name}: Spark assertion-to-Arrow bridge is missing\n$writer") + } + require( + collectAcross(plans) { + case p: CometScanExec => p + case p: CometNativeScanExec => p + }.nonEmpty, + s"${arm.name}: Comet scan is missing\n$plans") + require( + collectAcross(plans) { case p if p.nodeName.contains("Exchange") => p }.isEmpty, + s"${arm.name}: an exchange changed the source partitions\n$plans") + val sorted = collectAcross(plans) { case p if p.nodeName.contains("Sort") => p }.nonEmpty + require(workload.expectSort.contains(sorted), s"${arm.name}: unexpected sort\n$plans") + projections.flatMap { p => + p.projectList.map { + case a: Alias => QueryPlan.normalizeExpressions(a.child, p.child.output).toString + case e => QueryPlan.normalizeExpressions(e, p.child.output).toString + } + }.distinct + } + private def runWorkload(workload: Workload, values: Int): Unit = { // `warmupTime` is a deadline and `minTime` a total across the measured iterations, so both are // floors on a case's total cost rather than on one write. The floors are set below the fastest @@ -455,20 +649,38 @@ object CometIcebergWriteBenchmark extends CometBenchmarkBase { workload.properties // `WHERE false` takes the schema from the source without writing anything, which avoids // restating the corpus as DDL and keeps the two in step. + val schema = if (workload.requiredFields) { + spark + .table("parquetV1Table") + .schema + .fields + .map(f => s"${f.name} ${f.dataType.sql} NOT NULL") + .mkString("(", ", ", ")") + } else "" + val source = + if (workload.requiredFields) "" else "AS SELECT * FROM parquetV1Table WHERE false" spark.sql(s""" CREATE TABLE $targetTable + $schema USING iceberg ${workload.partitionSpec} TBLPROPERTIES (${properties.mkString(", ")}) - AS SELECT * FROM parquetV1Table WHERE false + $source """) workload.orderBy.foreach { col => // Give the write a required ordering so Iceberg selects the clustered writer instead of // defaulting to fanout. `DISTRIBUTED BY PARTITION` keeps the hash exchange so the case still // shuffles; `LOCALLY ORDERED BY` adds the per-task sort the clustered writer needs to keep a // single file open without erroring when a partition it already closed would reappear. - spark.sql( - s"ALTER TABLE $targetTable WRITE DISTRIBUTED BY PARTITION LOCALLY ORDERED BY $col") + if (workload.requiredFields) { + spark.sql(s"ALTER TABLE $targetTable WRITE ORDERED BY $col") + // WRITE ORDERED BY also switches distribution to range; keep only its local sort. + spark.sql( + s"ALTER TABLE $targetTable SET TBLPROPERTIES ('write.distribution-mode'='none')") + } else { + spark.sql( + s"ALTER TABLE $targetTable WRITE DISTRIBUTED BY PARTITION LOCALLY ORDERED BY $col") + } } if (workload.prePopulate) { // Always stock Spark, whichever arm is about to be timed. The two writers roll files at @@ -493,7 +705,7 @@ object CometIcebergWriteBenchmark extends CometBenchmarkBase { * column, because a row that is null in all ten columns is one definition-level pattern * repeated ten times, which both writers would compress better than a real table allows. */ - private def corpusQuery: String = { + private def corpusQuery(includeNulls: Boolean): String = { val columns = Seq( "c_bool" -> "PMOD(value, 2) = 0", "c_int" -> "CAST(PMOD(value, 2147483647) AS INT)", @@ -506,7 +718,8 @@ object CometIcebergWriteBenchmark extends CometBenchmarkBase { "c_date" -> "DATE_ADD(DATE '1970-01-01', CAST(PMOD(value, 20000) AS INT))", "c_ts" -> "TIMESTAMP_SECONDS(PMOD(value, 1600000000))") val projections = columns.zipWithIndex.map { case ((name, expr), position) => - s"IF(PMOD(value + $position, $nullStride) = 0, NULL, $expr) AS $name" + if (includeNulls) s"IF(PMOD(value + $position, $nullStride) = 0, NULL, $expr) AS $name" + else s"$expr AS $name" } s"SELECT ${projections.mkString(", ")} FROM $tbl" } diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala index 34bb32d908f..c7dade28f4f 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala @@ -21,7 +21,7 @@ package org.apache.spark.sql.comet.execution.arrow import java.nio.charset.StandardCharsets -import scala.collection.mutable.ArrayBuffer +import scala.collection.mutable.{ArrayBuffer, ListBuffer} import scala.jdk.CollectionConverters._ import scala.util.Using @@ -38,6 +38,7 @@ import org.apache.arrow.vector.types.pojo.{ArrowType, DictionaryEncoding, Field, import org.apache.spark.rdd.InputFileBlockHolder import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, SpecializedGetters} import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, GenericArrayData} +import org.apache.spark.sql.comet.CometSparkToColumnarExec import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution.vectorized.{ConstantColumnVector, Dictionary, OffHeapColumnVector, OnHeapColumnVector} import org.apache.spark.sql.types.{ArrayType, BooleanType, ByteType, CalendarIntervalType, DataType, DateType, DayTimeIntervalType, Decimal, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, ShortType, StringType, StructField, StructType, TimestampNTZType, TimestampType, YearMonthIntervalType} @@ -54,6 +55,25 @@ import org.apache.comet.vector.{CometPlainVector, CometVector, NativeUtil} */ class CometArrowStreamSuite extends AnyFunSuite with Matchers { + test("row conversion schema gate supports nested collections") { + val array = ArrayType(IntegerType) + val map = MapType(IntegerType, ArrayType(LongType)) + Seq(array, map, new StructType().add("a", array).add("m", map)).foreach { dataType => + val schema = new StructType().add("value", dataType) + CometSparkToColumnarExec.isRowSchemaSupported(schema, ListBuffer.empty) shouldBe true + CometSparkToColumnarExec.isSchemaSupported(schema, ListBuffer.empty) shouldBe false + } + + val duplicateFields = new StructType().add("x", IntegerType).add("x", LongType) + Seq( + ArrayType(duplicateFields), + MapType(StringType, new StructType()), + ArrayType(YearMonthIntervalType())).foreach { dataType => + val schema = new StructType().add("value", dataType) + CometSparkToColumnarExec.isRowSchemaSupported(schema, ListBuffer.empty) shouldBe false + } + } + private def expectedSchema(types: (String, ArrowType)*): Schema = { val fields = types.map { case (name, t) => new Field(name, new FieldType(true, t, null), java.util.Collections.emptyList[Field]())