[CALCITE-6396] Add ADD_MONTHS function (enabled in Spark library)

core/src/main/java/org/apache/calcite/adapter/enumerable/RexImpTable.java

@@ -116,6 +116,7 @@
 import static org.apache.calcite.sql.fun.SqlInternalOperators.LITERAL_AGG;
 import static org.apache.calcite.sql.fun.SqlInternalOperators.THROW_UNLESS;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.ACOSH;
+import static org.apache.calcite.sql.fun.SqlLibraryOperators.ADD_MONTHS;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.ARRAY;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.ARRAYS_OVERLAP;
 import static org.apache.calcite.sql.fun.SqlLibraryOperators.ARRAYS_ZIP;
@@ -863,6 +864,7 @@ Builder populate2() {
       defineMethod(MAP_VALUES, BuiltInMethod.MAP_VALUES.method, NullPolicy.STRICT);
       defineMethod(MAP_FROM_ARRAYS, BuiltInMethod.MAP_FROM_ARRAYS.method, NullPolicy.ANY);
       defineMethod(MAP_FROM_ENTRIES, BuiltInMethod.MAP_FROM_ENTRIES.method, NullPolicy.STRICT);
+      defineMethod(ADD_MONTHS, BuiltInMethod.ADD_MONTHS2.method, NullPolicy.ANY);
       map.put(STR_TO_MAP, new StringToMapImplementor());
       map.put(ARRAY_CONCAT, new ArrayConcatImplementor());
       map.put(SORT_ARRAY, new SortArrayImplementor());

core/src/main/java/org/apache/calcite/runtime/SqlFunctions.java

@@ -5556,6 +5556,12 @@ public static Map mapFromArrays(List keysArray, List valuesArray) {
     return map;
   }
 
+  public static String addMonths(String date, int months) {
+    LocalDate localDate = LocalDate.parse(date);
+    LocalDate result = localDate.plusMonths(months);
+    return result.toString();
+  }
+
   /** Support the MAP function.
    *
    * <p>odd-indexed elements are keys and even-indexed elements are values.

core/src/main/java/org/apache/calcite/sql/SqlKind.java

@@ -442,6 +442,9 @@ public enum SqlKind {
   /** {@code LEAST} function (Oracle). */
   LEAST,
 
+  /** {@code ADD_MONTHS} function (Spark).  */
+  ADD_MONTHS,
+
   /** {@code DATE_ADD} function (BigQuery Semantics). */
   DATE_ADD,
 

core/src/main/java/org/apache/calcite/sql/fun/SqlLibraryOperators.java

@@ -270,6 +270,13 @@ private static SqlCall transformConvert(SqlValidator validator, SqlCall call) {
               .andThen(SqlTypeTransforms.TO_NULLABLE_ALL),
           OperandTypes.SAME_SAME);
 
+  /** The "NVL(value, value)" function. */
+  @LibraryOperator(libraries = {SPARK})
+  public static final SqlBasicFunction ADD_MONTHS =
+      SqlBasicFunction.create(SqlKind.ADD_MONTHS,
+          ReturnTypes.VARCHAR_NULLABLE,
+          OperandTypes.STRING_INTEGER);
+
   /** The "IFNULL(value, value)" function. */
   @LibraryOperator(libraries = {BIG_QUERY, SPARK})
   public static final SqlFunction IFNULL = NVL.withName("IFNULL");

core/src/main/java/org/apache/calcite/util/BuiltInMethod.java

@@ -474,6 +474,7 @@ public enum BuiltInMethod {
   FLOOR_DIV(Math.class, "floorDiv", long.class, long.class),
   FLOOR_MOD(Math.class, "floorMod", long.class, long.class),
   ADD_MONTHS(DateTimeUtils.class, "addMonths", long.class, int.class),
+  ADD_MONTHS2(SqlFunctions.class, "addMonths", String.class, int.class),
   ADD_MONTHS_INT(DateTimeUtils.class, "addMonths", int.class, int.class),
   SUBTRACT_MONTHS(DateTimeUtils.class, "subtractMonths", long.class,
       long.class),

site/_docs/reference.md

@@ -2674,6 +2674,7 @@ In the following:
 | p | expr :: type                                   | Casts *expr* to *type*
 | m | expr1 <=> expr2                                | Whether two values are equal, treating null values as the same, and it's similar to `IS NOT DISTINCT FROM`
 | * | ACOSH(numeric)                                 | Returns the inverse hyperbolic cosine of *numeric*
+| s | ADD_MONTHS(string, months)                      | Returns the date that is *months* after *string*
 | s | ARRAY([expr [, expr ]*])                       | Construct an array in Apache Spark. The function allows users to use `ARRAY()` to create an empty array
 | s | ARRAY_APPEND(array, element)                   | Appends an *element* to the end of the *array* and returns the result. Type of *element* should be similar to type of the elements of the *array*. If the *array* is null, the function will return null. If an *element* that is null, the null *element* will be added to the end of the *array*
 | s | ARRAY_COMPACT(array)                           | Removes null values from the *array*

testkit/src/main/java/org/apache/calcite/test/SqlOperatorTest.java

@@ -10418,6 +10418,23 @@ void checkNvl(SqlOperatorFixture f0, FunctionAlias functionAlias) {
     f0.forEachLibrary(list(functionAlias.libraries), consumer);
   }
 
+  @Test void testAddMonthsFunc() {
+    final SqlOperatorFixture f0 = fixture();
+    f0.setFor(SqlLibraryOperators.ADD_MONTHS);
+    f0.checkFails("^add_months('2016-08-31', 1)^",
+        "No match found for function signature "
+            + "ADD_MONTHS\\(<CHARACTER>, <NUMERIC>\\)", false);
+
+    final SqlOperatorFixture f = f0.withLibrary(SqlLibrary.SPARK);
+    f.checkScalar("ADD_MONTHS('2016-08-31', 1)", "2016-09-30",
+        "VARCHAR NOT NULL");
+    f.checkScalar("ADD_MONTHS('2016-08-31', 5)", "2017-01-31",
+        "VARCHAR NOT NULL");
+    f.checkScalar("ADD_MONTHS('2016-08-31', 18)", "2018-02-28",
+        "VARCHAR NOT NULL");
+    f.checkNull("ADD_MONTHS(CAST(NULL AS VARCHAR(200)), 18)");
+  }
+
   @Test void testDecodeFunc() {
     checkDecodeFunc(fixture().withLibrary(SqlLibrary.ORACLE));
     checkDecodeFunc(fixture().withLibrary(SqlLibrary.SPARK));