chore: Add type parameter to CometAggregateExpressionSerde (#2249)

Author: andygrove

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

@@ -174,7 +174,7 @@ object QueryPlanSerde extends Logging with CometExprShim {
   /**
    * Mapping of Spark aggregate expression class to Comet expression handler.
    */
-  private val aggrSerdeMap: Map[Class[_], CometAggregateExpressionSerde] = Map(
+  private val aggrSerdeMap: Map[Class[_], CometAggregateExpressionSerde[_]] = Map(
     classOf[Sum] -> CometSum,
     classOf[Average] -> CometAverage,
     classOf[Count] -> CometCount,
@@ -498,7 +498,9 @@ object QueryPlanSerde extends Logging with CometExprShim {
     val cometExpr = aggrSerdeMap.get(fn.getClass)
     cometExpr match {
       case Some(handler) =>
-        handler.convert(aggExpr, fn, inputs, binding, conf)
+        handler
+          .asInstanceOf[CometAggregateExpressionSerde[AggregateFunction]]
+          .convert(aggExpr, fn, inputs, binding, conf)
       case _ =>
         withInfo(
           aggExpr,
@@ -2456,7 +2458,7 @@ trait CometExpressionSerde[T <: Expression] {
 /**
  * Trait for providing serialization logic for aggregate expressions.
  */
-trait CometAggregateExpressionSerde {
+trait CometAggregateExpressionSerde[T <: AggregateFunction] {
 
   /**
    * Convert a Spark expression into a protocol buffer representation that can be passed into
@@ -2479,7 +2481,7 @@ trait CometAggregateExpressionSerde {
    */
   def convert(
       aggExpr: AggregateExpression,
-      expr: Expression,
+      expr: T,
       inputs: Seq[Attribute],
       binding: Boolean,
       conf: SQLConf): Option[ExprOuterClass.AggExpr]