Add TIMESTAMPADD Function To OpenSearch SQL Plugin

core/src/main/java/org/opensearch/sql/expression/DSL.java

@@ -919,6 +919,15 @@ public static FunctionExpression time_format(FunctionProperties functionProperti
     return compile(functionProperties, BuiltinFunctionName.TIME_FORMAT, expressions);
   }
 
+  public static FunctionExpression timestampadd(Expression... expressions) {
+    return timestampadd(FunctionProperties.None, expressions);
+  }
+
+  public static FunctionExpression timestampadd(FunctionProperties functionProperties,
+                                               Expression... expressions) {
+    return compile(functionProperties, BuiltinFunctionName.TIMESTAMPADD, expressions);
+  }
+
   public static FunctionExpression utc_date(FunctionProperties functionProperties,
                                                 Expression... args) {
     return compile(functionProperties, BuiltinFunctionName.UTC_DATE, args);

core/src/main/java/org/opensearch/sql/expression/datetime/DateTimeFunction.java

@@ -9,9 +9,12 @@
 
 import static java.time.temporal.ChronoUnit.DAYS;
 import static java.time.temporal.ChronoUnit.HOURS;
+import static java.time.temporal.ChronoUnit.MICROS;
 import static java.time.temporal.ChronoUnit.MINUTES;
 import static java.time.temporal.ChronoUnit.MONTHS;
 import static java.time.temporal.ChronoUnit.SECONDS;
+import static java.time.temporal.ChronoUnit.WEEKS;
+import static java.time.temporal.ChronoUnit.YEARS;
 import static org.opensearch.sql.data.type.ExprCoreType.DATE;
 import static org.opensearch.sql.data.type.ExprCoreType.DATETIME;
 import static org.opensearch.sql.data.type.ExprCoreType.DOUBLE;
@@ -62,6 +65,7 @@
 import java.time.format.DateTimeFormatter;
 import java.time.format.DateTimeParseException;
 import java.time.format.TextStyle;
+import java.time.temporal.ChronoUnit;
 import java.time.temporal.TemporalAmount;
 import java.util.Arrays;
 import java.util.Locale;
@@ -84,6 +88,7 @@
 import org.opensearch.sql.data.model.ExprValue;
 import org.opensearch.sql.data.type.ExprCoreType;
 import org.opensearch.sql.exception.ExpressionEvaluationException;
+import org.opensearch.sql.exception.SemanticCheckException;
 import org.opensearch.sql.expression.function.BuiltinFunctionName;
 import org.opensearch.sql.expression.function.BuiltinFunctionRepository;
 import org.opensearch.sql.expression.function.DefaultFunctionResolver;
@@ -229,6 +234,7 @@ public void register(BuiltinFunctionRepository repository) {
     repository.register(time_to_sec());
     repository.register(timediff());
     repository.register(timestamp());
+    repository.register(timestampadd());
     repository.register(utc_date());
     repository.register(utc_time());
     repository.register(utc_timestamp());
@@ -893,6 +899,22 @@ private DefaultFunctionResolver timestamp() {
             TIMESTAMP, TIMESTAMP, TIMESTAMP));
   }
 
+  private DefaultFunctionResolver timestampadd() {
+    return define(BuiltinFunctionName.TIMESTAMPADD.getName(),
+        impl(nullMissingHandling(DateTimeFunction::exprTimestampAdd),
+            DATETIME, STRING, INTEGER, DATETIME),
+        impl(nullMissingHandling(DateTimeFunction::exprTimestampAdd),
+            DATETIME, STRING, INTEGER, TIMESTAMP),
+        implWithProperties(
+            nullMissingHandlingWithProperties(
+                (functionProperties, part, amount, time) -> exprTimestampAddForTimeType(
+                    functionProperties.getQueryStartClock(),
+                    part,
+                    amount,
+                    time)),
+            DATETIME, STRING, INTEGER, TIME));
+  }
+
   /**
    * TO_DAYS(STRING/DATE/DATETIME/TIMESTAMP). return the day number of the given date.
    */
@@ -1796,6 +1818,59 @@ private ExprValue exprTimeToSec(ExprValue time) {
     return new ExprLongValue(time.timeValue().toSecondOfDay());
   }
 
+  private ExprValue exprTimestampAdd(ExprValue partExpr,
+                                     ExprValue amountExpr,
+                                     ExprValue datetimeExpr) {
+    String part = partExpr.stringValue();
+    int amount = amountExpr.integerValue();
+    LocalDateTime datetime = datetimeExpr.datetimeValue();
+    ChronoUnit temporalUnit;
+
+    switch (part) {
+      case "MICROSECOND":
+        temporalUnit = MICROS;
+        break;
+      case "SECOND":
+        temporalUnit = SECONDS;
+        break;
+      case "MINUTE":
+        temporalUnit = MINUTES;
+        break;
+      case "HOUR":
+        temporalUnit = HOURS;
+        break;
+      case "DAY":
+        temporalUnit = DAYS;
+        break;
+      case "WEEK":
+        temporalUnit = WEEKS;
+        break;
+      case "MONTH":
+        temporalUnit = MONTHS;
+        break;
+      case "QUARTER":
+        temporalUnit = MONTHS;
+        amount *= 3;
+        break;
+      case "YEAR":
+        temporalUnit = YEARS;
+        break;
+      default:
+        return ExprNullValue.of();
+    }
+    return new ExprDatetimeValue(datetime.plus(amount, temporalUnit));
+  }
+
+  private ExprValue exprTimestampAddForTimeType(Clock clock,
+                                                ExprValue partExpr,
+                                                ExprValue amountExpr,
+                                                ExprValue timeExpr) {
+    LocalDateTime datetime = LocalDateTime.of(
+        formatNow(clock).toLocalDate(),
+        timeExpr.timeValue());
+    return exprTimestampAdd(partExpr, amountExpr, new ExprDatetimeValue(datetime));
+  }
+
   /**
    * UTC_DATE implementation for ExprValue.
    *

core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java

@@ -108,6 +108,7 @@ public enum BuiltinFunctionName {
   TIMEDIFF(FunctionName.of("timediff")),
   TIME_TO_SEC(FunctionName.of("time_to_sec")),
   TIMESTAMP(FunctionName.of("timestamp")),
+  TIMESTAMPADD(FunctionName.of("timestampadd")),
   TIME_FORMAT(FunctionName.of("time_format")),
   TO_DAYS(FunctionName.of("to_days")),
   TO_SECONDS(FunctionName.of("to_seconds")),

core/src/main/java/org/opensearch/sql/expression/function/FunctionDSL.java

@@ -183,6 +183,60 @@ public String toString() {
     };
   }
 
+  /**
+   * Implementation of a function that takes three arguments, returns a value, and
+   * requires FunctionProperties to complete.
+   *
+   * @param function   {@link ExprValue} based Binary function.
+   * @param returnType return type.
+   * @param args1Type first argument type.
+   * @param args2Type second argument type.
+   * @param args3Type third argument type.
+   * @return Binary Function Implementation.
+   */
+  public static SerializableFunction<FunctionName, Pair<FunctionSignature, FunctionBuilder>>
+      implWithProperties(
+      SerializableQuadFunction<
+          FunctionProperties,
+          ExprValue,
+          ExprValue,
+          ExprValue,
+          ExprValue> function,
+      ExprType returnType,
+      ExprType args1Type,
+      ExprType args2Type,
+      ExprType args3Type) {
+
+    return functionName -> {
+      FunctionSignature functionSignature =
+          new FunctionSignature(functionName, Arrays.asList(args1Type, args2Type, args3Type));
+      FunctionBuilder functionBuilder =
+          (functionProperties, arguments) -> new FunctionExpression(functionName, arguments) {
+            @Override
+            public ExprValue valueOf(Environment<Expression, ExprValue> valueEnv) {
+              ExprValue arg1 = arguments.get(0).valueOf(valueEnv);
+              ExprValue arg2 = arguments.get(1).valueOf(valueEnv);
+              ExprValue arg3 = arguments.get(2).valueOf(valueEnv);
+              return function.apply(functionProperties, arg1, arg2, arg3);
+            }
+
+            @Override
+            public ExprType type() {
+              return returnType;
+            }
+
+            @Override
+            public String toString() {
+              return String.format("%s(%s)", functionName,
+                  arguments.stream()
+                      .map(Object::toString)
+                      .collect(Collectors.joining(", ")));
+            }
+          };
+      return Pair.of(functionSignature, functionBuilder);
+    };
+  }
+
   /**
    * No Arg Function Implementation.
    *
@@ -275,6 +329,59 @@ public String toString() {
     };
   }
 
+  /**
+   * Quadruple Function Implementation.
+   *
+   * @param function   {@link ExprValue} based unary function.
+   * @param returnType return type.
+   * @param args1Type  argument type.
+   * @param args2Type  argument type.
+   * @param args3Type  argument type.
+   * @return Quadruple Function Implementation.
+   */
+  public static SerializableFunction<FunctionName, Pair<FunctionSignature, FunctionBuilder>> impl(
+      SerializableQuadFunction<ExprValue, ExprValue, ExprValue, ExprValue, ExprValue> function,
+      ExprType returnType,
+      ExprType args1Type,
+      ExprType args2Type,
+      ExprType args3Type,
+      ExprType args4Type) {
+
+    return functionName -> {
+      FunctionSignature functionSignature =
+          new FunctionSignature(functionName, Arrays.asList(
+              args1Type,
+              args2Type,
+              args3Type,
+              args4Type));
+      FunctionBuilder functionBuilder =
+          (functionProperties, arguments) -> new FunctionExpression(functionName, arguments) {
+            @Override
+            public ExprValue valueOf(Environment<Expression, ExprValue> valueEnv) {
+              ExprValue arg1 = arguments.get(0).valueOf(valueEnv);
+              ExprValue arg2 = arguments.get(1).valueOf(valueEnv);
+              ExprValue arg3 = arguments.get(2).valueOf(valueEnv);
+              ExprValue arg4 = arguments.get(3).valueOf(valueEnv);
+              return function.apply(arg1, arg2, arg3, arg4);
+            }
+
+            @Override
+            public ExprType type() {
+              return returnType;
+            }
+
+            @Override
+            public String toString() {
+              return String.format("%s(%s, %s, %s, %s)", functionName, arguments.get(0).toString(),
+                  arguments.get(1).toString(),
+                  arguments.get(2).toString(),
+                  arguments.get(3).toString());
+            }
+          };
+      return Pair.of(functionSignature, functionBuilder);
+    };
+  }
+
   /**
    * Wrapper the unary ExprValue function with default NULL and MISSING handling.
    */
@@ -358,4 +465,34 @@ public SerializableTriFunction<ExprValue, ExprValue, ExprValue, ExprValue> nullM
       }
     };
   }
+
+  /**
+   * Wrapper for the ExprValue function that takes 3 arguments and is aware of FunctionProperties,
+   * with default NULL and MISSING handling.
+   */
+  public static SerializableQuadFunction<
+      FunctionProperties,
+      ExprValue,
+      ExprValue,
+      ExprValue,
+      ExprValue>
+      nullMissingHandlingWithProperties(
+      SerializableQuadFunction<
+          FunctionProperties,
+          ExprValue,
+          ExprValue,
+          ExprValue,
+          ExprValue> implementation) {
+    return (functionProperties, v1, v2, v3) -> {
+      if (v1.isMissing() || v2.isMissing() || v3.isMissing()) {
+        return ExprValueUtils.missingValue();
+      }
+
+      if (v1.isNull() || v2.isNull() || v3.isNull()) {
+        return ExprValueUtils.nullValue();
+      }
+
+      return implementation.apply(functionProperties, v1, v2, v3);
+    };
+  }
 }

core/src/main/java/org/opensearch/sql/expression/function/SerializableQuadFunction.java

@@ -0,0 +1,31 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+
+package org.opensearch.sql.expression.function;
+
+import java.io.Serializable;
+
+/**
+ * Serializable Triple Function.
+ *
+ * @param <T> the type of the first argument to the function
+ * @param <U> the type of the second argument to the function
+ * @param <V> the type of the third argument to the function
+ * @param <W> the type of the fourth argument to the function
+ * @param <R> the type of the result of the function
+ */
+public interface SerializableQuadFunction<T, U, V, W, R> extends Serializable {
+  /**
+   * Applies this function to the given arguments.
+   *
+   * @param t the first function argument
+   * @param u the second function argument
+   * @param v the third function argument
+   * @param w the fourth function argument
+   * @return the function result
+   */
+  R apply(T t, U u, V v, W w);
+}