From e365ac31574b943356848e68eaa6d5c364ca3bab Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 22 Oct 2025 09:34:53 -0600 Subject: [PATCH 1/5] add version of generateDataFrame that accepts a custom schema --- .../org/apache/comet/testing/FuzzDataGenerator.scala | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) 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 7c7a6727fb0..a90d1797006 100644 --- a/spark/src/main/scala/org/apache/comet/testing/FuzzDataGenerator.scala +++ b/spark/src/main/scala/org/apache/comet/testing/FuzzDataGenerator.scala @@ -78,7 +78,6 @@ object FuzzDataGenerator { spark: SparkSession, numRows: Int, options: DataGenOptions): DataFrame = { - val filteredPrimitiveTypes = filteredPrimitives(options.excludeTypes) val dataTypes = ListBuffer[DataType]() dataTypes.appendAll(filteredPrimitiveTypes) @@ -119,6 +118,16 @@ object FuzzDataGenerator { .map(i => StructField(s"c${i._2}", i._1, nullable = true)) val schema = StructType(fields.toSeq) + generateDataFrame(r, spark, schema, numRows, options) + } + + def generateDataFrame( + r: Random, + spark: SparkSession, + schema: StructType, + numRows: Int, + options: DataGenOptions): DataFrame = { + // generate columnar data val cols: Seq[Seq[Any]] = schema.fields.map(f => generateColumn(r, f.dataType, numRows, options)).toSeq From 10c960e805eb75996a3d842cb35277b9a3404b7e Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 22 Oct 2025 09:49:41 -0600 Subject: [PATCH 2/5] refactor --- .../scala/org/apache/comet/fuzz/Main.scala | 4 +- .../comet/testing/FuzzDataGenerator.scala | 89 +------------------ .../comet/testing/ParquetGenerator.scala | 77 +++++++++++++++- .../comet/CometArrayExpressionSuite.scala | 22 ++--- .../comet/CometBitwiseExpressionSuite.scala | 6 +- .../org/apache/comet/CometFuzzTestBase.scala | 4 +- .../org/apache/comet/CometFuzzTestSuite.scala | 4 +- .../comet/CometMapExpressionSuite.scala | 4 +- .../comet/exec/CometAggregateSuite.scala | 4 +- .../apache/comet/exec/CometExecSuite.scala | 4 +- 10 files changed, 105 insertions(+), 113 deletions(-) diff --git a/fuzz-testing/src/main/scala/org/apache/comet/fuzz/Main.scala b/fuzz-testing/src/main/scala/org/apache/comet/fuzz/Main.scala index 9b9a4b6f3e5..74edd940089 100644 --- a/fuzz-testing/src/main/scala/org/apache/comet/fuzz/Main.scala +++ b/fuzz-testing/src/main/scala/org/apache/comet/fuzz/Main.scala @@ -26,7 +26,7 @@ import org.rogach.scallop.ScallopOption import org.apache.spark.sql.SparkSession -import org.apache.comet.testing.{DataGenOptions, ParquetGenerator} +import org.apache.comet.testing.{ParquetDataGenOptions, ParquetGenerator} class Conf(arguments: Seq[String]) extends ScallopConf(arguments) { object generateData extends Subcommand("data") { @@ -78,7 +78,7 @@ object Main { case Some(seed) => new Random(seed) case None => new Random() } - val options = DataGenOptions( + val options = ParquetDataGenOptions( allowNull = true, generateArray = conf.generateData.generateArrays(), generateStruct = conf.generateData.generateStructs(), 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 a90d1797006..d75c53259f0 100644 --- a/spark/src/main/scala/org/apache/comet/testing/FuzzDataGenerator.scala +++ b/spark/src/main/scala/org/apache/comet/testing/FuzzDataGenerator.scala @@ -44,89 +44,12 @@ object FuzzDataGenerator { val defaultBaseDate: Long = new SimpleDateFormat("YYYY-MM-DD hh:mm:ss").parse("3333-05-25 12:34:56").getTime - private val primitiveTypes = Seq( - DataTypes.BooleanType, - DataTypes.ByteType, - DataTypes.ShortType, - DataTypes.IntegerType, - DataTypes.LongType, - DataTypes.FloatType, - DataTypes.DoubleType, - DataTypes.createDecimalType(10, 2), - DataTypes.createDecimalType(36, 18), - DataTypes.DateType, - DataTypes.TimestampType, - DataTypes.TimestampNTZType, - DataTypes.StringType, - DataTypes.BinaryType) - - private def filteredPrimitives(excludeTypes: Seq[DataType]) = { - - primitiveTypes.filterNot { dataType => - excludeTypes.exists { - case _: DecimalType => - // For DecimalType, match if the type is also a DecimalType (ignore precision/scale) - dataType.isInstanceOf[DecimalType] - case excludeType => - dataType == excludeType - } - } - } - - def generateDataFrame( - r: Random, - spark: SparkSession, - numRows: Int, - options: DataGenOptions): DataFrame = { - val filteredPrimitiveTypes = filteredPrimitives(options.excludeTypes) - val dataTypes = ListBuffer[DataType]() - dataTypes.appendAll(filteredPrimitiveTypes) - - val arraysOfPrimitives = filteredPrimitiveTypes.map(DataTypes.createArrayType) - - if (options.generateStruct) { - dataTypes += StructType(filteredPrimitiveTypes.zipWithIndex.map(x => - StructField(s"c${x._2}", x._1, nullable = true))) - - if (options.generateArray) { - dataTypes += StructType(arraysOfPrimitives.zipWithIndex.map(x => - StructField(s"c${x._2}", x._1, nullable = true))) - } - } - - if (options.generateMap) { - dataTypes += MapType(DataTypes.IntegerType, DataTypes.StringType) - } - - if (options.generateArray) { - dataTypes.appendAll(arraysOfPrimitives) - - if (options.generateStruct) { - dataTypes += DataTypes.createArrayType( - StructType(filteredPrimitiveTypes.zipWithIndex.map(x => - StructField(s"c${x._2}", x._1, nullable = true)))) - } - - if (options.generateMap) { - dataTypes += DataTypes.createArrayType( - MapType(DataTypes.IntegerType, DataTypes.StringType)) - } - } - - // generate schema using random data types - val fields = dataTypes.zipWithIndex - .map(i => StructField(s"c${i._2}", i._1, nullable = true)) - val schema = StructType(fields.toSeq) - - generateDataFrame(r, spark, schema, numRows, options) - } - def generateDataFrame( r: Random, spark: SparkSession, schema: StructType, numRows: Int, - options: DataGenOptions): DataFrame = { + options: DataGenOptions2): DataFrame = { // generate columnar data val cols: Seq[Seq[Any]] = @@ -144,7 +67,7 @@ object FuzzDataGenerator { r: Random, dataType: DataType, numRows: Int, - options: DataGenOptions): Seq[Any] = { + options: DataGenOptions2): Seq[Any] = { dataType match { case ArrayType(elementType, _) => val values = generateColumn(r, elementType, numRows, options) @@ -256,11 +179,7 @@ object FuzzDataGenerator { } } -case class DataGenOptions( +case class DataGenOptions2( allowNull: Boolean = true, generateNegativeZero: Boolean = true, - baseDate: Long = FuzzDataGenerator.defaultBaseDate, - generateArray: Boolean = false, - generateStruct: Boolean = false, - generateMap: Boolean = false, - excludeTypes: Seq[DataType] = Seq.empty) + baseDate: Long = FuzzDataGenerator.defaultBaseDate) diff --git a/spark/src/main/scala/org/apache/comet/testing/ParquetGenerator.scala b/spark/src/main/scala/org/apache/comet/testing/ParquetGenerator.scala index 27e40c9d747..8d63bf02df2 100644 --- a/spark/src/main/scala/org/apache/comet/testing/ParquetGenerator.scala +++ b/spark/src/main/scala/org/apache/comet/testing/ParquetGenerator.scala @@ -19,21 +19,94 @@ package org.apache.comet.testing +import scala.collection.mutable.ListBuffer import scala.util.Random import org.apache.spark.sql.{SaveMode, SparkSession} +import org.apache.spark.sql.types.{DataType, DataTypes, MapType, StructField, StructType} object ParquetGenerator { + def makeParquetSchema(options: ParquetDataGenOptions): StructType = { + val primitiveTypes = options.primitiveTypes + val dataTypes = ListBuffer[DataType]() + dataTypes.appendAll(primitiveTypes) + + val arraysOfPrimitives = primitiveTypes.map(DataTypes.createArrayType) + + if (options.generateStruct) { + dataTypes += StructType( + primitiveTypes.zipWithIndex.map(x => StructField(s"c${x._2}", x._1, nullable = true))) + + if (options.generateArray) { + dataTypes += StructType(arraysOfPrimitives.zipWithIndex.map(x => + StructField(s"c${x._2}", x._1, nullable = true))) + } + } + + if (options.generateMap) { + dataTypes += MapType(DataTypes.IntegerType, DataTypes.StringType) + } + + if (options.generateArray) { + dataTypes.appendAll(arraysOfPrimitives) + + if (options.generateStruct) { + dataTypes += DataTypes.createArrayType(StructType(primitiveTypes.zipWithIndex.map(x => + StructField(s"c${x._2}", x._1, nullable = true)))) + } + + if (options.generateMap) { + dataTypes += DataTypes.createArrayType( + MapType(DataTypes.IntegerType, DataTypes.StringType)) + } + } + + // generate schema using random data types + val fields = dataTypes.zipWithIndex + .map(i => StructField(s"c${i._2}", i._1, nullable = true)) + StructType(fields.toSeq) + } + def makeParquetFile( r: Random, spark: SparkSession, filename: String, numRows: Int, - options: DataGenOptions): Unit = { + options: ParquetDataGenOptions): Unit = { - val df = FuzzDataGenerator.generateDataFrame(r, spark, numRows, options) + val schema = makeParquetSchema(options) + + val x = DataGenOptions2( + allowNull = options.allowNull, + generateNegativeZero = options.generateNegativeZero, + baseDate = options.baseDate) + + val df = FuzzDataGenerator.generateDataFrame(r, spark, schema, numRows, x) df.write.mode(SaveMode.Overwrite).parquet(filename) } } + +case class ParquetDataGenOptions( + allowNull: Boolean = true, + generateNegativeZero: Boolean = true, + baseDate: Long = FuzzDataGenerator.defaultBaseDate, + generateArray: Boolean = false, + generateStruct: Boolean = false, + generateMap: Boolean = false, + primitiveTypes: Seq[DataType] = Seq( + DataTypes.BooleanType, + DataTypes.ByteType, + DataTypes.ShortType, + DataTypes.IntegerType, + DataTypes.LongType, + DataTypes.FloatType, + DataTypes.DoubleType, + DataTypes.createDecimalType(10, 2), + DataTypes.createDecimalType(36, 18), + DataTypes.DateType, + DataTypes.TimestampType, + DataTypes.TimestampNTZType, + DataTypes.StringType, + DataTypes.BinaryType)) diff --git a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala index 2adb7a9ed61..1313e3e6b8d 100644 --- a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala @@ -30,7 +30,7 @@ import org.apache.spark.sql.functions._ import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus} import org.apache.comet.DataTypeSupport.isComplexType import org.apache.comet.serde.{CometArrayExcept, CometArrayRemove, CometArrayReverse, CometFlatten} -import org.apache.comet.testing.{DataGenOptions, ParquetGenerator} +import org.apache.comet.testing.{ParquetDataGenOptions, ParquetGenerator} class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { @@ -64,7 +64,7 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp spark, filename, 100, - DataGenOptions( + ParquetDataGenOptions( allowNull = true, generateNegativeZero = true, generateArray = false, @@ -95,7 +95,7 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp val filename = path.toString val random = new Random(42) withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - val options = DataGenOptions( + val options = ParquetDataGenOptions( allowNull = true, generateNegativeZero = true, generateArray = true, @@ -266,7 +266,7 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp spark, filename, 100, - DataGenOptions( + ParquetDataGenOptions( allowNull = true, generateNegativeZero = true, generateArray = true, @@ -310,7 +310,7 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp spark, filename, 100, - DataGenOptions( + ParquetDataGenOptions( allowNull = true, generateNegativeZero = true, generateArray = false, @@ -340,7 +340,7 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp spark, filename, 100, - DataGenOptions( + ParquetDataGenOptions( allowNull = true, generateNegativeZero = true, generateArray = true, @@ -588,7 +588,7 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp spark, filename, 100, - DataGenOptions( + ParquetDataGenOptions( allowNull = true, generateNegativeZero = true, generateArray = false, @@ -622,7 +622,7 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp val filename = path.toString val random = new Random(42) withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - val options = DataGenOptions( + val options = ParquetDataGenOptions( allowNull = true, generateNegativeZero = true, generateArray = true, @@ -692,7 +692,7 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp spark, filename, 100, - DataGenOptions( + ParquetDataGenOptions( allowNull = true, generateNegativeZero = true, generateArray = false, @@ -720,7 +720,7 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp val filename = path.toString val random = new Random(42) withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - val options = DataGenOptions( + val options = ParquetDataGenOptions( allowNull = true, generateNegativeZero = true, generateArray = true, @@ -773,7 +773,7 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp val filename = path.toString val random = new Random(42) withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - val options = DataGenOptions( + val options = ParquetDataGenOptions( allowNull = true, generateNegativeZero = true, generateArray = true, diff --git a/spark/src/test/scala/org/apache/comet/CometBitwiseExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometBitwiseExpressionSuite.scala index d89e81b0fd2..8eaec24aaeb 100644 --- a/spark/src/test/scala/org/apache/comet/CometBitwiseExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometBitwiseExpressionSuite.scala @@ -25,7 +25,7 @@ import org.apache.hadoop.fs.Path import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper -import org.apache.comet.testing.{DataGenOptions, ParquetGenerator} +import org.apache.comet.testing.{ParquetDataGenOptions, ParquetGenerator} class CometBitwiseExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { @@ -99,7 +99,7 @@ class CometBitwiseExpressionSuite extends CometTestBase with AdaptiveSparkPlanHe spark, filename, 100, - DataGenOptions( + ParquetDataGenOptions( allowNull = true, generateNegativeZero = true, generateArray = false, @@ -174,7 +174,7 @@ class CometBitwiseExpressionSuite extends CometTestBase with AdaptiveSparkPlanHe spark, filename, 10, - DataGenOptions( + ParquetDataGenOptions( allowNull = true, generateNegativeZero = true, generateArray = false, diff --git a/spark/src/test/scala/org/apache/comet/CometFuzzTestBase.scala b/spark/src/test/scala/org/apache/comet/CometFuzzTestBase.scala index a69080e446b..9068e5b74d3 100644 --- a/spark/src/test/scala/org/apache/comet/CometFuzzTestBase.scala +++ b/spark/src/test/scala/org/apache/comet/CometFuzzTestBase.scala @@ -35,7 +35,7 @@ import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.internal.SQLConf -import org.apache.comet.testing.{DataGenOptions, ParquetGenerator} +import org.apache.comet.testing.{ParquetDataGenOptions, ParquetGenerator} class CometFuzzTestBase extends CometTestBase with AdaptiveSparkPlanHelper { @@ -59,7 +59,7 @@ class CometFuzzTestBase extends CometTestBase with AdaptiveSparkPlanHelper { CometConf.COMET_ENABLED.key -> "false", SQLConf.SESSION_LOCAL_TIMEZONE.key -> defaultTimezone) { val options = - DataGenOptions( + ParquetDataGenOptions( generateArray = true, generateStruct = true, generateNegativeZero = false, diff --git a/spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala b/spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala index 398d9601363..a6e55446b49 100644 --- a/spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala @@ -28,7 +28,7 @@ import org.apache.spark.sql.internal.SQLConf.ParquetOutputTimestampType import org.apache.spark.sql.types._ import org.apache.comet.DataTypeSupport.isComplexType -import org.apache.comet.testing.{DataGenOptions, ParquetGenerator} +import org.apache.comet.testing.{ParquetDataGenOptions, ParquetGenerator} class CometFuzzTestSuite extends CometFuzzTestBase { @@ -262,7 +262,7 @@ class CometFuzzTestSuite extends CometFuzzTestBase { generateStruct: Boolean = true): Unit = { val options = - DataGenOptions( + ParquetDataGenOptions( generateArray = generateArray, generateStruct = generateStruct, generateNegativeZero = false) diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index 3557ad348bc..7283b696b4f 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -26,7 +26,7 @@ import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -import org.apache.comet.testing.{DataGenOptions, ParquetGenerator} +import org.apache.comet.testing.{ParquetDataGenOptions, ParquetGenerator} class CometMapExpressionSuite extends CometTestBase { @@ -108,7 +108,7 @@ class CometMapExpressionSuite extends CometTestBase { val filename = path.toString val random = new Random(42) withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - val options = DataGenOptions( + val options = ParquetDataGenOptions( allowNull = false, generateNegativeZero = false, generateArray = true, diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala index d0b1dfb3625..1bf817bd396 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -31,7 +31,7 @@ import org.apache.spark.sql.functions.{avg, count_distinct, sum} import org.apache.spark.sql.internal.SQLConf import org.apache.comet.CometConf -import org.apache.comet.testing.{DataGenOptions, ParquetGenerator} +import org.apache.comet.testing.{ParquetDataGenOptions, ParquetGenerator} /** * Test suite dedicated to Comet native aggregate operator @@ -45,7 +45,7 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { val filename = path.toString val random = new Random(42) withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - ParquetGenerator.makeParquetFile(random, spark, filename, 10000, DataGenOptions()) + ParquetGenerator.makeParquetFile(random, spark, filename, 10000, ParquetDataGenOptions()) } val tableName = "avg_decimal" withTable(tableName) { diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index ab7081e10eb..747bcff79a5 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -50,7 +50,7 @@ import org.apache.spark.unsafe.types.UTF8String import org.apache.comet.{CometConf, ExtendedExplainInfo} import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus} -import org.apache.comet.testing.{DataGenOptions, ParquetGenerator} +import org.apache.comet.testing.{ParquetDataGenOptions, ParquetGenerator} class CometExecSuite extends CometTestBase { @@ -2052,7 +2052,7 @@ class CometExecSuite extends CometTestBase { val filename = path.toString val random = new Random(42) withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - val options = DataGenOptions( + val options = ParquetDataGenOptions( allowNull = true, generateNegativeZero = true, generateArray = true, From 0ebd1439d4dea7b778e216fefd9ebd242a566662 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 22 Oct 2025 09:54:55 -0600 Subject: [PATCH 3/5] refactor --- docs/source/user-guide/latest/configs.md | 1 + .../scala/org/apache/comet/fuzz/Main.scala | 4 +- .../comet/testing/FuzzDataGenerator.scala | 6 +-- .../comet/testing/ParquetGenerator.scala | 53 +++++++++++-------- .../comet/CometArrayExpressionSuite.scala | 22 ++++---- .../comet/CometBitwiseExpressionSuite.scala | 6 +-- .../org/apache/comet/CometFuzzTestBase.scala | 4 +- .../org/apache/comet/CometFuzzTestSuite.scala | 4 +- .../comet/CometMapExpressionSuite.scala | 4 +- .../comet/exec/CometAggregateSuite.scala | 9 +++- .../apache/comet/exec/CometExecSuite.scala | 4 +- 11 files changed, 67 insertions(+), 50 deletions(-) diff --git a/docs/source/user-guide/latest/configs.md b/docs/source/user-guide/latest/configs.md index c4c33437227..bc8fdcb9823 100644 --- a/docs/source/user-guide/latest/configs.md +++ b/docs/source/user-guide/latest/configs.md @@ -268,6 +268,7 @@ These settings can be used to determine which parts of the plan are accelerated | `spark.comet.expression.Reverse.enabled` | Enable Comet acceleration for `Reverse` | true | | `spark.comet.expression.Round.enabled` | Enable Comet acceleration for `Round` | true | | `spark.comet.expression.Second.enabled` | Enable Comet acceleration for `Second` | true | +| `spark.comet.expression.Sha1.enabled` | Enable Comet acceleration for `Sha1` | true | | `spark.comet.expression.Sha2.enabled` | Enable Comet acceleration for `Sha2` | true | | `spark.comet.expression.ShiftLeft.enabled` | Enable Comet acceleration for `ShiftLeft` | true | | `spark.comet.expression.ShiftRight.enabled` | Enable Comet acceleration for `ShiftRight` | true | diff --git a/fuzz-testing/src/main/scala/org/apache/comet/fuzz/Main.scala b/fuzz-testing/src/main/scala/org/apache/comet/fuzz/Main.scala index 74edd940089..633a10ba95b 100644 --- a/fuzz-testing/src/main/scala/org/apache/comet/fuzz/Main.scala +++ b/fuzz-testing/src/main/scala/org/apache/comet/fuzz/Main.scala @@ -26,7 +26,7 @@ import org.rogach.scallop.ScallopOption import org.apache.spark.sql.SparkSession -import org.apache.comet.testing.{ParquetDataGenOptions, ParquetGenerator} +import org.apache.comet.testing.{ParquetGenerator, ParquetGeneratorOptions} class Conf(arguments: Seq[String]) extends ScallopConf(arguments) { object generateData extends Subcommand("data") { @@ -78,7 +78,7 @@ object Main { case Some(seed) => new Random(seed) case None => new Random() } - val options = ParquetDataGenOptions( + val options = ParquetGeneratorOptions( allowNull = true, generateArray = conf.generateData.generateArrays(), generateStruct = conf.generateData.generateStructs(), 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 d75c53259f0..bba38964d65 100644 --- a/spark/src/main/scala/org/apache/comet/testing/FuzzDataGenerator.scala +++ b/spark/src/main/scala/org/apache/comet/testing/FuzzDataGenerator.scala @@ -49,7 +49,7 @@ object FuzzDataGenerator { spark: SparkSession, schema: StructType, numRows: Int, - options: DataGenOptions2): DataFrame = { + options: DataGenOptions): DataFrame = { // generate columnar data val cols: Seq[Seq[Any]] = @@ -67,7 +67,7 @@ object FuzzDataGenerator { r: Random, dataType: DataType, numRows: Int, - options: DataGenOptions2): Seq[Any] = { + options: DataGenOptions): Seq[Any] = { dataType match { case ArrayType(elementType, _) => val values = generateColumn(r, elementType, numRows, options) @@ -179,7 +179,7 @@ object FuzzDataGenerator { } } -case class DataGenOptions2( +case class DataGenOptions( allowNull: Boolean = true, generateNegativeZero: Boolean = true, baseDate: Long = FuzzDataGenerator.defaultBaseDate) diff --git a/spark/src/main/scala/org/apache/comet/testing/ParquetGenerator.scala b/spark/src/main/scala/org/apache/comet/testing/ParquetGenerator.scala index 8d63bf02df2..279c91a1e29 100644 --- a/spark/src/main/scala/org/apache/comet/testing/ParquetGenerator.scala +++ b/spark/src/main/scala/org/apache/comet/testing/ParquetGenerator.scala @@ -27,7 +27,36 @@ import org.apache.spark.sql.types.{DataType, DataTypes, MapType, StructField, St object ParquetGenerator { - def makeParquetSchema(options: ParquetDataGenOptions): StructType = { + /** Generate a Parquet file using a generated schema */ + def makeParquetFile( + r: Random, + spark: SparkSession, + filename: String, + numRows: Int, + options: ParquetGeneratorOptions): Unit = { + val schema = generateSchema(options) + + val dataGenOptions = DataGenOptions( + allowNull = options.allowNull, + generateNegativeZero = options.generateNegativeZero, + baseDate = options.baseDate) + + makeParquetFile(r, spark, filename, schema, numRows, dataGenOptions) + } + + /** Generate a Parquet file using the provided schema */ + def makeParquetFile( + r: Random, + spark: SparkSession, + filename: String, + schema: StructType, + numRows: Int, + options: DataGenOptions): Unit = { + val df = FuzzDataGenerator.generateDataFrame(r, spark, schema, numRows, options) + df.write.mode(SaveMode.Overwrite).parquet(filename) + } + + private def generateSchema(options: ParquetGeneratorOptions): StructType = { val primitiveTypes = options.primitiveTypes val dataTypes = ListBuffer[DataType]() dataTypes.appendAll(primitiveTypes) @@ -67,28 +96,10 @@ object ParquetGenerator { .map(i => StructField(s"c${i._2}", i._1, nullable = true)) StructType(fields.toSeq) } - - def makeParquetFile( - r: Random, - spark: SparkSession, - filename: String, - numRows: Int, - options: ParquetDataGenOptions): Unit = { - - val schema = makeParquetSchema(options) - - val x = DataGenOptions2( - allowNull = options.allowNull, - generateNegativeZero = options.generateNegativeZero, - baseDate = options.baseDate) - - val df = FuzzDataGenerator.generateDataFrame(r, spark, schema, numRows, x) - - df.write.mode(SaveMode.Overwrite).parquet(filename) - } } -case class ParquetDataGenOptions( +/** Schema and Data generation options */ +case class ParquetGeneratorOptions( allowNull: Boolean = true, generateNegativeZero: Boolean = true, baseDate: Long = FuzzDataGenerator.defaultBaseDate, diff --git a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala index 1313e3e6b8d..bd764008578 100644 --- a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala @@ -30,7 +30,7 @@ import org.apache.spark.sql.functions._ import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus} import org.apache.comet.DataTypeSupport.isComplexType import org.apache.comet.serde.{CometArrayExcept, CometArrayRemove, CometArrayReverse, CometFlatten} -import org.apache.comet.testing.{ParquetDataGenOptions, ParquetGenerator} +import org.apache.comet.testing.{ParquetGenerator, ParquetGeneratorOptions} class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { @@ -64,7 +64,7 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp spark, filename, 100, - ParquetDataGenOptions( + ParquetGeneratorOptions( allowNull = true, generateNegativeZero = true, generateArray = false, @@ -95,7 +95,7 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp val filename = path.toString val random = new Random(42) withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - val options = ParquetDataGenOptions( + val options = ParquetGeneratorOptions( allowNull = true, generateNegativeZero = true, generateArray = true, @@ -266,7 +266,7 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp spark, filename, 100, - ParquetDataGenOptions( + ParquetGeneratorOptions( allowNull = true, generateNegativeZero = true, generateArray = true, @@ -310,7 +310,7 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp spark, filename, 100, - ParquetDataGenOptions( + ParquetGeneratorOptions( allowNull = true, generateNegativeZero = true, generateArray = false, @@ -340,7 +340,7 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp spark, filename, 100, - ParquetDataGenOptions( + ParquetGeneratorOptions( allowNull = true, generateNegativeZero = true, generateArray = true, @@ -588,7 +588,7 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp spark, filename, 100, - ParquetDataGenOptions( + ParquetGeneratorOptions( allowNull = true, generateNegativeZero = true, generateArray = false, @@ -622,7 +622,7 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp val filename = path.toString val random = new Random(42) withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - val options = ParquetDataGenOptions( + val options = ParquetGeneratorOptions( allowNull = true, generateNegativeZero = true, generateArray = true, @@ -692,7 +692,7 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp spark, filename, 100, - ParquetDataGenOptions( + ParquetGeneratorOptions( allowNull = true, generateNegativeZero = true, generateArray = false, @@ -720,7 +720,7 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp val filename = path.toString val random = new Random(42) withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - val options = ParquetDataGenOptions( + val options = ParquetGeneratorOptions( allowNull = true, generateNegativeZero = true, generateArray = true, @@ -773,7 +773,7 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp val filename = path.toString val random = new Random(42) withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - val options = ParquetDataGenOptions( + val options = ParquetGeneratorOptions( allowNull = true, generateNegativeZero = true, generateArray = true, diff --git a/spark/src/test/scala/org/apache/comet/CometBitwiseExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometBitwiseExpressionSuite.scala index 8eaec24aaeb..d2e214e29e7 100644 --- a/spark/src/test/scala/org/apache/comet/CometBitwiseExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometBitwiseExpressionSuite.scala @@ -25,7 +25,7 @@ import org.apache.hadoop.fs.Path import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper -import org.apache.comet.testing.{ParquetDataGenOptions, ParquetGenerator} +import org.apache.comet.testing.{ParquetGenerator, ParquetGeneratorOptions} class CometBitwiseExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { @@ -99,7 +99,7 @@ class CometBitwiseExpressionSuite extends CometTestBase with AdaptiveSparkPlanHe spark, filename, 100, - ParquetDataGenOptions( + ParquetGeneratorOptions( allowNull = true, generateNegativeZero = true, generateArray = false, @@ -174,7 +174,7 @@ class CometBitwiseExpressionSuite extends CometTestBase with AdaptiveSparkPlanHe spark, filename, 10, - ParquetDataGenOptions( + ParquetGeneratorOptions( allowNull = true, generateNegativeZero = true, generateArray = false, diff --git a/spark/src/test/scala/org/apache/comet/CometFuzzTestBase.scala b/spark/src/test/scala/org/apache/comet/CometFuzzTestBase.scala index 9068e5b74d3..9ab35c822fc 100644 --- a/spark/src/test/scala/org/apache/comet/CometFuzzTestBase.scala +++ b/spark/src/test/scala/org/apache/comet/CometFuzzTestBase.scala @@ -35,7 +35,7 @@ import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.internal.SQLConf -import org.apache.comet.testing.{ParquetDataGenOptions, ParquetGenerator} +import org.apache.comet.testing.{ParquetGenerator, ParquetGeneratorOptions} class CometFuzzTestBase extends CometTestBase with AdaptiveSparkPlanHelper { @@ -59,7 +59,7 @@ class CometFuzzTestBase extends CometTestBase with AdaptiveSparkPlanHelper { CometConf.COMET_ENABLED.key -> "false", SQLConf.SESSION_LOCAL_TIMEZONE.key -> defaultTimezone) { val options = - ParquetDataGenOptions( + ParquetGeneratorOptions( generateArray = true, generateStruct = true, generateNegativeZero = false, diff --git a/spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala b/spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala index a6e55446b49..1a4196710a9 100644 --- a/spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala @@ -28,7 +28,7 @@ import org.apache.spark.sql.internal.SQLConf.ParquetOutputTimestampType import org.apache.spark.sql.types._ import org.apache.comet.DataTypeSupport.isComplexType -import org.apache.comet.testing.{ParquetDataGenOptions, ParquetGenerator} +import org.apache.comet.testing.{ParquetGenerator, ParquetGeneratorOptions} class CometFuzzTestSuite extends CometFuzzTestBase { @@ -262,7 +262,7 @@ class CometFuzzTestSuite extends CometFuzzTestBase { generateStruct: Boolean = true): Unit = { val options = - ParquetDataGenOptions( + ParquetGeneratorOptions( generateArray = generateArray, generateStruct = generateStruct, generateNegativeZero = false) diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index 7283b696b4f..a9eb057f821 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -26,7 +26,7 @@ import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -import org.apache.comet.testing.{ParquetDataGenOptions, ParquetGenerator} +import org.apache.comet.testing.{ParquetGenerator, ParquetGeneratorOptions} class CometMapExpressionSuite extends CometTestBase { @@ -108,7 +108,7 @@ class CometMapExpressionSuite extends CometTestBase { val filename = path.toString val random = new Random(42) withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - val options = ParquetDataGenOptions( + val options = ParquetGeneratorOptions( allowNull = false, generateNegativeZero = false, generateArray = true, diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala index 1bf817bd396..f64d8b02cd7 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -31,7 +31,7 @@ import org.apache.spark.sql.functions.{avg, count_distinct, sum} import org.apache.spark.sql.internal.SQLConf import org.apache.comet.CometConf -import org.apache.comet.testing.{ParquetDataGenOptions, ParquetGenerator} +import org.apache.comet.testing.{ParquetGenerator, ParquetGeneratorOptions} /** * Test suite dedicated to Comet native aggregate operator @@ -45,7 +45,12 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { val filename = path.toString val random = new Random(42) withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - ParquetGenerator.makeParquetFile(random, spark, filename, 10000, ParquetDataGenOptions()) + ParquetGenerator.makeParquetFile( + random, + spark, + filename, + 10000, + ParquetGeneratorOptions()) } val tableName = "avg_decimal" withTable(tableName) { diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index 747bcff79a5..0fcc9c68e38 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -50,7 +50,7 @@ import org.apache.spark.unsafe.types.UTF8String import org.apache.comet.{CometConf, ExtendedExplainInfo} import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus} -import org.apache.comet.testing.{ParquetDataGenOptions, ParquetGenerator} +import org.apache.comet.testing.{ParquetGenerator, ParquetGeneratorOptions} class CometExecSuite extends CometTestBase { @@ -2052,7 +2052,7 @@ class CometExecSuite extends CometTestBase { val filename = path.toString val random = new Random(42) withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - val options = ParquetDataGenOptions( + val options = ParquetGeneratorOptions( allowNull = true, generateNegativeZero = true, generateArray = true, From 3ffa0dfdca3039155b2a46642194946c7ffa70c1 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 22 Oct 2025 10:43:02 -0600 Subject: [PATCH 4/5] Refactor --- .../comet/testing/FuzzDataGenerator.scala | 65 ++++++++++++++++++ .../comet/testing/ParquetGenerator.scala | 67 +++---------------- 2 files changed, 74 insertions(+), 58 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 bba38964d65..087221e1a30 100644 --- a/spark/src/main/scala/org/apache/comet/testing/FuzzDataGenerator.scala +++ b/spark/src/main/scala/org/apache/comet/testing/FuzzDataGenerator.scala @@ -44,6 +44,47 @@ object FuzzDataGenerator { val defaultBaseDate: Long = new SimpleDateFormat("YYYY-MM-DD hh:mm:ss").parse("3333-05-25 12:34:56").getTime + def generateSchema(options: SchemaGenOptions): StructType = { + val primitiveTypes = options.primitiveTypes + val dataTypes = ListBuffer[DataType]() + dataTypes.appendAll(primitiveTypes) + + val arraysOfPrimitives = primitiveTypes.map(DataTypes.createArrayType) + + if (options.generateStruct) { + dataTypes += StructType( + primitiveTypes.zipWithIndex.map(x => StructField(s"c${x._2}", x._1, nullable = true))) + + if (options.generateArray) { + dataTypes += StructType(arraysOfPrimitives.zipWithIndex.map(x => + StructField(s"c${x._2}", x._1, nullable = true))) + } + } + + if (options.generateMap) { + dataTypes += MapType(DataTypes.IntegerType, DataTypes.StringType) + } + + if (options.generateArray) { + dataTypes.appendAll(arraysOfPrimitives) + + if (options.generateStruct) { + dataTypes += DataTypes.createArrayType(StructType(primitiveTypes.zipWithIndex.map(x => + StructField(s"c${x._2}", x._1, nullable = true)))) + } + + if (options.generateMap) { + dataTypes += DataTypes.createArrayType( + MapType(DataTypes.IntegerType, DataTypes.StringType)) + } + } + + // generate schema using random data types + val fields = dataTypes.zipWithIndex + .map(i => StructField(s"c${i._2}", i._1, nullable = true)) + StructType(fields.toSeq) + } + def generateDataFrame( r: Random, spark: SparkSession, @@ -179,6 +220,30 @@ object FuzzDataGenerator { } } +object SchemaGenOptions { + val defaultPrimitiveTypes: Seq[DataType] = Seq( + DataTypes.BooleanType, + DataTypes.ByteType, + DataTypes.ShortType, + DataTypes.IntegerType, + DataTypes.LongType, + DataTypes.FloatType, + DataTypes.DoubleType, + DataTypes.createDecimalType(10, 2), + DataTypes.createDecimalType(36, 18), + DataTypes.DateType, + DataTypes.TimestampType, + DataTypes.TimestampNTZType, + DataTypes.StringType, + DataTypes.BinaryType) +} + +case class SchemaGenOptions( + generateArray: Boolean = false, + generateStruct: Boolean = false, + generateMap: Boolean = false, + primitiveTypes: Seq[DataType] = SchemaGenOptions.defaultPrimitiveTypes) + case class DataGenOptions( allowNull: Boolean = true, generateNegativeZero: Boolean = true, diff --git a/spark/src/main/scala/org/apache/comet/testing/ParquetGenerator.scala b/spark/src/main/scala/org/apache/comet/testing/ParquetGenerator.scala index 279c91a1e29..75912169918 100644 --- a/spark/src/main/scala/org/apache/comet/testing/ParquetGenerator.scala +++ b/spark/src/main/scala/org/apache/comet/testing/ParquetGenerator.scala @@ -19,11 +19,10 @@ package org.apache.comet.testing -import scala.collection.mutable.ListBuffer import scala.util.Random import org.apache.spark.sql.{SaveMode, SparkSession} -import org.apache.spark.sql.types.{DataType, DataTypes, MapType, StructField, StructType} +import org.apache.spark.sql.types.{DataType, StructType} object ParquetGenerator { @@ -34,7 +33,13 @@ object ParquetGenerator { filename: String, numRows: Int, options: ParquetGeneratorOptions): Unit = { - val schema = generateSchema(options) + + val schemaGenOptions = SchemaGenOptions( + generateArray = options.generateArray, + generateStruct = options.generateStruct, + generateMap = options.generateMap, + primitiveTypes = options.primitiveTypes) + val schema = FuzzDataGenerator.generateSchema(schemaGenOptions) val dataGenOptions = DataGenOptions( allowNull = options.allowNull, @@ -56,46 +61,6 @@ object ParquetGenerator { df.write.mode(SaveMode.Overwrite).parquet(filename) } - private def generateSchema(options: ParquetGeneratorOptions): StructType = { - val primitiveTypes = options.primitiveTypes - val dataTypes = ListBuffer[DataType]() - dataTypes.appendAll(primitiveTypes) - - val arraysOfPrimitives = primitiveTypes.map(DataTypes.createArrayType) - - if (options.generateStruct) { - dataTypes += StructType( - primitiveTypes.zipWithIndex.map(x => StructField(s"c${x._2}", x._1, nullable = true))) - - if (options.generateArray) { - dataTypes += StructType(arraysOfPrimitives.zipWithIndex.map(x => - StructField(s"c${x._2}", x._1, nullable = true))) - } - } - - if (options.generateMap) { - dataTypes += MapType(DataTypes.IntegerType, DataTypes.StringType) - } - - if (options.generateArray) { - dataTypes.appendAll(arraysOfPrimitives) - - if (options.generateStruct) { - dataTypes += DataTypes.createArrayType(StructType(primitiveTypes.zipWithIndex.map(x => - StructField(s"c${x._2}", x._1, nullable = true)))) - } - - if (options.generateMap) { - dataTypes += DataTypes.createArrayType( - MapType(DataTypes.IntegerType, DataTypes.StringType)) - } - } - - // generate schema using random data types - val fields = dataTypes.zipWithIndex - .map(i => StructField(s"c${i._2}", i._1, nullable = true)) - StructType(fields.toSeq) - } } /** Schema and Data generation options */ @@ -106,18 +71,4 @@ case class ParquetGeneratorOptions( generateArray: Boolean = false, generateStruct: Boolean = false, generateMap: Boolean = false, - primitiveTypes: Seq[DataType] = Seq( - DataTypes.BooleanType, - DataTypes.ByteType, - DataTypes.ShortType, - DataTypes.IntegerType, - DataTypes.LongType, - DataTypes.FloatType, - DataTypes.DoubleType, - DataTypes.createDecimalType(10, 2), - DataTypes.createDecimalType(36, 18), - DataTypes.DateType, - DataTypes.TimestampType, - DataTypes.TimestampNTZType, - DataTypes.StringType, - DataTypes.BinaryType)) + primitiveTypes: Seq[DataType] = SchemaGenOptions.defaultPrimitiveTypes) From fccb7e697b261be7425e0a1e471e19a826907aa2 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 22 Oct 2025 11:01:50 -0600 Subject: [PATCH 5/5] simplify --- .../scala/org/apache/comet/fuzz/Main.scala | 16 +-- .../comet/testing/ParquetGenerator.scala | 27 +---- .../comet/CometArrayExpressionSuite.scala | 106 +++++++----------- .../comet/CometBitwiseExpressionSuite.scala | 18 +-- .../org/apache/comet/CometFuzzTestBase.scala | 25 +++-- .../org/apache/comet/CometFuzzTestSuite.scala | 19 ++-- .../comet/CometMapExpressionSuite.scala | 19 ++-- .../comet/exec/CometAggregateSuite.scala | 5 +- .../apache/comet/exec/CometExecSuite.scala | 19 ++-- 9 files changed, 109 insertions(+), 145 deletions(-) diff --git a/fuzz-testing/src/main/scala/org/apache/comet/fuzz/Main.scala b/fuzz-testing/src/main/scala/org/apache/comet/fuzz/Main.scala index 633a10ba95b..1f81dc77919 100644 --- a/fuzz-testing/src/main/scala/org/apache/comet/fuzz/Main.scala +++ b/fuzz-testing/src/main/scala/org/apache/comet/fuzz/Main.scala @@ -26,7 +26,7 @@ import org.rogach.scallop.ScallopOption import org.apache.spark.sql.SparkSession -import org.apache.comet.testing.{ParquetGenerator, ParquetGeneratorOptions} +import org.apache.comet.testing.{DataGenOptions, ParquetGenerator, SchemaGenOptions} class Conf(arguments: Seq[String]) extends ScallopConf(arguments) { object generateData extends Subcommand("data") { @@ -78,19 +78,19 @@ object Main { case Some(seed) => new Random(seed) case None => new Random() } - val options = ParquetGeneratorOptions( - allowNull = true, - generateArray = conf.generateData.generateArrays(), - generateStruct = conf.generateData.generateStructs(), - generateMap = conf.generateData.generateMaps(), - generateNegativeZero = !conf.generateData.excludeNegativeZero()) for (i <- 0 until conf.generateData.numFiles()) { ParquetGenerator.makeParquetFile( r, spark, s"test$i.parquet", numRows = conf.generateData.numRows(), - options) + SchemaGenOptions( + generateArray = conf.generateData.generateArrays(), + generateStruct = conf.generateData.generateStructs(), + generateMap = conf.generateData.generateMaps()), + DataGenOptions( + allowNull = true, + generateNegativeZero = !conf.generateData.excludeNegativeZero())) } case Some(conf.generateQueries) => val r = conf.generateQueries.randomSeed.toOption match { diff --git a/spark/src/main/scala/org/apache/comet/testing/ParquetGenerator.scala b/spark/src/main/scala/org/apache/comet/testing/ParquetGenerator.scala index 75912169918..a43100a69d4 100644 --- a/spark/src/main/scala/org/apache/comet/testing/ParquetGenerator.scala +++ b/spark/src/main/scala/org/apache/comet/testing/ParquetGenerator.scala @@ -22,7 +22,7 @@ package org.apache.comet.testing import scala.util.Random import org.apache.spark.sql.{SaveMode, SparkSession} -import org.apache.spark.sql.types.{DataType, StructType} +import org.apache.spark.sql.types.StructType object ParquetGenerator { @@ -32,20 +32,9 @@ object ParquetGenerator { spark: SparkSession, filename: String, numRows: Int, - options: ParquetGeneratorOptions): Unit = { - - val schemaGenOptions = SchemaGenOptions( - generateArray = options.generateArray, - generateStruct = options.generateStruct, - generateMap = options.generateMap, - primitiveTypes = options.primitiveTypes) + schemaGenOptions: SchemaGenOptions, + dataGenOptions: DataGenOptions): Unit = { val schema = FuzzDataGenerator.generateSchema(schemaGenOptions) - - val dataGenOptions = DataGenOptions( - allowNull = options.allowNull, - generateNegativeZero = options.generateNegativeZero, - baseDate = options.baseDate) - makeParquetFile(r, spark, filename, schema, numRows, dataGenOptions) } @@ -62,13 +51,3 @@ object ParquetGenerator { } } - -/** Schema and Data generation options */ -case class ParquetGeneratorOptions( - allowNull: Boolean = true, - generateNegativeZero: Boolean = true, - baseDate: Long = FuzzDataGenerator.defaultBaseDate, - generateArray: Boolean = false, - generateStruct: Boolean = false, - generateMap: Boolean = false, - primitiveTypes: Seq[DataType] = SchemaGenOptions.defaultPrimitiveTypes) diff --git a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala index bd764008578..c346dc2e951 100644 --- a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala @@ -30,7 +30,7 @@ import org.apache.spark.sql.functions._ import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus} import org.apache.comet.DataTypeSupport.isComplexType import org.apache.comet.serde.{CometArrayExcept, CometArrayRemove, CometArrayReverse, CometFlatten} -import org.apache.comet.testing.{ParquetGenerator, ParquetGeneratorOptions} +import org.apache.comet.testing.{DataGenOptions, ParquetGenerator, SchemaGenOptions} class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { @@ -64,12 +64,8 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp spark, filename, 100, - ParquetGeneratorOptions( - allowNull = true, - generateNegativeZero = true, - generateArray = false, - generateStruct = false, - generateMap = false)) + SchemaGenOptions(generateArray = false, generateStruct = false, generateMap = false), + DataGenOptions(allowNull = true, generateNegativeZero = true)) } val table = spark.read.parquet(filename) table.createOrReplaceTempView("t1") @@ -95,13 +91,13 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp val filename = path.toString val random = new Random(42) withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - val options = ParquetGeneratorOptions( - allowNull = true, - generateNegativeZero = true, - generateArray = true, - generateStruct = true, - generateMap = false) - ParquetGenerator.makeParquetFile(random, spark, filename, 100, options) + ParquetGenerator.makeParquetFile( + random, + spark, + filename, + 100, + SchemaGenOptions(generateArray = true, generateStruct = true, generateMap = false), + DataGenOptions(allowNull = true, generateNegativeZero = true)) } withSQLConf( CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "false", @@ -266,12 +262,8 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp spark, filename, 100, - ParquetGeneratorOptions( - allowNull = true, - generateNegativeZero = true, - generateArray = true, - generateStruct = true, - generateMap = false)) + SchemaGenOptions(generateArray = true, generateStruct = true, generateMap = false), + DataGenOptions(allowNull = true, generateNegativeZero = true)) } val table = spark.read.parquet(filename) table.createOrReplaceTempView("t1") @@ -310,12 +302,8 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp spark, filename, 100, - ParquetGeneratorOptions( - allowNull = true, - generateNegativeZero = true, - generateArray = false, - generateStruct = false, - generateMap = false)) + SchemaGenOptions(generateArray = false, generateStruct = false, generateMap = false), + DataGenOptions(allowNull = true, generateNegativeZero = true)) } val table = spark.read.parquet(filename) table.createOrReplaceTempView("t2") @@ -340,12 +328,8 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp spark, filename, 100, - ParquetGeneratorOptions( - allowNull = true, - generateNegativeZero = true, - generateArray = true, - generateStruct = true, - generateMap = false)) + SchemaGenOptions(generateArray = true, generateStruct = true, generateMap = false), + DataGenOptions(allowNull = true, generateNegativeZero = true)) } withSQLConf( CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "false", @@ -588,12 +572,8 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp spark, filename, 100, - ParquetGeneratorOptions( - allowNull = true, - generateNegativeZero = true, - generateArray = false, - generateStruct = false, - generateMap = false)) + SchemaGenOptions(generateArray = false, generateStruct = false, generateMap = false), + DataGenOptions(allowNull = true, generateNegativeZero = true)) } withSQLConf(CometConf.COMET_EXPR_ALLOW_INCOMPATIBLE.key -> "true") { withTempView("t1", "t2") { @@ -622,13 +602,13 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp val filename = path.toString val random = new Random(42) withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - val options = ParquetGeneratorOptions( - allowNull = true, - generateNegativeZero = true, - generateArray = true, - generateStruct = true, - generateMap = false) - ParquetGenerator.makeParquetFile(random, spark, filename, 100, options) + ParquetGenerator.makeParquetFile( + random, + spark, + filename, + 100, + SchemaGenOptions(generateArray = true, generateStruct = true, generateMap = false), + DataGenOptions(allowNull = true, generateNegativeZero = true)) } withSQLConf( CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "false", @@ -692,12 +672,8 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp spark, filename, 100, - ParquetGeneratorOptions( - allowNull = true, - generateNegativeZero = true, - generateArray = false, - generateStruct = false, - generateMap = false)) + SchemaGenOptions(generateArray = false, generateStruct = false, generateMap = false), + DataGenOptions(allowNull = true, generateNegativeZero = true)) } val table = spark.read.parquet(filename) table.createOrReplaceTempView("t1") @@ -720,13 +696,13 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp val filename = path.toString val random = new Random(42) withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - val options = ParquetGeneratorOptions( - allowNull = true, - generateNegativeZero = true, - generateArray = true, - generateStruct = true, - generateMap = false) - ParquetGenerator.makeParquetFile(random, spark, filename, 100, options) + ParquetGenerator.makeParquetFile( + random, + spark, + filename, + 100, + SchemaGenOptions(generateArray = true, generateStruct = true, generateMap = false), + DataGenOptions(allowNull = true, generateNegativeZero = true)) } withSQLConf( CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "false", @@ -773,13 +749,13 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp val filename = path.toString val random = new Random(42) withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - val options = ParquetGeneratorOptions( - allowNull = true, - generateNegativeZero = true, - generateArray = true, - generateStruct = true, - generateMap = false) - ParquetGenerator.makeParquetFile(random, spark, filename, 100, options) + ParquetGenerator.makeParquetFile( + random, + spark, + filename, + 100, + SchemaGenOptions(generateArray = true, generateStruct = true, generateMap = false), + DataGenOptions(allowNull = true, generateNegativeZero = true)) } withSQLConf( CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "false", diff --git a/spark/src/test/scala/org/apache/comet/CometBitwiseExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometBitwiseExpressionSuite.scala index d2e214e29e7..02c003ede80 100644 --- a/spark/src/test/scala/org/apache/comet/CometBitwiseExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometBitwiseExpressionSuite.scala @@ -25,7 +25,7 @@ import org.apache.hadoop.fs.Path import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper -import org.apache.comet.testing.{ParquetGenerator, ParquetGeneratorOptions} +import org.apache.comet.testing.{DataGenOptions, ParquetGenerator, SchemaGenOptions} class CometBitwiseExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { @@ -99,12 +99,8 @@ class CometBitwiseExpressionSuite extends CometTestBase with AdaptiveSparkPlanHe spark, filename, 100, - ParquetGeneratorOptions( - allowNull = true, - generateNegativeZero = true, - generateArray = false, - generateStruct = false, - generateMap = false)) + SchemaGenOptions(generateArray = false, generateStruct = false, generateMap = false), + DataGenOptions(allowNull = true, generateNegativeZero = true)) } val table = spark.read.parquet(filename) checkSparkAnswerAndOperator( @@ -174,12 +170,8 @@ class CometBitwiseExpressionSuite extends CometTestBase with AdaptiveSparkPlanHe spark, filename, 10, - ParquetGeneratorOptions( - allowNull = true, - generateNegativeZero = true, - generateArray = false, - generateStruct = false, - generateMap = false)) + SchemaGenOptions(generateArray = false, generateStruct = false, generateMap = false), + DataGenOptions(allowNull = true, generateNegativeZero = true)) } val table = spark.read.parquet(filename) val df = diff --git a/spark/src/test/scala/org/apache/comet/CometFuzzTestBase.scala b/spark/src/test/scala/org/apache/comet/CometFuzzTestBase.scala index 9ab35c822fc..1c0636780e3 100644 --- a/spark/src/test/scala/org/apache/comet/CometFuzzTestBase.scala +++ b/spark/src/test/scala/org/apache/comet/CometFuzzTestBase.scala @@ -35,7 +35,7 @@ import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.internal.SQLConf -import org.apache.comet.testing.{ParquetGenerator, ParquetGeneratorOptions} +import org.apache.comet.testing.{DataGenOptions, ParquetGenerator, SchemaGenOptions} class CometFuzzTestBase extends CometTestBase with AdaptiveSparkPlanHelper { @@ -58,15 +58,20 @@ class CometFuzzTestBase extends CometTestBase with AdaptiveSparkPlanHelper { withSQLConf( CometConf.COMET_ENABLED.key -> "false", SQLConf.SESSION_LOCAL_TIMEZONE.key -> defaultTimezone) { - val options = - ParquetGeneratorOptions( - generateArray = true, - generateStruct = true, - generateNegativeZero = false, - // override base date due to known issues with experimental scans - baseDate = - new SimpleDateFormat("YYYY-MM-DD hh:mm:ss").parse("2024-05-25 12:34:56").getTime) - ParquetGenerator.makeParquetFile(random, spark, filename, 1000, options) + val schemaGenOptions = + SchemaGenOptions(generateArray = true, generateStruct = true) + val dataGenOptions = DataGenOptions( + generateNegativeZero = false, + // override base date due to known issues with experimental scans + baseDate = + new SimpleDateFormat("YYYY-MM-DD hh:mm:ss").parse("2024-05-25 12:34:56").getTime) + ParquetGenerator.makeParquetFile( + random, + spark, + filename, + 1000, + schemaGenOptions, + dataGenOptions) } } diff --git a/spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala b/spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala index 1a4196710a9..006112d2b08 100644 --- a/spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala @@ -28,7 +28,7 @@ import org.apache.spark.sql.internal.SQLConf.ParquetOutputTimestampType import org.apache.spark.sql.types._ import org.apache.comet.DataTypeSupport.isComplexType -import org.apache.comet.testing.{ParquetGenerator, ParquetGeneratorOptions} +import org.apache.comet.testing.{DataGenOptions, ParquetGenerator, SchemaGenOptions} class CometFuzzTestSuite extends CometFuzzTestBase { @@ -261,11 +261,10 @@ class CometFuzzTestSuite extends CometFuzzTestBase { generateArray: Boolean = true, generateStruct: Boolean = true): Unit = { - val options = - ParquetGeneratorOptions( - generateArray = generateArray, - generateStruct = generateStruct, - generateNegativeZero = false) + val schemaGenOptions = + SchemaGenOptions(generateArray = generateArray, generateStruct = generateStruct) + + val dataGenOptions = DataGenOptions(generateNegativeZero = false) withTempPath { filename => val random = new Random(42) @@ -273,7 +272,13 @@ class CometFuzzTestSuite extends CometFuzzTestBase { CometConf.COMET_ENABLED.key -> "false", SQLConf.PARQUET_OUTPUT_TIMESTAMP_TYPE.key -> outputTimestampType.toString, SQLConf.SESSION_LOCAL_TIMEZONE.key -> defaultTimezone) { - ParquetGenerator.makeParquetFile(random, spark, filename.toString, 100, options) + ParquetGenerator.makeParquetFile( + random, + spark, + filename.toString, + 100, + schemaGenOptions, + dataGenOptions) } Seq(defaultTimezone, "UTC", "America/Denver").foreach { tz => diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index a9eb057f821..88c13391a67 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -26,7 +26,7 @@ import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -import org.apache.comet.testing.{ParquetGenerator, ParquetGeneratorOptions} +import org.apache.comet.testing.{DataGenOptions, ParquetGenerator, SchemaGenOptions} class CometMapExpressionSuite extends CometTestBase { @@ -108,13 +108,16 @@ class CometMapExpressionSuite extends CometTestBase { val filename = path.toString val random = new Random(42) withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - val options = ParquetGeneratorOptions( - allowNull = false, - generateNegativeZero = false, - generateArray = true, - generateStruct = false, - generateMap = false) - ParquetGenerator.makeParquetFile(random, spark, filename, 100, options) + val schemaGenOptions = + SchemaGenOptions(generateArray = true, generateStruct = false, generateMap = false) + val dataGenOptions = DataGenOptions(allowNull = false, generateNegativeZero = false) + ParquetGenerator.makeParquetFile( + random, + spark, + filename, + 100, + schemaGenOptions, + dataGenOptions) } spark.read.parquet(filename).createOrReplaceTempView("t1") val df = spark.sql("SELECT map_from_arrays(array(c12), array(c3)) FROM t1") diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala index f64d8b02cd7..211cc16d05a 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -31,7 +31,7 @@ import org.apache.spark.sql.functions.{avg, count_distinct, sum} import org.apache.spark.sql.internal.SQLConf import org.apache.comet.CometConf -import org.apache.comet.testing.{ParquetGenerator, ParquetGeneratorOptions} +import org.apache.comet.testing.{DataGenOptions, ParquetGenerator, SchemaGenOptions} /** * Test suite dedicated to Comet native aggregate operator @@ -50,7 +50,8 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { spark, filename, 10000, - ParquetGeneratorOptions()) + SchemaGenOptions(), + DataGenOptions()) } val tableName = "avg_decimal" withTable(tableName) { diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index 0fcc9c68e38..1b15c39cac5 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -50,7 +50,7 @@ import org.apache.spark.unsafe.types.UTF8String import org.apache.comet.{CometConf, ExtendedExplainInfo} import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus} -import org.apache.comet.testing.{ParquetGenerator, ParquetGeneratorOptions} +import org.apache.comet.testing.{DataGenOptions, ParquetGenerator, SchemaGenOptions} class CometExecSuite extends CometTestBase { @@ -2052,13 +2052,16 @@ class CometExecSuite extends CometTestBase { val filename = path.toString val random = new Random(42) withSQLConf(CometConf.COMET_ENABLED.key -> "false") { - val options = ParquetGeneratorOptions( - allowNull = true, - generateNegativeZero = true, - generateArray = true, - generateStruct = true, - generateMap = true) - ParquetGenerator.makeParquetFile(random, spark, filename, 100, options) + val schemaGenOptions = + SchemaGenOptions(generateArray = true, generateStruct = true, generateMap = true) + val dataGenOptions = DataGenOptions(allowNull = true, generateNegativeZero = true) + ParquetGenerator.makeParquetFile( + random, + spark, + filename, + 100, + schemaGenOptions, + dataGenOptions) } withSQLConf( CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "false",