Skip to content

Commit

Permalink
Add TIMESTAMPADD Function To OpenSearch SQL Plugin (opensearch-proj…
Browse files Browse the repository at this point in the history 
…ect#1451)

* Add `TIMESTAMPADD` Function To OpenSearch SQL Plugin

Added Testing And Implementation for TIMESTAMPADD

Signed-off-by: GabeFernandez310 <[email protected]>

* Fixed Formatting

Signed-off-by: GabeFernandez310 <[email protected]>

---------

Signed-off-by: GabeFernandez310 <[email protected]>
  • Loading branch information
@GabeFernandez310
GabeFernandez310 authored Mar 20, 2023
1 parent bc39346 commit 7b83e50
Show file tree
Hide file tree
Showing 17 changed files with 719 additions and 2 deletions.
9 changes: 9 additions & 0 deletions core/src/main/java/org/opensearch/sql/expression/DSL.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
75 changes: 75 additions & 0 deletions core/src/main/java/org/opensearch/sql/expression/datetime/DateTimeFunction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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());
Expand Down Expand Up @@ -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.
*/
Expand Down Expand Up @@ -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.
*
Expand Down
1 change: 1 addition & 0 deletions core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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")),
Expand Down
137 changes: 137 additions & 0 deletions core/src/main/java/org/opensearch/sql/expression/function/FunctionDSL.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
*
Expand Down Expand Up @@ -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.
*/
Expand Down Expand Up @@ -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);
};
}
}
31 changes: 31 additions & 0 deletions core/src/main/java/org/opensearch/sql/expression/function/SerializableQuadFunction.java
View file
Original file line number Diff line number Diff line change
@@ -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);
}
Loading

0 comments on commit 7b83e50

Please sign in to comment.