[SPARK-37960][SQL] A new framework to represent catalyst expressions in DS v2 APIs

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

@@ -0,0 +1,52 @@
+/*
+ * 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 java.io.Serializable;
+
+import org.apache.spark.annotation.Evolving;
+
+/**
+ * The general SQL string corresponding to expression.
+ * <p>
+ * The currently supported SQL expressions:
+ * <ol>
+ *  <li><pre>Literal</pre> Since 3.3.0</li>
+ *  <li><pre>Attribute</pre> Since 3.3.0</li>
+ *  <li><pre>IsNull</pre> Since 3.3.0</li>
+ *  <li><pre>IsNotNull</pre> Since 3.3.0</li>
+ *  <li><pre>BinaryOperator</pre> Since 3.3.0</li>
+ *  <li><pre>Not</pre> Since 3.3.0</li>
+ *  <li><pre>CaseWhen</pre> Since 3.3.0</li>
+ * </ol>
+ *
+ * @since 3.3.0
+ */
+@Evolving
+public class GeneralSQLExpression implements Expression, Serializable {
+    private String sql;
+
+    public GeneralSQLExpression(String sql) {
+        this.sql = sql;
+    }
+
+    public String sql() { return sql; }
+
+    @Override
+    public String toString() { return sql; }
+}

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

@@ -18,7 +18,7 @@
 package org.apache.spark.sql.connector.expressions.aggregate;
 
 import org.apache.spark.annotation.Evolving;
-import org.apache.spark.sql.connector.expressions.NamedReference;
+import org.apache.spark.sql.connector.expressions.Expression;
 
 /**
  * An aggregate function that returns the mean of all the values in a group.
@@ -27,23 +27,23 @@
  */
 @Evolving
 public final class Avg implements AggregateFunc {
-  private final NamedReference column;
+  private final Expression input;
   private final boolean isDistinct;
 
-  public Avg(NamedReference column, boolean isDistinct) {
-    this.column = column;
+  public Avg(Expression column, boolean isDistinct) {
+    this.input = column;
     this.isDistinct = isDistinct;
   }
 
-  public NamedReference column() { return column; }
+  public Expression column() { return input; }
   public boolean isDistinct() { return isDistinct; }
 
   @Override
   public String toString() {
     if (isDistinct) {
-      return "AVG(DISTINCT " + column.describe() + ")";
+      return "AVG(DISTINCT " + input.describe() + ")";
     } else {
-      return "AVG(" + column.describe() + ")";
+      return "AVG(" + input.describe() + ")";
     }
   }
 }

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

@@ -18,7 +18,7 @@
 package org.apache.spark.sql.connector.expressions.aggregate;
 
 import org.apache.spark.annotation.Evolving;
-import org.apache.spark.sql.connector.expressions.NamedReference;
+import org.apache.spark.sql.connector.expressions.Expression;
 
 /**
  * An aggregate function that returns the number of the specific row in a group.
@@ -27,23 +27,23 @@
  */
 @Evolving
 public final class Count implements AggregateFunc {
-  private final NamedReference column;
+  private final Expression input;
   private final boolean isDistinct;
 
-  public Count(NamedReference column, boolean isDistinct) {
-    this.column = column;
+  public Count(Expression column, boolean isDistinct) {
+    this.input = column;
     this.isDistinct = isDistinct;
   }
 
-  public NamedReference column() { return column; }
+  public Expression column() { return input; }
   public boolean isDistinct() { return isDistinct; }
 
   @Override
   public String toString() {
     if (isDistinct) {
-      return "COUNT(DISTINCT " + column.describe() + ")";
+      return "COUNT(DISTINCT " + input.describe() + ")";
     } else {
-      return "COUNT(" + column.describe() + ")";
+      return "COUNT(" + input.describe() + ")";
     }
   }
 }

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

@@ -18,7 +18,7 @@
 package org.apache.spark.sql.connector.expressions.aggregate;
 
 import org.apache.spark.annotation.Evolving;
-import org.apache.spark.sql.connector.expressions.NamedReference;
+import org.apache.spark.sql.connector.expressions.Expression;
 
 /**
  * An aggregate function that returns the maximum value in a group.
@@ -27,12 +27,12 @@
  */
 @Evolving
 public final class Max implements AggregateFunc {
-  private final NamedReference column;
+  private final Expression input;
 
-  public Max(NamedReference column) { this.column = column; }
+  public Max(Expression column) { this.input = column; }
 
-  public NamedReference column() { return column; }
+  public Expression column() { return input; }
 
   @Override
-  public String toString() { return "MAX(" + column.describe() + ")"; }
+  public String toString() { return "MAX(" + input.describe() + ")"; }
 }

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

@@ -18,7 +18,7 @@
 package org.apache.spark.sql.connector.expressions.aggregate;
 
 import org.apache.spark.annotation.Evolving;
-import org.apache.spark.sql.connector.expressions.NamedReference;
+import org.apache.spark.sql.connector.expressions.Expression;
 
 /**
  * An aggregate function that returns the minimum value in a group.
@@ -27,12 +27,12 @@
  */
 @Evolving
 public final class Min implements AggregateFunc {
-  private final NamedReference column;
+  private final Expression input;
 
-  public Min(NamedReference column) { this.column = column; }
+  public Min(Expression column) { this.input = column; }
 
-  public NamedReference column() { return column; }
+  public Expression column() { return input; }
 
   @Override
-  public String toString() { return "MIN(" + column.describe() + ")"; }
+  public String toString() { return "MIN(" + input.describe() + ")"; }
 }

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

@@ -18,7 +18,7 @@
 package org.apache.spark.sql.connector.expressions.aggregate;
 
 import org.apache.spark.annotation.Evolving;
-import org.apache.spark.sql.connector.expressions.NamedReference;
+import org.apache.spark.sql.connector.expressions.Expression;
 
 /**
  * An aggregate function that returns the summation of all the values in a group.
@@ -27,23 +27,23 @@
  */
 @Evolving
 public final class Sum implements AggregateFunc {
-  private final NamedReference column;
+  private final Expression input;
   private final boolean isDistinct;
 
-  public Sum(NamedReference column, boolean isDistinct) {
-    this.column = column;
+  public Sum(Expression column, boolean isDistinct) {
+    this.input = column;
     this.isDistinct = isDistinct;
   }
 
-  public NamedReference column() { return column; }
+  public Expression column() { return input; }
   public boolean isDistinct() { return isDistinct; }
 
   @Override
   public String toString() {
     if (isDistinct) {
-      return "SUM(DISTINCT " + column.describe() + ")";
+      return "SUM(DISTINCT " + input.describe() + ")";
     } else {
-      return "SUM(" + column.describe() + ")";
+      return "SUM(" + input.describe() + ")";
     }
   }
 }

sql/core/src/main/scala/org/apache/spark/sql/catalyst/util/ExpressionSQLBuilder.scala

@@ -0,0 +1,65 @@
+/*
+ * 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.catalyst.util
+
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BinaryOperator, CaseWhen, EqualTo, Expression, IsNotNull, IsNull, Literal, Not}
+import org.apache.spark.sql.connector.expressions.LiteralValue
+
+/**
+ * The builder to build expression and generate sql string.
+ */
+class ExpressionSQLBuilder(e: Expression) {
+
+  def build(): Option[String] = generateSQL(e)
+
+  private def generateSQL(expr: Expression): Option[String] = expr match {
+    case Literal(value, dataType) => Some(LiteralValue(value, dataType).toString)
+    case a: Attribute => Some(quoteIfNeeded(a.name))
+    case IsNull(col) => generateSQL(col).map(c => s"$c IS NULL")
+    case IsNotNull(col) => generateSQL(col).map(c => s"$c IS NOT NULL")
+    case b: BinaryOperator =>
+      val l = generateSQL(b.left)
+      val r = generateSQL(b.right)
+      if (l.isDefined && r.isDefined) {
+        Some(s"${l.get} ${b.sqlOperator} ${r.get}")
+      } else {
+        None
+      }
+    case Not(EqualTo(left, right)) =>
+      val l = generateSQL(left)
+      val r = generateSQL(right)
+      if (l.isDefined && r.isDefined) {
+        Some(s"${l.get} != ${r.get}")
+      } else {
+        None
+      }
+    case CaseWhen(branches, elseValue) =>
+      val conditionsSQL = branches.map(_._1).flatMap(generateSQL)
+      val valuesSQL = branches.map(_._2).flatMap(generateSQL)
+      if (conditionsSQL.length == branches.length && valuesSQL.length == branches.length) {
+        val branchSQL =
+          conditionsSQL.zip(valuesSQL).map { case (c, v) => s" WHEN $c THEN $v" }.mkString
+        val elseSQL = elseValue.flatMap(generateSQL).map(v => s" ELSE $v").getOrElse("")
+        Some(s"CASE$branchSQL$elseSQL END")
+      } else {
+        None
+      }
+    // TODO supports other expressions
+    case _ => None
+  }
+}

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

@@ -52,10 +52,11 @@ object AggregatePushDownUtils {
 
     def processMinOrMax(agg: AggregateFunc): Boolean = {
       val (column, aggType) = agg match {
-        case max: Max => (max.column, "max")
-        case min: Min => (min.column, "min")
-        case _ =>
-          throw new IllegalArgumentException(s"Unexpected type of AggregateFunc ${agg.describe}")
+        case max: Max if max.column().isInstanceOf[NamedReference] =>
+          (max.column.asInstanceOf[NamedReference], "max")
+        case min: Min if min.column().isInstanceOf[NamedReference] =>
+          (min.column.asInstanceOf[NamedReference], "min")
+        case _ => return false
       }
 
       if (isPartitionCol(column)) {
@@ -117,10 +118,11 @@ object AggregatePushDownUtils {
         if (!processMinOrMax(max)) return None
       case min: Min =>
         if (!processMinOrMax(min)) return None
-      case count: Count =>
-        if (count.column.fieldNames.length != 1 || count.isDistinct) return None
+      case count: Count if count.column().isInstanceOf[NamedReference] =>
+        val reference = count.column.asInstanceOf[NamedReference]
+        if (reference.fieldNames.length != 1 || count.isDistinct) return None
         finalSchema =
-          finalSchema.add(StructField(s"count(" + count.column.fieldNames.head + ")", LongType))
+          finalSchema.add(StructField(s"count(" + reference.fieldNames.head + ")", LongType))
       case _: CountStar =>
         finalSchema = finalSchema.add(StructField("count(*)", LongType))
       case _ =>

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

@@ -38,9 +38,10 @@ 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.catalyst.util.ExpressionSQLBuilder
 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.{FieldReference, GeneralSQLExpression, NullOrdering, SortDirection, SortOrder => SortOrderV2, SortValue}
 import org.apache.spark.sql.connector.expressions.aggregate.{AggregateFunc, Aggregation, Avg, Count, CountStar, GeneralAggregateFunc, Max, Min, Sum}
 import org.apache.spark.sql.errors.QueryCompilationErrors
 import org.apache.spark.sql.execution.{InSubqueryExec, RowDataSourceScanExec, SparkPlan}
@@ -707,20 +708,30 @@ object DataSourceStrategy
       agg.aggregateFunction match {
         case aggregate.Min(PushableColumnWithoutNestedColumn(name)) =>
           Some(new Min(FieldReference.column(name)))
+        case aggregate.Min(PushableExpression(expr)) =>
+          Some(new Min(expr))
         case aggregate.Max(PushableColumnWithoutNestedColumn(name)) =>
           Some(new Max(FieldReference.column(name)))
+        case aggregate.Max(PushableExpression(expr)) =>
+          Some(new Max(expr))
         case count: aggregate.Count if count.children.length == 1 =>
           count.children.head match {
             // COUNT(any literal) is the same as COUNT(*)
             case Literal(_, _) => Some(new CountStar())
             case PushableColumnWithoutNestedColumn(name) =>
               Some(new Count(FieldReference.column(name), agg.isDistinct))
+            case PushableExpression(expr) =>
+              Some(new Count(expr, agg.isDistinct))
             case _ => None
           }
         case aggregate.Sum(PushableColumnWithoutNestedColumn(name), _) =>
           Some(new Sum(FieldReference.column(name), agg.isDistinct))
+        case aggregate.Sum(PushableExpression(expr), _) =>
+          Some(new Sum(expr, agg.isDistinct))
         case aggregate.Average(PushableColumnWithoutNestedColumn(name), _) =>
           Some(new Avg(FieldReference.column(name), agg.isDistinct))
+        case aggregate.Average(PushableExpression(expr), _) =>
+          Some(new Avg(expr, agg.isDistinct))
         case aggregate.VariancePop(PushableColumnWithoutNestedColumn(name), _) =>
           Some(new GeneralAggregateFunc(
             "VAR_POP", agg.isDistinct, Array(FieldReference.column(name))))
@@ -860,3 +871,12 @@ object PushableColumnAndNestedColumn extends PushableColumnBase {
 object PushableColumnWithoutNestedColumn extends PushableColumnBase {
   override val nestedPredicatePushdownEnabled = false
 }
+
+/**
+ * Get the SQL string of an expression that can be pushed down.
+ */
+object PushableExpression {
+  def unapply(e: Expression): Option[GeneralSQLExpression] = {
+    new ExpressionSQLBuilder(e).build().map(new GeneralSQLExpression(_))
+  }
+}