Support ScalarSubquery PPL

Signed-off-by: Lantao Jin <ltjin@amazon.com>
opensearch-project · Oct 8, 2024 · ef5a634 · ef5a634
1 parent e3a19dd
commit ef5a634
Show file tree

Hide file tree

Showing 6 changed files with 530 additions and 0 deletions.
diff --git a/ppl-spark-integration/src/main/antlr4/OpenSearchPPLParser.g4 b/ppl-spark-integration/src/main/antlr4/OpenSearchPPLParser.g4
@@ -359,6 +359,7 @@ valueExpression
    | positionFunction                                                                           # positionFunctionCall
    | caseFunction                                                                               # caseExpr
    | LT_PRTHS valueExpression RT_PRTHS                                                          # parentheticValueExpr
+   | LT_SQR_PRTHS subSearch RT_SQR_PRTHS                                                        # scalarSubqueryExpr
    ;
 
 primaryExpression

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
@@ -28,6 +28,7 @@
 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.ScalarSubquery;
 import org.opensearch.sql.ast.expression.Span;
 import org.opensearch.sql.ast.expression.UnresolvedArgument;
 import org.opensearch.sql.ast.expression.UnresolvedAttribute;
@@ -293,4 +294,8 @@ public T visitExplain(Explain node, C context) {
   public T visitInSubquery(InSubquery node, C context) {
     return visitChildren(node, context);
   }
+
+  public T visitScalarSubquery(ScalarSubquery node, C context) {
+    return visitChildren(node, context);
+  }
 }
diff --git a/ppl-spark-integration/src/main/java/org/opensearch/sql/ast/expression/ScalarSubquery.java b/ppl-spark-integration/src/main/java/org/opensearch/sql/ast/expression/ScalarSubquery.java
@@ -0,0 +1,26 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.ast.expression;
+
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+import lombok.ToString;
+import org.opensearch.sql.ast.AbstractNodeVisitor;
+import org.opensearch.sql.ast.tree.UnresolvedPlan;
+
+@Getter
+@ToString
+@EqualsAndHashCode(callSuper = false)
+@RequiredArgsConstructor
+public class ScalarSubquery extends UnresolvedExpression {
+    private final UnresolvedPlan query;
+
+    @Override
+    public <R, C> R accept(AbstractNodeVisitor<R, C> nodeVisitor, C context) {
+        return nodeVisitor.visitScalarSubquery(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
@@ -17,6 +17,7 @@
 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.plans.logical.*;
@@ -47,6 +48,7 @@
 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.ScalarSubquery;
 import org.opensearch.sql.ast.expression.Span;
 import org.opensearch.sql.ast.expression.UnresolvedExpression;
 import org.opensearch.sql.ast.expression.When;
@@ -774,5 +776,20 @@ public Expression visitInSubquery(InSubquery node, CatalystPlanContext outerCont
                     Option.empty()));
             return outerContext.getNamedParseExpressions().push(inSubQuery);
         }
+
+        @Override
+        public Expression visitScalarSubquery(ScalarSubquery node, CatalystPlanContext context) {
+            CatalystPlanContext innerContext = new CatalystPlanContext();
+            UnresolvedPlan outerPlan = node.getQuery();
+            LogicalPlan subSearch = CatalystQueryPlanVisitor.this.visitSubSearch(outerPlan, innerContext);
+            Expression scalarSubQuery = ScalarSubquery$.MODULE$.apply(
+                subSearch,
+                seq(new java.util.ArrayList<Expression>()),
+                NamedExpression.newExprId(),
+                seq(new java.util.ArrayList<Expression>()),
+                Option.empty(),
+                Option.empty());
+            return context.getNamedParseExpressions().push(scalarSubQuery);
+        }
     }
 }
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
@@ -31,6 +31,7 @@
 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.ScalarSubquery;
 import org.opensearch.sql.ast.expression.Span;
 import org.opensearch.sql.ast.expression.SpanUnit;
 import org.opensearch.sql.ast.expression.UnresolvedArgument;
@@ -387,6 +388,11 @@ public UnresolvedExpression visitInSubqueryExpr(OpenSearchPPLParser.InSubqueryEx
         return ctx.NOT() != null ? new Not(expr) : expr;
     }
 
+    @Override
+    public UnresolvedExpression visitScalarSubqueryExpr(OpenSearchPPLParser.ScalarSubqueryExprContext ctx) {
+        return new ScalarSubquery(astBuilder.visitSubSearch(ctx.subSearch()));
+    }
+
     private QualifiedName visitIdentifiers(List<? extends ParserRuleContext> ctx) {
         return new QualifiedName(
                 ctx.stream()