This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-5855-db790673d0a2025534fdd80c0319ace18d6d6770 in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit e103a7a731193fdf04b8905611df400cfff83699 Author: Han-Yin Chang <[email protected]> AuthorDate: Fri Sep 11 23:12:19 2026 +0000 fix: preserve nulls in generated integral columns (#5855) Co-authored-by: Andy Grove <[email protected]> --- .../apache/comet/testing/FuzzDataGenerator.scala | 27 ++++++--- .../org/apache/comet/DataGeneratorSuite.scala | 66 +++++++++++++++++++++- 2 files changed, 82 insertions(+), 11 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/testing/FuzzDataGenerator.scala b/spark/src/main/scala/org/apache/comet/testing/FuzzDataGenerator.scala index cf9d8a7568..c3d1078097 100644 --- a/spark/src/main/scala/org/apache/comet/testing/FuzzDataGenerator.scala +++ b/spark/src/main/scala/org/apache/comet/testing/FuzzDataGenerator.scala @@ -197,19 +197,28 @@ object FuzzDataGenerator { val k = generateColumn(r, keyType, numRows, mapOptions) val v = generateColumn(r, valueType, numRows, mapOptions) k.zip(v).map(x => Map(x._1 -> x._2)) + // A cast to `Long` here would unbox a null element to 0 instead of throwing, so match on + // the element. case DataTypes.BooleanType => - generateColumn(r, DataTypes.LongType, numRows, options) - .map(_.asInstanceOf[Long].toShort) - .map(s => s % 2 == 0) + generateColumn(r, DataTypes.LongType, numRows, options).map { + case x: Long => x.toShort % 2 == 0 + case null => null + } case DataTypes.ByteType => - generateColumn(r, DataTypes.LongType, numRows, options) - .map(_.asInstanceOf[Long].toByte) + generateColumn(r, DataTypes.LongType, numRows, options).map { + case x: Long => x.toByte + case null => null + } case DataTypes.ShortType => - generateColumn(r, DataTypes.LongType, numRows, options) - .map(_.asInstanceOf[Long].toShort) + generateColumn(r, DataTypes.LongType, numRows, options).map { + case x: Long => x.toShort + case null => null + } case DataTypes.IntegerType => - generateColumn(r, DataTypes.LongType, numRows, options) - .map(_.asInstanceOf[Long].toInt) + generateColumn(r, DataTypes.LongType, numRows, options).map { + case x: Long => x.toInt + case null => null + } case DataTypes.LongType => Range(0, numRows).map(_ => { r.nextInt(50) match { diff --git a/spark/src/test/scala/org/apache/comet/DataGeneratorSuite.scala b/spark/src/test/scala/org/apache/comet/DataGeneratorSuite.scala index eb04df6f3a..d1a321099c 100644 --- a/spark/src/test/scala/org/apache/comet/DataGeneratorSuite.scala +++ b/spark/src/test/scala/org/apache/comet/DataGeneratorSuite.scala @@ -22,12 +22,74 @@ package org.apache.comet import scala.util.Random import org.apache.spark.sql.CometTestBase -import org.apache.spark.sql.types.{ArrayType, DataType, MapType, StructType} +import org.apache.spark.sql.types.{ArrayType, DataType, DataTypes, MapType, StructField, StructType} -import org.apache.comet.testing.{FuzzDataGenerator, SchemaGenOptions} +import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, SchemaGenOptions} class DataGeneratorSuite extends CometTestBase { + test("allowNull produces nulls in every type that honours it") { + // Decimal, Date, Timestamp and TimestampNTZ never consult `allowNull`, so they are excluded. + val types = SchemaGenOptions.defaultPrimitiveTypes.filterNot { dataType => + dataType.isInstanceOf[org.apache.spark.sql.types.DecimalType] || + dataType == DataTypes.DateType || + dataType == DataTypes.TimestampType || + dataType == DataTypes.TimestampNTZType + } + val schema = StructType(types.zipWithIndex.map { case (dataType, i) => + StructField(s"c$i", dataType, nullable = true) + }) + + def rowsFor(allowNull: Boolean) = + FuzzDataGenerator + .generateDataFrame( + new Random(42), + spark, + schema, + numRows = 1000, + DataGenOptions(allowNull = allowNull)) + .collect() + + val withNulls = rowsFor(allowNull = true) + schema.fields.zipWithIndex.foreach { case (field, i) => + assert( + withNulls.exists(_.isNullAt(i)), + s"no nulls generated for nullable ${field.dataType.catalogString} column") + } + + val withoutNulls = rowsFor(allowNull = false) + schema.fields.zipWithIndex.foreach { case (field, i) => + assert( + !withoutNulls.exists(_.isNullAt(i)), + s"nulls generated for ${field.dataType.catalogString} column with allowNull disabled") + } + + val nestedSchema = StructType( + Seq( + StructField("array", ArrayType(DataTypes.IntegerType), nullable = true), + StructField( + "struct", + StructType(Seq(StructField("integer", DataTypes.IntegerType, nullable = true))), + nullable = true))) + val nestedRows = FuzzDataGenerator + .generateDataFrame( + new Random(42), + spark, + nestedSchema, + numRows = 1000, + DataGenOptions(allowNull = true)) + .collect() + assert( + nestedRows.exists { row => + val values = row.getAs[scala.collection.Seq[Any]]("array") + values != null && values.contains(null) + }, + "no null generated in an array element") + assert( + nestedRows.exists(row => row.getStruct(row.fieldIndex("struct")).isNullAt(0)), + "no null generated in a struct field") + } + test("generate nested schema has at least minDepth levels") { val minDepth = 3 val numCols = 4 --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
