initial commit

Bit-Quill · Oct 18, 2024 · 55e899d · 55e899d
1 parent e838e46
commit 55e899d
Show file tree

Hide file tree

Showing 8 changed files with 117 additions and 2 deletions.
diff --git a/core/build.gradle b/core/build.gradle
@@ -46,6 +46,8 @@ pitest {
 }
 
 dependencies {
+    implementation "org.opensearch:opensearch-geospatial-spi:${opensearch_build}"
+
     api group: 'com.google.guava', name: 'guava', version: "${guava_version}"
     api group: 'org.apache.commons', name: 'commons-lang3', version: '3.12.0'
     api group: 'org.apache.commons', name: 'commons-text', version: '1.10.0'

diff --git a/core/src/main/java/org/opensearch/sql/expression/DSL.java b/core/src/main/java/org/opensearch/sql/expression/DSL.java
@@ -655,6 +655,14 @@ public static FunctionExpression lte(Expression... expressions) {
     return lte(FunctionProperties.None, expressions);
   }
 
+  public static FunctionExpression geoip(FunctionProperties fp, Expression... expressions) {
+    return compile(fp, BuiltinFunctionName.GEOIP, expressions);
+  }
+
+  public static FunctionExpression geoip(Expression... expressions) {
+    return geoip(FunctionProperties.None, expressions);
+  }
+
   public static FunctionExpression greater(FunctionProperties fp, Expression... expressions) {
     return compile(fp, BuiltinFunctionName.GREATER, expressions);
   }

diff --git a/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java b/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java
@@ -252,7 +252,11 @@ public enum BuiltinFunctionName {
   MULTIMATCH(FunctionName.of("multimatch")),
   MULTIMATCHQUERY(FunctionName.of("multimatchquery")),
   WILDCARDQUERY(FunctionName.of("wildcardquery")),
-  WILDCARD_QUERY(FunctionName.of("wildcard_query"));
+  WILDCARD_QUERY(FunctionName.of("wildcard_query")),
+
+  /** Geospatial Function. */
+  GEOIP(FunctionName.of("geoip")),
+  IPLOCATION(FunctionName.of("iplocation"));
 
   private final FunctionName name;
 

diff --git a/core/src/main/java/org/opensearch/sql/expression/geospatial/GeospatialFunction.java b/core/src/main/java/org/opensearch/sql/expression/geospatial/GeospatialFunction.java
@@ -0,0 +1,43 @@
+package org.opensearch.sql.expression.geospatial;
+
+import static org.opensearch.sql.expression.function.FunctionDSL.define;
+import static org.opensearch.sql.expression.function.FunctionDSL.impl;
+
+import lombok.Builder;
+import lombok.experimental.UtilityClass;
+import org.opensearch.geospatial.spl
+import org.opensearch.sql.expression.function.BuiltinFunctionName;
+import org.opensearch.sql.expression.function.BuiltinFunctionRepository;
+import org.opensearch.sql.expression.function.DefaultFunctionResolver;
+import org.opensearch.sql.expression.function.FunctionName;
+
+@UtilityClass
+public class GeospatialFunction {
+    /**
+     * Register Geospatial Functions.
+     *
+     * @param repository {@link BuiltinFunctionRepository}.
+     */
+    public void register(BuiltinFunctionRepository repository) {
+        repository.register(geoip());
+    }
+
+    private DefaultFunctionResolver geoipIplocation(FunctionName functionName) {
+        return define(
+            functionName,
+            impl(
+                ,
+                STRING,
+                STRING
+            )
+        );
+    }
+
+    private DefaultFunctionResolver geoip() {
+        return geoipIplocation(BuiltinFunctionName.GEOIP.getName());
+    }
+
+    private DefaultFunctionResolver iplocation() {
+        return geoipIplocation(BuiltinFunctionName.IPLOCATION.getName());
+    }
+}
diff --git a/plugin/build.gradle b/plugin/build.gradle
@@ -50,7 +50,7 @@ opensearchplugin {
     name 'opensearch-sql'
     description 'OpenSearch SQL'
     classname 'org.opensearch.sql.plugin.SQLPlugin'
-    extendedPlugins = ['opensearch-job-scheduler']
+    extendedPlugins = ['opensearch-job-scheduler', 'opensearch-geospatial']
     licenseFile rootProject.file("LICENSE.txt")
     noticeFile rootProject.file("NOTICE")
 }
@@ -149,6 +149,7 @@ spotless {
 
 dependencies {
     compileOnly "org.opensearch:opensearch-job-scheduler-spi:${opensearch_build}"
+    compileOnly "org.opensearch:opensearch-geospatial-spi:${opensearch_build}"
     compileOnly "com.google.guava:guava:${guava_version}"
     compileOnly 'com.google.guava:failureaccess:1.0.2'
 
@@ -171,6 +172,7 @@ dependencies {
     testImplementation 'org.junit.jupiter:junit-jupiter:5.9.3'
 
     zipArchive group: 'org.opensearch.plugin', name:'opensearch-job-scheduler', version: "${opensearch_build}"
+    zipArchive group: 'org.opensearch.plugin', name:'opensearch-geospatial', version: "${opensearch_build}"
 }
 
 test {
@@ -301,8 +303,26 @@ def getJobSchedulerPlugin() {
     })
 }
 
+def getGeospatialPlugin() {
+    provider(new Callable<RegularFile>() {
+        @Override
+        RegularFile call() throws Exception {
+            return new RegularFile() {
+                @Override
+                File getAsFile() {
+                    return configurations.zipArchive.asFileTree.matching {
+                        include '**/opensearch-geospatial*'
+                    }.singleFile
+                }
+            }
+        }
+    })
+
+}
+
 testClusters.integTest {
     plugin(getJobSchedulerPlugin())
+    plugin(getGeospatialPlugin())
     plugin(project.tasks.bundlePlugin.archiveFile)
     testDistribution = "ARCHIVE"
 

diff --git a/ppl/src/main/antlr/OpenSearchPPLLexer.g4 b/ppl/src/main/antlr/OpenSearchPPLLexer.g4
@@ -329,6 +329,10 @@ NULLIF:                             'NULLIF';
 IF:                                 'IF';
 TYPEOF:                             'TYPEOF';
 
+// GEOSPATIAL FUNCTIONS
+GEOIP:                              'GEOIP';
+IPLOCATION:                         'IPLOCATION';
+
 // RELEVANCE FUNCTIONS AND PARAMETERS
 MATCH:                              'MATCH';
 MATCH_PHRASE:                       'MATCH_PHRASE';

diff --git a/ppl/src/main/antlr/OpenSearchPPLParser.g4 b/ppl/src/main/antlr/OpenSearchPPLParser.g4
@@ -387,6 +387,7 @@ evalFunctionName
    | flowControlFunctionName
    | systemFunctionName
    | positionFunctionName
+   | geospatialFunctionName
    ;
 
 functionArgs
@@ -666,6 +667,11 @@ positionFunctionName
    : POSITION
    ;
 
+geospatialFunctionName
+   : GEOIP
+   | IPLOCATION
+   ;
+
 // operators
  comparisonOperator
    : EQUAL

diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstExpressionBuilderTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/parser/AstExpressionBuilderTest.java
@@ -816,4 +816,32 @@ public void testTimeStampDiffFunctionExpr() {
                     stringLiteral("1997-01-01 00:00:00"),
                     stringLiteral("2001-03-06 00:00:00")))));
   }
+
+  @Test
+  public void testGeoIpFunctionExpr() {
+    assertEqual(
+            "source=t | eval f=geoip('127.0.0.1', 'lat,lon')",
+            eval(
+                    relation("t"),
+                    let(
+                            field("f"),
+                            function(
+                                    "geoip",
+                                    stringLiteral("127.0.0.1"),
+                                    stringLiteral("lat,lon")))));
+  }
+
+  @Test
+  public void testIpLocationFunctionExpr() {
+    assertEqual(
+            "source=t | eval f=iplocation('127.0.0.1', 'lat,lon')",
+            eval(
+                    relation("t"),
+                    let(
+                            field("f"),
+                            function(
+                                    "iplocation",
+                                    stringLiteral("127.0.0.1"),
+                                    stringLiteral("lat,lon")))));
+  }
 }