[SPARK-37789][SQL] Add a class to represent general aggregate functions in DS V2

### What changes were proposed in this pull request?

There are a lot of aggregate functions in SQL and it's a lot of work to add them one by one in the DS v2 API. This PR proposes to add a new `GeneralAggregateFunc` class to represent all the general SQL aggregate functions. Since it's general, Spark doesn't know its aggregation buffer and can only push down the aggregation to the source completely.

As an example, this PR also translates `AVG` to `GeneralAggregateFunc` and pushes it to JDBC V2.

### Why are the changes needed?

To add aggregate functions in DS v2 easier.

### Does this PR introduce _any_ user-facing change?

No

### How was this patch tested?

JDBC v2 test

Closes apache#35070 from cloud-fan/agg.

Lead-authored-by: Wenchen Fan <[email protected]>
Co-authored-by: Wenchen Fan <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>

Author: cloud-fan

sql/catalyst/src/main/java/org/apache/spark/sql/connector/expressions/Expression.java

@@ -29,5 +29,5 @@ public interface Expression {
   /**
    * Format the expression as a human readable SQL-like string.
    */
-  String describe();
+  default String describe() { return this.toString(); }
 }

sql/catalyst/src/main/java/org/apache/spark/sql/connector/expressions/aggregate/Count.java

@@ -46,7 +46,4 @@ public String toString() {
       return "COUNT(" + column.describe() + ")";
     }
   }
-
-  @Override
-  public String describe() { return this.toString(); }
 }

sql/catalyst/src/main/java/org/apache/spark/sql/connector/expressions/aggregate/CountStar.java

@@ -32,7 +32,4 @@ public CountStar() {
 
   @Override
   public String toString() { return "COUNT(*)"; }
-
-  @Override
-  public String describe() { return this.toString(); }
 }

sql/catalyst/src/main/java/org/apache/spark/sql/connector/expressions/aggregate/GeneralAggregateFunc.java

@@ -0,0 +1,66 @@
+/*
+ * 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.
+ */
+
+package org.apache.spark.sql.connector.expressions.aggregate;
+
+import java.util.Arrays;
+import java.util.stream.Collectors;
+
+import org.apache.spark.annotation.Evolving;
+import org.apache.spark.sql.connector.expressions.Expression;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * The general implementation of {@link AggregateFunc}, which contains the upper-cased function
+ * name, the `isDistinct` flag and all the inputs. Note that Spark cannot push down partial
+ * aggregate with this function to the source, but can only push down the entire aggregate.
+ * <p>
+ * The currently supported SQL aggregate functions:
+ * <ol>
+ *  <li><pre>AVG(input1)</pre> Since 3.3.0</li>
+ * </ol>
+ *
+ * @since 3.3.0
+ */
+@Evolving
+public final class GeneralAggregateFunc implements AggregateFunc {
+  private final String name;
+  private final boolean isDistinct;
+  private final NamedReference[] inputs;
+
+  public String name() { return name; }
+  public boolean isDistinct() { return isDistinct; }
+  public NamedReference[] inputs() { return inputs; }
+
+  public GeneralAggregateFunc(String name, boolean isDistinct, NamedReference[] inputs) {
+    this.name = name;
+    this.isDistinct = isDistinct;
+    this.inputs = inputs;
+  }
+
+  @Override
+  public String toString() {
+    String inputsString = Arrays.stream(inputs)
+      .map(Expression::describe)
+      .collect(Collectors.joining(", "));
+    if (isDistinct) {
+      return name + "(DISTINCT " + inputsString + ")";
+    } else {
+      return name + "(" + inputsString + ")";
+    }
+  }
+}

sql/catalyst/src/main/java/org/apache/spark/sql/connector/expressions/aggregate/Max.java

@@ -35,7 +35,4 @@ public final class Max implements AggregateFunc {
 
   @Override
   public String toString() { return "MAX(" + column.describe() + ")"; }
-
-  @Override
-  public String describe() { return this.toString(); }
 }

sql/catalyst/src/main/java/org/apache/spark/sql/connector/expressions/aggregate/Min.java

@@ -35,7 +35,4 @@ public final class Min implements AggregateFunc {
 
   @Override
   public String toString() { return "MIN(" + column.describe() + ")"; }
-
-  @Override
-  public String describe() { return this.toString(); }
 }

sql/catalyst/src/main/java/org/apache/spark/sql/connector/expressions/aggregate/Sum.java

@@ -46,7 +46,4 @@ public String toString() {
       return "SUM(" + column.describe() + ")";
     }
   }
-
-  @Override
-  public String describe() { return this.toString(); }
 }

sql/catalyst/src/main/java/org/apache/spark/sql/connector/expressions/filter/Filter.java

@@ -37,7 +37,4 @@ public abstract class Filter implements Expression, Serializable {
    * Returns list of columns that are referenced by this filter.
    */
   public abstract NamedReference[] references();
-
-  @Override
-  public String describe() { return this.toString(); }
 }

sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsPushDownAggregates.java

@@ -22,18 +22,19 @@
 
 /**
  * A mix-in interface for {@link ScanBuilder}. Data sources can implement this interface to
- * push down aggregates. Spark assumes that the data source can't fully complete the
- * grouping work, and will group the data source output again. For queries like
- * "SELECT min(value) AS m FROM t GROUP BY key", after pushing down the aggregate
- * to the data source, the data source can still output data with duplicated keys, which is OK
- * as Spark will do GROUP BY key again. The final query plan can be something like this:
+ * push down aggregates.
+ * <p>
+ * If the data source can't fully complete the grouping work, then
+ * {@link #supportCompletePushDown()} should return false, and Spark will group the data source
+ * output again. For queries like "SELECT min(value) AS m FROM t GROUP BY key", after pushing down
+ * the aggregate to the data source, the data source can still output data with duplicated keys,
+ * which is OK as Spark will do GROUP BY key again. The final query plan can be something like this:
  * <pre>
- *   Aggregate [key#1], [min(min(value)#2) AS m#3]
- *     +- RelationV2[key#1, min(value)#2]
+ *   Aggregate [key#1], [min(min_value#2) AS m#3]
+ *     +- RelationV2[key#1, min_value#2]
  * </pre>
  * Similarly, if there is no grouping expression, the data source can still output more than one
  * rows.
- *
  * <p>
  * When pushing down operators, Spark pushes down filters to the data source first, then push down
  * aggregates or apply column pruning. Depends on data source implementation, aggregates may or
@@ -46,8 +47,8 @@
 public interface SupportsPushDownAggregates extends ScanBuilder {
 
   /**
-   * Whether the datasource support complete aggregation push-down. Spark could avoid partial-agg
-   * and final-agg when the aggregation operation can be pushed down to the datasource completely.
+   * Whether the datasource support complete aggregation push-down. Spark will do grouping again
+   * if this method returns false.
    *
    * @return true if the aggregation can be pushed down to datasource completely, false otherwise.
    */

sql/catalyst/src/main/scala/org/apache/spark/sql/connector/expressions/expressions.scala

@@ -88,9 +88,7 @@ private[sql] abstract class SingleColumnTransform(ref: NamedReference) extends R
 
   override def arguments: Array[Expression] = Array(ref)
 
-  override def describe: String = name + "(" + reference.describe + ")"
-
-  override def toString: String = describe
+  override def toString: String = name + "(" + reference.describe + ")"
 
   protected def withNewRef(ref: NamedReference): Transform
 
@@ -114,16 +112,14 @@ private[sql] final case class BucketTransform(
 
   override def arguments: Array[Expression] = numBuckets +: columns.toArray
 
-  override def describe: String =
+  override def toString: String =
     if (sortedColumns.nonEmpty) {
       s"bucket(${arguments.map(_.describe).mkString(", ")}," +
         s" ${sortedColumns.map(_.describe).mkString(", ")})"
     } else {
       s"bucket(${arguments.map(_.describe).mkString(", ")})"
     }
 
-  override def toString: String = describe
-
   override def withReferences(newReferences: Seq[NamedReference]): Transform = {
     this.copy(columns = newReferences)
   }
@@ -169,9 +165,7 @@ private[sql] final case class ApplyTransform(
     arguments.collect { case named: NamedReference => named }
   }
 
-  override def describe: String = s"$name(${arguments.map(_.describe).mkString(", ")})"
-
-  override def toString: String = describe
+  override def toString: String = s"$name(${arguments.map(_.describe).mkString(", ")})"
 }
 
 /**
@@ -338,21 +332,19 @@ private[sql] object HoursTransform {
 }
 
 private[sql] final case class LiteralValue[T](value: T, dataType: DataType) extends Literal[T] {
-  override def describe: String = {
+  override def toString: String = {
     if (dataType.isInstanceOf[StringType]) {
       s"'$value'"
     } else {
       s"$value"
     }
   }
-  override def toString: String = describe
 }
 
 private[sql] final case class FieldReference(parts: Seq[String]) extends NamedReference {
   import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.MultipartIdentifierHelper
   override def fieldNames: Array[String] = parts.toArray
-  override def describe: String = parts.quoted
-  override def toString: String = describe
+  override def toString: String = parts.quoted
 }
 
 private[sql] object FieldReference {
@@ -366,7 +358,7 @@ private[sql] final case class SortValue(
     direction: SortDirection,
     nullOrdering: NullOrdering) extends SortOrder {
 
-  override def describe(): String = s"$expression $direction $nullOrdering"
+  override def toString(): String = s"$expression $direction $nullOrdering"
 }
 
 private[sql] object SortValue {

sql/catalyst/src/test/scala/org/apache/spark/sql/connector/expressions/TransformExtractorSuite.scala

@@ -28,15 +28,15 @@ class TransformExtractorSuite extends SparkFunSuite {
   private def lit[T](literal: T): Literal[T] = new Literal[T] {
     override def value: T = literal
     override def dataType: DataType = catalyst.expressions.Literal(literal).dataType
-    override def describe: String = literal.toString
+    override def toString: String = literal.toString
   }
 
   /**
    * Creates a NamedReference using an anonymous class.
    */
   private def ref(names: String*): NamedReference = new NamedReference {
     override def fieldNames: Array[String] = names.toArray
-    override def describe: String = names.mkString(".")
+    override def toString: String = names.mkString(".")
   }
 
   /**
@@ -46,7 +46,7 @@ class TransformExtractorSuite extends SparkFunSuite {
     override def name: String = func
     override def references: Array[NamedReference] = Array(ref)
     override def arguments: Array[Expression] = Array(ref)
-    override def describe: String = ref.describe
+    override def toString: String = ref.describe
   }
 
   test("Identity extractor") {
@@ -135,7 +135,7 @@ class TransformExtractorSuite extends SparkFunSuite {
       override def name: String = "bucket"
       override def references: Array[NamedReference] = Array(col)
       override def arguments: Array[Expression] = Array(lit(16), col)
-      override def describe: String = s"bucket(16, ${col.describe})"
+      override def toString: String = s"bucket(16, ${col.describe})"
     }
 
     bucketTransform match {

sql/core/src/main/scala/org/apache/spark/sql/execution/DataSourceScanExec.scala

@@ -156,8 +156,8 @@ case class RowDataSourceScanExec(
       "ReadSchema" -> requiredSchema.catalogString,
       "PushedFilters" -> seqToString(markedFilters.toSeq)) ++
       pushedDownOperators.aggregation.fold(Map[String, String]()) { v =>
-        Map("PushedAggregates" -> seqToString(v.aggregateExpressions),
-          "PushedGroupByColumns" -> seqToString(v.groupByColumns))} ++
+        Map("PushedAggregates" -> seqToString(v.aggregateExpressions.map(_.describe())),
+          "PushedGroupByColumns" -> seqToString(v.groupByColumns.map(_.describe())))} ++
       topNOrLimitInfo ++
       pushedDownOperators.sample.map(v => "PushedSample" ->
         s"SAMPLE (${(v.upperBound - v.lowerBound) * 100}) ${v.withReplacement} SEED(${v.seed})"

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala

@@ -41,7 +41,7 @@ import org.apache.spark.sql.catalyst.streaming.StreamingRelationV2
 import org.apache.spark.sql.connector.catalog.SupportsRead
 import org.apache.spark.sql.connector.catalog.TableCapability._
 import org.apache.spark.sql.connector.expressions.{FieldReference, NullOrdering, SortDirection, SortOrder => SortOrderV2, SortValue}
-import org.apache.spark.sql.connector.expressions.aggregate.{AggregateFunc, Count, CountStar, Max, Min, Sum}
+import org.apache.spark.sql.connector.expressions.aggregate.{AggregateFunc, Count, CountStar, GeneralAggregateFunc, Max, Min, Sum}
 import org.apache.spark.sql.errors.QueryCompilationErrors
 import org.apache.spark.sql.execution.{InSubqueryExec, RowDataSourceScanExec, SparkPlan}
 import org.apache.spark.sql.execution.command._
@@ -714,8 +714,10 @@ object DataSourceStrategy
               Some(new Count(FieldReference(name), aggregates.isDistinct))
             case _ => None
           }
-        case sum @ aggregate.Sum(PushableColumnWithoutNestedColumn(name), _) =>
+        case aggregate.Sum(PushableColumnWithoutNestedColumn(name), _) =>
           Some(new Sum(FieldReference(name), aggregates.isDistinct))
+        case aggregate.Average(PushableColumnWithoutNestedColumn(name), _) =>
+          Some(new GeneralAggregateFunc("AVG", aggregates.isDistinct, Array(FieldReference(name))))
         case _ => None
       }
     } else {

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala

@@ -26,7 +26,7 @@ import org.apache.spark.sql.catalyst.planning.ScanOperation
 import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, LeafNode, Limit, LocalLimit, LogicalPlan, Project, Sample, Sort}
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.connector.expressions.SortOrder
-import org.apache.spark.sql.connector.expressions.aggregate.Aggregation
+import org.apache.spark.sql.connector.expressions.aggregate.{Aggregation, GeneralAggregateFunc}
 import org.apache.spark.sql.connector.read.{Scan, ScanBuilder, SupportsPushDownAggregates, SupportsPushDownFilters, V1Scan}
 import org.apache.spark.sql.execution.datasources.DataSourceStrategy
 import org.apache.spark.sql.sources
@@ -109,6 +109,9 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper {
                 sHolder.builder, normalizedAggregates, normalizedGroupingExpressions)
               if (pushedAggregates.isEmpty) {
                 aggNode // return original plan node
+              } else if (!supportPartialAggPushDown(pushedAggregates.get) &&
+                !r.supportCompletePushDown()) {
+                aggNode // return original plan node
               } else {
                 // No need to do column pruning because only the aggregate columns are used as
                 // DataSourceV2ScanRelation output columns. All the other columns are not
@@ -145,9 +148,7 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper {
                       """.stripMargin)
 
                 val wrappedScan = getWrappedScan(scan, sHolder, pushedAggregates)
-
                 val scanRelation = DataSourceV2ScanRelation(sHolder.relation, wrappedScan, output)
-
                 if (r.supportCompletePushDown()) {
                   val projectExpressions = resultExpressions.map { expr =>
                     // TODO At present, only push down group by attribute is supported.
@@ -209,6 +210,11 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper {
       }
   }
 
+  private def supportPartialAggPushDown(agg: Aggregation): Boolean = {
+    // We don't know the agg buffer of `GeneralAggregateFunc`, so can't do partial agg push down.
+    agg.aggregateExpressions().forall(!_.isInstanceOf[GeneralAggregateFunc])
+  }
+
   private def addCastIfNeeded(aggAttribute: AttributeReference, aggDataType: DataType) =
     if (aggAttribute.dataType == aggDataType) {
       aggAttribute
@@ -256,7 +262,7 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper {
 
   def pushDownSample(plan: LogicalPlan): LogicalPlan = plan.transform {
     case sample: Sample => sample.child match {
-      case ScanOperation(_, filter, sHolder: ScanBuilderHolder) if filter.length == 0 =>
+      case ScanOperation(_, filter, sHolder: ScanBuilderHolder) if filter.isEmpty =>
         val tableSample = TableSampleInfo(
           sample.lowerBound,
           sample.upperBound,
@@ -282,7 +288,7 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper {
       }
       operation
     case s @ Sort(order, _, operation @ ScanOperation(_, filter, sHolder: ScanBuilderHolder))
-      if filter.isEmpty =>
+        if filter.isEmpty =>
       val orders = DataSourceStrategy.translateSortOrders(order)
       if (orders.length == order.length) {
         val topNPushed = PushDownUtils.pushTopN(sHolder.builder, orders.toArray, limit)

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/jdbc/JDBCScanBuilder.scala

@@ -79,7 +79,7 @@ case class JDBCScanBuilder(
     if (!jdbcOptions.pushDownAggregate) return false
 
     val dialect = JdbcDialects.get(jdbcOptions.url)
-    val compiledAggs = aggregation.aggregateExpressions.flatMap(dialect.compileAggregate(_))
+    val compiledAggs = aggregation.aggregateExpressions.flatMap(dialect.compileAggregate)
     if (compiledAggs.length != aggregation.aggregateExpressions.length) return false
 
     val groupByCols = aggregation.groupByColumns.map { col =>

sql/core/src/main/scala/org/apache/spark/sql/jdbc/JdbcDialects.scala

@@ -30,7 +30,7 @@ import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalyst.util.{DateFormatter, DateTimeUtils, TimestampFormatter}
 import org.apache.spark.sql.connector.catalog.TableChange
 import org.apache.spark.sql.connector.catalog.TableChange._
-import org.apache.spark.sql.connector.expressions.aggregate.{AggregateFunc, Count, CountStar, Max, Min, Sum}
+import org.apache.spark.sql.connector.expressions.aggregate.{AggregateFunc, Count, CountStar, GeneralAggregateFunc, Max, Min, Sum}
 import org.apache.spark.sql.errors.QueryCompilationErrors
 import org.apache.spark.sql.execution.datasources.jdbc.JdbcUtils
 import org.apache.spark.sql.execution.datasources.v2.TableSampleInfo
@@ -216,7 +216,11 @@ abstract class JdbcDialect extends Serializable with Logging{
         val column = quoteIdentifier(sum.column.fieldNames.head)
         Some(s"SUM($distinct$column)")
       case _: CountStar =>
-        Some(s"COUNT(*)")
+        Some("COUNT(*)")
+      case f: GeneralAggregateFunc if f.name() == "AVG" =>
+        assert(f.inputs().length == 1)
+        val distinct = if (f.isDistinct) "DISTINCT " else ""
+        Some(s"AVG($distinct${f.inputs().head})")
       case _ => None
     }
   }