Implement ppl scalar subquery command with Calcite (#3392)

* Implement ppl scalar subquery command with Calcite

Signed-off-by: Lantao Jin <ltjin@amazon.com>

* more general subquery checker

Signed-off-by: Lantao Jin <ltjin@amazon.com>

* support correlated IN subquery

Signed-off-by: Lantao Jin <ltjin@amazon.com>

---------

Signed-off-by: Lantao Jin <ltjin@amazon.com>
Signed-off-by: xinyual <xinyual@amazon.com>

Author: LantaoJin

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

@@ -38,6 +38,7 @@
 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;
@@ -348,6 +349,10 @@ public T visitSubqueryAlias(SubqueryAlias node, C context) {
     return visitChildren(node, context);
   }
 
+  public T visitScalarSubquery(ScalarSubquery node, C context) {
+    return visitChildren(node, context);
+  }
+
   public T visitExistsSubquery(ExistsSubquery node, C context) {
     return visitChildren(node, context);
   }

core/src/main/java/org/opensearch/sql/ast/expression/subquery/ExistsSubquery.java

@@ -16,9 +16,9 @@
 import org.opensearch.sql.common.utils.StringUtils;
 
 @Getter
-@EqualsAndHashCode(callSuper = false)
+@EqualsAndHashCode(callSuper = true)
 @RequiredArgsConstructor
-public class ExistsSubquery extends UnresolvedExpression {
+public class ExistsSubquery extends SubqueryExpression {
   private final UnresolvedPlan query;
 
   @Override

core/src/main/java/org/opensearch/sql/ast/expression/subquery/InSubquery.java

@@ -15,9 +15,9 @@
 import org.opensearch.sql.common.utils.StringUtils;
 
 @Getter
-@EqualsAndHashCode(callSuper = false)
+@EqualsAndHashCode(callSuper = true)
 @RequiredArgsConstructor
-public class InSubquery extends UnresolvedExpression {
+public class InSubquery extends SubqueryExpression {
   private final List<UnresolvedExpression> value;
   private final UnresolvedPlan query;
 

core/src/main/java/org/opensearch/sql/ast/expression/subquery/ScalarSubquery.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.ast.expression.subquery;
+
+import com.google.common.collect.ImmutableList;
+import java.util.List;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+import lombok.ToString;
+import org.opensearch.sql.ast.AbstractNodeVisitor;
+import org.opensearch.sql.ast.expression.UnresolvedExpression;
+import org.opensearch.sql.ast.tree.UnresolvedPlan;
+
+@Getter
+@ToString
+@EqualsAndHashCode(callSuper = true)
+@RequiredArgsConstructor
+public class ScalarSubquery extends SubqueryExpression {
+  private final UnresolvedPlan query;
+
+  @Override
+  public <R, C> R accept(AbstractNodeVisitor<R, C> nodeVisitor, C context) {
+    return nodeVisitor.visitScalarSubquery(this, context);
+  }
+
+  @Override
+  public List<UnresolvedExpression> getChild() {
+    return ImmutableList.of();
+  }
+}

core/src/main/java/org/opensearch/sql/ast/expression/subquery/SubqueryExpression.java

@@ -0,0 +1,13 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.ast.expression.subquery;
+
+import lombok.EqualsAndHashCode;
+import org.opensearch.sql.ast.expression.UnresolvedExpression;
+
+/** Basic class of subquery expression */
+@EqualsAndHashCode(callSuper = false)
+public abstract class SubqueryExpression extends UnresolvedExpression {}

core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java

@@ -13,6 +13,7 @@
 import static org.opensearch.sql.ast.tree.Sort.SortOrder.DESC;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterables;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
@@ -31,14 +32,14 @@
 import org.apache.calcite.util.Holder;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.opensearch.sql.ast.AbstractNodeVisitor;
+import org.opensearch.sql.ast.Node;
 import org.opensearch.sql.ast.expression.AllFields;
 import org.opensearch.sql.ast.expression.Argument;
-import org.opensearch.sql.ast.expression.Compare;
 import org.opensearch.sql.ast.expression.Field;
-import org.opensearch.sql.ast.expression.Not;
+import org.opensearch.sql.ast.expression.Let;
 import org.opensearch.sql.ast.expression.QualifiedName;
 import org.opensearch.sql.ast.expression.UnresolvedExpression;
-import org.opensearch.sql.ast.expression.subquery.ExistsSubquery;
+import org.opensearch.sql.ast.expression.subquery.SubqueryExpression;
 import org.opensearch.sql.ast.tree.Aggregation;
 import org.opensearch.sql.ast.tree.Eval;
 import org.opensearch.sql.ast.tree.Filter;
@@ -91,14 +92,14 @@ private RelBuilder scan(RelOptTable tableSchema, CalcitePlanContext context) {
   @Override
   public RelNode visitFilter(Filter node, CalcitePlanContext context) {
     visitChildren(node, context);
-    boolean containsExistsSubquery = containsExistsSubquery(node.getCondition());
+    boolean containsSubqueryExpression = containsSubqueryExpression(node.getCondition());
     final Holder<@Nullable RexCorrelVariable> v = Holder.empty();
-    if (containsExistsSubquery) {
+    if (containsSubqueryExpression) {
       context.relBuilder.variable(v::set);
       context.pushCorrelVar(v.get());
     }
     RexNode condition = rexVisitor.analyze(node.getCondition(), context);
-    if (containsExistsSubquery) {
+    if (containsSubqueryExpression) {
       context.relBuilder.filter(ImmutableList.of(v.get().id), condition);
       context.popCorrelVar();
     } else {
@@ -107,15 +108,20 @@ public RelNode visitFilter(Filter node, CalcitePlanContext context) {
     return context.relBuilder.peek();
   }
 
-  private boolean containsExistsSubquery(Object condition) {
-    if (condition instanceof ExistsSubquery) {
+  private boolean containsSubqueryExpression(Node expr) {
+    if (expr == null) {
+      return false;
+    }
+    if (expr instanceof SubqueryExpression) {
       return true;
     }
-    if (condition instanceof Not n) {
-      return containsExistsSubquery(n.getExpression());
+    if (expr instanceof Let l) {
+      return containsSubqueryExpression(l.getExpression());
     }
-    if (condition instanceof Compare c) {
-      return containsExistsSubquery(c.getLeft()) || containsExistsSubquery(c.getRight());
+    for (Node child : expr.getChild()) {
+      if (containsSubqueryExpression(child)) {
+        return true;
+      }
     }
     return false;
   }
@@ -187,8 +193,25 @@ public RelNode visitEval(Eval node, CalcitePlanContext context) {
         node.getExpressionList().stream()
             .map(
                 expr -> {
+                  boolean containsSubqueryExpression = containsSubqueryExpression(expr);
+                  final Holder<@Nullable RexCorrelVariable> v = Holder.empty();
+                  if (containsSubqueryExpression) {
+                    context.relBuilder.variable(v::set);
+                    context.pushCorrelVar(v.get());
+                  }
                   RexNode eval = rexVisitor.analyze(expr, context);
-                  context.relBuilder.projectPlus(eval);
+                  if (containsSubqueryExpression) {
+                    // RelBuilder.projectPlus doesn't have a parameter with variablesSet:
+                    // projectPlus(Iterable<CorrelationId> variablesSet, RexNode... nodes)
+                    context.relBuilder.project(
+                        Iterables.concat(context.relBuilder.fields(), ImmutableList.of(eval)),
+                        ImmutableList.of(),
+                        false,
+                        ImmutableList.of(v.get().id));
+                    context.popCorrelVar();
+                  } else {
+                    context.relBuilder.projectPlus(eval);
+                  }
                   return eval;
                 })
             .collect(Collectors.toList());

core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java

@@ -18,17 +18,14 @@
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeFactory;
 import org.apache.calcite.rex.RexBuilder;
-import org.apache.calcite.rex.RexCorrelVariable;
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.sql.SqlIntervalQualifier;
 import org.apache.calcite.sql.fun.SqlStdOperatorTable;
 import org.apache.calcite.sql.parser.SqlParserUtil;
 import org.apache.calcite.sql.type.SqlTypeName;
 import org.apache.calcite.util.DateString;
-import org.apache.calcite.util.Holder;
 import org.apache.calcite.util.TimeString;
 import org.apache.calcite.util.TimestampString;
-import org.checkerframework.checker.nullness.qual.Nullable;
 import org.opensearch.sql.ast.AbstractNodeVisitor;
 import org.opensearch.sql.ast.expression.Alias;
 import org.opensearch.sql.ast.expression.And;
@@ -46,6 +43,7 @@
 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.tree.UnresolvedPlan;
 import org.opensearch.sql.calcite.utils.BuiltinFunctionUtils;
 import org.opensearch.sql.exception.SemanticCheckException;
@@ -283,7 +281,7 @@ public RexNode visitFunction(Function node, CalcitePlanContext context) {
   public RexNode visitInSubquery(InSubquery node, CalcitePlanContext context) {
     List<RexNode> nodes = node.getChild().stream().map(child -> analyze(child, context)).toList();
     UnresolvedPlan subquery = node.getQuery();
-    RelNode subqueryRel = resolveSubqueryPlan(subquery, false, context);
+    RelNode subqueryRel = resolveSubqueryPlan(subquery, context);
     try {
       return context.relBuilder.in(subqueryRel, nodes);
       // TODO
@@ -303,18 +301,25 @@ public RexNode visitInSubquery(InSubquery node, CalcitePlanContext context) {
     }
   }
 
+  @Override
+  public RexNode visitScalarSubquery(ScalarSubquery node, CalcitePlanContext context) {
+    return context.relBuilder.scalarQuery(
+        b -> {
+          UnresolvedPlan subquery = node.getQuery();
+          return resolveSubqueryPlan(subquery, context);
+        });
+  }
+
   @Override
   public RexNode visitExistsSubquery(ExistsSubquery node, CalcitePlanContext context) {
-    final Holder<@Nullable RexCorrelVariable> v = Holder.empty();
     return context.relBuilder.exists(
         b -> {
           UnresolvedPlan subquery = node.getQuery();
-          return resolveSubqueryPlan(subquery, true, context);
+          return resolveSubqueryPlan(subquery, context);
         });
   }
 
-  private RelNode resolveSubqueryPlan(
-      UnresolvedPlan subquery, boolean isExists, CalcitePlanContext context) {
+  private RelNode resolveSubqueryPlan(UnresolvedPlan subquery, CalcitePlanContext context) {
     // clear and store the outer state
     boolean isResolvingJoinConditionOuter = context.isResolvingJoinCondition();
     if (isResolvingJoinConditionOuter) {