Skip to content

Commit 8555944

Browse files
yuancuyuancu
committed
WIP: adding operand types to UDFs
Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>
1 parent 5205a4e commit 8555944

13 files changed

Lines changed: 97 additions & 18 deletions

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

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -19,17 +19,17 @@
1919
import org.apache.calcite.linq4j.tree.Expression;
2020
import org.apache.calcite.rex.RexCall;
2121
import org.apache.calcite.sql.SqlOperator;
22-
import org.apache.calcite.sql.type.ReturnTypes;
23-
import org.apache.calcite.sql.type.SqlTypeTransforms;
2422
import org.apache.calcite.sql.type.CompositeOperandTypeChecker;
2523
import org.apache.calcite.sql.type.OperandTypes;
24+
import org.apache.calcite.sql.type.ReturnTypes;
2625
import org.apache.calcite.sql.type.SqlTypeFamily;
26+
import org.apache.calcite.sql.type.SqlTypeTransforms;
2727
import org.apache.calcite.sql.util.ReflectiveSqlOperatorTable;
2828
import org.apache.calcite.util.BuiltInMethod;
2929
import org.opensearch.sql.calcite.utils.PPLReturnTypes;
3030
import org.opensearch.sql.data.type.ExprCoreType;
3131
import org.opensearch.sql.expression.datetime.DateTimeFunctions;
32-
import org.opensearch.sql.expression.function.udf.CryptographicFunction;
32+
import org.opensearch.sql.expression.function.udf.Sha2Function;
3333
import org.opensearch.sql.expression.function.udf.datetime.AddSubDateFunction;
3434
import org.opensearch.sql.expression.function.udf.datetime.CurrentFunction;
3535
import org.opensearch.sql.expression.function.udf.datetime.DateAddSubFunction;
@@ -55,8 +55,6 @@
5555
import org.opensearch.sql.expression.function.udf.math.DivideFunction;
5656
import org.opensearch.sql.expression.function.udf.math.EulerFunction;
5757
import org.opensearch.sql.expression.function.udf.math.ModFunction;
58-
import org.opensearch.sql.calcite.utils.UserDefinedFunctionUtils;
59-
import org.opensearch.sql.expression.datetime.DateTimeFunctions;
6058

6159
/** Defines functions and operators that are implemented only by PPL */
6260
public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable {
@@ -68,7 +66,7 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable {
6866
public static final SqlOperator MOD = new ModFunction().toUDF("MOD");
6967
public static final SqlOperator CRC32 = new CRC32Function().toUDF("CRC32");
7068
public static final SqlOperator DIVIDE = new DivideFunction().toUDF("DIVIDE");
71-
public static final SqlOperator SHA2 = CryptographicFunction.sha2().toUDF("SHA2");
69+
public static final SqlOperator SHA2 = Sha2Function.sha2().toUDF("SHA2");
7270

7371
// Datetime function
7472
public static final SqlOperator TIMESTAMP = new TimestampFunction().toUDF("TIMESTAMP");
@@ -223,7 +221,11 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable {
223221
DateTimeFunctions.class,
224222
"exprTimeDiff",
225223
PPLReturnTypes.TIME_FORCE_NULLABLE,
226-
NullPolicy.ANY)
224+
NullPolicy.ANY,
225+
UDFOperandMetadata.wrap(
226+
(CompositeOperandTypeChecker)
227+
OperandTypes.family(SqlTypeFamily.TIME, SqlTypeFamily.TIME)
228+
.or(OperandTypes.STRING_STRING)))
227229
.toUDF("TIME_DIFF");
228230
public static final SqlOperator TIMESTAMPADD = new TimestampAddFunction().toUDF("TIMESTAMPADD");
229231
public static final SqlOperator TO_DAYS =

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
/*
2+
*
3+
* * Copyright OpenSearch Contributors
4+
* * SPDX-License-Identifier: Apache-2.0
5+
*
6+
*/
7+
8+
package org.opensearch.sql.expression.function;
9+
10+
public class PPLOperandTypes {
11+
// Avoid instantiation
12+
private PPLOperandTypes() {}
13+
14+
}

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,12 @@
1818
import org.apache.calcite.sql.type.SqlOperandMetadata;
1919
import org.apache.calcite.sql.type.SqlOperandTypeChecker;
2020
import org.apache.calcite.sql.type.SqlTypeFamily;
21+
import org.apache.calcite.sql.validate.SqlUserDefinedFunction;
2122

2223
/**
2324
* This class is created for the compatibility with SqlUserDefinedFunction constructors when
2425
* creating UDFs, so that a type checker can be passed to the constructor of {@link
25-
* org.apache.calcite.sql.validate.SqlUserDefinedFunction} as an OperandMetadata.
26+
* SqlUserDefinedFunction} as an OperandMetadata.
2627
*/
2728
public interface UDFOperandMetadata extends SqlOperandMetadata, ImplicitCastOperandTypeChecker {
2829
SqlOperandTypeChecker getInnerTypeChecker();

core/src/main/java/org/opensearch/sql/expression/function/udf/CryptographicFunction.java renamed to core/src/main/java/org/opensearch/sql/expression/function/udf/Sha2Function.java

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,28 +12,35 @@
1212
import org.apache.calcite.linq4j.tree.Expression;
1313
import org.apache.calcite.linq4j.tree.Expressions;
1414
import org.apache.calcite.rex.RexCall;
15+
import org.apache.calcite.sql.type.OperandTypes;
1516
import org.apache.calcite.sql.type.ReturnTypes;
1617
import org.apache.calcite.sql.type.SqlReturnTypeInference;
1718
import org.apache.calcite.sql.type.SqlTypeTransforms;
1819
import org.apache.commons.codec.binary.Hex;
1920
import org.apache.commons.codec.digest.DigestUtils;
2021
import org.apache.commons.codec.digest.MessageDigestAlgorithms;
2122
import org.opensearch.sql.expression.function.ImplementorUDF;
23+
import org.opensearch.sql.expression.function.UDFOperandMetadata;
2224

23-
public class CryptographicFunction extends ImplementorUDF {
24-
private CryptographicFunction(NotNullImplementor implementor, NullPolicy nullPolicy) {
25+
public class Sha2Function extends ImplementorUDF {
26+
private Sha2Function(NotNullImplementor implementor, NullPolicy nullPolicy) {
2527
super(implementor, nullPolicy);
2628
}
2729

28-
public static CryptographicFunction sha2() {
29-
return new CryptographicFunction(new Sha2Implementor(), NullPolicy.ANY);
30+
public static Sha2Function sha2() {
31+
return new Sha2Function(new Sha2Implementor(), NullPolicy.ANY);
3032
}
3133

3234
@Override
3335
public SqlReturnTypeInference getReturnTypeInference() {
3436
return ReturnTypes.VARCHAR.andThen(SqlTypeTransforms.FORCE_NULLABLE);
3537
}
3638

39+
@Override
40+
public UDFOperandMetadata getOperandMetadata() {
41+
return UDFOperandMetadata.wrap(OperandTypes.STRING_INTEGER);
42+
}
43+
3744
public static class Sha2Implementor implements NotNullImplementor {
3845
@Override
3946
public Expression implement(

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@
2121
import org.apache.calcite.linq4j.tree.Expressions;
2222
import org.apache.calcite.rel.type.RelDataType;
2323
import org.apache.calcite.rex.RexCall;
24+
import org.apache.calcite.sql.type.CompositeOperandTypeChecker;
25+
import org.apache.calcite.sql.type.OperandTypes;
2426
import org.apache.calcite.sql.type.SqlReturnTypeInference;
2527
import org.apache.calcite.sql.type.SqlTypeFamily;
2628
import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory;
@@ -33,6 +35,7 @@
3335
import org.opensearch.sql.data.type.ExprCoreType;
3436
import org.opensearch.sql.expression.function.FunctionProperties;
3537
import org.opensearch.sql.expression.function.ImplementorUDF;
38+
import org.opensearch.sql.expression.function.UDFOperandMetadata;
3639

3740
/**
3841
* Adds or sub an interval or a number of days to a date or time.
@@ -69,6 +72,11 @@ public SqlReturnTypeInference getReturnTypeInference() {
6972
};
7073
}
7174

75+
@Override
76+
public UDFOperandMetadata getOperandMetadata() {
77+
return UDFOperandMetadata.wrap((CompositeOperandTypeChecker) OperandTypes.family(SqlTypeFamily.DATETIME, SqlTypeFamily.DATETIME).or(OperandTypes.NUMERIC_NUMERIC));
78+
}
79+
7280
@RequiredArgsConstructor
7381
public static class AddSubDateImplementor implements NotNullImplementor {
7482
private final boolean isAdd;

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

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import org.opensearch.sql.expression.datetime.DateTimeFunctions;
2525
import org.opensearch.sql.expression.function.FunctionProperties;
2626
import org.opensearch.sql.expression.function.ImplementorUDF;
27+
import org.opensearch.sql.expression.function.UDFOperandMetadata;
2728

2829
/**
2930
* Implementation of the now-like functions:
@@ -55,7 +56,12 @@ public SqlReturnTypeInference getReturnTypeInference() {
5556
};
5657
}
5758

58-
@RequiredArgsConstructor
59+
@Override
60+
public UDFOperandMetadata getOperandMetadata() {
61+
return null;
62+
}
63+
64+
@RequiredArgsConstructor
5965
public static class CurrentFunctionImplementor implements NotNullImplementor {
6066
private final ExprType returnType;
6167

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,13 +15,16 @@
1515
import org.apache.calcite.linq4j.tree.Expressions;
1616
import org.apache.calcite.rel.type.RelDataType;
1717
import org.apache.calcite.rex.RexCall;
18+
import org.apache.calcite.sql.type.FamilyOperandTypeChecker;
19+
import org.apache.calcite.sql.type.OperandTypes;
1820
import org.apache.calcite.sql.type.SqlReturnTypeInference;
1921
import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory;
2022
import org.opensearch.sql.calcite.utils.PPLReturnTypes;
2123
import org.opensearch.sql.calcite.utils.UserDefinedFunctionUtils;
2224
import org.opensearch.sql.calcite.utils.datetime.DateTimeConversionUtils;
2325
import org.opensearch.sql.data.model.ExprValueUtils;
2426
import org.opensearch.sql.expression.function.ImplementorUDF;
27+
import org.opensearch.sql.expression.function.UDFOperandMetadata;
2528

2629
/**
2730
* Add or sub a specified interval to a date or time. If the first argument is TIME, today's date is
@@ -44,6 +47,11 @@ public SqlReturnTypeInference getReturnTypeInference() {
4447
return PPLReturnTypes.TIMESTAMP_FORCE_NULLABLE;
4548
}
4649

50+
@Override
51+
public UDFOperandMetadata getOperandMetadata() {
52+
return UDFOperandMetadata.wrap((FamilyOperandTypeChecker) OperandTypes.DATETIME_INTERVAL);
53+
}
54+
4755
@RequiredArgsConstructor
4856
public static class DateAddSubImplementor implements NotNullImplementor {
4957
private final boolean isAdd;

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
import org.opensearch.sql.expression.datetime.DateTimeFunctions;
3030
import org.opensearch.sql.expression.function.FunctionProperties;
3131
import org.opensearch.sql.expression.function.ImplementorUDF;
32+
import org.opensearch.sql.expression.function.UDFOperandMetadata;
3233

3334
/**
3435
* Implementations of date-part-related functions:
@@ -68,6 +69,11 @@ public SqlReturnTypeInference getReturnTypeInference() {
6869
return PPLReturnTypes.INTEGER_FORCE_NULLABLE;
6970
}
7071

72+
@Override
73+
public UDFOperandMetadata getOperandMetadata() {
74+
return null;
75+
}
76+
7177
@RequiredArgsConstructor
7278
public static class DatePartImplementor implements NotNullImplementor {
7379
private final TimeUnit timeUnit;

‎core/src/main/java/org/opensearch/sql/expression/function/udf/math/CRC32Function.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,12 @@
1313
import org.apache.calcite.linq4j.tree.Expression;
1414
import org.apache.calcite.linq4j.tree.Expressions;
1515
import org.apache.calcite.rex.RexCall;
16+
import org.apache.calcite.sql.type.FamilyOperandTypeChecker;
17+
import org.apache.calcite.sql.type.OperandTypes;
1618
import org.apache.calcite.sql.type.ReturnTypes;
1719
import org.apache.calcite.sql.type.SqlReturnTypeInference;
1820
import org.opensearch.sql.expression.function.ImplementorUDF;
21+
import org.opensearch.sql.expression.function.UDFOperandMetadata;
1922

2023
/**
2124
* <code>CRC32(value)</code> returns a 32-bit cyclic redundancy check (CRC) checksum
@@ -36,6 +39,11 @@ public SqlReturnTypeInference getReturnTypeInference() {
3639
return ReturnTypes.BIGINT_FORCE_NULLABLE;
3740
}
3841

42+
@Override
43+
public UDFOperandMetadata getOperandMetadata() {
44+
return UDFOperandMetadata.wrap((FamilyOperandTypeChecker) OperandTypes.STRING);
45+
}
46+
3947
public static class Crc32Implementor implements NotNullImplementor {
4048
@Override
4149
public Expression implement(

‎core/src/main/java/org/opensearch/sql/expression/function/udf/math/ConvFunction.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@
1414
import org.apache.calcite.linq4j.tree.Expressions;
1515
import org.apache.calcite.rel.type.RelDataType;
1616
import org.apache.calcite.rex.RexCall;
17+
import org.apache.calcite.sql.type.FamilyOperandTypeChecker;
18+
import org.apache.calcite.sql.type.OperandTypes;
1719
import org.apache.calcite.sql.type.SqlReturnTypeInference;
1820
import org.apache.calcite.sql.type.SqlTypeFamily;
1921
import org.opensearch.sql.calcite.utils.PPLReturnTypes;
@@ -38,7 +40,7 @@ public SqlReturnTypeInference getReturnTypeInference() {
3840

3941
@Override
4042
public UDFOperandMetadata getOperandMetadata() {
41-
return null;
43+
return UDFOperandMetadata.wrap((FamilyOperandTypeChecker) OperandTypes.STRING_NUMERIC_NUMERIC);
4244
}
4345

4446
public static class ConvImplementor implements NotNullImplementor {

0 commit comments

Comments
 (0)