Skip to content

Commit 7f2a4c1

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 7f2a4c1

15 files changed

Lines changed: 723 additions & 458 deletions

File tree

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

Lines changed: 43 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,10 @@
88
import com.google.common.collect.ImmutableList;
99
import java.util.Arrays;
1010
import java.util.List;
11+
import java.util.Set;
1112
import java.util.stream.Collectors;
1213

14+
import java.util.Locale;
1315
import org.apache.calcite.avatica.util.TimeUnit;
1416
import org.apache.calcite.rel.type.RelDataType;
1517
import org.apache.calcite.rex.RexBuilder;
@@ -21,6 +23,18 @@
2123
import org.apache.calcite.sql.type.SqlTypeName;
2224
import org.apache.calcite.sql.type.SqlTypeUtil;
2325
import org.opensearch.sql.ast.expression.SpanUnit;
26+
import org.opensearch.sql.calcite.type.AbstractExprRelDataType;
27+
import org.opensearch.sql.calcite.type.ExprSqlType;
28+
import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory;
29+
import org.opensearch.sql.data.type.ExprCoreType;
30+
import org.opensearch.sql.exception.ExpressionEvaluationException;
31+
import org.opensearch.sql.exception.SemanticCheckException;
32+
import org.opensearch.sql.expression.function.PPLBuiltinOperators;
33+
34+
import static org.opensearch.sql.calcite.utils.OpenSearchTypeFactory.ExprUDT.EXPR_DATE;
35+
import static org.opensearch.sql.calcite.utils.OpenSearchTypeFactory.ExprUDT.EXPR_IP;
36+
import static org.opensearch.sql.calcite.utils.OpenSearchTypeFactory.ExprUDT.EXPR_TIME;
37+
import static org.opensearch.sql.calcite.utils.OpenSearchTypeFactory.ExprUDT.EXPR_TIMESTAMP;
2438

2539
public class ExtendedRexBuilder extends RexBuilder {
2640

@@ -119,7 +133,35 @@ public RexNode makeCast(
119133
// SqlStdOperatorTable.NOT_EQUALS,
120134
// ImmutableList.of(exp, makeZeroLiteral(exp.getType())));
121135
}
122-
}
136+
} else if (OpenSearchTypeFactory.isUserDefinedType(type)) {
137+
var udt = ((AbstractExprRelDataType<?>) type).getUdt();
138+
var argExprType = OpenSearchTypeFactory.convertRelDataTypeToExprType(exp.getType());
139+
switch (udt) {
140+
case EXPR_DATE:
141+
return makeCall(type, PPLBuiltinOperators.DATE, List.of(exp));
142+
case EXPR_TIME:
143+
return makeCall(type, PPLBuiltinOperators.TIME, List.of(exp));
144+
case EXPR_TIMESTAMP:
145+
return makeCall(type, PPLBuiltinOperators.TIMESTAMP, List.of(exp));
146+
case EXPR_IP:
147+
if (argExprType == ExprCoreType.IP) {
148+
return exp;
149+
} else if (argExprType == ExprCoreType.STRING) {
150+
return makeCall(type, PPLBuiltinOperators.IP, List.of(exp));
151+
}
152+
// Throwing error inside implementation will be suppressed by Calcite, thus
153+
// throwing 500 error. Therefore, we throw error here to ensure the error
154+
// information is displayed properly.
155+
throw new ExpressionEvaluationException(
156+
String.format(
157+
Locale.ROOT,
158+
"Cannot convert %s to IP, only STRING and IP types are supported",
159+
argExprType));
160+
default:
161+
throw new SemanticCheckException(
162+
String.format(Locale.ROOT, "Cannot cast from %s to %s", argExprType, udt.name()));
163+
}
164+
}
123165
return super.makeCast(pos, type, exp, matchNullability, safe, format);
124166
}
125167
}

‎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: 29 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -516,38 +516,30 @@ void registerOperator(BuiltinFunctionName functionName, SqlOperator... operators
516516
// Only the composite operand type checker for UDFs are concerned here.
517517
if (operator instanceof SqlUserDefinedFunction
518518
&& typeChecker instanceof CompositeOperandTypeChecker) {
519-
CompositeOperandTypeChecker compositeTypeChecker = (CompositeOperandTypeChecker) typeChecker;
520519
// UDFs implement their own composite type checkers, which always use OR logic for argument
521520
// types. Verifying the composition type would require accessing a protected field in
522521
// CompositeOperandTypeChecker. If access to this field is not allowed, type checking will
523522
// be skipped, so we avoid checking the composition type here.
523+
CompositeOperandTypeChecker compositeTypeChecker = (CompositeOperandTypeChecker) typeChecker;
524524
register(functionName, wrapWithCompositeTypeChecker(operator, compositeTypeChecker, false));
525525
} else if (typeChecker instanceof ImplicitCastOperandTypeChecker) {
526526
ImplicitCastOperandTypeChecker implicitCastTypeChecker = (ImplicitCastOperandTypeChecker) typeChecker;
527527
register(functionName, wrapWithImplicitCastTypeChecker(operator, implicitCastTypeChecker));
528528
} else if (typeChecker instanceof CompositeOperandTypeChecker) {
529-
CompositeOperandTypeChecker compositeTypeChecker = (CompositeOperandTypeChecker) typeChecker;
530529
// If compositeTypeChecker contains operand checkers other than family type checkers or
531530
// other than OR compositions, the function with be registered with a null type checker,
532531
// which means the function will not be type checked.
532+
CompositeOperandTypeChecker compositeTypeChecker = (CompositeOperandTypeChecker) typeChecker;
533533
register(functionName, wrapWithCompositeTypeChecker(operator, compositeTypeChecker, true));
534534
} else if (typeChecker instanceof SameOperandTypeChecker) {
535-
SameOperandTypeChecker comparableTypeChecker = (SameOperandTypeChecker) typeChecker;
536535
// Comparison operators like EQUAL, GREATER_THAN, LESS_THAN, etc.
537536
// SameOperandTypeCheckers like COALESCE, IFNULL, etc.
537+
SameOperandTypeChecker comparableTypeChecker = (SameOperandTypeChecker) typeChecker;
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) {
541+
UDFOperandMetadata.UDTOperandMetadata udtOperandMetadata = (UDFOperandMetadata.UDTOperandMetadata) typeChecker;
542+
register(functionName, wrapWithUdtTypeChecker(operator, udtOperandMetadata));
551543
} else {
552544
logger.info(
553545
"Cannot create type checker for function: {}. Will skip its type checking",
@@ -566,6 +558,13 @@ private static SqlOperandTypeChecker extractTypeCheckerFromUDF(
566558
return (udfOperandMetadata == null) ? null : udfOperandMetadata.getInnerTypeChecker();
567559
}
568560

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

634+
private static FunctionImp wrapWithUdtTypeChecker(
635+
SqlOperator operator, UDFOperandMetadata.UDTOperandMetadata udtOperandMetadata) {
636+
return new FunctionImp() {
637+
@Override
638+
public RexNode resolve(RexBuilder builder, RexNode... args) {
639+
return builder.makeCall(operator, args);
640+
}
641+
642+
@Override
643+
public PPLTypeChecker getTypeChecker() {
644+
return PPLTypeChecker.wrapUDT(udtOperandMetadata.getAllowSignatures());
645+
}
646+
};
647+
}
648+
635649
private static FunctionImp createFunctionImpWithTypeChecker(
636650
BiFunction<RexBuilder, RexNode, RexNode> resolver, PPLTypeChecker typeChecker) {
637651
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).collect(Collectors.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: 12 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@
77

88
import java.util.Collections;
99
import java.util.List;
10+
11+
import lombok.Getter;
1012
import org.apache.calcite.rel.type.RelDataType;
1113
import org.apache.calcite.rel.type.RelDataTypeFactory;
1214
import org.apache.calcite.sql.SqlCallBinding;
@@ -17,6 +19,7 @@
1719
import org.apache.calcite.sql.type.SqlOperandMetadata;
1820
import org.apache.calcite.sql.type.SqlOperandTypeChecker;
1921
import org.apache.calcite.sql.validate.SqlUserDefinedFunction;
22+
import org.opensearch.sql.data.type.ExprType;
2023

2124
/**
2225
* This class is created for the compatibility with {@link SqlUserDefinedFunction} constructors when
@@ -105,48 +108,18 @@ public String getAllowedSignatures(SqlOperator op, String opName) {
105108
};
106109
}
107110

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-
}
111+
static UDFOperandMetadata wrapUDT(List<List<ExprType>> allowSignatures) {
112+
return new UDTOperandMetadata(allowSignatures);
113+
}
132114

133-
@Override
134-
public SqlOperandCountRange getOperandCountRange() {
135-
return null;
136-
}
115+
class UDTOperandMetadata implements UDFOperandMetadata {
116+
@Getter private final List<List<ExprType>> allowSignatures;
137117

138-
@Override
139-
public String getAllowedSignatures(SqlOperator op, String opName) {
140-
return "";
141-
}
142-
}
118+
public UDTOperandMetadata(List<List<ExprType>> allowSignatures) {
119+
this.allowSignatures = allowSignatures;
120+
}
143121

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 {
149-
@Override
122+
@Override
150123
public SqlOperandTypeChecker getInnerTypeChecker() {
151124
return this;
152125
}

0 commit comments

Comments
 (0)