Skip to content

Commit e1e4deb

Browse files
yuancuyuancu
committed
Define type checkers for UDFs
Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>
1 parent a887e2d commit e1e4deb

32 files changed

Lines changed: 395 additions & 59 deletions

‎core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -338,7 +338,7 @@ public RexNode visitFunction(Function node, CalcitePlanContext context) {
338338
List<RexNode> arguments =
339339
node.getFuncArgs().stream().map(arg -> analyze(arg, context)).toList();
340340
RexNode resolvedNode =
341-
PPLFuncImpTable.INSTANCE.resolveSafe(
341+
PPLFuncImpTable.INSTANCE.resolve(
342342
context.rexBuilder, node.getFuncName(), arguments.toArray(new RexNode[0]));
343343
if (resolvedNode != null) {
344344
return resolvedNode;

‎core/src/main/java/org/opensearch/sql/calcite/utils/PPLOperandTypes.java‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
/*
2+
*
3+
* * Copyright OpenSearch Contributors
4+
* * SPDX-License-Identifier: Apache-2.0
5+
*
6+
*/
7+
8+
package org.opensearch.sql.calcite.utils;
9+
10+
import org.apache.calcite.sql.type.CompositeOperandTypeChecker;
11+
import org.apache.calcite.sql.type.FamilyOperandTypeChecker;
12+
import org.apache.calcite.sql.type.OperandTypes;
13+
import org.apache.calcite.sql.type.SqlTypeFamily;
14+
import org.opensearch.sql.expression.function.UDFOperandMetadata;
15+
16+
public class PPLOperandTypes {
17+
private PPLOperandTypes() {}
18+
19+
public static final UDFOperandMetadata DATETIME_OR_STRING =
20+
UDFOperandMetadata.wrap(
21+
(CompositeOperandTypeChecker) OperandTypes.DATETIME.or(OperandTypes.STRING));
22+
public static final UDFOperandMetadata DATETIME_DATETIME =
23+
UDFOperandMetadata.wrap(OperandTypes.family(SqlTypeFamily.DATETIME, SqlTypeFamily.DATETIME));
24+
public static final UDFOperandMetadata DATETIME_OR_STRING_DATETIME_OR_STRING =
25+
UDFOperandMetadata.wrap(
26+
(CompositeOperandTypeChecker)
27+
OperandTypes.STRING_STRING.or(
28+
OperandTypes.family(SqlTypeFamily.DATETIME, SqlTypeFamily.DATETIME)));
29+
public static final UDFOperandMetadata TIME_OR_TIMESTAMP_OR_STRING =
30+
UDFOperandMetadata.wrap(
31+
(CompositeOperandTypeChecker)
32+
OperandTypes.STRING.or(OperandTypes.TIME).or(OperandTypes.TIMESTAMP));
33+
public static final UDFOperandMetadata DATE_OR_TIMESTAMP_OR_STRING =
34+
UDFOperandMetadata.wrap(
35+
(CompositeOperandTypeChecker) OperandTypes.DATE_OR_TIMESTAMP.or(OperandTypes.STRING));
36+
37+
public static final UDFOperandMetadata STRING =
38+
UDFOperandMetadata.wrap((FamilyOperandTypeChecker) OperandTypes.STRING);
39+
public static final UDFOperandMetadata INTEGER =
40+
UDFOperandMetadata.wrap((FamilyOperandTypeChecker) OperandTypes.INTEGER);
41+
public static final UDFOperandMetadata NUMERIC =
42+
UDFOperandMetadata.wrap((FamilyOperandTypeChecker) OperandTypes.NUMERIC);
43+
public static final UDFOperandMetadata INTEGER_INTEGER =
44+
UDFOperandMetadata.wrap((FamilyOperandTypeChecker) OperandTypes.INTEGER_INTEGER);
45+
public static final UDFOperandMetadata STRING_STRING =
46+
UDFOperandMetadata.wrap(OperandTypes.STRING_STRING);
47+
public static final UDFOperandMetadata NUMERIC_NUMERIC =
48+
UDFOperandMetadata.wrap((FamilyOperandTypeChecker) OperandTypes.NUMERIC_NUMERIC);
49+
public static final UDFOperandMetadata NUMERIC_NUMERIC_NUMERIC =
50+
UDFOperandMetadata.wrap(
51+
OperandTypes.family(SqlTypeFamily.NUMERIC, SqlTypeFamily.NUMERIC, SqlTypeFamily.NUMERIC));
52+
public static final UDFOperandMetadata OPTIONAL_INTEGER =
53+
UDFOperandMetadata.wrap(
54+
(CompositeOperandTypeChecker) OperandTypes.INTEGER.or(OperandTypes.family()));
55+
public static final UDFOperandMetadata EMPTY = UDFOperandMetadata.wrap(OperandTypes.family());
56+
}

‎core/src/main/java/org/opensearch/sql/calcite/utils/UserDefinedFunctionUtils.java‎

Lines changed: 16 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import java.util.Locale;
2424
import java.util.Objects;
2525
import java.util.TimeZone;
26+
import javax.annotation.Nullable;
2627
import org.apache.calcite.DataContext;
2728
import org.apache.calcite.adapter.enumerable.NotNullImplementor;
2829
import org.apache.calcite.adapter.enumerable.NullPolicy;
@@ -238,33 +239,6 @@ public static SqlTypeName convertRelDataTypeToSqlTypeName(RelDataType type) {
238239
return type.getSqlTypeName();
239240
}
240241

241-
public static ImplementorUDF adaptExprMethodToUDF(
242-
java.lang.reflect.Type type,
243-
String methodName,
244-
SqlReturnTypeInference returnTypeInference,
245-
NullPolicy nullPolicy,
246-
UDFOperandMetadata typeChecker) {
247-
NotNullImplementor implementor =
248-
(translator, call, translatedOperands) -> {
249-
List<Expression> operands =
250-
convertToExprValues(
251-
translatedOperands, call.getOperands().stream().map(RexNode::getType).toList());
252-
Expression exprResult = Expressions.call(type, methodName, operands);
253-
return Expressions.call(exprResult, "valueForCalcite");
254-
};
255-
return new ImplementorUDF(implementor, nullPolicy) {
256-
@Override
257-
public SqlReturnTypeInference getReturnTypeInference() {
258-
return returnTypeInference;
259-
}
260-
261-
@Override
262-
public UDFOperandMetadata getOperandMetadata() {
263-
return typeChecker;
264-
}
265-
};
266-
}
267-
268242
// TODO: pass the function properties directly to the UDF instead of string
269243
public static FunctionProperties restoreFunctionProperties(DataContext dataContext) {
270244
long currentTimeInNanos = DataContext.Variable.UTC_TIMESTAMP.get(dataContext);
@@ -325,13 +299,15 @@ public static List<Expression> convertToExprValues(
325299
* @param methodName the name of the method
326300
* @param returnTypeInference the return type inference of the UDF
327301
* @param nullPolicy the null policy of the UDF
302+
* @param operandMetadata type checker
328303
* @return an adapted ImplementorUDF with the expr method, which is a UserDefinedFunctionBuilder
329304
*/
330305
public static ImplementorUDF adaptExprMethodToUDF(
331306
java.lang.reflect.Type type,
332307
String methodName,
333308
SqlReturnTypeInference returnTypeInference,
334-
NullPolicy nullPolicy) {
309+
NullPolicy nullPolicy,
310+
@Nullable UDFOperandMetadata operandMetadata) {
335311
NotNullImplementor implementor =
336312
(translator, call, translatedOperands) -> {
337313
List<Expression> operands =
@@ -345,6 +321,11 @@ public static ImplementorUDF adaptExprMethodToUDF(
345321
public SqlReturnTypeInference getReturnTypeInference() {
346322
return returnTypeInference;
347323
}
324+
325+
@Override
326+
public UDFOperandMetadata getOperandMetadata() {
327+
return operandMetadata;
328+
}
348329
};
349330
}
350331

@@ -362,7 +343,8 @@ public static ImplementorUDF adaptExprMethodWithPropertiesToUDF(
362343
java.lang.reflect.Type type,
363344
String methodName,
364345
SqlReturnTypeInference returnTypeInference,
365-
NullPolicy nullPolicy) {
346+
NullPolicy nullPolicy,
347+
UDFOperandMetadata operandMetadata) {
366348
NotNullImplementor implementor =
367349
(translator, call, translatedOperands) -> {
368350
List<Expression> operands =
@@ -377,6 +359,11 @@ public static ImplementorUDF adaptExprMethodWithPropertiesToUDF(
377359
public SqlReturnTypeInference getReturnTypeInference() {
378360
return returnTypeInference;
379361
}
362+
363+
@Override
364+
public UDFOperandMetadata getOperandMetadata() {
365+
return operandMetadata;
366+
}
380367
};
381368
}
382369
}

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

Lines changed: 46 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import org.apache.calcite.sql.type.SqlTypeTransforms;
2727
import org.apache.calcite.sql.util.ReflectiveSqlOperatorTable;
2828
import org.apache.calcite.util.BuiltInMethod;
29+
import org.opensearch.sql.calcite.utils.PPLOperandTypes;
2930
import org.opensearch.sql.calcite.utils.PPLReturnTypes;
3031
import org.opensearch.sql.calcite.utils.UserDefinedFunctionUtils;
3132
import org.opensearch.sql.data.type.ExprCoreType;
@@ -76,7 +77,8 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable {
7677
DateTimeFunctions.class,
7778
"exprDate",
7879
PPLReturnTypes.DATE_FORCE_NULLABLE,
79-
NullPolicy.ARG0)
80+
NullPolicy.ARG0,
81+
PPLOperandTypes.DATE_OR_TIMESTAMP_OR_STRING)
8082
.toUDF("DATE");
8183
public static final SqlOperator YEARWEEK = new YearweekFunction().toUDF("YEARWEEK");
8284
public static final SqlOperator WEEKDAY = new WeekdayFunction().toUDF("WEEKDAY");
@@ -88,14 +90,16 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable {
8890
DateTimeFunctions.class,
8991
"exprAddTime",
9092
PPLReturnTypes.TIME_APPLY_RETURN_TYPE,
91-
NullPolicy.ANY)
93+
NullPolicy.ANY,
94+
PPLOperandTypes.DATETIME_DATETIME)
9295
.toUDF("ADDTIME");
9396
public static final SqlOperator SUBTIME =
9497
adaptExprMethodWithPropertiesToUDF(
9598
DateTimeFunctions.class,
9699
"exprSubTime",
97100
PPLReturnTypes.TIME_APPLY_RETURN_TYPE,
98-
NullPolicy.ANY)
101+
NullPolicy.ANY,
102+
PPLOperandTypes.DATETIME_DATETIME)
99103
.toUDF("SUBTIME");
100104
public static final SqlOperator ADDDATE = new AddSubDateFunction(true).toUDF("ADDDATE");
101105
public static final SqlOperator SUBDATE = new AddSubDateFunction(false).toUDF("SUBDATE");
@@ -117,7 +121,8 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable {
117121
DateTimeFunctions.class,
118122
"exprMinuteOfDay",
119123
PPLReturnTypes.INTEGER_FORCE_NULLABLE,
120-
NullPolicy.ARG0)
124+
NullPolicy.ARG0,
125+
PPLOperandTypes.TIME_OR_TIMESTAMP_OR_STRING)
121126
.toUDF("MINUTE_OF_DAY");
122127
public static final SqlOperator SECOND = new DatePartFunction(TimeUnit.SECOND).toUDF("SECOND");
123128
public static final SqlOperator MICROSECOND =
@@ -139,14 +144,22 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable {
139144
DateTimeFunctions.class,
140145
"exprConvertTZ",
141146
PPLReturnTypes.TIMESTAMP_FORCE_NULLABLE,
142-
NullPolicy.ANY)
147+
NullPolicy.ANY,
148+
UDFOperandMetadata.wrap(
149+
(CompositeOperandTypeChecker)
150+
OperandTypes.STRING_STRING_STRING.or(
151+
OperandTypes.family(
152+
SqlTypeFamily.TIMESTAMP,
153+
SqlTypeFamily.STRING,
154+
SqlTypeFamily.STRING))))
143155
.toUDF("CONVERT_TZ");
144156
public static final SqlOperator DATEDIFF =
145157
adaptExprMethodWithPropertiesToUDF(
146158
DateTimeFunctions.class,
147159
"exprDateDiff",
148160
ReturnTypes.BIGINT_FORCE_NULLABLE,
149-
NullPolicy.ANY)
161+
NullPolicy.ANY,
162+
PPLOperandTypes.DATETIME_DATETIME)
150163
.toUDF("DATEDIFF");
151164
public static final SqlOperator TIMESTAMPDIFF =
152165
new TimestampDiffFunction().toUDF("TIMESTAMPDIFF");
@@ -156,50 +169,57 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable {
156169
DateTimeFunctions.class,
157170
"exprFromDays",
158171
PPLReturnTypes.DATE_FORCE_NULLABLE,
159-
NullPolicy.ANY)
172+
NullPolicy.ANY,
173+
PPLOperandTypes.INTEGER)
160174
.toUDF("FROM_DAYS");
161175
public static final SqlOperator FROM_UNIXTIME = new FromUnixTimeFunction().toUDF("FROM_UNIXTIME");
162176
public static final SqlOperator GET_FORMAT =
163177
adaptExprMethodToUDF(
164178
DateTimeFunctions.class,
165179
"exprGetFormat",
166180
ReturnTypes.VARCHAR.andThen(SqlTypeTransforms.FORCE_NULLABLE),
167-
NullPolicy.ANY)
181+
NullPolicy.ANY,
182+
PPLOperandTypes.STRING_STRING)
168183
.toUDF("GET_FORMAT");
169184
public static final SqlOperator MAKEDATE =
170185
adaptExprMethodToUDF(
171186
DateTimeFunctions.class,
172187
"exprMakeDate",
173188
PPLReturnTypes.DATE_FORCE_NULLABLE,
174-
NullPolicy.ANY)
189+
NullPolicy.ANY,
190+
PPLOperandTypes.NUMERIC_NUMERIC)
175191
.toUDF("MAKEDATE");
176192
public static final SqlOperator MAKETIME =
177193
adaptExprMethodToUDF(
178194
DateTimeFunctions.class,
179195
"exprMakeTime",
180196
PPLReturnTypes.TIME_FORCE_NULLABLE,
181-
NullPolicy.ANY)
197+
NullPolicy.ANY,
198+
PPLOperandTypes.NUMERIC_NUMERIC_NUMERIC)
182199
.toUDF("MAKETIME");
183200
public static final SqlOperator PERIOD_DIFF =
184201
adaptExprMethodToUDF(
185202
DateTimeFunctions.class,
186203
"exprPeriodDiff",
187204
PPLReturnTypes.INTEGER_FORCE_NULLABLE,
188-
NullPolicy.ANY)
205+
NullPolicy.ANY,
206+
PPLOperandTypes.INTEGER_INTEGER)
189207
.toUDF("PERIOD_DIFF");
190208
public static final SqlOperator PERIOD_ADD =
191209
adaptExprMethodToUDF(
192210
DateTimeFunctions.class,
193211
"exprPeriodAdd",
194212
PPLReturnTypes.INTEGER_FORCE_NULLABLE,
195-
NullPolicy.ANY)
213+
NullPolicy.ANY,
214+
PPLOperandTypes.INTEGER_INTEGER)
196215
.toUDF("PERIOD_ADD");
197216
public static final SqlOperator STR_TO_DATE =
198217
adaptExprMethodWithPropertiesToUDF(
199218
DateTimeFunctions.class,
200219
"exprStrToDate",
201220
PPLReturnTypes.TIMESTAMP_FORCE_NULLABLE,
202-
NullPolicy.ANY)
221+
NullPolicy.ANY,
222+
PPLOperandTypes.STRING_STRING)
203223
.toUDF("STR_TO_DATE");
204224
public static final SqlOperator SYSDATE = new SysdateFunction().toUDF("SYSDATE");
205225
public static final SqlOperator SEC_TO_TIME = new SecToTimeFunction().toUDF("SEC_TO_TIME");
@@ -208,55 +228,58 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable {
208228
DateTimeFunctions.class,
209229
"exprTime",
210230
PPLReturnTypes.TIME_FORCE_NULLABLE,
211-
NullPolicy.ARG0)
231+
NullPolicy.ARG0,
232+
PPLOperandTypes.DATETIME_OR_STRING)
212233
.toUDF("TIME");
213234
public static final SqlOperator TIME_TO_SEC =
214235
adaptExprMethodToUDF(
215236
DateTimeFunctions.class,
216237
"exprTimeToSec",
217238
ReturnTypes.BIGINT_FORCE_NULLABLE,
218-
NullPolicy.ARG0)
239+
NullPolicy.ARG0,
240+
PPLOperandTypes.TIME_OR_TIMESTAMP_OR_STRING)
219241
.toUDF("TIME_TO_SEC");
220242
public static final SqlOperator TIMEDIFF =
221243
UserDefinedFunctionUtils.adaptExprMethodToUDF(
222244
DateTimeFunctions.class,
223245
"exprTimeDiff",
224246
PPLReturnTypes.TIME_FORCE_NULLABLE,
225247
NullPolicy.ANY,
226-
UDFOperandMetadata.wrap(
227-
(CompositeOperandTypeChecker)
228-
OperandTypes.family(SqlTypeFamily.TIME, SqlTypeFamily.TIME)
229-
.or(OperandTypes.family(SqlTypeFamily.STRING, SqlTypeFamily.STRING))))
248+
PPLOperandTypes.DATETIME_OR_STRING_DATETIME_OR_STRING)
230249
.toUDF("TIME_DIFF");
231250
public static final SqlOperator TIMESTAMPADD = new TimestampAddFunction().toUDF("TIMESTAMPADD");
232251
public static final SqlOperator TO_DAYS =
233252
adaptExprMethodToUDF(
234253
DateTimeFunctions.class,
235254
"exprToDays",
236255
ReturnTypes.BIGINT_FORCE_NULLABLE,
237-
NullPolicy.ARG0)
256+
NullPolicy.ARG0,
257+
PPLOperandTypes.DATE_OR_TIMESTAMP_OR_STRING)
238258
.toUDF("TO_DAYS");
239259
public static final SqlOperator DATETIME = new DatetimeFunction().toUDF("DATETIME");
240260
public static final SqlOperator UTC_DATE =
241261
adaptExprMethodWithPropertiesToUDF(
242262
DateTimeFunctions.class,
243263
"exprUtcDate",
244264
PPLReturnTypes.DATE_FORCE_NULLABLE,
245-
NullPolicy.NONE)
265+
NullPolicy.NONE,
266+
PPLOperandTypes.EMPTY)
246267
.toUDF("UTC_DATE");
247268
public static final SqlOperator UTC_TIME =
248269
adaptExprMethodWithPropertiesToUDF(
249270
DateTimeFunctions.class,
250271
"exprUtcTime",
251272
PPLReturnTypes.TIME_FORCE_NULLABLE,
252-
NullPolicy.NONE)
273+
NullPolicy.NONE,
274+
PPLOperandTypes.EMPTY)
253275
.toUDF("UTC_TIME");
254276
public static final SqlOperator UTC_TIMESTAMP =
255277
adaptExprMethodWithPropertiesToUDF(
256278
DateTimeFunctions.class,
257279
"exprUtcTimestamp",
258280
PPLReturnTypes.TIMESTAMP_FORCE_NULLABLE,
259-
NullPolicy.NONE)
281+
NullPolicy.NONE,
282+
PPLOperandTypes.EMPTY)
260283
.toUDF("UTC_TIMESTAMP");
261284
public static final SqlOperator WEEK = new WeekFunction().toUDF("WEEK");
262285

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@
55

66
package org.opensearch.sql.expression.function;
77

8-
import static java.lang.Math.E;
98
import static org.apache.calcite.sql.type.SqlTypeFamily.IGNORE;
109
import static org.opensearch.sql.calcite.utils.OpenSearchTypeFactory.getLegacyTypeName;
1110
import static org.opensearch.sql.expression.function.BuiltinFunctionName.*;

0 commit comments

Comments
 (0)