fix(expr): add spark.auron.enableNativeDateTimeFunctions
diff --git a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/ShimsImpl.scala b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/ShimsImpl.scala index dff3628..b4dca22 100644 --- a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/ShimsImpl.scala +++ b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/ShimsImpl.scala
@@ -691,7 +691,7 @@ .build() } - case e if isUnixDateConversion(e) => + case e if NativeConverters.nativeDateTimeFunctionsEnabled && isUnixDateConversion(e) => // Date32 and Int32 share the same epoch-day representation. Some( pb.PhysicalExprNode
diff --git a/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronFunctionSuite.scala b/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronFunctionSuite.scala index 8a06374..8d3e9af 100644 --- a/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronFunctionSuite.scala +++ b/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronFunctionSuite.scala
@@ -18,6 +18,7 @@ import java.text.SimpleDateFormat +import org.apache.spark.SparkConf import org.apache.spark.sql.{AuronQueryTest, Row} import org.apache.spark.sql.auron.NativeConverters import org.apache.spark.sql.catalyst.expressions.Alias @@ -30,62 +31,8 @@ class AuronFunctionSuite extends AuronQueryTest with BaseAuronSQLSuite { - test("float comparisons follow Spark NaN and signed zero ordering") { - withTable("t_float_ordering") { - sql( - "create table t_float_ordering (d1 double, d2 double, f1 float, f2 float) using parquet") - sql("""insert into t_float_ordering values - | (0.0, -0.0, cast(0.0 as float), cast(-0.0 as float)), - | (cast('NaN' as double), cast('NaN' as double), - | cast('NaN' as float), cast('NaN' as float)), - | (cast('NaN' as double), cast('Infinity' as double), - | cast('NaN' as float), cast('Infinity' as float)), - | (null, 1.0, null, cast(1.0 as float)) - |""".stripMargin) - - checkSparkAnswerAndOperator("""select - | d1 = d2, not (d1 = d2), d1 < d2, d1 <= d2, d1 > d2, d1 >= d2, d1 <=> d2, - | f1 = f2, not (f1 = f2), f1 < f2, f1 <= f2, f1 > f2, f1 >= f2, f1 <=> f2 - |from t_float_ordering - |""".stripMargin) - } - } - - test("float scan filters preserve NaN and signed zero results") { - Seq("parquet", "orc").foreach { format => - withTable("t_float_scan") { - sql(s"create table t_float_scan (d double) using $format") - sql("""insert into t_float_scan values - | (0.0), (-0.0), (cast('NaN' as double)) - |""".stripMargin) - - // Spark 3.0/3.1 ORC pushdown can drop NaN before evaluating this filter. - val greaterThan = sql("select d from t_float_scan where d > 100.0") - checkAnswer(greaterThan, Seq(Row(Double.NaN))) - assertPlanIsNative(greaterThan) - checkSparkAnswerAndOperator("select d from t_float_scan where d = 0.0") - - val notIn = sql("select count(*) from t_float_scan where d not in (100.0, 200.0)") - checkAnswer(notIn, Seq(Row(3L))) - assertPlanIsNative(notIn) - } - } - } - - test("float min and max follow Spark NaN ordering") { - withTable("t_float_ordering") { - sql("create table t_float_ordering (d double, f float) using parquet") - sql("""insert into t_float_ordering values - | (cast('Infinity' as double), cast('Infinity' as float)), - | (cast('NaN' as double), cast('NaN' as float)), - | (-1.0, cast(-1.0 as float)), - | (null, null) - |""".stripMargin) - - checkSparkAnswerAndOperator( - "select isnan(max(d)), min(d), isnan(max(f)), min(f) from t_float_ordering") - } - } + override protected def sparkConf: SparkConf = + super.sparkConf.set("spark.auron.enableNativeDateTimeFunctions", "true") test("sum function with float input") { if (AuronTestUtils.isSparkV31OrGreater) {
diff --git a/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronQuerySuite.scala b/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronQuerySuite.scala index c961283..475ebd4 100644 --- a/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronQuerySuite.scala +++ b/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronQuerySuite.scala
@@ -163,7 +163,9 @@ } test("test select multiple spark ext functions with the same signature") { - withSQLConf("spark.auron.udf.singleChildFallback.enabled" -> "true") { + withSQLConf( + "spark.auron.udf.singleChildFallback.enabled" -> "true", + "spark.auron.enableNativeDateTimeFunctions" -> "true") { withTable("t1") { sql("create table t1 using parquet as select '2024-12-18' as event_time") checkSparkAnswerAndOperator("select year(event_time), month(event_time) from t1")
diff --git a/spark-extension-shims-spark/src/test/scala/org/apache/auron/NativeConvertersSuite.scala b/spark-extension-shims-spark/src/test/scala/org/apache/auron/NativeConvertersSuite.scala index 7aac711..5c4bc22 100644 --- a/spark-extension-shims-spark/src/test/scala/org/apache/auron/NativeConvertersSuite.scala +++ b/spark-extension-shims-spark/src/test/scala/org/apache/auron/NativeConvertersSuite.scala
@@ -18,16 +18,105 @@ import org.apache.spark.sql.AuronQueryTest import org.apache.spark.sql.auron.NativeConverters -import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Cast, GreaterThan, Literal} -import org.apache.spark.sql.types.{BooleanType, DataType, DoubleType, FloatType, IntegerType, StringType} +import org.apache.spark.sql.catalyst.expressions.{Cast, Expression, Literal} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{BooleanType, DataType, IntegerType, StringType} import org.apache.auron.protobuf.ScalarFunction +import org.apache.auron.spark.configuration.SparkAuronConfiguration +import org.apache.auron.util.AuronTestUtils class NativeConvertersSuite extends AuronQueryTest with BaseAuronSQLSuite with AuronSQLTestHelper { + private def dateTimeExpressions: Seq[Expression] = { + val functions = Seq( + "year(date'2024-03-15')", + "month(date'2024-03-15')", + "dayofmonth(date'2024-03-15')", + "dayofyear(date'2024-03-15')", + "dayofweek(date'2024-03-15')", + "weekofyear(date'2024-03-15')", + "quarter(date'2024-03-15')", + "hour(timestamp'2024-03-15 14:30:45')", + "minute(timestamp'2024-03-15 14:30:45')", + "second(timestamp'2024-03-15 14:30:45')", + "last_day(date'2024-03-15')", + "next_day(date'2024-03-15', 'MON')", + "datediff(date'2024-03-15', date'2024-03-01')", + "date_add(date'2024-03-15', 1)", + "date_sub(date'2024-03-15', 1)", + "add_months(date'2024-03-15', 1)", + "months_between(timestamp'2024-03-15 14:30:45', timestamp'2024-02-15 14:30:45')", + "date_trunc('day', timestamp'2024-03-15 14:30:45')") ++ + (if (AuronTestUtils.isSparkV30OrGreater) Seq("make_date(2024, 3, 15)") else Nil) ++ + (if (AuronTestUtils.isSparkV31OrGreater) { + Seq("unix_date(date'2024-03-15')", "date_from_unix_date(1)") + } else Nil) + functions.map { function => + sql(s"select $function").queryExecution.analyzed.expressions.head.children.head + } + } + + test("日期时间 native 翻译默认关闭") { + assert(!SparkAuronConfiguration.ENABLE_NATIVE_DATETIME_FUNCTIONS.defaultValue()) + assert(!NativeConverters.nativeDateTimeFunctionsEnabled) + withSQLConf( + "spark.auron.udf.singleChildFallback.enabled" -> "true", + "spark.auron.datetime.extract.enabled" -> "true") { + dateTimeExpressions.foreach { expr => + assert(NativeConverters.convertExpr(expr).hasSparkUdfWrapperExpr) + } + } + } + + test("关闭日期时间 native 翻译时普通表达式和裁剪表达式均回退") { + withSQLConf( + "spark.auron.enableNativeDateTimeFunctions" -> "false", + "spark.auron.datetime.extract.enabled" -> "true") { + for (expr <- dateTimeExpressions; isPruningExpr <- Seq(false, true)) { + var fallbackCalled = false + NativeConverters.convertExprWithFallback( + expr, + isPruningExpr, + e => { + assert(e == expr) + fallbackCalled = true + org.apache.auron.protobuf.PhysicalExprNode.getDefaultInstance + }) + assert(fallbackCalled) + } + assert(NativeConverters.convertExpr(Literal(1)).hasLiteral) + } + } + + test("开启日期时间 native 翻译后恢复现有转换") { + withSQLConf( + "spark.auron.enableNativeDateTimeFunctions" -> "true", + "spark.auron.datetime.extract.enabled" -> "true", + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC") { + dateTimeExpressions.foreach { expr => + val converted = NativeConverters.convertExpr(expr) + assert(converted.hasScalarFunction || converted.hasTryCast) + } + } + } + + test("日期时间总开关不覆盖时分秒提取开关") { + withSQLConf( + "spark.auron.enableNativeDateTimeFunctions" -> "true", + "spark.auron.datetime.extract.enabled" -> "false", + "spark.auron.udf.singleChildFallback.enabled" -> "true") { + dateTimeExpressions + .filter(e => Seq("Hour", "Minute", "Second").contains(e.getClass.getSimpleName)) + .foreach { expr => + assert(NativeConverters.convertExpr(expr).hasSparkUdfWrapperExpr) + } + } + } + private def assertTrimmedCast(rawValue: String, targetType: DataType): Unit = { val expr = Cast(Literal.create(rawValue, StringType), targetType) val nativeExpr = NativeConverters.convertExpr(expr)
diff --git a/spark-extension/src/main/java/org/apache/auron/spark/configuration/SparkAuronConfiguration.java b/spark-extension/src/main/java/org/apache/auron/spark/configuration/SparkAuronConfiguration.java index ac4ba38..14550c3 100644 --- a/spark-extension/src/main/java/org/apache/auron/spark/configuration/SparkAuronConfiguration.java +++ b/spark-extension/src/main/java/org/apache/auron/spark/configuration/SparkAuronConfiguration.java
@@ -577,6 +577,12 @@ .withDescription("Enable decimal arithmetic operations conversion to native Auron implementations.") .withDefaultValue(false); + public static final ConfigOption<Boolean> ENABLE_NATIVE_DATETIME_FUNCTIONS = new SQLConfOption<>(Boolean.class) + .withKey("auron.enableNativeDateTimeFunctions") + .withCategory("Expression/Function Supports") + .withDescription("启用日期时间函数到 native Auron 实现的转换，默认回退到 Spark。") + .withDefaultValue(false); + public static final ConfigOption<Boolean> DATETIME_EXTRACT_ENABLED = new SQLConfOption<>(Boolean.class) .withKey("auron.datetime.extract.enabled") .withCategory("Expression/Function Supports")
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/auron/NativeConverters.scala b/spark-extension/src/main/scala/org/apache/spark/sql/auron/NativeConverters.scala index 2975499..d9b8434 100644 --- a/spark-extension/src/main/scala/org/apache/spark/sql/auron/NativeConverters.scala +++ b/spark-extension/src/main/scala/org/apache/spark/sql/auron/NativeConverters.scala
@@ -89,6 +89,8 @@ def udfJsonEnabled: Boolean = SparkAuronConfiguration.UDF_JSON_ENABLED.get() def udfBrickHouseEnabled: Boolean = SparkAuronConfiguration.UDF_BRICKHOUSE_ENABLED.get() def decimalArithOpEnabled: Boolean = SparkAuronConfiguration.DECIMAL_ARITH_OP_ENABLED.get() + def nativeDateTimeFunctionsEnabled: Boolean = + SparkAuronConfiguration.ENABLE_NATIVE_DATETIME_FUNCTIONS.get() def datetimeExtractEnabled: Boolean = SparkAuronConfiguration.DATETIME_EXTRACT_ENABLED.get() def castTrimStringEnabled: Boolean = SparkAuronConfiguration.CAST_STRING_TRIM_ENABLE.get() def singleChildFallbackEnabled: Boolean = @@ -415,6 +417,11 @@ this.buildExtScalarFunctionNode(_, _, _, isPruningExpr, fallback) sparkExpr match { + case e @ (_: Year | _: Month | _: DayOfMonth | _: DayOfYear | _: DayOfWeek | + _: WeekOfYear | _: Quarter | _: Hour | _: Minute | _: Second | _: LastDay | + _: NextDay | _: DateDiff | _: DateAdd | _: DateSub | _: AddMonths | + _: MonthsBetween | _: TruncTimestamp) if !nativeDateTimeFunctionsEnabled => + fallback(e) case e: NativeExprWrapperBase => e.wrapped case e: Literal => val schema = StructType(Seq(StructField("", e.dataType, e.nullable))) @@ -842,7 +849,8 @@ case e: Acosh => buildScalarFunction(pb.ScalarFunction.Acosh, e.children, e.dataType) case e: Atan => buildScalarFunction(pb.ScalarFunction.Atan, e.children, e.dataType) case e: Exp => buildScalarFunction(pb.ScalarFunction.Exp, e.children, e.dataType) - case e: MakeDate => + case e: MakeDate + if nativeDateTimeFunctionsEnabled => buildExtScalarFunction( "Spark_MakeDate", e.children :+ Literal(Shims.get.getMakeDateFailOnError(e)),