Skip to content

Commit

Permalink
[Enhancement] limit clause support user variables (StarRocks#55257)
Browse files Browse the repository at this point in the history 
Signed-off-by: Seaven <seaven_7@qq.com>
  • Loading branch information
@Seaven
Seaven authored Jan 21, 2025
1 parent ea3ff40 commit 0a580fd
Show file tree
Hide file tree
Showing 6 changed files with 201 additions and 28 deletions.
49 changes: 31 additions & 18 deletions fe/fe-core/src/main/java/com/starrocks/analysis/LimitElement.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@

package com.starrocks.analysis;

import com.google.common.base.Preconditions;
import com.starrocks.sql.ast.AstVisitor;
import com.starrocks.sql.parser.NodePosition;

Expand All @@ -46,8 +47,8 @@ public class LimitElement implements ParseNode {
/////////////////////////////////////////
// BEGIN: Members that need to be reset()

private long limit;
private long offset;
private final Expr limit;
private final Expr offset;

// END: Members that need to be reset()
/////////////////////////////////////////
Expand All @@ -56,19 +57,21 @@ public class LimitElement implements ParseNode {

public LimitElement() {
pos = NodePosition.ZERO;
limit = -1;
offset = 0;
}

public LimitElement(long limit) {
this(0, limit, NodePosition.ZERO);
limit = new IntLiteral(-1);
offset = new IntLiteral(0);
}

public LimitElement(long offset, long limit) {
this(offset, limit, NodePosition.ZERO);
}

public LimitElement(long offset, long limit, NodePosition pos) {
this.pos = pos;
this.offset = new IntLiteral(offset);
this.limit = new IntLiteral(limit);
}

public LimitElement(Expr offset, Expr limit, NodePosition pos) {
this.pos = pos;
this.offset = offset;
this.limit = limit;
Expand All @@ -90,48 +93,58 @@ public LimitElement clone() {
* first. If no limit was set, then -1 is returned.
*/
public long getLimit() {
return limit;
Preconditions.checkState(limit instanceof LiteralExpr);
return ((LiteralExpr) limit).getLongValue();
}

public boolean hasLimit() {
return limit != -1;
return getLimit() != -1;
}

/**
* Returns the integer offset, evaluated from the offset expression. Must call
* analyze() first. If no offsetExpr exists, then 0 (the default offset) is returned.
*/
public long getOffset() {
return offset;
Preconditions.checkState(offset instanceof LiteralExpr);
return ((LiteralExpr) offset).getLongValue();
}

public boolean hasOffset() {
return offset != 0;
return getOffset() != 0;
}

public String toSql() {
if (limit == -1) {
if (getLimit() == -1) {
return "";
}
StringBuilder sb = new StringBuilder(" LIMIT ");
if (offset != 0) {
sb.append(offset).append(", ");
if (getOffset() != 0) {
sb.append(getOffset()).append(", ");
}
sb.append("").append(limit);
sb.append("").append(getLimit());
return sb.toString();
}

public Expr getLimitExpr() {
return limit;
}

public Expr getOffsetExpr() {
return offset;
}

@Override
public NodePosition getPos() {
return pos;
}

public String toDigest() {
if (limit == -1) {
if (getLimit() == -1) {
return "";
}
StringBuilder sb = new StringBuilder(" limit ");
if (offset != 0) {
if (getOffset() != 0) {
sb.append(" ?, ");
}
sb.append("").append(" ? ");
Expand Down
39 changes: 36 additions & 3 deletions fe/fe-core/src/main/java/com/starrocks/sql/analyzer/SelectAnalyzer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
import com.starrocks.analysis.OrderByElement;
import com.starrocks.analysis.ParseNode;
import com.starrocks.analysis.SlotRef;
import com.starrocks.analysis.UserVariableExpr;
import com.starrocks.catalog.FunctionSet;
import com.starrocks.catalog.PrimitiveType;
import com.starrocks.catalog.Type;
Expand Down Expand Up @@ -189,9 +190,7 @@ public void analyze(AnalyzeState analyzeState,
analyzeState.setOrderSourceExpressions(orderSourceExpressions);
}

if (limitElement != null && limitElement.hasLimit()) {
analyzeState.setLimit(new LimitElement(limitElement.getOffset(), limitElement.getLimit()));
}
analyzeState.setLimit(analyzeLimit(limitElement, analyzeState, sourceScope));
}

private List<Expr> analyzeSelect(SelectList selectList, Relation fromRelation, boolean hasGroupByClause,
Expand Down Expand Up @@ -574,6 +573,40 @@ private void analyzeHaving(Expr havingClause, AnalyzeState analyzeState,
}
}

private LimitElement analyzeLimit(LimitElement limitElement, AnalyzeState analyzeState, Scope scope) {
if (limitElement == null) {
return null;
}

Expr limitExpr = limitElement.getLimitExpr();
Expr offsetExpr = limitElement.getOffsetExpr();
long limit;
long offset;
analyzeExpression(limitExpr, analyzeState, scope);
analyzeExpression(offsetExpr, analyzeState, scope);
if (limitExpr.isLiteral()) {
limit = limitElement.getLimit();
} else if (limitExpr instanceof UserVariableExpr &&
((UserVariableExpr) limitExpr).getValue() instanceof IntLiteral) {
limit = ((IntLiteral) ((UserVariableExpr) limitExpr).getValue()).getLongValue();
} else {
throw new SemanticException("LIMIT clause %s must be number", limitExpr.toMySql());
}
if (limit == -1) {
return null;
}

if (offsetExpr.isLiteral()) {
offset = limitElement.getOffset();
} else if (offsetExpr instanceof UserVariableExpr &&
((UserVariableExpr) offsetExpr).getValue() instanceof IntLiteral) {
offset = ((IntLiteral) ((UserVariableExpr) offsetExpr).getValue()).getLongValue();
} else {
throw new SemanticException("OFFSET clause %s must be number", offsetExpr.toMySql());
}
return new LimitElement(offset, limit, limitElement.getPos());
}

// If alias is same with table column name, we directly use table name.
// otherwise, we use output expression according to the alias
public static class RewriteAliasVisitor implements AstVisitor<Expr, Void> {
Expand Down
24 changes: 20 additions & 4 deletions fe/fe-core/src/main/java/com/starrocks/sql/parser/AstBuilder.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5572,13 +5572,29 @@ private static boolean getOrderingType(Token token) {

@Override
public ParseNode visitLimitElement(StarRocksParser.LimitElementContext context) {
if (context.limit.getText().equals("?") || (context.offset != null && context.offset.getText().equals("?"))) {
if (context.limit.PARAMETER() != null || (context.offset != null && context.offset.PARAMETER() != null)) {
throw new ParsingException("using parameter(?) as limit or offset not supported");
}
long limit = Long.parseLong(context.limit.getText());
long offset = 0;

Expr limit;
Expr offset = new IntLiteral(0);

if (context.limit.INTEGER_VALUE() != null) {
limit = new IntLiteral(Long.parseLong(context.limit.INTEGER_VALUE().getText()));
} else if (context.limit.userVariable() != null) {
limit = (UserVariableExpr) visit(context.limit.userVariable());
} else {
throw new ParsingException("unsupported invalid limit value", createPos(context.limit));
}

if (context.offset != null) {
offset = Long.parseLong(context.offset.getText());
if (context.offset.INTEGER_VALUE() != null) {
offset = new IntLiteral(Long.parseLong(context.offset.INTEGER_VALUE().getText()));
} else if (context.offset.userVariable() != null) {
offset = (UserVariableExpr) visit(context.offset.userVariable());
} else {
throw new ParsingException("unsupported invalid offset value", createPos(context.offset));
}
}
return new LimitElement(offset, limit, createPos(context));
}
Expand Down
10 changes: 8 additions & 2 deletions fe/fe-core/src/main/java/com/starrocks/sql/parser/StarRocks.g4
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2160,9 +2160,15 @@ sortItem
: expression ordering = (ASC | DESC)? (NULLS nullOrdering=(FIRST | LAST))?
;

limitConstExpr
: INTEGER_VALUE
| PARAMETER
| userVariable
;

limitElement
: LIMIT limit =(INTEGER_VALUE|PARAMETER) (OFFSET offset=(INTEGER_VALUE|PARAMETER))?
| LIMIT offset =(INTEGER_VALUE|PARAMETER) ',' limit=(INTEGER_VALUE|PARAMETER)
: LIMIT limit=limitConstExpr (OFFSET offset=limitConstExpr)?
| LIMIT offset=limitConstExpr ',' limit=limitConstExpr
;

querySpecification
Expand Down
2 changes: 1 addition & 1 deletion fe/fe-core/src/test/java/com/starrocks/common/proc/OptimizeProcDirTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -212,7 +212,7 @@ public void testFetchResultByFilter() throws AnalysisException {
ArrayList<OrderByPair> orderByPairs = Lists.newArrayList();
orderByPairs.add(new OrderByPair(0));

LimitElement limitElement = new LimitElement(1);
LimitElement limitElement = new LimitElement(0, 1);

BaseProcResult result = (BaseProcResult) optimizeProcDir.fetchResultByFilter(
filter, orderByPairs, limitElement);
Expand Down
105 changes: 105 additions & 0 deletions fe/fe-core/src/test/java/com/starrocks/sql/plan/LimitTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,11 @@
import com.starrocks.catalog.Replica;
import com.starrocks.common.FeConstants;
import com.starrocks.qe.SessionVariable;
import com.starrocks.qe.SetExecutor;
import com.starrocks.server.GlobalStateMgr;
import com.starrocks.sql.analyzer.SemanticException;
import com.starrocks.sql.ast.SetStmt;
import com.starrocks.utframe.UtFrameUtils;
import mockit.Expectations;
import org.junit.Assert;
import org.junit.Test;
Expand Down Expand Up @@ -954,4 +958,105 @@ public void testTransformGroupByToLimit() throws Exception {
" 2:EXCHANGE\n" +
" limit: 1");
}

@Test
public void testLimitUserVariable() throws Exception {
{
String sql = "set @var = 123";
SetStmt stmt = (SetStmt) UtFrameUtils.parseStmtWithNewParser(sql, connectContext);
SetExecutor executor = new SetExecutor(connectContext, stmt);
executor.execute();

sql = "select * from t0 limit @var";
String plan = getFragmentPlan(sql);
assertContains(plan, "EXCHANGE\n" +
" limit: 123");
}
{
String sql = "set @var = 123";
SetStmt stmt = (SetStmt) UtFrameUtils.parseStmtWithNewParser(sql, connectContext);
SetExecutor executor = new SetExecutor(connectContext, stmt);
executor.execute();

sql = "select * from t0 limit @var, @var";
String plan = getFragmentPlan(sql);
assertContains(plan, "offset: 123\n" +
" limit: 123");
}
{
String sql = "set @var = 123";
SetStmt stmt = (SetStmt) UtFrameUtils.parseStmtWithNewParser(sql, connectContext);
SetExecutor executor = new SetExecutor(connectContext, stmt);
executor.execute();

sql = "select * from t0 limit @var OFFSET @var";
String plan = getFragmentPlan(sql);
assertContains(plan, "offset: 123\n" +
" limit: 123");
}
{
String sql = "set @var = 123";
SetStmt stmt = (SetStmt) UtFrameUtils.parseStmtWithNewParser(sql, connectContext);
SetExecutor executor = new SetExecutor(connectContext, stmt);
executor.execute();

sql = "select * from t0 limit @var, 11";
String plan = getFragmentPlan(sql);
assertContains(plan, "offset: 123\n" +
" limit: 11");
}
{
String sql = "set @var = 123";
SetStmt stmt = (SetStmt) UtFrameUtils.parseStmtWithNewParser(sql, connectContext);
SetExecutor executor = new SetExecutor(connectContext, stmt);
executor.execute();

sql = "select * from t0 limit 12, @var";
String plan = getFragmentPlan(sql);
assertContains(plan, "offset: 12\n" +
" limit: 123");
}
{
String sql = "set @var = 31 + 16";
SetStmt stmt = (SetStmt) UtFrameUtils.parseStmtWithNewParser(sql, connectContext);
SetExecutor executor = new SetExecutor(connectContext, stmt);
executor.execute();

sql = "select * from t0 limit 12, @var";
String plan = getFragmentPlan(sql);
assertContains(plan, "offset: 12\n" +
" limit: 47");
}
}

@Test
public void testLimitUserVariableError() throws Exception {
{
String sql = "set @var = '123'";
SetStmt stmt = (SetStmt) UtFrameUtils.parseStmtWithNewParser(sql, connectContext);
SetExecutor executor = new SetExecutor(connectContext, stmt);
executor.execute();

String tql = "select * from t0 limit @var";
Assert.assertThrows(SemanticException.class, () -> getFragmentPlan(tql));
}
{
String sql = "set @var = 'abc'";
SetStmt stmt = (SetStmt) UtFrameUtils.parseStmtWithNewParser(sql, connectContext);
SetExecutor executor = new SetExecutor(connectContext, stmt);
executor.execute();

String tql = "select * from t0 limit @var";
Assert.assertThrows(SemanticException.class, () -> getFragmentPlan(tql));
}
{
String sql = "set @var = 'abc'";
SetStmt stmt = (SetStmt) UtFrameUtils.parseStmtWithNewParser(sql, connectContext);
SetExecutor executor = new SetExecutor(connectContext, stmt);
executor.execute();

String tql = "select * from t0 limit @var, 2";
Assert.assertThrows(SemanticException.class, () -> getFragmentPlan(tql));
}
}
}

0 comments on commit 0a580fd

Please sign in to comment.