Skip to content

Commit 6b78b84

Browse files
yuancuyuancu
committed
Refactor: unify types in udf to use expr types (deprecating SqlTypeName)
Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>
1 parent cc152b5 commit 6b78b84

3 files changed

Lines changed: 26 additions & 27 deletions

File tree

‎core/src/main/java/org/opensearch/sql/expression/function/PPLBuiltinOperators.java‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,11 +20,11 @@
2020
import org.apache.calcite.rex.RexCall;
2121
import org.apache.calcite.sql.SqlOperator;
2222
import org.apache.calcite.sql.type.ReturnTypes;
23-
import org.apache.calcite.sql.type.SqlTypeName;
2423
import org.apache.calcite.sql.type.SqlTypeTransforms;
2524
import org.apache.calcite.sql.util.ReflectiveSqlOperatorTable;
2625
import org.apache.calcite.util.BuiltInMethod;
2726
import org.opensearch.sql.calcite.utils.PPLReturnTypes;
27+
import org.opensearch.sql.data.type.ExprCoreType;
2828
import org.opensearch.sql.expression.datetime.DateTimeFunctions;
2929
import org.opensearch.sql.expression.function.udf.datetime.AddSubDateFunction;
3030
import org.opensearch.sql.expression.function.udf.datetime.CurrentFunction;
@@ -122,15 +122,15 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable {
122122
public static final SqlOperator SECOND = new DatePartFunction(TimeUnit.SECOND).toUDF("SECOND");
123123
public static final SqlOperator MICROSECOND =
124124
new DatePartFunction(TimeUnit.MICROSECOND).toUDF("MICROSECOND");
125-
public static final SqlOperator NOW = new CurrentFunction(SqlTypeName.TIMESTAMP).toUDF("NOW");
125+
public static final SqlOperator NOW = new CurrentFunction(ExprCoreType.TIMESTAMP).toUDF("NOW");
126126
public static final SqlOperator CURRENT_TIME =
127-
new CurrentFunction(SqlTypeName.TIME).toUDF("CURRENT_TIME");
127+
new CurrentFunction(ExprCoreType.TIME).toUDF("CURRENT_TIME");
128128
public static final SqlOperator CURRENT_DATE =
129-
new CurrentFunction(SqlTypeName.DATE).toUDF("CURRENT_DATE");
129+
new CurrentFunction(ExprCoreType.DATE).toUDF("CURRENT_DATE");
130130
public static final SqlOperator DATE_FORMAT =
131-
new FormatFunction(SqlTypeName.DATE).toUDF("DATE_FORMAT");
131+
new FormatFunction(ExprCoreType.DATE).toUDF("DATE_FORMAT");
132132
public static final SqlOperator TIME_FORMAT =
133-
new FormatFunction(SqlTypeName.TIME).toUDF("TIME_FORMAT");
133+
new FormatFunction(ExprCoreType.TIME).toUDF("TIME_FORMAT");
134134
public static final SqlOperator DAYNAME = new PeriodNameFunction(TimeUnit.DAY).toUDF("DAYNAME");
135135
public static final SqlOperator MONTHNAME =
136136
new PeriodNameFunction(TimeUnit.MONTH).toUDF("MONTHNAME");

‎core/src/main/java/org/opensearch/sql/expression/function/udf/datetime/CurrentFunction.java‎

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -8,19 +8,19 @@
88
import java.time.LocalDateTime;
99
import java.util.List;
1010
import lombok.RequiredArgsConstructor;
11-
import org.apache.calcite.DataContext;
1211
import org.apache.calcite.adapter.enumerable.NotNullImplementor;
1312
import org.apache.calcite.adapter.enumerable.NullPolicy;
1413
import org.apache.calcite.adapter.enumerable.RexToLixTranslator;
1514
import org.apache.calcite.linq4j.tree.Expression;
1615
import org.apache.calcite.linq4j.tree.Expressions;
1716
import org.apache.calcite.rex.RexCall;
1817
import org.apache.calcite.sql.type.SqlReturnTypeInference;
19-
import org.apache.calcite.sql.type.SqlTypeName;
2018
import org.opensearch.sql.calcite.utils.UserDefinedFunctionUtils;
2119
import org.opensearch.sql.data.model.ExprDateValue;
2220
import org.opensearch.sql.data.model.ExprTimeValue;
2321
import org.opensearch.sql.data.model.ExprTimestampValue;
22+
import org.opensearch.sql.data.type.ExprCoreType;
23+
import org.opensearch.sql.data.type.ExprType;
2424
import org.opensearch.sql.expression.datetime.DateTimeFunctions;
2525
import org.opensearch.sql.expression.function.FunctionProperties;
2626
import org.opensearch.sql.expression.function.ImplementorUDF;
@@ -37,9 +37,9 @@
3737
* <p>It returns the current date, time, or timestamp based on the specified return type.
3838
*/
3939
public class CurrentFunction extends ImplementorUDF {
40-
private final SqlTypeName returnType;
40+
private final ExprType returnType;
4141

42-
public CurrentFunction(SqlTypeName returnType) {
42+
public CurrentFunction(ExprType returnType) {
4343
super(new CurrentFunctionImplementor(returnType), NullPolicy.NONE);
4444
this.returnType = returnType;
4545
}
@@ -48,39 +48,39 @@ public CurrentFunction(SqlTypeName returnType) {
4848
public SqlReturnTypeInference getReturnTypeInference() {
4949
return opBinding ->
5050
switch (returnType) {
51-
case DATE -> UserDefinedFunctionUtils.NULLABLE_DATE_UDT;
52-
case TIME -> UserDefinedFunctionUtils.NULLABLE_TIME_UDT;
53-
case TIMESTAMP -> UserDefinedFunctionUtils.NULLABLE_TIMESTAMP_UDT;
51+
case ExprCoreType.DATE -> UserDefinedFunctionUtils.NULLABLE_DATE_UDT;
52+
case ExprCoreType.TIME -> UserDefinedFunctionUtils.NULLABLE_TIME_UDT;
53+
case ExprCoreType.TIMESTAMP -> UserDefinedFunctionUtils.NULLABLE_TIMESTAMP_UDT;
5454
default -> throw new IllegalArgumentException("Unsupported return type: " + returnType);
5555
};
5656
}
5757

5858
@RequiredArgsConstructor
5959
public static class CurrentFunctionImplementor implements NotNullImplementor {
60-
private final SqlTypeName returnType;
60+
private final ExprType returnType;
6161

6262
@Override
6363
public Expression implement(
6464
RexToLixTranslator translator, RexCall call, List<Expression> translatedOperands) {
6565

6666
String functionName =
6767
switch (returnType) {
68-
case DATE -> "currentDate";
69-
case TIME -> "currentTime";
70-
case TIMESTAMP -> "currentTimestamp";
68+
case ExprCoreType.DATE -> "currentDate";
69+
case ExprCoreType.TIME -> "currentTime";
70+
case ExprCoreType.TIMESTAMP -> "currentTimestamp";
7171
default -> throw new IllegalArgumentException("Unsupported return type: " + returnType);
7272
};
7373

74-
Expression now =
74+
Expression properties =
7575
Expressions.call(
76-
CurrentFunctionImplementor.class, "getNowFromProperties", translator.getRoot());
76+
UserDefinedFunctionUtils.class, "restoreFunctionProperties", translator.getRoot());
77+
Expression now =
78+
Expressions.call(CurrentFunctionImplementor.class, "getNowFromProperties", properties);
7779

7880
return Expressions.call(CurrentFunctionImplementor.class, functionName, now);
7981
}
8082

81-
public static LocalDateTime getNowFromProperties(DataContext propertyContext) {
82-
FunctionProperties functionProperties =
83-
UserDefinedFunctionUtils.restoreFunctionProperties(propertyContext);
83+
public static LocalDateTime getNowFromProperties(FunctionProperties functionProperties) {
8484
return DateTimeFunctions.formatNow(functionProperties.getQueryStartClock());
8585
}
8686

‎core/src/main/java/org/opensearch/sql/expression/function/udf/datetime/FormatFunction.java‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@
1818
import org.apache.calcite.linq4j.tree.Expressions;
1919
import org.apache.calcite.rex.RexCall;
2020
import org.apache.calcite.sql.type.SqlReturnTypeInference;
21-
import org.apache.calcite.sql.type.SqlTypeName;
2221
import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory;
2322
import org.opensearch.sql.calcite.utils.PPLReturnTypes;
2423
import org.opensearch.sql.calcite.utils.UserDefinedFunctionUtils;
@@ -41,9 +40,9 @@
4140
* </ul>
4241
*/
4342
public class FormatFunction extends ImplementorUDF {
44-
public FormatFunction(SqlTypeName functionType) {
43+
public FormatFunction(ExprType functionType) {
4544
super(new DataFormatImplementor(functionType), NullPolicy.ANY);
46-
if (!functionType.equals(SqlTypeName.DATE) && !functionType.equals(SqlTypeName.TIME)) {
45+
if (!functionType.equals(ExprCoreType.DATE) && !functionType.equals(ExprCoreType.TIME)) {
4746
throw new IllegalArgumentException(
4847
"Function type can only be DATE or TIME, but got: " + functionType);
4948
}
@@ -56,7 +55,7 @@ public SqlReturnTypeInference getReturnTypeInference() {
5655

5756
@RequiredArgsConstructor
5857
public static class DataFormatImplementor implements NotNullImplementor {
59-
private final SqlTypeName functionType;
58+
private final ExprType functionType;
6059

6160
@Override
6261
public Expression implement(
@@ -75,7 +74,7 @@ public Expression implement(
7574
Expressions.constant(type));
7675
Expression format = Expressions.new_(ExprStringValue.class, translatedOperands.get(1));
7776

78-
if (SqlTypeName.TIME.equals(functionType)) {
77+
if (ExprCoreType.TIME.equals(functionType)) {
7978
return Expressions.call(DataFormatImplementor.class, "timeFormat", datetime, format);
8079
} else {
8180
if (ExprCoreType.TIME.equals(type)) {

0 commit comments

Comments
 (0)