Skip to content

Implement ppl scalar subquery command with Calcite #3392

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
LantaoJin merged 3 commits into from
Mar 7, 2025
Merged

Implement ppl scalar subquery command with Calcite #3392

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
991b793
Implement ppl scalar subquery command with Calcite
LantaoJin Mar 7, 2025
92a9d53
more general subquery checker
LantaoJin Mar 7, 2025
d321975
support correlated IN subquery
LantaoJin Mar 7, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
}
Expand Down
34 changes: 34 additions & 0 deletions core/src/main/java/org/opensearch/sql/ast/expression/subquery/ScalarSubquery.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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 = 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);
}

@Override
public List<UnresolvedExpression> getChild() {
return ImmutableList.of();
}
}
66 changes: 55 additions & 11 deletions core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -31,14 +32,15 @@
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.ScalarSubquery;
import org.opensearch.sql.ast.tree.Aggregation;
import org.opensearch.sql.ast.tree.Eval;
import org.opensearch.sql.ast.tree.Filter;
Expand Down Expand Up @@ -92,13 +94,14 @@ private RelBuilder scan(RelOptTable tableSchema, CalcitePlanContext context) {
public RelNode visitFilter(Filter node, CalcitePlanContext context) {
visitChildren(node, context);
boolean containsExistsSubquery = containsExistsSubquery(node.getCondition());
boolean containsScalarSubquery = containsScalarSubquery(node.getCondition());
final Holder<@Nullable RexCorrelVariable> v = Holder.empty();
if (containsExistsSubquery) {
if (containsExistsSubquery || containsScalarSubquery) {
context.relBuilder.variable(v::set);
context.pushCorrelVar(v.get());
}
RexNode condition = rexVisitor.analyze(node.getCondition(), context);
if (containsExistsSubquery) {
if (containsExistsSubquery || containsScalarSubquery) {
context.relBuilder.filter(ImmutableList.of(v.get().id), condition);
context.popCorrelVar();
} else {
Expand All @@ -107,15 +110,38 @@ public RelNode visitFilter(Filter node, CalcitePlanContext context) {
return context.relBuilder.peek();
}

private boolean containsExistsSubquery(Object condition) {
if (condition instanceof ExistsSubquery) {
private boolean containsExistsSubquery(Node expr) {
Copy link
Collaborator

@qianheng-aws qianheng-aws Mar 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How about combining this 2 flags together? Will we use them separately in specific cases?

Copy link
Member Author

@LantaoJin LantaoJin Mar 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

if (expr == null) {
return false;
}
if (expr instanceof ExistsSubquery) {
return true;
}
if (expr instanceof Let l) {
return containsExistsSubquery(l.getExpression());
}
for (Node child : expr.getChild()) {
if (containsExistsSubquery(child)) {
return true;
}
}
return false;
}

private boolean containsScalarSubquery(Node expr) {
if (expr == null) {
return false;
}
if (expr instanceof ScalarSubquery) {
return true;
}
if (condition instanceof Not n) {
return containsExistsSubquery(n.getExpression());
if (expr instanceof Let l) {
return containsScalarSubquery(l.getExpression());
}
if (condition instanceof Compare c) {
return containsExistsSubquery(c.getLeft()) || containsExistsSubquery(c.getRight());
for (Node child : expr.getChild()) {
if (containsScalarSubquery(child)) {
return true;
}
}
return false;
Copy link
Collaborator

@qianheng-aws qianheng-aws Mar 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we support more cases than Not, Let and Compare?

Copy link
Member Author

@LantaoJin LantaoJin Mar 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure, addressed in latest commit, please check it again

}
Expand Down Expand Up @@ -187,8 +213,26 @@ public RelNode visitEval(Eval node, CalcitePlanContext context) {
node.getExpressionList().stream()
.map(
expr -> {
boolean containsExistsSubquery = containsExistsSubquery(expr);
boolean containsScalarSubquery = containsScalarSubquery(expr);
final Holder<@Nullable RexCorrelVariable> v = Holder.empty();
if (containsExistsSubquery || containsScalarSubquery) {
context.relBuilder.variable(v::set);
context.pushCorrelVar(v.get());
}
RexNode eval = rexVisitor.analyze(expr, context);
context.relBuilder.projectPlus(eval);
if (containsExistsSubquery || containsScalarSubquery) {
// 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());
Expand Down
21 changes: 13 additions & 8 deletions core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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
Expand All @@ -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) {
Expand Down
Loading
Loading