Skip to content

Commit

Permalink
[Feature] Support multi expression partition table
Browse files Browse the repository at this point in the history 
Signed-off-by: meegoo <[email protected]>
  • Loading branch information
@meegoo
meegoo committed Nov 10, 2024
1 parent 868f415 commit b72ec81
Show file tree
Hide file tree
Showing 9 changed files with 291 additions and 20 deletions.
2 changes: 2 additions & 0 deletions fe/fe-core/src/main/java/com/starrocks/common/FeConstants.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -95,6 +95,8 @@ public class FeConstants {
// the raw data of one tablet equals to 10GB approximately
public static final long AUTO_DISTRIBUTION_UNIT = 3221225472L;

public static final String GENERATED_PARTITION_COLUMN_PREFIX = "__generated_partition_column_";

// Max counter num of TOP K function
public static final int MAX_COUNTER_NUM_OF_TOP_K = 100000;

Expand Down
11 changes: 5 additions & 6 deletions fe/fe-core/src/main/java/com/starrocks/sql/analyzer/CreateTableAnalyzer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,6 @@
import com.starrocks.qe.ConnectContext;
import com.starrocks.server.CatalogMgr;
import com.starrocks.server.GlobalStateMgr;
import com.starrocks.server.RunMode;
import com.starrocks.server.TemporaryTableMgr;
import com.starrocks.sql.ast.ColumnDef;
import com.starrocks.sql.ast.CreateTableStmt;
Expand Down Expand Up @@ -294,6 +293,10 @@ private static void analyzeKeysDesc(CreateTableStmt stmt) {
} else {
int keyLength = 0;
for (ColumnDef columnDef : columnDefs) {
// generated column should not be key
if (columnDef.isGeneratedColumn()) {
break;
}
keyLength += columnDef.getType().getIndexSize();
if (keysColumnNames.size() >= FeConstants.SHORTKEY_MAX_COLUMN_COUNT
|| keyLength > FeConstants.SHORTKEY_MAXSIZE_BYTES) {
Expand Down Expand Up @@ -571,10 +574,6 @@ public static void analyzeGeneratedColumn(CreateTableStmt stmt, ConnectContext c
throw new SemanticException("Generated Column does not support AGG table");
}

if (RunMode.isSharedDataMode()) {
throw new SemanticException("Does not support generated column in shared data cluster yet");
}

final TableName tableNameObject = stmt.getDbTbl();

List<Column> columns = stmt.getColumns();
Expand All @@ -591,7 +590,7 @@ public static void analyzeGeneratedColumn(CreateTableStmt stmt, ConnectContext c

if (column.isGeneratedColumn()) {
if (keysDesc.containsCol(column.getName())) {
throw new SemanticException("Generated Column can not be KEY");
throw new SemanticException("Generated Column " + column.getName() + " can not be KEY");
}

Expr expr = column.getGeneratedColumnExpr(columns);
Expand Down
6 changes: 4 additions & 2 deletions fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rewrite/ScalarOperatorEvaluator.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -222,15 +222,17 @@ public boolean isMonotonicFunction(CallOperator call) {

FunctionInvoker invoker = functions.get(signature);

return invoker != null && invoker.isMonotonic;
return invoker != null && isMonotonicFunc(invoker, call);
}

private boolean isMonotonicFunc(FunctionInvoker invoker, CallOperator operator) {
if (!invoker.isMonotonic) {
return false;
}

if (FunctionSet.DATE_FORMAT.equalsIgnoreCase(invoker.getSignature().getName())) {
if (FunctionSet.DATE_FORMAT.equalsIgnoreCase(invoker.getSignature().getName())
|| (FunctionSet.FROM_UNIXTIME.equalsIgnoreCase(invoker.getSignature().getName())
&& operator.getChildren().size() == 2)) {
String pattern = operator.getChild(1).toString();
if (pattern.isEmpty()) {
return true;
Expand Down
12 changes: 6 additions & 6 deletions fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rewrite/ScalarOperatorFunctions.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -577,8 +577,8 @@ public static ConstantOperator unixTimestampNow() {
}

@ConstantFunction.List(list = {
@ConstantFunction(name = "unix_timestamp", argTypes = {DATETIME}, returnType = BIGINT),
@ConstantFunction(name = "unix_timestamp", argTypes = {DATE}, returnType = BIGINT)
@ConstantFunction(name = "unix_timestamp", argTypes = {DATETIME}, returnType = BIGINT, isMonotonic = true),
@ConstantFunction(name = "unix_timestamp", argTypes = {DATE}, returnType = BIGINT, isMonotonic = true)
})
public static ConstantOperator unixTimestamp(ConstantOperator arg) {
LocalDateTime dt = arg.getDatetime();
Expand All @@ -591,8 +591,8 @@ public static ConstantOperator unixTimestamp(ConstantOperator arg) {
}

@ConstantFunction.List(list = {
@ConstantFunction(name = "from_unixtime", argTypes = {INT}, returnType = VARCHAR),
@ConstantFunction(name = "from_unixtime", argTypes = {BIGINT}, returnType = VARCHAR)
@ConstantFunction(name = "from_unixtime", argTypes = {INT}, returnType = VARCHAR, isMonotonic = true),
@ConstantFunction(name = "from_unixtime", argTypes = {BIGINT}, returnType = VARCHAR, isMonotonic = true)
})
public static ConstantOperator fromUnixTime(ConstantOperator unixTime) throws AnalysisException {
long value = 0;
Expand Down Expand Up @@ -627,8 +627,8 @@ public static ConstantOperator fromUnixTimeMs(ConstantOperator unixTime) throws
}

@ConstantFunction.List(list = {
@ConstantFunction(name = "from_unixtime", argTypes = {INT, VARCHAR}, returnType = VARCHAR),
@ConstantFunction(name = "from_unixtime", argTypes = {BIGINT, VARCHAR}, returnType = VARCHAR)
@ConstantFunction(name = "from_unixtime", argTypes = {INT, VARCHAR}, returnType = VARCHAR, isMonotonic = true),
@ConstantFunction(name = "from_unixtime", argTypes = {BIGINT, VARCHAR}, returnType = VARCHAR, isMonotonic = true)
})
public static ConstantOperator fromUnixTime(ConstantOperator unixTime, ConstantOperator fmtLiteral)
throws AnalysisException {
Expand Down
86 changes: 81 additions & 5 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 @@ -101,6 +101,7 @@
import com.starrocks.common.Config;
import com.starrocks.common.CsvFormat;
import com.starrocks.common.DdlException;
import com.starrocks.common.FeConstants;
import com.starrocks.common.NotImplementedException;
import com.starrocks.common.Pair;
import com.starrocks.common.profile.Tracers;
Expand All @@ -115,9 +116,14 @@
import com.starrocks.scheduler.persist.TaskSchedule;
import com.starrocks.server.WarehouseManager;
import com.starrocks.sql.ShowTemporaryTableStmt;
import com.starrocks.sql.analyzer.AnalyzeState;
import com.starrocks.sql.analyzer.AnalyzerUtils;
import com.starrocks.sql.analyzer.ExpressionAnalyzer;
import com.starrocks.sql.analyzer.Field;
import com.starrocks.sql.analyzer.FunctionAnalyzer;
import com.starrocks.sql.analyzer.RelationFields;
import com.starrocks.sql.analyzer.RelationId;
import com.starrocks.sql.analyzer.Scope;
import com.starrocks.sql.analyzer.SemanticException;
import com.starrocks.sql.ast.AddBackendBlackListStmt;
import com.starrocks.sql.ast.AddBackendClause;
Expand Down Expand Up @@ -777,7 +783,7 @@ public ParseNode visitCreateTableStatement(StarRocksParser.CreateTableStatementC
context.charsetDesc() == null ? null :
((Identifier) visit(context.charsetDesc().identifierOrString())).getValue(),
context.keyDesc() == null ? null : getKeysDesc(context.keyDesc()),
context.partitionDesc() == null ? null : getPartitionDesc(context.partitionDesc(), columnDefs),
context.partitionDesc() == null ? null : getPartitionDesc(context.partitionDesc(), columnDefs, tableName),
context.distributionDesc() == null ? null : (DistributionDesc) visit(context.distributionDesc()),
properties,
extProperties,
Expand All @@ -803,7 +809,7 @@ public ParseNode visitCreateTableStatement(StarRocksParser.CreateTableStatementC
context.charsetDesc() == null ? null :
((Identifier) visit(context.charsetDesc().identifierOrString())).getValue(),
context.keyDesc() == null ? null : getKeysDesc(context.keyDesc()),
context.partitionDesc() == null ? null : getPartitionDesc(context.partitionDesc(), columnDefs),
context.partitionDesc() == null ? null : getPartitionDesc(context.partitionDesc(), columnDefs, tableName),
context.distributionDesc() == null ? null : (DistributionDesc) visit(context.distributionDesc()),
properties,
extProperties,
Expand All @@ -815,10 +821,64 @@ public ParseNode visitCreateTableStatement(StarRocksParser.CreateTableStatementC
.stream().map(Identifier::getValue).collect(toList()));
}

private PartitionDesc getPartitionDesc(StarRocksParser.PartitionDescContext context, List<ColumnDef> columnDefs) {
private PartitionDesc generateMulitListPartitionDesc(StarRocksParser.PartitionDescContext context,
List<ParseNode> multiDescList, List<ColumnDef> columnDefs,
TableName tableName) {
List<String> columnList = Lists.newArrayList();
List<PartitionDesc> partitionDescList = new ArrayList<>();
int placeHolderSlotId = 0;
for (ParseNode partitionExpr : multiDescList) {
if (partitionExpr instanceof Identifier) {
Identifier identifier = (Identifier) partitionExpr;
columnList.add(identifier.getValue());
}
if (partitionExpr instanceof FunctionCallExpr) {
FunctionCallExpr expr = (FunctionCallExpr) partitionExpr;
ExpressionAnalyzer.analyzeExpression(expr, new AnalyzeState(), new Scope(RelationId.anonymous(),
new RelationFields(columnDefs.stream().map(col -> new Field(col.getName(),
col.getType(), tableName, null)).collect(Collectors.toList()))),
new ConnectContext());
String columnName = FeConstants.GENERATED_PARTITION_COLUMN_PREFIX + placeHolderSlotId++;
columnList.add(columnName);
Type type = expr.getType();
if (type.isScalarType()) {
ScalarType scalarType = (ScalarType) type;
if (scalarType.isWildcardChar()) {
type = ScalarType.createCharType(ScalarType.getOlapMaxVarcharLength());
} else if (scalarType.isWildcardVarchar()) {
type = ScalarType.createVarcharType(ScalarType.getOlapMaxVarcharLength());
}
}
TypeDef typeDef = new TypeDef(type);
try {
typeDef.analyze();
} catch (Exception e) {
throw new ParsingException("Generate partition column " + columnName + " error: "
+ e.getMessage(), createPos(context));
}
ColumnDef generatedPartitionColumn = new ColumnDef(
columnName, typeDef, null, false, null, null, true,
ColumnDef.DefaultValueDef.NOT_SET, null, expr, "");
columnDefs.add(generatedPartitionColumn);
}
}
ListPartitionDesc listPartitionDesc = new ListPartitionDesc(columnList, partitionDescList);
listPartitionDesc.setAutoPartitionTable(true);
return listPartitionDesc;
}

private PartitionDesc getPartitionDesc(StarRocksParser.PartitionDescContext context,
List<ColumnDef> columnDefs, TableName tableName) {
List<PartitionDesc> partitionDescList = new ArrayList<>();
// for automatic partition
if (context.functionCall() != null) {
FunctionCallExpr functionCallExpr = (FunctionCallExpr) visit(context.functionCall());
String functionName = functionCallExpr.getFnName().getFunction();
// except date_trunc, time_slice, str_to_date use generated column as partition column
if (!FunctionSet.DATE_TRUNC.equals(functionName) && !FunctionSet.TIME_SLICE.equals(functionName)
&& !FunctionSet.STR2DATE.equals(functionName)) {
return generateMulitListPartitionDesc(context, Lists.newArrayList(functionCallExpr), columnDefs, tableName);
}
String currentGranularity = null;
for (StarRocksParser.RangePartitionDescContext rangePartitionDescContext : context.rangePartitionDesc()) {
final PartitionDesc rangePartitionDesc = (PartitionDesc) visit(rangePartitionDescContext);
Expand All @@ -836,7 +896,6 @@ private PartitionDesc getPartitionDesc(StarRocksParser.PartitionDescContext cont
}
partitionDescList.add(rangePartitionDesc);
}
FunctionCallExpr functionCallExpr = (FunctionCallExpr) visit(context.functionCall());
List<String> columnList = AnalyzerUtils.checkAndExtractPartitionCol(functionCallExpr, columnDefs);
AnalyzerUtils.checkAutoPartitionTableLimit(functionCallExpr, currentGranularity);
RangePartitionDesc rangePartitionDesc = new RangePartitionDesc(columnList, partitionDescList);
Expand Down Expand Up @@ -865,6 +924,23 @@ private PartitionDesc getPartitionDesc(StarRocksParser.PartitionDescContext cont
}
return new ExpressionPartitionDesc(rangePartitionDesc, primaryExpression);
}
if (context.identifierList() == null) {
if (context.partitionExpr() != null) {
List<ParseNode> multiDescList = Lists.newArrayList();
for (StarRocksParser.PartitionExprContext partitionExpr : context.partitionExpr()) {
if (partitionExpr.identifier() != null) {
Identifier identifier = (Identifier) visit(partitionExpr.identifier());
multiDescList.add(identifier);
} else if (partitionExpr.functionCall() != null) {
FunctionCallExpr expr = (FunctionCallExpr) visit(partitionExpr.functionCall());
multiDescList.add(expr);
} else {
throw new ParsingException("Partition column list is empty", createPos(context));
}
}
return generateMulitListPartitionDesc(context, multiDescList, columnDefs, tableName);
}
}
List<Identifier> identifierList = visit(context.identifierList().identifier(), Identifier.class);
List<String> columnList = identifierList.stream().map(Identifier::getValue).collect(toList());
if (context.RANGE() != null) {
Expand Down Expand Up @@ -4274,7 +4350,7 @@ public ParseNode visitModifyPropertiesClause(StarRocksParser.ModifyPropertiesCla
public ParseNode visitOptimizeClause(StarRocksParser.OptimizeClauseContext context) {
return new OptimizeClause(
context.keyDesc() == null ? null : getKeysDesc(context.keyDesc()),
context.partitionDesc() == null ? null : getPartitionDesc(context.partitionDesc(), null),
context.partitionDesc() == null ? null : getPartitionDesc(context.partitionDesc(), null, null),
context.distributionDesc() == null ? null : (DistributionDesc) visit(context.distributionDesc()),
context.orderByDesc() == null ? null :
visit(context.orderByDesc().identifierList().identifier(), Identifier.class)
Expand Down
6 changes: 6 additions & 0 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 @@ -2524,13 +2524,19 @@ optimizerTrace
: TRACE (ALL | LOGS | TIMES | VALUES | REASON) identifier?
;

partitionExpr
: identifier
| functionCall
;

partitionDesc
: PARTITION BY RANGE identifierList '(' (rangePartitionDesc (',' rangePartitionDesc)*)? ')'
| PARTITION BY RANGE primaryExpression '(' (rangePartitionDesc (',' rangePartitionDesc)*)? ')'
| PARTITION BY LIST? identifierList '(' (listPartitionDesc (',' listPartitionDesc)*)? ')'
| PARTITION BY LIST? identifierList
| PARTITION BY functionCall '(' (rangePartitionDesc (',' rangePartitionDesc)*)? ')'
| PARTITION BY functionCall
| PARTITION BY partitionExpr (',' partitionExpr)*
;

listPartitionDesc
Expand Down
Loading

0 comments on commit b72ec81

Please sign in to comment.