Skip to content

Commit fe0385c

Browse files
yuancuyuancu
committed
Display expected signatures when type checking fails & support type checking for non-operator registrated udf
Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>
1 parent e1e4deb commit fe0385c

15 files changed

Lines changed: 231 additions & 80 deletions

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

Lines changed: 23 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -14,26 +14,13 @@
1414
import org.opensearch.sql.expression.function.UDFOperandMetadata;
1515

1616
public class PPLOperandTypes {
17+
// This class is not meant to be instantiated.
1718
private PPLOperandTypes() {}
1819

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 =
20+
public static final UDFOperandMetadata NONE = UDFOperandMetadata.wrap(OperandTypes.family());
21+
public static final UDFOperandMetadata OPTIONAL_INTEGER =
3422
UDFOperandMetadata.wrap(
35-
(CompositeOperandTypeChecker) OperandTypes.DATE_OR_TIMESTAMP.or(OperandTypes.STRING));
36-
23+
(CompositeOperandTypeChecker) OperandTypes.INTEGER.or(OperandTypes.family()));
3724
public static final UDFOperandMetadata STRING =
3825
UDFOperandMetadata.wrap((FamilyOperandTypeChecker) OperandTypes.STRING);
3926
public static final UDFOperandMetadata INTEGER =
@@ -49,8 +36,24 @@ private PPLOperandTypes() {}
4936
public static final UDFOperandMetadata NUMERIC_NUMERIC_NUMERIC =
5037
UDFOperandMetadata.wrap(
5138
OperandTypes.family(SqlTypeFamily.NUMERIC, SqlTypeFamily.NUMERIC, SqlTypeFamily.NUMERIC));
52-
public static final UDFOperandMetadata OPTIONAL_INTEGER =
39+
40+
public static final UDFOperandMetadata DATETIME_OR_STRING =
5341
UDFOperandMetadata.wrap(
54-
(CompositeOperandTypeChecker) OperandTypes.INTEGER.or(OperandTypes.family()));
55-
public static final UDFOperandMetadata EMPTY = UDFOperandMetadata.wrap(OperandTypes.family());
42+
(CompositeOperandTypeChecker) OperandTypes.DATETIME.or(OperandTypes.STRING));
43+
public static final UDFOperandMetadata DATETIME_DATETIME =
44+
UDFOperandMetadata.wrap(OperandTypes.family(SqlTypeFamily.DATETIME, SqlTypeFamily.DATETIME));
45+
public static final UDFOperandMetadata DATETIME_OR_STRING_DATETIME_OR_STRING =
46+
UDFOperandMetadata.wrap(
47+
(CompositeOperandTypeChecker)
48+
OperandTypes.STRING_STRING
49+
.or(OperandTypes.family(SqlTypeFamily.DATETIME, SqlTypeFamily.DATETIME))
50+
.or(OperandTypes.family(SqlTypeFamily.DATETIME, SqlTypeFamily.STRING))
51+
.or(OperandTypes.family(SqlTypeFamily.STRING, SqlTypeFamily.DATETIME)));
52+
public static final UDFOperandMetadata TIME_OR_TIMESTAMP_OR_STRING =
53+
UDFOperandMetadata.wrap(
54+
(CompositeOperandTypeChecker)
55+
OperandTypes.STRING.or(OperandTypes.TIME).or(OperandTypes.TIMESTAMP));
56+
public static final UDFOperandMetadata DATE_OR_TIMESTAMP_OR_STRING =
57+
UDFOperandMetadata.wrap(
58+
(CompositeOperandTypeChecker) OperandTypes.DATE_OR_TIMESTAMP.or(OperandTypes.STRING));
5659
}

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

Lines changed: 8 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -91,15 +91,15 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable {
9191
"exprAddTime",
9292
PPLReturnTypes.TIME_APPLY_RETURN_TYPE,
9393
NullPolicy.ANY,
94-
PPLOperandTypes.DATETIME_DATETIME)
94+
PPLOperandTypes.DATETIME_OR_STRING_DATETIME_OR_STRING)
9595
.toUDF("ADDTIME");
9696
public static final SqlOperator SUBTIME =
9797
adaptExprMethodWithPropertiesToUDF(
9898
DateTimeFunctions.class,
9999
"exprSubTime",
100100
PPLReturnTypes.TIME_APPLY_RETURN_TYPE,
101101
NullPolicy.ANY,
102-
PPLOperandTypes.DATETIME_DATETIME)
102+
PPLOperandTypes.DATETIME_OR_STRING_DATETIME_OR_STRING)
103103
.toUDF("SUBTIME");
104104
public static final SqlOperator ADDDATE = new AddSubDateFunction(true).toUDF("ADDDATE");
105105
public static final SqlOperator SUBDATE = new AddSubDateFunction(false).toUDF("SUBDATE");
@@ -149,17 +149,15 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable {
149149
(CompositeOperandTypeChecker)
150150
OperandTypes.STRING_STRING_STRING.or(
151151
OperandTypes.family(
152-
SqlTypeFamily.TIMESTAMP,
153-
SqlTypeFamily.STRING,
154-
SqlTypeFamily.STRING))))
152+
SqlTypeFamily.DATETIME, SqlTypeFamily.STRING, SqlTypeFamily.STRING))))
155153
.toUDF("CONVERT_TZ");
156154
public static final SqlOperator DATEDIFF =
157155
adaptExprMethodWithPropertiesToUDF(
158156
DateTimeFunctions.class,
159157
"exprDateDiff",
160158
ReturnTypes.BIGINT_FORCE_NULLABLE,
161159
NullPolicy.ANY,
162-
PPLOperandTypes.DATETIME_DATETIME)
160+
PPLOperandTypes.DATETIME_OR_STRING_DATETIME_OR_STRING)
163161
.toUDF("DATEDIFF");
164162
public static final SqlOperator TIMESTAMPDIFF =
165163
new TimestampDiffFunction().toUDF("TIMESTAMPDIFF");
@@ -237,7 +235,7 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable {
237235
"exprTimeToSec",
238236
ReturnTypes.BIGINT_FORCE_NULLABLE,
239237
NullPolicy.ARG0,
240-
PPLOperandTypes.TIME_OR_TIMESTAMP_OR_STRING)
238+
PPLOperandTypes.DATETIME_OR_STRING)
241239
.toUDF("TIME_TO_SEC");
242240
public static final SqlOperator TIMEDIFF =
243241
UserDefinedFunctionUtils.adaptExprMethodToUDF(
@@ -263,23 +261,23 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable {
263261
"exprUtcDate",
264262
PPLReturnTypes.DATE_FORCE_NULLABLE,
265263
NullPolicy.NONE,
266-
PPLOperandTypes.EMPTY)
264+
PPLOperandTypes.NONE)
267265
.toUDF("UTC_DATE");
268266
public static final SqlOperator UTC_TIME =
269267
adaptExprMethodWithPropertiesToUDF(
270268
DateTimeFunctions.class,
271269
"exprUtcTime",
272270
PPLReturnTypes.TIME_FORCE_NULLABLE,
273271
NullPolicy.NONE,
274-
PPLOperandTypes.EMPTY)
272+
PPLOperandTypes.NONE)
275273
.toUDF("UTC_TIME");
276274
public static final SqlOperator UTC_TIMESTAMP =
277275
adaptExprMethodWithPropertiesToUDF(
278276
DateTimeFunctions.class,
279277
"exprUtcTimestamp",
280278
PPLReturnTypes.TIMESTAMP_FORCE_NULLABLE,
281279
NullPolicy.NONE,
282-
PPLOperandTypes.EMPTY)
280+
PPLOperandTypes.NONE)
283281
.toUDF("UTC_TIMESTAMP");
284282
public static final SqlOperator WEEK = new WeekFunction().toUDF("WEEK");
285283

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

Lines changed: 81 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,10 @@
1818
import java.util.HashMap;
1919
import java.util.List;
2020
import java.util.Map;
21+
import java.util.Objects;
2122
import java.util.Optional;
23+
import java.util.StringJoiner;
24+
import java.util.stream.Collectors;
2225
import org.apache.calcite.rel.type.RelDataType;
2326
import org.apache.calcite.rex.RexBuilder;
2427
import org.apache.calcite.rex.RexNode;
@@ -34,6 +37,8 @@
3437
import org.apache.calcite.sql.type.SqlTypeName;
3538
import org.apache.calcite.sql.validate.SqlUserDefinedFunction;
3639
import org.checkerframework.checker.nullness.qual.Nullable;
40+
import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory;
41+
import org.opensearch.sql.exception.ExpressionEvaluationException;
3742
import org.opensearch.sql.executor.QueryType;
3843

3944
public class PPLFuncImpTable {
@@ -143,8 +148,20 @@ public RexNode resolve(
143148
functionName, argTypes, e.getMessage()),
144149
e);
145150
}
146-
throw new IllegalArgumentException(
147-
String.format("Cannot resolve function: %s, arguments: %s", functionName, argTypes));
151+
StringJoiner joiner = new StringJoiner(",");
152+
for (var implement : implementList) {
153+
joiner.add(implement.getKey().typeChecker().getAllowedSignatures());
154+
}
155+
String actualSignature =
156+
"["
157+
+ argTypes.stream()
158+
.map(OpenSearchTypeFactory::convertRelDataTypeToExprType)
159+
.map(Objects::toString)
160+
.collect(Collectors.joining(","))
161+
+ "]";
162+
throw new ExpressionEvaluationException(
163+
String.format(
164+
"%s function expects {%s}, but got %s", functionName, joiner, actualSignature));
148165
}
149166

150167
@SuppressWarnings({"UnusedReturnValue", "SameParameterValue"})
@@ -364,57 +381,65 @@ void populate() {
364381
// Note, make the implementation an individual class if too complex.
365382
register(
366383
TRIM,
367-
((FunctionImp1)
384+
createFunctionImpWithTypeChecker(
368385
(builder, arg) ->
369386
builder.makeCall(
370387
SqlStdOperatorTable.TRIM,
371388
builder.makeFlag(Flag.BOTH),
372389
builder.makeLiteral(" "),
373-
arg)));
390+
arg),
391+
family(SqlTypeFamily.STRING)));
392+
374393
register(
375394
LTRIM,
376-
((FunctionImp1)
395+
createFunctionImpWithTypeChecker(
377396
(builder, arg) ->
378397
builder.makeCall(
379398
SqlStdOperatorTable.TRIM,
380399
builder.makeFlag(Flag.LEADING),
381400
builder.makeLiteral(" "),
382-
arg)));
401+
arg),
402+
family(SqlTypeFamily.STRING)));
383403
register(
384404
RTRIM,
385-
((FunctionImp1)
405+
createFunctionImpWithTypeChecker(
386406
(builder, arg) ->
387407
builder.makeCall(
388408
SqlStdOperatorTable.TRIM,
389409
builder.makeFlag(Flag.TRAILING),
390410
builder.makeLiteral(" "),
391-
arg)));
411+
arg),
412+
family(SqlTypeFamily.STRING)));
392413
register(
393414
STRCMP,
394-
((FunctionImp2)
395-
(builder, arg1, arg2) -> builder.makeCall(SqlLibraryOperators.STRCMP, arg2, arg1)));
415+
createFunctionImpWithTypeChecker(
416+
(builder, arg1, arg2) -> builder.makeCall(SqlLibraryOperators.STRCMP, arg2, arg1),
417+
family(SqlTypeFamily.STRING, SqlTypeFamily.STRING)));
396418
register(
397419
LOG,
398-
((FunctionImp2)
399-
(builder, arg1, arg2) -> builder.makeCall(SqlLibraryOperators.LOG, arg2, arg1)));
420+
createFunctionImpWithTypeChecker(
421+
(builder, arg1, arg2) -> builder.makeCall(SqlLibraryOperators.LOG, arg2, arg1),
422+
family(SqlTypeFamily.NUMERIC, SqlTypeFamily.NUMERIC)));
400423
register(
401424
LOG,
402-
((FunctionImp1)
425+
createFunctionImpWithTypeChecker(
403426
(builder, arg) ->
404427
builder.makeCall(
405428
SqlLibraryOperators.LOG,
406429
arg,
407-
builder.makeApproxLiteral(BigDecimal.valueOf(Math.E)))));
430+
builder.makeApproxLiteral(BigDecimal.valueOf(Math.E))),
431+
family(SqlTypeFamily.NUMERIC)));
408432
// SqlStdOperatorTable.SQRT is declared but not implemented. The call to SQRT in Calcite is
409433
// converted to POWER(x, 0.5).
410434
register(
411435
SQRT,
412-
((FunctionImp1)
436+
createFunctionImpWithTypeChecker(
413437
(builder, arg) ->
414438
builder.makeCall(
415439
SqlStdOperatorTable.POWER,
416440
arg,
417-
builder.makeApproxLiteral(BigDecimal.valueOf(0.5)))));
441+
builder.makeApproxLiteral(BigDecimal.valueOf(0.5))),
442+
family(SqlTypeFamily.NUMERIC)));
418443
register(
419444
TYPEOF,
420445
(FunctionImp1)
@@ -465,4 +490,44 @@ public PPLTypeChecker getTypeChecker() {
465490
return family(booleanFamily, booleanFamily);
466491
}
467492
}
493+
494+
@FunctionalInterface
495+
private interface RexUnaryResolver {
496+
RexNode apply(RexBuilder builder, RexNode node);
497+
}
498+
499+
@FunctionalInterface
500+
private interface RexBinaryResolver {
501+
RexNode apply(RexBuilder builder, RexNode arg1, RexNode arg2);
502+
}
503+
504+
private static FunctionImp createFunctionImpWithTypeChecker(
505+
RexUnaryResolver resolver, PPLTypeChecker typeChecker) {
506+
return new FunctionImp1() {
507+
@Override
508+
public RexNode resolve(RexBuilder builder, RexNode arg1) {
509+
return resolver.apply(builder, arg1);
510+
}
511+
512+
@Override
513+
public PPLTypeChecker getTypeChecker() {
514+
return typeChecker;
515+
}
516+
};
517+
}
518+
519+
private static FunctionImp createFunctionImpWithTypeChecker(
520+
RexBinaryResolver resolver, PPLTypeChecker typeChecker) {
521+
return new FunctionImp2() {
522+
@Override
523+
public RexNode resolve(RexBuilder builder, RexNode arg1, RexNode arg2) {
524+
return resolver.apply(builder, arg1, arg2);
525+
}
526+
527+
@Override
528+
public PPLTypeChecker getTypeChecker() {
529+
return typeChecker;
530+
}
531+
};
532+
}
468533
}

0 commit comments

Comments
 (0)