Merge branch 'main' into feature/json-extract

Signed-off-by: kenrickyap <[email protected]>

core/src/main/java/org/opensearch/sql/data/type/ExprType.java

@@ -9,6 +9,7 @@
 
 import java.util.Arrays;
 import java.util.List;
+import java.util.Optional;
 import org.opensearch.sql.data.model.ExprValue;
 import org.opensearch.sql.expression.Expression;
 
@@ -54,4 +55,17 @@ default List<ExprType> getParent() {
   default String legacyTypeName() {
     return typeName();
   }
+
+  /** Get the original path. Types like alias type will set the actual path in field property. */
+  default Optional<String> getOriginalPath() {
+    return Optional.empty();
+  }
+
+  /**
+   * Get the original path. Types like alias type should be derived from the type of the original
+   * field.
+   */
+  default ExprType getOriginalExprType() {
+    return this;
+  }
 }

core/src/main/java/org/opensearch/sql/expression/DSL.java

@@ -985,6 +985,10 @@ public static FunctionExpression utc_timestamp(
     return compile(functionProperties, BuiltinFunctionName.UTC_TIMESTAMP, args);
   }
 
+  public static FunctionExpression geoip(Expression... args) {
+    return compile(FunctionProperties.None, BuiltinFunctionName.GEOIP, args);
+  }
+
   @SuppressWarnings("unchecked")
   private static <T extends FunctionImplementation> T compile(
       FunctionProperties functionProperties, BuiltinFunctionName bfn, Expression... args) {

core/src/main/java/org/opensearch/sql/expression/ReferenceExpression.java

@@ -23,6 +23,8 @@
 public class ReferenceExpression implements Expression {
   @Getter private final String attr;
 
+  @Getter private final String rawPath;
+
   @Getter private final List<String> paths;
 
   private final ExprType type;
@@ -36,8 +38,11 @@ public class ReferenceExpression implements Expression {
   public ReferenceExpression(String ref, ExprType type) {
     this.attr = ref;
     // Todo. the define of paths need to be redefined after adding multiple index/variable support.
-    this.paths = Arrays.asList(ref.split("\\."));
-    this.type = type;
+    // For AliasType, the actual path is set in the property of `path` and the type is derived
+    // from the type of field on that path; Otherwise, use ref itself as the path
+    this.rawPath = type.getOriginalPath().orElse(ref);
+    this.paths = Arrays.asList(rawPath.split("\\."));
+    this.type = type.getOriginalExprType();
   }
 
   @Override

core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java

@@ -209,6 +209,9 @@ public enum BuiltinFunctionName {
   JSON(FunctionName.of("json")),
   JSON_EXTRACT(FunctionName.of("json_extract")),
 
+  /** GEOSPATIAL Functions. */
+  GEOIP(FunctionName.of("geoip")),
+
   /** NULL Test. */
   IS_NULL(FunctionName.of("is null")),
   IS_NOT_NULL(FunctionName.of("is not null")),

core/src/main/java/org/opensearch/sql/expression/function/OpenSearchFunctions.java

@@ -173,4 +173,54 @@ public String toString() {
       return String.format("%s(%s)", functionName, String.join(", ", args));
     }
   }
+
+  /**
+   * Static class to identify functional Expression which specifically designed for OpenSearch
+   * storage runtime
+   */
+  public static class OpenSearchExecutableFunction extends FunctionExpression {
+    private final FunctionName functionName;
+    private final List<Expression> arguments;
+    private final ExprType returnType;
+
+    public OpenSearchExecutableFunction(
+        FunctionName functionName, List<Expression> arguments, ExprType returnType) {
+      super(functionName, arguments);
+      this.functionName = functionName;
+      this.arguments = arguments;
+      this.returnType = returnType;
+    }
+
+    @Override
+    public ExprValue valueOf(Environment<Expression, ExprValue> valueEnv) {
+      throw new UnsupportedOperationException(
+          String.format(
+              "OpenSearch defined function [%s] is only supported in Eval operation.",
+              functionName));
+    }
+
+    @Override
+    public ExprType type() {
+      return returnType;
+    }
+
+    /**
+     * Util method to generate probe implementation with given list of argument types, with marker
+     * class `OpenSearchFunction` to annotate this is an OpenSearch specific expression.
+     *
+     * @param returnType return type.
+     * @return Binary Function Implementation.
+     */
+    public static SerializableFunction<FunctionName, Pair<FunctionSignature, FunctionBuilder>>
+        openSearchImpl(ExprType returnType, List<ExprType> args) {
+      return functionName -> {
+        FunctionSignature functionSignature = new FunctionSignature(functionName, args);
+        FunctionBuilder functionBuilder =
+            (functionProperties, arguments) ->
+                new OpenSearchFunctions.OpenSearchExecutableFunction(
+                    functionName, arguments, returnType);
+        return Pair.of(functionSignature, functionBuilder);
+      };
+    }
+  }
 }

core/src/main/java/org/opensearch/sql/expression/ip/IPFunctions.java

@@ -11,8 +11,10 @@
 import static org.opensearch.sql.expression.function.FunctionDSL.define;
 import static org.opensearch.sql.expression.function.FunctionDSL.impl;
 import static org.opensearch.sql.expression.function.FunctionDSL.nullMissingHandling;
+import static org.opensearch.sql.expression.function.OpenSearchFunctions.OpenSearchExecutableFunction.openSearchImpl;
 
 import inet.ipaddr.IPAddress;
+import java.util.Arrays;
 import lombok.experimental.UtilityClass;
 import org.opensearch.sql.data.model.ExprValue;
 import org.opensearch.sql.data.model.ExprValueUtils;
@@ -28,6 +30,7 @@ public class IPFunctions {
 
   public void register(BuiltinFunctionRepository repository) {
     repository.register(cidrmatch());
+    repository.register(geoIp());
   }
 
   private DefaultFunctionResolver cidrmatch() {
@@ -57,4 +60,16 @@ private ExprValue exprCidrMatch(ExprValue addressExprValue, ExprValue rangeExprV
         ? ExprValueUtils.LITERAL_FALSE
         : ExprValueUtils.LITERAL_TRUE;
   }
+
+  /**
+   * To register all method signatures related to geoip( ) expression under eval.
+   *
+   * @return Resolver for geoip( ) expression.
+   */
+  private DefaultFunctionResolver geoIp() {
+    return define(
+        BuiltinFunctionName.GEOIP.getName(),
+        openSearchImpl(BOOLEAN, Arrays.asList(STRING, STRING)),
+        openSearchImpl(BOOLEAN, Arrays.asList(STRING, STRING, STRING)));
+  }
 }

core/src/main/java/org/opensearch/sql/planner/physical/EvalOperator.java

@@ -87,7 +87,7 @@ public ExprValue next() {
    * @param env {@link Environment}
    * @return The mapping of reference and {@link ExprValue} for each expression.
    */
-  private Map<String, ExprValue> eval(Environment<Expression, ExprValue> env) {
+  protected Map<String, ExprValue> eval(Environment<Expression, ExprValue> env) {
     Map<String, ExprValue> evalResultMap = new LinkedHashMap<>();
     for (Pair<ReferenceExpression, Expression> pair : expressionList) {
       ReferenceExpression var = pair.getKey();

core/src/test/java/org/opensearch/sql/data/type/ExprTypeTest.java

@@ -24,6 +24,7 @@
 import static org.opensearch.sql.data.type.ExprCoreType.UNDEFINED;
 import static org.opensearch.sql.data.type.ExprCoreType.UNKNOWN;
 
+import java.util.Optional;
 import org.hamcrest.Matchers;
 import org.junit.jupiter.api.Test;
 
@@ -85,4 +86,16 @@ void defaultLegacyTypeName() {
     final ExprType exprType = () -> "dummy";
     assertEquals("dummy", exprType.legacyTypeName());
   }
+
+  @Test
+  void getOriginalPath() {
+    final ExprType exprType = () -> "dummy";
+    assertEquals(Optional.empty(), exprType.getOriginalPath());
+  }
+
+  @Test
+  void getOriginalExprType() {
+    final ExprType exprType = () -> "dummy";
+    assertEquals(exprType, exprType.getOriginalExprType());
+  }
 }

core/src/test/java/org/opensearch/sql/expression/function/OpenSearchFunctionsTest.java

@@ -22,6 +22,7 @@
 import org.opensearch.sql.expression.ExpressionTestBase;
 import org.opensearch.sql.expression.FunctionExpression;
 import org.opensearch.sql.expression.NamedArgumentExpression;
+import org.opensearch.sql.expression.ReferenceExpression;
 import org.opensearch.sql.expression.env.Environment;
 
 public class OpenSearchFunctionsTest extends ExpressionTestBase {
@@ -309,4 +310,17 @@ void nested_query() {
     assertEquals(expr.valueOf(nestedTuple), ExprValueUtils.stringValue("result"));
     assertEquals(expr.type(), STRING);
   }
+
+  @Test
+  void opensearchExecutableFunction_valueOf() {
+    var ipInStr =
+        new OpenSearchFunctions.OpenSearchExecutableFunction(
+            BuiltinFunctionName.GEOIP.getName(),
+            List.of(DSL.literal("my-datasource"), new ReferenceExpression("ipInStr", STRING)),
+            BOOLEAN);
+    assertThrows(
+        UnsupportedOperationException.class,
+        () -> ipInStr.valueOf(valueEnv()),
+        "OpenSearch defined function [geoip] is only supported in Eval operation.");
+  }
 }

core/src/test/java/org/opensearch/sql/expression/ip/GeoIPFunctionTest.java

@@ -0,0 +1,49 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.expression.ip;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.opensearch.sql.data.type.ExprCoreType.BOOLEAN;
+import static org.opensearch.sql.data.type.ExprCoreType.STRING;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+import org.opensearch.sql.data.model.ExprValue;
+import org.opensearch.sql.expression.DSL;
+import org.opensearch.sql.expression.Expression;
+import org.opensearch.sql.expression.env.Environment;
+
+@ExtendWith(MockitoExtension.class)
+public class GeoIPFunctionTest {
+
+  // Mock value environment for testing.
+  @Mock private Environment<Expression, ExprValue> env;
+
+  @Test
+  public void testGeoipFunctionSignature() {
+    var geoip = DSL.geoip(DSL.literal("HARDCODED_DATASOURCE_NAME"), DSL.ref("ip_address", STRING));
+    assertEquals(BOOLEAN, geoip.type());
+  }
+
+  /** To make sure no logic being evaluated when no environment being passed. */
+  @Test
+  public void testDefaultValueOf() {
+    UnsupportedOperationException exception =
+        assertThrows(
+            UnsupportedOperationException.class,
+            () ->
+                DSL.geoip(DSL.literal("HARDCODED_DATASOURCE_NAME"), DSL.ref("ip_address", STRING))
+                    .valueOf(env));
+    assertTrue(
+        exception
+            .getMessage()
+            .contains("OpenSearch defined function [geoip] is only supported in Eval operation."));
+  }
+}

docs/dev/datasource-prometheus.md

@@ -175,7 +175,7 @@ Since we don't have column names similar to relational databases, `@timestamp` a
 
 ### Passthrough PromQL Support.
 * PromQL is easy to use powerful language over metrics and we want to support promQL queries from catalogs with prometheus connectors.
-* Function signature: source = ``promcatalog.nativeQuery(`promQLCommand`, startTime = "{{startTime}}", endTime="{{endTime}", step="{{resolution}}")``
+* Function signature: source = ``promcatalog.nativeQuery(`promQLCommand`, startTime = "{{startTime}}", endTime="{{endTime}}", step="{{resolution}}")``
 
 
 

docs/user/ppl/functions/ip.rst

@@ -36,3 +36,34 @@ Note:
  - `cidr` can be an IPv4 or IPv6 block
  - `ip` and `cidr` must both be valid and non-missing/non-null
 
+
+GEOIP
+---------
+
+Description
+>>>>>>>>>>>
+
+Usage: `geoip(dataSourceName, ipAddress[, options])` to lookup location information from given IP addresses via OpenSearch GeoSpatial plugin API.
+
+Argument type: STRING, STRING, STRING
+
+Return type: OBJECT
+
+.. The execution of below example is being excluded, as this requires a standalone Geo-Spatial dataSource setup, which is not yet supported by docTest.
+
+Example:
+
+    > source=weblogs | eval LookupResult = geoip("dataSourceName", "50.68.18.229", "country_iso_code,city_name")
+    fetched rows / total rows = 1/1
+    +-------------------------------------------------------------+
+    | LookupResult                                                |
+    |-------------------------------------------------------------|
+    | {'city_name': 'Vancouver', 'country_iso_code': 'CA'}        |
+    +-------------------------------------------------------------+
+
+
+Note:
+ - `dataSourceName` must be an established dataSource on OpenSearch GeoSpatial plugin, detail of configuration can be found: https://opensearch.org/docs/latest/ingest-pipelines/processors/ip2geo/
+ - `ip` can be an IPv4 or an IPv6 address
+ - `options` is an optional String of comma separated fields to output: the selection of fields is subject to dataSourceProvider's schema.  For example, the list of fields in the provided `geolite2-city` dataset includes: "country_iso_code", "country_name", "continent_name", "region_iso_code", "region_name", "city_name", "time_zone", "location"
+

integ-test/build.gradle

@@ -199,6 +199,9 @@ dependencies {
     // Needed for BWC tests
     zipArchive group: 'org.opensearch.plugin', name:'opensearch-job-scheduler', version: "${opensearch_build}"
     zipArchive group: 'org.opensearch.plugin', name:'opensearch-sql-plugin', version: "${bwcVersion}-SNAPSHOT"
+
+    // For GeoIP PPL functions
+    zipArchive group: 'org.opensearch.plugin', name:'geospatial', version: "${opensearch_build}"
 }
 
 java {
@@ -242,16 +245,27 @@ def getJobSchedulerPlugin() {
     })
 }
 
+def getGeoSpatialPlugin() {
+    provider { (RegularFile) (() ->
+            configurations.zipArchive.asFileTree.matching {
+                include '**/geospatia*'
+            }.singleFile )
+    }
+}
+
+
 testClusters {
     integTest {
         testDistribution = 'archive'
         plugin(getJobSchedulerPlugin())
+        plugin(getGeoSpatialPlugin())
         plugin ":opensearch-sql-plugin"
         setting "plugins.query.datasources.encryption.masterkey", "1234567812345678"
     }
     remoteCluster {
         testDistribution = 'archive'
         plugin(getJobSchedulerPlugin())
+        plugin(getGeoSpatialPlugin())
         plugin ":opensearch-sql-plugin"
     }
     integTestWithSecurity {

integ-test/src/test/java/org/opensearch/sql/legacy/SQLIntegTestCase.java

@@ -8,6 +8,7 @@
 import static com.google.common.base.Strings.isNullOrEmpty;
 import static org.opensearch.sql.legacy.TestUtils.createIndexByRestClient;
 import static org.opensearch.sql.legacy.TestUtils.getAccountIndexMapping;
+import static org.opensearch.sql.legacy.TestUtils.getAliasIndexMapping;
 import static org.opensearch.sql.legacy.TestUtils.getBankIndexMapping;
 import static org.opensearch.sql.legacy.TestUtils.getBankWithNullValuesIndexMapping;
 import static org.opensearch.sql.legacy.TestUtils.getDataTypeNonnumericIndexMapping;
@@ -20,6 +21,7 @@
 import static org.opensearch.sql.legacy.TestUtils.getDogs3IndexMapping;
 import static org.opensearch.sql.legacy.TestUtils.getEmployeeNestedTypeIndexMapping;
 import static org.opensearch.sql.legacy.TestUtils.getGameOfThronesIndexMapping;
+import static org.opensearch.sql.legacy.TestUtils.getGeoIpIndexMapping;
 import static org.opensearch.sql.legacy.TestUtils.getGeopointIndexMapping;
 import static org.opensearch.sql.legacy.TestUtils.getJoinTypeIndexMapping;
 import static org.opensearch.sql.legacy.TestUtils.getJsonTestIndexMapping;
@@ -626,6 +628,11 @@ public enum Index {
         "unexpandedObject",
         getUnexpandedObjectIndexMapping(),
         "src/test/resources/unexpanded_objects.json"),
+    GEOIP(
+        TestsConstants.TEST_INDEX_GEOIP,
+        "geoip",
+        getGeoIpIndexMapping(),
+        "src/test/resources/geoip.json"),
     BANK(
         TestsConstants.TEST_INDEX_BANK,
         "account",
@@ -751,7 +758,12 @@ public enum Index {
         TestsConstants.TEST_INDEX_JSON_TEST,
         "json",
         getJsonTestIndexMapping(),
-        "src/test/resources/json_test.json");
+        "src/test/resources/json_test.json"),
+    DATA_TYPE_ALIAS(
+        TestsConstants.TEST_INDEX_ALIAS,
+        "alias",
+        getAliasIndexMapping(),
+        "src/test/resources/alias.json");
 
     private final String name;
     private final String type;