Skip to content

Commit 73c9156

Browse files
andygroveclaude
andcommitted
feat: add support for last_day expression
Adds native Comet support for Spark's last_day function, which returns the last day of the month for a given date. Uses the SparkLastDay implementation from datafusion-spark crate. Closes #3090 Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
1 parent c3f22eb commit 73c9156

File tree

5 files changed

+33
-1
lines changed

5 files changed

+33
-1
lines changed

docs/source/user-guide/latest/configs.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -266,6 +266,7 @@ These settings can be used to determine which parts of the plan are accelerated
266266
| `spark.comet.expression.IsNull.enabled` | Enable Comet acceleration for `IsNull` | true |
267267
| `spark.comet.expression.JsonToStructs.enabled` | Enable Comet acceleration for `JsonToStructs` | true |
268268
| `spark.comet.expression.KnownFloatingPointNormalized.enabled` | Enable Comet acceleration for `KnownFloatingPointNormalized` | true |
269+
| `spark.comet.expression.LastDay.enabled` | Enable Comet acceleration for `LastDay` | true |
269270
| `spark.comet.expression.Length.enabled` | Enable Comet acceleration for `Length` | true |
270271
| `spark.comet.expression.LessThan.enabled` | Enable Comet acceleration for `LessThan` | true |
271272
| `spark.comet.expression.LessThanOrEqual.enabled` | Enable Comet acceleration for `LessThanOrEqual` | true |

native/core/src/execution/jni_api.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ use datafusion_spark::function::bitwise::bit_get::SparkBitGet;
4444
use datafusion_spark::function::bitwise::bitwise_not::SparkBitwiseNot;
4545
use datafusion_spark::function::datetime::date_add::SparkDateAdd;
4646
use datafusion_spark::function::datetime::date_sub::SparkDateSub;
47+
use datafusion_spark::function::datetime::last_day::SparkLastDay;
4748
use datafusion_spark::function::hash::sha1::SparkSha1;
4849
use datafusion_spark::function::hash::sha2::SparkSha2;
4950
use datafusion_spark::function::math::expm1::SparkExpm1;
@@ -345,6 +346,7 @@ fn register_datafusion_spark_function(session_ctx: &SessionContext) {
345346
session_ctx.register_udf(ScalarUDF::new_from_impl(SparkBitGet::default()));
346347
session_ctx.register_udf(ScalarUDF::new_from_impl(SparkDateAdd::default()));
347348
session_ctx.register_udf(ScalarUDF::new_from_impl(SparkDateSub::default()));
349+
session_ctx.register_udf(ScalarUDF::new_from_impl(SparkLastDay::default()));
348350
session_ctx.register_udf(ScalarUDF::new_from_impl(SparkSha1::default()));
349351
session_ctx.register_udf(ScalarUDF::new_from_impl(SparkConcat::default()));
350352
session_ctx.register_udf(ScalarUDF::new_from_impl(SparkBitwiseNot::default()));

spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -187,6 +187,7 @@ object QueryPlanSerde extends Logging with CometExprShim {
187187
classOf[DateAdd] -> CometDateAdd,
188188
classOf[DateSub] -> CometDateSub,
189189
classOf[FromUnixTime] -> CometFromUnixTime,
190+
classOf[LastDay] -> CometLastDay,
190191
classOf[Hour] -> CometHour,
191192
classOf[Minute] -> CometMinute,
192193
classOf[Second] -> CometSecond,

spark/src/main/scala/org/apache/comet/serde/datetime.scala

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ package org.apache.comet.serde
2121

2222
import java.util.Locale
2323

24-
import org.apache.spark.sql.catalyst.expressions.{Attribute, DateAdd, DateSub, DayOfMonth, DayOfWeek, DayOfYear, GetDateField, Hour, Literal, Minute, Month, Quarter, Second, TruncDate, TruncTimestamp, WeekDay, WeekOfYear, Year}
24+
import org.apache.spark.sql.catalyst.expressions.{Attribute, DateAdd, DateSub, DayOfMonth, DayOfWeek, DayOfYear, GetDateField, Hour, LastDay, Literal, Minute, Month, Quarter, Second, TruncDate, TruncTimestamp, WeekDay, WeekOfYear, Year}
2525
import org.apache.spark.sql.types.{DateType, IntegerType}
2626
import org.apache.spark.unsafe.types.UTF8String
2727

@@ -258,6 +258,8 @@ object CometDateAdd extends CometScalarFunction[DateAdd]("date_add")
258258

259259
object CometDateSub extends CometScalarFunction[DateSub]("date_sub")
260260

261+
object CometLastDay extends CometScalarFunction[LastDay]("last_day")
262+
261263
object CometTruncDate extends CometExpressionSerde[TruncDate] {
262264

263265
val supportedFormats: Seq[String] =

spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,4 +122,30 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH
122122
StructField("fmt", DataTypes.StringType, true)))
123123
FuzzDataGenerator.generateDataFrame(r, spark, schema, 1000, DataGenOptions())
124124
}
125+
126+
test("last_day") {
127+
val r = new Random(42)
128+
val schema = StructType(Seq(StructField("c0", DataTypes.DateType, true)))
129+
val df = FuzzDataGenerator.generateDataFrame(r, spark, schema, 1000, DataGenOptions())
130+
df.createOrReplaceTempView("tbl")
131+
132+
// Basic test with random dates
133+
checkSparkAnswerAndOperator("SELECT c0, last_day(c0) FROM tbl ORDER BY c0")
134+
135+
// Disable constant folding to ensure literal expressions are executed by Comet
136+
withSQLConf(
137+
SQLConf.OPTIMIZER_EXCLUDED_RULES.key ->
138+
"org.apache.spark.sql.catalyst.optimizer.ConstantFolding") {
139+
// Test with literal dates - various months
140+
checkSparkAnswerAndOperator(
141+
"SELECT last_day(DATE('2024-01-15')), last_day(DATE('2024-02-15')), last_day(DATE('2024-12-01'))")
142+
143+
// Test leap year handling (February)
144+
checkSparkAnswerAndOperator(
145+
"SELECT last_day(DATE('2024-02-01')), last_day(DATE('2023-02-01'))")
146+
147+
// Test null handling
148+
checkSparkAnswerAndOperator("SELECT last_day(NULL)")
149+
}
150+
}
125151
}

0 commit comments

Comments
 (0)