Skip to content

Commit 86fd1bd

Browse files
yuancuyuancu
committed
Refactor registerExternalFunction to registerExternalOperator
Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>
1 parent 0092724 commit 86fd1bd

2 files changed

Lines changed: 71 additions & 58 deletions

File tree

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

Lines changed: 70 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -352,24 +352,24 @@ private PPLFuncImpTable(Builder builder, AggBuilder aggBuilder) {
352352
}
353353

354354
/**
355-
* Register a function implementation from external services dynamically.
355+
* Register an operator from external services dynamically.
356356
*
357357
* @param functionName the name of the function, has to be defined in BuiltinFunctionName
358-
* @param functionImp the implementation of the function
359-
* @param typeChecker the type checker of the function. If null, type checking will be skipped for
360-
* this function.
358+
* @param operator a SqlOperator representing an externally implemented function
361359
*/
362-
public void registerExternalFunction(
363-
BuiltinFunctionName functionName,
364-
FunctionImp functionImp,
365-
@Nullable PPLTypeChecker typeChecker) {
360+
public void registerExternalOperator(BuiltinFunctionName functionName, SqlOperator operator) {
361+
PPLTypeChecker typeChecker =
362+
wrapSqlOperandTypeChecker(
363+
operator.getOperandTypeChecker(),
364+
operator.getName(),
365+
operator instanceof SqlUserDefinedFunction);
366366
CalciteFuncSignature signature = new CalciteFuncSignature(functionName.getName(), typeChecker);
367367
externalFunctionRegistry.compute(
368368
functionName,
369369
(name, existingList) -> {
370370
List<Pair<CalciteFuncSignature, FunctionImp>> list =
371371
existingList == null ? new ArrayList<>() : new ArrayList<>(existingList);
372-
list.add(Pair.of(signature, functionImp));
372+
list.add(Pair.of(signature, (builder, args) -> builder.makeCall(operator, args)));
373373
return list;
374374
});
375375
}
@@ -519,6 +519,12 @@ private void compulsoryCast(
519519
return null;
520520
}
521521

522+
/**
523+
* Get a string representation of the argument types expressed in ExprType for error messages.
524+
*
525+
* @param argTypes the list of argument types as {@link RelDataType}
526+
* @return a string in the format [type1,type2,...] representing the argument types
527+
*/
522528
private static String getActualSignature(List<RelDataType> argTypes) {
523529
return "["
524530
+ argTypes.stream()
@@ -528,6 +534,58 @@ private static String getActualSignature(List<RelDataType> argTypes) {
528534
+ "]";
529535
}
530536

537+
/**
538+
* Wraps a {@link SqlOperandTypeChecker} into a {@link PPLTypeChecker} for use in function
539+
* signature validation.
540+
*
541+
* @param typeChecker the original SQL operand type checker
542+
* @param functionName the name of the function for error reporting
543+
* @param isUserDefinedFunction true if the function is user-defined, false otherwise
544+
* @return a {@link PPLTypeChecker} that delegates to the provided {@code typeChecker}
545+
*/
546+
private static PPLTypeChecker wrapSqlOperandTypeChecker(
547+
SqlOperandTypeChecker typeChecker, String functionName, boolean isUserDefinedFunction) {
548+
PPLTypeChecker pplTypeChecker;
549+
// Only the composite operand type checker for UDFs are concerned here.
550+
if (isUserDefinedFunction
551+
&& typeChecker instanceof CompositeOperandTypeChecker compositeTypeChecker) {
552+
// UDFs implement their own composite type checkers, which always use OR logic for
553+
// argument
554+
// types. Verifying the composition type would require accessing a protected field in
555+
// CompositeOperandTypeChecker. If access to this field is not allowed, type checking will
556+
// be skipped, so we avoid checking the composition type here.
557+
pplTypeChecker = PPLTypeChecker.wrapComposite(compositeTypeChecker, false);
558+
} else if (typeChecker instanceof ImplicitCastOperandTypeChecker implicitCastTypeChecker) {
559+
pplTypeChecker = PPLTypeChecker.wrapFamily(implicitCastTypeChecker);
560+
} else if (typeChecker instanceof CompositeOperandTypeChecker compositeTypeChecker) {
561+
// If compositeTypeChecker contains operand checkers other than family type checkers or
562+
// other than OR compositions, the function with be registered with a null type checker,
563+
// which means the function will not be type checked.
564+
try {
565+
pplTypeChecker = PPLTypeChecker.wrapComposite(compositeTypeChecker, true);
566+
} catch (IllegalArgumentException | UnsupportedOperationException e) {
567+
logger.debug(
568+
String.format(
569+
"Failed to create composite type checker for operator: %s. Will skip its type"
570+
+ " checking",
571+
functionName),
572+
e);
573+
pplTypeChecker = null;
574+
}
575+
} else if (typeChecker instanceof SameOperandTypeChecker comparableTypeChecker) {
576+
// Comparison operators like EQUAL, GREATER_THAN, LESS_THAN, etc.
577+
// SameOperandTypeCheckers like COALESCE, IFNULL, etc.
578+
pplTypeChecker = PPLTypeChecker.wrapComparable(comparableTypeChecker);
579+
} else if (typeChecker instanceof UDFOperandMetadata.UDTOperandMetadata udtOperandMetadata) {
580+
pplTypeChecker = PPLTypeChecker.wrapUDT(udtOperandMetadata.allowedParamTypes());
581+
} else {
582+
logger.info(
583+
"Cannot create type checker for function: {}. Will skip its type checking", functionName);
584+
pplTypeChecker = null;
585+
}
586+
return pplTypeChecker;
587+
}
588+
531589
@SuppressWarnings({"UnusedReturnValue", "SameParameterValue"})
532590
private abstract static class AbstractBuilder {
533591

@@ -556,46 +614,9 @@ public void registerOperator(BuiltinFunctionName functionName, SqlOperator... op
556614
typeChecker = operator.getOperandTypeChecker();
557615
}
558616

559-
PPLTypeChecker pplTypeChecker;
560-
// Only the composite operand type checker for UDFs are concerned here.
561-
if (operator instanceof SqlUserDefinedFunction
562-
&& typeChecker instanceof CompositeOperandTypeChecker compositeTypeChecker) {
563-
// UDFs implement their own composite type checkers, which always use OR logic for
564-
// argument
565-
// types. Verifying the composition type would require accessing a protected field in
566-
// CompositeOperandTypeChecker. If access to this field is not allowed, type checking will
567-
// be skipped, so we avoid checking the composition type here.
568-
pplTypeChecker = PPLTypeChecker.wrapComposite(compositeTypeChecker, false);
569-
} else if (typeChecker instanceof ImplicitCastOperandTypeChecker implicitCastTypeChecker) {
570-
pplTypeChecker = PPLTypeChecker.wrapFamily(implicitCastTypeChecker);
571-
} else if (typeChecker instanceof CompositeOperandTypeChecker compositeTypeChecker) {
572-
// If compositeTypeChecker contains operand checkers other than family type checkers or
573-
// other than OR compositions, the function with be registered with a null type checker,
574-
// which means the function will not be type checked.
575-
try {
576-
pplTypeChecker = PPLTypeChecker.wrapComposite(compositeTypeChecker, true);
577-
} catch (IllegalArgumentException | UnsupportedOperationException e) {
578-
logger.debug(
579-
String.format(
580-
"Failed to create composite type checker for operator: %s. Will skip its type"
581-
+ " checking",
582-
operator.getName()),
583-
e);
584-
pplTypeChecker = null;
585-
}
586-
} else if (typeChecker instanceof SameOperandTypeChecker comparableTypeChecker) {
587-
// Comparison operators like EQUAL, GREATER_THAN, LESS_THAN, etc.
588-
// SameOperandTypeCheckers like COALESCE, IFNULL, etc.
589-
pplTypeChecker = PPLTypeChecker.wrapComparable(comparableTypeChecker);
590-
} else if (typeChecker
591-
instanceof UDFOperandMetadata.UDTOperandMetadata udtOperandMetadata) {
592-
pplTypeChecker = PPLTypeChecker.wrapUDT(udtOperandMetadata.allowedParamTypes());
593-
} else {
594-
logger.info(
595-
"Cannot create type checker for function: {}. Will skip its type checking",
596-
functionName);
597-
pplTypeChecker = null;
598-
}
617+
PPLTypeChecker pplTypeChecker =
618+
wrapSqlOperandTypeChecker(
619+
typeChecker, operator.getName(), operator instanceof SqlUserDefinedFunction);
599620
register(
600621
functionName,
601622
(RexBuilder builder, RexNode... args) -> builder.makeCall(operator, args),

‎opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java‎

Lines changed: 1 addition & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,6 @@
2727
import org.apache.calcite.rel.type.RelDataTypeField;
2828
import org.apache.calcite.runtime.Hook;
2929
import org.apache.calcite.sql.SqlExplainLevel;
30-
import org.apache.calcite.sql.type.CompositeOperandTypeChecker;
3130
import org.apache.calcite.sql.type.ReturnTypes;
3231
import org.apache.calcite.sql.type.SqlTypeName;
3332
import org.apache.calcite.sql.validate.SqlUserDefinedFunction;
@@ -46,7 +45,6 @@
4645
import org.opensearch.sql.executor.pagination.PlanSerializer;
4746
import org.opensearch.sql.expression.function.BuiltinFunctionName;
4847
import org.opensearch.sql.expression.function.PPLFuncImpTable;
49-
import org.opensearch.sql.expression.function.PPLTypeChecker;
5048
import org.opensearch.sql.opensearch.client.OpenSearchClient;
5149
import org.opensearch.sql.opensearch.executor.protector.ExecutionProtector;
5250
import org.opensearch.sql.opensearch.functions.DistinctCountApproxAggFunction;
@@ -265,13 +263,7 @@ private void buildResultSet(
265263
/** Registers opensearch-dependent functions */
266264
private void registerOpenSearchFunctions() {
267265
SqlUserDefinedFunction geoIpFunction = new GeoIpFunction(client.getNodeClient()).toUDF("GEOIP");
268-
PPLFuncImpTable.FunctionImp geoIpImpl =
269-
(builder, args) -> builder.makeCall(geoIpFunction, args);
270-
PPLFuncImpTable.INSTANCE.registerExternalFunction(
271-
BuiltinFunctionName.GEOIP,
272-
geoIpImpl,
273-
PPLTypeChecker.wrapComposite(
274-
(CompositeOperandTypeChecker) geoIpFunction.getOperandTypeChecker(), false));
266+
PPLFuncImpTable.INSTANCE.registerExternalOperator(BuiltinFunctionName.GEOIP, geoIpFunction);
275267

276268
PPLFuncImpTable.INSTANCE.registerExternalAggFunction(
277269
DISTINCT_COUNT_APPROX,

0 commit comments

Comments
 (0)