Update GEOIP function to support IP types as input (#4613)

Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>

Author: yuancu

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
 

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

@@ -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")));
   }
 }

integ-test/src/test/java/org/opensearch/sql/ppl/GeoIpFunctionsIT.java

@@ -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";
 
@@ -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(
@@ -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(
@@ -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(

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.sql.expression.function.UDFOperandMetadata;
 import org.opensearch.transport.client.node.NodeClient;
@@ -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
  * </ul>
  */
 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<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(