[SPARK-38899][SQL]DS V2 supports push down datetime functions

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

@@ -0,0 +1,120 @@
+/*
+ * 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;
+
+/**
+ * Represent an extract function, which extracts and returns the value of a
+ * and a source expression where the field should be extracted.
+ * <p>
+ * The currently supported field names:
+ * <ol>
+ *  <li>Field Name: <code>SECOND</code>
+ *   <ul>
+ *    <li>SQL semantic: <code>EXTRACT(SECOND FROM source)</code></li>
+ *    <li>Since version: 3.4.0</li>
+ *   </ul>
+ *  </li>
+ *  <li>Field Name: <code>MINUTE</code>
+ *   <ul>
+ *    <li>SQL semantic: <code>EXTRACT(MINUTE FROM source)</code></li>
+ *    <li>Since version: 3.4.0</li>
+ *   </ul>
+ *  </li>
+ *  <li>Field Name: <code>HOUR</code>
+ *   <ul>
+ *    <li>SQL semantic: <code>EXTRACT(HOUR FROM source)</code></li>
+ *    <li>Since version: 3.4.0</li>
+ *   </ul>
+ *  </li>
+ *  <li>Field Name: <code>MONTH</code>
+ *   <ul>
+ *    <li>SQL semantic: <code>EXTRACT(MONTH FROM source)</code></li>
+ *    <li>Since version: 3.4.0</li>
+ *   </ul>
+ *  </li>
+ *  <li>Field Name: <code>QUARTER</code>
+ *   <ul>
+ *    <li>SQL semantic: <code>EXTRACT(QUARTER FROM source)</code></li>
+ *    <li>Since version: 3.4.0</li>
+ *   </ul>
+ *  </li>
+ *  <li>Field Name: <code>YEAR</code>
+ *   <ul>
+ *    <li>SQL semantic: <code>EXTRACT(YEAR FROM source)</code></li>
+ *    <li>Since version: 3.4.0</li>
+ *   </ul>
+ *  </li>
+ *  <li>Field Name: <code>ISO_DAY_OF_WEEK</code>
+ *   <ul>
+ *    <li>SQL semantic: <code>EXTRACT(ISO_DAY_OF_WEEK FROM source)</code></li>
+ *    <li>Database dialects need to follow ISO semantics when handling ISO_DAY_OF_WEEK.</li>
+ *    <li>Since version: 3.4.0</li>
+ *   </ul>
+ *  </li>
+ *  <li>Field Name: <code>DAY</code>
+ *   <ul>
+ *    <li>SQL semantic: <code>EXTRACT(DAY FROM source)</code></li>
+ *    <li>Since version: 3.4.0</li>
+ *   </ul>
+ *  </li>
+ *  <li>Field Name: <code>DOY</code>
+ *   <ul>
+ *    <li>SQL semantic: <code>EXTRACT(DOY FROM source)</code></li>
+ *    <li>Since version: 3.4.0</li>
+ *   </ul>
+ *  </li>
+ *  <li>Field Name: <code>WEEK</code>
+ *   <ul>
+ *    <li>SQL semantic: <code>EXTRACT(WEEK FROM source)</code></li>
+ *    <li>Database dialects need to follow ISO semantics when handling WEEK.</li>
+ *    <li>Since version: 3.4.0</li>
+ *   </ul>
+ *  </li>
+ *  <li>Field Name: <code>YEAR_OF_WEEK</code>
+ *   <ul>
+ *    <li>SQL semantic: <code>EXTRACT(YEAR_OF_WEEK FROM source)</code></li>
+ *    <li>Database dialects need to follow ISO semantics when handling YEAR_OF_WEEK.</li>
+ *    <li>Since version: 3.4.0</li>
+ *   </ul>
+ *  </li>
+ * </ol>
+ *
+ * @since 3.4.0
+ */
+
+@Evolving
+public class Extract implements Expression, Serializable {
+
+  private String field;
+  private Expression source;
+
+  public Extract(String field, Expression source) {
+    this.field = field;
+    this.source = source;
+  }
+
+  public String field() { return field; }
+  public Expression source() { return source; }
+
+  @Override
+  public Expression[] children() { return new Expression[]{ source() }; }
+}

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

@@ -346,6 +346,24 @@
  *    <li>Since version: 3.4.0</li>
  *   </ul>
  *  </li>
+ *  <li>Name: <code>DATE_ADD</code>
+ *   <ul>
+ *    <li>SQL semantic: <code>DATE_ADD(start_date, num_days)</code></li>
+ *    <li>Since version: 3.4.0</li>
+ *   </ul>
+ *  </li>
+ *  <li>Name: <code>DATE_DIFF</code>
+ *   <ul>
+ *    <li>SQL semantic: <code>DATE_DIFF(end_date, start_date)</code></li>
+ *    <li>Since version: 3.4.0</li>
+ *   </ul>
+ *  </li>
+ *  <li>Name: <code>TRUNC</code>
+ *   <ul>
+ *    <li>SQL semantic: <code>TRUNC(date, format)</code></li>
+ *    <li>Since version: 3.4.0</li>
+ *   </ul>
+ *  </li>
  * </ol>
  * Note: SQL semantic conforms ANSI standard, so some expressions are not supported when ANSI off,
  * including: add, subtract, multiply, divide, remainder, pmod.

sql/catalyst/src/main/java/org/apache/spark/sql/connector/util/V2ExpressionSQLBuilder.java

@@ -23,6 +23,7 @@
 
 import org.apache.spark.sql.connector.expressions.Cast;
 import org.apache.spark.sql.connector.expressions.Expression;
+import org.apache.spark.sql.connector.expressions.Extract;
 import org.apache.spark.sql.connector.expressions.NamedReference;
 import org.apache.spark.sql.connector.expressions.GeneralScalarExpression;
 import org.apache.spark.sql.connector.expressions.Literal;
@@ -46,6 +47,9 @@ public String build(Expression expr) {
     } else if (expr instanceof Cast) {
       Cast cast = (Cast) expr;
       return visitCast(build(cast.expression()), cast.dataType());
+    } else if (expr instanceof Extract) {
+      Extract extract = (Extract) expr;
+      return visitExtract(extract.field(), build(extract.source()));
     } else if (expr instanceof GeneralScalarExpression) {
       GeneralScalarExpression e = (GeneralScalarExpression) expr;
       String name = e.name();
@@ -136,6 +140,9 @@ public String build(Expression expr) {
         case "UPPER":
         case "LOWER":
         case "TRANSLATE":
+        case "DATE_ADD":
+        case "DATE_DIFF":
+        case "TRUNC":
           return visitSQLFunction(name,
             Arrays.stream(e.children()).map(c -> build(c)).toArray(String[]::new));
         case "CASE_WHEN": {
@@ -327,4 +334,15 @@ protected String visitTrim(String direction, String[] inputs) {
       return "TRIM(" + direction + " " + inputs[1] + " FROM " + inputs[0] + ")";
     }
   }
+
+  protected String visitExtract(String field, String source) {
+    switch (field) {
+      case "DAY_OF_WEEK":
+        return "(EXTRACT(ISO_DAY_OF_WEEK FROM " + source + ") % 7)+ 1";
+      case "WEEK_DAY":
+        return  "EXTRACT(ISO_DAY_OF_WEEK FROM " + source + ") -1";
+      default:
+        return "EXTRACT(" + field + " FROM " + source + ")";
+    }
+  }
 }

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

@@ -18,7 +18,7 @@
 package org.apache.spark.sql.catalyst.util
 
 import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.connector.expressions.{Cast => V2Cast, Expression => V2Expression, FieldReference, GeneralScalarExpression, LiteralValue, UserDefinedScalarFunc}
+import org.apache.spark.sql.connector.expressions.{Cast => V2Cast, Expression => V2Expression, Extract => V2Extract, FieldReference, GeneralScalarExpression, LiteralValue, UserDefinedScalarFunc}
 import org.apache.spark.sql.connector.expressions.filter.{AlwaysFalse, AlwaysTrue, And => V2And, Not => V2Not, Or => V2Or, Predicate => V2Predicate}
 import org.apache.spark.sql.types.BooleanType
 
@@ -344,6 +344,51 @@ class V2ExpressionBuilder(e: Expression, isPredicate: Boolean = false) {
       } else {
         None
       }
+    case date: DateAdd =>
+      val childrenExpressions = date.children.flatMap(generateExpression(_))
+      if (childrenExpressions.length == date.children.length) {
+        Some(new GeneralScalarExpression("DATE_ADD", childrenExpressions.toArray[V2Expression]))
+      } else {
+        None
+      }
+    case date: DateDiff =>
+      val childrenExpressions = date.children.flatMap(generateExpression(_))
+      if (childrenExpressions.length == date.children.length) {
+        Some(new GeneralScalarExpression("DATE_DIFF", childrenExpressions.toArray[V2Expression]))
+      } else {
+        None
+      }
+    case date: TruncDate =>
+      val childrenExpressions = date.children.flatMap(generateExpression(_))
+      if (childrenExpressions.length == date.children.length) {
+        Some(new GeneralScalarExpression("TRUNC", childrenExpressions.toArray[V2Expression]))
+      } else {
+        None
+      }
+    case Second(child, _) =>
+      generateExpression(child).map(v => new V2Extract("SECOND", v))
+    case Minute(child, _) =>
+      generateExpression(child).map(v => new V2Extract("MINUTE", v))
+    case Hour(child, _) =>
+      generateExpression(child).map(v => new V2Extract("HOUR", v))
+    case Month(child) =>
+      generateExpression(child).map(v => new V2Extract("MONTH", v))
+    case Quarter(child) =>
+      generateExpression(child).map(v => new V2Extract("QUARTER", v))
+    case Year(child) =>
+      generateExpression(child).map(v => new V2Extract("YEAR", v))
+    case DayOfWeek(child) =>
+      generateExpression(child).map(v => new V2Extract("DAY_OF_WEEK", v))
+    case WeekDay(child) =>
+      generateExpression(child).map(v => new V2Extract("WEEK_DAY", v))
+    case DayOfMonth(child) =>
+      generateExpression(child).map(v => new V2Extract("DAY", v))
+    case DayOfYear(child) =>
+      generateExpression(child).map(v => new V2Extract("DOY", v))
+    case WeekOfYear(child) =>
+      generateExpression(child).map(v => new V2Extract("WEEK", v))
+    case YearOfWeek(child) =>
+      generateExpression(child).map(v => new V2Extract("YEAR_OF_WEEK", v))
     // TODO supports other expressions
     case ApplyFunctionExpression(function, children) =>
       val childrenExpressions = children.flatMap(generateExpression(_))

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

@@ -22,10 +22,12 @@ import java.util.Locale
 import java.util.concurrent.ConcurrentHashMap
 
 import scala.collection.JavaConverters._
+import scala.util.control.NonFatal
 
 import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalyst.analysis.{NoSuchNamespaceException, NoSuchTableException, TableAlreadyExistsException}
 import org.apache.spark.sql.connector.catalog.functions.UnboundFunction
+import org.apache.spark.sql.connector.expressions.Expression
 import org.apache.spark.sql.connector.expressions.aggregate.{AggregateFunc, GeneralAggregateFunc}
 import org.apache.spark.sql.execution.datasources.jdbc.JdbcUtils
 import org.apache.spark.sql.types.{BooleanType, ByteType, DataType, DecimalType, ShortType, StringType}
@@ -38,7 +40,9 @@ private[sql] object H2Dialect extends JdbcDialect {
     Set("ABS", "COALESCE", "GREATEST", "LEAST", "RAND", "LOG", "LOG10", "LN", "EXP",
       "POWER", "SQRT", "FLOOR", "CEIL", "ROUND", "SIN", "SINH", "COS", "COSH", "TAN",
       "TANH", "COT", "ASIN", "ACOS", "ATAN", "ATAN2", "DEGREES", "RADIANS", "SIGN",
-      "PI", "SUBSTRING", "UPPER", "LOWER", "TRANSLATE", "TRIM")
+      "PI", "SUBSTRING", "UPPER", "LOWER", "TRANSLATE", "TRIM", "SECOND", "MINUTE",
+      "HOUR", "MONTH", "QUARTER", "YEAR", "DAY", "DOY", "DAY_OF_WEEK", "WEEK_DAY",
+      "WEEK", "YEAR_OF_WEEK")
 
   override def isSupportedFunction(funcName: String): Boolean =
     supportedFunctions.contains(funcName)
@@ -123,4 +127,26 @@ private[sql] object H2Dialect extends JdbcDialect {
     }
     super.classifyException(message, e)
   }
+
+  override def compileExpression(expr: Expression): Option[String] = {
+    val jdbcSQLBuilder = new H2JDBCSQLBuilder()
+    try {
+      Some(jdbcSQLBuilder.build(expr))
+    } catch {
+      case NonFatal(e) =>
+        logWarning("Error occurs while compiling V2 expression", e)
+        None
+    }
+  }
+
+  class H2JDBCSQLBuilder extends JDBCSQLBuilder {
+
+    override def visitExtract(field: String, source: String): String = {
+      field match {
+        case "WEEK" => s"EXTRACT(ISO_WEEK FROM $source)"
+        case "YEAR_OF_WEEK" => s"EXTRACT(ISO_WEEK_YEAR FROM $source)"
+        case _ => super.visitExtract(field, source)
+      }
+    }
+  }
 }

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

@@ -270,6 +270,17 @@ abstract class JdbcDialect extends Serializable with Logging{
           s"${this.getClass.getSimpleName} does not support function: TRIM")
       }
     }
+
+    override def visitExtract(field: String, source: String): String = {
+      if (isSupportedFunction(field)) {
+        super.visitExtract(field, source)
+      } else {
+        // The framework will catch the error and give up the push-down.
+        // Please see `JdbcDialect.compileExpression(expr: Expression)` for more details.
+        throw new UnsupportedOperationException(
+          s"${this.getClass.getSimpleName} does not support function: EXTRACT")
+      }
+    }
   }
 
   /**