Skip to content

Commit 92cb089

Browse files
yuancuyuancu
authored
Implement geoip udf with Calcite (opensearch-project#3604)
* Implement GEOIP Signed-off-by: Yuanchun Shen <yuanchu@amazon.com> * Add class documentation for GeoIpFunction udf Signed-off-by: Yuanchun Shen <yuanchu@amazon.com> * Alter node client passing: OpenSearchPluginModule -> QueryService -> CalcitePlanContext Signed-off-by: Yuanchun Shen <yuanchu@amazon.com> * Refactor: separate geoip function implementation into w/ and w/o options versions Signed-off-by: Yuanchun Shen <yuanchu@amazon.com> * Move geoip function to opensearch package & register it in OpenSearchExecutionEngine Signed-off-by: Yuanchun Shen <yuanchu@amazon.com> * Refactor: remove stale comments Signed-off-by: Yuanchun Shen <yuanchu@amazon.com> * Add javadoc for function registries in PPLFuncImpTable Signed-off-by: Yuanchun Shen <yuanchu@amazon.com> * Refactor: replace Pairlist<L,R> with List<Pair<L,R>> Signed-off-by: Yuanchun Shen <yuanchu@amazon.com> * Retrieve UDF implementations from external registry first to allow data-storage-dependent overriding Signed-off-by: Yuanchun Shen <yuanchu@amazon.com> * Remove stale resolveSafe function Signed-off-by: Yuanchun Shen <yuanchu@amazon.com> * Use Locale.ROOT in toLowercase() when building AST Signed-off-by: Yuanchun Shen <yuanchu@amazon.com> --------- Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>
1 parent 66e230e commit 92cb089

7 files changed

Lines changed: 195 additions & 35 deletions

File tree

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -342,7 +342,7 @@ public RexNode visitFunction(Function node, CalcitePlanContext context) {
342342
List<RexNode> arguments =
343343
node.getFuncArgs().stream().map(arg -> analyze(arg, context)).toList();
344344
RexNode resolvedNode =
345-
PPLFuncImpTable.INSTANCE.resolveSafe(
345+
PPLFuncImpTable.INSTANCE.resolve(
346346
context.rexBuilder, node.getFuncName(), arguments.toArray(new RexNode[0]));
347347
if (resolvedNode != null) {
348348
return resolvedNode;

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

Lines changed: 46 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111

1212
import com.google.common.collect.ImmutableMap;
1313
import java.math.BigDecimal;
14+
import java.util.ArrayList;
1415
import java.util.Arrays;
1516
import java.util.HashMap;
1617
import java.util.List;
@@ -19,13 +20,12 @@
1920
import org.apache.calcite.rel.type.RelDataType;
2021
import org.apache.calcite.rex.RexBuilder;
2122
import org.apache.calcite.rex.RexNode;
22-
import org.apache.calcite.runtime.PairList;
2323
import org.apache.calcite.sql.SqlOperator;
2424
import org.apache.calcite.sql.fun.SqlLibraryOperators;
2525
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
2626
import org.apache.calcite.sql.fun.SqlTrimFunction.Flag;
2727
import org.apache.calcite.sql.type.SqlTypeName;
28-
import org.checkerframework.checker.nullness.qual.Nullable;
28+
import org.apache.commons.lang3.tuple.Pair;
2929
import org.opensearch.sql.executor.QueryType;
3030

3131
public class PPLFuncImpTable {
@@ -91,21 +91,44 @@ default List<RelDataType> getParams() {
9191
INSTANCE = new PPLFuncImpTable(builder);
9292
}
9393

94-
private final ImmutableMap<BuiltinFunctionName, PairList<CalciteFuncSignature, FunctionImp>> map;
94+
/**
95+
* The registry for built-in functions. Functions defined by the PPL specification, whose
96+
* implementations are independent of any specific data storage, should be registered here
97+
* internally.
98+
*/
99+
private final ImmutableMap<BuiltinFunctionName, List<Pair<CalciteFuncSignature, FunctionImp>>>
100+
functionRegistry;
101+
102+
/**
103+
* The external function registry. Functions whose implementations depend on a specific data
104+
* engine should be registered here. This reduces coupling between the core module and particular
105+
* storage backends.
106+
*/
107+
private final Map<BuiltinFunctionName, List<Pair<CalciteFuncSignature, FunctionImp>>>
108+
externalFunctionRegistry;
95109

96110
private PPLFuncImpTable(Builder builder) {
97-
final ImmutableMap.Builder<BuiltinFunctionName, PairList<CalciteFuncSignature, FunctionImp>>
111+
final ImmutableMap.Builder<BuiltinFunctionName, List<Pair<CalciteFuncSignature, FunctionImp>>>
98112
mapBuilder = ImmutableMap.builder();
99-
builder.map.forEach((k, v) -> mapBuilder.put(k, v.immutable()));
100-
this.map = ImmutableMap.copyOf(mapBuilder.build());
113+
builder.map.forEach((k, v) -> mapBuilder.put(k, List.copyOf(v)));
114+
this.functionRegistry = ImmutableMap.copyOf(mapBuilder.build());
115+
this.externalFunctionRegistry = new HashMap<>();
101116
}
102117

103-
public @Nullable RexNode resolveSafe(
104-
final RexBuilder builder, final String functionName, RexNode... args) {
105-
try {
106-
return resolve(builder, functionName, args);
107-
} catch (Exception e) {
108-
return null;
118+
/**
119+
* Register a function implementation from external services dynamically.
120+
*
121+
* @param functionName the name of the function, has to be defined in BuiltinFunctionName
122+
* @param functionImp the implementation of the function
123+
*/
124+
public void registerExternalFunction(BuiltinFunctionName functionName, FunctionImp functionImp) {
125+
CalciteFuncSignature signature =
126+
new CalciteFuncSignature(functionName.getName(), functionImp.getParams());
127+
if (externalFunctionRegistry.containsKey(functionName)) {
128+
externalFunctionRegistry.get(functionName).add(Pair.of(signature, functionImp));
129+
} else {
130+
externalFunctionRegistry.put(
131+
functionName, new ArrayList<>(List.of(Pair.of(signature, functionImp))));
109132
}
110133
}
111134

@@ -119,7 +142,14 @@ public RexNode resolve(final RexBuilder builder, final String functionName, RexN
119142

120143
public RexNode resolve(
121144
final RexBuilder builder, final BuiltinFunctionName functionName, RexNode... args) {
122-
final PairList<CalciteFuncSignature, FunctionImp> implementList = map.get(functionName);
145+
// Check the external function registry first. This allows the data-storage-dependent
146+
// function implementations to override the internal ones with the same name.
147+
List<Pair<CalciteFuncSignature, FunctionImp>> implementList =
148+
externalFunctionRegistry.get(functionName);
149+
// If the function is not part of the external registry, check the internal registry.
150+
if (implementList == null) {
151+
implementList = functionRegistry.get(functionName);
152+
}
123153
if (implementList == null || implementList.isEmpty()) {
124154
throw new IllegalStateException(String.format("Cannot resolve function: %s", functionName));
125155
}
@@ -401,17 +431,17 @@ void populate() {
401431
}
402432

403433
private static class Builder extends AbstractBuilder {
404-
private final Map<BuiltinFunctionName, PairList<CalciteFuncSignature, FunctionImp>> map =
434+
private final Map<BuiltinFunctionName, List<Pair<CalciteFuncSignature, FunctionImp>>> map =
405435
new HashMap<>();
406436

407437
@Override
408438
void register(BuiltinFunctionName functionName, FunctionImp implement) {
409439
CalciteFuncSignature signature =
410440
new CalciteFuncSignature(functionName.getName(), implement.getParams());
411441
if (map.containsKey(functionName)) {
412-
map.get(functionName).add(signature, implement);
442+
map.get(functionName).add(Pair.of(signature, implement));
413443
} else {
414-
map.put(functionName, PairList.of(signature, implement));
444+
map.put(functionName, new ArrayList<>(List.of(Pair.of(signature, implement))));
415445
}
416446
}
417447
}

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@ public class CalciteGeoIpFunctionsIT extends GeoIpFunctionsIT {
1212
public void init() throws Exception {
1313
super.init();
1414
enableCalcite();
15-
// TODO: "https://github.com/opensearch-project/sql/issues/3506"
16-
// disallowCalciteFallback();
15+
disallowCalciteFallback();
1716
}
1817
}

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

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@
1818
import java.util.List;
1919
import java.util.Map;
2020
import java.util.concurrent.atomic.AtomicReference;
21-
import lombok.RequiredArgsConstructor;
2221
import org.apache.calcite.plan.RelOptUtil;
2322
import org.apache.calcite.rel.RelNode;
2423
import org.apache.calcite.rel.RelRoot;
@@ -38,21 +37,33 @@
3837
import org.opensearch.sql.executor.ExecutionEngine.Schema.Column;
3938
import org.opensearch.sql.executor.Explain;
4039
import org.opensearch.sql.executor.pagination.PlanSerializer;
40+
import org.opensearch.sql.expression.function.BuiltinFunctionName;
41+
import org.opensearch.sql.expression.function.PPLFuncImpTable;
4142
import org.opensearch.sql.opensearch.client.OpenSearchClient;
4243
import org.opensearch.sql.opensearch.executor.protector.ExecutionProtector;
44+
import org.opensearch.sql.opensearch.functions.GeoIpFunction;
4345
import org.opensearch.sql.opensearch.util.JdbcOpenSearchDataTypeConvertor;
4446
import org.opensearch.sql.planner.physical.PhysicalPlan;
4547
import org.opensearch.sql.storage.TableScanOperator;
4648

4749
/** OpenSearch execution engine implementation. */
48-
@RequiredArgsConstructor
4950
public class OpenSearchExecutionEngine implements ExecutionEngine {
5051

5152
private final OpenSearchClient client;
5253

5354
private final ExecutionProtector executionProtector;
5455
private final PlanSerializer planSerializer;
5556

57+
public OpenSearchExecutionEngine(
58+
OpenSearchClient client,
59+
ExecutionProtector executionProtector,
60+
PlanSerializer planSerializer) {
61+
this.client = client;
62+
this.executionProtector = executionProtector;
63+
this.planSerializer = planSerializer;
64+
registerOpenSearchFunctions();
65+
}
66+
5667
@Override
5768
public void execute(PhysicalPlan physicalPlan, ResponseListener<QueryResponse> listener) {
5869
execute(physicalPlan, ExecutionContext.emptyExecutionContext(), listener);
@@ -228,4 +239,12 @@ private void buildResultSet(
228239
QueryResponse response = new QueryResponse(schema, values, null);
229240
listener.onResponse(response);
230241
}
242+
243+
/** Registers opensearch-dependent functions */
244+
private void registerOpenSearchFunctions() {
245+
PPLFuncImpTable.FunctionImp geoIpImpl =
246+
(builder, args) ->
247+
builder.makeCall(new GeoIpFunction(client.getNodeClient()).toUDF("GEOIP"), args);
248+
PPLFuncImpTable.INSTANCE.registerExternalFunction(BuiltinFunctionName.GEOIP, geoIpImpl);
249+
}
231250
}

‎opensearch/src/main/java/org/opensearch/sql/opensearch/functions/GeoIpFunction.java‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.opensearch.functions;
7+
8+
import java.util.*;
9+
import java.util.stream.Collectors;
10+
import lombok.Getter;
11+
import org.apache.calcite.adapter.enumerable.NotNullImplementor;
12+
import org.apache.calcite.adapter.enumerable.NullPolicy;
13+
import org.apache.calcite.adapter.enumerable.RexToLixTranslator;
14+
import org.apache.calcite.linq4j.tree.Expression;
15+
import org.apache.calcite.linq4j.tree.Expressions;
16+
import org.apache.calcite.rel.type.RelDataType;
17+
import org.apache.calcite.rel.type.RelDataTypeFactory;
18+
import org.apache.calcite.rex.RexCall;
19+
import org.apache.calcite.sql.type.SqlReturnTypeInference;
20+
import org.apache.calcite.sql.type.SqlTypeName;
21+
import org.opensearch.geospatial.action.IpEnrichmentActionClient;
22+
import org.opensearch.sql.common.utils.StringUtils;
23+
import org.opensearch.sql.data.model.ExprStringValue;
24+
import org.opensearch.sql.data.model.ExprTupleValue;
25+
import org.opensearch.sql.data.model.ExprValue;
26+
import org.opensearch.sql.expression.function.ImplementorUDF;
27+
import org.opensearch.transport.client.node.NodeClient;
28+
29+
/**
30+
* {@code GEOIP(dataSourceName, ipAddress[, options])} looks up location information from given IP
31+
* addresses via OpenSearch GeoSpatial plugin API. The options is a comma-separated list of fields
32+
* to be returned. If not specified, all fields are returned.
33+
*
34+
* <p>Signatures:
35+
*
36+
* <ul>
37+
* <li>(STRING, STRING) -> MAP
38+
* <li>(STRING, STRING, STRING) -> MAP
39+
* </ul>
40+
*/
41+
public class GeoIpFunction extends ImplementorUDF {
42+
public GeoIpFunction(NodeClient nodeClient) {
43+
super(new GeoIPImplementor(nodeClient), NullPolicy.ANY);
44+
}
45+
46+
@Override
47+
public SqlReturnTypeInference getReturnTypeInference() {
48+
return op -> {
49+
RelDataTypeFactory typeFactory = op.getTypeFactory();
50+
RelDataType varcharType = typeFactory.createSqlType(SqlTypeName.VARCHAR);
51+
RelDataType anyType = typeFactory.createSqlType(SqlTypeName.ANY);
52+
return typeFactory.createMapType(varcharType, anyType);
53+
};
54+
}
55+
56+
public static class GeoIPImplementor implements NotNullImplementor {
57+
@Getter private static NodeClient nodeClient;
58+
59+
public GeoIPImplementor(NodeClient nodeClient) {
60+
GeoIPImplementor.nodeClient = nodeClient;
61+
}
62+
63+
@Override
64+
public Expression implement(
65+
RexToLixTranslator translator, RexCall call, List<Expression> translatedOperands) {
66+
if (getNodeClient() == null) {
67+
throw new IllegalStateException("nodeClient is null.");
68+
}
69+
List<Expression> operandsWithClient = new ArrayList<>(translatedOperands);
70+
// Since a NodeClient cannot be passed as a parameter using Expressions.constant,
71+
// it is instead provided through a function call.
72+
operandsWithClient.add(Expressions.call(GeoIPImplementor.class, "getNodeClient"));
73+
return Expressions.call(GeoIPImplementor.class, "fetchIpEnrichment", operandsWithClient);
74+
}
75+
76+
public static Map<String, ?> fetchIpEnrichment(
77+
String dataSource, String ipAddress, NodeClient nodeClient) {
78+
return fetchIpEnrichment(dataSource, ipAddress, Collections.emptySet(), nodeClient);
79+
}
80+
81+
public static Map<String, ?> fetchIpEnrichment(
82+
String dataSource, String ipAddress, String commaSeparatedOptions, NodeClient nodeClient) {
83+
String unquotedOptions = StringUtils.unquoteText(commaSeparatedOptions);
84+
final Set<String> options =
85+
Arrays.stream(unquotedOptions.split(",")).map(String::trim).collect(Collectors.toSet());
86+
return fetchIpEnrichment(dataSource, ipAddress, options, nodeClient);
87+
}
88+
89+
private static Map<String, ?> fetchIpEnrichment(
90+
String dataSource, String ipAddress, Set<String> options, NodeClient nodeClient) {
91+
IpEnrichmentActionClient ipClient = new IpEnrichmentActionClient(nodeClient);
92+
dataSource = StringUtils.unquoteText(dataSource);
93+
try {
94+
Map<String, Object> geoLocationData = ipClient.getGeoLocationData(ipAddress, dataSource);
95+
Map<String, ExprValue> enrichmentResult =
96+
geoLocationData.entrySet().stream()
97+
.filter(entry -> options.isEmpty() || options.contains(entry.getKey()))
98+
.collect(
99+
Collectors.toMap(
100+
Map.Entry::getKey, v -> new ExprStringValue(v.getValue().toString())));
101+
@SuppressWarnings("unchecked")
102+
Map<String, ?> result =
103+
(Map<String, ?>) ExprTupleValue.fromExprValueMap(enrichmentResult).valueForCalcite();
104+
return result;
105+
} catch (Exception e) {
106+
throw new RuntimeException(e);
107+
}
108+
}
109+
}
110+
}

‎ppl/src/main/java/org/opensearch/sql/ppl/parser/AstExpressionBuilder.java‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -233,7 +233,7 @@ public UnresolvedExpression visitTakeAggFunctionCall(
233233
/** Eval function. */
234234
@Override
235235
public UnresolvedExpression visitBooleanFunctionCall(BooleanFunctionCallContext ctx) {
236-
final String functionName = ctx.conditionFunctionName().getText().toLowerCase();
236+
final String functionName = ctx.conditionFunctionName().getText().toLowerCase(Locale.ROOT);
237237
return buildFunction(
238238
FUNCTION_NAME_MAPPING.getOrDefault(functionName, functionName),
239239
ctx.functionArgs().functionArg());
@@ -287,15 +287,15 @@ private Function buildFunction(
287287
public UnresolvedExpression visitSingleFieldRelevanceFunction(
288288
SingleFieldRelevanceFunctionContext ctx) {
289289
return new Function(
290-
ctx.singleFieldRelevanceFunctionName().getText().toLowerCase(),
290+
ctx.singleFieldRelevanceFunctionName().getText().toLowerCase(Locale.ROOT),
291291
singleFieldRelevanceArguments(ctx));
292292
}
293293

294294
@Override
295295
public UnresolvedExpression visitMultiFieldRelevanceFunction(
296296
MultiFieldRelevanceFunctionContext ctx) {
297297
return new Function(
298-
ctx.multiFieldRelevanceFunctionName().getText().toLowerCase(),
298+
ctx.multiFieldRelevanceFunctionName().getText().toLowerCase(Locale.ROOT),
299299
multiFieldRelevanceArguments(ctx));
300300
}
301301

@@ -506,7 +506,7 @@ private List<UnresolvedExpression> singleFieldRelevanceArguments(
506506
v ->
507507
builder.add(
508508
new UnresolvedArgument(
509-
v.relevanceArgName().getText().toLowerCase(),
509+
v.relevanceArgName().getText().toLowerCase(Locale.ROOT),
510510
new Literal(
511511
StringUtils.unquoteText(v.relevanceArgValue().getText()),
512512
DataType.STRING))));
@@ -534,7 +534,7 @@ private List<UnresolvedExpression> multiFieldRelevanceArguments(
534534
v ->
535535
builder.add(
536536
new UnresolvedArgument(
537-
v.relevanceArgName().getText().toLowerCase(),
537+
v.relevanceArgName().getText().toLowerCase(Locale.ROOT),
538538
new Literal(
539539
StringUtils.unquoteText(v.relevanceArgValue().getText()),
540540
DataType.STRING))));

0 commit comments

Comments
 (0)