Merge remote-tracking branch 'apache/main' into cast-to-int-perf

Author: andygrove

native/core/src/execution/operators/iceberg_scan.rs

@@ -369,11 +369,7 @@ impl IcebergFileStream {
                         }
                     }
 
-                    match ready!(self
-                        .metrics
-                        .baseline
-                        .record_poll(current.poll_next_unpin(cx)))
-                    {
+                    match ready!(current.poll_next_unpin(cx)) {
                         Some(result) => {
                             // Stop time_scanning_until_data on first batch (idempotent)
                             self.metrics.file_stream.time_scanning_until_data.stop();
@@ -428,7 +424,7 @@ impl Stream for IcebergFileStream {
         self.metrics.file_stream.time_processing.start();
         let result = self.poll_inner(cx);
         self.metrics.file_stream.time_processing.stop();
-        result
+        self.metrics.baseline.record_poll(result)
     }
 }
 

spark/src/main/scala/org/apache/comet/serde/structs.scala

@@ -111,26 +111,6 @@ object CometStructsToJson extends CometExpressionSerde[StructsToJson] {
       withInfo(expr, "StructsToJson with options is not supported")
       None
     } else {
-
-      def isSupportedType(dt: DataType): Boolean = {
-        dt match {
-          case StructType(fields) =>
-            fields.forall(f => isSupportedType(f.dataType))
-          case DataTypes.BooleanType | DataTypes.ByteType | DataTypes.ShortType |
-              DataTypes.IntegerType | DataTypes.LongType | DataTypes.FloatType |
-              DataTypes.DoubleType | DataTypes.StringType =>
-            true
-          case DataTypes.DateType | DataTypes.TimestampType =>
-            // TODO implement these types with tests for formatting options and timezone
-            false
-          case _: MapType | _: ArrayType =>
-            // Spark supports map and array in StructsToJson but this is not yet
-            // implemented in Comet
-            false
-          case _ => false
-        }
-      }
-
       val isSupported = expr.child.dataType match {
         case s: StructType =>
           s.fields.forall(f => isSupportedType(f.dataType))
@@ -166,6 +146,25 @@ object CometStructsToJson extends CometExpressionSerde[StructsToJson] {
       }
     }
   }
+
+  def isSupportedType(dt: DataType): Boolean = {
+    dt match {
+      case StructType(fields) =>
+        fields.forall(f => isSupportedType(f.dataType))
+      case DataTypes.BooleanType | DataTypes.ByteType | DataTypes.ShortType |
+          DataTypes.IntegerType | DataTypes.LongType | DataTypes.FloatType |
+          DataTypes.DoubleType | DataTypes.StringType =>
+        true
+      case DataTypes.DateType | DataTypes.TimestampType =>
+        // TODO implement these types with tests for formatting options and timezone
+        false
+      case _: MapType | _: ArrayType =>
+        // Spark supports map and array in StructsToJson but this is not yet
+        // implemented in Comet
+        false
+      case _ => false
+    }
+  }
 }
 
 object CometJsonToStructs extends CometExpressionSerde[JsonToStructs] {

spark/src/main/scala/org/apache/comet/testing/FuzzDataGenerator.scala

@@ -229,8 +229,8 @@ object FuzzDataGenerator {
         Range(0, numRows).map(_ => {
           r.nextInt(20) match {
             case 0 if options.allowNull => null
-            case 1 => Float.NegativeInfinity
-            case 2 => Float.PositiveInfinity
+            case 1 if options.generateInfinity => Float.NegativeInfinity
+            case 2 if options.generateInfinity => Float.PositiveInfinity
             case 3 => Float.MinValue
             case 4 => Float.MaxValue
             case 5 => 0.0f
@@ -243,8 +243,8 @@ object FuzzDataGenerator {
         Range(0, numRows).map(_ => {
           r.nextInt(20) match {
             case 0 if options.allowNull => null
-            case 1 => Double.NegativeInfinity
-            case 2 => Double.PositiveInfinity
+            case 1 if options.generateInfinity => Double.NegativeInfinity
+            case 2 if options.generateInfinity => Double.PositiveInfinity
             case 3 => Double.MinValue
             case 4 => Double.MaxValue
             case 5 => 0.0
@@ -329,4 +329,5 @@ case class DataGenOptions(
     generateNaN: Boolean = true,
     baseDate: Long = FuzzDataGenerator.defaultBaseDate,
     customStrings: Seq[String] = Seq.empty,
-    maxStringLength: Int = 8)
+    maxStringLength: Int = 8,
+    generateInfinity: Boolean = true)

spark/src/test/scala/org/apache/comet/CometJsonExpressionSuite.scala

@@ -19,24 +19,59 @@
 
 package org.apache.comet
 
+import scala.util.Random
+
 import org.scalactic.source.Position
 import org.scalatest.Tag
 
+import org.apache.hadoop.fs.Path
 import org.apache.spark.sql.CometTestBase
-import org.apache.spark.sql.catalyst.expressions.JsonToStructs
+import org.apache.spark.sql.catalyst.expressions.{JsonToStructs, StructsToJson}
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.functions._
+
+import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus
+import org.apache.comet.serde.CometStructsToJson
+import org.apache.comet.testing.{DataGenOptions, ParquetGenerator, SchemaGenOptions}
 
 class CometJsonExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
 
   override protected def test(testName: String, testTags: Tag*)(testFun: => Any)(implicit
       pos: Position): Unit = {
     super.test(testName, testTags: _*) {
-      withSQLConf(CometConf.getExprAllowIncompatConfigKey(classOf[JsonToStructs]) -> "true") {
+      withSQLConf(
+        CometConf.getExprAllowIncompatConfigKey(classOf[JsonToStructs]) -> "true",
+        CometConf.getExprAllowIncompatConfigKey(classOf[StructsToJson]) -> "true") {
         testFun
       }
     }
   }
 
+  test("to_json - all supported types") {
+    assume(!isSpark40Plus)
+    withTempDir { dir =>
+      val path = new Path(dir.toURI.toString, "test.parquet")
+      val filename = path.toString
+      val random = new Random(42)
+      withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+        ParquetGenerator.makeParquetFile(
+          random,
+          spark,
+          filename,
+          100,
+          SchemaGenOptions(generateArray = false, generateStruct = false, generateMap = false),
+          DataGenOptions(generateNaN = false, generateInfinity = false))
+      }
+      val table = spark.read.parquet(filename)
+      val fieldsNames = table.schema.fields
+        .filter(sf => CometStructsToJson.isSupportedType(sf.dataType))
+        .map(sf => col(sf.name))
+        .toSeq
+      val df = table.select(to_json(struct(fieldsNames: _*)))
+      checkSparkAnswerAndOperator(df)
+    }
+  }
+
   test("from_json - basic primitives") {
     Seq(true, false).foreach { dictionaryEnabled =>
       withParquetTable(

spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala

@@ -37,7 +37,7 @@ import org.apache.parquet.hadoop.example.{ExampleParquetWriter, GroupWriteSuppor
 import org.apache.parquet.schema.{MessageType, MessageTypeParser}
 import org.apache.spark._
 import org.apache.spark.internal.config.{MEMORY_OFFHEAP_ENABLED, MEMORY_OFFHEAP_SIZE, SHUFFLE_MANAGER}
-import org.apache.spark.sql.comet._
+import org.apache.spark.sql.comet.CometPlanChecker
 import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec}
 import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
@@ -58,7 +58,8 @@ abstract class CometTestBase
     with BeforeAndAfterEach
     with AdaptiveSparkPlanHelper
     with ShimCometSparkSessionExtensions
-    with ShimCometTestBase {
+    with ShimCometTestBase
+    with CometPlanChecker {
   import testImplicits._
 
   protected val shuffleManager: String =
@@ -396,26 +397,6 @@ abstract class CometTestBase
     checkPlanNotMissingInput(plan)
   }
 
-  protected def findFirstNonCometOperator(
-      plan: SparkPlan,
-      excludedClasses: Class[_]*): Option[SparkPlan] = {
-    val wrapped = wrapCometSparkToColumnar(plan)
-    wrapped.foreach {
-      case _: CometNativeScanExec | _: CometScanExec | _: CometBatchScanExec |
-          _: CometIcebergNativeScanExec =>
-      case _: CometSinkPlaceHolder | _: CometScanWrapper =>
-      case _: CometColumnarToRowExec =>
-      case _: CometSparkToColumnarExec =>
-      case _: CometExec | _: CometShuffleExchangeExec =>
-      case _: CometBroadcastExchangeExec =>
-      case _: WholeStageCodegenExec | _: ColumnarToRowExec | _: InputAdapter =>
-      case op if !excludedClasses.exists(c => c.isAssignableFrom(op.getClass)) =>
-        return Some(op)
-      case _ =>
-    }
-    None
-  }
-
   // checks the plan node has no missing inputs
   // such nodes represented in plan with exclamation mark !
   // example: !CometWindowExec
@@ -449,14 +430,6 @@ abstract class CometTestBase
     }
   }
 
-  /** Wraps the CometRowToColumn as ScanWrapper, so the child operators will not be checked */
-  private def wrapCometSparkToColumnar(plan: SparkPlan): SparkPlan = {
-    plan.transformDown {
-      // don't care the native operators
-      case p: CometSparkToColumnarExec => CometScanWrapper(null, p)
-    }
-  }
-
   private var _spark: SparkSessionType = _
   override protected implicit def spark: SparkSessionType = _spark
   protected implicit def sqlContext: SQLContext = _spark.sqlContext

spark/src/test/scala/org/apache/spark/sql/benchmark/CometBenchmarkBase.scala

@@ -31,14 +31,19 @@ import org.apache.parquet.crypto.keytools.mocks.InMemoryKMS
 import org.apache.spark.SparkConf
 import org.apache.spark.benchmark.Benchmark
 import org.apache.spark.sql.{DataFrame, DataFrameWriter, Row, SparkSession}
+import org.apache.spark.sql.comet.CometPlanChecker
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
 import org.apache.spark.sql.execution.benchmark.SqlBasedBenchmark
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.DecimalType
 
 import org.apache.comet.CometConf
 import org.apache.comet.CometSparkSessionExtensions
 
-trait CometBenchmarkBase extends SqlBasedBenchmark {
+trait CometBenchmarkBase
+    extends SqlBasedBenchmark
+    with AdaptiveSparkPlanHelper
+    with CometPlanChecker {
   override def getSparkSession: SparkSession = {
     val conf = new SparkConf()
       .setAppName("CometReadBenchmark")
@@ -88,28 +93,6 @@ trait CometBenchmarkBase extends SqlBasedBenchmark {
     }
   }
 
-  /** Runs function `f` with Comet on and off. */
-  final def runWithComet(name: String, cardinality: Long)(f: => Unit): Unit = {
-    val benchmark = new Benchmark(name, cardinality, output = output)
-
-    benchmark.addCase(s"$name - Spark ") { _ =>
-      withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
-        f
-      }
-    }
-
-    benchmark.addCase(s"$name - Comet") { _ =>
-      withSQLConf(
-        CometConf.COMET_ENABLED.key -> "true",
-        CometConf.COMET_EXEC_ENABLED.key -> "true",
-        SQLConf.ANSI_ENABLED.key -> "false") {
-        f
-      }
-    }
-
-    benchmark.run()
-  }
-
   /**
    * Runs an expression benchmark with standard cases: Spark, Comet (Scan), Comet (Scan + Exec).
    * This provides a consistent benchmark structure for expression evaluation.
@@ -149,6 +132,29 @@ trait CometBenchmarkBase extends SqlBasedBenchmark {
       CometConf.COMET_EXEC_ENABLED.key -> "true",
       "spark.sql.optimizer.constantFolding.enabled" -> "false") ++ extraCometConfigs
 
+    // Check that the plan is fully Comet native before running the benchmark
+    withSQLConf(cometExecConfigs.toSeq: _*) {
+      val df = spark.sql(query)
+      df.noop()
+      val plan = stripAQEPlan(df.queryExecution.executedPlan)
+      findFirstNonCometOperator(plan) match {
+        case Some(op) =>
+          // scalastyle:off println
+          println()
+          println("=" * 80)
+          println("WARNING: Benchmark plan is NOT fully Comet native!")
+          println(s"First non-Comet operator: ${op.nodeName}")
+          println("=" * 80)
+          println("Query plan:")
+          println(plan.treeString)
+          println("=" * 80)
+          println()
+        // scalastyle:on println
+        case None =>
+        // All operators are Comet native, no warning needed
+      }
+    }
+
     benchmark.addCase("Comet (Scan + Exec)") { _ =>
       withSQLConf(cometExecConfigs.toSeq: _*) {
         spark.sql(query).noop()