[feature](nereids) add scalar subquery expression (#11332)

scalar subquery:
A subquery that will return only one row and one column.

A limit has been added, where a = subquery returns a result of 1 row and 1 column.
Here, the first limit is to return 1 column.

TODO: the subsequent return will limit the return to 1 row
This commit is contained in:
zhengshiJ
2022-07-29 21:05:19 +08:00
committed by GitHub
parent 8483660fe7
commit e7fae413dd
9 changed files with 169 additions and 49 deletions

View File

@ -89,6 +89,7 @@ import org.apache.doris.nereids.trees.expressions.IntervalLiteral;
import org.apache.doris.nereids.trees.expressions.LessThan;
import org.apache.doris.nereids.trees.expressions.LessThanEqual;
import org.apache.doris.nereids.trees.expressions.Like;
import org.apache.doris.nereids.trees.expressions.ListQuery;
import org.apache.doris.nereids.trees.expressions.Literal;
import org.apache.doris.nereids.trees.expressions.Mod;
import org.apache.doris.nereids.trees.expressions.Multiply;
@ -98,8 +99,8 @@ import org.apache.doris.nereids.trees.expressions.NullLiteral;
import org.apache.doris.nereids.trees.expressions.NullSafeEqual;
import org.apache.doris.nereids.trees.expressions.Or;
import org.apache.doris.nereids.trees.expressions.Regexp;
import org.apache.doris.nereids.trees.expressions.ScalarSubquery;
import org.apache.doris.nereids.trees.expressions.StringLiteral;
import org.apache.doris.nereids.trees.expressions.SubqueryExpr;
import org.apache.doris.nereids.trees.expressions.Subtract;
import org.apache.doris.nereids.trees.expressions.TimestampArithmetic;
import org.apache.doris.nereids.trees.expressions.WhenClause;
@ -804,7 +805,7 @@ public class LogicalPlanBuilder extends DorisParserBaseVisitor<Object> {
} else {
outExpression = new InSubquery(
valueExpression,
typedVisit(ctx.query())
new ListQuery(typedVisit(ctx.query()))
);
}
break;
@ -831,7 +832,7 @@ public class LogicalPlanBuilder extends DorisParserBaseVisitor<Object> {
@Override
public Expression visitSubqueryExpression(SubqueryExpressionContext subqueryExprCtx) {
return ParserUtils.withOrigin(subqueryExprCtx, () -> new SubqueryExpr(typedVisit(subqueryExprCtx.query())));
return ParserUtils.withOrigin(subqueryExprCtx, () -> new ScalarSubquery(typedVisit(subqueryExprCtx.query())));
}
@Override

View File

@ -23,15 +23,12 @@ import org.apache.doris.nereids.trees.plans.logical.LogicalPlan;
import org.apache.doris.nereids.types.BooleanType;
import org.apache.doris.nereids.types.DataType;
import com.google.common.base.Preconditions;
import java.util.List;
import java.util.Objects;
/**
* Exists subquery expression.
*/
public class Exists extends SubqueryExpr {
public class Exists extends SubqueryExpr implements LeafExpression {
public Exists(LogicalPlan subquery) {
super(Objects.requireNonNull(subquery, "subquery can not be null"));
@ -55,28 +52,4 @@ public class Exists extends SubqueryExpr {
public <R, C> R accept(ExpressionVisitor<R, C> visitor, C context) {
return visitor.visitExistsSubquery(this, context);
}
@Override
public Expression withChildren(List<Expression> children) {
Preconditions.checkArgument(children.size() == 1);
Preconditions.checkArgument(children.get(0) instanceof SubqueryExpr);
return new Exists(((SubqueryExpr) children.get(0)).getQueryPlan());
}
@Override
public boolean equals(Object o) {
if (this == o) {
return true;
}
if (o == null || getClass() != o.getClass()) {
return false;
}
Exists exists = (Exists) o;
return Objects.equals(this.queryPlan, exists.getQueryPlan());
}
@Override
public int hashCode() {
return Objects.hash(this.queryPlan);
}
}

View File

@ -19,8 +19,6 @@ package org.apache.doris.nereids.trees.expressions;
import org.apache.doris.nereids.exceptions.UnboundException;
import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;
import org.apache.doris.nereids.trees.plans.logical.LogicalPlan;
import org.apache.doris.nereids.types.BooleanType;
import org.apache.doris.nereids.types.DataType;
import com.google.common.base.Preconditions;
@ -33,15 +31,17 @@ import java.util.Objects;
*/
public class InSubquery extends SubqueryExpr implements BinaryExpression {
private Expression compareExpr;
private ListQuery listQuery;
public InSubquery(Expression compareExpression, LogicalPlan subquery) {
super(Objects.requireNonNull(subquery, "subquery can not be null"));
public InSubquery(Expression compareExpression, ListQuery listQuery) {
super(Objects.requireNonNull(listQuery.getQueryPlan(), "subquery can not be null"));
this.compareExpr = compareExpression;
this.listQuery = listQuery;
}
@Override
public DataType getDataType() throws UnboundException {
return BooleanType.INSTANCE;
return listQuery.getDataType();
}
@Override
@ -71,8 +71,8 @@ public class InSubquery extends SubqueryExpr implements BinaryExpression {
public Expression withChildren(List<Expression> children) {
Preconditions.checkArgument(children.size() == 2);
Preconditions.checkArgument(children.get(0) instanceof Expression);
Preconditions.checkArgument(children.get(1) instanceof SubqueryExpr);
return new InSubquery(children.get(0), ((SubqueryExpr) children.get(1)).getQueryPlan());
Preconditions.checkArgument(children.get(1) instanceof ListQuery);
return new InSubquery(children.get(0), (ListQuery) children.get(1));
}
@Override
@ -85,11 +85,11 @@ public class InSubquery extends SubqueryExpr implements BinaryExpression {
}
InSubquery inSubquery = (InSubquery) o;
return Objects.equals(this.compareExpr, inSubquery.getCompareExpr())
&& Objects.equals(this.queryPlan, inSubquery.getQueryPlan());
&& checkEquals(this.queryPlan, inSubquery.queryPlan);
}
@Override
public int hashCode() {
return Objects.hash(this.compareExpr, this.queryPlan);
return Objects.hash(this.compareExpr, this.listQuery);
}
}

View File

@ -0,0 +1,57 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
package org.apache.doris.nereids.trees.expressions;
import org.apache.doris.nereids.exceptions.UnboundException;
import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;
import org.apache.doris.nereids.trees.plans.logical.LogicalPlan;
import org.apache.doris.nereids.types.DataType;
import java.util.Objects;
/**
* Encapsulate LogicalPlan as Expression.
* just for subquery.
*/
public class ListQuery extends SubqueryExpr implements LeafExpression {
public ListQuery(LogicalPlan subquery) {
super(Objects.requireNonNull(subquery, "subquery can not be null"));
}
@Override
public DataType getDataType() throws UnboundException {
// TODO:
// For multiple lines, struct type is returned, in the form of splicing,
// but it seems that struct type is not currently supported
throw new UnboundException("not support");
}
@Override
public String toSql() {
return " (LISTQUERY) " + super.toSql();
}
@Override
public String toString() {
return " (LISTQUERY) " + super.toString();
}
public <R, C> R accept(ExpressionVisitor<R, C> visitor, C context) {
return visitor.visitListQuery(this, context);
}
}

View File

@ -0,0 +1,56 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
package org.apache.doris.nereids.trees.expressions;
import org.apache.doris.nereids.exceptions.UnboundException;
import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;
import org.apache.doris.nereids.trees.plans.logical.LogicalPlan;
import org.apache.doris.nereids.types.DataType;
import com.google.common.base.Preconditions;
import java.util.Objects;
/**
* A subquery that will return only one row and one column.
*/
public class ScalarSubquery extends SubqueryExpr implements LeafExpression {
public ScalarSubquery(LogicalPlan subquery) {
super(Objects.requireNonNull(subquery, "subquery can not be null"));
}
@Override
public DataType getDataType() throws UnboundException {
Preconditions.checkArgument(queryPlan.getOutput().size() == 1);
return queryPlan.getOutput().get(0).getDataType();
}
@Override
public String toSql() {
return " (SCALARSUBQUERY) " + super.toSql();
}
@Override
public String toString() {
return " (SCALARSUBQUERY) " + super.toString();
}
public <R, C> R accept(ExpressionVisitor<R, C> visitor, C context) {
return visitor.visitScalarSubquery(this, context);
}
}

View File

@ -39,10 +39,6 @@ public class SubqueryExpr extends Expression {
@Override
public DataType getDataType() throws UnboundException {
// TODO:
// Returns the data type of the row on a single line
// For multiple lines, struct type is returned, in the form of splicing,
// but it seems that struct type is not currently supported
throw new UnboundException("not support");
}
@ -95,7 +91,7 @@ public class SubqueryExpr extends Expression {
* @param o compared query.
* @return equal ? true : false;
*/
private boolean checkEquals(Object i, Object o) {
protected boolean checkEquals(Object i, Object o) {
if (!(i instanceof LogicalPlan) || !(o instanceof LogicalPlan)) {
return false;
}

View File

@ -18,9 +18,12 @@
package org.apache.doris.nereids.trees.expressions.visitor;
import org.apache.doris.nereids.analyzer.NereidsAnalyzer;
import org.apache.doris.nereids.exceptions.AnalysisException;
import org.apache.doris.nereids.rules.analysis.Scope;
import org.apache.doris.nereids.trees.expressions.Expression;
import org.apache.doris.nereids.trees.expressions.InSubquery;
import org.apache.doris.nereids.trees.expressions.ListQuery;
import org.apache.doris.nereids.trees.expressions.ScalarSubquery;
import org.apache.doris.nereids.trees.expressions.SubqueryExpr;
import org.apache.doris.nereids.trees.plans.logical.LogicalPlan;
import org.apache.doris.qe.ConnectContext;
@ -44,7 +47,17 @@ public class DefaultSubExprRewriter<C> extends DefaultExpressionRewriter<C> {
@Override
public Expression visitInSubquery(InSubquery expr, C context) {
return new InSubquery(expr.getCompareExpr(), analyzeSubquery(expr));
return new InSubquery(expr.getCompareExpr(), new ListQuery(analyzeSubquery(expr)));
}
@Override
public Expression visitScalarSubquery(ScalarSubquery scalar, C context) {
LogicalPlan analyzed = analyzeSubquery(scalar);
if (analyzed.getOutput().size() != 1) {
throw new AnalysisException("Multiple columns returned by subquery are not yet supported. Found "
+ analyzed.getOutput().size());
}
return new ScalarSubquery(analyzed);
}
private LogicalPlan analyzeSubquery(SubqueryExpr expr) {

View File

@ -45,6 +45,7 @@ import org.apache.doris.nereids.trees.expressions.IntegerLiteral;
import org.apache.doris.nereids.trees.expressions.LessThan;
import org.apache.doris.nereids.trees.expressions.LessThanEqual;
import org.apache.doris.nereids.trees.expressions.Like;
import org.apache.doris.nereids.trees.expressions.ListQuery;
import org.apache.doris.nereids.trees.expressions.Literal;
import org.apache.doris.nereids.trees.expressions.Mod;
import org.apache.doris.nereids.trees.expressions.Multiply;
@ -54,6 +55,7 @@ import org.apache.doris.nereids.trees.expressions.NullLiteral;
import org.apache.doris.nereids.trees.expressions.NullSafeEqual;
import org.apache.doris.nereids.trees.expressions.Or;
import org.apache.doris.nereids.trees.expressions.Regexp;
import org.apache.doris.nereids.trees.expressions.ScalarSubquery;
import org.apache.doris.nereids.trees.expressions.Slot;
import org.apache.doris.nereids.trees.expressions.SlotReference;
import org.apache.doris.nereids.trees.expressions.StringLiteral;
@ -240,6 +242,14 @@ public abstract class ExpressionVisitor<R, C> {
return visit(arithmetic, context);
}
public R visitScalarSubquery(ScalarSubquery scalar, C context) {
return visitSubqueryExpr(scalar, context);
}
public R visitListQuery(ListQuery listQuery, C context) {
return visitSubqueryExpr(listQuery, context);
}
/* ********************************************************************************************
* Unbound expressions
* ********************************************************************************************/

View File

@ -83,16 +83,30 @@ public class SubqueryTest extends AnalyzeCheckTestBase {
@Test
public void scalarTest() {
String sql = "select * from t0 where t0.id = "
// min will be rewritten as as, it needs to be bound
/*String sql = "select * from t0 where t0.id = "
+ "(select min(t1.id) from t1 where t0.k1 = t1.k1)";
checkAnalyze(sql);
checkAnalyze(sql);*/
// Require that the return value in the where subquery must have only 1 column.
String sql1 = "select * from t0 where t0.id = "
+ "(select t1.k1, t1.id from t1 where t0.k1 = t1.k1)";
assert sql1 != null;
String sql2 = "select * from t0 where t0.id = "
+ "(select t1.k1 from t1 where t0.k1 = t1.k1)";
checkAnalyze(sql2);
String sql3 = "select * from t0 where t0.id = "
+ "(select * from t1 where t0.k1 = t1.k1)";
assert sql3 != null;
}
@Test
public void inScalarTest() {
String sql = "select * from t0 where t0.id in "
+ "(select * from t1 where t1.k1 = "
+ "(select * from t2 where t0.id = t2.id));";
+ "(select t2.id from t2 where t0.id = t2.id));";
checkAnalyze(sql);
}
}