[WIP] FEAT(#3038): iplocation command

core/src/main/java/org/opensearch/sql/analysis/ExpressionAnalyzer.java

@@ -21,33 +21,7 @@
 import org.opensearch.sql.analysis.symbol.Namespace;
 import org.opensearch.sql.analysis.symbol.Symbol;
 import org.opensearch.sql.ast.AbstractNodeVisitor;
-import org.opensearch.sql.ast.expression.AggregateFunction;
-import org.opensearch.sql.ast.expression.AllFields;
-import org.opensearch.sql.ast.expression.And;
-import org.opensearch.sql.ast.expression.Between;
-import org.opensearch.sql.ast.expression.Case;
-import org.opensearch.sql.ast.expression.Cast;
-import org.opensearch.sql.ast.expression.Compare;
-import org.opensearch.sql.ast.expression.DataType;
-import org.opensearch.sql.ast.expression.EqualTo;
-import org.opensearch.sql.ast.expression.Field;
-import org.opensearch.sql.ast.expression.Function;
-import org.opensearch.sql.ast.expression.HighlightFunction;
-import org.opensearch.sql.ast.expression.In;
-import org.opensearch.sql.ast.expression.Interval;
-import org.opensearch.sql.ast.expression.Literal;
-import org.opensearch.sql.ast.expression.Not;
-import org.opensearch.sql.ast.expression.Or;
-import org.opensearch.sql.ast.expression.QualifiedName;
-import org.opensearch.sql.ast.expression.RelevanceFieldList;
-import org.opensearch.sql.ast.expression.ScoreFunction;
-import org.opensearch.sql.ast.expression.Span;
-import org.opensearch.sql.ast.expression.UnresolvedArgument;
-import org.opensearch.sql.ast.expression.UnresolvedAttribute;
-import org.opensearch.sql.ast.expression.UnresolvedExpression;
-import org.opensearch.sql.ast.expression.When;
-import org.opensearch.sql.ast.expression.WindowFunction;
-import org.opensearch.sql.ast.expression.Xor;
+import org.opensearch.sql.ast.expression.*;
 import org.opensearch.sql.data.model.ExprValueUtils;
 import org.opensearch.sql.data.type.ExprCoreType;
 import org.opensearch.sql.data.type.ExprType;
@@ -67,6 +41,7 @@
 import org.opensearch.sql.expression.function.BuiltinFunctionRepository;
 import org.opensearch.sql.expression.function.FunctionName;
 import org.opensearch.sql.expression.function.OpenSearchFunctions;
+import org.opensearch.sql.expression.ip.GeoipExpression;
 import org.opensearch.sql.expression.parse.ParseExpression;
 import org.opensearch.sql.expression.span.SpanExpression;
 import org.opensearch.sql.expression.window.aggregation.AggregateWindowFunction;
@@ -203,6 +178,14 @@ public Expression visitFunction(Function node, AnalysisContext context) {
         repository.compile(context.getFunctionProperties(), functionName, arguments);
   }
 
+  @Override
+  public Expression visitGeoip(Geoip node, AnalysisContext context) {
+    return new GeoipExpression(
+            node.getDatasource().accept(this, context),
+            node.getIpAddress().accept(this ,context),
+            node.getProperties());
+  }
+
   @SuppressWarnings("unchecked")
   @Override
   public Expression visitWindowFunction(WindowFunction node, AnalysisContext context) {

core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java

@@ -5,37 +5,7 @@
 
 package org.opensearch.sql.ast;
 
-import org.opensearch.sql.ast.expression.AggregateFunction;
-import org.opensearch.sql.ast.expression.Alias;
-import org.opensearch.sql.ast.expression.AllFields;
-import org.opensearch.sql.ast.expression.And;
-import org.opensearch.sql.ast.expression.Argument;
-import org.opensearch.sql.ast.expression.AttributeList;
-import org.opensearch.sql.ast.expression.Between;
-import org.opensearch.sql.ast.expression.Case;
-import org.opensearch.sql.ast.expression.Cast;
-import org.opensearch.sql.ast.expression.Compare;
-import org.opensearch.sql.ast.expression.EqualTo;
-import org.opensearch.sql.ast.expression.Field;
-import org.opensearch.sql.ast.expression.Function;
-import org.opensearch.sql.ast.expression.HighlightFunction;
-import org.opensearch.sql.ast.expression.In;
-import org.opensearch.sql.ast.expression.Interval;
-import org.opensearch.sql.ast.expression.Let;
-import org.opensearch.sql.ast.expression.Literal;
-import org.opensearch.sql.ast.expression.Map;
-import org.opensearch.sql.ast.expression.NestedAllTupleFields;
-import org.opensearch.sql.ast.expression.Not;
-import org.opensearch.sql.ast.expression.Or;
-import org.opensearch.sql.ast.expression.QualifiedName;
-import org.opensearch.sql.ast.expression.RelevanceFieldList;
-import org.opensearch.sql.ast.expression.ScoreFunction;
-import org.opensearch.sql.ast.expression.Span;
-import org.opensearch.sql.ast.expression.UnresolvedArgument;
-import org.opensearch.sql.ast.expression.UnresolvedAttribute;
-import org.opensearch.sql.ast.expression.When;
-import org.opensearch.sql.ast.expression.WindowFunction;
-import org.opensearch.sql.ast.expression.Xor;
+import org.opensearch.sql.ast.expression.*;
 import org.opensearch.sql.ast.statement.Explain;
 import org.opensearch.sql.ast.statement.Query;
 import org.opensearch.sql.ast.statement.Statement;
@@ -165,6 +135,8 @@ public T visitFunction(Function node, C context) {
     return visitChildren(node, context);
   }
 
+  public T visitGeoip(Geoip node, C context) { return visitChildren(node, context); }
+
   public T visitWindowFunction(WindowFunction node, C context) {
     return visitChildren(node, context);
   }

core/src/main/java/org/opensearch/sql/ast/dsl/AstDSL.java

@@ -10,39 +10,7 @@
 import java.util.stream.Collectors;
 import lombok.experimental.UtilityClass;
 import org.apache.commons.lang3.tuple.Pair;
-import org.opensearch.sql.ast.expression.AggregateFunction;
-import org.opensearch.sql.ast.expression.Alias;
-import org.opensearch.sql.ast.expression.AllFields;
-import org.opensearch.sql.ast.expression.And;
-import org.opensearch.sql.ast.expression.Argument;
-import org.opensearch.sql.ast.expression.Between;
-import org.opensearch.sql.ast.expression.Case;
-import org.opensearch.sql.ast.expression.Cast;
-import org.opensearch.sql.ast.expression.Compare;
-import org.opensearch.sql.ast.expression.DataType;
-import org.opensearch.sql.ast.expression.EqualTo;
-import org.opensearch.sql.ast.expression.Field;
-import org.opensearch.sql.ast.expression.Function;
-import org.opensearch.sql.ast.expression.HighlightFunction;
-import org.opensearch.sql.ast.expression.In;
-import org.opensearch.sql.ast.expression.Interval;
-import org.opensearch.sql.ast.expression.Let;
-import org.opensearch.sql.ast.expression.Literal;
-import org.opensearch.sql.ast.expression.Map;
-import org.opensearch.sql.ast.expression.NestedAllTupleFields;
-import org.opensearch.sql.ast.expression.Not;
-import org.opensearch.sql.ast.expression.Or;
-import org.opensearch.sql.ast.expression.ParseMethod;
-import org.opensearch.sql.ast.expression.QualifiedName;
-import org.opensearch.sql.ast.expression.ScoreFunction;
-import org.opensearch.sql.ast.expression.Span;
-import org.opensearch.sql.ast.expression.SpanUnit;
-import org.opensearch.sql.ast.expression.UnresolvedArgument;
-import org.opensearch.sql.ast.expression.UnresolvedAttribute;
-import org.opensearch.sql.ast.expression.UnresolvedExpression;
-import org.opensearch.sql.ast.expression.When;
-import org.opensearch.sql.ast.expression.WindowFunction;
-import org.opensearch.sql.ast.expression.Xor;
+import org.opensearch.sql.ast.expression.*;
 import org.opensearch.sql.ast.tree.Aggregation;
 import org.opensearch.sql.ast.tree.Dedupe;
 import org.opensearch.sql.ast.tree.Eval;
@@ -429,6 +397,10 @@ public static List<Argument> defaultSortFieldArgs() {
     return exprList(argument("asc", booleanLiteral(true)), argument("type", nullLiteral()));
   }
 
+  public static Geoip geoip(UnresolvedExpression datasource, UnresolvedExpression ipAddress, String properties) {
+    return new Geoip(datasource, ipAddress, properties);
+  }
+
   public static Span span(UnresolvedExpression field, UnresolvedExpression value, SpanUnit unit) {
     return new Span(field, value, unit);
   }

core/src/main/java/org/opensearch/sql/ast/expression/Geoip.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.ast.expression;
+
+import com.google.common.collect.ImmutableList;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+import org.opensearch.sql.ast.AbstractNodeVisitor;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.stream.Collectors;
+
+/**
+ * Expression node of scalar function. Params include function name (@funcName) and function
+ * arguments (@funcArgs)
+ */
+@Getter
+@EqualsAndHashCode(callSuper = false)
+@RequiredArgsConstructor
+public class Geoip extends UnresolvedExpression {
+  private final UnresolvedExpression datasource;
+  private final UnresolvedExpression ipAddress;
+  private final String properties;
+
+  @Override
+  public List<UnresolvedExpression> getChild() {
+    return ImmutableList.of(datasource, ipAddress);
+  }
+
+  @Override
+  public <R, C> R accept(AbstractNodeVisitor<R, C> nodeVisitor, C context) {
+    return nodeVisitor.visitGeoip(this, context);
+  }
+}

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

@@ -6,6 +6,8 @@
 package org.opensearch.sql.expression;
 
 import java.util.Arrays;
+
+import org.opensearch.sql.ast.expression.Geoip;
 import org.opensearch.sql.ast.expression.SpanUnit;
 import org.opensearch.sql.data.model.ExprShortValue;
 import org.opensearch.sql.data.model.ExprValue;
@@ -19,6 +21,7 @@
 import org.opensearch.sql.expression.function.BuiltinFunctionRepository;
 import org.opensearch.sql.expression.function.FunctionImplementation;
 import org.opensearch.sql.expression.function.FunctionProperties;
+import org.opensearch.sql.expression.ip.GeoipExpression;
 import org.opensearch.sql.expression.parse.GrokExpression;
 import org.opensearch.sql.expression.parse.ParseExpression;
 import org.opensearch.sql.expression.parse.PatternsExpression;
@@ -142,6 +145,10 @@ public static SpanExpression span(Expression field, Expression value, String uni
     return new SpanExpression(field, value, SpanUnit.of(unit));
   }
 
+  public static GeoipExpression geoip(Expression datasource, Expression ipAddress, String properties) {
+    return new GeoipExpression(datasource, ipAddress, properties);
+  }
+
   public static FunctionExpression abs(Expression... expressions) {
     return compile(FunctionProperties.None, BuiltinFunctionName.ABS, expressions);
   }
@@ -965,7 +972,7 @@ public static FunctionExpression utc_timestamp(
   private static <T extends FunctionImplementation> T compile(
       FunctionProperties functionProperties, BuiltinFunctionName bfn, Expression... args) {
     return (T)
-        BuiltinFunctionRepository.getInstance()
-            .compile(functionProperties, bfn.getName(), Arrays.asList(args));
+            BuiltinFunctionRepository.getInstance()
+                    .compile(functionProperties, bfn.getName(), Arrays.asList(args));
   }
 }

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

@@ -201,6 +201,9 @@ public enum BuiltinFunctionName {
   TRIM(FunctionName.of("trim")),
   UPPER(FunctionName.of("upper")),
 
+  /** 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/ip/GeoipExpression.java

@@ -0,0 +1,52 @@
+package org.opensearch.sql.expression.ip;
+
+
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.ToString;
+import org.opensearch.sql.common.grok.Converter;
+import org.opensearch.sql.data.model.ExprValue;
+import org.opensearch.sql.data.type.ExprType;
+import org.opensearch.sql.expression.Expression;
+import org.opensearch.sql.expression.ExpressionNodeVisitor;
+import org.opensearch.sql.expression.env.Environment;
+import org.opensearch.sql.planner.physical.collector.Rounding;
+
+import java.util.List;
+
+import static org.opensearch.sql.data.type.ExprCoreType.ARRAY;
+import static org.opensearch.sql.data.type.ExprCoreType.STRING;
+
+@Getter
+@ToString
+@EqualsAndHashCode
+public class GeoipExpression implements Expression {
+    private final Expression datasource;
+    private final Expression ipAddress;
+    private final String[] properties;
+
+    public GeoipExpression(Expression datasource, Expression ipAddress, String properties) {
+        this.datasource = datasource;
+        this.ipAddress = ipAddress;
+        this.properties = properties.split(",");
+    }
+
+    @Override
+    public ExprValue valueOf(Environment<Expression, ExprValue> valueEnv) {
+        return ipAddress.valueOf(valueEnv);
+    }
+
+    @Override
+    public ExprType type() {
+        if (this.properties.length <= 1) {
+            return STRING;
+        } else {
+            return ARRAY;
+        }
+    }
+
+    @Override
+    public <T, C> T accept(ExpressionNodeVisitor<T, C> visitor, C context) {
+        return visitor.visitNode(this, context);
+    }
+}

ppl/src/main/antlr/OpenSearchPPLLexer.g4

@@ -329,6 +329,9 @@ NULLIF:                             'NULLIF';
 IF:                                 'IF';
 TYPEOF:                             'TYPEOF';
 
+// GEOLOCATION FUNCTIONS
+GEOIP:                              'GEOIP';
+
 // RELEVANCE FUNCTIONS AND PARAMETERS
 MATCH:                              'MATCH';
 MATCH_PHRASE:                       'MATCH_PHRASE';

ppl/src/main/antlr/OpenSearchPPLParser.g4

@@ -285,6 +285,7 @@ valueExpression
 primaryExpression
    : evalFunctionCall
    | dataTypeFunctionCall
+   | geoipFunctionCall
    | fieldExpression
    | literalValue
    ;
@@ -356,6 +357,11 @@ evalFunctionCall
    : evalFunctionName LT_PRTHS functionArgs RT_PRTHS
    ;
 
+// geoip function
+geoipFunctionCall
+   : GEOIP LT_PRTHS (datasource = functionArg COMMA)? ipAddress = functionArg (COMMA properties = stringLiteral)? RT_PRTHS
+   ;
+
 // cast function
 dataTypeFunctionCall
    : CAST LT_PRTHS expression AS convertedDataType RT_PRTHS
@@ -666,6 +672,7 @@ positionFunctionName
    : POSITION
    ;
 
+
 // operators
  comparisonOperator
    : EQUAL

ppl/src/main/java/org/opensearch/sql/ppl/parser/AstExpressionBuilder.java

@@ -21,6 +21,7 @@
 import static org.opensearch.sql.ppl.antlr.parser.OpenSearchPPLParser.EvalClauseContext;
 import static org.opensearch.sql.ppl.antlr.parser.OpenSearchPPLParser.EvalFunctionCallContext;
 import static org.opensearch.sql.ppl.antlr.parser.OpenSearchPPLParser.FieldExpressionContext;
+import static org.opensearch.sql.ppl.antlr.parser.OpenSearchPPLParser.GeoipFunctionCallContext;
 import static org.opensearch.sql.ppl.antlr.parser.OpenSearchPPLParser.IdentsAsQualifiedNameContext;
 import static org.opensearch.sql.ppl.antlr.parser.OpenSearchPPLParser.IdentsAsTableQualifiedNameContext;
 import static org.opensearch.sql.ppl.antlr.parser.OpenSearchPPLParser.IdentsAsWildcardQualifiedNameContext;
@@ -199,6 +200,14 @@ public UnresolvedExpression visitEvalFunctionCall(EvalFunctionCallContext ctx) {
     return buildFunction(ctx.evalFunctionName().getText(), ctx.functionArgs().functionArg());
   }
 
+  /** Eval function. */
+  @Override
+  public UnresolvedExpression visitGeoipFunctionCall(GeoipFunctionCallContext ctx) {
+    String properties = ctx.properties != null ? visit(ctx.properties).toString() : "";
+    UnresolvedExpression datasource = ctx.datasource != null ? visit(ctx.datasource) : AstDSL.stringLiteral("");
+    return new Geoip(datasource, visit(ctx.ipAddress), properties);
+  }
+
   /** Cast function. */
   @Override
   public UnresolvedExpression visitDataTypeFunctionCall(DataTypeFunctionCallContext ctx) {