Skip to content

Support IP-typed fields as arguments of GEOIP function with Calcite #4613

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yuancu merged 1 commit into from
Oct 21, 2025
Merged

Support IP-typed fields as arguments of GEOIP function with Calcite #4613

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion docs/user/ppl/functions/ip.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ Description

Usage: `geoip(dataSourceName, ipAddress[, options])` to lookup location information from given IP addresses via OpenSearch GeoSpatial plugin API.

Argument type: STRING, STRING, STRING
Argument type: STRING, STRING/IP, STRING

Return type: OBJECT

Expand Down
40 changes: 40 additions & 0 deletions integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteGeoIpFunctionsIT.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,12 +5,52 @@

package org.opensearch.sql.calcite.remote;

import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_WEBLOGS;
import static org.opensearch.sql.util.MatcherUtils.rows;
import static org.opensearch.sql.util.MatcherUtils.schema;
import static org.opensearch.sql.util.MatcherUtils.verifyDataRows;
import static org.opensearch.sql.util.MatcherUtils.verifySchema;

import java.io.IOException;
import java.util.Map;
import org.json.JSONObject;
import org.junit.jupiter.api.Test;
import org.opensearch.client.Request;
import org.opensearch.sql.ppl.GeoIpFunctionsIT;

public class CalciteGeoIpFunctionsIT extends GeoIpFunctionsIT {
@Override
public void init() throws Exception {
super.init();
loadIndex(Index.WEBLOG);
enableCalcite();

// Only limited IPs are loaded into geospatial data sources. Therefore, we insert IPs that match
// those known ones for test purpose
Request bulkRequest = new Request("POST", "/_bulk?refresh=true");
bulkRequest.setJsonEntity(
String.format(
"{\"index\":{\"_index\":\"%s\",\"_id\":6}}\n"
+ "{\"host\":\"10.0.0.1\",\"method\":\"POST\"}\n"
+ "{\"index\":{\"_index\":\"%s\",\"_id\":7}}\n"
+ "{\"host\":\"fd12:2345:6789:1:a1b2:c3d4:e5f6:789a\",\"method\":\"POST\"}\n",
TEST_INDEX_WEBLOGS, TEST_INDEX_WEBLOGS));
client().performRequest(bulkRequest);
}

// In v2 it supports only string as IP inputs
@Test
public void testGeoIpEnrichmentWithIpFieldAsInput() throws IOException {
JSONObject result =
executeQuery(
String.format(
"source=%s | where method='POST' | eval ip_to_country = geoip('%s', host,"
+ " 'country') | fields host, ip_to_country",
TEST_INDEX_WEBLOGS, DATASOURCE_NAME));
verifySchema(result, schema("host", "ip"), schema("ip_to_country", "struct"));
verifyDataRows(
result,
rows("10.0.0.1", Map.of("country", "USA")),
rows("fd12:2345:6789:1:a1b2:c3d4:e5f6:789a", Map.of("country", "India")));
}
}
8 changes: 4 additions & 4 deletions integ-test/src/test/java/org/opensearch/sql/ppl/GeoIpFunctionsIT.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ public class GeoIpFunctionsIT extends PPLIntegTestCase {
"endpoint",
"https://raw.githubusercontent.com/opensearch-project/geospatial/main/src/test/resources/ip2geo/server/city/manifest.json");

private static String DATASOURCE_NAME = "dummycityindex";
protected static String DATASOURCE_NAME = "dummycityindex";

private static String PLUGIN_NAME = "opensearch-geospatial";

Expand Down Expand Up @@ -83,7 +83,7 @@ public void testGeoIpEnrichment() {
String.format(
"search source=%s | eval enrichmentResult = geoip(\\\"%s\\\",%s) | fields name, ip,"
+ " enrichmentResult",
TEST_INDEX_GEOIP, "dummycityindex", "ip"));
TEST_INDEX_GEOIP, DATASOURCE_NAME, "ip"));

verifyColumn(resultGeoIp, columnName("name"), columnName("ip"), columnName("enrichmentResult"));
verifyDataRows(
Expand All @@ -101,7 +101,7 @@ public void testGeoIpEnrichmentWithSingleOption() {
String.format(
"search source=%s | eval enrichmentResult = geoip(\\\"%s\\\",%s,\\\"%s\\\") |"
+ " fields name, ip, enrichmentResult",
TEST_INDEX_GEOIP, "dummycityindex", "ip", "city"));
TEST_INDEX_GEOIP, DATASOURCE_NAME, "ip", "city"));

verifyColumn(resultGeoIp, columnName("name"), columnName("ip"), columnName("enrichmentResult"));
verifyDataRows(
Expand All @@ -119,7 +119,7 @@ public void testGeoIpEnrichmentWithSpaceSeparatedMultipleOptions() {
String.format(
"search source=%s | eval enrichmentResult = geoip(\\\"%s\\\",%s,\\\"%s\\\") |"
+ " fields name, ip, enrichmentResult",
TEST_INDEX_GEOIP, "dummycityindex", "ip", "city , country"));
TEST_INDEX_GEOIP, DATASOURCE_NAME, "ip", "city , country"));

verifyColumn(resultGeoIp, columnName("name"), columnName("ip"), columnName("enrichmentResult"));
verifyDataRows(
Expand Down
30 changes: 16 additions & 14 deletions opensearch/src/main/java/org/opensearch/sql/opensearch/functions/GeoIpFunction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,16 +16,15 @@
import org.apache.calcite.rel.type.RelDataType;
import org.apache.calcite.rel.type.RelDataTypeFactory;
import org.apache.calcite.rex.RexCall;
import org.apache.calcite.sql.type.CompositeOperandTypeChecker;
import org.apache.calcite.sql.type.OperandTypes;
import org.apache.calcite.sql.type.SqlReturnTypeInference;
import org.apache.calcite.sql.type.SqlTypeFamily;
import org.apache.calcite.sql.type.SqlTypeName;
import org.opensearch.geospatial.action.IpEnrichmentActionClient;
import org.opensearch.sql.common.utils.StringUtils;
import org.opensearch.sql.data.model.ExprIpValue;
import org.opensearch.sql.data.model.ExprStringValue;
import org.opensearch.sql.data.model.ExprTupleValue;
import org.opensearch.sql.data.model.ExprValue;
import org.opensearch.sql.data.type.ExprCoreType;
import org.opensearch.sql.expression.function.ImplementorUDF;
import org.opensearch.sql.expression.function.UDFOperandMetadata;
import org.opensearch.transport.client.node.NodeClient;
Expand All @@ -38,8 +37,8 @@
* <p>Signatures:
*
* <ul>
* <li>(STRING, STRING) -> MAP
* <li>(STRING, STRING, STRING) -> MAP
* <li>(STRING, IP) -> MAP
* <li>(STRING, IP, STRING) -> MAP
Comment on lines +40 to +41
Copy link
Member

@LantaoJin LantaoJin Oct 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should be following ? Seems not a breaking change from the ip.rst, but why the description here is different.

 *.  <li>(STRING, STRING) -> MAP
 *   <li>(STRING, STRING, STRING) -> MAP
 *   <li>(STRING, IP) -> MAP
 *   <li>(STRING, IP, STRING) -> MAP

Copy link
Collaborator Author

@yuancu yuancu Oct 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Now I modified the underlying function to accept IP only.

It still accepts strings as inputs with implicit type coercion.

If it gets string as input, it will try to cast it to IP first. This guarantees that the input is a valid IP address.

* </ul>
*/
public class GeoIpFunction extends ImplementorUDF {
Expand All @@ -59,11 +58,10 @@ public SqlReturnTypeInference getReturnTypeInference() {

@Override
public UDFOperandMetadata getOperandMetadata() {
return UDFOperandMetadata.wrap(
(CompositeOperandTypeChecker)
OperandTypes.CHARACTER_CHARACTER.or(
OperandTypes.family(
SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER)));
return UDFOperandMetadata.wrapUDT(
List.of(
List.of(ExprCoreType.STRING, ExprCoreType.IP),
List.of(ExprCoreType.STRING, ExprCoreType.IP, ExprCoreType.STRING)));
}

public static class GeoIPImplementor implements NotNullImplementor {
Expand All @@ -87,16 +85,20 @@ public Expression implement(
}

public static Map<String, ?> fetchIpEnrichment(
String dataSource, String ipAddress, NodeClient nodeClient) {
return fetchIpEnrichment(dataSource, ipAddress, Collections.emptySet(), nodeClient);
String dataSource, ExprIpValue ipAddress, NodeClient nodeClient) {
return fetchIpEnrichment(
dataSource, ipAddress.toString(), Collections.emptySet(), nodeClient);
}

public static Map<String, ?> fetchIpEnrichment(
String dataSource, String ipAddress, String commaSeparatedOptions, NodeClient nodeClient) {
String dataSource,
ExprIpValue ipAddress,
String commaSeparatedOptions,
NodeClient nodeClient) {
String unquotedOptions = StringUtils.unquoteText(commaSeparatedOptions);
final Set<String> options =
Arrays.stream(unquotedOptions.split(",")).map(String::trim).collect(Collectors.toSet());
return fetchIpEnrichment(dataSource, ipAddress, options, nodeClient);
return fetchIpEnrichment(dataSource, ipAddress.toString(), options, nodeClient);
}

private static Map<String, ?> fetchIpEnrichment(
Expand Down
Loading