feat: Add Comet windows function support. fmt

spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala

@@ -20,12 +20,12 @@
 package org.apache.comet
 
 import java.nio.ByteOrder
+
 import org.apache.spark.SparkConf
 import org.apache.spark.internal.Logging
 import org.apache.spark.network.util.ByteUnit
 import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.SparkSessionExtensions
-import org.apache.spark.sql.catalyst.expressions.AttributeReference
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.comet._
 import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle}
@@ -37,15 +37,16 @@ import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat
 import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
 import org.apache.spark.sql.execution.datasources.v2.parquet.ParquetScan
 import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ShuffleExchangeExec}
+import org.apache.spark.sql.execution.window.WindowExec
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types._
+
 import org.apache.comet.CometConf._
 import org.apache.comet.CometSparkSessionExtensions.{isANSIEnabled, isCometBroadCastEnabled, isCometColumnarShuffleEnabled, isCometEnabled, isCometExecEnabled, isCometOperatorEnabled, isCometScan, isCometScanEnabled, isCometShuffleEnabled, isSchemaSupported}
 import org.apache.comet.parquet.{CometParquetScan, SupportsComet}
 import org.apache.comet.serde.OperatorOuterClass.Operator
 import org.apache.comet.serde.QueryPlanSerde
 import org.apache.comet.shims.ShimCometSparkSessionExtensions
-import org.apache.spark.sql.execution.window.WindowExec
 
 class CometSparkSessionExtensions
     extends (SparkSessionExtensions => Unit)

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

@@ -20,6 +20,7 @@
 package org.apache.comet.serde
 
 import scala.collection.JavaConverters._
+
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Average, Count, Final, First, Last, Max, Min, Partial, Sum}
@@ -32,15 +33,16 @@ import org.apache.spark.sql.execution
 import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.aggregate.HashAggregateExec
 import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ShuffleExchangeExec}
+import org.apache.spark.sql.execution.window.WindowExec
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types._
 import org.apache.spark.unsafe.types.UTF8String
+
 import org.apache.comet.CometSparkSessionExtensions.{isCometOperatorEnabled, isCometScan, isSpark32, isSpark34Plus}
-import org.apache.comet.serde.ExprOuterClass.{AggExpr, Expr, ScalarFunc, DataType => ProtoDataType}
+import org.apache.comet.serde.ExprOuterClass.{AggExpr, DataType => ProtoDataType, Expr, ScalarFunc}
 import org.apache.comet.serde.ExprOuterClass.DataType.{DataTypeInfo, DecimalInfo, ListInfo, MapInfo, StructInfo}
-import org.apache.comet.serde.OperatorOuterClass.{Operator, AggregateMode => CometAggregateMode}
+import org.apache.comet.serde.OperatorOuterClass.{AggregateMode => CometAggregateMode, Operator}
 import org.apache.comet.shims.ShimQueryPlanSerde
-import org.apache.spark.sql.execution.window.WindowExec
 
 /**
  * An utility object for query plan and expression serialization.
@@ -186,8 +188,8 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde {
   }
 
   def windowExprToProto(
-                         windowExpr: WindowExpression,
-                         inputs: Seq[Attribute]): Option[OperatorOuterClass.WindowExpr] = {
+      windowExpr: WindowExpression,
+      inputs: Seq[Attribute]): Option[OperatorOuterClass.WindowExpr] = {
     val func = exprToProto(windowExpr.windowFunction, inputs).getOrElse(return None)
 
     val f = windowExpr.windowSpec.frameSpecification

spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala

@@ -1,10 +1,26 @@
-package org.apache.spark.sql.comet
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
 
-import org.apache.comet.serde.OperatorOuterClass
-import org.apache.comet.serde.OperatorOuterClass.Operator
-import org.apache.comet.serde.QueryPlanSerde.{exprToProto, serializeDataType, windowExprToProto}
+package org.apache.spark.sql.comet
 
 import scala.collection.JavaConverters.asJavaIterableConverter
+
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, Expression, NamedExpression, SortOrder, WindowExpression}
 import org.apache.spark.sql.catalyst.plans.physical.Partitioning
@@ -13,6 +29,10 @@ import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode}
 import org.apache.spark.sql.execution.metric.{SQLMetrics, SQLShuffleReadMetricsReporter, SQLShuffleWriteMetricsReporter}
 import org.apache.spark.sql.vectorized.ColumnarBatch
 
+import org.apache.comet.serde.OperatorOuterClass
+import org.apache.comet.serde.OperatorOuterClass.Operator
+import org.apache.comet.serde.QueryPlanSerde.{exprToProto, serializeDataType, windowExprToProto}
+
 /**
  * Comet physical plan node for Spark `WindowsExec`.
  *
@@ -21,12 +41,12 @@ import org.apache.spark.sql.vectorized.ColumnarBatch
  * executions separated by a Comet shuffle exchange.
  */
 case class CometWindowExec(
-                            override val originalPlan: SparkPlan,
-                            windowExpression: Seq[NamedExpression],
-                            partitionSpec: Seq[Expression],
-                            orderSpec: Seq[SortOrder],
-                            child: SparkPlan)
-  extends CometExec
+    override val originalPlan: SparkPlan,
+    windowExpression: Seq[NamedExpression],
+    partitionSpec: Seq[Expression],
+    orderSpec: Seq[SortOrder],
+    child: SparkPlan)
+    extends CometExec
     with UnaryExecNode {
 
   override def nodeName: String = "CometWindowExec"
@@ -68,11 +88,11 @@ case class CometWindowExec(
 
 object CometWindowExec {
   def getNativePlan(
-                     outputAttributes: Seq[Attribute],
-                     windowExpression: Seq[NamedExpression],
-                     partitionSpec: Seq[Expression],
-                     orderSpec: Seq[SortOrder],
-                     child: SparkPlan): Option[Operator] = {
+      outputAttributes: Seq[Attribute],
+      windowExpression: Seq[NamedExpression],
+      partitionSpec: Seq[Expression],
+      orderSpec: Seq[SortOrder],
+      child: SparkPlan): Option[Operator] = {
 
     val orderSpecs = orderSpec.map(exprToProto(_, child.output))
     val partitionSpecs = partitionSpec.map(exprToProto(_, child.output))
@@ -110,4 +130,4 @@ object CometWindowExec {
       Some(opBuilder.setWindow(windowBuilder).build())
     } else None
   }
-}
+}

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

@@ -90,13 +90,13 @@ class CometCastSuite extends CometTestBase with AdaptiveSparkPlanHelper {
     Range(0, len).map(_ => chars.charAt(r.nextInt(chars.length))).mkString
   }
 
-  private def fuzzCastFromString(chars: String, maxLen: Int, toType: DataType) {
+  private def fuzzCastFromString(chars: String, maxLen: Int, toType: DataType): Unit = {
     val r = new Random(0)
     val inputs = Range(0, 10000).map(_ => genString(r, chars, maxLen))
     castTest(inputs.toDF("a"), toType)
   }
 
-  private def castTest(input: DataFrame, toType: DataType) {
+  private def castTest(input: DataFrame, toType: DataType): Unit = {
     withTempPath { dir =>
       val df = roundtripParquet(input, dir)
         .withColumn("converted", col("a").cast(toType))