feat: date_add and date_sub functions

native/core/src/execution/datafusion/expressions/comet_scalar_funcs.rs

@@ -20,9 +20,9 @@ use datafusion_comet_spark_expr::scalar_funcs::hash_expressions::{
     spark_sha224, spark_sha256, spark_sha384, spark_sha512,
 };
 use datafusion_comet_spark_expr::scalar_funcs::{
-    spark_ceil, spark_decimal_div, spark_floor, spark_hex, spark_isnan, spark_make_decimal,
-    spark_murmur3_hash, spark_read_side_padding, spark_round, spark_unhex, spark_unscaled_value,
-    spark_xxhash64, SparkChrFunc,
+    spark_ceil, spark_date_add, spark_decimal_div, spark_floor, spark_hex, spark_isnan,
+    spark_make_decimal, spark_murmur3_hash, spark_read_side_padding, spark_round, spark_unhex,
+    spark_unscaled_value, spark_xxhash64, SparkChrFunc,
 };
 use datafusion_common::{DataFusionError, Result as DataFusionResult};
 use datafusion_expr::registry::FunctionRegistry;
@@ -121,6 +121,10 @@ pub fn create_comet_physical_fun(
             let func = Arc::new(spark_sha512);
             make_comet_scalar_udf!("sha512", func, without data_type)
         }
+        "date_add" => {
+            let func = Arc::new(spark_date_add);
+            make_comet_scalar_udf!("date_add", func, without data_type)
+        }
         _ => registry.udf(fun_name).map_err(|e| {
             DataFusionError::Execution(format!(
                 "Function {fun_name} not found in the registry: {e}",

native/spark-expr/src/scalar_funcs.rs

@@ -15,16 +15,20 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use arrow::compute::kernels::numeric::add;
+use arrow::datatypes::IntervalDayTime;
 use arrow::{
     array::{
         ArrayRef, AsArray, Decimal128Builder, Float32Array, Float64Array, Int16Array, Int32Array,
         Int64Array, Int64Builder, Int8Array, OffsetSizeTrait,
     },
     datatypes::{validate_decimal_precision, Decimal128Type, Int64Type},
 };
-use arrow_array::builder::GenericStringBuilder;
+use arrow_array::builder::{GenericStringBuilder, IntervalDayTimeBuilder};
+use arrow_array::types::Int32Type;
 use arrow_array::{Array, ArrowNativeTypeOp, BooleanArray, Decimal128Array};
 use arrow_schema::{DataType, DECIMAL128_MAX_PRECISION};
+use datafusion::physical_expr_common::datum;
 use datafusion::{functions::math::round::round, physical_plan::ColumnarValue};
 use datafusion_common::{
     cast::as_generic_string_array, exec_err, internal_err, DataFusionError,
@@ -547,3 +551,32 @@ pub fn spark_isnan(args: &[ColumnarValue]) -> Result<ColumnarValue, DataFusionEr
         },
     }
 }
+
+/// Spark-compatible `date_add` expression, which assumes days for the second argument, but we
+/// cannot directly add that to a Date32. We generate an IntervalDayTime from the second argument,
+/// and use DataFusion's interface to apply Arrow's add operator.
+pub fn spark_date_add(args: &[ColumnarValue]) -> Result<ColumnarValue, DataFusionError> {
+    let start = &args[0];
+    if let ColumnarValue::Scalar(ScalarValue::Int32(Some(days))) = &args[1] {
+        let interval = IntervalDayTime::new(*days, 0);
+        let interval_cv = ColumnarValue::Scalar(ScalarValue::IntervalDayTime(Some(interval)));
+        let result = datum::apply(start, &interval_cv, add)?;
+        return Ok(result);
+    } else if let ColumnarValue::Array(days) = &args[1] {
+        let mut interval_builder = IntervalDayTimeBuilder::new();
+        for day in days.as_primitive::<Int32Type>().into_iter() {
+            if let Some(non_null_day) = day {
+                interval_builder.append_value(IntervalDayTime::new(non_null_day, 0));
+            } else {
+                interval_builder.append_null();
+            }
+        }
+        let interval_cv = ColumnarValue::Array(Arc::new(interval_builder.finish()));
+        let result = datum::apply(start, &interval_cv, add)?;
+        return Ok(result);
+    }
+    Err(DataFusionError::Internal(format!(
+        "Unsupported data type {:?} for function date_add",
+        args,
+    )))
+}

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

@@ -1492,6 +1492,12 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
             None
           }
 
+        case DateAdd(left, right) =>
+          val leftExpr = exprToProtoInternal(left, inputs)
+          val rightExpr = exprToProtoInternal(right, inputs)
+          val optExpr = scalarExprToProtoWithReturnType("date_add", DateType, leftExpr, rightExpr)
+          optExprWithInfo(optExpr, expr, left, right)
+
         case TruncDate(child, format) =>
           val childExpr = exprToProtoInternal(child, inputs)
           val formatExpr = exprToProtoInternal(format, inputs)

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

@@ -146,6 +146,32 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
     }
   }
 
+  test("date_add with scalar") {
+    Seq(true, false).foreach { dictionaryEnabled =>
+      withTempDir { dir =>
+        val path = new Path(dir.toURI.toString, "test.parquet")
+        makeParquetFileAllTypes(path, dictionaryEnabled = dictionaryEnabled, 10000)
+        withParquetTable(path.toString, "tbl") {
+          checkSparkAnswerAndOperator(
+            "SELECT _2 FROM tbl WHERE _20 + 2 > CAST('1970-01-02' AS DATE)")
+        }
+      }
+    }
+  }
+
+  test("date_add with array") {
+    Seq(true, false).foreach { dictionaryEnabled =>
+      withTempDir { dir =>
+        val path = new Path(dir.toURI.toString, "test.parquet")
+        makeParquetFileAllTypes(path, dictionaryEnabled = dictionaryEnabled, 10000)
+        withParquetTable(path.toString, "tbl") {
+          checkSparkAnswerAndOperator(
+            "SELECT _2 FROM tbl WHERE _20 + _4 > CAST('1970-01-02' AS DATE)")
+        }
+      }
+    }
+  }
+
   test("dictionary arithmetic") {
     // TODO: test ANSI mode
     withSQLConf(SQLConf.ANSI_ENABLED.key -> "false", "parquet.enable.dictionary" -> "true") {