From 7dcf6b00b44182ade4d036face720787655891e6 Mon Sep 17 00:00:00 2001 From: Sudarshan Date: Mon, 25 May 2026 22:25:49 +0530 Subject: [PATCH] feat: adding add_months datetime expression --- .../spark_expressions_support.md | 2 +- docs/source/user-guide/latest/expressions.md | 1 + native/core/src/execution/jni_api.rs | 2 + .../apache/comet/serde/QueryPlanSerde.scala | 1 + .../org/apache/comet/serde/datetime.scala | 4 +- .../expressions/datetime/add_months.sql | 38 +++++++++++++++++++ 6 files changed, 46 insertions(+), 2 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/datetime/add_months.sql diff --git a/docs/source/contributor-guide/spark_expressions_support.md b/docs/source/contributor-guide/spark_expressions_support.md index 65f941210a..7fd2302199 100644 --- a/docs/source/contributor-guide/spark_expressions_support.md +++ b/docs/source/contributor-guide/spark_expressions_support.md @@ -214,7 +214,7 @@ ### datetime_funcs -- [ ] add_months +- [x] add_months - [x] convert_timezone - [ ] curdate - [ ] current_date diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index a610a83cea..a83d3b9311 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -101,6 +101,7 @@ of expressions that be disabled. | Expression | SQL | | ---------------- | ---------------------------- | +| AddMonths | `add_months` | | ConvertTimezone | `convert_timezone` | | CurrentTimeZone | `current_timezone` | | DateAdd | `date_add` | diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 97e3f851c5..3c908eac78 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -46,6 +46,7 @@ use datafusion_spark::function::array::array_contains::SparkArrayContains; use datafusion_spark::function::bitwise::bit_count::SparkBitCount; use datafusion_spark::function::bitwise::bit_get::SparkBitGet; use datafusion_spark::function::bitwise::bitwise_not::SparkBitwiseNot; +use datafusion_spark::function::datetime::add_months::SparkAddMonths; use datafusion_spark::function::datetime::date_add::SparkDateAdd; use datafusion_spark::function::datetime::date_sub::SparkDateSub; use datafusion_spark::function::datetime::from_utc_timestamp::SparkFromUtcTimestamp; @@ -579,6 +580,7 @@ fn register_datafusion_spark_function(session_ctx: &SessionContext) { session_ctx.register_udf(ScalarUDF::new_from_impl(SparkSha2::default())); session_ctx.register_udf(ScalarUDF::new_from_impl(CharFunc::default())); session_ctx.register_udf(ScalarUDF::new_from_impl(SparkBitGet::default())); + session_ctx.register_udf(ScalarUDF::new_from_impl(SparkAddMonths::default())); session_ctx.register_udf(ScalarUDF::new_from_impl(SparkDateAdd::default())); session_ctx.register_udf(ScalarUDF::new_from_impl(SparkDateSub::default())); session_ctx.register_udf(ScalarUDF::new_from_impl(SparkFromUtcTimestamp::default())); diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index 8d48239e76..6ac8393bc6 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -219,6 +219,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { private[comet] val temporalExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] = Map( + classOf[AddMonths] -> CometAddMonths, classOf[ConvertTimezone] -> CometConvertTimezone, classOf[DateAdd] -> CometDateAdd, classOf[DateDiff] -> CometDateDiff, diff --git a/spark/src/main/scala/org/apache/comet/serde/datetime.scala b/spark/src/main/scala/org/apache/comet/serde/datetime.scala index b57b1e4e56..abf26f88f3 100644 --- a/spark/src/main/scala/org/apache/comet/serde/datetime.scala +++ b/spark/src/main/scala/org/apache/comet/serde/datetime.scala @@ -21,7 +21,7 @@ package org.apache.comet.serde import java.util.Locale -import org.apache.spark.sql.catalyst.expressions.{Attribute, ConvertTimezone, DateAdd, DateDiff, DateFormatClass, DateFromUnixDate, DateSub, DayOfMonth, DayOfWeek, DayOfYear, Days, FromUTCTimestamp, GetDateField, Hour, Hours, LastDay, Literal, MakeDate, Minute, Month, NextDay, Quarter, Second, SecondsToTimestamp, ToUTCTimestamp, TruncDate, TruncTimestamp, UnixDate, UnixTimestamp, WeekDay, WeekOfYear, Year} +import org.apache.spark.sql.catalyst.expressions.{AddMonths, Attribute, ConvertTimezone, DateAdd, DateDiff, DateFormatClass, DateFromUnixDate, DateSub, DayOfMonth, DayOfWeek, DayOfYear, Days, FromUTCTimestamp, GetDateField, Hour, Hours, LastDay, Literal, MakeDate, Minute, Month, NextDay, Quarter, Second, SecondsToTimestamp, ToUTCTimestamp, TruncDate, TruncTimestamp, UnixDate, UnixTimestamp, WeekDay, WeekOfYear, Year} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DateType, DoubleType, FloatType, IntegerType, LongType, StringType, TimestampNTZType, TimestampType} import org.apache.spark.unsafe.types.UTF8String @@ -360,6 +360,8 @@ object CometUnixTimestamp extends CometExpressionSerde[UnixTimestamp] { } } +object CometAddMonths extends CometScalarFunction[AddMonths]("add_months") + object CometDateAdd extends CometScalarFunction[DateAdd]("date_add") object CometDateSub extends CometScalarFunction[DateSub]("date_sub") diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/add_months.sql b/spark/src/test/resources/sql-tests/expressions/datetime/add_months.sql new file mode 100644 index 0000000000..caf1f3ce90 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/datetime/add_months.sql @@ -0,0 +1,38 @@ +-- 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. + +statement +CREATE TABLE test_add_months(d date, m int) USING parquet + +statement +INSERT INTO test_add_months VALUES + (date('2024-01-15'), 1), + (date('2024-12-15'), 2), + (date('2024-03-15'), -2), + (date('2024-06-15'), 0), + (date('2024-01-31'), 1), + (date('2023-01-31'), 1), + (date('2024-02-29'), 12), + (date('2024-01-15'), 24), + (date('2024-01-15'), -24), + (date('1970-01-01'), 600), + (date('2024-06-15'), NULL), + (NULL, 1), + (NULL, NULL) + +query +SELECT d, m, add_months(d, m) FROM test_add_months