database/doris
2
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

[feature](Nereids) support add constraint on table (#27627)

Browse Source 
support add constraint on the table including
- primary key constraint
- unique constrain
- foreign key constraint
This commit is contained in:
谢健
2023-12-01 13:28:48 +08:00
committed by GitHub
parent 48d7df205f
commit e868c990ff
16 changed files with 821 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisLexer.g4
View File

@ -161,6 +161,7 @@ CONFIG: 'CONFIG';
CONNECTION: 'CONNECTION';
CONNECTION_ID: 'CONNECTION_ID';
CONSISTENT: 'CONSISTENT';
CONSTRAINT: 'CONSTRAINT';
CONVERT: 'CONVERT';
COPY: 'COPY';
COUNT: 'COUNT';
@ -251,6 +252,7 @@ FLOAT: 'FLOAT';
FOLLOWER: 'FOLLOWER';
FOLLOWING: 'FOLLOWING';
FOR: 'FOR';
FOREIGN: 'FOREIGN';
FORCE: 'FORCE';
FORMAT: 'FORMAT';
FREE: 'FREE';
@ -401,6 +403,7 @@ PLUGINS: 'PLUGINS';
POLICY: 'POLICY';
PRECEDING: 'PRECEDING';
PREPARE: 'PREPARE';
PRIMARY: 'PRIMARY';
PROC: 'PROC';
PROCEDURE: 'PROCEDURE';
PROCESSLIST: 'PROCESSLIST';
@ -419,6 +422,7 @@ REBALANCE: 'REBALANCE';
RECOVER: 'RECOVER';
RECYCLE: 'RECYCLE';
REFRESH: 'REFRESH';
REFERENCES: 'REFERENCES';
REGEXP: 'REGEXP';
RELEASE: 'RELEASE';
RENAME: 'RENAME';

13
fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4
View File

@ -96,6 +96,19 @@ statement
| (REFRESH (refreshMethod | refreshTrigger | refreshMethod refreshTrigger))
| (SET LEFT_PAREN fileProperties=propertyItemList RIGHT_PAREN)) #alterMTMV
| DROP MATERIALIZED VIEW (IF EXISTS)? mvName=multipartIdentifier #dropMTMV
| ALTER TABLE table=relation
ADD CONSTRAINT constraintName=errorCapturingIdentifier
constraint #addConstraint
| ALTER TABLE table=relation
DROP CONSTRAINT constraintName=errorCapturingIdentifier #dropConstraint
;
constraint
: PRIMARY KEY LEFT_PAREN slots+=errorCapturingIdentifier (COMMA slots+=errorCapturingIdentifier)* RIGHT_PAREN
| UNIQUE LEFT_PAREN slots+=errorCapturingIdentifier (COMMA slots+=errorCapturingIdentifier)* RIGHT_PAREN
| FOREIGN KEY LEFT_PAREN slots+=errorCapturingIdentifier (COMMA slots+=errorCapturingIdentifier)* RIGHT_PAREN
REFERENCES referenceTable=multipartIdentifier
LEFT_PAREN referenceSlots+=errorCapturingIdentifier (COMMA referenceSlots+=errorCapturingIdentifier)* RIGHT_PAREN
;
dataDesc

14
fe/fe-core/src/main/java/org/apache/doris/catalog/Table.java
View File

@ -19,6 +19,7 @@ package org.apache.doris.catalog;
import org.apache.doris.alter.AlterCancelException;
import org.apache.doris.analysis.CreateTableStmt;
import org.apache.doris.catalog.constraint.Constraint;
import org.apache.doris.common.DdlException;
import org.apache.doris.common.MetaNotFoundException;
import org.apache.doris.common.io.Text;
@ -48,6 +49,7 @@ import java.io.DataOutput;
import java.io.IOException;
import java.time.Instant;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
@ -115,6 +117,9 @@ public abstract class Table extends MetaObject implements Writable, TableIf {
// sql for creating this table, default is "";
protected String ddlSql = "";
@SerializedName(value = "constraints")
private HashMap<String, Constraint> constraintsMap = new HashMap<>();
public Table(TableType type) {
this.type = type;
this.fullSchema = Lists.newArrayList();
@ -291,6 +296,15 @@ public abstract class Table extends MetaObject implements Writable, TableIf {
}
}
public Constraint getConstraint(String name) {
return constraintsMap.get(name);
}
@Override
public Map<String, Constraint> getConstraintsMap() {
return constraintsMap;
}
public TableType getType() {
return type;
}

83
fe/fe-core/src/main/java/org/apache/doris/catalog/TableIf.java
View File

@ -18,14 +18,21 @@
package org.apache.doris.catalog;
import org.apache.doris.alter.AlterCancelException;
import org.apache.doris.catalog.constraint.Constraint;
import org.apache.doris.catalog.constraint.ForeignKeyConstraint;
import org.apache.doris.catalog.constraint.PrimaryKeyConstraint;
import org.apache.doris.catalog.constraint.UniqueConstraint;
import org.apache.doris.common.DdlException;
import org.apache.doris.common.MetaNotFoundException;
import org.apache.doris.nereids.exceptions.AnalysisException;
import org.apache.doris.statistics.AnalysisInfo;
import org.apache.doris.statistics.BaseAnalysisTask;
import org.apache.doris.statistics.ColumnStatistic;
import org.apache.doris.statistics.TableStatsMeta;
import org.apache.doris.thrift.TTableDescriptor;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import org.apache.logging.log4j.LogManager;
@ -153,6 +160,82 @@ public interface TableIf {
void write(DataOutput out) throws IOException;
default Map<String, Constraint> getConstraintsMap() {
throw new RuntimeException(String.format("Not implemented constraint for table %s", this));
}
// Note this function is not thread safe
default void checkConstraintNotExistence(String name, Constraint primaryKeyConstraint,
Map<String, Constraint> constraintMap) {
if (constraintMap.containsKey(name)) {
throw new RuntimeException(String.format("Constraint name %s has existed", name));
}
for (Map.Entry<String, Constraint> entry : constraintMap.entrySet()) {
if (entry.getValue().equals(primaryKeyConstraint)) {
throw new RuntimeException(String.format(
"Constraint %s has existed, named %s", primaryKeyConstraint, entry.getKey()));
}
}
}
default void addUniqueConstraint(String name, ImmutableList<String> columns) {
writeLock();
try {
Map<String, Constraint> constraintMap = getConstraintsMap();
UniqueConstraint uniqueConstraint = new UniqueConstraint(name, ImmutableSet.copyOf(columns));
checkConstraintNotExistence(name, uniqueConstraint, constraintMap);
constraintMap.put(name, uniqueConstraint);
} finally {
writeUnlock();
}
}
default void addPrimaryKeyConstraint(String name, ImmutableList<String> columns) {
writeLock();
try {
Map<String, Constraint> constraintMap = getConstraintsMap();
PrimaryKeyConstraint primaryKeyConstraint = new PrimaryKeyConstraint(name, ImmutableSet.copyOf(columns));
checkConstraintNotExistence(name, primaryKeyConstraint, constraintMap);
constraintMap.put(name, primaryKeyConstraint);
} finally {
writeUnlock();
}
}
default void updatePrimaryKeyForForeignKey(PrimaryKeyConstraint requirePrimaryKey, TableIf referencedTable) {
referencedTable.writeLock();
try {
Optional<Constraint> primaryKeyConstraint = referencedTable.getConstraintsMap().values().stream()
.filter(requirePrimaryKey::equals)
.findFirst();
if (!primaryKeyConstraint.isPresent()) {
throw new AnalysisException(String.format(
"Foreign key constraint requires a primary key constraint %s in %s",
requirePrimaryKey.getPrimaryKeyNames(), referencedTable.getName()));
}
((PrimaryKeyConstraint) (primaryKeyConstraint.get())).addForeignTable(this);
} finally {
referencedTable.writeUnlock();
}
}
default void addForeignConstraint(String name, ImmutableList<String> columns,
TableIf referencedTable, ImmutableList<String> referencedColumns) {
writeLock();
try {
Map<String, Constraint> constraintMap = getConstraintsMap();
ForeignKeyConstraint foreignKeyConstraint =
new ForeignKeyConstraint(name, columns, referencedTable, referencedColumns);
checkConstraintNotExistence(name, foreignKeyConstraint, constraintMap);
PrimaryKeyConstraint requirePrimaryKey = new PrimaryKeyConstraint(name,
foreignKeyConstraint.getReferencedColumnNames());
updatePrimaryKeyForForeignKey(requirePrimaryKey, referencedTable);
constraintMap.put(name, foreignKeyConstraint);
} finally {
writeUnlock();
}
}
/**
* return true if this kind of table need read lock when doing query plan.
*

39
fe/fe-core/src/main/java/org/apache/doris/catalog/constraint/Constraint.java Normal file
View File

@ -0,0 +1,39 @@
// 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.catalog.constraint;
public abstract class Constraint {
public enum ConstraintType {
FOREIGN_KEY,
PRIMARY_KEY,
UNIQUE
}
private final String name;
private final ConstraintType type;
protected Constraint(ConstraintType type, String name) {
this.name = name;
this.type = type;
}
public String getName() {
return name;
}
}

89
fe/fe-core/src/main/java/org/apache/doris/catalog/constraint/ForeignKeyConstraint.java Normal file
View File

@ -0,0 +1,89 @@
// 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.catalog.constraint;
import org.apache.doris.catalog.Column;
import org.apache.doris.catalog.TableIf;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableMap.Builder;
import com.google.common.collect.ImmutableSet;
import java.util.List;
import java.util.Objects;
public class ForeignKeyConstraint extends Constraint {
private final ImmutableMap<String, String> foreignToReference;
private final TableIdentifier referencedTable;
public ForeignKeyConstraint(String name, List<String> columns,
TableIf refTable, List<String> referencedColumns) {
super(ConstraintType.FOREIGN_KEY, name);
ImmutableMap.Builder<String, String> builder = new Builder<>();
Preconditions.checkArgument(columns.size() == referencedColumns.size(),
"Foreign keys' size must be same as the size of reference keys");
Preconditions.checkArgument(ImmutableSet.copyOf(columns).size() == columns.size(),
"Foreign keys contains duplicate slots.");
Preconditions.checkArgument(ImmutableSet.copyOf(referencedColumns).size() == referencedColumns.size(),
"Reference keys contains duplicate slots.");
this.referencedTable = new TableIdentifier(refTable);
for (int i = 0; i < columns.size(); i++) {
builder.put(columns.get(i), referencedColumns.get(i));
}
this.foreignToReference = builder.build();
}
public ImmutableSet<String> getForeignKeyNames() {
return foreignToReference.keySet();
}
public ImmutableSet<String> getReferencedColumnNames() {
return ImmutableSet.copyOf(foreignToReference.values());
}
public String getReferencedColumnName(String column) {
return foreignToReference.get(column);
}
public Column getReferencedColumn(String column) {
return getReferencedTable().getColumn(getReferencedColumnName(column));
}
public TableIf getReferencedTable() {
return referencedTable.toTableIf();
}
@Override
public int hashCode() {
return Objects.hash(foreignToReference, referencedTable);
}
@Override
public boolean equals(Object obj) {
if (this == obj) {
return true;
}
if (obj == null || getClass() != obj.getClass()) {
return false;
}
ForeignKeyConstraint other = (ForeignKeyConstraint) obj;
return Objects.equals(foreignToReference, other.foreignToReference)
&& Objects.equals(referencedTable, other.referencedTable);
}
}

71
fe/fe-core/src/main/java/org/apache/doris/catalog/constraint/PrimaryKeyConstraint.java Normal file
View File

@ -0,0 +1,71 @@
// 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.catalog.constraint;
import org.apache.doris.catalog.TableIf;
import com.google.common.base.Objects;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
public class PrimaryKeyConstraint extends Constraint {
private final ImmutableSet<String> columns;
// record the foreign table which references the primary key
private final Set<TableIdentifier> foreignTables = new HashSet<>();
public PrimaryKeyConstraint(String name, Set<String> columns) {
super(ConstraintType.PRIMARY_KEY, name);
this.columns = ImmutableSet.copyOf(columns);
}
public ImmutableSet<String> getPrimaryKeyNames() {
return columns;
}
public void addForeignTable(TableIf table) {
foreignTables.add(new TableIdentifier(table));
}
public List<TableIf> getReferenceTables() {
return foreignTables.stream()
.map(TableIdentifier::toTableIf)
.collect(ImmutableList.toImmutableList());
}
@Override
public boolean equals(Object o) {
if (this == o) {
return true;
}
if (o == null || getClass() != o.getClass()) {
return false;
}
PrimaryKeyConstraint that = (PrimaryKeyConstraint) o;
return columns.equals(that.columns);
}
@Override
public int hashCode() {
return Objects.hashCode(columns);
}
}

68
fe/fe-core/src/main/java/org/apache/doris/catalog/constraint/TableIdentifier.java Normal file
View File

@ -0,0 +1,68 @@
// 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.catalog.constraint;
import org.apache.doris.catalog.DatabaseIf;
import org.apache.doris.catalog.Env;
import org.apache.doris.catalog.TableIf;
import com.google.common.base.Preconditions;
import java.util.Objects;
class TableIdentifier {
private final long databaseId;
private final long tableId;
TableIdentifier(TableIf tableIf) {
Preconditions.checkArgument(tableIf != null,
"Table can not be null in constraint");
databaseId = tableIf.getDatabase().getId();
tableId = tableIf.getId();
}
TableIf toTableIf() {
DatabaseIf databaseIf = Env.getCurrentEnv().getCurrentCatalog().getDbNullable(databaseId);
if (databaseIf == null) {
throw new RuntimeException(String.format("Can not find database %s in constraint", databaseId));
}
TableIf tableIf = databaseIf.getTableNullable(tableId);
if (tableIf == null) {
throw new RuntimeException(String.format("Can not find table %s in constraint", databaseId));
}
return tableIf;
}
@Override
public boolean equals(Object o) {
if (this == o) {
return true;
}
if (o == null || getClass() != o.getClass()) {
return false;
}
TableIdentifier that = (TableIdentifier) o;
return databaseId == that.databaseId
&& tableId == that.tableId;
}
@Override
public int hashCode() {
return Objects.hash(databaseId, tableId);
}
}

53
fe/fe-core/src/main/java/org/apache/doris/catalog/constraint/UniqueConstraint.java Normal file
View File

@ -0,0 +1,53 @@
// 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.catalog.constraint;
import com.google.common.base.Objects;
import com.google.common.collect.ImmutableSet;
import java.util.Set;
public class UniqueConstraint extends Constraint {
private final ImmutableSet<String> columns;
public UniqueConstraint(String name, Set<String> columns) {
super(ConstraintType.UNIQUE, name);
this.columns = ImmutableSet.copyOf(columns);
}
public ImmutableSet<String> getUniqueColumnNames() {
return columns;
}
@Override
public boolean equals(Object o) {
if (this == o) {
return true;
}
if (o == null || getClass() != o.getClass()) {
return false;
}
UniqueConstraint that = (UniqueConstraint) o;
return columns.equals(that.columns);
}
@Override
public int hashCode() {
return Objects.hashCode(columns);
}
}

30
fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java
View File

@ -37,6 +37,7 @@ import org.apache.doris.mtmv.MTMVRefreshInfo;
import org.apache.doris.mtmv.MTMVRefreshSchedule;
import org.apache.doris.mtmv.MTMVRefreshTriggerInfo;
import org.apache.doris.nereids.DorisParser;
import org.apache.doris.nereids.DorisParser.AddConstraintContext;
import org.apache.doris.nereids.DorisParser.AggClauseContext;
import org.apache.doris.nereids.DorisParser.AliasQueryContext;
import org.apache.doris.nereids.DorisParser.AliasedQueryContext;
@ -232,6 +233,7 @@ import org.apache.doris.nereids.trees.expressions.OrderExpression;
import org.apache.doris.nereids.trees.expressions.Properties;
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.StatementScopeIdGenerator;
import org.apache.doris.nereids.trees.expressions.Subtract;
import org.apache.doris.nereids.trees.expressions.TimestampArithmetic;
@ -309,8 +311,10 @@ import org.apache.doris.nereids.trees.plans.LimitPhase;
import org.apache.doris.nereids.trees.plans.Plan;
import org.apache.doris.nereids.trees.plans.algebra.Aggregate;
import org.apache.doris.nereids.trees.plans.algebra.SetOperation.Qualifier;
import org.apache.doris.nereids.trees.plans.commands.AddConstraintCommand;
import org.apache.doris.nereids.trees.plans.commands.AlterMTMVCommand;
import org.apache.doris.nereids.trees.plans.commands.Command;
import org.apache.doris.nereids.trees.plans.commands.Constraint;
import org.apache.doris.nereids.trees.plans.commands.CreateMTMVCommand;
import org.apache.doris.nereids.trees.plans.commands.CreatePolicyCommand;
import org.apache.doris.nereids.trees.plans.commands.CreateTableCommand;
@ -614,6 +618,32 @@ public class LogicalPlanBuilder extends DorisParserBaseVisitor<Object> {
return new AlterMTMVCommand(alterMTMVInfo);
}
@Override
public LogicalPlan visitAddConstraint(AddConstraintContext ctx) {
LogicalPlan curTable = visitRelation(ctx.table);
ImmutableList<Slot> slots = ctx.constraint().slots.stream()
.map(RuleContext::getText)
.map(UnboundSlot::new)
.collect(ImmutableList.toImmutableList());
Constraint constraint;
if (ctx.constraint().UNIQUE() != null) {
constraint = Constraint.newUniqueConstraint(curTable, slots);
} else if (ctx.constraint().PRIMARY() != null) {
constraint = Constraint.newPrimaryKeyConstraint(curTable, slots);
} else if (ctx.constraint().FOREIGN() != null) {
ImmutableList<Slot> referenceSlots = ctx.constraint().referenceSlots.stream()
.map(RuleContext::getText)
.map(UnboundSlot::new)
.collect(ImmutableList.toImmutableList());
List<String> nameParts = visitMultipartIdentifier(ctx.constraint().referenceTable);
LogicalPlan referenceTable = new UnboundRelation(StatementScopeIdGenerator.newRelationId(), nameParts);
constraint = Constraint.newForeignKeyConstraint(curTable, slots, referenceTable, referenceSlots);
} else {
throw new AnalysisException("Unsupported constraint " + ctx.getText());
}
return new AddConstraintCommand(ctx.constraintName.getText().toLowerCase(), constraint);
}
@Override
public LogicalPlan visitUpdate(UpdateContext ctx) {
LogicalPlan query = LogicalPlanBuilderAssistant.withCheckPolicy(new UnboundRelation(

1
fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/PlanType.java
View File

@ -126,6 +126,7 @@ public enum PlanType {
UPDATE_COMMAND,
CREATE_MTMV_COMMAND,
ALTER_MTMV_COMMAND,
ADD_CONSTRAINT_COMMAND,
REFRESH_MTMV_COMMAND,
DROP_MTMV_COMMAND
}

101
fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/AddConstraintCommand.java Normal file
View File

@ -0,0 +1,101 @@
// 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.plans.commands;
import org.apache.doris.catalog.Table;
import org.apache.doris.catalog.TableIf;
import org.apache.doris.common.Pair;
import org.apache.doris.nereids.NereidsPlanner;
import org.apache.doris.nereids.exceptions.AnalysisException;
import org.apache.doris.nereids.properties.PhysicalProperties;
import org.apache.doris.nereids.trees.expressions.SlotReference;
import org.apache.doris.nereids.trees.plans.Plan;
import org.apache.doris.nereids.trees.plans.PlanType;
import org.apache.doris.nereids.trees.plans.commands.ExplainCommand.ExplainLevel;
import org.apache.doris.nereids.trees.plans.logical.LogicalCatalogRelation;
import org.apache.doris.nereids.trees.plans.logical.LogicalPlan;
import org.apache.doris.nereids.trees.plans.visitor.PlanVisitor;
import org.apache.doris.qe.ConnectContext;
import org.apache.doris.qe.StmtExecutor;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.Set;
/**
* create multi table materialized view
*/
public class AddConstraintCommand extends Command implements ForwardWithSync {
public static final Logger LOG = LogManager.getLogger(AddConstraintCommand.class);
private final String name;
private final Constraint constraint;
/**
* constructor
*/
public AddConstraintCommand(String name, Constraint constraint) {
super(PlanType.ADD_CONSTRAINT_COMMAND);
this.constraint = constraint;
this.name = name;
}
@Override
public void run(ConnectContext ctx, StmtExecutor executor) throws Exception {
Pair<ImmutableList<String>, TableIf> columnsAndTable = extractColumnsAndTable(ctx, constraint.toProject());
if (constraint.isForeignKey()) {
Pair<ImmutableList<String>, TableIf> foreignColumnsAndTable
= extractColumnsAndTable(ctx, constraint.toReferenceProject());
columnsAndTable.second.addForeignConstraint(name, columnsAndTable.first,
foreignColumnsAndTable.second, foreignColumnsAndTable.first);
} else if (constraint.isPrimaryKey()) {
columnsAndTable.second.addPrimaryKeyConstraint(name, columnsAndTable.first);
} else if (constraint.isUnique()) {
columnsAndTable.second.addUniqueConstraint(name, columnsAndTable.first);
}
}
private Pair<ImmutableList<String>, TableIf> extractColumnsAndTable(ConnectContext ctx, LogicalPlan plan) {
NereidsPlanner planner = new NereidsPlanner(ctx.getStatementContext());
Plan analyzedPlan = planner.plan(plan, PhysicalProperties.ANY, ExplainLevel.ANALYZED_PLAN);
Set<LogicalCatalogRelation> logicalCatalogRelationSet = analyzedPlan
.collect(LogicalCatalogRelation.class::isInstance);
if (logicalCatalogRelationSet.size() != 1) {
throw new AnalysisException("Can not found table in constraint " + constraint.toString());
}
LogicalCatalogRelation catalogRelation = logicalCatalogRelationSet.iterator().next();
Preconditions.checkArgument(catalogRelation.getTable() instanceof Table,
"We only support table now but we meet ", catalogRelation.getTable());
ImmutableList<String> columns = analyzedPlan.getOutput().stream()
.map(s -> {
Preconditions.checkArgument(s instanceof SlotReference
&& ((SlotReference) s).getColumn().isPresent(),
"Constraint contains a invalid slot ", s);
return ((SlotReference) s).getColumn().get().getName();
}).collect(ImmutableList.toImmutableList());
return Pair.of(columns, catalogRelation.getTable());
}
@Override
public <R, C> R accept(PlanVisitor<R, C> visitor, C context) {
return visitor.visitAddConstraintCommand(this, context);
}
}

115
fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/Constraint.java Normal file
View File

@ -0,0 +1,115 @@
// 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.plans.commands;
import org.apache.doris.catalog.constraint.Constraint.ConstraintType;
import org.apache.doris.nereids.trees.expressions.Slot;
import org.apache.doris.nereids.trees.plans.logical.LogicalPlan;
import org.apache.doris.nereids.trees.plans.logical.LogicalProject;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import java.util.Objects;
import javax.annotation.Nullable;
/**
* Constraint
*/
public class Constraint {
private final ImmutableList<Slot> slots;
private final LogicalPlan curTable;
private final @Nullable LogicalPlan referenceTable;
private final @Nullable ImmutableList<Slot> referenceSlots;
private final ConstraintType type;
Constraint(ConstraintType type, LogicalPlan curTable, ImmutableList<Slot> slots) {
Preconditions.checkArgument(slots != null && !slots.isEmpty(),
"slots of constraint can't be null or empty");
this.type = type;
this.slots = slots;
this.curTable = Objects.requireNonNull(curTable,
"table of constraint can't be null");
this.referenceTable = null;
this.referenceSlots = null;
}
Constraint(LogicalPlan curTable, ImmutableList<Slot> slots,
LogicalPlan referenceTable, ImmutableList<Slot> referenceSlotSet) {
Preconditions.checkArgument(slots != null && !slots.isEmpty(),
"slots of constraint can't be null or empty");
this.type = ConstraintType.FOREIGN_KEY;
this.slots = slots;
this.curTable = Objects.requireNonNull(curTable,
"table of constraint can't be null");
this.referenceTable = Objects.requireNonNull(referenceTable,
"reference table in foreign key can not be null");
this.referenceSlots = Objects.requireNonNull(referenceSlotSet,
"reference slots in foreign key can not be null");
Preconditions.checkArgument(referenceSlots.size() == slots.size(),
"Foreign key's size must be same as the size of reference slots");
}
public static Constraint newUniqueConstraint(LogicalPlan curTable, ImmutableList<Slot> slotSet) {
return new Constraint(ConstraintType.UNIQUE, curTable, slotSet);
}
public static Constraint newPrimaryKeyConstraint(LogicalPlan curTable, ImmutableList<Slot> slotSet) {
return new Constraint(ConstraintType.PRIMARY_KEY, curTable, slotSet);
}
public static Constraint newForeignKeyConstraint(
LogicalPlan curTable, ImmutableList<Slot> slotSet,
LogicalPlan referenceTable, ImmutableList<Slot> referenceSlotSet) {
return new Constraint(curTable, slotSet, referenceTable, referenceSlotSet);
}
public boolean isForeignKey() {
return type == ConstraintType.FOREIGN_KEY;
}
public boolean isUnique() {
return type == ConstraintType.UNIQUE;
}
public boolean isPrimaryKey() {
return type == ConstraintType.PRIMARY_KEY;
}
public LogicalPlan toProject() {
return new LogicalProject<>(ImmutableList.copyOf(slots), curTable);
}
public LogicalPlan toReferenceProject() {
Preconditions.checkArgument(referenceSlots != null, "Reference slot set of foreign key cannot be null");
return new LogicalProject<>(ImmutableList.copyOf(referenceSlots), referenceTable);
}
@Override
public String toString() {
StringBuilder sb = new StringBuilder();
sb.append("Constraint Type: ").append(type).append("\n");
sb.append("Slot Set: ").append(slots).append("\n");
sb.append("Current Table: ").append(curTable).append("\n");
if (type.equals(ConstraintType.FOREIGN_KEY)) {
sb.append("Reference Table: ").append(referenceTable).append("\n");
sb.append("Reference Slot Set: ").append(referenceSlots);
}
return sb.toString();
}
}

5
fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/visitor/CommandVisitor.java
View File

@ -17,6 +17,7 @@
package org.apache.doris.nereids.trees.plans.visitor;
import org.apache.doris.nereids.trees.plans.commands.AddConstraintCommand;
import org.apache.doris.nereids.trees.plans.commands.AlterMTMVCommand;
import org.apache.doris.nereids.trees.plans.commands.Command;
import org.apache.doris.nereids.trees.plans.commands.CreateMTMVCommand;
@ -77,6 +78,10 @@ public interface CommandVisitor<R, C> {
return visitCommand(alterMTMVCommand, context);
}
default R visitAddConstraintCommand(AddConstraintCommand addConstraintCommand, C context) {
return visitCommand(addConstraintCommand, context);
}
default R visitRefreshMTMVCommand(RefreshMTMVCommand refreshMTMVCommand, C context) {
return visitCommand(refreshMTMVCommand, context);
}

134
fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/AddConstraintTest.java Normal file
View File

@ -0,0 +1,134 @@
// 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.plans;
import org.apache.doris.catalog.Column;
import org.apache.doris.catalog.TableIf;
import org.apache.doris.catalog.constraint.Constraint;
import org.apache.doris.catalog.constraint.ForeignKeyConstraint;
import org.apache.doris.catalog.constraint.PrimaryKeyConstraint;
import org.apache.doris.catalog.constraint.UniqueConstraint;
import org.apache.doris.nereids.parser.NereidsParser;
import org.apache.doris.nereids.trees.expressions.SlotReference;
import org.apache.doris.nereids.trees.plans.commands.AddConstraintCommand;
import org.apache.doris.nereids.util.PlanChecker;
import org.apache.doris.nereids.util.PlanPatternMatchSupported;
import org.apache.doris.utframe.TestWithFeService;
import com.google.common.collect.Sets;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
import java.util.List;
import java.util.Set;
class AddConstraintTest extends TestWithFeService implements PlanPatternMatchSupported {
@Override
public void runBeforeAll() throws Exception {
createDatabase("test");
connectContext.setDatabase("default_cluster:test");
createTable("create table t1 (\n"
+ " k1 int,\n"
+ " k2 int\n"
+ ")\n"
+ "unique key(k1, k2)\n"
+ "distributed by hash(k1) buckets 4\n"
+ "properties(\n"
+ " \"replication_num\"=\"1\"\n"
+ ")");
createTable("create table t2 (\n"
+ " k1 int,\n"
+ " k2 int\n"
+ ")\n"
+ "unique key(k1, k2)\n"
+ "distributed by hash(k1) buckets 4\n"
+ "properties(\n"
+ " \"replication_num\"=\"1\"\n"
+ ")");
}
@Test
void addPrimaryKeyConstraintTest() throws Exception {
AddConstraintCommand command = (AddConstraintCommand) new NereidsParser().parseSingle(
"alter table t1 add constraint pk primary key (k1)");
command.run(connectContext, null);
PlanChecker.from(connectContext).parse("select * from t1").analyze().matches(logicalOlapScan().when(o -> {
Constraint c = o.getTable().getConstraint("pk");
if (c instanceof PrimaryKeyConstraint) {
Set<String> columns = ((PrimaryKeyConstraint) c).getPrimaryKeyNames();
return columns.size() == 1 && columns.iterator().next().equals("k1");
}
return false;
}));
}
@Test
void addUniqueConstraintTest() throws Exception {
AddConstraintCommand command = (AddConstraintCommand) new NereidsParser().parseSingle(
"alter table t1 add constraint un unique (k1)");
command.run(connectContext, null);
PlanChecker.from(connectContext).parse("select * from t1").analyze().matches(logicalOlapScan().when(o -> {
Constraint c = o.getTable().getConstraint("un");
if (c instanceof UniqueConstraint) {
Set<String> columns = ((UniqueConstraint) c).getUniqueColumnNames();
return columns.size() == 1 && columns.iterator().next().equals("k1");
}
return false;
}));
}
@Test
void addForeignKeyConstraintTest() throws Exception {
AddConstraintCommand command = (AddConstraintCommand) new NereidsParser().parseSingle(
"alter table t1 add constraint fk foreign key (k1) references t2(k1)");
try {
command.run(connectContext, null);
} catch (Exception e) {
Assertions.assertEquals("Foreign key constraint requires a primary key constraint [k1] in t2",
e.getMessage());
}
((AddConstraintCommand) new NereidsParser().parseSingle(
"alter table t2 add constraint pk primary key (k1, k2)")).run(connectContext, null);
((AddConstraintCommand) new NereidsParser().parseSingle(
"alter table t1 add constraint fk foreign key (k1, k2) references t2(k1, k2)")).run(connectContext,
null);
PlanChecker.from(connectContext).parse("select * from t1").analyze().matches(logicalOlapScan().when(o -> {
Constraint c = o.getTable().getConstraint("fk");
if (c instanceof ForeignKeyConstraint) {
ForeignKeyConstraint f = (ForeignKeyConstraint) c;
Column ref1 = f.getReferencedColumn(((SlotReference) o.getOutput().get(0)).getColumn().get().getName());
Column ref2 = f.getReferencedColumn(((SlotReference) o.getOutput().get(1)).getColumn().get().getName());
return ref1.getName().equals("k1") && ref2.getName().equals("k2");
}
return false;
}));
PlanChecker.from(connectContext).parse("select * from t2").analyze().matches(logicalOlapScan().when(o -> {
Constraint c = o.getTable().getConstraint("pk");
if (c instanceof PrimaryKeyConstraint) {
Set<String> columnNames = ((PrimaryKeyConstraint) c).getPrimaryKeyNames();
List<TableIf> referenceTable = ((PrimaryKeyConstraint) c).getReferenceTables();
return columnNames.size() == 2
&& columnNames.equals(Sets.newHashSet("k1", "k2"))
&& referenceTable.size() == 1 && referenceTable.get(0).getName().equals("t1");
}
return false;
}));
}
}

1
fe/fe-core/src/test/java/org/apache/doris/utframe/TestWithFeService.java
View File

@ -272,6 +272,7 @@ public abstract class TestWithFeService {
ctx.setRemoteIP(host);
ctx.setEnv(Env.getCurrentEnv());
ctx.setThreadLocalInfo();
ctx.setStatementContext(new StatementContext());
return ctx;
}