Skip to content

Commit c5ad14d

Browse files
yuancuyuancu
committed
Support casting to IP with Calcite (opensearch-project#3919)
* Support casting ip Signed-off-by: Yuanchun Shen <yuanchu@amazon.com> * Refactor type checking for UDT (specifically, IP) Signed-off-by: Yuanchun Shen <yuanchu@amazon.com> * Remove IPCastFunctionTest Signed-off-by: Yuanchun Shen <yuanchu@amazon.com> * Remove unused IPUtil function Signed-off-by: Yuanchun Shen <yuanchu@amazon.com> * Fix error message of comparing with different types Signed-off-by: Yuanchun Shen <yuanchu@amazon.com> * Correct casting ip to string & update conversion doc Signed-off-by: Yuanchun Shen <yuanchu@amazon.com> * Change string representation form IP [ip] to [ip] to accomadate casting ip to string Signed-off-by: Yuanchun Shen <yuanchu@amazon.com> --------- Signed-off-by: Yuanchun Shen <yuanchu@amazon.com> (cherry picked from commit 411395d)
1 parent 0878a3b commit c5ad14d

15 files changed

Lines changed: 711 additions & 454 deletions

File tree

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

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import java.util.List;
1111
import java.util.stream.Collectors;
1212

13+
import java.util.Locale;
1314
import org.apache.calcite.avatica.util.TimeUnit;
1415
import org.apache.calcite.rel.type.RelDataType;
1516
import org.apache.calcite.rex.RexBuilder;
@@ -21,6 +22,12 @@
2122
import org.apache.calcite.sql.type.SqlTypeName;
2223
import org.apache.calcite.sql.type.SqlTypeUtil;
2324
import org.opensearch.sql.ast.expression.SpanUnit;
25+
import org.opensearch.sql.calcite.type.AbstractExprRelDataType;
26+
import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory;
27+
import org.opensearch.sql.data.type.ExprCoreType;
28+
import org.opensearch.sql.exception.ExpressionEvaluationException;
29+
import org.opensearch.sql.exception.SemanticCheckException;
30+
import org.opensearch.sql.expression.function.PPLBuiltinOperators;
2431

2532
public class ExtendedRexBuilder extends RexBuilder {
2633

@@ -119,6 +126,34 @@ public RexNode makeCast(
119126
// SqlStdOperatorTable.NOT_EQUALS,
120127
// ImmutableList.of(exp, makeZeroLiteral(exp.getType())));
121128
}
129+
} else if (type instanceof ExprSqlType exprSqlType
130+
&& Set.of(EXPR_DATE, EXPR_TIME, EXPR_TIMESTAMP).contains(exprSqlType.getUdt())) {
131+
switch (exprSqlType.getUdt()) {
132+
case EXPR_DATE:
133+
return makeCall(type, PPLBuiltinOperators.DATE, List.of(exp));
134+
case EXPR_TIME:
135+
return makeCall(type, PPLBuiltinOperators.TIME, List.of(exp));
136+
case EXPR_TIMESTAMP:
137+
return makeCall(type, PPLBuiltinOperators.TIMESTAMP, List.of(exp));
138+
case EXPR_IP:
139+
if (argExprType == ExprCoreType.IP) {
140+
return exp;
141+
} else if (argExprType == ExprCoreType.STRING) {
142+
return makeCall(type, PPLBuiltinOperators.IP, List.of(exp));
143+
}
144+
// Throwing error inside implementation will be suppressed by Calcite, thus
145+
// throwing 500 error. Therefore, we throw error here to ensure the error
146+
// information is displayed properly.
147+
throw new ExpressionEvaluationException(
148+
String.format(
149+
Locale.ROOT,
150+
"Cannot convert %s to IP, only STRING and IP types are supported",
151+
argExprType));
152+
break;
153+
throw new SemanticCheckException(
154+
String.format(Locale.ROOT, "Cannot cast from %s to %s", argExprType, udt.name()));
155+
}
156+
}
122157
}
123158
return super.makeCast(pos, type, exp, matchNullability, safe, format);
124159
}

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

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -313,7 +313,7 @@ public static String getLegacyTypeName(RelDataType relDataType, QueryType queryT
313313

314314
/** Converts a Calcite data type to OpenSearch ExprCoreType. */
315315
public static ExprType convertRelDataTypeToExprType(RelDataType type) {
316-
if (type instanceof AbstractExprRelDataType<?>) {
316+
if (isUserDefinedType(type)) {
317317
AbstractExprRelDataType<?> udt = (AbstractExprRelDataType<?>) type;
318318
return udt.getExprType();
319319
}
@@ -346,4 +346,14 @@ public Type getJavaClass(RelDataType type) {
346346
}
347347
return super.getJavaClass(type);
348348
}
349+
350+
/**
351+
* Whether a given RelDataType is a user-defined type (UDT)
352+
*
353+
* @param type the RelDataType to check
354+
* @return true if the type is a user-defined type, false otherwise
355+
*/
356+
public static boolean isUserDefinedType(RelDataType type) {
357+
return type instanceof AbstractExprRelDataType<?>;
358+
}
349359
}

‎core/src/main/java/org/opensearch/sql/data/model/ExprIpValue.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,8 @@ public boolean equal(ExprValue other) {
4545

4646
@Override
4747
public String toString() {
48-
return String.format("IP %s", value());
48+
// used for casting to string
49+
return value();
4950
}
5051

5152
@Override

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,7 @@
7575
import org.opensearch.sql.expression.function.udf.datetime.YearweekFunction;
7676
import org.opensearch.sql.expression.function.udf.ip.CidrMatchFunction;
7777
import org.opensearch.sql.expression.function.udf.ip.CompareIpFunction;
78+
import org.opensearch.sql.expression.function.udf.ip.IPFunction;
7879
import org.opensearch.sql.expression.function.udf.math.CRC32Function;
7980
import org.opensearch.sql.expression.function.udf.math.ConvFunction;
8081
import org.opensearch.sql.expression.function.udf.math.DivideFunction;
@@ -281,6 +282,9 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable {
281282
NullPolicy.ARG0,
282283
PPLOperandTypes.DATETIME_OR_STRING)
283284
.toUDF("TIME");
285+
286+
// IP cast function
287+
public static final SqlOperator IP = new IPFunction().toUDF("IP");
284288
public static final SqlOperator TIME_TO_SEC =
285289
adaptExprMethodToUDF(
286290
DateTimeFunctions.class,

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

Lines changed: 25 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -536,18 +536,9 @@ void registerOperator(BuiltinFunctionName functionName, SqlOperator... operators
536536
// Comparison operators like EQUAL, GREATER_THAN, LESS_THAN, etc.
537537
// SameOperandTypeCheckers like COALESCE, IFNULL, etc.
538538
register(functionName, wrapWithComparableTypeChecker(operator, comparableTypeChecker));
539-
} else if (typeChecker instanceof UDFOperandMetadata.IPOperandMetadata) {
540-
register(
541-
functionName,
542-
createFunctionImpWithTypeChecker(
543-
(builder, arg1, arg2) -> builder.makeCall(operator, arg1, arg2),
544-
new PPLTypeChecker.PPLIPCompareTypeChecker()));
545-
} else if (typeChecker instanceof UDFOperandMetadata.CidrOperandMetadata) {
546-
register(
547-
functionName,
548-
createFunctionImpWithTypeChecker(
549-
(builder, arg1, arg2) -> builder.makeCall(operator, arg1, arg2),
550-
new PPLTypeChecker.PPLCidrTypeChecker()));
539+
} else if (typeChecker
540+
instanceof UDFOperandMetadata.UDTOperandMetadata udtOperandMetadata) {
541+
register(functionName, wrapWithUdtTypeChecker(operator, udtOperandMetadata));
551542
} else {
552543
logger.info(
553544
"Cannot create type checker for function: {}. Will skip its type checking",
@@ -566,6 +557,13 @@ private static SqlOperandTypeChecker extractTypeCheckerFromUDF(
566557
return (udfOperandMetadata == null) ? null : udfOperandMetadata.getInnerTypeChecker();
567558
}
568559

560+
// Such wrapWith*TypeChecker methods are useful in that we don't have to create explicit
561+
// overrides of resolve function for different number of operands.
562+
// I.e. we don't have to explicitly call
563+
// (FuncImp1) (builder, arg1) -> builder.makeCall(operator, arg1);
564+
// (FuncImp2) (builder, arg1, arg2) -> builder.makeCall(operator, arg1, arg2);
565+
// etc.
566+
569567
/**
570568
* Wrap a SqlOperator into a FunctionImp with a composite type checker.
571569
*
@@ -632,6 +630,21 @@ public PPLTypeChecker getTypeChecker() {
632630
};
633631
}
634632

633+
private static FunctionImp wrapWithUdtTypeChecker(
634+
SqlOperator operator, UDFOperandMetadata.UDTOperandMetadata udtOperandMetadata) {
635+
return new FunctionImp() {
636+
@Override
637+
public RexNode resolve(RexBuilder builder, RexNode... args) {
638+
return builder.makeCall(operator, args);
639+
}
640+
641+
@Override
642+
public PPLTypeChecker getTypeChecker() {
643+
return PPLTypeChecker.wrapUDT(udtOperandMetadata.allowedParamTypes());
644+
}
645+
};
646+
}
647+
635648
private static FunctionImp createFunctionImpWithTypeChecker(
636649
BiFunction<RexBuilder, RexNode, RexNode> resolver, PPLTypeChecker typeChecker) {
637650
return new FunctionImp1() {

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

Lines changed: 40 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@
2323
import org.apache.calcite.sql.type.SqlTypeFamily;
2424
import org.apache.calcite.sql.type.SqlTypeName;
2525
import org.apache.calcite.sql.type.SqlTypeUtil;
26-
import org.opensearch.sql.calcite.type.ExprIPType;
2726
import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory;
2827
import org.opensearch.sql.calcite.utils.UserDefinedFunctionUtils;
2928
import org.opensearch.sql.data.type.ExprCoreType;
@@ -270,53 +269,6 @@ public String getAllowedSignatures() {
270269
}
271270
}
272271

273-
class PPLIPCompareTypeChecker implements PPLTypeChecker {
274-
@Override
275-
public boolean checkOperandTypes(List<RelDataType> types) {
276-
if (types.size() != 2) {
277-
return false;
278-
}
279-
RelDataType type1 = types.get(0);
280-
RelDataType type2 = types.get(1);
281-
return areIpAndStringTypes(type1, type2)
282-
|| areIpAndStringTypes(type2, type1)
283-
|| (type1 instanceof ExprIPType && type2 instanceof ExprIPType);
284-
}
285-
286-
@Override
287-
public String getAllowedSignatures() {
288-
// Will be merged with the allowed signatures of comparable type checker,
289-
// shown as [COMPARABLE_TYPE,COMPARABLE_TYPE]
290-
return "";
291-
}
292-
293-
private static boolean areIpAndStringTypes(RelDataType typeIp, RelDataType typeString) {
294-
return typeIp instanceof ExprIPType && typeString.getFamily() == SqlTypeFamily.CHARACTER;
295-
}
296-
}
297-
298-
class PPLCidrTypeChecker implements PPLTypeChecker {
299-
@Override
300-
public boolean checkOperandTypes(List<RelDataType> types) {
301-
if (types.size() != 2) {
302-
return false;
303-
}
304-
RelDataType type1 = types.get(0);
305-
RelDataType type2 = types.get(1);
306-
307-
// accept (STRING, STRING) or (IP, STRING)
308-
if (type2.getFamily() != SqlTypeFamily.CHARACTER) {
309-
return false;
310-
}
311-
return type1 instanceof ExprIPType || type1.getFamily() == SqlTypeFamily.CHARACTER;
312-
}
313-
314-
@Override
315-
public String getAllowedSignatures() {
316-
return "[STRING,STRING],[IP,STRING]";
317-
}
318-
}
319-
320272
/**
321273
* Creates a {@link PPLFamilyTypeChecker} with a fixed operand count, validating that each operand
322274
* belongs to its corresponding {@link SqlTypeFamily}.
@@ -392,6 +344,42 @@ static PPLComparableTypeChecker wrapComparable(SameOperandTypeChecker typeChecke
392344
return new PPLComparableTypeChecker(typeChecker);
393345
}
394346

347+
/**
348+
* Create a {@link PPLTypeChecker} from a list of allowed signatures consisted of {@link
349+
* ExprType}. This is useful to validate arguments against user-defined types (UDT) that does not
350+
* match any Calcite {@link SqlTypeFamily}.
351+
*
352+
* @param allowedSignatures a list of allowed signatures, where each signature is a list of {@link
353+
* ExprType} representing the expected types of the function arguments.
354+
* @return a {@link PPLTypeChecker} that checks if the operand types match any of the allowed
355+
* signatures
356+
*/
357+
static PPLTypeChecker wrapUDT(List<List<ExprType>> allowedSignatures) {
358+
return new PPLTypeChecker() {
359+
@Override
360+
public boolean checkOperandTypes(List<RelDataType> types) {
361+
List<ExprType> argExprTypes =
362+
types.stream().map(OpenSearchTypeFactory::convertRelDataTypeToExprType).toList();
363+
for (var allowedSignature : allowedSignatures) {
364+
if (allowedSignature.size() != types.size()) {
365+
continue; // Skip signatures that do not match the operand count
366+
}
367+
// Check if the argument types match the allowed signature
368+
if (IntStream.range(0, allowedSignature.size())
369+
.allMatch(i -> allowedSignature.get(i).equals(argExprTypes.get(i)))) {
370+
return true;
371+
}
372+
}
373+
return false;
374+
}
375+
376+
@Override
377+
public String getAllowedSignatures() {
378+
return PPLTypeChecker.getExprFamilySignature(allowedSignatures);
379+
}
380+
};
381+
}
382+
395383
// Util Functions
396384
/**
397385
* Generates a list of allowed function signatures based on the provided {@link
@@ -487,6 +475,10 @@ private static String getFamilySignature(List<SqlTypeFamily> families) {
487475
List<List<ExprType>> signatures = Lists.cartesianProduct(exprTypes);
488476

489477
// Convert each signature to a string representation and then concatenate them
478+
return getExprFamilySignature(signatures);
479+
}
480+
481+
private static String getExprFamilySignature(List<List<ExprType>> signatures) {
490482
return signatures.stream()
491483
.map(
492484
types ->

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

Lines changed: 4 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import org.apache.calcite.sql.type.SqlOperandMetadata;
1818
import org.apache.calcite.sql.type.SqlOperandTypeChecker;
1919
import org.apache.calcite.sql.validate.SqlUserDefinedFunction;
20+
import org.opensearch.sql.data.type.ExprType;
2021

2122
/**
2223
* This class is created for the compatibility with {@link SqlUserDefinedFunction} constructors when
@@ -105,47 +106,11 @@ public String getAllowedSignatures(SqlOperator op, String opName) {
105106
};
106107
}
107108

108-
/**
109-
* A named class that serves as an identifier for IP comparator's operand metadata. It does not
110-
* implement any actual type checking logic.
111-
*/
112-
class IPOperandMetadata implements UDFOperandMetadata {
113-
@Override
114-
public SqlOperandTypeChecker getInnerTypeChecker() {
115-
return this;
116-
}
117-
118-
@Override
119-
public List<RelDataType> paramTypes(RelDataTypeFactory typeFactory) {
120-
return List.of();
121-
}
122-
123-
@Override
124-
public List<String> paramNames() {
125-
return List.of();
126-
}
127-
128-
@Override
129-
public boolean checkOperandTypes(SqlCallBinding callBinding, boolean throwOnFailure) {
130-
return false;
131-
}
132-
133-
@Override
134-
public SqlOperandCountRange getOperandCountRange() {
135-
return null;
136-
}
137-
138-
@Override
139-
public String getAllowedSignatures(SqlOperator op, String opName) {
140-
return "";
141-
}
109+
static UDFOperandMetadata wrapUDT(List<List<ExprType>> allowSignatures) {
110+
return new UDTOperandMetadata(allowSignatures);
142111
}
143112

144-
/**
145-
* A named class that serves as an identifier for cidr's operand metadata. It does not implement
146-
* any actual type checking logic.
147-
*/
148-
class CidrOperandMetadata implements UDFOperandMetadata {
113+
record UDTOperandMetadata(List<List<ExprType>> allowedParamTypes) implements UDFOperandMetadata {
149114
@Override
150115
public SqlOperandTypeChecker getInnerTypeChecker() {
151116
return this;

‎core/src/main/java/org/opensearch/sql/expression/function/udf/ip/CidrMatchFunction.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import org.opensearch.sql.data.model.ExprIpValue;
1818
import org.opensearch.sql.data.model.ExprValue;
1919
import org.opensearch.sql.data.model.ExprValueUtils;
20+
import org.opensearch.sql.data.type.ExprCoreType;
2021
import org.opensearch.sql.expression.function.ImplementorUDF;
2122
import org.opensearch.sql.expression.function.UDFOperandMetadata;
2223
import org.opensearch.sql.expression.ip.IPFunctions;
@@ -46,7 +47,10 @@ public UDFOperandMetadata getOperandMetadata() {
4647
// EXPR_IP is mapped to SqlTypeFamily.OTHER in
4748
// UserDefinedFunctionUtils.convertRelDataTypeToSqlTypeName
4849
// We use a specific type checker to serve
49-
return new UDFOperandMetadata.CidrOperandMetadata();
50+
return UDFOperandMetadata.wrapUDT(
51+
List.of(
52+
List.of(ExprCoreType.IP, ExprCoreType.STRING),
53+
List.of(ExprCoreType.STRING, ExprCoreType.STRING)));
5054
}
5155

5256
public static class CidrMatchImplementor implements NotNullImplementor {

‎core/src/main/java/org/opensearch/sql/expression/function/udf/ip/CompareIpFunction.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import org.apache.calcite.sql.type.ReturnTypes;
1616
import org.apache.calcite.sql.type.SqlReturnTypeInference;
1717
import org.opensearch.sql.data.model.ExprIpValue;
18+
import org.opensearch.sql.data.type.ExprCoreType;
1819
import org.opensearch.sql.expression.function.ImplementorUDF;
1920
import org.opensearch.sql.expression.function.UDFOperandMetadata;
2021

@@ -66,7 +67,11 @@ public SqlReturnTypeInference getReturnTypeInference() {
6667

6768
@Override
6869
public UDFOperandMetadata getOperandMetadata() {
69-
return new UDFOperandMetadata.IPOperandMetadata();
70+
return UDFOperandMetadata.wrapUDT(
71+
List.of(
72+
List.of(ExprCoreType.IP, ExprCoreType.IP),
73+
List.of(ExprCoreType.IP, ExprCoreType.STRING),
74+
List.of(ExprCoreType.STRING, ExprCoreType.IP)));
7075
}
7176

7277
public static class CompareImplementor implements NotNullImplementor {

0 commit comments

Comments
 (0)