diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometSink.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometSink.scala index b1834c5083..ed69f82e2a 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometSink.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometSink.scala @@ -26,6 +26,7 @@ import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.adaptive.ShuffleQueryStageExec import org.apache.spark.sql.execution.exchange.ReusedExchangeExec +import org.apache.spark.sql.types.DataType import org.apache.comet.CometConf import org.apache.comet.CometSparkSessionExtensions.withFallbackReason @@ -42,6 +43,13 @@ abstract class CometSink[T <: SparkPlan] extends CometOperatorSerde[T] { override def enabledConfig: Option[ConfigEntry[Boolean]] = None + /** + * The data type to declare for a scan output field. Overridden by sinks whose source carries + * non-null nested child fields that must be widened to match the planned kernel output types + * (see [[org.apache.spark.sql.comet.CometLocalTableScanExec]] and issue #4789). + */ + protected def scanFieldType(dt: DataType): DataType = dt + override def convert( op: T, builder: Operator.Builder, @@ -64,7 +72,7 @@ abstract class CometSink[T <: SparkPlan] extends CometOperatorSerde[T] { } val scanTypes = op.output.flatten { attr => - serializeDataType(attr.dataType) + serializeDataType(scanFieldType(attr.dataType)) } if (scanTypes.length == op.output.length) { diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometLocalTableScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometLocalTableScanExec.scala index 161a4e0553..28ff6c85c2 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometLocalTableScanExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometLocalTableScanExec.scala @@ -31,7 +31,7 @@ import org.apache.spark.sql.comet.execution.arrow.{CometArrowStream, CometNative import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution.{LeafExecNode, LocalTableScanExec} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} -import org.apache.spark.sql.types.{DataType, NullType} +import org.apache.spark.sql.types.{DataType, NullType, StructType} import com.google.common.base.Objects @@ -76,7 +76,15 @@ case class CometLocalTableScanExec( consume: (String, BufferAllocator => ArrowReader) => Iterator[T]): RDD[T] = { val numOutputRows = longMetric("numOutputRows") val maxRecordsPerBatch = CometConf.COMET_BATCH_SIZE.get(conf) - val sparkSchema = originalPlan.schema + // Normalize nested child-field nullability. An in-memory Seq/Map encodes array elements and + // map values as non-null (containsNull=false / valueContainsNull=false), but Comet expression + // serdes promise nullable children, so feeding a non-null child from the local scan into a + // native kernel fails the planned-vs-actual type assertion (issue #4789). Widening children to + // nullable is always safe: the row data contains no nulls, and Arrow map keys stay non-null via + // Utils.toArrowField. This must agree with the widened scan schema declared in + // CometLocalTableScanExec.scanFieldType so the native ScanExec does not cast the batch back. + val sparkSchema = + StructType(originalPlan.schema.map(f => f.copy(dataType = f.dataType.asNullable))) rdd.mapPartitionsInternal { rowIter => val arrowSchema = Utils.toArrowSchema(sparkSchema, CometArrowStream.NATIVE_TIMEZONE) consume( @@ -118,6 +126,11 @@ object CometLocalTableScanExec extends CometSink[LocalTableScanExec] with DataTy override def enabledConfig: Option[ConfigEntry[Boolean]] = Some( CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED) + // Widen nested child-field nullability in the declared scan schema so it matches both the widened + // Arrow batch produced by CometLocalTableScanExec and the nullable children promised by + // downstream expression serdes (issue #4789). + override protected def scanFieldType(dt: DataType): DataType = dt.asNullable + // ArrowWriter (used by RowArrowReader) handles NullType via Utils.toArrowType + NullWriter; // other types off DataTypeSupport's allow list (TimeType, intervals, ...) have no ArrowWriter // coverage and must fall back to Spark. diff --git a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala index 87788c0760..7377f4fe4f 100644 --- a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala @@ -1146,6 +1146,37 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp } } + // Local table scan carries non-null array child fields (an in-memory Seq encodes + // containsNull=false) into native kernels that promise nullable elements. ConvertToLocalRelation + // must be disabled or the optimizer folds the expression at plan time and nothing runs natively. + // https://github.com/apache/datafusion-comet/issues/4789 + private def withLocalTableScanNoFold(f: => Unit): Unit = { + withSQLConf( + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true", + "spark.sql.optimizer.excludedRules" -> + "org.apache.spark.sql.catalyst.optimizer.ConvertToLocalRelation") { + f + } + } + + test("slice on non-null element array from local table scan (#4789)") { + withLocalTableScanNoFold { + import testImplicits._ + val df = Seq(Seq(1, 2, 3), Seq(4, 5)).toDF("x") + checkSparkAnswerAndOperator(df.selectExpr("slice(x, 2, 2)")) + } + } + + test("array_insert on non-null element array from local table scan (#4789)") { + assume(isSpark35Plus) + withLocalTableScanNoFold { + import testImplicits._ + val df = Seq(Seq(1, 2, 3), Seq(4, 5)).toDF("x") + // SPARK-41233 array prepend lowers to array_insert at position 1 + checkSparkAnswerAndOperator(df.selectExpr("array_insert(x, 1, 0)")) + } + } + // https://issues.apache.org/jira/browse/SPARK-55747 test("(ansi) GetArrayItem on null array from split()") { withSQLConf( diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index 9904919b74..da11053720 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -247,4 +247,18 @@ class CometMapExpressionSuite extends CometTestBase { } } + test("map_entries on non-null value map from local table scan (#4789)") { + // An in-memory Map encodes valueContainsNull=false; the local scan must widen the map value + // to nullable so map_entries' native ListArray/Struct build does not fail on the child type. + // ConvertToLocalRelation must be disabled or the expression folds at plan time. + withSQLConf( + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true", + "spark.sql.optimizer.excludedRules" -> + "org.apache.spark.sql.catalyst.optimizer.ConvertToLocalRelation") { + import testImplicits._ + val df = Seq(Map(1 -> 100, 2 -> 200)).toDF("m") + checkSparkAnswerAndOperator(df.selectExpr("map_entries(m)")) + } + } + }