From 3d0ba56e31c356fd0b04a15e103953225b07cb80 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 7 Jul 2026 08:44:59 -0600 Subject: [PATCH] fix: widen local table scan child nullability to match native kernels An in-memory Seq/Map column read through CometLocalTableScanExec encodes its array element and map value as non-null (containsNull=false / valueContainsNull=false), but Comet expression serdes promise nullable children. Feeding a non-null child into a native array/map kernel fails the planned-vs-actual type assertion (e.g. spark_array_slice returning List(non-null Int32) while List(Int32) was promised) or panics building a ListArray for map_entries. Widen nested child-field nullability on both sides so they agree: the declared scan schema (serialized from op.output in CometSink) and the Arrow batch produced by the local scan. Without widening the declared schema too, the native ScanExec casts the batch back to the non-null type in build_record_batch. The widening is scoped to the local table scan via an overridable scanFieldType hook, leaving other sinks untouched. Map keys stay non-null, matching Arrow. This mirrors the Parquet scan path, which already reads children as nullable. Closes #4789 --- .../comet/serde/operator/CometSink.scala | 10 +++++- .../sql/comet/CometLocalTableScanExec.scala | 17 ++++++++-- .../comet/CometArrayExpressionSuite.scala | 31 +++++++++++++++++++ .../comet/CometMapExpressionSuite.scala | 14 +++++++++ 4 files changed, 69 insertions(+), 3 deletions(-) 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)")) + } + } + }