[SPARK-34952][SQL] DSv2 Aggregate push down APIs

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

@@ -22,7 +22,8 @@
 /**
  * An interface for building the {@link Scan}. Implementations can mixin SupportsPushDownXYZ
  * interfaces to do operator pushdown, and keep the operator pushdown result in the returned
- * {@link Scan}.
+ * {@link Scan}. When pushing down operators, Spark pushes down filters first, then pushes down
+ * aggregates or applies column pruning.
  *
  * @since 3.0.0
  */

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

@@ -0,0 +1,51 @@
+/*
+ * 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.read;
+
+import org.apache.spark.annotation.Evolving;
+import org.apache.spark.sql.connector.expressions.Aggregation;
+
+/**
+ * A mix-in interface for {@link ScanBuilder}. Data source can implement this interface to
+ * push down aggregates. Depends on the data source implementation, the aggregates may not
+ * be able to push down, or partially push down and have a final aggregate at Spark.
+ * For example, "SELECT min(_1) FROM t GROUP BY _2" can be pushed down to data source,
+ * the partially aggregated result min(_1) grouped by _2 will be returned to Spark, and
+ * then have a final aggregation.
+ * {{{
+ *   Aggregate [_2#10], [min(_2#10) AS min(_1)#16]
+ *     +- RelationV2[_2#10, min(_1)#18]
+ * }}}
+ *
+ * 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
+ * may not be able to be pushed down with filters. If pushed filters still need to be evaluated
+ * after scanning, aggregates can't be pushed down.
+ *
+ * @since 3.2.0
+ */
+@Evolving
+public interface SupportsPushDownAggregates extends ScanBuilder {
+
+  /**
+   * Pushes down Aggregation to datasource. The order of the datasource scan output columns should
+   * be: grouping columns, aggregate columns (in the same order as the aggregate functions in
+   * the given Aggregation).
+   */
+  boolean pushAggregation(Aggregation aggregation);
+}

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

@@ -0,0 +1,35 @@
+/*
+ * 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
+
+import org.apache.spark.sql.types.DataType
+
+// Aggregate Functions in SQL statement.
+// e.g. SELECT COUNT(EmployeeID), Max(salary) FROM dept GROUP BY deptID
+// aggregateExpressions are (COUNT(EmployeeID), Max(salary)), groupByColumns are (deptID)
+case class Aggregation(aggregateExpressions: Seq[AggregateFunc],
+                       groupByColumns: Seq[Expression])
+
+abstract class AggregateFunc
+
+case class Min(column: Expression, dataType: DataType) extends AggregateFunc
+case class Max(column: Expression, dataType: DataType) extends AggregateFunc
+case class Sum(column: Expression, dataType: DataType, isDistinct: Boolean)
+  extends AggregateFunc
+case class Count(column: Expression, dataType: DataType, isDistinct: Boolean)
+  extends AggregateFunc

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -870,6 +870,12 @@ object SQLConf {
       .checkValue(threshold => threshold >= 0, "The threshold must not be negative.")
       .createWithDefault(10)
 
+  val PARQUET_AGGREGATE_PUSHDOWN_ENABLED = buildConf("spark.sql.parquet.aggregatePushdown")
+    .doc("Enables Parquet aggregate push-down optimization when set to true.")
+    .version("3.2.0")
+    .booleanConf
+    .createWithDefault(false)
+
   val PARQUET_WRITE_LEGACY_FORMAT = buildConf("spark.sql.parquet.writeLegacyFormat")
     .doc("If true, data will be written in a way of Spark 1.4 and earlier. For example, decimal " +
       "values will be written in Apache Parquet's fixed-length byte array format, which other " +
@@ -3619,6 +3625,8 @@ class SQLConf extends Serializable with Logging {
   def parquetFilterPushDownInFilterThreshold: Int =
     getConf(PARQUET_FILTER_PUSHDOWN_INFILTERTHRESHOLD)
 
+  def parquetAggregatePushDown: Boolean = getConf(PARQUET_AGGREGATE_PUSHDOWN_ENABLED)
+
   def orcFilterPushDown: Boolean = getConf(ORC_FILTER_PUSHDOWN_ENABLED)
 
   def isOrcSchemaMergingEnabled: Boolean = getConf(ORC_SCHEMA_MERGING_ENABLED)

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

@@ -33,12 +33,14 @@ import org.apache.spark.sql.catalyst.catalog._
 import org.apache.spark.sql.catalyst.encoders.RowEncoder
 import org.apache.spark.sql.catalyst.expressions
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
 import org.apache.spark.sql.catalyst.planning.ScanOperation
 import org.apache.spark.sql.catalyst.plans.logical.{InsertIntoDir, InsertIntoStatement, LogicalPlan, Project}
 import org.apache.spark.sql.catalyst.rules.Rule
 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.{AggregateFunc, Count, FieldReference, LiteralValue, Max, Min}
 import org.apache.spark.sql.errors.QueryCompilationErrors
 import org.apache.spark.sql.execution.{InSubqueryExec, RowDataSourceScanExec, SparkPlan}
 import org.apache.spark.sql.execution.command._
@@ -692,6 +694,29 @@ object DataSourceStrategy
     (nonconvertiblePredicates ++ unhandledPredicates, pushedFilters, handledFilters)
   }
 
+  protected[sql] def translateAggregate(aggregates: AggregateExpression): Option[AggregateFunc] = {
+    if (aggregates.filter.isEmpty) {
+      aggregates.aggregateFunction match {
+        case min@aggregate.Min(PushableColumnAndNestedColumn(name)) =>
+          Some(Min(FieldReference(Seq(name)), min.dataType))
+        case max@aggregate.Max(PushableColumnAndNestedColumn(name)) =>
+          Some(Max(FieldReference(Seq(name)), max.dataType))
+        case count: aggregate.Count if count.children.length == 1 =>
+          count.children.head match {
+            // SELECT COUNT(*) FROM table is translated to SELECT 1 FROM table
+            case Literal(_, _) =>
+              Some(Count(LiteralValue(1L, LongType), LongType, aggregates.isDistinct))
+            case PushableColumnAndNestedColumn(name) =>
+              Some(Count(FieldReference(Seq(name)), LongType, aggregates.isDistinct))
+            case _ => None
+          }
+        case _ => None
+      }
+    } else {
+      None
+    }
+  }
+
   /**
    * Convert RDD of Row into RDD of InternalRow with objects in catalyst types
    */

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

@@ -20,7 +20,10 @@ package org.apache.spark.sql.execution.datasources.v2
 import scala.collection.mutable
 
 import org.apache.spark.sql.catalyst.expressions.{AttributeReference, AttributeSet, Expression, NamedExpression, PredicateHelper, SchemaPruning}
+import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
 import org.apache.spark.sql.catalyst.util.CharVarcharUtils
+import org.apache.spark.sql.connector.expressions.{Aggregation, FieldReference}
+import org.apache.spark.sql.connector.read.{Scan, ScanBuilder, SupportsPushDownAggregates, SupportsPushDownFilters, SupportsPushDownRequiredColumns}
 import org.apache.spark.sql.connector.read.{Scan, ScanBuilder, SupportsPushDownFilters, SupportsPushDownRequiredColumns}
 import org.apache.spark.sql.execution.datasources.DataSourceStrategy
 import org.apache.spark.sql.internal.SQLConf
@@ -70,6 +73,41 @@ object PushDownUtils extends PredicateHelper {
     }
   }
 
+  /**
+   * Pushes down aggregates to the data source reader
+   *
+   * @return pushed aggregation.
+   */
+  def pushAggregates(
+      scanBuilder: ScanBuilder,
+      aggregates: Seq[AggregateExpression],
+      groupBy: Seq[Expression]): Option[Aggregation] = {
+
+    def columnAsString(e: Expression): Option[FieldReference] = e match {
+      case AttributeReference(name, _, _, _) => Some(FieldReference(Seq(name)))
+      case _ => None
+    }
+
+    scanBuilder match {
+      case r: SupportsPushDownAggregates =>
+        val translatedAggregates = aggregates.map(DataSourceStrategy.translateAggregate).flatten
+        val translatedGroupBys = groupBy.map(columnAsString).flatten
+
+        if (translatedAggregates.length != aggregates.length ||
+          translatedGroupBys.length != groupBy.length) {
+          return None
+        }
+
+        val agg = Aggregation(translatedAggregates, translatedGroupBys)
+        if (r.pushAggregation(agg)) {
+          Some(agg)
+        } else {
+          None
+        }
+      case _ => None
+    }
+  }
+
   /**
    * Applies column pruning to the data source, w.r.t. the references of the given expressions.
    *