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

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

@@ -0,0 +1,31 @@
+/*
+ * 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.annotation.Evolving;
+
+import java.io.Serializable;
+
+/**
+ * Base class of the Aggregate Functions.
+ *
+ * @since 3.2.0
+ */
+@Evolving
+public interface AggregateFunc extends Expression, Serializable {
+}

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

@@ -0,0 +1,46 @@
+/*
+ * 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.annotation.Evolving;
+
+import java.io.Serializable;
+
+/**
+ * Aggregation in SQL statement.
+ *
+ * @since 3.2.0
+ */
+@Evolving
+public final class Aggregation implements Serializable {
+  private AggregateFunc[] aggregateExpressions;
+  private FieldReference[] groupByColumns;
+
+  public Aggregation(AggregateFunc[] aggregateExpressions, FieldReference[] groupByColumns) {
+    this.aggregateExpressions = aggregateExpressions;
+    this.groupByColumns = groupByColumns;
+  }
+
+  public AggregateFunc[] aggregateExpressions() {
+    return aggregateExpressions;
+  }
+
+  public FieldReference[] groupByColumns() {
+    return groupByColumns;
+  }
+}

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

@@ -0,0 +1,49 @@
+/*
+ * 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.annotation.Evolving;
+
+/**
+ * An aggregate function that returns the number of the specific row in a group.
+ *
+ * @since 3.2.0
+ */
+@Evolving
+public final class Count implements AggregateFunc {
+    private FieldReference column;
+    private boolean isDistinct;
+
+    public Count(FieldReference column, boolean isDistinct) {
+        this.column = column;
+        this.isDistinct = isDistinct;
+    }
+
+    public FieldReference column() {
+        return column;
+    }
+    public boolean isDinstinct() {
+        return isDistinct;
+    }
+
+    @Override
+    public String toString() { return "Count(" + column.describe() + "," + isDistinct + ")"; }
+
+    @Override
+    public String describe() { return this.toString(); }
+}

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

@@ -0,0 +1,40 @@
+/*
+ * 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.annotation.Evolving;
+
+/**
+ * An aggregate function that returns the number of rows in a group.
+ *
+ * @since 3.2.0
+ */
+@Evolving
+public final class CountStar implements AggregateFunc {
+
+    public CountStar() {
+    }
+
+    @Override
+    public String toString() {
+        return "CountStar()";
+    }
+
+    @Override
+    public String describe() { return this.toString(); }
+}

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

@@ -0,0 +1,44 @@
+/*
+ * 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.annotation.Evolving;
+
+/**
+ * An aggregate function that returns the maximum value in a group.
+ *
+ * @since 3.2.0
+ */
+@Evolving
+public final class Max implements AggregateFunc {
+    private FieldReference column;
+
+    public Max(FieldReference column) {
+        this.column = column;
+    }
+
+    public FieldReference column() { return column; }
+
+    @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/Min.java

@@ -0,0 +1,46 @@
+/*
+ * 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.annotation.Evolving;
+
+/**
+ * An aggregate function that returns the minimum value in a group.
+ *
+ * @since 3.2.0
+ */
+@Evolving
+public final class Min implements AggregateFunc {
+    private FieldReference column;
+
+    public Min(FieldReference column) {
+        this.column = column;
+    }
+
+    public FieldReference column() {
+        return column;
+    }
+
+    @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/Sum.java

@@ -0,0 +1,57 @@
+/*
+ * 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.annotation.Evolving;
+import org.apache.spark.sql.types.DataType;
+
+/**
+ * An aggregate function that returns the summation of all the values in a group.
+ *
+ * @since 3.2.0
+ */
+@Evolving
+public final class Sum implements AggregateFunc {
+    private FieldReference column;
+    private DataType dataType;
+    private boolean isDistinct;
+
+    public Sum(FieldReference column, DataType dataType, boolean isDistinct) {
+        this.column = column;
+        this.dataType = dataType;
+        this.isDistinct = isDistinct;
+    }
+
+    public FieldReference column() {
+        return column;
+    }
+    public DataType dataType() {
+        return dataType;
+    }
+    public boolean isDinstinct() {
+        return isDistinct;
+    }
+
+    @Override
+    public String toString() {
+        return "Sum(" + column.describe() + "," + dataType + "," + isDistinct + ")";
+    }
+
+    @Override
+    public String describe() { return this.toString(); }
+}

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,56 @@
+/*
+ * 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 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:
+ * {{{
+ *   Aggregate [key#1], [min(min(value)#2) AS m#3]
+ *     +- RelationV2[key#1, min(value)#2]
+ * }}}
+ *
+ * <p>
+ * 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
+ * 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);
+}