Skip to content

Commit

Permalink
Add Support For TIME Type in "*_OF_YEAR" Functions (#199)
Browse files Browse the repository at this point in the history 
Added Support And Tests For Time Type in day_of_year, week_of_year, month_of_year Functions
Signed-off-by: GabeFernandez310 <[email protected]>
  • Loading branch information
@GabeFernandez310
GabeFernandez310 authored Jan 5, 2023
1 parent 91ef34d commit c586115
Show file tree
Hide file tree
Showing 10 changed files with 602 additions and 232 deletions.
20 changes: 12 additions & 8 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 @@ -334,8 +334,9 @@ public static FunctionExpression dayofyear(Expression... expressions) {
return compile(FunctionProperties.None, BuiltinFunctionName.DAYOFYEAR, expressions);
}

public static FunctionExpression day_of_year(Expression... expressions) {
return compile(FunctionProperties.None, BuiltinFunctionName.DAY_OF_YEAR, expressions);
public static FunctionExpression day_of_year(
FunctionProperties functionProperties, Expression... expressions) {
return compile(functionProperties, BuiltinFunctionName.DAY_OF_YEAR, expressions);
}

public static FunctionExpression from_days(Expression... expressions) {
Expand All @@ -358,8 +359,9 @@ public static FunctionExpression month(Expression... expressions) {
return compile(FunctionProperties.None, BuiltinFunctionName.MONTH, expressions);
}

public static FunctionExpression month_of_year(Expression... expressions) {
return compile(FunctionProperties.None, BuiltinFunctionName.MONTH_OF_YEAR, expressions);
public static FunctionExpression month_of_year(
FunctionProperties functionProperties, Expression... expressions) {
return compile(functionProperties, BuiltinFunctionName.MONTH_OF_YEAR, expressions);
}

public static FunctionExpression monthname(Expression... expressions) {
Expand Down Expand Up @@ -398,12 +400,14 @@ public static FunctionExpression to_days(Expression... expressions) {
return compile(FunctionProperties.None, BuiltinFunctionName.TO_DAYS, expressions);
}

public static FunctionExpression week(Expression... expressions) {
return compile(FunctionProperties.None, BuiltinFunctionName.WEEK, expressions);
public static FunctionExpression week(
FunctionProperties functionProperties, Expression... expressions) {
return compile(functionProperties, BuiltinFunctionName.WEEK, expressions);
}

public static FunctionExpression week_of_year(Expression... expressions) {
return compile(FunctionProperties.None, BuiltinFunctionName.WEEK_OF_YEAR, expressions);
public static FunctionExpression week_of_year(
FunctionProperties functionProperties, Expression... expressions) {
return compile(functionProperties, BuiltinFunctionName.WEEK_OF_YEAR, expressions);
}

public static FunctionExpression year(Expression... expressions) {
Expand Down
31 changes: 31 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 @@ -20,6 +20,7 @@
import static org.opensearch.sql.expression.function.FunctionDSL.impl;
import static org.opensearch.sql.expression.function.FunctionDSL.implWithProperties;
import static org.opensearch.sql.expression.function.FunctionDSL.nullMissingHandling;
import static org.opensearch.sql.expression.function.FunctionDSL.nullMissingHandlingWithProperties;
import static org.opensearch.sql.utils.DateTimeFormatters.DATE_FORMATTER_LONG_YEAR;
import static org.opensearch.sql.utils.DateTimeFormatters.DATE_FORMATTER_SHORT_YEAR;
import static org.opensearch.sql.utils.DateTimeFormatters.DATE_TIME_FORMATTER_LONG_YEAR;
Expand Down Expand Up @@ -82,6 +83,9 @@ public class DateTimeFunction {
// 32536771199.999999, or equivalent '3001-01-18 23:59:59.999999' UTC
private static final Double MYSQL_MAX_TIMESTAMP = 32536771200d;

// Mode used for week/week_of_year function by default when no argument is provided
private static final ExprIntegerValue DEFAULT_WEEK_OF_YEAR_MODE = new ExprIntegerValue(0);

/**
* Register Date and Time Functions.
*
Expand Down Expand Up @@ -368,6 +372,9 @@ private DefaultFunctionResolver dayOfWeek() {
*/
private DefaultFunctionResolver dayOfYear(BuiltinFunctionName dayOfYear) {
return define(dayOfYear.getName(),
implWithProperties(nullMissingHandlingWithProperties((functionProperties, arg)
-> DateTimeFunction.dayOfYearToday(
functionProperties.getQueryStartClock())), INTEGER, TIME),
impl(nullMissingHandling(DateTimeFunction::exprDayOfYear), INTEGER, DATE),
impl(nullMissingHandling(DateTimeFunction::exprDayOfYear), INTEGER, DATETIME),
impl(nullMissingHandling(DateTimeFunction::exprDayOfYear), INTEGER, TIMESTAMP),
Expand Down Expand Up @@ -441,6 +448,9 @@ private DefaultFunctionResolver minute() {
*/
private DefaultFunctionResolver month(BuiltinFunctionName month) {
return define(month.getName(),
implWithProperties(nullMissingHandlingWithProperties((functionProperties, arg)
-> DateTimeFunction.monthOfYearToday(
functionProperties.getQueryStartClock())), INTEGER, TIME),
impl(nullMissingHandling(DateTimeFunction::exprMonth), INTEGER, DATE),
impl(nullMissingHandling(DateTimeFunction::exprMonth), INTEGER, DATETIME),
impl(nullMissingHandling(DateTimeFunction::exprMonth), INTEGER, TIMESTAMP),
Expand Down Expand Up @@ -602,10 +612,18 @@ private DefaultFunctionResolver utc_timestamp() {
*/
private DefaultFunctionResolver week(BuiltinFunctionName week) {
return define(week.getName(),
implWithProperties(nullMissingHandlingWithProperties((functionProperties, arg)
-> DateTimeFunction.weekOfYearToday(
DEFAULT_WEEK_OF_YEAR_MODE,
functionProperties.getQueryStartClock())), INTEGER, TIME),
impl(nullMissingHandling(DateTimeFunction::exprWeekWithoutMode), INTEGER, DATE),
impl(nullMissingHandling(DateTimeFunction::exprWeekWithoutMode), INTEGER, DATETIME),
impl(nullMissingHandling(DateTimeFunction::exprWeekWithoutMode), INTEGER, TIMESTAMP),
impl(nullMissingHandling(DateTimeFunction::exprWeekWithoutMode), INTEGER, STRING),
implWithProperties(nullMissingHandlingWithProperties((functionProperties, time, modeArg)
-> DateTimeFunction.weekOfYearToday(
modeArg,
functionProperties.getQueryStartClock())), INTEGER, TIME, INTEGER),
impl(nullMissingHandling(DateTimeFunction::exprWeek), INTEGER, DATE, INTEGER),
impl(nullMissingHandling(DateTimeFunction::exprWeek), INTEGER, DATETIME, INTEGER),
impl(nullMissingHandling(DateTimeFunction::exprWeek), INTEGER, TIMESTAMP, INTEGER),
Expand Down Expand Up @@ -646,6 +664,15 @@ private DefaultFunctionResolver date_format() {
);
}

private ExprValue dayOfYearToday(Clock clock) {
return new ExprIntegerValue(LocalDateTime.now(clock).getDayOfYear());
}

private ExprValue weekOfYearToday(ExprValue mode, Clock clock) {
return new ExprIntegerValue(
CalendarLookup.getWeekNumber(mode.integerValue(), LocalDateTime.now(clock).toLocalDate()));
}

/**
* ADDDATE function implementation for ExprValue.
*
Expand Down Expand Up @@ -1241,6 +1268,10 @@ private ExprValue exprYear(ExprValue date) {
return new ExprIntegerValue(date.dateValue().getYear());
}

private ExprValue monthOfYearToday(Clock clock) {
return new ExprIntegerValue(LocalDateTime.now(clock).getMonthValue());
}

private LocalDateTime formatNow(Clock clock) {
return formatNow(clock, 0);
}
Expand Down
91 changes: 66 additions & 25 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 @@ -137,6 +137,52 @@ public String toString() {
};
}

/**
* Implementation of a function that takes two 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.
* @return Binary Function Implementation.
*/
public static SerializableFunction<FunctionName, Pair<FunctionSignature, FunctionBuilder>>
implWithProperties(
SerializableTriFunction<FunctionProperties, ExprValue, ExprValue, ExprValue> function,
ExprType returnType,
ExprType args1Type,
ExprType args2Type) {

return functionName -> {
FunctionSignature functionSignature =
new FunctionSignature(functionName, Arrays.asList(args1Type, args2Type));
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);
return function.apply(functionProperties, arg1, arg2);
}

@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 @@ -181,31 +227,8 @@ public static SerializableFunction<FunctionName, Pair<FunctionSignature, Functio
ExprType args1Type,
ExprType args2Type) {

return functionName -> {
FunctionSignature functionSignature =
new FunctionSignature(functionName, Arrays.asList(args1Type, args2Type));
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);
return function.apply(arg1, arg2);
}

@Override
public ExprType type() {
return returnType;
}

@Override
public String toString() {
return String.format("%s(%s, %s)", functionName, arguments.get(0).toString(),
arguments.get(1).toString());
}
};
return Pair.of(functionSignature, functionBuilder);
};
return implWithProperties((fp, arg1, arg2) ->
function.apply(arg1, arg2), returnType, args1Type, args2Type);
}

/**
Expand Down Expand Up @@ -317,4 +340,22 @@ public SerializableTriFunction<ExprValue, ExprValue, ExprValue, ExprValue> nullM
}
};
}

/**
* Wrapper for the ExprValue function that takes 2 arguments and is aware of FunctionProperties,
* with default NULL and MISSING handling.
*/
public static SerializableTriFunction<FunctionProperties, ExprValue, ExprValue, ExprValue>
nullMissingHandlingWithProperties(
SerializableTriFunction<FunctionProperties, ExprValue, ExprValue,ExprValue> implementation) {
return (functionProperties, v1, v2) -> {
if (v1.isMissing() || v2.isMissing()) {
return ExprValueUtils.missingValue();
} else if (v1.isNull() || v2.isNull()) {
return ExprValueUtils.nullValue();
} else {
return implementation.apply(functionProperties, v1, v2);
}
};
}
}
Loading

0 comments on commit c586115

Please sign in to comment.