diff --git a/docs/user/ppl/functions/ip.rst b/docs/user/ppl/functions/ip.rst
index 897d6ccfad0..ec853c27093 100644
--- a/docs/user/ppl/functions/ip.rst
+++ b/docs/user/ppl/functions/ip.rst
@@ -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
diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java
index b2991b92017..9f2970c1c22 100644
--- a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java
+++ b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java
@@ -1091,9 +1091,9 @@ public void testExplainPushDownScriptsContainingUDT() throws IOException {
"source=%s | where cidrmatch(host, '0.0.0.0/24') | fields host",
TEST_INDEX_WEBLOGS)));
- assertYamlEqualsJsonIgnoreId(
+ assertYamlEqualsIgnoreId(
loadExpectedPlan("explain_agg_script_timestamp_push.yaml"),
- explainQueryToString(
+ explainQueryYaml(
String.format(
"source=%s | eval t = unix_timestamp(birthdate) | stats count() by t | sort t |"
+ " head 3",
diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteGeoIpFunctionsIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteGeoIpFunctionsIT.java
index f5bc073cbe7..9047b497b02 100644
--- a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteGeoIpFunctionsIT.java
+++ b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteGeoIpFunctionsIT.java
@@ -5,14 +5,54 @@
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;
import java.io.IOException;
public class CalciteGeoIpFunctionsIT extends GeoIpFunctionsIT {
- @Override
- public void init() throws IOException {
- super.init();
- enableCalcite();
- }
+ @Override
+ public void init() throws IOException {
+ 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")));
+ }
}
\ No newline at end of file
diff --git a/integ-test/src/test/java/org/opensearch/sql/ppl/GeoIpFunctionsIT.java b/integ-test/src/test/java/org/opensearch/sql/ppl/GeoIpFunctionsIT.java
index b9a72c5c9f0..6bee5916393 100644
--- a/integ-test/src/test/java/org/opensearch/sql/ppl/GeoIpFunctionsIT.java
+++ b/integ-test/src/test/java/org/opensearch/sql/ppl/GeoIpFunctionsIT.java
@@ -42,7 +42,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";
@@ -82,8 +82,9 @@ public void testGeoIpEnrichment() {
JSONObject resultGeoIp =
executeQuery(
String.format(
- "search source=%s | eval enrichmentResult = geoip(\\\"%s\\\",%s)",
- TEST_INDEX_GEOIP, "dummycityindex", "ip"));
+ "search source=%s | eval enrichmentResult = geoip(\\\"%s\\\",%s) | fields name, ip,"
+ + " enrichmentResult",
+ TEST_INDEX_GEOIP, DATASOURCE_NAME, "ip"));
verifyColumn(resultGeoIp, columnName("name"), columnName("ip"), columnName("enrichmentResult"));
verifyDataRows(
@@ -99,8 +100,9 @@ public void testGeoIpEnrichmentWithSingleOption() {
JSONObject resultGeoIp =
executeQuery(
String.format(
- "search source=%s | eval enrichmentResult = geoip(\\\"%s\\\",%s,\\\"%s\\\")",
- TEST_INDEX_GEOIP, "dummycityindex", "ip", "city"));
+ "search source=%s | eval enrichmentResult = geoip(\\\"%s\\\",%s,\\\"%s\\\") |"
+ + " fields name, ip, enrichmentResult",
+ TEST_INDEX_GEOIP, DATASOURCE_NAME, "ip", "city"));
verifyColumn(resultGeoIp, columnName("name"), columnName("ip"), columnName("enrichmentResult"));
verifyDataRows(
@@ -116,8 +118,9 @@ public void testGeoIpEnrichmentWithSpaceSeparatedMultipleOptions() {
JSONObject resultGeoIp =
executeQuery(
String.format(
- "search source=%s | eval enrichmentResult = geoip(\\\"%s\\\",%s,\\\"%s\\\")",
- TEST_INDEX_GEOIP, "dummycityindex", "ip", "city , country"));
+ "search source=%s | eval enrichmentResult = geoip(\\\"%s\\\",%s,\\\"%s\\\") |"
+ + " fields name, ip, enrichmentResult",
+ TEST_INDEX_GEOIP, DATASOURCE_NAME, "ip", "city , country"));
verifyColumn(resultGeoIp, columnName("name"), columnName("ip"), columnName("enrichmentResult"));
verifyDataRows(
diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/functions/GeoIpFunction.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/functions/GeoIpFunction.java
index 245b5b84693..77d99507adf 100644
--- a/opensearch/src/main/java/org/opensearch/sql/opensearch/functions/GeoIpFunction.java
+++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/functions/GeoIpFunction.java
@@ -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.client.node.NodeClient;
import org.opensearch.sql.expression.function.UDFOperandMetadata;
@@ -38,8 +37,8 @@
* Signatures:
*
*
- * - (STRING, STRING) -> MAP
- *
- (STRING, STRING, STRING) -> MAP
+ *
- (STRING, IP) -> MAP
+ *
- (STRING, IP, STRING) -> MAP
*
*/
public class GeoIpFunction extends ImplementorUDF {
@@ -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 {
@@ -87,16 +85,20 @@ public Expression implement(
}
public static Map 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 fetchIpEnrichment(
- String dataSource, String ipAddress, String commaSeparatedOptions, NodeClient nodeClient) {
+ String dataSource,
+ ExprIpValue ipAddress,
+ String commaSeparatedOptions,
+ NodeClient nodeClient) {
String unquotedOptions = StringUtils.unquoteText(commaSeparatedOptions);
final Set 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 fetchIpEnrichment(