Skip to content

Commit fe5eb74

Browse files
ishaoxyyuancu
ishaoxy
and
yuancu
authored
[Backport 2.19-dev] Add compare_ip operator udfs (opensearch-project#3821) (opensearch-project#3874)
* Add compare_ip operator udfs (opensearch-project#3821) * ip_compare operator added Signed-off-by: Xinyu Hao <haoxinyu@amazon.com> * only type checker issue left Signed-off-by: Xinyu Hao <haoxinyu@amazon.com> * fix by modifying ip.sqlTypeName from OTHER to NULL in type checker Signed-off-by: Xinyu Hao <haoxinyu@amazon.com> * fix less Signed-off-by: Xinyu Hao <haoxinyu@amazon.com> * modify the CalcitePPLFunctionTypeTest text Signed-off-by: Xinyu Hao <haoxinyu@amazon.com> * allow CalciteIPComparisonIT in CalciteNoPushdownIT Signed-off-by: Xinyu Hao <haoxinyu@amazon.com> * Modify the signature description in udf Signed-off-by: Xinyu Hao <haoxinyu@amazon.com> * fix some typing errors Signed-off-by: Xinyu Hao <haoxinyu@amazon.com> * modify the udfs for better style Signed-off-by: Xinyu Hao <haoxinyu@amazon.com> * Make IpComparisonOperators an inner enum of CompareIPFunction Signed-off-by: Yuanchun Shen <yuanchu@amazon.com> * modify registerOperator Signed-off-by: Xinyu Hao <haoxinyu@amazon.com> * modify registerOperator Signed-off-by: Xinyu Hao <haoxinyu@amazon.com> * add type checker for cidr Signed-off-by: Xinyu Hao <haoxinyu@amazon.com> * add javadoc Signed-off-by: Xinyu Hao <haoxinyu@amazon.com> * move switch case to the implement method Signed-off-by: Xinyu Hao <haoxinyu@amazon.com> --------- Signed-off-by: Xinyu Hao <haoxinyu@amazon.com> Signed-off-by: Yuanchun Shen <yuanchu@amazon.com> Co-authored-by: Yuanchun Shen <yuanchu@amazon.com> (cherry picked from commit 6c3efa1) * backport to java11 Signed-off-by: Xinyu Hao <haoxinyu@amazon.com> * backport to java11 Signed-off-by: Xinyu Hao <haoxinyu@amazon.com> * backport to java11 Signed-off-by: Xinyu Hao <haoxinyu@amazon.com> * backport to java11 Signed-off-by: Xinyu Hao <haoxinyu@amazon.com> * backport to java11 Signed-off-by: Xinyu Hao <haoxinyu@amazon.com> * backport to java11 Signed-off-by: Xinyu Hao <haoxinyu@amazon.com> --------- Signed-off-by: Xinyu Hao <haoxinyu@amazon.com> Signed-off-by: Yuanchun Shen <yuanchu@amazon.com> Co-authored-by: Yuanchun Shen <yuanchu@amazon.com>
1 parent 5e36b38 commit fe5eb74

9 files changed

Lines changed: 353 additions & 93 deletions

File tree

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -127,8 +127,10 @@ public static SqlTypeName convertRelDataTypeToSqlTypeName(RelDataType type) {
127127
return SqlTypeName.TIME;
128128
case EXPR_TIMESTAMP:
129129
return SqlTypeName.TIMESTAMP;
130+
// EXPR_IP is mapped to SqlTypeName.OTHER since there is no
131+
// corresponding SqlTypeName in Calcite.
130132
case EXPR_IP:
131-
return SqlTypeName.VARCHAR;
133+
return SqlTypeName.OTHER;
132134
case EXPR_BINARY:
133135
return SqlTypeName.VARBINARY;
134136
default:

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@
7272
import org.opensearch.sql.expression.function.udf.datetime.WeekdayFunction;
7373
import org.opensearch.sql.expression.function.udf.datetime.YearweekFunction;
7474
import org.opensearch.sql.expression.function.udf.ip.CidrMatchFunction;
75+
import org.opensearch.sql.expression.function.udf.ip.CompareIpFunction;
7576
import org.opensearch.sql.expression.function.udf.math.CRC32Function;
7677
import org.opensearch.sql.expression.function.udf.math.ConvFunction;
7778
import org.opensearch.sql.expression.function.udf.math.DivideFunction;
@@ -103,6 +104,15 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable {
103104
public static final SqlOperator SHA2 = CryptographicFunction.sha2().toUDF("SHA2");
104105
public static final SqlOperator CIDRMATCH = new CidrMatchFunction().toUDF("CIDRMATCH");
105106

107+
// IP comparing functions
108+
public static final SqlOperator NOT_EQUALS_IP =
109+
CompareIpFunction.notEquals().toUDF("NOT_EQUALS_IP");
110+
public static final SqlOperator EQUALS_IP = CompareIpFunction.equals().toUDF("EQUALS_IP");
111+
public static final SqlOperator GREATER_IP = CompareIpFunction.greater().toUDF("GREATER_IP");
112+
public static final SqlOperator GTE_IP = CompareIpFunction.greaterOrEquals().toUDF("GTE_IP");
113+
public static final SqlOperator LESS_IP = CompareIpFunction.less().toUDF("LESS_IP");
114+
public static final SqlOperator LTE_IP = CompareIpFunction.lessOrEquals().toUDF("LTE_IP");
115+
106116
// Condition function
107117
public static final SqlOperator EARLIEST = new EarliestFunction().toUDF("EARLIEST");
108118
public static final SqlOperator LATEST = new LatestFunction().toUDF("LATEST");

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

Lines changed: 76 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -465,7 +465,10 @@ functionName, getActualSignature(argTypes), e.getMessage()),
465465
}
466466
StringJoiner allowedSignatures = new StringJoiner(",");
467467
for (var implement : implementList) {
468-
allowedSignatures.add(implement.getKey().getTypeChecker().getAllowedSignatures());
468+
String signature = implement.getKey().getTypeChecker().getAllowedSignatures();
469+
if (!signature.isEmpty()) {
470+
allowedSignatures.add(signature);
471+
}
469472
}
470473
throw new ExpressionEvaluationException(
471474
String.format(
@@ -488,45 +491,71 @@ private abstract static class AbstractBuilder {
488491
/** Maps an operator to an implementation. */
489492
abstract void register(BuiltinFunctionName functionName, FunctionImp functionImp);
490493

491-
void registerOperator(BuiltinFunctionName functionName, SqlOperator operator) {
492-
SqlOperandTypeChecker typeChecker;
493-
if (operator instanceof SqlUserDefinedFunction) {
494-
SqlUserDefinedFunction udfOperator = (SqlUserDefinedFunction) operator;
495-
typeChecker = extractTypeCheckerFromUDF(udfOperator);
496-
} else {
497-
typeChecker = operator.getOperandTypeChecker();
498-
}
494+
/**
495+
* Register one or multiple operators under a single function name. This allows function
496+
* overloading based on operand types.
497+
*
498+
* <p>When a function is called, the system will try each registered operator in sequence,
499+
* checking if the provided arguments match the operator's type requirements. The first operator
500+
* whose type checker accepts the arguments will be used to execute the function.
501+
*
502+
* @param functionName the built-in function name under which to register the operators
503+
* @param operators the operators to associate with this function name, tried in sequence until
504+
* one matches the argument types during resolution
505+
*/
506+
void registerOperator(BuiltinFunctionName functionName, SqlOperator... operators) {
507+
for (SqlOperator operator : operators) {
508+
SqlOperandTypeChecker typeChecker;
509+
if (operator instanceof SqlUserDefinedFunction) {
510+
SqlUserDefinedFunction udfOperator = (SqlUserDefinedFunction) operator;
511+
typeChecker = extractTypeCheckerFromUDF(udfOperator);
512+
} else {
513+
typeChecker = operator.getOperandTypeChecker();
514+
}
499515

500-
// Only the composite operand type checker for UDFs are concerned here.
501-
if (operator instanceof SqlUserDefinedFunction
502-
&& typeChecker instanceof CompositeOperandTypeChecker) {
503-
CompositeOperandTypeChecker compositeTypeChecker = (CompositeOperandTypeChecker) typeChecker;
504-
// UDFs implement their own composite type checkers, which always use OR logic for argument
505-
// types. Verifying the composition type would require accessing a protected field in
506-
// CompositeOperandTypeChecker. If access to this field is not allowed, type checking will
507-
// be skipped, so we avoid checking the composition type here.
508-
register(functionName, wrapWithCompositeTypeChecker(operator, compositeTypeChecker, false));
509-
} else if (typeChecker instanceof ImplicitCastOperandTypeChecker) {
510-
ImplicitCastOperandTypeChecker implicitCastTypeChecker = (ImplicitCastOperandTypeChecker) typeChecker;
511-
register(functionName, wrapWithImplicitCastTypeChecker(operator, implicitCastTypeChecker));
512-
} else if (typeChecker instanceof CompositeOperandTypeChecker) {
513-
CompositeOperandTypeChecker compositeTypeChecker = (CompositeOperandTypeChecker) typeChecker;
514-
// If compositeTypeChecker contains operand checkers other than family type checkers or
515-
// other than OR compositions, the function with be registered with a null type checker,
516-
// which means the function will not be type checked.
517-
register(functionName, wrapWithCompositeTypeChecker(operator, compositeTypeChecker, true));
518-
} else if (typeChecker instanceof SameOperandTypeChecker) {
519-
SameOperandTypeChecker comparableTypeChecker = (SameOperandTypeChecker) typeChecker;
520-
// Comparison operators like EQUAL, GREATER_THAN, LESS_THAN, etc.
521-
// SameOperandTypeCheckers like COALESCE, IFNULL, etc.
522-
register(functionName, wrapWithComparableTypeChecker(operator, comparableTypeChecker));
523-
} else {
524-
logger.info(
525-
"Cannot create type checker for function: {}. Will skip its type checking",
526-
functionName);
527-
register(
528-
functionName,
529-
(RexBuilder builder, RexNode... node) -> builder.makeCall(operator, node));
516+
// Only the composite operand type checker for UDFs are concerned here.
517+
if (operator instanceof SqlUserDefinedFunction
518+
&& typeChecker instanceof CompositeOperandTypeChecker) {
519+
CompositeOperandTypeChecker compositeTypeChecker = (CompositeOperandTypeChecker) typeChecker;
520+
// UDFs implement their own composite type checkers, which always use OR logic for argument
521+
// types. Verifying the composition type would require accessing a protected field in
522+
// CompositeOperandTypeChecker. If access to this field is not allowed, type checking will
523+
// be skipped, so we avoid checking the composition type here.
524+
register(functionName, wrapWithCompositeTypeChecker(operator, compositeTypeChecker, false));
525+
} else if (typeChecker instanceof ImplicitCastOperandTypeChecker) {
526+
ImplicitCastOperandTypeChecker implicitCastTypeChecker = (ImplicitCastOperandTypeChecker) typeChecker;
527+
register(functionName, wrapWithImplicitCastTypeChecker(operator, implicitCastTypeChecker));
528+
} else if (typeChecker instanceof CompositeOperandTypeChecker) {
529+
CompositeOperandTypeChecker compositeTypeChecker = (CompositeOperandTypeChecker) typeChecker;
530+
// If compositeTypeChecker contains operand checkers other than family type checkers or
531+
// other than OR compositions, the function with be registered with a null type checker,
532+
// which means the function will not be type checked.
533+
register(functionName, wrapWithCompositeTypeChecker(operator, compositeTypeChecker, true));
534+
} else if (typeChecker instanceof SameOperandTypeChecker) {
535+
SameOperandTypeChecker comparableTypeChecker = (SameOperandTypeChecker) typeChecker;
536+
// Comparison operators like EQUAL, GREATER_THAN, LESS_THAN, etc.
537+
// SameOperandTypeCheckers like COALESCE, IFNULL, etc.
538+
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()));
551+
} else {
552+
logger.info(
553+
"Cannot create type checker for function: {}. Will skip its type checking",
554+
functionName);
555+
register(
556+
functionName,
557+
(RexBuilder builder, RexNode... node) -> builder.makeCall(operator, node));
558+
}
530559
}
531560
}
532561

@@ -634,16 +663,18 @@ public PPLTypeChecker getTypeChecker() {
634663
}
635664

636665
void populate() {
666+
// register operators for comparison
667+
registerOperator(NOTEQUAL, PPLBuiltinOperators.NOT_EQUALS_IP, SqlStdOperatorTable.NOT_EQUALS);
668+
registerOperator(EQUAL, PPLBuiltinOperators.EQUALS_IP, SqlStdOperatorTable.EQUALS);
669+
registerOperator(GREATER, PPLBuiltinOperators.GREATER_IP, SqlStdOperatorTable.GREATER_THAN);
670+
registerOperator(GTE, PPLBuiltinOperators.GTE_IP, SqlStdOperatorTable.GREATER_THAN_OR_EQUAL);
671+
registerOperator(LESS, PPLBuiltinOperators.LESS_IP, SqlStdOperatorTable.LESS_THAN);
672+
registerOperator(LTE, PPLBuiltinOperators.LTE_IP, SqlStdOperatorTable.LESS_THAN_OR_EQUAL);
673+
637674
// Register std operator
638675
registerOperator(AND, SqlStdOperatorTable.AND);
639676
registerOperator(OR, SqlStdOperatorTable.OR);
640677
registerOperator(NOT, SqlStdOperatorTable.NOT);
641-
registerOperator(NOTEQUAL, SqlStdOperatorTable.NOT_EQUALS);
642-
registerOperator(EQUAL, SqlStdOperatorTable.EQUALS);
643-
registerOperator(GREATER, SqlStdOperatorTable.GREATER_THAN);
644-
registerOperator(GTE, SqlStdOperatorTable.GREATER_THAN_OR_EQUAL);
645-
registerOperator(LESS, SqlStdOperatorTable.LESS_THAN);
646-
registerOperator(LTE, SqlStdOperatorTable.LESS_THAN_OR_EQUAL);
647678
registerOperator(ADD, SqlStdOperatorTable.PLUS);
648679
registerOperator(SUBTRACT, SqlStdOperatorTable.MINUS);
649680
registerOperator(MULTIPLY, SqlStdOperatorTable.MULTIPLY);

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

Lines changed: 48 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
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.AbstractExprRelDataType;
26+
import org.opensearch.sql.calcite.type.ExprIPType;
2727
import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory;
2828
import org.opensearch.sql.calcite.utils.UserDefinedFunctionUtils;
2929
import org.opensearch.sql.data.type.ExprCoreType;
@@ -222,10 +222,6 @@ public boolean checkOperandTypes(List<RelDataType> types) {
222222
RelDataType type_l = types.get(i);
223223
RelDataType type_r = types.get(i + 1);
224224
if (!SqlTypeUtil.isComparable(type_l, type_r)) {
225-
if (areIpAndStringTypes(type_l, type_r) || areIpAndStringTypes(type_r, type_l)) {
226-
// Allow IP and string comparison
227-
continue;
228-
}
229225
return false;
230226
}
231227
// Disallow coercing between strings and numeric, boolean
@@ -252,15 +248,6 @@ private static boolean cannotConvertStringInCompare(SqlTypeFamily typeFamily) {
252248
}
253249
}
254250

255-
private static boolean areIpAndStringTypes(RelDataType typeIp, RelDataType typeString) {
256-
if (typeIp instanceof AbstractExprRelDataType<?>) {
257-
AbstractExprRelDataType<?> exprRelDataType = (AbstractExprRelDataType<?>) typeIp;
258-
return exprRelDataType.getExprType() == ExprCoreType.IP
259-
&& typeString.getFamily() == SqlTypeFamily.CHARACTER;
260-
}
261-
return false;
262-
}
263-
264251
@Override
265252
public String getAllowedSignatures() {
266253
int min = innerTypeChecker.getOperandCountRange().getMin();
@@ -283,6 +270,53 @@ public String getAllowedSignatures() {
283270
}
284271
}
285272

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+
286320
/**
287321
* Creates a {@link PPLFamilyTypeChecker} with a fixed operand count, validating that each operand
288322
* belongs to its corresponding {@link SqlTypeFamily}.

0 commit comments

Comments
 (0)