From e65e664739fa166e13731d30c941ecec656fa50d Mon Sep 17 00:00:00 2001 From: "KENRICK-THINKPA\\kenricky" Date: Mon, 4 Nov 2024 18:38:38 -0800 Subject: [PATCH] added stub method for geoip --- .../src/main/antlr4/OpenSearchPPLParser.g4 | 4 +- .../sql/ast/AbstractNodeVisitor.java | 31 ++---------- .../opensearch/sql/ast/expression/GeoIp.java | 22 +++++++++ .../sql/ppl/CatalystQueryPlanVisitor.java | 47 ++++++------------- .../sql/ppl/parser/AstExpressionBuilder.java | 39 ++++++--------- ...PLLogicalPlanEvalTranslatorTestSuite.scala | 4 ++ 6 files changed, 60 insertions(+), 87 deletions(-) create mode 100644 ppl-spark-integration/src/main/java/org/opensearch/sql/ast/expression/GeoIp.java diff --git a/ppl-spark-integration/src/main/antlr4/OpenSearchPPLParser.g4 b/ppl-spark-integration/src/main/antlr4/OpenSearchPPLParser.g4 index 0c010a51f..8b84f0348 100644 --- a/ppl-spark-integration/src/main/antlr4/OpenSearchPPLParser.g4 +++ b/ppl-spark-integration/src/main/antlr4/OpenSearchPPLParser.g4 @@ -385,6 +385,7 @@ valueExpression primaryExpression : evalFunctionCall + | geoIpFunctionCall | fieldExpression | literalValue ; @@ -465,7 +466,6 @@ wcFieldExpression // functions evalFunctionCall : evalFunctionName LT_PRTHS functionArgs RT_PRTHS - | geoipFunction ; // cast function @@ -782,7 +782,7 @@ coalesceFunctionName : COALESCE ; -geoipFunction +geoIpFunctionCall : GEOIP LT_PRTHS (datasource = functionArg COMMA)? ipAddress = functionArg (COMMA properties = stringLiteral)? RT_PRTHS ; diff --git a/ppl-spark-integration/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java b/ppl-spark-integration/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java index c361ded08..66303b2a6 100644 --- a/ppl-spark-integration/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java +++ b/ppl-spark-integration/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java @@ -5,37 +5,10 @@ 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.Compare; -import org.opensearch.sql.ast.expression.EqualTo; -import org.opensearch.sql.ast.expression.Field; -import org.opensearch.sql.ast.expression.FieldsMapping; -import org.opensearch.sql.ast.expression.Function; -import org.opensearch.sql.ast.expression.In; +import org.opensearch.sql.ast.expression.*; import org.opensearch.sql.ast.expression.subquery.ExistsSubquery; import org.opensearch.sql.ast.expression.subquery.InSubquery; -import org.opensearch.sql.ast.expression.Interval; -import org.opensearch.sql.ast.expression.IsEmpty; -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.Not; -import org.opensearch.sql.ast.expression.Or; -import org.opensearch.sql.ast.expression.QualifiedName; import org.opensearch.sql.ast.expression.subquery.ScalarSubquery; -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.statement.Explain; import org.opensearch.sql.ast.statement.Query; import org.opensearch.sql.ast.statement.Statement; @@ -307,4 +280,6 @@ public T visitScalarSubquery(ScalarSubquery node, C context) { public T visitExistsSubquery(ExistsSubquery node, C context) { return visitChildren(node, context); } + + public T visitGeoIp(GeoIp node, C context) { return visitGeoip(node, context); } } diff --git a/ppl-spark-integration/src/main/java/org/opensearch/sql/ast/expression/GeoIp.java b/ppl-spark-integration/src/main/java/org/opensearch/sql/ast/expression/GeoIp.java new file mode 100644 index 000000000..b3edf86ab --- /dev/null +++ b/ppl-spark-integration/src/main/java/org/opensearch/sql/ast/expression/GeoIp.java @@ -0,0 +1,22 @@ +package org.opensearch.sql.ast.expression; + +import org.opensearch.sql.ast.AbstractNodeVisitor; + +import java.util.Arrays; +import java.util.List; + +public class GeoIp extends UnresolvedExpression { + private UnresolvedExpression datasource; + private UnresolvedExpression ipAddress; + private Literal properties; + + @Override + public List getChild() { + return Arrays.asList(datasource, ipAddress); + } + + @Override + public T accept(AbstractNodeVisitor nodeVisitor, C context) { + return nodeVisitor.visitGeoIp(this, context); + } +} diff --git a/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/CatalystQueryPlanVisitor.java b/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/CatalystQueryPlanVisitor.java index 902fc72e3..22b53d638 100644 --- a/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/CatalystQueryPlanVisitor.java +++ b/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/CatalystQueryPlanVisitor.java @@ -10,18 +10,7 @@ import org.apache.spark.sql.catalyst.analysis.UnresolvedFunction; import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation; import org.apache.spark.sql.catalyst.analysis.UnresolvedStar$; -import org.apache.spark.sql.catalyst.expressions.Ascending$; -import org.apache.spark.sql.catalyst.expressions.CaseWhen; -import org.apache.spark.sql.catalyst.expressions.Descending$; -import org.apache.spark.sql.catalyst.expressions.Exists$; -import org.apache.spark.sql.catalyst.expressions.Expression; -import org.apache.spark.sql.catalyst.expressions.InSubquery$; -import org.apache.spark.sql.catalyst.expressions.ListQuery$; -import org.apache.spark.sql.catalyst.expressions.NamedExpression; -import org.apache.spark.sql.catalyst.expressions.Predicate; -import org.apache.spark.sql.catalyst.expressions.ScalarSubquery$; -import org.apache.spark.sql.catalyst.expressions.SortDirection; -import org.apache.spark.sql.catalyst.expressions.SortOrder; +import org.apache.spark.sql.catalyst.expressions.*; import org.apache.spark.sql.catalyst.plans.logical.*; import org.apache.spark.sql.execution.ExplainMode; import org.apache.spark.sql.execution.command.DescribeTableCommand; @@ -29,34 +18,18 @@ import org.apache.spark.sql.types.DataTypes; import org.apache.spark.sql.util.CaseInsensitiveStringMap; import org.opensearch.sql.ast.AbstractNodeVisitor; -import org.opensearch.sql.ast.expression.AggregateFunction; +import org.opensearch.sql.ast.expression.*; 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.BinaryExpression; -import org.opensearch.sql.ast.expression.Case; -import org.opensearch.sql.ast.expression.Compare; -import org.opensearch.sql.ast.expression.Field; -import org.opensearch.sql.ast.expression.FieldsMapping; -import org.opensearch.sql.ast.expression.Function; import org.opensearch.sql.ast.expression.In; -import org.opensearch.sql.ast.expression.subquery.ExistsSubquery; -import org.opensearch.sql.ast.expression.subquery.InSubquery; -import org.opensearch.sql.ast.expression.Interval; -import org.opensearch.sql.ast.expression.IsEmpty; -import org.opensearch.sql.ast.expression.Let; 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.ParseMethod; -import org.opensearch.sql.ast.expression.QualifiedName; -import org.opensearch.sql.ast.expression.subquery.ScalarSubquery; -import org.opensearch.sql.ast.expression.Span; -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.subquery.ExistsSubquery; +import org.opensearch.sql.ast.expression.subquery.InSubquery; +import org.opensearch.sql.ast.expression.subquery.ScalarSubquery; import org.opensearch.sql.ast.statement.Explain; import org.opensearch.sql.ast.statement.Query; import org.opensearch.sql.ast.statement.Statement; @@ -93,6 +66,8 @@ import scala.collection.Seq; import java.util.*; +import java.util.Map; +import java.util.Stack; import java.util.function.BiFunction; import java.util.stream.Collectors; @@ -829,5 +804,13 @@ public Expression visitExistsSubquery(ExistsSubquery node, CatalystPlanContext c Option.empty()); return context.getNamedParseExpressions().push(existsSubQuery); } + + @Override + public Expression visitGeoIp(GeoIp node, CatalystPlanContext context) { + + ScalaUDF udf = new ScalaUDF(); + + return context.getNamedParseExpressions().push(udf); + } } } diff --git a/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/parser/AstExpressionBuilder.java b/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/parser/AstExpressionBuilder.java index 3b98edd77..394049651 100644 --- a/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/parser/AstExpressionBuilder.java +++ b/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/parser/AstExpressionBuilder.java @@ -11,34 +11,10 @@ import org.antlr.v4.runtime.RuleContext; import org.opensearch.flint.spark.ppl.OpenSearchPPLParser; import org.opensearch.flint.spark.ppl.OpenSearchPPLParserBaseVisitor; -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.Case; -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.*; import org.opensearch.sql.ast.expression.subquery.ExistsSubquery; import org.opensearch.sql.ast.expression.subquery.InSubquery; -import org.opensearch.sql.ast.expression.Interval; -import org.opensearch.sql.ast.expression.IntervalUnit; -import org.opensearch.sql.ast.expression.IsEmpty; -import org.opensearch.sql.ast.expression.Let; -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.subquery.ScalarSubquery; -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.UnresolvedExpression; -import org.opensearch.sql.ast.expression.When; -import org.opensearch.sql.ast.expression.Xor; import org.opensearch.sql.common.utils.StringUtils; import org.opensearch.sql.ppl.utils.ArgumentFactory; @@ -399,6 +375,19 @@ public UnresolvedExpression visitExistsSubqueryExpr(OpenSearchPPLParser.ExistsSu return new ExistsSubquery(astBuilder.visitSubSearch(ctx.subSearch())); } + @Override + public UnresolvedExpression visitGeoIpFunctionCall(OpenSearchPPLParser.GeoIpFunctionCallContext ctx) { + UnresolvedExpression datasource = visit(ctx.datasource); + UnresolvedExpression ipAddress = visit(ctx.ipAddress); + Literal properties = (Literal) visit(ctx.properties); + return new GeoIp(datasource, ipAddress, properties); + } + + Literal pattern = (Literal) internalVisitExpression(ctx.pattern); + + return new Parse(ParseMethod.REGEX, sourceField, pattern, ImmutableMap.of()); +} + private QualifiedName visitIdentifiers(List ctx) { return new QualifiedName( ctx.stream() diff --git a/ppl-spark-integration/src/test/scala/org/opensearch/flint/spark/ppl/PPLLogicalPlanEvalTranslatorTestSuite.scala b/ppl-spark-integration/src/test/scala/org/opensearch/flint/spark/ppl/PPLLogicalPlanEvalTranslatorTestSuite.scala index 3e2b3cc30..b8cc9776d 100644 --- a/ppl-spark-integration/src/test/scala/org/opensearch/flint/spark/ppl/PPLLogicalPlanEvalTranslatorTestSuite.scala +++ b/ppl-spark-integration/src/test/scala/org/opensearch/flint/spark/ppl/PPLLogicalPlanEvalTranslatorTestSuite.scala @@ -177,6 +177,10 @@ class PPLLogicalPlanEvalTranslatorTestSuite comparePlans(expectedPlan, logPlan, checkAnalysis = false) } + test("test eval expression - geoip function") { + + } + // Todo fields-excluded command not supported ignore("test eval expressions with fields-excluded command") { val context = new CatalystPlanContext