[AURON #2485] Support last_day function natively

Author: Sigma-Ma

native-engine/datafusion-ext-functions/src/lib.rs

@@ -94,6 +94,7 @@ pub fn create_auron_ext_function(
         "Spark_DayOfWeek" => shared_function!(spark_dates::spark_dayofweek),
         "Spark_WeekOfYear" => shared_function!(spark_dates::spark_weekofyear),
         "Spark_Quarter" => shared_function!(spark_dates::spark_quarter),
+        "Spark_LastDay" => shared_function!(spark_dates::spark_last_day),
         "Spark_MakeDate" => shared_function!(spark_dates::spark_make_date),
         "Spark_Hour" => shared_function!(spark_dates::spark_hour),
         "Spark_Minute" => shared_function!(spark_dates::spark_minute),

native-engine/datafusion-ext-functions/src/spark_dates.rs

@@ -277,6 +277,18 @@ pub fn spark_day(args: &[ColumnarValue]) -> Result<ColumnarValue> {
     )?))
 }
 
+pub fn spark_last_day(args: &[ColumnarValue]) -> Result<ColumnarValue> {
+    let input = resolve_local_date32(args)?;
+    let last_day = Date32Array::from_iter(input.iter().map(|opt_days| {
+        opt_days
+            .and_then(NaiveDate::from_epoch_days)
+            .and_then(|date| date.with_day(days_in_month(date.year(), date.month())))
+            .map(|date| date.to_epoch_days())
+    }));
+
+    Ok(ColumnarValue::Array(Arc::new(last_day)))
+}
+
 pub fn spark_make_date(args: &[ColumnarValue]) -> Result<ColumnarValue> {
     if args.len() != 4 {
         return Err(DataFusionError::Execution(
@@ -595,6 +607,32 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_spark_last_day() -> Result<()> {
+        let date = |year, month, day| {
+            NaiveDate::from_ymd_opt(year, month, day)
+                .expect("test date must be valid")
+                .to_epoch_days()
+        };
+        let input = Arc::new(Date32Array::from(vec![
+            Some(date(2009, 1, 12)),
+            Some(date(2024, 2, 10)),
+            Some(date(2023, 2, 10)),
+            Some(date(1969, 12, 1)),
+            None,
+        ]));
+        let args = vec![ColumnarValue::Array(input)];
+        let expected_ret: ArrayRef = Arc::new(Date32Array::from(vec![
+            Some(date(2009, 1, 31)),
+            Some(date(2024, 2, 29)),
+            Some(date(2023, 2, 28)),
+            Some(date(1969, 12, 31)),
+            None,
+        ]));
+        assert_eq!(&spark_last_day(&args)?.into_array(1)?, &expected_ret);
+        Ok(())
+    }
+
     #[test]
     fn test_spark_make_date_null_and_invalid_inputs() -> Result<()> {
         let result = spark_make_date(&[

spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronFunctionSuite.scala

@@ -170,6 +170,21 @@ class AuronFunctionSuite extends AuronQueryTest with BaseAuronSQLSuite {
     }
   }
 
+  test("last_day function") {
+    withTable("t1") {
+      sql("create table t1(c1 date) using parquet")
+      sql("""insert into t1 values
+          |  (date'2009-01-12'),
+          |  (date'2024-02-10'),
+          |  (date'2023-02-10'),
+          |  (date'1969-12-01'),
+          |  (null)
+          |""".stripMargin)
+
+      checkSparkAnswerAndOperator("select last_day(c1) from t1")
+    }
+  }
+
   test("date-part functions with non-UTC timezone") {
     withTable("t1") {
       sql("create table t1(c1 timestamp) using parquet")

spark-extension/src/main/scala/org/apache/spark/sql/auron/NativeConverters.scala

@@ -998,6 +998,8 @@ object NativeConverters extends Logging {
         buildTimePartExt("Spark_WeekOfYear", child, isPruningExpr, fallback)
       case Quarter(child) =>
         buildTimePartExt("Spark_Quarter", child, isPruningExpr, fallback)
+      case e: LastDay =>
+        buildExtScalarFunction("Spark_LastDay", e.children, e.dataType)
 
       case e: Levenshtein =>
         buildScalarFunction(pb.ScalarFunction.Levenshtein, e.children, e.dataType)