Skip to content

Commit 04881a0

Browse files
LantaoJinxinyual
LantaoJin
authored and
xinyual
committed
Support TYPEOF function with Calcite (#3446)
--------- Signed-off-by: Lantao Jin <ltjin@amazon.com> Signed-off-by: xinyual <xinyual@amazon.com>
1 parent 4e85f04 commit 04881a0

File tree

12 files changed

+174
-16
lines changed

12 files changed

+174
-16
lines changed

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

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,20 +19,23 @@
1919
import org.apache.calcite.tools.RelBuilder;
2020
import org.opensearch.sql.ast.expression.UnresolvedExpression;
2121
import org.opensearch.sql.calcite.utils.CalciteToolsHelper;
22+
import org.opensearch.sql.executor.QueryType;
2223

2324
public class CalcitePlanContext {
2425

2526
public FrameworkConfig config;
2627
public final Connection connection;
2728
public final RelBuilder relBuilder;
2829
public final ExtendedRexBuilder rexBuilder;
30+
public final QueryType queryType;
2931

3032
@Getter @Setter private boolean isResolvingJoinCondition = false;
3133
@Getter @Setter private boolean isResolvingExistsSubquery = false;
3234
private final Stack<RexCorrelVariable> correlVar = new Stack<>();
3335

34-
private CalcitePlanContext(FrameworkConfig config) {
36+
private CalcitePlanContext(FrameworkConfig config, QueryType queryType) {
3537
this.config = config;
38+
this.queryType = queryType;
3639
this.connection = CalciteToolsHelper.connect(config, TYPE_FACTORY);
3740
this.relBuilder = CalciteToolsHelper.create(config, TYPE_FACTORY, connection);
3841
this.rexBuilder = new ExtendedRexBuilder(relBuilder.getRexBuilder());
@@ -67,7 +70,7 @@ public Optional<RexCorrelVariable> peekCorrelVar() {
6770
}
6871
}
6972

70-
public static CalcitePlanContext create(FrameworkConfig config) {
71-
return new CalcitePlanContext(config);
73+
public static CalcitePlanContext create(FrameworkConfig config, QueryType queryType) {
74+
return new CalcitePlanContext(config, queryType);
7275
}
7376
}

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import static org.opensearch.sql.ast.expression.SpanUnit.NONE;
99
import static org.opensearch.sql.ast.expression.SpanUnit.UNKNOWN;
1010
import static org.opensearch.sql.calcite.utils.BuiltinFunctionUtils.translateArgument;
11+
import static org.opensearch.sql.calcite.utils.PlanUtils.intervalUnitToSpanUnit;
1112

1213
import java.math.BigDecimal;
1314
import java.util.List;
@@ -19,6 +20,7 @@
1920
import org.apache.calcite.rel.type.RelDataTypeFactory;
2021
import org.apache.calcite.rex.RexBuilder;
2122
import org.apache.calcite.rex.RexNode;
23+
import org.apache.calcite.sql.SqlIntervalQualifier;
2224
import org.apache.calcite.sql.SqlOperator;
2325
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
2426
import org.apache.calcite.sql.type.SqlTypeName;
@@ -35,6 +37,7 @@
3537
import org.opensearch.sql.ast.expression.EqualTo;
3638
import org.opensearch.sql.ast.expression.Function;
3739
import org.opensearch.sql.ast.expression.In;
40+
import org.opensearch.sql.ast.expression.Interval;
3841
import org.opensearch.sql.ast.expression.Let;
3942
import org.opensearch.sql.ast.expression.Literal;
4043
import org.opensearch.sql.ast.expression.Not;
@@ -116,6 +119,15 @@ public RexNode visitLiteral(Literal node, CalcitePlanContext context) {
116119
}
117120
}
118121

122+
@Override
123+
public RexNode visitInterval(Interval node, CalcitePlanContext context) {
124+
RexNode value = analyze(node.getValue(), context);
125+
SqlIntervalQualifier intervalQualifier =
126+
context.rexBuilder.createIntervalUntil(intervalUnitToSpanUnit(node.getUnit()));
127+
return context.rexBuilder.makeIntervalLiteral(
128+
new BigDecimal(value.toString()), intervalQualifier);
129+
}
130+
119131
@Override
120132
public RexNode visitAnd(And node, CalcitePlanContext context) {
121133
final RelDataType booleanType =

‎core/src/main/java/org/opensearch/sql/calcite/udf/systemUDF/TypeOfFunction.java‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.calcite.udf.systemUDF;
7+
8+
import org.opensearch.sql.calcite.udf.UserDefinedFunction;
9+
10+
public class TypeOfFunction implements UserDefinedFunction {
11+
12+
@Override
13+
public Object eval(Object... args) {
14+
return args[0];
15+
}
16+
}

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
package org.opensearch.sql.calcite.utils;
77

88
import static java.lang.Math.E;
9+
import static org.opensearch.sql.calcite.utils.OpenSearchTypeFactory.getLegacyTypeName;
910
import static org.opensearch.sql.calcite.utils.UserDefinedFunctionUtils.*;
1011

1112
import java.math.BigDecimal;
@@ -30,6 +31,7 @@
3031
import org.opensearch.sql.calcite.udf.mathUDF.EulerFunction;
3132
import org.opensearch.sql.calcite.udf.mathUDF.ModFunction;
3233
import org.opensearch.sql.calcite.udf.mathUDF.SqrtFunction;
34+
import org.opensearch.sql.calcite.udf.systemUDF.TypeOfFunction;
3335
import org.opensearch.sql.calcite.udf.textUDF.LocateFunction;
3436
import org.opensearch.sql.calcite.udf.textUDF.ReplaceFunction;
3537

@@ -202,6 +204,10 @@ static SqlOperator translate(String op) {
202204
return SqlStdOperatorTable.IS_NOT_NULL;
203205
case "IS NULL":
204206
return SqlStdOperatorTable.IS_NULL;
207+
case "TYPEOF":
208+
// TODO optimize this function to ImplementableFunction
209+
return TransferUserDefinedFunction(
210+
TypeOfFunction.class, "typeof", ReturnTypes.VARCHAR_2000_NULLABLE);
205211
// TODO Add more, ref RexImpTable
206212
default:
207213
throw new IllegalArgumentException("Unsupported operator: " + op);
@@ -273,6 +279,11 @@ static List<RexNode> translateArgument(
273279
throw new IllegalArgumentException("Log cannot accept argument list: " + argList);
274280
}
275281
return LogArgs;
282+
case "TYPEOF":
283+
return List.of(
284+
context.rexBuilder.makeLiteral(
285+
getLegacyTypeName(
286+
argList.getFirst().getType().getSqlTypeName(), context.queryType)));
276287
default:
277288
return argList;
278289
}

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

Lines changed: 50 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,14 @@
2121
import static org.opensearch.sql.data.type.ExprCoreType.TIME;
2222
import static org.opensearch.sql.data.type.ExprCoreType.TIMESTAMP;
2323
import static org.opensearch.sql.data.type.ExprCoreType.UNDEFINED;
24+
import static org.opensearch.sql.data.type.ExprCoreType.UNKNOWN;
25+
import static org.opensearch.sql.executor.QueryType.PPL;
26+
import static org.opensearch.sql.lang.PPLLangSpec.PPL_SPEC;
2427

2528
import java.lang.reflect.Type;
2629
import java.util.ArrayList;
2730
import java.util.List;
31+
import java.util.Locale;
2832
import java.util.Map;
2933
import org.apache.calcite.jdbc.JavaTypeFactoryImpl;
3034
import org.apache.calcite.rel.type.RelDataType;
@@ -35,6 +39,7 @@
3539
import org.opensearch.sql.data.type.ExprCoreType;
3640
import org.opensearch.sql.data.type.ExprType;
3741
import org.opensearch.sql.executor.OpenSearchTypeSystem;
42+
import org.opensearch.sql.executor.QueryType;
3843
import org.opensearch.sql.storage.Table;
3944

4045
/** This class is used to create RelDataType and map RelDataType to Java data type */
@@ -132,9 +137,13 @@ public static RelDataType convertExprTypeToRelDataType(ExprType fieldType, boole
132137
}
133138
}
134139

135-
/** Converts a Calcite data type to OpenSearch ExprCoreType. */
136-
public static ExprType convertRelDataTypeToExprType(RelDataType type) {
137-
switch (type.getSqlTypeName()) {
140+
/**
141+
* Usually, {@link this#createSqlType(SqlTypeName, boolean)} is used to create RelDataType, then
142+
* convert it to ExprType. This is a util to convert when you don't have typeFactory. So they are
143+
* all ExprCoreType.
144+
*/
145+
public static ExprType convertSqlTypeNameToExprType(SqlTypeName sqlTypeName) {
146+
switch (sqlTypeName) {
138147
case TINYINT:
139148
return BYTE;
140149
case SMALLINT:
@@ -143,6 +152,7 @@ public static ExprType convertRelDataTypeToExprType(RelDataType type) {
143152
return INTEGER;
144153
case BIGINT:
145154
return LONG;
155+
case FLOAT:
146156
case REAL:
147157
return FLOAT;
148158
case DOUBLE:
@@ -155,16 +165,25 @@ public static ExprType convertRelDataTypeToExprType(RelDataType type) {
155165
case DATE:
156166
return DATE;
157167
case TIME:
168+
case TIME_TZ:
169+
case TIME_WITH_LOCAL_TIME_ZONE:
158170
return TIME;
159171
case TIMESTAMP:
172+
case TIMESTAMP_WITH_LOCAL_TIME_ZONE:
173+
case TIMESTAMP_TZ:
160174
return TIMESTAMP;
161-
case GEOMETRY:
162-
return IP;
163175
case INTERVAL_YEAR:
176+
case INTERVAL_YEAR_MONTH:
164177
case INTERVAL_MONTH:
165178
case INTERVAL_DAY:
179+
case INTERVAL_DAY_HOUR:
180+
case INTERVAL_DAY_MINUTE:
181+
case INTERVAL_DAY_SECOND:
166182
case INTERVAL_HOUR:
183+
case INTERVAL_HOUR_MINUTE:
184+
case INTERVAL_HOUR_SECOND:
167185
case INTERVAL_MINUTE:
186+
case INTERVAL_MINUTE_SECOND:
168187
case INTERVAL_SECOND:
169188
return INTERVAL;
170189
case ARRAY:
@@ -174,9 +193,33 @@ public static ExprType convertRelDataTypeToExprType(RelDataType type) {
174193
case NULL:
175194
return UNDEFINED;
176195
default:
177-
throw new IllegalArgumentException(
178-
"Unsupported conversion for Relational Data type: " + type.getSqlTypeName());
196+
return UNKNOWN;
197+
}
198+
}
199+
200+
/** Get legacy name for a SqlTypeName. */
201+
public static String getLegacyTypeName(SqlTypeName sqlTypeName, QueryType queryType) {
202+
switch (sqlTypeName) {
203+
case BINARY:
204+
case VARBINARY:
205+
return "BINARY";
206+
case GEOMETRY:
207+
return "GEO_POINT";
208+
default:
209+
ExprType type = convertSqlTypeNameToExprType(sqlTypeName);
210+
return (queryType == PPL ? PPL_SPEC.typeName(type) : type.legacyTypeName())
211+
.toUpperCase(Locale.ROOT);
212+
}
213+
}
214+
215+
/** Converts a Calcite data type to OpenSearch ExprCoreType. */
216+
public static ExprType convertRelDataTypeToExprType(RelDataType type) {
217+
ExprType exprType = convertSqlTypeNameToExprType(type.getSqlTypeName());
218+
if (exprType == UNKNOWN) {
219+
throw new IllegalArgumentException(
220+
"Unsupported conversion for Relational Data type: " + type.getSqlTypeName());
179221
}
222+
return exprType;
180223
}
181224

182225
public static RelDataType convertSchema(Table table) {

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

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.calcite.utils;
7+
8+
import org.opensearch.sql.ast.expression.IntervalUnit;
9+
import org.opensearch.sql.ast.expression.SpanUnit;
10+
11+
public interface PlanUtils {
12+
13+
static SpanUnit intervalUnitToSpanUnit(IntervalUnit unit) {
14+
return switch (unit) {
15+
case MICROSECOND -> SpanUnit.MILLISECOND;
16+
case SECOND -> SpanUnit.SECOND;
17+
case MINUTE -> SpanUnit.MINUTE;
18+
case HOUR -> SpanUnit.HOUR;
19+
case DAY -> SpanUnit.DAY;
20+
case WEEK -> SpanUnit.WEEK;
21+
case MONTH -> SpanUnit.MONTH;
22+
case QUARTER -> SpanUnit.QUARTER;
23+
case YEAR -> SpanUnit.YEAR;
24+
case UNKNOWN -> SpanUnit.UNKNOWN;
25+
default -> throw new UnsupportedOperationException("Unsupported interval unit: " + unit);
26+
};
27+
}
28+
}

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

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import org.apache.calcite.sql.SqlKind;
2222
import org.apache.calcite.sql.SqlOperator;
2323
import org.apache.calcite.sql.parser.SqlParserPos;
24+
import org.apache.calcite.sql.type.InferTypes;
2425
import org.apache.calcite.sql.type.SqlReturnTypeInference;
2526
import org.apache.calcite.sql.type.SqlTypeName;
2627
import org.apache.calcite.sql.validate.SqlUserDefinedAggFunction;
@@ -63,7 +64,12 @@ public static SqlOperator TransferUserDefinedFunction(
6364
SqlIdentifier udfLtrimIdentifier =
6465
new SqlIdentifier(Collections.singletonList(functionName), null, SqlParserPos.ZERO, null);
6566
return new SqlUserDefinedFunction(
66-
udfLtrimIdentifier, SqlKind.OTHER_FUNCTION, returnType, null, null, udfFunction);
67+
udfLtrimIdentifier,
68+
SqlKind.OTHER_FUNCTION,
69+
returnType,
70+
InferTypes.ANY_NULLABLE,
71+
null,
72+
udfFunction);
6773
}
6874

6975
public static SqlReturnTypeInference getReturnTypeInferenceForArray() {

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ public void execute(
8787
(PrivilegedAction<Void>)
8888
() -> {
8989
final FrameworkConfig config = buildFrameworkConfig();
90-
final CalcitePlanContext context = CalcitePlanContext.create(config);
90+
final CalcitePlanContext context = CalcitePlanContext.create(config, queryType);
9191
executePlanByCalcite(analyze(plan, context), context, listener);
9292
return null;
9393
});

‎integ-test/src/test/java/org/opensearch/sql/calcite/remote/nonfallback/NonFallbackCalciteSystemFunctionIT.java‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,8 @@
55

66
package org.opensearch.sql.calcite.remote.nonfallback;
77

8-
import org.junit.Ignore;
98
import org.opensearch.sql.calcite.remote.fallback.CalciteSystemFunctionIT;
109

11-
@Ignore("https://github.com/opensearch-project/sql/issues/3418")
1210
public class NonFallbackCalciteSystemFunctionIT extends CalciteSystemFunctionIT {
1311
@Override
1412
public void init() throws Exception {

‎integ-test/src/test/java/org/opensearch/sql/calcite/standalone/CalcitePPLBuiltinFunctionIT.java‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
import java.io.IOException;
1919
import org.json.JSONObject;
20+
import org.junit.Ignore;
2021
import org.junit.jupiter.api.Test;
2122

2223
public class CalcitePPLBuiltinFunctionIT extends CalcitePPLIntegTestCase {
@@ -106,6 +107,43 @@ public void testAtanAndAtan2WithSort() {
106107
verifyDataRowsInOrder(actual, rows("Hello", 30, 4), rows("Jake", 70, 4));
107108
}
108109

110+
@Test
111+
public void testTypeOfBasic() {
112+
JSONObject result =
113+
executeQuery(
114+
String.format(
115+
"""
116+
source=%s
117+
| eval `typeof(1)` = typeof(1)
118+
| eval `typeof(true)` = typeof(true)
119+
| eval `typeof(2.0)` = typeof(2.0)
120+
| eval `typeof("2.0")` = typeof("2.0")
121+
| eval `typeof(name)` = typeof(name)
122+
| eval `typeof(country)` = typeof(country)
123+
| eval `typeof(age)` = typeof(age)
124+
| eval `typeof(interval)` = typeof(INTERVAL 2 DAY)
125+
| fields `typeof(1)`, `typeof(true)`, `typeof(2.0)`, `typeof("2.0")`, `typeof(name)`, `typeof(country)`, `typeof(age)`, `typeof(interval)`
126+
| head 1
127+
""",
128+
TEST_INDEX_STATE_COUNTRY));
129+
verifyDataRows(
130+
result, rows("INT", "BOOLEAN", "DOUBLE", "STRING", "STRING", "STRING", "INT", "INTERVAL"));
131+
}
132+
133+
@Ignore("https://github.com/opensearch-project/sql/issues/3400")
134+
public void testTypeOfDateTime() {
135+
JSONObject result =
136+
executeQuery(
137+
String.format(
138+
"""
139+
source=%s
140+
| eval `typeof(date)` = typeof(DATE('2008-04-14'))
141+
| eval `typeof(now())` = typeof(now())
142+
| fields `typeof(date)`, `typeof(now())`
143+
""",
144+
TEST_INDEX_STATE_COUNTRY));
145+
}
146+
109147
@Test
110148
public void testCeilingAndFloor() {
111149
JSONObject actual =

0 commit comments

Comments
 (0)