Skip to content

Commit baa5ee4

Browse files
xinyualxinyual
committed
merge from main
Signed-off-by: xinyual <xinyual@amazon.com>
2 parents 7dc2e88 + 270aa0d commit baa5ee4

56 files changed

Lines changed: 1817 additions & 155 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -245,7 +245,7 @@ public static ExprType convertSqlTypeNameToExprType(SqlTypeName sqlTypeName) {
245245
case ARRAY -> ARRAY;
246246
case MAP -> STRUCT;
247247
case GEOMETRY -> GEO_POINT;
248-
case NULL -> UNDEFINED;
248+
case NULL, ANY -> UNDEFINED;
249249
default -> UNKNOWN;
250250
};
251251
}

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

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
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+
/**
17+
* This class contains common operand types for PPL functions. They are created by either wrapping a
18+
* {@link FamilyOperandTypeChecker} or a {@link CompositeOperandTypeChecker} with a {@link
19+
* UDFOperandMetadata}.
20+
*/
21+
public class PPLOperandTypes {
22+
// This class is not meant to be instantiated.
23+
private PPLOperandTypes() {}
24+
25+
public static final UDFOperandMetadata NONE = UDFOperandMetadata.wrap(OperandTypes.family());
26+
public static final UDFOperandMetadata OPTIONAL_INTEGER =
27+
UDFOperandMetadata.wrap(
28+
(CompositeOperandTypeChecker) OperandTypes.INTEGER.or(OperandTypes.family()));
29+
public static final UDFOperandMetadata STRING =
30+
UDFOperandMetadata.wrap((FamilyOperandTypeChecker) OperandTypes.STRING);
31+
public static final UDFOperandMetadata INTEGER =
32+
UDFOperandMetadata.wrap((FamilyOperandTypeChecker) OperandTypes.INTEGER);
33+
public static final UDFOperandMetadata NUMERIC =
34+
UDFOperandMetadata.wrap((FamilyOperandTypeChecker) OperandTypes.NUMERIC);
35+
public static final UDFOperandMetadata INTEGER_INTEGER =
36+
UDFOperandMetadata.wrap((FamilyOperandTypeChecker) OperandTypes.INTEGER_INTEGER);
37+
public static final UDFOperandMetadata STRING_STRING =
38+
UDFOperandMetadata.wrap(OperandTypes.STRING_STRING);
39+
public static final UDFOperandMetadata NUMERIC_NUMERIC =
40+
UDFOperandMetadata.wrap((FamilyOperandTypeChecker) OperandTypes.NUMERIC_NUMERIC);
41+
public static final UDFOperandMetadata NUMERIC_NUMERIC_NUMERIC =
42+
UDFOperandMetadata.wrap(
43+
OperandTypes.family(SqlTypeFamily.NUMERIC, SqlTypeFamily.NUMERIC, SqlTypeFamily.NUMERIC));
44+
45+
public static final UDFOperandMetadata DATETIME_OR_STRING =
46+
UDFOperandMetadata.wrap(
47+
(CompositeOperandTypeChecker) OperandTypes.DATETIME.or(OperandTypes.STRING));
48+
public static final UDFOperandMetadata DATETIME_DATETIME =
49+
UDFOperandMetadata.wrap(OperandTypes.family(SqlTypeFamily.DATETIME, SqlTypeFamily.DATETIME));
50+
public static final UDFOperandMetadata DATETIME_OR_STRING_DATETIME_OR_STRING =
51+
UDFOperandMetadata.wrap(
52+
(CompositeOperandTypeChecker)
53+
OperandTypes.STRING_STRING
54+
.or(OperandTypes.family(SqlTypeFamily.DATETIME, SqlTypeFamily.DATETIME))
55+
.or(OperandTypes.family(SqlTypeFamily.DATETIME, SqlTypeFamily.STRING))
56+
.or(OperandTypes.family(SqlTypeFamily.STRING, SqlTypeFamily.DATETIME)));
57+
public static final UDFOperandMetadata TIME_OR_TIMESTAMP_OR_STRING =
58+
UDFOperandMetadata.wrap(
59+
(CompositeOperandTypeChecker)
60+
OperandTypes.STRING.or(OperandTypes.TIME).or(OperandTypes.TIMESTAMP));
61+
public static final UDFOperandMetadata DATE_OR_TIMESTAMP_OR_STRING =
62+
UDFOperandMetadata.wrap(
63+
(CompositeOperandTypeChecker) OperandTypes.DATE_OR_TIMESTAMP.or(OperandTypes.STRING));
64+
}

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

Lines changed: 32 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import java.util.Collections;
1616
import java.util.List;
1717
import java.util.TimeZone;
18+
import javax.annotation.Nullable;
1819
import org.apache.calcite.DataContext;
1920
import org.apache.calcite.adapter.enumerable.NotNullImplementor;
2021
import org.apache.calcite.adapter.enumerable.NullPolicy;
@@ -34,12 +35,14 @@
3435
import org.apache.calcite.sql.validate.SqlUserDefinedAggFunction;
3536
import org.apache.calcite.tools.RelBuilder;
3637
import org.apache.calcite.util.Optionality;
38+
import org.opensearch.sql.calcite.type.AbstractExprRelDataType;
3739
import org.opensearch.sql.calcite.udf.UserDefinedAggFunction;
3840
import org.opensearch.sql.data.model.ExprValueUtils;
3941
import org.opensearch.sql.data.type.ExprType;
4042
import org.opensearch.sql.executor.QueryType;
4143
import org.opensearch.sql.expression.function.FunctionProperties;
4244
import org.opensearch.sql.expression.function.ImplementorUDF;
45+
import org.opensearch.sql.expression.function.UDFOperandMetadata;
4346

4447
public class UserDefinedFunctionUtils {
4548

@@ -88,7 +91,20 @@ static SqlReturnTypeInference getReturnTypeInferenceForArray() {
8891
};
8992
}
9093

91-
// TODO: pass the function properties directly to the UDF instead of string
94+
public static SqlTypeName convertRelDataTypeToSqlTypeName(RelDataType type) {
95+
if (type instanceof AbstractExprRelDataType<?> exprType) {
96+
return switch (exprType.getUdt()) {
97+
case EXPR_DATE -> SqlTypeName.DATE;
98+
case EXPR_TIME -> SqlTypeName.TIME;
99+
case EXPR_TIMESTAMP -> SqlTypeName.TIMESTAMP;
100+
case EXPR_IP -> SqlTypeName.VARCHAR;
101+
case EXPR_BINARY -> SqlTypeName.VARBINARY;
102+
default -> type.getSqlTypeName();
103+
};
104+
}
105+
return type.getSqlTypeName();
106+
}
107+
92108
public static FunctionProperties restoreFunctionProperties(DataContext dataContext) {
93109
long currentTimeInNanos = DataContext.Variable.UTC_TIMESTAMP.get(dataContext);
94110
Instant instant =
@@ -148,13 +164,15 @@ public static List<Expression> convertToExprValues(
148164
* @param methodName the name of the method
149165
* @param returnTypeInference the return type inference of the UDF
150166
* @param nullPolicy the null policy of the UDF
167+
* @param operandMetadata type checker
151168
* @return an adapted ImplementorUDF with the expr method, which is a UserDefinedFunctionBuilder
152169
*/
153170
public static ImplementorUDF adaptExprMethodToUDF(
154171
java.lang.reflect.Type type,
155172
String methodName,
156173
SqlReturnTypeInference returnTypeInference,
157-
NullPolicy nullPolicy) {
174+
NullPolicy nullPolicy,
175+
@Nullable UDFOperandMetadata operandMetadata) {
158176
NotNullImplementor implementor =
159177
(translator, call, translatedOperands) -> {
160178
List<Expression> operands =
@@ -168,6 +186,11 @@ public static ImplementorUDF adaptExprMethodToUDF(
168186
public SqlReturnTypeInference getReturnTypeInference() {
169187
return returnTypeInference;
170188
}
189+
190+
@Override
191+
public UDFOperandMetadata getOperandMetadata() {
192+
return operandMetadata;
193+
}
171194
};
172195
}
173196

@@ -185,7 +208,8 @@ public static ImplementorUDF adaptExprMethodWithPropertiesToUDF(
185208
java.lang.reflect.Type type,
186209
String methodName,
187210
SqlReturnTypeInference returnTypeInference,
188-
NullPolicy nullPolicy) {
211+
NullPolicy nullPolicy,
212+
UDFOperandMetadata operandMetadata) {
189213
NotNullImplementor implementor =
190214
(translator, call, translatedOperands) -> {
191215
List<Expression> operands =
@@ -200,6 +224,11 @@ public static ImplementorUDF adaptExprMethodWithPropertiesToUDF(
200224
public SqlReturnTypeInference getReturnTypeInference() {
201225
return returnTypeInference;
202226
}
227+
228+
@Override
229+
public UDFOperandMetadata getOperandMetadata() {
230+
return operandMetadata;
231+
}
203232
};
204233
}
205234
}

‎core/src/main/java/org/opensearch/sql/exception/CalciteUnsupportedException.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,4 +10,8 @@ public class CalciteUnsupportedException extends QueryEngineException {
1010
public CalciteUnsupportedException(String message) {
1111
super(message);
1212
}
13+
14+
public CalciteUnsupportedException(String message, Throwable cause) {
15+
super(message, cause);
16+
}
1317
}

‎core/src/main/java/org/opensearch/sql/exception/ExpressionEvaluationException.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,12 @@
77

88
/** Exception for Expression Evaluation. */
99
public class ExpressionEvaluationException extends QueryEngineException {
10+
1011
public ExpressionEvaluationException(String message) {
1112
super(message);
1213
}
14+
15+
public ExpressionEvaluationException(String message, Throwable cause) {
16+
super(message, cause);
17+
}
1318
}

‎core/src/main/java/org/opensearch/sql/exception/NonFallbackCalciteException.java‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.exception;
7+
8+
/** Non-fallback to v2 exception for Calcite. */
9+
public class NonFallbackCalciteException extends QueryEngineException {
10+
11+
public NonFallbackCalciteException(String message) {
12+
super(message);
13+
}
14+
}

‎core/src/main/java/org/opensearch/sql/executor/QueryService.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@
4141
import org.opensearch.sql.common.setting.Settings.Key;
4242
import org.opensearch.sql.datasource.DataSourceService;
4343
import org.opensearch.sql.exception.CalciteUnsupportedException;
44+
import org.opensearch.sql.exception.NonFallbackCalciteException;
4445
import org.opensearch.sql.planner.PlanContext;
4546
import org.opensearch.sql.planner.Planner;
4647
import org.opensearch.sql.planner.logical.LogicalPaginate;
@@ -107,13 +108,13 @@ public void executeWithCalcite(
107108
return null;
108109
});
109110
} catch (Throwable t) {
110-
if (isCalciteFallbackAllowed()) {
111+
if (isCalciteFallbackAllowed() && !(t instanceof NonFallbackCalciteException)) {
111112
log.warn("Fallback to V2 query engine since got exception", t);
112113
executeWithLegacy(plan, queryType, listener, Optional.of(t));
113114
} else {
114115
if (t instanceof Error) {
115116
// Calcite may throw AssertError during query execution.
116-
listener.onFailure(new CalciteUnsupportedException(t.getMessage()));
117+
listener.onFailure(new CalciteUnsupportedException(t.getMessage(), t));
117118
} else {
118119
listener.onFailure((Exception) t);
119120
}

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

Lines changed: 7 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -5,26 +5,18 @@
55

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

8-
import static org.opensearch.sql.expression.function.PPLFuncImpTable.FunctionImp.ANY_TYPE;
9-
108
import java.util.List;
119
import org.apache.calcite.rel.type.RelDataType;
1210

1311
/** Function signature is composed by function name and arguments list. */
14-
public record CalciteFuncSignature(FunctionName functionName, List<RelDataType> funcArgTypes) {
12+
public record CalciteFuncSignature(FunctionName functionName, PPLTypeChecker typeChecker) {
1513

1614
public boolean match(FunctionName functionName, List<RelDataType> paramTypeList) {
17-
if (funcArgTypes == null) return true;
18-
if (!functionName.equals(this.functionName()) || paramTypeList.size() != funcArgTypes.size()) {
19-
return false;
20-
}
21-
for (int i = 0; i < paramTypeList.size(); i++) {
22-
RelDataType paramType = paramTypeList.get(i);
23-
RelDataType funcType = funcArgTypes.get(i);
24-
if (ANY_TYPE != funcType && paramType.getFamily() != funcType.getFamily()) {
25-
return false;
26-
}
27-
}
28-
return true;
15+
if (!functionName.equals(this.functionName())) return false;
16+
// For complex type checkers (e.g., OperandTypes.COMPARABLE_UNORDERED_COMPARABLE_UNORDERED),
17+
// the typeChecker will be null because only simple family-based type checks are currently
18+
// supported.
19+
if (typeChecker == null) return true;
20+
return typeChecker.checkOperandTypes(paramTypeList);
2921
}
3022
}

0 commit comments

Comments
 (0)