diff --git a/fe/fe-core/src/main/java/org/apache/doris/alter/SchemaChangeHandler.java b/fe/fe-core/src/main/java/org/apache/doris/alter/SchemaChangeHandler.java index 8396e36bdb7f16..44d2ca87c85126 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/alter/SchemaChangeHandler.java +++ b/fe/fe-core/src/main/java/org/apache/doris/alter/SchemaChangeHandler.java @@ -543,6 +543,17 @@ private void checkSeqMapConditionMet(OlapTable olapTable, Map pr } } + private String findConstraintWithColumn(OlapTable table, String columnName) { + String mappingConstraint = Env.getCurrentEnv().getConstraintManager() + .findDistributionMappingConstraintWithColumn(table, columnName); + if (mappingConstraint != null) { + return mappingConstraint; + } + return Env.getCurrentEnv().getConstraintManager() + .findConstraintWithColumn(TableNameInfoUtils.fromCatalogDb( + table.getDatabase().getCatalog(), table.getDatabase(), table), columnName); + } + private void processDropColumn(DropColumnOp dropColumnOp, Table externalTable, List newSchema) throws DdlException { String dropColName = dropColumnOp.getColName(); @@ -594,11 +605,14 @@ private boolean processDropColumn(DropColumnOp dropColumnOp, OlapTable olapTable throws DdlException { String dropColName = dropColumnOp.getColName(); + String targetIndexName = dropColumnOp.getRollupName(); - String constraintName = Env.getCurrentEnv().getConstraintManager() - .findConstraintWithColumn(TableNameInfoUtils.fromCatalogDb( - olapTable.getDatabase().getCatalog(), - olapTable.getDatabase(), olapTable), dropColName); + String constraintName = targetIndexName == null + ? findConstraintWithColumn(olapTable, dropColName) + : Env.getCurrentEnv().getConstraintManager() + .findConstraintWithColumn(TableNameInfoUtils.fromCatalogDb( + olapTable.getDatabase().getCatalog(), olapTable.getDatabase(), olapTable), + dropColName); if (constraintName != null) { throw new DdlException(String.format( "Cannot drop column '%s' because it is used by constraint '%s'. " @@ -606,7 +620,6 @@ private boolean processDropColumn(DropColumnOp dropColumnOp, OlapTable olapTable dropColName, constraintName)); } - String targetIndexName = dropColumnOp.getRollupName(); checkIndexExists(olapTable, targetIndexName); String baseIndexName = olapTable.getName(); @@ -975,6 +988,16 @@ private boolean processModifyColumn(ModifyColumnOp modifyColumnOp, OlapTable ola } ColumnPosition columnPos = modifyColumnOp.getColPos(); String targetIndexName = modifyColumnOp.getRollupName(); + if (targetIndexName == null) { + String mappingConstraint = Env.getCurrentEnv().getConstraintManager() + .findDistributionMappingConstraintWithColumn(olapTable, modColumn.getName()); + if (mappingConstraint != null) { + throw new DdlException(String.format( + "Cannot modify column '%s' because it is used by constraint '%s'. " + + "Drop the constraint first.", + modColumn.getName(), mappingConstraint)); + } + } checkIndexExists(olapTable, targetIndexName); String baseIndexName = olapTable.getName(); diff --git a/fe/fe-core/src/main/java/org/apache/doris/backup/RestoreJob.java b/fe/fe-core/src/main/java/org/apache/doris/backup/RestoreJob.java index fd73b2a669fb3d..d872c8ae390860 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/backup/RestoreJob.java +++ b/fe/fe-core/src/main/java/org/apache/doris/backup/RestoreJob.java @@ -600,6 +600,10 @@ private void checkAndPrepareMeta() { } Preconditions.checkNotNull(backupMeta); + if (!validateDistributionMappingConstraintsForRestore()) { + return; + } + // Check the olap table state. // // If isAtomicRestore is not set, set all restored tbls' state to RESTORE, @@ -638,6 +642,14 @@ private void checkAndPrepareMeta() { continue; } + if (!env.getConstraintManager().getDistributionMappingConstraints(olapTbl).isEmpty()) { + status = new Status(ErrCode.COMMON_ERROR, + "Cannot restore into existing table " + olapTbl.getName() + + " because it has distribution mapping constraints. " + + "Drop those constraints before using non-atomic restore."); + return; + } + olapTbl.setState(OlapTableState.RESTORE); // set restore status for partitions BackupOlapTableInfo tblInfo = jobInfo.backupOlapTableObjects.get(tableName); @@ -1063,6 +1075,35 @@ private void checkAndPrepareMeta() { setState(RestoreJobState.CREATING); } + private boolean validateDistributionMappingConstraintsForRestore() { + boolean featureCompatibilityValidated = false; + for (String tableName : jobInfo.backupOlapTableObjects.keySet()) { + OlapTable restoredTable = (OlapTable) backupMeta.getTable(tableName); + if (env.getConstraintManager().getDistributionMappingConstraints(restoredTable).isEmpty()) { + continue; + } + if (isAtomicRestore) { + status = new Status(ErrCode.COMMON_ERROR, + "Cannot atomically restore table " + tableName + + " because its backup contains distribution mapping constraints. " + + "Use a non-atomic restore or create a backup without those constraints."); + return false; + } + try { + if (!featureCompatibilityValidated) { + env.getConstraintManager().validateDistributionMappingFeatureCompatibility(); + featureCompatibilityValidated = true; + } + env.getConstraintManager().validateDistributionMappingConstraints(restoredTable); + } catch (org.apache.doris.nereids.exceptions.AnalysisException e) { + status = new Status(ErrCode.COMMON_ERROR, + "Cannot restore table " + tableName + ": " + e.getMessage()); + return false; + } + } + return true; + } + protected void doCreateReplicas() { // Send create replica task to BE outside the db lock int numBatchTasks = batchTaskPerTable.values() diff --git a/fe/fe-core/src/main/java/org/apache/doris/catalog/Env.java b/fe/fe-core/src/main/java/org/apache/doris/catalog/Env.java index 4a615be5c9ffe4..f4227f84b3da12 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/catalog/Env.java +++ b/fe/fe-core/src/main/java/org/apache/doris/catalog/Env.java @@ -5998,6 +5998,15 @@ private void renameColumn(Database db, OlapTable table, String colName, if (partitionInfo.getPartitionColumns().stream().anyMatch(c -> c.getName().equalsIgnoreCase(colName))) { throw new DdlException("Renaming partition columns has problems, forbidden in current Doris version"); } + if (!isReplay) { + String mappingConstraint = constraintManager.findDistributionMappingConstraintWithColumn(table, colName); + if (mappingConstraint != null) { + throw new DdlException(String.format( + "Cannot rename column '%s' because it is used by constraint '%s'. " + + "Drop the constraint first.", + colName, mappingConstraint)); + } + } Map indexIdToMeta = table.getIndexIdToMeta(); for (Map.Entry entry : indexIdToMeta.entrySet()) { @@ -6341,6 +6350,10 @@ public void updateBinlogConfig(Database db, OlapTable table, BinlogConfig newBin public void replayModifyTableProperty(short opCode, ModifyTablePropertyOperationLog info) throws MetaNotFoundException { + if (info.hasDistributionMappingConstraintMutation()) { + replayDistributionMappingConstraint(info); + return; + } String ctlName = info.getCtlName(); long dbId = info.getDbId(); long tableId = info.getTableId(); @@ -6391,6 +6404,19 @@ public void replayModifyTableProperty(short opCode, ModifyTablePropertyOperation } } + private void replayDistributionMappingConstraint(ModifyTablePropertyOperationLog info) + throws MetaNotFoundException { + Database db = getInternalCatalog().getDbOrMetaException(info.getDbId()); + OlapTable table = (OlapTable) db.getTableOrMetaException(info.getTableId(), TableType.OLAP); + table.writeLock(); + try { + constraintManager.replayDistributionMappingConstraints( + table, info.getProperties()); + } finally { + table.writeUnlock(); + } + } + private void setExternalTableAutoAnalyze(Map properties, ModifyTablePropertyOperationLog info) { if (properties.size() != 1) { LOG.warn("External table property should contain exactly 1 entry."); @@ -6939,6 +6965,10 @@ public void convertDistributionType(Database db, OlapTable tbl) throws DdlExcept } } } + if (!constraintManager.getDistributionMappingConstraints(tbl).isEmpty()) { + throw new DdlException("Cannot change distribution type of table with" + + " distribution mapping constraints. Drop the constraints first."); + } if (!tbl.convertHashDistributionToRandomDistribution()) { throw new DdlException("Table " + tbl.getName() + " is not hash distributed"); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/catalog/TableProperty.java b/fe/fe-core/src/main/java/org/apache/doris/catalog/TableProperty.java index a8817b48a53f76..bab898e4db920d 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/catalog/TableProperty.java +++ b/fe/fe-core/src/main/java/org/apache/doris/catalog/TableProperty.java @@ -18,11 +18,13 @@ package org.apache.doris.catalog; import org.apache.doris.analysis.DataSortInfo; +import org.apache.doris.catalog.constraint.DistributionMappingConstraint; import org.apache.doris.common.AnalysisException; import org.apache.doris.common.Config; import org.apache.doris.common.util.PropertyAnalyzer; import org.apache.doris.persist.OperationType; import org.apache.doris.persist.gson.GsonPostProcessable; +import org.apache.doris.persist.gson.GsonUtils; import org.apache.doris.proto.OlapFile.EncryptionAlgorithmPB; import org.apache.doris.thrift.TCompressionType; import org.apache.doris.thrift.TEncryptionAlgorithm; @@ -34,14 +36,18 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.base.Joiner; import com.google.common.base.Strings; +import com.google.common.collect.ImmutableMap; import com.google.common.collect.Lists; import com.google.common.collect.Maps; +import com.google.gson.JsonParseException; import com.google.gson.annotations.SerializedName; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import java.io.IOException; +import java.util.ArrayList; import java.util.Arrays; +import java.util.Comparator; import java.util.HashMap; import java.util.List; import java.util.Locale; @@ -59,10 +65,15 @@ public class TableProperty implements GsonPostProcessable { "default." + PropertyAnalyzer.PROPERTIES_REPLICATION_NUM; private static final String DEFAULT_REPLICATION_ALLOCATION = "default." + PropertyAnalyzer.PROPERTIES_REPLICATION_ALLOCATION; + public static final String DISTRIBUTION_MAPPING_CONSTRAINTS_PROPERTY = + "__distribution_mapping_constraints"; @SerializedName(value = "properties") private Map properties; + // Derived from DISTRIBUTION_MAPPING_CONSTRAINTS_PROPERTY. The property is the persistent source of truth. + private Map distributionMappingConstraints = ImmutableMap.of(); + // the follower variables are built from "properties" private DynamicPartitionProperty dynamicPartitionProperty = EnvFactory.getInstance().createDynamicPartitionProperty(Maps.newHashMap()); @@ -166,6 +177,7 @@ public TableProperty buildProperty(short opCode) { buildReplicaAllocation(); break; case OperationType.OP_MODIFY_TABLE_PROPERTIES: + buildDistributionMappingConstraints(); buildInMemory(); buildMinLoadReplicaNum(); buildStorageMedium(); @@ -586,6 +598,8 @@ public void removeInvalidProperties() { properties.remove(PropertyAnalyzer.PROPERTIES_STORAGE_POLICY); storagePolicy = ""; properties.remove(PropertyAnalyzer.PROPERTIES_COLOCATE_WITH); + properties.remove(DISTRIBUTION_MAPPING_CONSTRAINTS_PROPERTY); + distributionMappingConstraints = ImmutableMap.of(); properties.remove(DynamicPartitionProperty.STORAGE_POLICY); dynamicPartitionProperty.clearStoragePolicy(); } @@ -724,6 +738,74 @@ public Map getProperties() { return properties; } + public Map getDistributionMappingConstraints() { + return distributionMappingConstraints; + } + + public void addDistributionMappingConstraint(DistributionMappingConstraint constraint) { + Map updated = new HashMap<>(distributionMappingConstraints); + updated.put(constraint.getName(), constraint); + updateDistributionMappingConstraints(updated); + } + + public DistributionMappingConstraint removeDistributionMappingConstraint(String constraintName) { + Map updated = new HashMap<>(distributionMappingConstraints); + DistributionMappingConstraint removed = updated.remove(constraintName); + if (removed != null) { + updateDistributionMappingConstraints(updated); + } + return removed; + } + + public Map getDistributionMappingConstraintProperties() { + return ImmutableMap.of(DISTRIBUTION_MAPPING_CONSTRAINTS_PROPERTY, + properties.get(DISTRIBUTION_MAPPING_CONSTRAINTS_PROPERTY)); + } + + private void updateDistributionMappingConstraints( + Map updatedConstraints) { + List sortedConstraints = new ArrayList<>(updatedConstraints.values()); + sortedConstraints.sort(Comparator.comparing(DistributionMappingConstraint::getName)); + properties.put(DISTRIBUTION_MAPPING_CONSTRAINTS_PROPERTY, GsonUtils.GSON.toJson(sortedConstraints)); + distributionMappingConstraints = ImmutableMap.copyOf(updatedConstraints); + } + + private void buildDistributionMappingConstraints() { + try { + distributionMappingConstraints = deserializeDistributionMappingConstraints(); + } catch (IOException e) { + throw new IllegalStateException("Failed to deserialize distribution mapping constraints", e); + } + } + + private Map deserializeDistributionMappingConstraints() + throws IOException { + String serializedConstraints = properties.get(DISTRIBUTION_MAPPING_CONSTRAINTS_PROPERTY); + if (serializedConstraints == null) { + return ImmutableMap.of(); + } + DistributionMappingConstraint[] constraints; + try { + constraints = GsonUtils.GSON.fromJson(serializedConstraints, DistributionMappingConstraint[].class); + } catch (JsonParseException e) { + throw new IOException("Invalid distribution mapping constraints property", e); + } + if (constraints == null) { + throw new IOException("Distribution mapping constraints property must be a JSON array"); + } + Map constraintsByName = new HashMap<>(); + for (DistributionMappingConstraint constraint : constraints) { + if (constraint == null || constraint.getName() == null) { + throw new IOException("Distribution mapping constraints property contains an invalid constraint"); + } + if (constraintsByName.put(constraint.getName(), constraint) != null) { + throw new IOException("Distribution mapping constraints property contains duplicate constraint name: " + + constraint.getName()); + } + } + return ImmutableMap.copyOf(constraintsByName); + } + public DynamicPartitionProperty getDynamicPartitionProperty() { return dynamicPartitionProperty; } @@ -950,6 +1032,7 @@ public void buildReplicaAllocation() { } public void gsonPostProcess() throws IOException { + distributionMappingConstraints = deserializeDistributionMappingConstraints(); executeBuildDynamicProperty(); buildInMemory(); buildMinLoadReplicaNum(); diff --git a/fe/fe-core/src/main/java/org/apache/doris/catalog/constraint/Constraint.java b/fe/fe-core/src/main/java/org/apache/doris/catalog/constraint/Constraint.java index 21ef319a0759a7..cdeaa989502b74 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/catalog/constraint/Constraint.java +++ b/fe/fe-core/src/main/java/org/apache/doris/catalog/constraint/Constraint.java @@ -21,6 +21,7 @@ public abstract class Constraint { public enum ConstraintType { + DISTRIBUTION_MAPPING("DISTRIBUTION MAPPING"), FOREIGN_KEY("FOREIGN KEY"), PRIMARY_KEY("PRIMARY KEY"), UNIQUE("UNIQUE"); diff --git a/fe/fe-core/src/main/java/org/apache/doris/catalog/constraint/ConstraintManager.java b/fe/fe-core/src/main/java/org/apache/doris/catalog/constraint/ConstraintManager.java index 5044f35d92181d..a8e4763557a7fa 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/catalog/constraint/ConstraintManager.java +++ b/fe/fe-core/src/main/java/org/apache/doris/catalog/constraint/ConstraintManager.java @@ -19,19 +19,31 @@ import org.apache.doris.catalog.DatabaseIf; import org.apache.doris.catalog.Env; +import org.apache.doris.catalog.HashDistributionInfo; +import org.apache.doris.catalog.OlapTable; import org.apache.doris.catalog.TableIf; +import org.apache.doris.catalog.TableProperty; import org.apache.doris.catalog.info.TableNameInfo; import org.apache.doris.common.DdlException; +import org.apache.doris.common.Version; import org.apache.doris.common.io.Text; import org.apache.doris.common.io.Writable; import org.apache.doris.datasource.CatalogIf; +import org.apache.doris.datasource.InternalCatalog; +import org.apache.doris.info.TableNameInfoUtils; import org.apache.doris.nereids.exceptions.AnalysisException; import org.apache.doris.persist.AlterConstraintLog; +import org.apache.doris.persist.EditLog; +import org.apache.doris.persist.ModifyTablePropertyOperationLog; +import org.apache.doris.persist.OperationType; import org.apache.doris.persist.gson.GsonPostProcessable; import org.apache.doris.persist.gson.GsonUtils; +import org.apache.doris.system.Frontend; +import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.util.concurrent.RateLimiter; import com.google.gson.annotations.SerializedName; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; @@ -39,12 +51,15 @@ import java.io.DataInput; import java.io.DataOutput; import java.io.IOException; +import java.util.ArrayList; import java.util.Collection; +import java.util.Collections; import java.util.HashMap; import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Map.Entry; +import java.util.TreeSet; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.locks.ReentrantReadWriteLock; import java.util.stream.Collectors; @@ -56,6 +71,8 @@ public class ConstraintManager implements Writable, GsonPostProcessable { private static final Logger LOG = LogManager.getLogger(ConstraintManager.class); + private static final RateLimiter DISTRIBUTION_MAPPING_FALLBACK_LOG_LIMITER = + RateLimiter.create(1.0 / 60.0); @SerializedName("cm") private final ConcurrentHashMap> constraintsMap @@ -103,15 +120,22 @@ private void writeUnlock() { */ public void addConstraint(TableNameInfo tableNameInfo, String constraintName, Constraint constraint, boolean replay) { + Preconditions.checkArgument(!(constraint instanceof DistributionMappingConstraint), + "distribution mapping constraints must be stored on the table"); String key = toKey(tableNameInfo); writeLock(); try { + TableIf table = null; if (!replay) { - validateTableAndColumns(tableNameInfo, constraint); + table = validateTableAndColumns(tableNameInfo, constraint); } Map tableConstraints = constraintsMap.computeIfAbsent( key, k -> new HashMap<>()); checkConstraintNotExistence(constraintName, constraint, tableConstraints); + if (table != null) { + checkConstraintNotExistence(constraintName, constraint, + getDistributionMappingConstraintsMap(table)); + } if (constraint instanceof ForeignKeyConstraint) { registerForeignKeyReference( tableNameInfo, (ForeignKeyConstraint) constraint); @@ -126,6 +150,38 @@ public void addConstraint(TableNameInfo tableNameInfo, String constraintName, } } + /** Add a mapping owned by an internal OLAP table and return its pending journal item. */ + public EditLog.EditLogItem addDistributionMappingConstraint(TableNameInfo tableNameInfo, + OlapTable table, DistributionMappingConstraint constraint) { + Preconditions.checkState(table.isWriteLockHeldByCurrentThread(), + "table write lock is required when adding a distribution mapping constraint"); + validateDistributionMappingConstraint(tableNameInfo, table, constraint); + validateDistributionMappingFeatureCompatibility(); + DistributionMappingConstraint boundConstraint = constraint.bindTo(table); + writeLock(); + try { + Map mappings = + getDistributionMappingConstraintsMap(table); + checkConstraintNotExistence(boundConstraint.getName(), boundConstraint, mappings); + Map centralizedConstraints = constraintsMap.get(toKey(tableNameInfo)); + if (centralizedConstraints != null) { + checkConstraintNotExistence( + boundConstraint.getName(), boundConstraint, centralizedConstraints); + } + TableProperty tableProperty = getOrCreateTableProperty(table); + tableProperty.addDistributionMappingConstraint(boundConstraint); + ModifyTablePropertyOperationLog log = new ModifyTablePropertyOperationLog( + table.getDatabase().getId(), table.getId(), table.getName(), + tableProperty.getDistributionMappingConstraintProperties()); + LOG.info("Added distribution mapping constraint {} on table {}", + boundConstraint.getName(), toKey(tableNameInfo)); + return Env.getCurrentEnv().getEditLog() + .submitEdit(OperationType.OP_MODIFY_TABLE_PROPERTIES, log); + } finally { + writeUnlock(); + } + } + /** * Snapshot the tables whose foreign keys would be cascade-dropped along with the given primary * key constraint. Taken under the read lock: {@link PrimaryKeyConstraint#getForeignTableInfos()} @@ -180,6 +236,52 @@ public void dropConstraint(TableNameInfo tableNameInfo, String constraintName, } } + /** Drop a mapping owned by an internal OLAP table and return its pending journal item. */ + public EditLog.EditLogItem dropDistributionMappingConstraint(TableNameInfo tableNameInfo, + OlapTable table, String constraintName) { + Preconditions.checkState(table.isWriteLockHeldByCurrentThread(), + "table write lock is required when dropping a distribution mapping constraint"); + writeLock(); + try { + Map mappings = + getDistributionMappingConstraintsMap(table); + DistributionMappingConstraint constraint = mappings.get(constraintName); + if (constraint == null) { + throw new AnalysisException(String.format( + "Unknown constraint %s on table %s.", constraintName, tableNameInfo)); + } + TableProperty tableProperty = getOrCreateTableProperty(table); + tableProperty.removeDistributionMappingConstraint(constraintName); + ModifyTablePropertyOperationLog log = new ModifyTablePropertyOperationLog( + table.getDatabase().getId(), table.getId(), table.getName(), + tableProperty.getDistributionMappingConstraintProperties()); + LOG.info("Dropped distribution mapping constraint {} from table {}", + constraintName, toKey(tableNameInfo)); + return Env.getCurrentEnv().getEditLog() + .submitEdit(OperationType.OP_MODIFY_TABLE_PROPERTIES, log); + } finally { + writeUnlock(); + } + } + + /** Replay a complete table-local mapping snapshot from the legacy table-property envelope. */ + public void replayDistributionMappingConstraints(OlapTable table, Map properties) { + Preconditions.checkState(table.isWriteLockHeldByCurrentThread(), + "table write lock is required when replaying a distribution mapping constraint"); + Preconditions.checkState(properties.containsKey(TableProperty.DISTRIBUTION_MAPPING_CONSTRAINTS_PROPERTY), + "distribution mapping constraint snapshot is required"); + writeLock(); + try { + TableProperty tableProperty = getOrCreateTableProperty(table); + tableProperty.modifyTableProperties(properties); + tableProperty.buildProperty(OperationType.OP_MODIFY_TABLE_PROPERTIES); + LOG.info("Replayed {} distribution mapping constraints on table {}", + tableProperty.getDistributionMappingConstraints().size(), table.getName()); + } finally { + writeUnlock(); + } + } + /** Returns an immutable copy of all constraints for the given table. */ public Map getConstraints(TableNameInfo tableNameInfo) { String key = toKey(tableNameInfo); @@ -196,6 +298,24 @@ public Map getConstraints(TableNameInfo tableNameInfo) { } } + /** Returns centralized constraints together with mappings owned by the concrete table. */ + public Map getConstraints(TableNameInfo tableNameInfo, TableIf table) { + readLock(); + try { + Map constraints = new HashMap<>(); + Map centralizedConstraints = constraintsMap.get(toKey(tableNameInfo)); + if (centralizedConstraints != null) { + constraints.putAll(centralizedConstraints); + } + Map mappings = getDistributionMappingConstraintsMap(table); + validateNoDistributionMappingConstraintNameCollision(toKey(tableNameInfo), mappings); + constraints.putAll(mappings); + return ImmutableMap.copyOf(constraints); + } finally { + readUnlock(); + } + } + /** Get a single constraint by name, or null if not found. */ public Constraint getConstraint(TableNameInfo tableNameInfo, String constraintName) { @@ -213,6 +333,126 @@ public Constraint getConstraint(TableNameInfo tableNameInfo, } } + /** Return a table-owned mapping first, then a centralized constraint with the same name. */ + public Constraint getConstraint(TableNameInfo tableNameInfo, TableIf table, + String constraintName) { + readLock(); + try { + DistributionMappingConstraint mapping = + getDistributionMappingConstraintsMap(table).get(constraintName); + if (mapping != null) { + return mapping; + } + Map tableConstraints = constraintsMap.get(toKey(tableNameInfo)); + return tableConstraints == null ? null : tableConstraints.get(constraintName); + } finally { + readUnlock(); + } + } + + /** Returns all distribution mappings owned by the concrete table. */ + public ImmutableList getDistributionMappingConstraints(TableIf table) { + readLock(); + try { + return getDistributionMappingConstraintsMap(table).entrySet().stream() + .sorted(Map.Entry.comparingByKey()) + .map(Entry::getValue) + .collect(ImmutableList.toImmutableList()); + } finally { + readUnlock(); + } + } + + /** Return mappings that are safe to consume, or no mappings when the optional optimization is unavailable. */ + public ImmutableList getDistributionMappingConstraintsForPlanning( + OlapTable table) { + ImmutableList constraints = getDistributionMappingConstraints(table); + if (constraints.isEmpty()) { + return constraints; + } + if (table.isBeingSynced()) { + logDistributionMappingFallback(table, "table is being synchronized by CCR"); + return ImmutableList.of(); + } + TableNameInfo tableNameInfo = TableNameInfoUtils.fromCatalogDb( + table.getDatabase().getCatalog(), table.getDatabase(), table); + readLock(); + try { + validateNoDistributionMappingConstraintNameCollision(toKey(tableNameInfo), + getDistributionMappingConstraintsMap(table)); + } finally { + readUnlock(); + } + List incompatibleFrontends = getIncompatibleFrontendsForDistributionMapping(); + if (!incompatibleFrontends.isEmpty()) { + logDistributionMappingFallback(table, "frontend versions are mixed or unknown; current version: " + + getCurrentFrontendVersion() + ", incompatible frontends: " + incompatibleFrontends); + return ImmutableList.of(); + } + DistributionMappingConstraint incompatibleConstraint = + findIncompatibleDistributionMappingConstraint(table, constraints); + if (incompatibleConstraint != null) { + logDistributionMappingFallback(table, "constraint " + incompatibleConstraint.getName() + + " is incompatible with current schema version " + table.getBaseSchemaVersion() + + "; drop and recreate it to re-enable the optimization"); + return ImmutableList.of(); + } + return constraints; + } + + /** Reject persisted mappings that no longer describe the current table schema. */ + public void validateDistributionMappingConstraints(OlapTable table) { + validateDistributionMappingConstraints(table, getDistributionMappingConstraints(table)); + } + + private void validateDistributionMappingConstraints(OlapTable table, + List constraints) { + DistributionMappingConstraint incompatibleConstraint = + findIncompatibleDistributionMappingConstraint(table, constraints); + if (incompatibleConstraint != null) { + throw new AnalysisException(String.format( + "Distribution mapping constraint %s on table %s is incompatible with the current schema. " + + "Drop and recreate the constraint.", + incompatibleConstraint.getName(), table.getName())); + } + } + + private DistributionMappingConstraint findIncompatibleDistributionMappingConstraint( + OlapTable table, List constraints) { + for (DistributionMappingConstraint constraint : constraints) { + if (!constraint.isCompatibleWith(table)) { + return constraint; + } + } + return null; + } + + private void logDistributionMappingFallback(OlapTable table, String reason) { + if (DISTRIBUTION_MAPPING_FALLBACK_LOG_LIMITER.tryAcquire()) { + LOG.warn("Ignore distribution mapping constraints on table {} (id={}) during query planning and " + + "fall back to regular planning: {}", + table.getName(), table.getId(), reason); + } + } + + /** Return the mapping that uses the given column, if any. */ + public String findDistributionMappingConstraintWithColumn(TableIf table, String columnName) { + readLock(); + try { + for (Entry entry + : getDistributionMappingConstraintsMap(table).entrySet()) { + DistributionMappingConstraint mapping = entry.getValue(); + if (containsIgnoreCase(mapping.getDeterminantColumnNames(), columnName) + || containsIgnoreCase(mapping.getDistributionColumnNames(), columnName)) { + return entry.getKey(); + } + } + return null; + } finally { + readUnlock(); + } + } + /** Returns all PrimaryKeyConstraints for the given table. */ public ImmutableList getPrimaryKeyConstraints( TableNameInfo tableNameInfo) { @@ -661,12 +901,12 @@ public void dropAndRenameConstraints(TableNameInfo oldTable, // ==================== Private helpers ==================== private void checkConstraintNotExistence(String name, - Constraint constraint, Map constraintMap) { + Constraint constraint, Map constraintMap) { if (constraintMap.containsKey(name)) { throw new AnalysisException( String.format("Constraint name %s has existed", name)); } - for (Entry entry : constraintMap.entrySet()) { + for (Entry entry : constraintMap.entrySet()) { if (entry.getValue().equals(constraint)) { throw new AnalysisException(String.format( "Constraint %s has existed, named %s", @@ -675,6 +915,22 @@ private void checkConstraintNotExistence(String name, } } + private void validateNoDistributionMappingConstraintNameCollision(String tableKey, + Map mappings) { + Map centralizedConstraints = constraintsMap.get(tableKey); + if (centralizedConstraints == null) { + return; + } + for (String mappingName : mappings.keySet()) { + if (centralizedConstraints.containsKey(mappingName)) { + throw new AnalysisException(String.format( + "Distribution mapping constraint %s on table %s conflicts with another constraint " + + "of the same name. Drop one of the conflicting constraints.", + mappingName, tableKey)); + } + } + } + /** * For FK constraints: find the matching PK on the referenced table * (using FK's referencedTableInfo) and register the FK table in PK's @@ -867,7 +1123,7 @@ private ImmutableList getConstraintsByType( * Validate that the table and columns referenced by the constraint * actually exist. Only called for non-replay operations. */ - private void validateTableAndColumns(TableNameInfo tableNameInfo, + private TableIf validateTableAndColumns(TableNameInfo tableNameInfo, Constraint constraint) { TableIf table = resolveTableForValidation(tableNameInfo); if (constraint instanceof PrimaryKeyConstraint) { @@ -894,6 +1150,7 @@ private void validateTableAndColumns(TableNameInfo tableNameInfo, toKey(refTableInfo)); } } + return table; } private TableIf resolveTableForValidation( @@ -945,6 +1202,92 @@ private void validateColumnsExist(TableIf table, } } + private void validateDistributionMappingConstraint(TableNameInfo tableNameInfo, + OlapTable table, DistributionMappingConstraint constraint) { + if (table.getCatalogId() != InternalCatalog.INTERNAL_CATALOG_ID) { + throw new AnalysisException("Distribution mapping constraint only supports internal OLAP tables"); + } + if (table.isTemporary()) { + throw new AnalysisException("Distribution mapping constraint does not support temporary tables"); + } + if (table.isBeingSynced()) { + throw new AnalysisException( + "Distribution mapping constraint does not support tables being synchronized by CCR"); + } + validateColumnsExist(table, constraint.getDeterminantColumnNames(), toKey(tableNameInfo)); + validateColumnsExist(table, constraint.getDistributionColumnNames(), toKey(tableNameInfo)); + + TreeSet determinantColumns = new TreeSet<>(String.CASE_INSENSITIVE_ORDER); + determinantColumns.addAll(constraint.getDeterminantColumnNames()); + if (determinantColumns.size() != constraint.getDeterminantColumnNames().size()) { + throw new AnalysisException("Determinant columns in distribution mapping constraint must be unique"); + } + TreeSet distributionColumns = new TreeSet<>(String.CASE_INSENSITIVE_ORDER); + distributionColumns.addAll(constraint.getDistributionColumnNames()); + if (distributionColumns.size() != constraint.getDistributionColumnNames().size()) { + throw new AnalysisException("Distribution columns in distribution mapping constraint must be unique"); + } + + if (!(table.getDefaultDistributionInfo() instanceof HashDistributionInfo)) { + throw new AnalysisException("Distribution mapping constraint requires hash distribution"); + } + if (!constraint.hasCompatibleDistributionColumns(table)) { + throw new AnalysisException("Distribution columns in distribution mapping constraint" + + " must be an ordered subset of table distribution columns"); + } + } + + /** Reject ADD and restore until every registered FE reports this exact build. */ + public void validateDistributionMappingFeatureCompatibility() { + String currentVersion = getCurrentFrontendVersion(); + List incompatibleFrontends = getIncompatibleFrontendsForDistributionMapping(); + if (!incompatibleFrontends.isEmpty()) { + throw new AnalysisException("Distribution mapping constraints cannot be added or restored while" + + " frontend versions are mixed or unknown. Current version: " + currentVersion + + ", incompatible frontends: " + incompatibleFrontends); + } + } + + private String getCurrentFrontendVersion() { + return Version.DORIS_BUILD_VERSION + "-" + Version.DORIS_BUILD_SHORT_HASH; + } + + private List getIncompatibleFrontendsForDistributionMapping() { + String currentVersion = getCurrentFrontendVersion(); + List incompatibleFrontends = new ArrayList<>(); + for (Frontend frontend : Env.getCurrentEnv().getFrontends(null)) { + String frontendVersion = frontend.getVersion(); + if (!currentVersion.equals(frontendVersion)) { + incompatibleFrontends.add(frontend.getNodeName() + "(" + frontendVersion + ")"); + } + } + Collections.sort(incompatibleFrontends); + return incompatibleFrontends; + } + + private Map getDistributionMappingConstraintsMap(TableIf table) { + if (!(table instanceof OlapTable) + || ((OlapTable) table).getCatalogId() != InternalCatalog.INTERNAL_CATALOG_ID) { + return Collections.emptyMap(); + } + TableProperty tableProperty = ((OlapTable) table).getTableProperty(); + return tableProperty == null + ? Collections.emptyMap() : tableProperty.getDistributionMappingConstraints(); + } + + private TableProperty getOrCreateTableProperty(OlapTable table) { + TableProperty tableProperty = table.getTableProperty(); + if (tableProperty == null) { + tableProperty = new TableProperty(new HashMap<>()); + table.setTableProperty(tableProperty); + } + return tableProperty; + } + + private boolean containsIgnoreCase(Collection columnNames, String columnName) { + return columnNames.stream().anyMatch(column -> column.equalsIgnoreCase(columnName)); + } + // ==================== Swap helpers ==================== private void swapForeignKeyReference(ForeignKeyConstraint fk, diff --git a/fe/fe-core/src/main/java/org/apache/doris/catalog/constraint/DistributionMappingConstraint.java b/fe/fe-core/src/main/java/org/apache/doris/catalog/constraint/DistributionMappingConstraint.java new file mode 100644 index 00000000000000..66f871b0b22d08 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/catalog/constraint/DistributionMappingConstraint.java @@ -0,0 +1,222 @@ +// 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.HashDistributionInfo; +import org.apache.doris.catalog.OlapTable; + +import com.google.common.base.Objects; +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; +import com.google.gson.annotations.SerializedName; + +import java.util.List; + +/** + * Declares that determinant columns use the named cross-table mapping to determine distribution columns. + */ +public class DistributionMappingConstraint extends Constraint { + @SerializedName(value = "mi") + private final String mappingId; + @SerializedName(value = "dc") + private final List determinantColumns; + @SerializedName(value = "tc") + private final List distributionColumns; + @SerializedName(value = "sv") + private final Integer baseSchemaVersion; + @SerializedName(value = "di") + private final List determinantColumnUniqueIds; + @SerializedName(value = "ti") + private final List distributionColumnUniqueIds; + @SerializedName(value = "ds") + private final List determinantColumnTypeSignatures; + @SerializedName(value = "ts") + private final List distributionColumnTypeSignatures; + + /** Constructor. */ + public DistributionMappingConstraint(String name, String mappingId, + List determinantColumns, List distributionColumns) { + this(name, mappingId, determinantColumns, distributionColumns, + null, ImmutableList.of(), ImmutableList.of(), ImmutableList.of(), ImmutableList.of()); + } + + private DistributionMappingConstraint(String name, String mappingId, + List determinantColumns, List distributionColumns, + Integer baseSchemaVersion, List determinantColumnUniqueIds, + List distributionColumnUniqueIds, List determinantColumnTypeSignatures, + List distributionColumnTypeSignatures) { + super(ConstraintType.DISTRIBUTION_MAPPING, name); + this.mappingId = mappingId; + this.determinantColumns = ImmutableList.copyOf(determinantColumns); + this.distributionColumns = ImmutableList.copyOf(distributionColumns); + this.baseSchemaVersion = baseSchemaVersion; + this.determinantColumnUniqueIds = ImmutableList.copyOf(determinantColumnUniqueIds); + this.distributionColumnUniqueIds = ImmutableList.copyOf(distributionColumnUniqueIds); + this.determinantColumnTypeSignatures = ImmutableList.copyOf(determinantColumnTypeSignatures); + this.distributionColumnTypeSignatures = ImmutableList.copyOf(distributionColumnTypeSignatures); + } + + public String getMappingId() { + return mappingId; + } + + public List getDeterminantColumnNames() { + return determinantColumns; + } + + public List getDistributionColumnNames() { + return distributionColumns; + } + + public Integer getBaseSchemaVersion() { + return baseSchemaVersion; + } + + public List getDeterminantColumnUniqueIds() { + return determinantColumnUniqueIds; + } + + public List getDistributionColumnUniqueIds() { + return distributionColumnUniqueIds; + } + + public List getDeterminantColumnTypeSignatures() { + return determinantColumnTypeSignatures; + } + + public List getDistributionColumnTypeSignatures() { + return distributionColumnTypeSignatures; + } + + DistributionMappingConstraint bindTo(OlapTable table) { + return new DistributionMappingConstraint( + getName(), mappingId, determinantColumns, distributionColumns, + table.getBaseSchemaVersion(), getColumnUniqueIds(table, determinantColumns), + getColumnUniqueIds(table, distributionColumns), + getColumnTypeSignatures(table, determinantColumns), + getColumnTypeSignatures(table, distributionColumns)); + } + + boolean isCompatibleWith(OlapTable table) { + if (!hasCompatibleDistributionColumns(table) + || baseSchemaVersion == null + || determinantColumnUniqueIds == null + || distributionColumnUniqueIds == null + || determinantColumnTypeSignatures == null + || distributionColumnTypeSignatures == null + || determinantColumns.size() != determinantColumnUniqueIds.size() + || distributionColumns.size() != distributionColumnUniqueIds.size() + || determinantColumns.size() != determinantColumnTypeSignatures.size() + || distributionColumns.size() != distributionColumnTypeSignatures.size()) { + return false; + } + boolean sameSchemaVersion = baseSchemaVersion == table.getBaseSchemaVersion(); + return columnsMatch(table, determinantColumns, determinantColumnUniqueIds, + determinantColumnTypeSignatures, sameSchemaVersion) + && columnsMatch(table, distributionColumns, distributionColumnUniqueIds, + distributionColumnTypeSignatures, sameSchemaVersion); + } + + boolean hasCompatibleDistributionColumns(OlapTable table) { + if (!(table.getDefaultDistributionInfo() instanceof HashDistributionInfo)) { + return false; + } + List tableDistributionColumns = + ((HashDistributionInfo) table.getDefaultDistributionInfo()).getDistributionColumns(); + int previousIndex = -1; + for (String distributionColumn : distributionColumns) { + int index = -1; + for (int i = 0; i < tableDistributionColumns.size(); i++) { + if (tableDistributionColumns.get(i).getName().equalsIgnoreCase(distributionColumn)) { + index = i; + break; + } + } + if (index <= previousIndex) { + return false; + } + previousIndex = index; + } + return true; + } + + private static List getColumnUniqueIds(OlapTable table, List columnNames) { + ImmutableList.Builder uniqueIds = ImmutableList.builder(); + for (String columnName : columnNames) { + Column column = table.getColumn(columnName); + Preconditions.checkNotNull(column, "column %s does not exist", columnName); + uniqueIds.add(column.getUniqueId()); + } + return uniqueIds.build(); + } + + private static List getColumnTypeSignatures(OlapTable table, List columnNames) { + ImmutableList.Builder typeSignatures = ImmutableList.builder(); + for (String columnName : columnNames) { + Column column = table.getColumn(columnName); + Preconditions.checkNotNull(column, "column %s does not exist", columnName); + typeSignatures.add(column.getType().toSql()); + } + return typeSignatures.build(); + } + + private static boolean columnsMatch(OlapTable table, List columnNames, + List expectedUniqueIds, List expectedTypeSignatures, boolean sameSchemaVersion) { + for (int i = 0; i < columnNames.size(); i++) { + Column column = table.getColumn(columnNames.get(i)); + if (column == null || !column.getType().toSql().equals(expectedTypeSignatures.get(i))) { + return false; + } + int expectedUniqueId = expectedUniqueIds.get(i); + if (expectedUniqueId == Column.COLUMN_UNIQUE_ID_INIT_VALUE) { + if (!sameSchemaVersion) { + return false; + } + } else if (column.getUniqueId() != expectedUniqueId) { + return false; + } + } + return true; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + DistributionMappingConstraint that = (DistributionMappingConstraint) o; + return mappingId.equals(that.mappingId) + && determinantColumns.equals(that.determinantColumns) + && distributionColumns.equals(that.distributionColumns); + } + + @Override + public int hashCode() { + return Objects.hashCode(mappingId, determinantColumns, distributionColumns); + } + + @Override + public String toString() { + return String.format("COLOCATE MAPPING %s (%s) DETERMINES DISTRIBUTION KEY (%s) NOT ENFORCED", + mappingId, String.join(", ", determinantColumns), String.join(", ", distributionColumns)); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/cascades/CostAndEnforcerJob.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/cascades/CostAndEnforcerJob.java index 31206e85e38ad5..6359d082ab1edc 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/cascades/CostAndEnforcerJob.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/cascades/CostAndEnforcerJob.java @@ -296,6 +296,9 @@ private void enforce(PhysicalProperties outputProperty, List } return; } + if (!requiredProperties.isEnforceable()) { + return; + } if (context.getRequiredProperties().isDistributionOnlyProperties()) { // For properties without an orderSpec, enforceMissingPropertiesHelper always adds a distributor diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java index 8b7126d730a0a6..e9df9f511c14b9 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java @@ -2044,6 +2044,12 @@ public LogicalPlan visitAddConstraint(AddConstraintContext ctx) { nameParts ); constraint = Constraint.newForeignKeyConstraint(curTable, slots, referenceTable, referencedSlots); + } else if (ctx.constraint().COLOCATE() != null) { + ImmutableList distributionSlots = ctx.constraint().distributionSlots.identifierSeq().ident.stream() + .map(ident -> new UnboundSlot(ident.getText())) + .collect(ImmutableList.toImmutableList()); + constraint = Constraint.newDistributionMappingConstraint( + curTable, ctx.constraint().mappingId.getText().toLowerCase(Locale.ROOT), slots, distributionSlots); } else { throw new AnalysisException("Unsupported constraint " + ctx.getText()); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/ChildOutputPropertyDeriver.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/ChildOutputPropertyDeriver.java index 2df7723a7ab052..f431f136f298a1 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/ChildOutputPropertyDeriver.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/ChildOutputPropertyDeriver.java @@ -17,6 +17,8 @@ package org.apache.doris.nereids.properties; +import org.apache.doris.catalog.Column; +import org.apache.doris.catalog.HashDistributionInfo; import org.apache.doris.nereids.PlanContext; import org.apache.doris.nereids.exceptions.AnalysisException; import org.apache.doris.nereids.memo.GroupExpression; @@ -78,8 +80,10 @@ import java.util.HashSet; import java.util.LinkedHashMap; import java.util.List; +import java.util.Locale; import java.util.Map; import java.util.Objects; +import java.util.Optional; import java.util.Set; import java.util.stream.Collectors; @@ -164,7 +168,50 @@ public PhysicalProperties visitPhysicalOlapScan(PhysicalOlapScan olapScan, PlanC if (context.getStatementContext().isShortCircuitQuery() && olapScan.getSelectedTabletIds().size() == 1) { return PhysicalProperties.GATHER; } - return new PhysicalProperties(olapScan.getDistributionSpec()); + DistributionSpec distributionSpec = olapScan.getDistributionSpec(); + PhysicalProperties properties = new PhysicalProperties(distributionSpec); + if (distributionSpec instanceof DistributionSpecHash) { + return properties.withNaturalDistributionMappingSpec( + naturalMappingSpecFromOlapScan(olapScan, (DistributionSpecHash) distributionSpec)); + } + return properties; + } + + private Optional naturalMappingSpecFromOlapScan( + PhysicalOlapScan olapScan, DistributionSpecHash hashSpec) { + if (hashSpec.getShuffleType() != ShuffleType.NATURAL + || hashSpec.getDistributionMappings().isEmpty()) { + return Optional.empty(); + } + + HashDistributionInfo distributionInfo + = (HashDistributionInfo) olapScan.getTable().getDefaultDistributionInfo(); + List distributionColumns = distributionInfo.getDistributionColumns(); + Map distributionIndices = Maps.newHashMapWithExpectedSize(distributionColumns.size()); + for (int i = 0; i < distributionColumns.size(); i++) { + distributionIndices.put(distributionColumns.get(i).getName().toLowerCase(Locale.ROOT), i); + } + + Map visibleDistributionExprs = Maps.newHashMap(); + for (Slot slot : olapScan.getOutput()) { + SlotReference slotReference = (SlotReference) slot; + if (slotReference.getOriginalColumn().isPresent()) { + String columnName = slotReference.getOriginalColumn().get() + .tryGetBaseColumnName().toLowerCase(Locale.ROOT); + Integer index = distributionIndices.get(columnName); + if (index != null) { + visibleDistributionExprs.put(slot.getExprId(), index); + } + } + } + + return Optional.of(new NaturalDistributionMappingSpec( + hashSpec.getTableId(), + hashSpec.getSelectedIndexId(), + hashSpec.getPartitionIds(), + distributionColumns.size(), + visibleDistributionExprs, + hashSpec.getDistributionMappings())); } @Override @@ -204,12 +251,50 @@ public PhysicalProperties visitPhysicalHashAggregate( && isShuffleCompatible(childOutputProperty.getDistributionSpec())) { return PhysicalProperties.ANY; } + if (agg.isDistinctOrDeduplicate()) { + return withoutNaturalDistributionMapping(childOutputProperty) + .withOrderSpec(new OrderSpec()); + } + if (childOutputProperty.getNaturalDistributionMappingSpec().isPresent()) { + return computeAggregateOutputProperties(agg, childOutputProperty) + .withOrderSpec(new OrderSpec()); + } return new PhysicalProperties(childOutputProperty.getDistributionSpec()); default: throw new RuntimeException("Could not derive output properties for agg phase: " + agg.getAggPhase()); } } + private PhysicalProperties computeAggregateOutputProperties( + PhysicalHashAggregate agg, PhysicalProperties childOutputProperty) { + NaturalDistributionMappingSpec naturalMappingSpec = + childOutputProperty.getNaturalDistributionMappingSpec().get(); + if (agg.hasSourceRepeat()) { + return withoutNaturalDistributionMapping(childOutputProperty); + } + + Set groupByExprIds = Sets.newHashSet(); + for (Expression groupBy : agg.getGroupByExpressions()) { + if (!(groupBy instanceof SlotReference)) { + return withoutNaturalDistributionMapping(childOutputProperty); + } + groupByExprIds.add(((SlotReference) groupBy).getExprId()); + } + if (!naturalMappingSpec.distributionKeysCoveredByDirectOrMapping(groupByExprIds)) { + return withoutNaturalDistributionMapping(childOutputProperty); + } + + return computeProjectOutputProperties(agg.getOutputExpressions(), childOutputProperty); + } + + private PhysicalProperties withoutNaturalDistributionMapping(PhysicalProperties properties) { + DistributionSpec distributionSpec = properties.getDistributionSpec(); + if (distributionSpec instanceof DistributionSpecHash) { + distributionSpec = ((DistributionSpecHash) distributionSpec).withoutDistributionMappings(); + } + return new PhysicalProperties(distributionSpec, properties.getOrderSpec()); + } + /** * Returns true if the child's distribution is a shuffle-compatible hash that the * bucketed fusion pattern produces (ShuffleType.REQUIRE). EXECUTION_BUCKETED is @@ -262,7 +347,7 @@ public PhysicalProperties visitPhysicalFilter(PhysicalFilter fil @Override public PhysicalProperties visitPhysicalGenerate(PhysicalGenerate generate, PlanContext context) { Preconditions.checkState(childrenOutputProperties.size() == 1); - return childrenOutputProperties.get(0); + return withoutNaturalDistributionMapping(childrenOutputProperties.get(0)); } @Override @@ -289,17 +374,30 @@ public PhysicalProperties visitPhysicalHashJoin( DistributionSpecHash mockedRightHashSpec = mockAnotherSideSpecFromConjuncts( hashJoin, (DistributionSpecHash) leftDistributionSpec); if (SessionVariable.canUseNereidsDistributePlanner()) { - return computeShuffleJoinOutputProperties(hashJoin, - (DistributionSpecHash) leftDistributionSpec, mockedRightHashSpec); + return withoutNaturalDistributionMapping(computeShuffleJoinOutputProperties(hashJoin, + (DistributionSpecHash) leftDistributionSpec, mockedRightHashSpec)); } else { - return legacyComputeShuffleJoinOutputProperties(hashJoin, - (DistributionSpecHash) leftDistributionSpec, mockedRightHashSpec); + return withoutNaturalDistributionMapping(legacyComputeShuffleJoinOutputProperties(hashJoin, + (DistributionSpecHash) leftDistributionSpec, mockedRightHashSpec)); } } else { - return new PhysicalProperties(leftDistributionSpec); + return withoutNaturalDistributionMapping(new PhysicalProperties(leftDistributionSpec)); } } + if (!(leftOutputProperty.getDistributionSpec() instanceof DistributionSpecHash + && rightOutputProperty.getDistributionSpec() instanceof DistributionSpecHash) + && leftOutputProperty.getNaturalDistributionMappingSpec().isPresent() + && rightOutputProperty.getNaturalDistributionMappingSpec().isPresent() + && JoinUtils.couldColocateJoinByMapping( + leftOutputProperty.getNaturalDistributionMappingSpec().get(), + rightOutputProperty.getNaturalDistributionMappingSpec().get(), + hashJoin.getHashJoinConjuncts())) { + // The join is colocated, but hidden bucket positions are not executable output slots. + // Do not expose them as a reusable hash distribution above this join. + return PhysicalProperties.STORAGE_ANY; + } + // shuffle if (leftOutputProperty.getDistributionSpec() instanceof DistributionSpecHash && rightOutputProperty.getDistributionSpec() instanceof DistributionSpecHash) { @@ -327,7 +425,8 @@ public PhysicalProperties visitPhysicalNestedLoopJoin( PlanContext context) { Preconditions.checkState(childrenOutputProperties.size() == 2); PhysicalProperties leftOutputProperty = childrenOutputProperties.get(0); - return new PhysicalProperties(leftOutputProperty.getDistributionSpec()); + return withoutNaturalDistributionMapping( + new PhysicalProperties(leftOutputProperty.getDistributionSpec())); } /** @@ -338,31 +437,32 @@ public static PhysicalProperties computeProjectOutputProperties( PhysicalProperties childProperties) { DistributionSpec childDistributionSpec = childProperties.getDistributionSpec(); OrderSpec childOrderSpec = childProperties.getOrderSpec(); - if (childDistributionSpec instanceof DistributionSpecHash) { - Map projections = Maps.newHashMap(); - Set obstructions = Sets.newHashSet(); - for (NamedExpression namedExpression : projects) { - if (namedExpression instanceof Alias) { - Alias alias = (Alias) namedExpression; - Expression child = alias.child(); - if (child instanceof SlotReference) { - projections.put(((SlotReference) child).getExprId(), alias.getExprId()); - } else if (child instanceof Cast && child.child(0) instanceof Slot - && isSameHashValue(child.child(0).getDataType(), child.getDataType())) { - // cast(slot as varchar(10)) can do projection if slot is varchar(3) - projections.put(((Slot) child.child(0)).getExprId(), alias.getExprId()); - } else { - obstructions.addAll( - child.getInputSlots().stream() - .map(NamedExpression::getExprId) - .collect(Collectors.toSet())); - } + Map projections = Maps.newHashMap(); + Set obstructions = Sets.newHashSet(); + for (NamedExpression namedExpression : projects) { + if (namedExpression instanceof Alias) { + Alias alias = (Alias) namedExpression; + Expression child = alias.child(); + if (child instanceof SlotReference) { + projections.put(((SlotReference) child).getExprId(), alias.getExprId()); + } else if (child instanceof Cast && child.child(0) instanceof Slot + && isHashValuePreservingCast(child.child(0).getDataType(), child.getDataType())) { + // cast(slot as varchar(10)) can do projection if slot is varchar(3) + projections.put(((Slot) child.child(0)).getExprId(), alias.getExprId()); } else { - // namedExpression is slot - projections.put(namedExpression.getExprId(), namedExpression.getExprId()); + obstructions.addAll( + child.getInputSlots().stream() + .map(NamedExpression::getExprId) + .collect(Collectors.toSet())); } + } else { + // namedExpression is slot + projections.put(namedExpression.getExprId(), namedExpression.getExprId()); } + } + PhysicalProperties projectedProperties = childProperties; + if (childDistributionSpec instanceof DistributionSpecHash) { DistributionSpecHash childDistributionSpecHash = (DistributionSpecHash) childDistributionSpec; boolean canUseChildProperties = true; for (ExprId exprId : childDistributionSpecHash.getOrderedShuffledColumns()) { @@ -372,17 +472,21 @@ && isSameHashValue(child.child(0).getDataType(), child.getDataType())) { } } - if (canUseChildProperties) { - return childProperties; + if (!canUseChildProperties || !childDistributionSpecHash.getDistributionMappings().isEmpty()) { + DistributionSpec defaultAnySpec = childDistributionSpecHash.getShuffleType() == ShuffleType.NATURAL + ? DistributionSpecStorageAny.INSTANCE : DistributionSpecAny.INSTANCE; + boolean allDistributionKeysProjected = childDistributionSpecHash.getOrderedShuffledColumns().stream() + .allMatch(projections::containsKey); + DistributionSpec outputDistributionSpec = + childProperties.getNaturalDistributionMappingSpec().isPresent() + && !allDistributionKeysProjected + ? defaultAnySpec + : childDistributionSpecHash.project(projections, obstructions, defaultAnySpec); + projectedProperties = new PhysicalProperties(outputDistributionSpec, childOrderSpec); } - DistributionSpec defaultAnySpec = childDistributionSpecHash.getShuffleType() == ShuffleType.NATURAL - ? DistributionSpecStorageAny.INSTANCE : DistributionSpecAny.INSTANCE; - DistributionSpec outputDistributionSpec = childDistributionSpecHash.project( - projections, obstructions, defaultAnySpec); - return new PhysicalProperties(outputDistributionSpec, childOrderSpec); - } else { - return childProperties; } + return projectedProperties.withNaturalDistributionMappingSpec( + childProperties.getNaturalDistributionMappingSpec().flatMap(spec -> spec.project(projections))); } @Override @@ -416,7 +520,7 @@ public PhysicalProperties visitPhysicalRepeat(PhysicalRepeat rep intersectGroupingKeysId.add(((SlotReference) key).getExprId()); } if (intersectGroupingKeysId.containsAll(orderedShuffledColumns)) { - return childrenOutputProperties.get(0); + return withoutNaturalDistributionMapping(childrenOutputProperties.get(0)); } } output = PhysicalProperties.createAnyFromHash((DistributionSpecHash) childDistributionSpec); @@ -431,7 +535,7 @@ public PhysicalProperties visitPhysicalPartitionTopN(PhysicalPartitionTopN window, PlanContext context) { Preconditions.checkState(childrenOutputProperties.size() == 1); - return childrenOutputProperties.get(0); + return withoutNaturalDistributionMapping(childrenOutputProperties.get(0)); } private PhysicalProperties computeShuffleJoinOutputProperties( @@ -757,7 +863,7 @@ private DistributionSpecHash mockAnotherSideSpecFromConjuncts( return new DistributionSpecHash(anotherSideOrderedExprIds, oneSideSpec.getShuffleType()); } - private static boolean isSameHashValue(DataType originType, DataType castType) { + static boolean isHashValuePreservingCast(DataType originType, DataType castType) { if (originType.isStringLikeType() && (castType.isVarcharType() || castType.isStringType()) && (castType.width() >= originType.width() || castType.width() < 0)) { return true; diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/ChildrenPropertiesRegulator.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/ChildrenPropertiesRegulator.java index e94e4f7dfae20f..4427e8876b0263 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/ChildrenPropertiesRegulator.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/ChildrenPropertiesRegulator.java @@ -119,6 +119,14 @@ public List> visitPhysicalHashAggregate( return ImmutableList.of(); } PhysicalProperties requiredChildProperty = requiredProperties.get(0); + DistributionSpec requiredChildDistribution = requiredChildProperty.getDistributionSpec(); + if (requiredChildDistribution instanceof DistributionSpecHash + && ((DistributionSpecHash) requiredChildDistribution).getShuffleType() + == ShuffleType.COLOCATE_MAPPING_REQUIRE) { + return originChildrenProperties.get(0).satisfy(requiredChildProperty) + ? ImmutableList.of(originChildrenProperties) + : ImmutableList.of(); + } if (!agg.getAggregateParam().canBeBanned) { return visit(agg, context); } @@ -429,6 +437,32 @@ public List> visitPhysicalHashJoin( return ImmutableList.of(originChildrenProperties); } + DistributionSpec leftRequiredSpec = requiredProperties.get(0).getDistributionSpec(); + DistributionSpec rightRequiredSpec = requiredProperties.get(1).getDistributionSpec(); + boolean leftRequiresColocateMapping = leftRequiredSpec instanceof DistributionSpecHash + && ((DistributionSpecHash) leftRequiredSpec).getShuffleType() + == ShuffleType.COLOCATE_MAPPING_REQUIRE; + boolean rightRequiresColocateMapping = rightRequiredSpec instanceof DistributionSpecHash + && ((DistributionSpecHash) rightRequiredSpec).getShuffleType() + == ShuffleType.COLOCATE_MAPPING_REQUIRE; + Preconditions.checkState(leftRequiresColocateMapping == rightRequiresColocateMapping, + "colocate mapping join requires matching child properties"); + if (leftRequiresColocateMapping) { + Optional leftMappingSpec = + originChildrenProperties.get(0).getNaturalDistributionMappingSpec(); + Optional rightMappingSpec = + originChildrenProperties.get(1).getNaturalDistributionMappingSpec(); + if (!originChildrenProperties.get(0).satisfy(requiredProperties.get(0)) + || !originChildrenProperties.get(1).satisfy(requiredProperties.get(1)) + || !leftMappingSpec.isPresent() + || !rightMappingSpec.isPresent() + || !JoinUtils.couldColocateJoinByMapping( + leftMappingSpec.get(), rightMappingSpec.get(), hashJoin.getHashJoinConjuncts())) { + return ImmutableList.of(); + } + return ImmutableList.of(originChildrenProperties); + } + // shuffle if (!(leftDistributionSpec instanceof DistributionSpecHash) || !(rightDistributionSpec instanceof DistributionSpecHash)) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/DistributionMapping.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/DistributionMapping.java new file mode 100644 index 00000000000000..c2ac0e47447ab9 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/DistributionMapping.java @@ -0,0 +1,97 @@ +// 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.properties; + +import org.apache.doris.nereids.trees.expressions.ExprId; + +import com.google.common.collect.ImmutableList; + +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; + +/** + * A named cross-table mapping from determinant expressions to positions in the storage distribution key. + */ +public class DistributionMapping { + private final String mappingId; + private final List determinantExprIds; + private final List targetDistributionIndices; + + /** Constructor. */ + public DistributionMapping(String mappingId, List determinantExprIds, + List targetDistributionIndices) { + this.mappingId = Objects.requireNonNull(mappingId, "mappingId should not be null"); + this.determinantExprIds = ImmutableList.copyOf(determinantExprIds); + this.targetDistributionIndices = ImmutableList.copyOf(targetDistributionIndices); + } + + public String getMappingId() { + return mappingId; + } + + public List getDeterminantExprIds() { + return determinantExprIds; + } + + public List getTargetDistributionIndices() { + return targetDistributionIndices; + } + + /** Remap determinant expressions through a projection, or drop the mapping if any determinant is absent. */ + public Optional project(Map projections) { + ImmutableList.Builder projected = ImmutableList.builderWithExpectedSize(determinantExprIds.size()); + for (ExprId determinant : determinantExprIds) { + ExprId projectedExprId = projections.get(determinant); + if (projectedExprId == null) { + return Optional.empty(); + } + projected.add(projectedExprId); + } + return Optional.of(new DistributionMapping(mappingId, projected.build(), targetDistributionIndices)); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof DistributionMapping)) { + return false; + } + DistributionMapping that = (DistributionMapping) o; + return mappingId.equals(that.mappingId) + && determinantExprIds.equals(that.determinantExprIds) + && targetDistributionIndices.equals(that.targetDistributionIndices); + } + + @Override + public int hashCode() { + return Objects.hash(mappingId, determinantExprIds, targetDistributionIndices); + } + + @Override + public String toString() { + return "DistributionMapping{" + + "mappingId='" + mappingId + '\'' + + ", determinantExprIds=" + determinantExprIds + + ", targetDistributionIndices=" + targetDistributionIndices + + '}'; + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/DistributionSpecHash.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/DistributionSpecHash.java index ab96960684a154..38b49e182a9a3b 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/DistributionSpecHash.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/DistributionSpecHash.java @@ -48,6 +48,7 @@ public class DistributionSpecHash extends DistributionSpec { // use for satisfied judge private final List> equivalenceExprIds; private final Map exprIdToEquivalenceSet; + private final List distributionMappings; // below two attributes use for colocate join, only store one table info is enough private final long tableId; @@ -74,6 +75,13 @@ public DistributionSpecHash(List orderedShuffledColumns, ShuffleType shu */ public DistributionSpecHash(List orderedShuffledColumns, ShuffleType shuffleType, long tableId, long selectedIndexId, Set partitionIds) { + this(orderedShuffledColumns, shuffleType, tableId, selectedIndexId, partitionIds, ImmutableList.of()); + } + + /** Normal constructor with storage distribution mappings. */ + public DistributionSpecHash(List orderedShuffledColumns, ShuffleType shuffleType, + long tableId, long selectedIndexId, Set partitionIds, + List distributionMappings) { this.orderedShuffledColumns = ImmutableList.copyOf( Objects.requireNonNull(orderedShuffledColumns, "orderedShuffledColumns should not null")); this.shuffleType = Objects.requireNonNull(shuffleType, "shuffleType should not null"); @@ -81,6 +89,7 @@ public DistributionSpecHash(List orderedShuffledColumns, ShuffleType shu Objects.requireNonNull(partitionIds, "partitionIds should not null")); this.tableId = tableId; this.selectedIndexId = selectedIndexId; + this.distributionMappings = ImmutableList.copyOf(distributionMappings); ImmutableList.Builder> equivalenceExprIdsBuilder = ImmutableList.builderWithExpectedSize(orderedShuffledColumns.size()); ImmutableMap.Builder exprIdToEquivalenceSetBuilder @@ -110,6 +119,14 @@ public DistributionSpecHash(List orderedShuffledColumns, ShuffleType shu public DistributionSpecHash(List orderedShuffledColumns, ShuffleType shuffleType, long tableId, long selectedIndexId, Set partitionIds, List> equivalenceExprIds, Map exprIdToEquivalenceSet) { + this(orderedShuffledColumns, shuffleType, tableId, selectedIndexId, partitionIds, + equivalenceExprIds, exprIdToEquivalenceSet, ImmutableList.of()); + } + + /** Constructor with precomputed equivalence sets and storage distribution mappings. */ + public DistributionSpecHash(List orderedShuffledColumns, ShuffleType shuffleType, long tableId, + long selectedIndexId, Set partitionIds, List> equivalenceExprIds, + Map exprIdToEquivalenceSet, List distributionMappings) { this.orderedShuffledColumns = ImmutableList.copyOf(Objects.requireNonNull(orderedShuffledColumns, "orderedShuffledColumns should not null")); this.shuffleType = Objects.requireNonNull(shuffleType, "shuffleType should not null"); @@ -121,6 +138,7 @@ public DistributionSpecHash(List orderedShuffledColumns, ShuffleType shu Objects.requireNonNull(equivalenceExprIds, "equivalenceExprIds should not null")); this.exprIdToEquivalenceSet = ImmutableMap.copyOf( Objects.requireNonNull(exprIdToEquivalenceSet, "exprIdToEquivalenceSet should not null")); + this.distributionMappings = ImmutableList.copyOf(distributionMappings); } static DistributionSpecHash merge(DistributionSpecHash left, DistributionSpecHash right, ShuffleType shuffleType) { @@ -175,6 +193,18 @@ public Map getExprIdToEquivalenceSet() { return exprIdToEquivalenceSet; } + public List getDistributionMappings() { + return distributionMappings; + } + + public DistributionSpecHash withoutDistributionMappings() { + if (distributionMappings.isEmpty()) { + return this; + } + return new DistributionSpecHash(orderedShuffledColumns, shuffleType, tableId, selectedIndexId, partitionIds, + equivalenceExprIds, exprIdToEquivalenceSet, ImmutableList.of()); + } + public Set getEquivalenceExprIdsOf(ExprId exprId) { if (exprIdToEquivalenceSet.containsKey(exprId)) { return equivalenceExprIds.get(exprIdToEquivalenceSet.get(exprId)); @@ -228,13 +258,15 @@ private boolean equalsSatisfy(List required) { } public DistributionSpecHash withShuffleType(ShuffleType shuffleType) { + List mappings = this.shuffleType == ShuffleType.NATURAL + && shuffleType == ShuffleType.NATURAL ? distributionMappings : ImmutableList.of(); return new DistributionSpecHash(orderedShuffledColumns, shuffleType, tableId, selectedIndexId, partitionIds, - equivalenceExprIds, exprIdToEquivalenceSet); + equivalenceExprIds, exprIdToEquivalenceSet, mappings); } public DistributionSpecHash withShuffleTypeAndForbidColocateJoin(ShuffleType shuffleType) { return new DistributionSpecHash(orderedShuffledColumns, shuffleType, -1, -1, partitionIds, - equivalenceExprIds, exprIdToEquivalenceSet); + equivalenceExprIds, exprIdToEquivalenceSet, ImmutableList.of()); } /** @@ -303,8 +335,14 @@ public DistributionSpec project(Map projections, exprIdToEquivalenceSet.put(exprIdSetKV.getKey(), exprIdSetKV.getValue()); } } + ImmutableList.Builder projectedMappings = ImmutableList.builder(); + if (this.orderedShuffledColumns.stream().allMatch(projections::containsKey)) { + for (DistributionMapping distributionMapping : distributionMappings) { + distributionMapping.project(projections).ifPresent(projectedMappings::add); + } + } return new DistributionSpecHash(orderedShuffledColumns, shuffleType, tableId, selectedIndexId, partitionIds, - equivalenceExprIds, exprIdToEquivalenceSet); + equivalenceExprIds, exprIdToEquivalenceSet, projectedMappings.build()); } @Override @@ -313,12 +351,14 @@ public boolean equals(Object o) { return false; } DistributionSpecHash that = (DistributionSpecHash) o; - return shuffleType == that.shuffleType && orderedShuffledColumns.equals(that.orderedShuffledColumns); + return shuffleType == that.shuffleType + && orderedShuffledColumns.equals(that.orderedShuffledColumns) + && distributionMappings.equals(that.distributionMappings); } @Override public int hashCode() { - return Objects.hash(shuffleType, orderedShuffledColumns); + return Objects.hash(shuffleType, orderedShuffledColumns, distributionMappings); } @Override @@ -330,7 +370,8 @@ public String toString() { "selectedIndexId", selectedIndexId, "partitionIds", partitionIds, "equivalenceExprIds", equivalenceExprIds, - "exprIdToEquivalenceSet", exprIdToEquivalenceSet); + "exprIdToEquivalenceSet", exprIdToEquivalenceSet, + "distributionMappings", distributionMappings); } /** @@ -346,6 +387,8 @@ public enum ShuffleType { // output, for shuffle by storage hash method STORAGE_BUCKETED, // require, need to satisfy the distribution spec by equals. - REQUIRE_EQUAL + REQUIRE_EQUAL, + // Non-enforceable proof request allowing NATURAL distribution mappings to cover required keys. + COLOCATE_MAPPING_REQUIRE } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/NaturalDistributionMappingSpec.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/NaturalDistributionMappingSpec.java new file mode 100644 index 00000000000000..16349fedc69083 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/NaturalDistributionMappingSpec.java @@ -0,0 +1,167 @@ +// 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.properties; + +import org.apache.doris.nereids.trees.expressions.ExprId; + +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; + +import java.util.BitSet; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.Set; + +/** + * Describes storage bucket locality that remains valid after distribution-key slots are projected out. + * + *

This property is only a proof artifact for mapping-based colocate join. It must never be used to + * build an Exchange or a bucket-shuffle requirement because some distribution positions may have no + * materialized output slot. + */ +public class NaturalDistributionMappingSpec { + private final long tableId; + private final long selectedIndexId; + private final Set partitionIds; + private final int distributionKeyCount; + private final Map visibleDistributionExprToIndex; + private final List distributionMappings; + + /** Constructor. */ + public NaturalDistributionMappingSpec(long tableId, long selectedIndexId, Set partitionIds, + int distributionKeyCount, Map visibleDistributionExprToIndex, + List distributionMappings) { + Preconditions.checkArgument(distributionKeyCount > 0, "distributionKeyCount must be positive"); + this.tableId = tableId; + this.selectedIndexId = selectedIndexId; + this.partitionIds = ImmutableSet.copyOf(partitionIds); + this.distributionKeyCount = distributionKeyCount; + this.visibleDistributionExprToIndex = ImmutableMap.copyOf(visibleDistributionExprToIndex); + this.distributionMappings = ImmutableList.copyOf(distributionMappings); + } + + public long getTableId() { + return tableId; + } + + public long getSelectedIndexId() { + return selectedIndexId; + } + + public Set getPartitionIds() { + return partitionIds; + } + + public int getDistributionKeyCount() { + return distributionKeyCount; + } + + public Map getVisibleDistributionExprToIndex() { + return visibleDistributionExprToIndex; + } + + public List getDistributionMappings() { + return distributionMappings; + } + + /** Return whether direct slots and complete mapping determinants cover every bucket position. */ + public boolean distributionKeysCoveredByDirectOrMapping(Set exprIds) { + BitSet coveredIndices = new BitSet(distributionKeyCount); + for (ExprId exprId : exprIds) { + Integer index = visibleDistributionExprToIndex.get(exprId); + if (index != null) { + coveredIndices.set(index); + } + } + for (DistributionMapping mapping : distributionMappings) { + if (exprIds.containsAll(mapping.getDeterminantExprIds())) { + mapping.getTargetDistributionIndices().forEach(coveredIndices::set); + } + } + return coveredIndices.nextClearBit(0) >= distributionKeyCount; + } + + /** Return whether direct slots and mapping determinants cover every underlying bucket position. */ + public boolean satisfy(List requiredExprIds) { + return distributionKeysCoveredByDirectOrMapping(ImmutableSet.copyOf(requiredExprIds)); + } + + /** + * Remap visible distribution slots and determinants through a projection. + * Missing slots are intentionally omitted while the underlying bucket positions remain unchanged. + */ + public Optional project(Map projections) { + ImmutableMap.Builder visibleDistributionExprs = ImmutableMap.builder(); + for (Map.Entry entry : visibleDistributionExprToIndex.entrySet()) { + ExprId projected = projections.get(entry.getKey()); + if (projected != null) { + visibleDistributionExprs.put(projected, entry.getValue()); + } + } + + ImmutableList.Builder projectedMappings = ImmutableList.builder(); + for (DistributionMapping mapping : distributionMappings) { + mapping.project(projections).ifPresent(projectedMappings::add); + } + List mappings = projectedMappings.build(); + if (mappings.isEmpty()) { + return Optional.empty(); + } + return Optional.of(new NaturalDistributionMappingSpec(tableId, selectedIndexId, partitionIds, + distributionKeyCount, visibleDistributionExprs.buildKeepingLast(), mappings)); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof NaturalDistributionMappingSpec)) { + return false; + } + NaturalDistributionMappingSpec that = (NaturalDistributionMappingSpec) o; + return tableId == that.tableId + && selectedIndexId == that.selectedIndexId + && distributionKeyCount == that.distributionKeyCount + && partitionIds.equals(that.partitionIds) + && visibleDistributionExprToIndex.equals(that.visibleDistributionExprToIndex) + && distributionMappings.equals(that.distributionMappings); + } + + @Override + public int hashCode() { + return Objects.hash(tableId, selectedIndexId, partitionIds, distributionKeyCount, + visibleDistributionExprToIndex, distributionMappings); + } + + @Override + public String toString() { + return "NaturalDistributionMappingSpec{" + + "tableId=" + tableId + + ", selectedIndexId=" + selectedIndexId + + ", partitionIds=" + partitionIds + + ", distributionKeyCount=" + distributionKeyCount + + ", visibleDistributionExprToIndex=" + visibleDistributionExprToIndex + + ", distributionMappings=" + distributionMappings + + '}'; + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/PhysicalProperties.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/PhysicalProperties.java index c28d6ac3cb4d47..d026914b43ba76 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/PhysicalProperties.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/PhysicalProperties.java @@ -25,6 +25,7 @@ import java.util.Collection; import java.util.List; import java.util.Objects; +import java.util.Optional; import java.util.stream.Collectors; /** @@ -59,26 +60,32 @@ public class PhysicalProperties { private final DistributionSpec distributionSpec; + private final Optional naturalDistributionMappingSpec; + private Integer hashCode = null; private PhysicalProperties() { - this.orderSpec = new OrderSpec(); - this.distributionSpec = DistributionSpecAny.INSTANCE; + this(DistributionSpecAny.INSTANCE, new OrderSpec(), Optional.empty()); } public PhysicalProperties(DistributionSpec distributionSpec) { - this.distributionSpec = distributionSpec; - this.orderSpec = new OrderSpec(); + this(distributionSpec, new OrderSpec(), Optional.empty()); } public PhysicalProperties(OrderSpec orderSpec) { - this.orderSpec = orderSpec; - this.distributionSpec = DistributionSpecAny.INSTANCE; + this(DistributionSpecAny.INSTANCE, orderSpec, Optional.empty()); } public PhysicalProperties(DistributionSpec distributionSpec, OrderSpec orderSpec) { + this(distributionSpec, orderSpec, Optional.empty()); + } + + /** Constructor with mapping-based natural bucket locality. */ + public PhysicalProperties(DistributionSpec distributionSpec, OrderSpec orderSpec, + Optional naturalDistributionMappingSpec) { this.distributionSpec = distributionSpec; this.orderSpec = orderSpec; + this.naturalDistributionMappingSpec = naturalDistributionMappingSpec; } /** @@ -115,12 +122,28 @@ public static PhysicalProperties createAnyFromHash(DistributionSpecHash... child } public PhysicalProperties withOrderSpec(OrderSpec orderSpec) { - return new PhysicalProperties(distributionSpec, orderSpec); + return new PhysicalProperties(distributionSpec, orderSpec, naturalDistributionMappingSpec); + } + + public PhysicalProperties withNaturalDistributionMappingSpec( + Optional naturalDistributionMappingSpec) { + return new PhysicalProperties(distributionSpec, orderSpec, naturalDistributionMappingSpec); } - // Current properties satisfies other properties. + /** Return whether the current properties satisfy the required properties. */ public boolean satisfy(PhysicalProperties other) { - return orderSpec.satisfy(other.orderSpec) && distributionSpec.satisfy(other.distributionSpec); + if (!orderSpec.satisfy(other.orderSpec)) { + return false; + } + if (other.distributionSpec instanceof DistributionSpecHash + && ((DistributionSpecHash) other.distributionSpec).getShuffleType() + == ShuffleType.COLOCATE_MAPPING_REQUIRE) { + return naturalDistributionMappingSpec + .map(spec -> spec.satisfy( + ((DistributionSpecHash) other.distributionSpec).getOrderedShuffledColumns())) + .orElse(false); + } + return distributionSpec.satisfy(other.distributionSpec); } public OrderSpec getOrderSpec() { @@ -131,6 +154,17 @@ public DistributionSpec getDistributionSpec() { return distributionSpec; } + public Optional getNaturalDistributionMappingSpec() { + return naturalDistributionMappingSpec; + } + + /** Whether a missing property can be produced by adding physical enforcers. */ + public boolean isEnforceable() { + return !(distributionSpec instanceof DistributionSpecHash) + || ((DistributionSpecHash) distributionSpec).getShuffleType() + != ShuffleType.COLOCATE_MAPPING_REQUIRE; + } + public boolean isDistributionOnlyProperties() { return orderSpec.getOrderKeys().isEmpty(); } @@ -148,13 +182,14 @@ public boolean equals(Object o) { return false; } return orderSpec.equals(that.orderSpec) - && distributionSpec.equals(that.distributionSpec); + && distributionSpec.equals(that.distributionSpec) + && naturalDistributionMappingSpec.equals(that.naturalDistributionMappingSpec); } @Override public int hashCode() { if (hashCode == null) { - hashCode = Objects.hash(orderSpec, distributionSpec); + hashCode = Objects.hash(orderSpec, distributionSpec, naturalDistributionMappingSpec); } return hashCode; } @@ -170,7 +205,8 @@ public String toString() { if (this.equals(GATHER)) { return "GATHER"; } - return distributionSpec.toString() + " " + orderSpec.toString(); + return distributionSpec.toString() + " " + orderSpec.toString() + + naturalDistributionMappingSpec.map(spec -> " " + spec).orElse(""); } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/RequestPropertyDeriver.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/RequestPropertyDeriver.java index 2a94a6c971bb95..95822e87196114 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/RequestPropertyDeriver.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/RequestPropertyDeriver.java @@ -27,6 +27,7 @@ import org.apache.doris.nereids.rules.implementation.LogicalWindowToPhysicalWindow.WindowFrameGroup; import org.apache.doris.nereids.stats.StatsCalculator; import org.apache.doris.nereids.trees.expressions.Alias; +import org.apache.doris.nereids.trees.expressions.Cast; import org.apache.doris.nereids.trees.expressions.ExprId; import org.apache.doris.nereids.trees.expressions.Expression; import org.apache.doris.nereids.trees.expressions.NamedExpression; @@ -261,6 +262,11 @@ public Void visitPhysicalHashJoin(PhysicalHashJoin agg addRequestPropertyToChildren(PhysicalProperties.ANY); return null; } else if (agg.getAggPhase().isGlobal()) { + addColocateMappingRequestForAggregate(agg); // partition expressions already set by rule if (agg.getPartitionExpressions().isPresent() && !agg.getPartitionExpressions().get().isEmpty()) { addRequestPropertyToChildren( @@ -504,6 +515,57 @@ public Void visitPhysicalHashAggregate(PhysicalHashAggregate agg return null; } + private void addColocateMappingRequestForAggregate(PhysicalHashAggregate agg) { + DistributionSpec parentDistribution = requestPropertyFromParent.getDistributionSpec(); + if (connectContext == null + || !connectContext.getSessionVariable().isEnableColocateMappingConstraint() + || agg.hasSourceRepeat() + || agg.isDistinctOrDeduplicate() + || !(parentDistribution instanceof DistributionSpecHash) + || ((DistributionSpecHash) parentDistribution).getShuffleType() + != ShuffleType.COLOCATE_MAPPING_REQUIRE) { + return; + } + + Map outputByExprId = agg.getOutputExpressions().stream() + .collect(Collectors.toMap(NamedExpression::getExprId, output -> output, (left, right) -> left)); + Set groupByExprIds = Sets.newHashSet(); + for (Expression groupBy : agg.getGroupByExpressions()) { + if (!(groupBy instanceof SlotReference)) { + return; + } + groupByExprIds.add(((SlotReference) groupBy).getExprId()); + } + List childRequiredExprIds = Lists.newArrayList(); + for (ExprId requiredExprId + : ((DistributionSpecHash) parentDistribution).getOrderedShuffledColumns()) { + NamedExpression output = outputByExprId.get(requiredExprId); + ExprId childExprId; + if (output instanceof Alias && ((Alias) output).child() instanceof SlotReference) { + childExprId = ((SlotReference) ((Alias) output).child()).getExprId(); + } else if (output instanceof Alias && ((Alias) output).child() instanceof Cast + && ((Alias) output).child().child(0) instanceof Slot + && ChildOutputPropertyDeriver.isHashValuePreservingCast( + ((Alias) output).child().child(0).getDataType(), + ((Alias) output).child().getDataType())) { + childExprId = ((Slot) ((Alias) output).child().child(0)).getExprId(); + } else if (output instanceof SlotReference) { + childExprId = output.getExprId(); + } else { + continue; + } + if (!groupByExprIds.contains(childExprId)) { + return; + } + childRequiredExprIds.add(childExprId); + } + if (childRequiredExprIds.isEmpty()) { + return; + } + addRequestPropertyToChildren( + PhysicalProperties.createHash(childRequiredExprIds, ShuffleType.COLOCATE_MAPPING_REQUIRE)); + } + private boolean shouldUseParent(List parentHashExprIds, PhysicalHashAggregate agg, PlanContext context) { if (!context.getConnectContext().getSessionVariable().aggShuffleUseParentKey) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/implementation/LogicalOlapScanToPhysicalOlapScan.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/implementation/LogicalOlapScanToPhysicalOlapScan.java index 8448f14831cfa7..e69b4e309216cb 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/implementation/LogicalOlapScanToPhysicalOlapScan.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/implementation/LogicalOlapScanToPhysicalOlapScan.java @@ -19,8 +19,11 @@ import org.apache.doris.catalog.Column; import org.apache.doris.catalog.DistributionInfo; +import org.apache.doris.catalog.Env; import org.apache.doris.catalog.HashDistributionInfo; import org.apache.doris.catalog.OlapTable; +import org.apache.doris.catalog.constraint.DistributionMappingConstraint; +import org.apache.doris.nereids.properties.DistributionMapping; import org.apache.doris.nereids.properties.DistributionSpec; import org.apache.doris.nereids.properties.DistributionSpecHash; import org.apache.doris.nereids.properties.DistributionSpecHash.ShuffleType; @@ -33,11 +36,16 @@ import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan; import org.apache.doris.nereids.trees.plans.physical.PhysicalOlapScan; import org.apache.doris.nereids.util.Utils; +import org.apache.doris.qe.ConnectContext; +import com.google.common.collect.ImmutableList; import com.google.common.collect.Lists; import com.google.common.collect.Sets; +import java.util.HashMap; import java.util.List; +import java.util.Locale; +import java.util.Map; import java.util.Optional; /** @@ -114,8 +122,7 @@ public static DistributionSpec convertDistribution(LogicalOlapScan olapScan) { } } } - return new DistributionSpecHash(hashColumns, ShuffleType.NATURAL, olapScan.getTable().getId(), - olapScan.getSelectedIndexId(), Sets.newLinkedHashSet(olapScan.getSelectedPartitionIds())); + return createNaturalHashSpec(olapScan, hashDistributionInfo, hashColumns, output); } else { HashDistributionInfo hashDistributionInfo = (HashDistributionInfo) distributionInfo; List output = olapScan.getOutput(); @@ -132,12 +139,94 @@ public static DistributionSpec convertDistribution(LogicalOlapScan olapScan) { } } } - return new DistributionSpecHash(hashColumns, ShuffleType.NATURAL, olapScan.getTable().getId(), - olapScan.getSelectedIndexId(), Sets.newLinkedHashSet(olapScan.getSelectedPartitionIds())); + return createNaturalHashSpec(olapScan, hashDistributionInfo, hashColumns, output); } } else { - // RandomDistributionInfo + if (!(distributionInfo instanceof HashDistributionInfo)) { + validateDistributionMappingsForPlanning(olapTable); + } return DistributionSpecStorageAny.INSTANCE; } } + + private static void validateDistributionMappingsForPlanning(OlapTable table) { + ConnectContext context = ConnectContext.get(); + if (context != null && context.getSessionVariable().isEnableColocateMappingConstraint()) { + Env.getCurrentEnv().getConstraintManager().getDistributionMappingConstraintsForPlanning(table); + } + } + + private static DistributionSpecHash createNaturalHashSpec(LogicalOlapScan olapScan, + HashDistributionInfo hashDistributionInfo, List hashColumns, List output) { + return new DistributionSpecHash(hashColumns, ShuffleType.NATURAL, olapScan.getTable().getId(), + olapScan.getSelectedIndexId(), Sets.newLinkedHashSet(olapScan.getSelectedPartitionIds()), + buildDistributionMappings(olapScan.getTable(), hashDistributionInfo, output)); + } + + static List buildDistributionMappings(OlapTable table, + HashDistributionInfo hashDistributionInfo, List output) { + ConnectContext context = ConnectContext.get(); + if (context == null || !context.getSessionVariable().isEnableColocateMappingConstraint()) { + return ImmutableList.of(); + } + List mappings = buildDistributionMappings( + hashDistributionInfo, + output, + Env.getCurrentEnv().getConstraintManager() + .getDistributionMappingConstraintsForPlanning(table)); + if (!mappings.isEmpty()) { + context.getStatementContext().getSqlCacheContext() + .ifPresent(sqlCacheContext -> sqlCacheContext.setHasUnsupportedTables(true)); + } + return mappings; + } + + static List buildDistributionMappings(HashDistributionInfo hashDistributionInfo, + List output, List constraints) { + Map columnExprIds = new HashMap<>(); + for (Slot slot : output) { + SlotReference slotReference = (SlotReference) slot; + slotReference.getOriginalColumn().ifPresent( + column -> columnExprIds.put( + column.tryGetBaseColumnName().toLowerCase(Locale.ROOT), slot.getExprId())); + } + Map distributionIndices = new HashMap<>(); + List distributionColumns = hashDistributionInfo.getDistributionColumns(); + for (int i = 0; i < distributionColumns.size(); i++) { + distributionIndices.put(distributionColumns.get(i).getName().toLowerCase(Locale.ROOT), i); + } + + ImmutableList.Builder mappings = ImmutableList.builder(); + for (DistributionMappingConstraint constraint : constraints) { + ImmutableList.Builder determinants = ImmutableList.builder(); + boolean allDeterminantsAvailable = true; + for (String column : constraint.getDeterminantColumnNames()) { + ExprId exprId = columnExprIds.get(column.toLowerCase(Locale.ROOT)); + if (exprId == null) { + allDeterminantsAvailable = false; + break; + } + determinants.add(exprId); + } + if (!allDeterminantsAvailable) { + continue; + } + ImmutableList.Builder targetIndices = ImmutableList.builder(); + boolean allTargetsAvailable = true; + for (String column : constraint.getDistributionColumnNames()) { + Integer targetIndex = distributionIndices.get(column.toLowerCase(Locale.ROOT)); + if (targetIndex == null) { + allTargetsAvailable = false; + break; + } + targetIndices.add(targetIndex); + } + if (!allTargetsAvailable) { + continue; + } + mappings.add(new DistributionMapping( + constraint.getMappingId(), determinants.build(), targetIndices.build())); + } + return mappings.build(); + } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/AddConstraintCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/AddConstraintCommand.java index 5eeb3f80f4edc8..b988b8b1677c82 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/AddConstraintCommand.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/AddConstraintCommand.java @@ -19,7 +19,9 @@ import org.apache.doris.catalog.Env; import org.apache.doris.catalog.MTMV; +import org.apache.doris.catalog.OlapTable; import org.apache.doris.catalog.TableIf; +import org.apache.doris.catalog.constraint.DistributionMappingConstraint; import org.apache.doris.catalog.constraint.ForeignKeyConstraint; import org.apache.doris.catalog.constraint.PrimaryKeyConstraint; import org.apache.doris.catalog.constraint.UniqueConstraint; @@ -40,6 +42,7 @@ 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.persist.EditLog; import org.apache.doris.qe.ConnectContext; import org.apache.doris.qe.StmtExecutor; @@ -98,6 +101,15 @@ public void run(ConnectContext ctx, StmtExecutor executor) throws Exception { } else if (constraint.isUnique()) { addConstraintAndInvalidate( tableNameInfo, new UniqueConstraint(name, ImmutableSet.copyOf(columns))); + } else if (constraint.isDistributionMapping()) { + Pair, TableIf> distributionColumnsAndTable = + extractColumnsAndTable(ctx, constraint.toDistributionProject()); + if (table != distributionColumnsAndTable.second) { + throw new AnalysisException("Table changed while adding constraint on " + tableNameInfo); + } + addDistributionMapping(tableNameInfo, table, + new DistributionMappingConstraint( + name, constraint.getMappingId(), columns, distributionColumnsAndTable.first)); } else { throw new AnalysisException("Unsupported constraint type: " + constraint); } @@ -122,6 +134,34 @@ private void addConstraintAndInvalidate( String.format("after add constraint %s on table %s", constraint.getName(), tableNameInfo)); } + private void addDistributionMapping(TableNameInfo tableNameInfo, + TableIf analyzedTable, DistributionMappingConstraint constraint) throws Exception { + if (!(analyzedTable instanceof OlapTable)) { + throw new AnalysisException("Distribution mapping constraint only supports OLAP tables"); + } + EditLog.EditLogItem logItem; + analyzedTable.getDatabase().readLock(); + try { + if (analyzedTable.getDatabase().getCatalog().getDbNullable(tableNameInfo.getDb()) + != analyzedTable.getDatabase() + || analyzedTable.getDatabase().getTableNullable(tableNameInfo.getTbl()) != analyzedTable) { + throw new AnalysisException("Table changed while adding constraint on " + tableNameInfo); + } + analyzedTable.writeLock(); + try { + OlapTable table = (OlapTable) analyzedTable; + table.checkNormalStateForAlter(); + logItem = Env.getCurrentEnv().getConstraintManager() + .addDistributionMappingConstraint(tableNameInfo, table, constraint); + } finally { + analyzedTable.writeUnlock(); + } + } finally { + analyzedTable.getDatabase().readUnlock(); + } + logItem.await(); + } + private Pair, TableIf> extractColumnsAndTable(ConnectContext ctx, LogicalPlan plan) { NereidsPlanner planner = new NereidsPlanner(ctx.getStatementContext()); Plan analyzedPlan = planner.planWithLock( diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/Constraint.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/Constraint.java index f6fe938ec779e9..70dd9a7a043fa4 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/Constraint.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/Constraint.java @@ -36,6 +36,8 @@ public class Constraint { private final LogicalPlan curTable; private final @Nullable LogicalPlan referenceTable; private final @Nullable ImmutableList referenceSlots; + private final @Nullable String mappingId; + private final @Nullable ImmutableList distributionSlots; private final ConstraintType type; Constraint(ConstraintType type, LogicalPlan curTable, ImmutableList slots) { @@ -47,6 +49,8 @@ public class Constraint { "table of constraint can't be null"); this.referenceTable = null; this.referenceSlots = null; + this.mappingId = null; + this.distributionSlots = null; } Constraint(LogicalPlan curTable, ImmutableList slots, @@ -61,10 +65,27 @@ public class Constraint { "reference table in foreign key can not be null"); this.referenceSlots = Objects.requireNonNull(referenceSlotSet, "reference slots in foreign key can not be null"); + this.mappingId = null; + this.distributionSlots = null; Preconditions.checkArgument(referenceSlots.size() == slots.size(), "Foreign key's size must be same as the size of reference slots"); } + Constraint(LogicalPlan curTable, String mappingId, ImmutableList determinantSlots, + ImmutableList distributionSlots) { + Preconditions.checkArgument(determinantSlots != null && !determinantSlots.isEmpty(), + "determinant slots of distribution mapping constraint can't be null or empty"); + Preconditions.checkArgument(distributionSlots != null && !distributionSlots.isEmpty(), + "distribution slots of distribution mapping constraint can't be null or empty"); + this.type = ConstraintType.DISTRIBUTION_MAPPING; + this.slots = determinantSlots; + this.curTable = Objects.requireNonNull(curTable, "table of constraint can't be null"); + this.referenceTable = null; + this.referenceSlots = null; + this.mappingId = Objects.requireNonNull(mappingId, "mapping id can't be null"); + this.distributionSlots = distributionSlots; + } + public static Constraint newUniqueConstraint(LogicalPlan curTable, ImmutableList slotSet) { return new Constraint(ConstraintType.UNIQUE, curTable, slotSet); } @@ -79,6 +100,11 @@ public static Constraint newForeignKeyConstraint( return new Constraint(curTable, slotSet, referenceTable, referenceSlotSet); } + public static Constraint newDistributionMappingConstraint(LogicalPlan curTable, String mappingId, + ImmutableList determinantSlots, ImmutableList distributionSlots) { + return new Constraint(curTable, mappingId, determinantSlots, distributionSlots); + } + public boolean isForeignKey() { return type == ConstraintType.FOREIGN_KEY; } @@ -91,6 +117,14 @@ public boolean isPrimaryKey() { return type == ConstraintType.PRIMARY_KEY; } + public boolean isDistributionMapping() { + return type == ConstraintType.DISTRIBUTION_MAPPING; + } + + public String getMappingId() { + return Objects.requireNonNull(mappingId, "mapping id is only available for distribution mapping constraint"); + } + public LogicalPlan toProject() { return new LogicalProject<>(ImmutableList.copyOf(slots), curTable); } @@ -100,6 +134,12 @@ public LogicalPlan toReferenceProject() { return new LogicalProject<>(ImmutableList.copyOf(referenceSlots), referenceTable); } + public LogicalPlan toDistributionProject() { + Preconditions.checkArgument(distributionSlots != null, + "distribution slots are only available for distribution mapping constraint"); + return new LogicalProject<>(ImmutableList.copyOf(distributionSlots), curTable); + } + @Override public String toString() { StringBuilder sb = new StringBuilder(); @@ -109,6 +149,9 @@ public String toString() { if (type.equals(ConstraintType.FOREIGN_KEY)) { sb.append("Reference Table: ").append(referenceTable).append("\n"); sb.append("Reference Slot Set: ").append(referenceSlots); + } else if (type.equals(ConstraintType.DISTRIBUTION_MAPPING)) { + sb.append("Mapping Id: ").append(mappingId).append("\n"); + sb.append("Distribution Slot Set: ").append(distributionSlots); } return sb.toString(); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/DropConstraintCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/DropConstraintCommand.java index 15419eb20be5c5..fcd56397c64a73 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/DropConstraintCommand.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/DropConstraintCommand.java @@ -19,8 +19,10 @@ import org.apache.doris.catalog.Env; import org.apache.doris.catalog.MTMV; +import org.apache.doris.catalog.OlapTable; import org.apache.doris.catalog.TableIf; import org.apache.doris.catalog.constraint.Constraint; +import org.apache.doris.catalog.constraint.DistributionMappingConstraint; import org.apache.doris.catalog.info.TableNameInfo; import org.apache.doris.common.ErrorCode; import org.apache.doris.common.ErrorReport; @@ -37,6 +39,7 @@ 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.persist.EditLog; import org.apache.doris.qe.ConnectContext; import org.apache.doris.qe.StmtExecutor; @@ -67,8 +70,9 @@ public DropConstraintCommand(String name, LogicalPlan plan) { @Override public void run(ConnectContext ctx, StmtExecutor executor) throws Exception { TableNameInfo tableNameInfo; + TableIf table = null; try { - TableIf table = extractTable(ctx, plan); + table = extractTable(ctx, plan); tableNameInfo = TableNameInfoUtils.fromCatalogDb( table.getDatabase().getCatalog(), table.getDatabase(), table); } catch (Exception e) { @@ -81,6 +85,14 @@ public void run(ConnectContext ctx, StmtExecutor executor) throws Exception { // must be checked on both paths above: table resolution failing (which includes an // authorization failure) falls back to a name-only lookup that binds nothing. checkAlterPriv(ctx, tableNameInfo); + if (table != null) { + Constraint mapping = Env.getCurrentEnv().getConstraintManager() + .getConstraint(tableNameInfo, table, name); + if (mapping instanceof DistributionMappingConstraint) { + dropDistributionMapping(tableNameInfo, table); + return; + } + } Constraint constraint = Env.getCurrentEnv().getConstraintManager().getConstraint(tableNameInfo, name); if (constraint == null) { throw new AnalysisException( @@ -101,6 +113,32 @@ public void run(ConnectContext ctx, StmtExecutor executor) throws Exception { String.format("after drop constraint %s on table %s", constraint.getName(), tableNameInfo)); } + private void dropDistributionMapping(TableNameInfo tableNameInfo, + TableIf analyzedTable) throws Exception { + EditLog.EditLogItem logItem; + analyzedTable.getDatabase().readLock(); + try { + if (analyzedTable.getDatabase().getCatalog().getDbNullable(tableNameInfo.getDb()) + != analyzedTable.getDatabase() + || analyzedTable.getDatabase().getTableNullable(tableNameInfo.getTbl()) != analyzedTable) { + throw new AnalysisException("Table changed while dropping constraint on " + tableNameInfo); + } + analyzedTable.writeLock(); + try { + OlapTable table = (OlapTable) analyzedTable; + table.checkNormalStateForAlter(); + logItem = Env.getCurrentEnv().getConstraintManager() + .dropDistributionMappingConstraint( + tableNameInfo, table, name); + } finally { + analyzedTable.writeUnlock(); + } + } finally { + analyzedTable.getDatabase().readUnlock(); + } + logItem.await(); + } + private void checkAlterPriv(ConnectContext ctx, TableNameInfo tableNameInfo) throws org.apache.doris.common.AnalysisException { if (!Env.getCurrentEnv().getAccessManager().checkTblPriv(ctx, tableNameInfo.getCtl(), diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ShowConstraintsCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ShowConstraintsCommand.java index f19694f6ab1c1a..df254b4a0e69b0 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ShowConstraintsCommand.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ShowConstraintsCommand.java @@ -73,7 +73,7 @@ public ShowResultSet doRun(ConnectContext ctx, StmtExecutor executor) throws Exc TableNameInfo tableNameInfo = TableNameInfoUtils.fromCatalogDb( tableIf.getDatabase().getCatalog(), tableIf.getDatabase(), tableIf); Map constraints = Env.getCurrentEnv().getConstraintManager() - .getConstraints(tableNameInfo); + .getConstraints(tableNameInfo, tableIf); List> res = constraints.entrySet().stream() .map(e -> Lists.newArrayList(e.getKey(), e.getValue().getType().getName(), diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalHashAggregate.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalHashAggregate.java index 395ff52119880c..b6bf2d9e1bd1f0 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalHashAggregate.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalHashAggregate.java @@ -30,6 +30,7 @@ import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction; import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateParam; import org.apache.doris.nereids.trees.expressions.functions.agg.Count; +import org.apache.doris.nereids.trees.expressions.functions.agg.MultiDistinction; import org.apache.doris.nereids.trees.expressions.functions.agg.Ndv; import org.apache.doris.nereids.trees.expressions.functions.agg.NullableAggregateFunction; import org.apache.doris.nereids.trees.plans.AbstractPlan; @@ -191,6 +192,26 @@ public boolean hasSourceRepeat() { return hasSourceRepeat; } + /** Whether this node is part of DISTINCT processing or only removes duplicate grouping keys. */ + public boolean isDistinctOrDeduplicate() { + if (getAggPhase() == AggPhase.DISTINCT_LOCAL || getAggPhase() == AggPhase.DISTINCT_GLOBAL) { + return true; + } + boolean hasAggregateExpression = false; + for (NamedExpression outputExpression : outputExpressions) { + List aggregateExpressions = outputExpression.collectToList( + AggregateExpression.class::isInstance); + hasAggregateExpression |= !aggregateExpressions.isEmpty(); + for (AggregateExpression aggregateExpression : aggregateExpressions) { + AggregateFunction function = aggregateExpression.getFunction(); + if (function.isDistinct() || function instanceof MultiDistinction) { + return true; + } + } + } + return !hasAggregateExpression; + } + @Override public String toString() { TopnPushInfo topnPushInfo = (TopnPushInfo) getMutableState( diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/util/JoinUtils.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/util/JoinUtils.java index 20292df56ca7a5..5f3c28557e62c3 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/util/JoinUtils.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/util/JoinUtils.java @@ -21,10 +21,12 @@ import org.apache.doris.catalog.Env; import org.apache.doris.common.Pair; import org.apache.doris.nereids.properties.DataTrait; +import org.apache.doris.nereids.properties.DistributionMapping; import org.apache.doris.nereids.properties.DistributionSpec; import org.apache.doris.nereids.properties.DistributionSpecHash; import org.apache.doris.nereids.properties.DistributionSpecHash.ShuffleType; import org.apache.doris.nereids.properties.DistributionSpecReplicated; +import org.apache.doris.nereids.properties.NaturalDistributionMappingSpec; import org.apache.doris.nereids.properties.PhysicalProperties; import org.apache.doris.nereids.rules.rewrite.AdjustNullable; import org.apache.doris.nereids.rules.rewrite.ForeignKeyContext; @@ -49,6 +51,7 @@ import com.google.common.collect.ImmutableList.Builder; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Lists; import com.google.common.collect.Sets; import java.util.HashMap; @@ -239,8 +242,18 @@ public static boolean shouldColocateJoin(AbstractPhysicalJoin leftTablePartitions = leftHashSpec.getPartitionIds(); final Set rightTablePartitions = rightHashSpec.getPartitionIds(); - // For UT or no partition is selected, getSelectedIndexId() == -1, see selectMaterializedView() - boolean hitSameIndex = (leftTableId == rightTableId) - && (leftHashSpec.getSelectedIndexId() != -1 && rightHashSpec.getSelectedIndexId() != -1) - && (leftHashSpec.getSelectedIndexId() == rightHashSpec.getSelectedIndexId()); - boolean noNeedCheckColocateGroup = hitSameIndex && (leftTablePartitions.equals(rightTablePartitions)) - && (leftTablePartitions.size() <= 1); - ColocateTableIndex colocateIndex = Env.getCurrentColocateIndex(); - if (!noNeedCheckColocateGroup && (!colocateIndex.isSameGroup(leftTableId, rightTableId) - || colocateIndex.isGroupUnstable(colocateIndex.getGroup(leftTableId)))) { + if (!isSameStableColocateGroup(leftTableId, leftHashSpec.getSelectedIndexId(), leftTablePartitions, + rightTableId, rightHashSpec.getSelectedIndexId(), rightTablePartitions)) { + return false; + } + + if (couldColocateJoinOnDistributionColumns(leftHashSpec, rightHashSpec, conjuncts)) { + return true; + } + if (!ConnectContext.get().getSessionVariable().isEnableColocateMappingConstraint()) { + return false; + } + return couldColocateJoinOnDistributionMappings(leftHashSpec, rightHashSpec, conjuncts); + } + + /** Could use a mapping-based colocate join after distribution-key slots have been projected out. */ + public static boolean couldColocateJoinByMapping(NaturalDistributionMappingSpec leftSpec, + NaturalDistributionMappingSpec rightSpec, List conjuncts) { + if (ConnectContext.get() == null + || ConnectContext.get().getSessionVariable().isDisableColocatePlan() + || !ConnectContext.get().getSessionVariable().isEnableColocateMappingConstraint() + || leftSpec.getDistributionKeyCount() != rightSpec.getDistributionKeyCount()) { + return false; + } + if (!isSameStableColocateGroup( + leftSpec.getTableId(), leftSpec.getSelectedIndexId(), leftSpec.getPartitionIds(), + rightSpec.getTableId(), rightSpec.getSelectedIndexId(), rightSpec.getPartitionIds())) { return false; } + return couldColocateJoinOnDistributionMappings( + leftSpec.getVisibleDistributionExprToIndex(), leftSpec.getDistributionMappings(), + rightSpec.getVisibleDistributionExprToIndex(), rightSpec.getDistributionMappings(), + leftSpec.getDistributionKeyCount(), conjuncts); + } + + private static boolean isSameStableColocateGroup(long leftTableId, long leftSelectedIndexId, + Set leftTablePartitions, long rightTableId, long rightSelectedIndexId, + Set rightTablePartitions) { + // For UT or no partition is selected, selectedIndexId == -1, see selectMaterializedView(). + boolean hitSameIndex = leftTableId == rightTableId + && leftSelectedIndexId != -1 + && rightSelectedIndexId != -1 + && leftSelectedIndexId == rightSelectedIndexId; + boolean noNeedCheckColocateGroup = hitSameIndex + && leftTablePartitions.equals(rightTablePartitions) + && leftTablePartitions.size() <= 1; + ColocateTableIndex colocateIndex = Env.getCurrentColocateIndex(); + return noNeedCheckColocateGroup + || (colocateIndex.isSameGroup(leftTableId, rightTableId) + && !colocateIndex.isGroupUnstable(colocateIndex.getGroup(leftTableId))); + } + private static boolean couldColocateJoinOnDistributionColumns(DistributionSpecHash leftHashSpec, + DistributionSpecHash rightHashSpec, List conjuncts) { + if (!areAllSlotEqualPredicates(conjuncts)) { + return false; + } Set equalIndices = new HashSet<>(); for (Expression expr : conjuncts) { - // only simple equal predicate can use colocate join - if (!(expr instanceof EqualPredicate)) { - return false; - } Expression leftChild = ((EqualPredicate) expr).left(); Expression rightChild = ((EqualPredicate) expr).right(); - if (!(leftChild instanceof SlotReference) || !(rightChild instanceof SlotReference)) { - return false; - } - SlotReference leftSlot = (SlotReference) leftChild; SlotReference rightSlot = (SlotReference) rightChild; Integer leftIndex = leftHashSpec.getExprIdToEquivalenceSet().get(leftSlot.getExprId()); @@ -307,12 +356,120 @@ public static boolean couldColocateJoin(DistributionSpecHash leftHashSpec, Distr equalIndices.add(leftIndex); } } - // on conditions must contain all distributed columns - if (equalIndices.containsAll(leftHashSpec.getExprIdToEquivalenceSet().values())) { - return true; - } else { + return equalIndices.containsAll(leftHashSpec.getExprIdToEquivalenceSet().values()); + } + + private static boolean couldColocateJoinOnDistributionMappings(DistributionSpecHash leftHashSpec, + DistributionSpecHash rightHashSpec, List conjuncts) { + return couldColocateJoinOnDistributionMappings( + leftHashSpec.getExprIdToEquivalenceSet(), leftHashSpec.getDistributionMappings(), + rightHashSpec.getExprIdToEquivalenceSet(), rightHashSpec.getDistributionMappings(), + leftHashSpec.getOrderedShuffledColumns().size(), conjuncts); + } + + private static boolean couldColocateJoinOnDistributionMappings( + Map leftDistributionExprToIndex, List leftMappings, + Map rightDistributionExprToIndex, List rightMappings, + int distributionKeyCount, List conjuncts) { + if (!areAllSlotEqualPredicates(conjuncts)) { return false; } + List> equalExprIds = Lists.newArrayList(); + Set coveredIndices = new HashSet<>(); + for (Expression expr : conjuncts) { + ExprId first = ((SlotReference) ((EqualPredicate) expr).left()).getExprId(); + ExprId second = ((SlotReference) ((EqualPredicate) expr).right()).getExprId(); + equalExprIds.add(Pair.of(first, second)); + + Integer leftIndex = leftDistributionExprToIndex.get(first); + Integer rightIndex = rightDistributionExprToIndex.get(second); + if (leftIndex == null) { + leftIndex = leftDistributionExprToIndex.get(second); + rightIndex = rightDistributionExprToIndex.get(first); + } + if (leftIndex != null && Objects.equals(leftIndex, rightIndex)) { + coveredIndices.add(leftIndex); + } + } + + Map> rightMappingsByKey = + new HashMap<>(); + for (DistributionMapping rightMapping : rightMappings) { + rightMappingsByKey.computeIfAbsent( + new DistributionMappingKey(rightMapping), ignored -> Lists.newArrayList()) + .add(rightMapping); + } + for (DistributionMapping leftMapping : leftMappings) { + List compatibleRightMappings = rightMappingsByKey.get( + new DistributionMappingKey(leftMapping)); + if (compatibleRightMappings == null) { + continue; + } + for (DistributionMapping rightMapping : compatibleRightMappings) { + boolean determinantsEqual = true; + for (int i = 0; i < leftMapping.getDeterminantExprIds().size(); i++) { + if (!containsEqualPair(equalExprIds, leftMapping.getDeterminantExprIds().get(i), + rightMapping.getDeterminantExprIds().get(i))) { + determinantsEqual = false; + break; + } + } + if (determinantsEqual) { + coveredIndices.addAll(leftMapping.getTargetDistributionIndices()); + } + } + } + for (int i = 0; i < distributionKeyCount; i++) { + if (!coveredIndices.contains(i)) { + return false; + } + } + return true; + } + + /** Whether every conjunct is a simple slot-to-slot equality supported by colocate proof. */ + public static boolean areAllSlotEqualPredicates(List conjuncts) { + return conjuncts.stream().allMatch(expr -> + expr instanceof EqualPredicate + && ((EqualPredicate) expr).left() instanceof SlotReference + && ((EqualPredicate) expr).right() instanceof SlotReference); + } + + private static boolean containsEqualPair( + List> equalExprIds, ExprId left, ExprId right) { + return equalExprIds.stream().anyMatch(pair -> + (pair.first.equals(left) && pair.second.equals(right)) + || (pair.first.equals(right) && pair.second.equals(left))); + } + + private static final class DistributionMappingKey { + private final String mappingId; + private final List targetDistributionIndices; + private final int determinantCount; + + private DistributionMappingKey(DistributionMapping mapping) { + mappingId = mapping.getMappingId(); + targetDistributionIndices = mapping.getTargetDistributionIndices(); + determinantCount = mapping.getDeterminantExprIds().size(); + } + + @Override + public boolean equals(Object obj) { + if (!(obj instanceof DistributionMappingKey)) { + return false; + } + DistributionMappingKey other = (DistributionMappingKey) obj; + return determinantCount == other.determinantCount + && mappingId.equals(other.mappingId) + && targetDistributionIndices.equals( + other.targetDistributionIndices); + } + + @Override + public int hashCode() { + return Objects.hash( + mappingId, targetDistributionIndices, determinantCount); + } } public static Set getJoinOutputExprIdSet(Plan left, Plan right) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/persist/EditLog.java b/fe/fe-core/src/main/java/org/apache/doris/persist/EditLog.java index 1c7c9738415d4b..c4568c01053af5 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/persist/EditLog.java +++ b/fe/fe-core/src/main/java/org/apache/doris/persist/EditLog.java @@ -1004,7 +1004,10 @@ public static void loadJournal(Env env, Long logId, JournalEntity journal) { case OperationType.OP_MODIFY_REPLICATION_NUM: { ModifyTablePropertyOperationLog log = (ModifyTablePropertyOperationLog) journal.getData(); env.replayModifyTableProperty(opCode, log); - env.getBinlogManager().addModifyTableProperty(log, logId); + // Mapping constraints only reuse this legacy envelope for old-FE readability. + if (!log.hasDistributionMappingConstraintMutation()) { + env.getBinlogManager().addModifyTableProperty(log, logId); + } break; } case OperationType.OP_TABLE_STREAM_CLEANUP: { diff --git a/fe/fe-core/src/main/java/org/apache/doris/persist/ModifyTablePropertyOperationLog.java b/fe/fe-core/src/main/java/org/apache/doris/persist/ModifyTablePropertyOperationLog.java index 1d64029bcfdb00..c92b93065b4f42 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/persist/ModifyTablePropertyOperationLog.java +++ b/fe/fe-core/src/main/java/org/apache/doris/persist/ModifyTablePropertyOperationLog.java @@ -17,6 +17,7 @@ package org.apache.doris.persist; +import org.apache.doris.catalog.TableProperty; import org.apache.doris.common.io.Text; import org.apache.doris.common.io.Writable; import org.apache.doris.persist.gson.GsonUtils; @@ -95,6 +96,10 @@ public Map getProperties() { return properties; } + public boolean hasDistributionMappingConstraintMutation() { + return properties.containsKey(TableProperty.DISTRIBUTION_MAPPING_CONSTRAINTS_PROPERTY); + } + @Override public void write(DataOutput out) throws IOException { Text.writeString(out, GsonUtils.GSON.toJson(this)); diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java b/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java index c5b503e9f144b0..bf7f68e927d924 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java @@ -182,6 +182,7 @@ public class SessionVariable implements Serializable, Writable { public static final String ENABLE_LOCAL_EXCHANGE_BEFORE_STREAMING_AGG = "enable_local_exchange_before_streaming_agg"; public static final String DISABLE_COLOCATE_PLAN = "disable_colocate_plan"; + public static final String ENABLE_COLOCATE_MAPPING_CONSTRAINT = "enable_colocate_mapping_constraint"; public static final String COLOCATE_MAX_PARALLEL_NUM = "colocate_max_parallel_num"; public static final String ENABLE_BUCKET_SHUFFLE_JOIN = "enable_bucket_shuffle_join"; public static final String PARALLEL_FRAGMENT_EXEC_INSTANCE_NUM = "parallel_fragment_exec_instance_num"; @@ -1434,6 +1435,12 @@ public void checkQuerySlotCount(String slotCnt) { @VarAttrDef.VarAttr(name = DISABLE_COLOCATE_PLAN) public boolean disableColocatePlan = false; + @VarAttrDef.VarAttr(name = ENABLE_COLOCATE_MAPPING_CONSTRAINT, + varType = VariableAnnotation.EXPERIMENTAL_ONLINE, + affectQueryResultInPlan = true, + description = "Whether to derive colocate joins from distribution mapping constraints") + public boolean enableColocateMappingConstraint = false; + @VarAttrDef.VarAttr(name = ENABLE_BUCKET_SHUFFLE_JOIN, varType = VariableAnnotation.EXPERIMENTAL_ONLINE) public boolean enableBucketShuffleJoin = true; @@ -4364,6 +4371,10 @@ public boolean isDisableColocatePlan() { return disableColocatePlan; } + public boolean isEnableColocateMappingConstraint() { + return enableColocateMappingConstraint; + } + public boolean isEnableBucketShuffleJoin() { return enableBucketShuffleJoin; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/system/Frontend.java b/fe/fe-core/src/main/java/org/apache/doris/system/Frontend.java index 62970a9315cc6a..28b5a3865127a3 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/system/Frontend.java +++ b/fe/fe-core/src/main/java/org/apache/doris/system/Frontend.java @@ -49,7 +49,7 @@ public class Frontend implements Writable { @SerializedName("cloudUniqueId") private String cloudUniqueId; - private String version; + private volatile String version; private transient String localResourceGroup = ""; private int queryPort; diff --git a/fe/fe-core/src/test/java/org/apache/doris/alter/SchemaChangeJobV2Test.java b/fe/fe-core/src/test/java/org/apache/doris/alter/SchemaChangeJobV2Test.java index f184c73d62bfeb..7113d168566c63 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/alter/SchemaChangeJobV2Test.java +++ b/fe/fe-core/src/test/java/org/apache/doris/alter/SchemaChangeJobV2Test.java @@ -42,8 +42,10 @@ import org.apache.doris.catalog.PrimitiveType; import org.apache.doris.catalog.Replica; import org.apache.doris.catalog.ScalarType; +import org.apache.doris.catalog.TableProperty; import org.apache.doris.catalog.Tablet; import org.apache.doris.catalog.Type; +import org.apache.doris.catalog.constraint.DistributionMappingConstraint; import org.apache.doris.catalog.info.ColumnPosition; import org.apache.doris.catalog.info.IndexType; import org.apache.doris.common.DdlException; @@ -60,6 +62,7 @@ import org.apache.doris.nereids.trees.plans.commands.info.DropColumnOp; import org.apache.doris.nereids.trees.plans.commands.info.ModifyTablePropertiesOp; import org.apache.doris.nereids.types.DataType; +import org.apache.doris.persist.TableInfo; import org.apache.doris.qe.ConnectContext; import org.apache.doris.task.AgentBatchTask; import org.apache.doris.task.AgentTask; @@ -93,6 +96,7 @@ import java.lang.reflect.InvocationTargetException; import java.util.ArrayList; import java.util.HashMap; +import java.util.LinkedList; import java.util.List; import java.util.Map; import java.util.Optional; @@ -558,6 +562,85 @@ public void testModifyTableDistributionType() throws DdlException { Assert.assertTrue(partition1.getDistributionInfo().getType() == DistributionInfo.DistributionInfoType.RANDOM); } + @Test + public void testModifyTableDistributionTypeRejectsDistributionMapping() throws Exception { + if (fakeEnv != null) { + fakeEnv.close(); + } + fakeEnv = new FakeEnv(); + if (fakeEditLog != null) { + fakeEditLog.close(); + } + fakeEditLog = new FakeEditLog(); + FakeEnv.setEnv(masterEnv); + Database db = masterEnv.getInternalCatalog().getDb(CatalogTestUtil.testDbId1).get(); + OlapTable olapTable = (OlapTable) db.getTable(CatalogTestUtil.testTableId1).get(); + DistributionMappingConstraint mapping = new DistributionMappingConstraint( + "distribution_mapping", "distribution_mapping", + List.of("v"), List.of("k1")); + TableProperty tableProperty = new TableProperty(Maps.newHashMap()); + tableProperty.addDistributionMappingConstraint(mapping); + olapTable.setTableProperty(tableProperty); + + DdlException exception = Assert.assertThrows( + DdlException.class, + () -> Env.getCurrentEnv().convertDistributionType(db, olapTable)); + Assert.assertTrue(exception.getMessage().contains("Drop the constraints first")); + Assert.assertEquals( + DistributionInfo.DistributionInfoType.HASH, + olapTable.getDefaultDistributionInfo().getType()); + Assert.assertEquals( + List.of(mapping), + masterEnv.getConstraintManager().getDistributionMappingConstraints(olapTable)); + + TableInfo tableInfo = TableInfo.createForModifyDistribution( + db.getId(), olapTable.getId()); + Env.getCurrentEnv().replayConvertDistributionType(tableInfo); + Assert.assertEquals( + DistributionInfo.DistributionInfoType.RANDOM, + olapTable.getDefaultDistributionInfo().getType()); + Assert.assertEquals( + List.of(mapping), + masterEnv.getConstraintManager().getDistributionMappingConstraints(olapTable)); + } + + @Test + public void testDropMappingColumnFromRollupIsAllowed() throws Exception { + if (fakeEnv != null) { + fakeEnv.close(); + } + fakeEnv = new FakeEnv(); + FakeEnv.setEnv(masterEnv); + SchemaChangeHandler schemaChangeHandler = Env.getCurrentEnv().getSchemaChangeHandler(); + Database db = masterEnv.getInternalCatalog().getDb(CatalogTestUtil.testDbId1).get(); + OlapTable olapTable = (OlapTable) db.getTable(CatalogTestUtil.testTableId1).get(); + DistributionMappingConstraint mapping = new DistributionMappingConstraint( + "distribution_mapping", "distribution_mapping", List.of("v"), List.of("k1")); + TableProperty tableProperty = new TableProperty(Maps.newHashMap()); + tableProperty.addDistributionMappingConstraint(mapping); + olapTable.setTableProperty(tableProperty); + + long rollupIndexId = 1000L; + String rollupIndexName = "mapping_rollup"; + LinkedList baseSchema = new LinkedList<>( + olapTable.getSchemaByIndexId(olapTable.getBaseIndexId())); + LinkedList rollupSchema = new LinkedList<>(baseSchema); + olapTable.setIndexMeta(rollupIndexId, rollupIndexName, new ArrayList<>(rollupSchema), + 0, CatalogTestUtil.testSchemaHash1, (short) 1, TStorageType.COLUMN, KeysType.AGG_KEYS); + Map> indexSchemaMap = new HashMap<>(); + indexSchemaMap.put(olapTable.getBaseIndexId(), baseSchema); + indexSchemaMap.put(rollupIndexId, rollupSchema); + + Deencapsulation.invoke(schemaChangeHandler, "processDropColumn", + new DropColumnOp("v", rollupIndexName, Maps.newHashMap()), + olapTable, indexSchemaMap, new ArrayList()); + + Assert.assertFalse(rollupSchema.stream().anyMatch(column -> column.getName().equalsIgnoreCase("v"))); + Assert.assertEquals( + List.of(mapping), + masterEnv.getConstraintManager().getDistributionMappingConstraints(olapTable)); + } + @Test public void testAbnormalModifyTableDistributionType1() throws UserException { OlapTable table = Mockito.mock(OlapTable.class); diff --git a/fe/fe-core/src/test/java/org/apache/doris/backup/BackupJobTest.java b/fe/fe-core/src/test/java/org/apache/doris/backup/BackupJobTest.java index e1e79f5bbe0dba..16299c379a5fe5 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/backup/BackupJobTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/backup/BackupJobTest.java @@ -26,6 +26,7 @@ import org.apache.doris.catalog.MaterializedIndex.IndexExtState; import org.apache.doris.catalog.OlapTable; import org.apache.doris.catalog.TableProperty; +import org.apache.doris.catalog.constraint.DistributionMappingConstraint; import org.apache.doris.catalog.info.TableNameInfo; import org.apache.doris.common.AnalysisException; import org.apache.doris.common.Config; @@ -368,6 +369,21 @@ public void testBackupCopyTableWithDirtyDynamicPartitionStorageMedium() { Assert.assertTrue(copied.getTableProperty().hasInvalidDynamicPartition()); } + @Test + public void testBackupCopyPreservesDistributionMappingConstraint() { + DistributionMappingConstraint mapping = new DistributionMappingConstraint( + "mapping", "mapping_id", List.of("d1"), List.of("k1")); + TableProperty tableProperty = new TableProperty(Maps.newHashMap()); + tableProperty.addDistributionMappingConstraint(mapping); + table2.setTableProperty(tableProperty); + + OlapTable copied = table2.selectiveCopy(null, IndexExtState.VISIBLE, true); + + Assert.assertNotNull(copied); + Assert.assertEquals(mapping, copied.getTableProperty() + .getDistributionMappingConstraints().get(mapping.getName())); + } + @Test public void testBackupCopyTableWithDirtyDynamicPartitionStoragePolicy() { Map dirtyProperties = Maps.newHashMap(); diff --git a/fe/fe-core/src/test/java/org/apache/doris/backup/RestoreJobTest.java b/fe/fe-core/src/test/java/org/apache/doris/backup/RestoreJobTest.java index 40a6f72b2db2c2..4a48105c5abd57 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/backup/RestoreJobTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/backup/RestoreJobTest.java @@ -33,7 +33,10 @@ import org.apache.doris.catalog.ReplicaAllocation; import org.apache.doris.catalog.Resource; import org.apache.doris.catalog.Table; +import org.apache.doris.catalog.TableProperty; import org.apache.doris.catalog.Tablet; +import org.apache.doris.catalog.constraint.ConstraintManager; +import org.apache.doris.catalog.constraint.DistributionMappingConstraint; import org.apache.doris.common.AnalysisException; import org.apache.doris.common.FeConstants; import org.apache.doris.common.MarkedCountDownLatch; @@ -45,6 +48,7 @@ import org.apache.doris.system.SystemInfoService; import org.apache.doris.thrift.TStorageMedium; +import com.google.common.collect.ImmutableList; import com.google.common.collect.Lists; import com.google.common.collect.Maps; import org.junit.After; @@ -165,7 +169,7 @@ public void setUp() throws Exception { List metas = inv.getArgument(1); metas.add(backupMeta); return Status.OK; - }).when(repo).getSnapshotMetaFile(Mockito.eq(label), Mockito.anyList(), Mockito.eq(-1)); + }).when(repo).getSnapshotMetaFile(Mockito.eq(label), Mockito.anyList(), Mockito.anyInt()); mockedMarkedCountDownLatch = Mockito.mockConstruction(MarkedCountDownLatch.class, Mockito.withSettings().defaultAnswer(Mockito.CALLS_REAL_METHODS), @@ -275,6 +279,172 @@ public void testSerialization() throws IOException, AnalysisException { Files.delete(path); } + @Test + public void testRestoreMappingRejectsMixedFrontendVersionsWhenTargetExists() { + DistributionMappingConstraint mapping = new DistributionMappingConstraint( + "mapping", "mapping_id", List.of("k1"), List.of("k1")); + TableProperty tableProperty = new TableProperty(Maps.newHashMap()); + tableProperty.addDistributionMappingConstraint(mapping); + expectedRestoreTbl.setTableProperty(tableProperty); + Assert.assertTrue(db.registerTable(expectedRestoreTbl)); + ConstraintManager constraintManager = Mockito.mock(ConstraintManager.class); + Mockito.when(env.getConstraintManager()).thenReturn(constraintManager); + Mockito.when(constraintManager.getDistributionMappingConstraints(expectedRestoreTbl)) + .thenReturn(ImmutableList.of(mapping)); + Mockito.doThrow(new org.apache.doris.nereids.exceptions.AnalysisException("mixed versions")) + .when(constraintManager).validateDistributionMappingFeatureCompatibility(); + Deencapsulation.setField(job, "repo", repo); + + Deencapsulation.invoke(job, "checkAndPrepareMeta"); + + Assert.assertFalse(job.getStatus().ok()); + Assert.assertTrue(job.getStatus().getErrMsg().contains("Cannot restore table")); + Assert.assertEquals(OlapTable.OlapTableState.NORMAL, expectedRestoreTbl.getState()); + Mockito.verify(constraintManager).validateDistributionMappingFeatureCompatibility(); + Mockito.verify(constraintManager, Mockito.never()) + .validateDistributionMappingConstraints(expectedRestoreTbl); + } + + @Test + public void testNonAtomicRestoreRejectsExistingTableWithDistributionMapping() { + OlapTable localTable = expectedRestoreTbl.selectiveCopy(null, IndexExtState.VISIBLE, true); + DistributionMappingConstraint mapping = new DistributionMappingConstraint( + "mapping", "mapping_id", List.of("k1"), List.of("k1")); + TableProperty tableProperty = new TableProperty(Maps.newHashMap()); + tableProperty.addDistributionMappingConstraint(mapping); + localTable.setTableProperty(tableProperty); + Assert.assertTrue(db.registerTable(localTable)); + + ConstraintManager constraintManager = Mockito.mock(ConstraintManager.class); + Mockito.when(env.getConstraintManager()).thenReturn(constraintManager); + Mockito.when(constraintManager.getDistributionMappingConstraints(expectedRestoreTbl)) + .thenReturn(ImmutableList.of()); + Mockito.when(constraintManager.getDistributionMappingConstraints(localTable)) + .thenReturn(ImmutableList.of(mapping)); + Deencapsulation.setField(job, "repo", repo); + + Deencapsulation.invoke(job, "checkAndPrepareMeta"); + + Assert.assertFalse(job.getStatus().ok()); + Assert.assertTrue(job.getStatus().getErrMsg().contains( + "Cannot restore into existing table " + localTable.getName())); + Assert.assertEquals(OlapTable.OlapTableState.NORMAL, localTable.getState()); + } + + @Test + public void testNonAtomicRestoreRejectsMatchingDistributionMappingsOnExistingTable() { + DistributionMappingConstraint mapping = new DistributionMappingConstraint( + "mapping", "mapping_id", List.of("k1"), List.of("k1")); + TableProperty backupTableProperty = new TableProperty(Maps.newHashMap()); + backupTableProperty.addDistributionMappingConstraint(mapping); + expectedRestoreTbl.setTableProperty(backupTableProperty); + OlapTable localTable = expectedRestoreTbl.selectiveCopy(null, IndexExtState.VISIBLE, true); + Assert.assertTrue(db.registerTable(localTable)); + + ConstraintManager constraintManager = Mockito.mock(ConstraintManager.class); + Mockito.when(env.getConstraintManager()).thenReturn(constraintManager); + Mockito.when(constraintManager.getDistributionMappingConstraints(expectedRestoreTbl)) + .thenReturn(ImmutableList.of(mapping)); + Mockito.when(constraintManager.getDistributionMappingConstraints(localTable)) + .thenReturn(ImmutableList.of(mapping)); + Deencapsulation.setField(job, "repo", repo); + + Deencapsulation.invoke(job, "checkAndPrepareMeta"); + + Assert.assertFalse(job.getStatus().ok()); + Assert.assertTrue(job.getStatus().getErrMsg().contains( + "Cannot restore into existing table " + localTable.getName())); + Assert.assertEquals(OlapTable.OlapTableState.NORMAL, localTable.getState()); + Mockito.verify(constraintManager).validateDistributionMappingFeatureCompatibility(); + Mockito.verify(constraintManager).validateDistributionMappingConstraints(expectedRestoreTbl); + } + + @Test + public void testAtomicRestoreRejectsDistributionMappingBeforeStaging() { + DistributionMappingConstraint mapping = new DistributionMappingConstraint( + "mapping", "mapping_id", List.of("k1"), List.of("k1")); + TableProperty tableProperty = new TableProperty(Maps.newHashMap()); + tableProperty.addDistributionMappingConstraint(mapping); + expectedRestoreTbl.setTableProperty(tableProperty); + ConstraintManager constraintManager = Mockito.mock(ConstraintManager.class); + Mockito.when(env.getConstraintManager()).thenReturn(constraintManager); + Mockito.when(constraintManager.getDistributionMappingConstraints(expectedRestoreTbl)) + .thenReturn(ImmutableList.of(mapping)); + Deencapsulation.setField(job, "backupMeta", backupMeta); + Deencapsulation.setField(job, "isAtomicRestore", true); + + boolean valid = Deencapsulation.invoke(job, "validateDistributionMappingConstraintsForRestore"); + + Assert.assertFalse(valid); + Assert.assertTrue(job.getStatus().getErrMsg().contains( + "Cannot atomically restore table " + expectedRestoreTbl.getName())); + Assert.assertFalse(expectedRestoreTbl.isInAtomicRestore()); + Mockito.verify(constraintManager, Mockito.never()) + .validateDistributionMappingFeatureCompatibility(); + Mockito.verify(constraintManager, Mockito.never()) + .validateDistributionMappingConstraints(expectedRestoreTbl); + } + + @Test + public void testRestoreMappingRejectsIncompatibleSchema() { + DistributionMappingConstraint mapping = new DistributionMappingConstraint( + "mapping", "mapping_id", List.of("k1"), List.of("k1")); + TableProperty tableProperty = new TableProperty(Maps.newHashMap()); + tableProperty.addDistributionMappingConstraint(mapping); + expectedRestoreTbl.setTableProperty(tableProperty); + Assert.assertTrue(db.registerTable(expectedRestoreTbl)); + ConstraintManager constraintManager = Mockito.mock(ConstraintManager.class); + Mockito.when(env.getConstraintManager()).thenReturn(constraintManager); + Mockito.when(constraintManager.getDistributionMappingConstraints(expectedRestoreTbl)) + .thenReturn(ImmutableList.of(mapping)); + Mockito.doThrow(new org.apache.doris.nereids.exceptions.AnalysisException("incompatible schema")) + .when(constraintManager).validateDistributionMappingConstraints(expectedRestoreTbl); + Deencapsulation.setField(job, "repo", repo); + + Deencapsulation.invoke(job, "checkAndPrepareMeta"); + + Assert.assertFalse(job.getStatus().ok()); + Assert.assertTrue(job.getStatus().getErrMsg().contains("Cannot restore table")); + Assert.assertEquals(OlapTable.OlapTableState.NORMAL, expectedRestoreTbl.getState()); + Mockito.verify(constraintManager).validateDistributionMappingFeatureCompatibility(); + Mockito.verify(constraintManager).validateDistributionMappingConstraints(expectedRestoreTbl); + } + + @Test + public void testRestoreMappingValidatesEveryBackupTable() { + DistributionMappingConstraint firstMapping = new DistributionMappingConstraint( + "first_mapping", "mapping_id", List.of("k1"), List.of("k1")); + DistributionMappingConstraint secondMapping = new DistributionMappingConstraint( + "second_mapping", "mapping_id", List.of("k1"), List.of("k1")); + OlapTable secondRestoreTable = Mockito.mock(OlapTable.class); + Mockito.when(secondRestoreTable.getName()).thenReturn("second_restore_table"); + Mockito.when(secondRestoreTable.getId()).thenReturn(60000L); + + jobInfo.backupOlapTableObjects = Maps.newLinkedHashMap(); + jobInfo.backupOlapTableObjects.put(expectedRestoreTbl.getName(), new BackupOlapTableInfo()); + jobInfo.backupOlapTableObjects.put(secondRestoreTable.getName(), new BackupOlapTableInfo()); + backupMeta = new BackupMeta( + Lists.newArrayList(expectedRestoreTbl, secondRestoreTable), Lists.newArrayList()); + Deencapsulation.setField(job, "backupMeta", backupMeta); + + ConstraintManager constraintManager = Mockito.mock(ConstraintManager.class); + Mockito.when(env.getConstraintManager()).thenReturn(constraintManager); + Mockito.when(constraintManager.getDistributionMappingConstraints(expectedRestoreTbl)) + .thenReturn(ImmutableList.of(firstMapping)); + Mockito.when(constraintManager.getDistributionMappingConstraints(secondRestoreTable)) + .thenReturn(ImmutableList.of(secondMapping)); + Mockito.doThrow(new org.apache.doris.nereids.exceptions.AnalysisException("incompatible schema")) + .when(constraintManager).validateDistributionMappingConstraints(secondRestoreTable); + + boolean valid = Deencapsulation.invoke(job, "validateDistributionMappingConstraintsForRestore"); + + Assert.assertFalse(valid); + Assert.assertTrue(job.getStatus().getErrMsg().contains("second_restore_table")); + Mockito.verify(constraintManager).validateDistributionMappingFeatureCompatibility(); + Mockito.verify(constraintManager).validateDistributionMappingConstraints(expectedRestoreTbl); + Mockito.verify(constraintManager).validateDistributionMappingConstraints(secondRestoreTable); + } + @Test public void testResetPartitionVisibleAndNextVersionForRestore() throws Exception { long visibleVersion = 1234; diff --git a/fe/fe-core/src/test/java/org/apache/doris/catalog/OlapTableTest.java b/fe/fe-core/src/test/java/org/apache/doris/catalog/OlapTableTest.java index 81090d8d7ddc3d..6d719870177793 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/catalog/OlapTableTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/catalog/OlapTableTest.java @@ -18,6 +18,7 @@ package org.apache.doris.catalog; import org.apache.doris.catalog.TableIf.TableType; +import org.apache.doris.catalog.constraint.DistributionMappingConstraint; import org.apache.doris.catalog.info.IndexType; import org.apache.doris.cloud.common.util.CloudPropertyAnalyzer; import org.apache.doris.cloud.proto.Cloud; @@ -198,6 +199,23 @@ public void testResetPropertiesForRestore() { Assert.assertEquals((short) 3, olapTable.getDefaultReplicaAllocation().getTotalReplicaNum()); } + @Test + public void testBeingSyncedPropertiesRemoveDistributionMappings() { + TableProperty tableProperty = new TableProperty(Maps.newHashMap()); + tableProperty.addDistributionMappingConstraint(new DistributionMappingConstraint( + "mapping", "mapping_id", List.of("d1"), List.of("k1"))); + OlapTable olapTable = new OlapTable(); + olapTable.setTableProperty(tableProperty); + olapTable.setPartitionInfo(Mockito.mock(PartitionInfo.class)); + + olapTable.setBeingSyncedProperties(); + + Assert.assertTrue(olapTable.isBeingSynced()); + Assert.assertTrue(tableProperty.getDistributionMappingConstraints().isEmpty()); + Assert.assertFalse(tableProperty.getProperties().containsKey( + TableProperty.DISTRIBUTION_MAPPING_CONSTRAINTS_PROPERTY)); + } + @Test public void testBfIndexTableLevelFppDoesNotAffectSignature() throws IOException { try (MockedStatic mockedEnv = Mockito.mockStatic(Env.class, Mockito.CALLS_REAL_METHODS)) { diff --git a/fe/fe-core/src/test/java/org/apache/doris/catalog/constraint/ConstraintManagerTest.java b/fe/fe-core/src/test/java/org/apache/doris/catalog/constraint/ConstraintManagerTest.java index 628497f90f8d1a..90dc74f72edfe1 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/catalog/constraint/ConstraintManagerTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/catalog/constraint/ConstraintManagerTest.java @@ -17,9 +17,22 @@ package org.apache.doris.catalog.constraint; +import org.apache.doris.catalog.Column; +import org.apache.doris.catalog.DatabaseIf; +import org.apache.doris.catalog.Env; +import org.apache.doris.catalog.HashDistributionInfo; +import org.apache.doris.catalog.OlapTable; +import org.apache.doris.catalog.RandomDistributionInfo; +import org.apache.doris.catalog.TableProperty; +import org.apache.doris.catalog.Type; import org.apache.doris.catalog.info.TableNameInfo; import org.apache.doris.common.DdlException; +import org.apache.doris.common.Version; +import org.apache.doris.datasource.CatalogIf; +import org.apache.doris.datasource.InternalCatalog; import org.apache.doris.nereids.exceptions.AnalysisException; +import org.apache.doris.persist.gson.GsonUtils; +import org.apache.doris.system.Frontend; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; @@ -27,6 +40,8 @@ import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; +import org.mockito.Mockito; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; @@ -34,6 +49,8 @@ import java.io.DataInputStream; import java.io.DataOutput; import java.io.DataOutputStream; +import java.util.HashMap; +import java.util.List; import java.util.Map; /** @@ -123,6 +140,269 @@ void getConstraintsForNonExistentTableReturnsEmpty() { Assertions.assertTrue(mgr.getConstraints(T1).isEmpty()); } + @Test + void distributionMappingReplayUsesTableOwnedStorage() { + OlapTable table = Mockito.spy(new OlapTable()); + Mockito.doReturn(InternalCatalog.INTERNAL_CATALOG_ID).when(table).getCatalogId(); + DistributionMappingConstraint mapping = new DistributionMappingConstraint( + "mapping", "mapping_id", List.of("d1"), List.of("k1")); + + table.writeLock(); + try { + mgr.replayDistributionMappingConstraints(table, mappingProperties(mapping)); + } finally { + table.writeUnlock(); + } + + Assertions.assertTrue(mgr.getConstraints(T1).isEmpty()); + Assertions.assertEquals(mapping, mgr.getConstraint(T1, table, mapping.getName())); + Assertions.assertEquals(List.of(mapping), mgr.getDistributionMappingConstraints(table)); + + table.writeLock(); + try { + mgr.replayDistributionMappingConstraints(table, mappingProperties()); + } finally { + table.writeUnlock(); + } + Assertions.assertTrue(mgr.getDistributionMappingConstraints(table).isEmpty()); + } + + @Test + void distributionMappingReplayDefersConstraintNameCollisionToConsumers() { + OlapTable table = Mockito.spy(new OlapTable()); + Mockito.doReturn(InternalCatalog.INTERNAL_CATALOG_ID).when(table).getCatalogId(); + mgr.addConstraint(T1, "constraint", newPk("constraint", "k1"), true); + DistributionMappingConstraint mapping = new DistributionMappingConstraint( + "constraint", "mapping_id", List.of("d1"), List.of("k1")); + + table.writeLock(); + try { + mgr.replayDistributionMappingConstraints(table, mappingProperties(mapping)); + } finally { + table.writeUnlock(); + } + Assertions.assertEquals(List.of(mapping), mgr.getDistributionMappingConstraints(table)); + AnalysisException exception = Assertions.assertThrows( + AnalysisException.class, () -> mgr.getConstraints(T1, table)); + Assertions.assertTrue(exception.getMessage().contains("conflicts with another constraint")); + } + + @Test + void centralizedStorageRejectsDistributionMapping() { + DistributionMappingConstraint mapping = new DistributionMappingConstraint( + "mapping", "mapping_id", List.of("d1"), List.of("k1")); + + Assertions.assertThrows(IllegalArgumentException.class, + () -> mgr.addConstraint(T1, mapping.getName(), mapping, true)); + Assertions.assertTrue(mgr.isEmpty()); + } + + @Test + void distributionMappingsAreInternalOnlyAndReturnedInStableOrder() { + DistributionMappingConstraint second = new DistributionMappingConstraint( + "BB", "mapping_b", List.of("d1"), List.of("k1")); + DistributionMappingConstraint first = new DistributionMappingConstraint( + "Aa", "mapping_a", List.of("d1"), List.of("k1")); + TableProperty tableProperty = tablePropertyWithMappings(second, first); + + OlapTable internalTable = Mockito.mock(OlapTable.class); + Mockito.when(internalTable.getCatalogId()).thenReturn(InternalCatalog.INTERNAL_CATALOG_ID); + Mockito.when(internalTable.getTableProperty()).thenReturn(tableProperty); + Assertions.assertEquals(List.of(first, second), mgr.getDistributionMappingConstraints(internalTable)); + + OlapTable externalTable = Mockito.mock(OlapTable.class); + Mockito.when(externalTable.getCatalogId()).thenReturn(1L); + Mockito.when(externalTable.getTableProperty()).thenReturn(tableProperty); + Assertions.assertTrue(mgr.getDistributionMappingConstraints(externalTable).isEmpty()); + } + + @Test + void addDistributionMappingRejectsExternalOlapTable() { + OlapTable externalTable = Mockito.mock(OlapTable.class); + Mockito.when(externalTable.isWriteLockHeldByCurrentThread()).thenReturn(true); + Mockito.when(externalTable.getCatalogId()).thenReturn(1L); + DistributionMappingConstraint mapping = new DistributionMappingConstraint( + "mapping", "mapping_id", List.of("d1"), List.of("k1")); + + AnalysisException exception = Assertions.assertThrows(AnalysisException.class, + () -> mgr.addDistributionMappingConstraint(T1, externalTable, mapping)); + Assertions.assertTrue(exception.getMessage().contains("only supports internal OLAP tables")); + } + + @Test + void addDistributionMappingRejectsBeingSyncedTable() { + OlapTable table = Mockito.mock(OlapTable.class); + Mockito.when(table.isWriteLockHeldByCurrentThread()).thenReturn(true); + Mockito.when(table.getCatalogId()).thenReturn(InternalCatalog.INTERNAL_CATALOG_ID); + Mockito.when(table.isBeingSynced()).thenReturn(true); + DistributionMappingConstraint mapping = new DistributionMappingConstraint( + "mapping", "mapping_id", List.of("d1"), List.of("k1")); + + AnalysisException exception = Assertions.assertThrows(AnalysisException.class, + () -> mgr.addDistributionMappingConstraint(T1, table, mapping)); + + Assertions.assertTrue(exception.getMessage().contains("being synchronized by CCR")); + } + + @Test + void distributionMappingBindingUsesStableColumnIdsAcrossSchemaVersions() { + OlapTable table = Mockito.mock(OlapTable.class); + Column determinant = new Column("d1", Type.INT); + Column distribution = new Column("k1", Type.INT); + determinant.setUniqueId(10); + distribution.setUniqueId(20); + Mockito.when(table.getBaseSchemaVersion()).thenReturn(7); + Mockito.when(table.getColumn("d1")).thenReturn(determinant); + Mockito.when(table.getColumn("k1")).thenReturn(distribution); + HashDistributionInfo distributionInfo = Mockito.mock(HashDistributionInfo.class); + Mockito.when(distributionInfo.getDistributionColumns()).thenReturn(List.of(distribution)); + Mockito.when(table.getDefaultDistributionInfo()).thenReturn(distributionInfo); + + DistributionMappingConstraint unbound = new DistributionMappingConstraint( + "mapping", "mapping_id", List.of("d1"), List.of("k1")); + Assertions.assertFalse(unbound.isCompatibleWith(table)); + + DistributionMappingConstraint bound = unbound.bindTo(table); + Assertions.assertEquals(7, bound.getBaseSchemaVersion()); + Assertions.assertEquals(List.of(10), bound.getDeterminantColumnUniqueIds()); + Assertions.assertEquals(List.of(20), bound.getDistributionColumnUniqueIds()); + Assertions.assertEquals(List.of("int"), bound.getDeterminantColumnTypeSignatures()); + Assertions.assertEquals(List.of("int"), bound.getDistributionColumnTypeSignatures()); + Assertions.assertTrue(bound.isCompatibleWith(table)); + + Mockito.when(table.getBaseSchemaVersion()).thenReturn(8); + Assertions.assertTrue(bound.isCompatibleWith(table)); + + Column replacement = new Column("d1", Type.INT); + replacement.setUniqueId(11); + Mockito.when(table.getColumn("d1")).thenReturn(replacement); + Assertions.assertFalse(bound.isCompatibleWith(table)); + + Column typeChanged = new Column("d1", Type.BIGINT); + typeChanged.setUniqueId(10); + Mockito.when(table.getColumn("d1")).thenReturn(typeChanged); + Assertions.assertFalse(bound.isCompatibleWith(table)); + + Mockito.when(table.getColumn("d1")).thenReturn(determinant); + Mockito.when(table.getDefaultDistributionInfo()).thenReturn(new RandomDistributionInfo(4)); + Assertions.assertFalse(bound.isCompatibleWith(table)); + } + + @Test + void distributionMappingBindingUsesSchemaVersionWithoutStableColumnIds() { + OlapTable table = Mockito.mock(OlapTable.class); + Column determinant = new Column("d1", Type.INT); + Column distribution = new Column("k1", Type.INT); + Mockito.when(table.getBaseSchemaVersion()).thenReturn(7); + Mockito.when(table.getColumn("d1")).thenReturn(determinant); + Mockito.when(table.getColumn("k1")).thenReturn(distribution); + HashDistributionInfo distributionInfo = Mockito.mock(HashDistributionInfo.class); + Mockito.when(distributionInfo.getDistributionColumns()).thenReturn(List.of(distribution)); + Mockito.when(table.getDefaultDistributionInfo()).thenReturn(distributionInfo); + + DistributionMappingConstraint bound = new DistributionMappingConstraint( + "mapping", "mapping_id", List.of("d1"), List.of("k1")).bindTo(table); + Assertions.assertTrue(bound.isCompatibleWith(table)); + + Mockito.when(table.getBaseSchemaVersion()).thenReturn(8); + Assertions.assertFalse(bound.isCompatibleWith(table)); + } + + @Test + void validateDistributionMappingFeatureCompatibilityRejectsMixedOrUnknownFrontendVersions() { + String currentVersion = Version.DORIS_BUILD_VERSION + "-" + Version.DORIS_BUILD_SHORT_HASH; + Env env = Mockito.mock(Env.class); + Frontend currentFrontend = Mockito.mock(Frontend.class); + Frontend oldFrontend = Mockito.mock(Frontend.class); + Frontend unknownFrontend = Mockito.mock(Frontend.class); + + Mockito.when(env.getFrontends(null)) + .thenReturn(List.of(currentFrontend, oldFrontend, unknownFrontend)); + Mockito.when(currentFrontend.getVersion()).thenReturn(currentVersion); + Mockito.when(oldFrontend.getNodeName()).thenReturn("old-fe"); + Mockito.when(oldFrontend.getVersion()).thenReturn("old-version"); + Mockito.when(unknownFrontend.getNodeName()).thenReturn("unknown-fe"); + + try (MockedStatic mockedEnv = Mockito.mockStatic(Env.class)) { + mockedEnv.when(Env::getCurrentEnv).thenReturn(env); + AnalysisException exception = Assertions.assertThrows(AnalysisException.class, + mgr::validateDistributionMappingFeatureCompatibility); + Assertions.assertTrue(exception.getMessage().contains("old-fe(old-version)")); + Assertions.assertTrue(exception.getMessage().contains("unknown-fe(null)")); + Assertions.assertTrue(exception.getMessage().contains("cannot be added or restored")); + } + } + + @Test + void distributionMappingPlanningFallsBackForMixedOrUnknownFrontendVersions() { + TableProperty tableProperty = new TableProperty(new HashMap<>()); + OlapTable table = mockInternalMappingTable(tableProperty); + DistributionMappingConstraint mapping = new DistributionMappingConstraint( + "mapping", "mapping_id", List.of("d1"), List.of("k1")).bindTo(table); + tableProperty.addDistributionMappingConstraint(mapping); + + Env env = Mockito.mock(Env.class); + Frontend currentFrontend = Mockito.mock(Frontend.class); + Frontend oldFrontend = Mockito.mock(Frontend.class); + Frontend unknownFrontend = Mockito.mock(Frontend.class); + Mockito.when(currentFrontend.getVersion()) + .thenReturn(Version.DORIS_BUILD_VERSION + "-" + Version.DORIS_BUILD_SHORT_HASH); + Mockito.when(oldFrontend.getNodeName()).thenReturn("old-fe"); + Mockito.when(oldFrontend.getVersion()).thenReturn("old-version"); + Mockito.when(unknownFrontend.getNodeName()).thenReturn("unknown-fe"); + Mockito.when(env.getFrontends(null)) + .thenReturn(List.of(oldFrontend), List.of(unknownFrontend), List.of(currentFrontend)); + + try (MockedStatic mockedEnv = Mockito.mockStatic(Env.class)) { + mockedEnv.when(Env::getCurrentEnv).thenReturn(env); + Assertions.assertTrue(mgr.getDistributionMappingConstraintsForPlanning(table).isEmpty()); + Assertions.assertTrue(mgr.getDistributionMappingConstraintsForPlanning(table).isEmpty()); + Assertions.assertEquals(List.of(mapping), + mgr.getDistributionMappingConstraintsForPlanning(table)); + } + } + + @Test + void distributionMappingPlanningFallsBackForBeingSyncedTable() { + TableProperty tableProperty = new TableProperty(new HashMap<>()); + OlapTable table = mockInternalMappingTable(tableProperty); + DistributionMappingConstraint mapping = new DistributionMappingConstraint( + "mapping", "mapping_id", List.of("d1"), List.of("k1")).bindTo(table); + tableProperty.addDistributionMappingConstraint(mapping); + Mockito.when(table.isBeingSynced()).thenReturn(true); + + Assertions.assertEquals(List.of(mapping), mgr.getDistributionMappingConstraints(table)); + Assertions.assertTrue(mgr.getDistributionMappingConstraintsForPlanning(table).isEmpty()); + } + + @Test + void distributionMappingPlanningFallsBackForIncompatibleSchema() { + TableProperty tableProperty = new TableProperty(new HashMap<>()); + OlapTable table = mockInternalMappingTable(tableProperty); + DistributionMappingConstraint mapping = new DistributionMappingConstraint( + "mapping", "mapping_id", List.of("d1"), List.of("k1")).bindTo(table); + DistributionMappingConstraint validMapping = new DistributionMappingConstraint( + "valid_mapping", "valid_mapping_id", List.of("k1"), List.of("k1")).bindTo(table); + tableProperty.addDistributionMappingConstraint(mapping); + tableProperty.addDistributionMappingConstraint(validMapping); + + Env env = Mockito.mock(Env.class); + Mockito.when(env.getFrontends(null)).thenReturn(List.of()); + + try (MockedStatic mockedEnv = Mockito.mockStatic(Env.class)) { + mockedEnv.when(Env::getCurrentEnv).thenReturn(env); + Assertions.assertEquals(List.of(mapping, validMapping), + mgr.getDistributionMappingConstraintsForPlanning(table)); + + Column replacement = new Column("d1", Type.INT); + replacement.setUniqueId(11); + Mockito.when(table.getColumn("d1")).thenReturn(replacement); + Assertions.assertTrue(mgr.getDistributionMappingConstraintsForPlanning(table).isEmpty()); + } + Assertions.assertThrows(AnalysisException.class, + () -> mgr.validateDistributionMappingConstraints(table)); + } + // ==================== Type-specific getters ==================== @Test @@ -636,4 +916,44 @@ private static ForeignKeyConstraint newFk(String name, TableNameInfo refTable, return new ForeignKeyConstraint(name, ImmutableList.of(fkCol), refTable, ImmutableList.of(pkCol)); } + + private static Map mappingProperties(DistributionMappingConstraint... mappings) { + return ImmutableMap.of(TableProperty.DISTRIBUTION_MAPPING_CONSTRAINTS_PROPERTY, + GsonUtils.GSON.toJson(mappings)); + } + + private static TableProperty tablePropertyWithMappings(DistributionMappingConstraint... mappings) { + TableProperty tableProperty = new TableProperty(new HashMap<>()); + for (DistributionMappingConstraint mapping : mappings) { + tableProperty.addDistributionMappingConstraint(mapping); + } + return tableProperty; + } + + private static OlapTable mockInternalMappingTable(TableProperty tableProperty) { + Column determinant = new Column("d1", Type.INT); + determinant.setUniqueId(10); + Column distribution = new Column("k1", Type.INT); + distribution.setUniqueId(20); + HashDistributionInfo distributionInfo = Mockito.mock(HashDistributionInfo.class); + Mockito.when(distributionInfo.getDistributionColumns()).thenReturn(List.of(distribution)); + + CatalogIf catalog = Mockito.mock(CatalogIf.class); + Mockito.when(catalog.getName()).thenReturn("ctl"); + DatabaseIf database = Mockito.mock(DatabaseIf.class); + Mockito.when(database.getCatalog()).thenReturn(catalog); + Mockito.when(database.getFullName()).thenReturn("db"); + + OlapTable table = Mockito.mock(OlapTable.class); + Mockito.when(table.getId()).thenReturn(1L); + Mockito.when(table.getName()).thenReturn("t1"); + Mockito.when(table.getCatalogId()).thenReturn(InternalCatalog.INTERNAL_CATALOG_ID); + Mockito.when(table.getDatabase()).thenReturn(database); + Mockito.when(table.getTableProperty()).thenReturn(tableProperty); + Mockito.when(table.getBaseSchemaVersion()).thenReturn(7); + Mockito.when(table.getColumn("d1")).thenReturn(determinant); + Mockito.when(table.getColumn("k1")).thenReturn(distribution); + Mockito.when(table.getDefaultDistributionInfo()).thenReturn(distributionInfo); + return table; + } } diff --git a/fe/fe-core/src/test/java/org/apache/doris/catalog/constraint/ConstraintPersistTest.java b/fe/fe-core/src/test/java/org/apache/doris/catalog/constraint/ConstraintPersistTest.java index 7a40236d1334b8..403af94cb5b7c6 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/catalog/constraint/ConstraintPersistTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/catalog/constraint/ConstraintPersistTest.java @@ -17,11 +17,14 @@ package org.apache.doris.catalog.constraint; +import org.apache.doris.binlog.BinlogManager; import org.apache.doris.catalog.Column; import org.apache.doris.catalog.Env; import org.apache.doris.catalog.MTMV; +import org.apache.doris.catalog.OlapTable; import org.apache.doris.catalog.PrimitiveType; import org.apache.doris.catalog.TableIf; +import org.apache.doris.catalog.TableProperty; import org.apache.doris.catalog.info.TableNameInfo; import org.apache.doris.common.AnalysisException; import org.apache.doris.common.Config; @@ -34,6 +37,7 @@ import org.apache.doris.nereids.util.RelationUtil; import org.apache.doris.persist.AlterConstraintLog; import org.apache.doris.persist.EditLog; +import org.apache.doris.persist.ModifyTablePropertyOperationLog; import org.apache.doris.persist.OperationType; import org.apache.doris.utframe.TestWithFeService; @@ -207,6 +211,72 @@ void externalTableTest() throws Exception { Assertions.assertEquals(1, loadedMgr.getConstraints(extTni).size()); } + @Test + void distributionMappingTablePropertyJournalReplayTest() throws Exception { + OlapTable table = (OlapTable) RelationUtil.getTable( + RelationUtil.getQualifierName(connectContext, Lists.newArrayList("test", "t1")), + connectContext.getEnv(), Optional.empty()); + TableNameInfo tableNameInfo = new TableNameInfo(table.getNameWithFullQualifiers()); + DistributionMappingConstraint mapping = new DistributionMappingConstraint( + "mapping_replay", "mapping_replay", List.of("k2"), List.of("k1")); + ConstraintManager manager = Env.getCurrentEnv().getConstraintManager(); + TableProperty mappingProperty = new TableProperty(Maps.newHashMap()); + mappingProperty.addDistributionMappingConstraint(mapping); + + JournalEntity addJournal = new JournalEntity(); + addJournal.setData(new ModifyTablePropertyOperationLog( + table.getDatabase().getId(), table.getId(), table.getName(), + mappingProperty.getDistributionMappingConstraintProperties())); + addJournal.setOpCode(OperationType.OP_MODIFY_TABLE_PROPERTIES); + EditLog.loadJournal(Env.getCurrentEnv(), 0L, addJournal); + + Assertions.assertNull(manager.getConstraint(tableNameInfo, mapping.getName())); + Assertions.assertEquals(mapping, + manager.getConstraint(tableNameInfo, table, mapping.getName())); + + JournalEntity dropJournal = new JournalEntity(); + mappingProperty.removeDistributionMappingConstraint(mapping.getName()); + dropJournal.setData(new ModifyTablePropertyOperationLog( + table.getDatabase().getId(), table.getId(), table.getName(), + mappingProperty.getDistributionMappingConstraintProperties())); + dropJournal.setOpCode(OperationType.OP_MODIFY_TABLE_PROPERTIES); + EditLog.loadJournal(Env.getCurrentEnv(), 0L, dropJournal); + + Assertions.assertTrue(manager.getDistributionMappingConstraints(table).isEmpty()); + } + + @Test + void distributionMappingReplayDoesNotPublishTablePropertyBinlog() throws Exception { + Env env = Mockito.mock(Env.class); + BinlogManager binlogManager = Mockito.mock(BinlogManager.class); + Mockito.when(env.getBinlogManager()).thenReturn(binlogManager); + DistributionMappingConstraint mapping = new DistributionMappingConstraint( + "mapping", "mapping_id", List.of("k2"), List.of("k1")); + TableProperty mappingProperty = new TableProperty(Maps.newHashMap()); + mappingProperty.addDistributionMappingConstraint(mapping); + ModifyTablePropertyOperationLog mappingLog = new ModifyTablePropertyOperationLog( + 1L, 2L, "table", mappingProperty.getDistributionMappingConstraintProperties()); + JournalEntity mappingJournal = new JournalEntity(); + mappingJournal.setData(mappingLog); + mappingJournal.setOpCode(OperationType.OP_MODIFY_TABLE_PROPERTIES); + + EditLog.loadJournal(env, 3L, mappingJournal); + + Mockito.verify(env).replayModifyTableProperty( + OperationType.OP_MODIFY_TABLE_PROPERTIES, mappingLog); + Mockito.verifyNoInteractions(binlogManager); + + ModifyTablePropertyOperationLog propertyLog = new ModifyTablePropertyOperationLog( + 1L, 2L, "table", Map.of("in_memory", "false")); + JournalEntity propertyJournal = new JournalEntity(); + propertyJournal.setData(propertyLog); + propertyJournal.setOpCode(OperationType.OP_MODIFY_TABLE_PROPERTIES); + + EditLog.loadJournal(env, 4L, propertyJournal); + + Mockito.verify(binlogManager).addModifyTableProperty(propertyLog, 4L); + } + @Test void addConstraintLogPersistForExternalTableTest() throws Exception { Config.edit_log_type = "local"; diff --git a/fe/fe-core/src/test/java/org/apache/doris/catalog/constraint/DistributionMappingConstraintPersistTest.java b/fe/fe-core/src/test/java/org/apache/doris/catalog/constraint/DistributionMappingConstraintPersistTest.java new file mode 100644 index 00000000000000..8316d1b0431f54 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/catalog/constraint/DistributionMappingConstraintPersistTest.java @@ -0,0 +1,138 @@ +// 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.OlapTable; +import org.apache.doris.catalog.TableProperty; +import org.apache.doris.catalog.Type; +import org.apache.doris.persist.ModifyTablePropertyOperationLog; +import org.apache.doris.persist.gson.GsonUtils; + +import com.google.gson.annotations.SerializedName; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +class DistributionMappingConstraintPersistTest { + + @Test + void tablePropertyRoundTripPreservesMappingsInStableOrder() { + DistributionMappingConstraint second = newBoundMapping("mapping_b", "mapping_b_id"); + DistributionMappingConstraint first = newBoundMapping("mapping_a", "mapping_a_id"); + TableProperty tableProperty = new TableProperty(new HashMap<>()); + tableProperty.addDistributionMappingConstraint(second); + tableProperty.addDistributionMappingConstraint(first); + + String snapshot = tableProperty.getProperties() + .get(TableProperty.DISTRIBUTION_MAPPING_CONSTRAINTS_PROPERTY); + Assertions.assertTrue(snapshot.indexOf("mapping_a") < snapshot.indexOf("mapping_b")); + + String json = GsonUtils.GSON.toJson(tableProperty); + Assertions.assertFalse(json.contains("\"distributionMappingConstraints\"")); + TableProperty restored = GsonUtils.GSON.fromJson(json, TableProperty.class); + DistributionMappingConstraint restoredMapping = restored.getDistributionMappingConstraints().get("mapping_a"); + Assertions.assertEquals(List.of(first, second), restored.getDistributionMappingConstraints().entrySet().stream() + .sorted(Map.Entry.comparingByKey()) + .map(Map.Entry::getValue) + .toList()); + Assertions.assertEquals(first, restoredMapping); + Assertions.assertEquals(first.getBaseSchemaVersion(), restoredMapping.getBaseSchemaVersion()); + Assertions.assertEquals( + first.getDeterminantColumnUniqueIds(), restoredMapping.getDeterminantColumnUniqueIds()); + Assertions.assertEquals( + first.getDistributionColumnUniqueIds(), restoredMapping.getDistributionColumnUniqueIds()); + Assertions.assertEquals( + first.getDeterminantColumnTypeSignatures(), restoredMapping.getDeterminantColumnTypeSignatures()); + Assertions.assertEquals( + first.getDistributionColumnTypeSignatures(), restoredMapping.getDistributionColumnTypeSignatures()); + } + + @Test + void oldFrontendReplayAndCheckpointPreserveAddAndDropSnapshots() { + DistributionMappingConstraint mapping = newBoundMapping("mapping", "mapping_id"); + TableProperty currentProperty = new TableProperty(new HashMap<>()); + currentProperty.addDistributionMappingConstraint(mapping); + ModifyTablePropertyOperationLog addLog = new ModifyTablePropertyOperationLog( + 1L, 2L, "table", currentProperty.getDistributionMappingConstraintProperties()); + + LegacyModifyTablePropertyOperationLog legacyAddLog = GsonUtils.GSON.fromJson( + addLog.toJson(), LegacyModifyTablePropertyOperationLog.class); + Assertions.assertEquals(1L, legacyAddLog.dbId); + Assertions.assertEquals(2L, legacyAddLog.tableId); + Assertions.assertEquals("table", legacyAddLog.tableName); + Assertions.assertTrue(legacyAddLog.properties.containsKey( + TableProperty.DISTRIBUTION_MAPPING_CONSTRAINTS_PROPERTY)); + + LegacyTableProperty legacyTableProperty = new LegacyTableProperty(); + legacyTableProperty.properties.putAll(legacyAddLog.properties); + legacyTableProperty.properties.put("in_memory", "false"); + TableProperty restoredAfterOldCheckpoint = GsonUtils.GSON.fromJson( + GsonUtils.GSON.toJson(legacyTableProperty), TableProperty.class); + Assertions.assertEquals(mapping, + restoredAfterOldCheckpoint.getDistributionMappingConstraints().get(mapping.getName())); + + currentProperty.removeDistributionMappingConstraint(mapping.getName()); + ModifyTablePropertyOperationLog dropLog = new ModifyTablePropertyOperationLog( + 1L, 2L, "table", currentProperty.getDistributionMappingConstraintProperties()); + LegacyModifyTablePropertyOperationLog legacyDropLog = GsonUtils.GSON.fromJson( + dropLog.toJson(), LegacyModifyTablePropertyOperationLog.class); + Assertions.assertEquals("[]", legacyDropLog.properties.get( + TableProperty.DISTRIBUTION_MAPPING_CONSTRAINTS_PROPERTY)); + + legacyTableProperty.properties.putAll(legacyDropLog.properties); + TableProperty restoredAfterOldDropCheckpoint = GsonUtils.GSON.fromJson( + GsonUtils.GSON.toJson(legacyTableProperty), TableProperty.class); + Assertions.assertTrue(restoredAfterOldDropCheckpoint.getDistributionMappingConstraints().isEmpty()); + Assertions.assertEquals("[]", restoredAfterOldDropCheckpoint.getProperties().get( + TableProperty.DISTRIBUTION_MAPPING_CONSTRAINTS_PROPERTY)); + } + + private DistributionMappingConstraint newBoundMapping(String name, String mappingId) { + OlapTable table = Mockito.mock(OlapTable.class); + Column determinant = new Column("d1", Type.INT); + Column distribution = new Column("k1", Type.BIGINT); + determinant.setUniqueId(10); + distribution.setUniqueId(20); + Mockito.when(table.getBaseSchemaVersion()).thenReturn(7); + Mockito.when(table.getColumn("d1")).thenReturn(determinant); + Mockito.when(table.getColumn("k1")).thenReturn(distribution); + return new DistributionMappingConstraint( + name, mappingId, List.of("d1"), List.of("k1")).bindTo(table); + } + + private static class LegacyTableProperty { + @SerializedName("properties") + private Map properties = new HashMap<>(); + } + + private static class LegacyModifyTablePropertyOperationLog { + @SerializedName("dbId") + private long dbId; + @SerializedName("tableId") + private long tableId; + @SerializedName("tableName") + private String tableName; + @SerializedName("properties") + private Map properties = new HashMap<>(); + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/properties/ChildOutputPropertyDeriverTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/properties/ChildOutputPropertyDeriverTest.java index fd23de7a2f46d6..73881c081ee25c 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/properties/ChildOutputPropertyDeriverTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/properties/ChildOutputPropertyDeriverTest.java @@ -26,14 +26,19 @@ import org.apache.doris.nereids.memo.GroupExpression; import org.apache.doris.nereids.memo.GroupId; import org.apache.doris.nereids.properties.DistributionSpecHash.ShuffleType; +import org.apache.doris.nereids.rules.implementation.LogicalWindowToPhysicalWindow.WindowFrameGroup; +import org.apache.doris.nereids.trees.expressions.AggregateExpression; import org.apache.doris.nereids.trees.expressions.Alias; import org.apache.doris.nereids.trees.expressions.AssertNumRowsElement; +import org.apache.doris.nereids.trees.expressions.Cast; import org.apache.doris.nereids.trees.expressions.EqualTo; import org.apache.doris.nereids.trees.expressions.ExprId; import org.apache.doris.nereids.trees.expressions.NamedExpression; import org.apache.doris.nereids.trees.expressions.Slot; import org.apache.doris.nereids.trees.expressions.SlotReference; import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateParam; +import org.apache.doris.nereids.trees.expressions.functions.agg.MultiDistinctCount; +import org.apache.doris.nereids.trees.expressions.functions.agg.Sum; import org.apache.doris.nereids.trees.expressions.functions.scalar.Abs; import org.apache.doris.nereids.trees.expressions.literal.Literal; import org.apache.doris.nereids.trees.plans.AggMode; @@ -42,26 +47,33 @@ import org.apache.doris.nereids.trees.plans.GroupPlan; import org.apache.doris.nereids.trees.plans.JoinType; import org.apache.doris.nereids.trees.plans.LimitPhase; +import org.apache.doris.nereids.trees.plans.PartitionTopnPhase; import org.apache.doris.nereids.trees.plans.RelationId; import org.apache.doris.nereids.trees.plans.SortPhase; +import org.apache.doris.nereids.trees.plans.WindowFuncType; import org.apache.doris.nereids.trees.plans.logical.LogicalOneRowRelation; import org.apache.doris.nereids.trees.plans.physical.AbstractPhysicalPlan; import org.apache.doris.nereids.trees.plans.physical.PhysicalAssertNumRows; +import org.apache.doris.nereids.trees.plans.physical.PhysicalGenerate; import org.apache.doris.nereids.trees.plans.physical.PhysicalHashAggregate; import org.apache.doris.nereids.trees.plans.physical.PhysicalHashJoin; import org.apache.doris.nereids.trees.plans.physical.PhysicalLimit; import org.apache.doris.nereids.trees.plans.physical.PhysicalNestedLoopJoin; +import org.apache.doris.nereids.trees.plans.physical.PhysicalPartitionTopN; import org.apache.doris.nereids.trees.plans.physical.PhysicalQuickSort; import org.apache.doris.nereids.trees.plans.physical.PhysicalRepeat; import org.apache.doris.nereids.trees.plans.physical.PhysicalTopN; +import org.apache.doris.nereids.trees.plans.physical.PhysicalWindow; import org.apache.doris.nereids.types.BigIntType; import org.apache.doris.nereids.types.IntegerType; import org.apache.doris.nereids.types.TinyIntType; +import org.apache.doris.nereids.types.VarcharType; import org.apache.doris.nereids.util.ExpressionUtils; import org.apache.doris.qe.ConnectContext; import org.apache.doris.qe.SessionVariable; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; import com.google.common.collect.Lists; import com.google.common.collect.Maps; import com.google.common.collect.Sets; @@ -600,14 +612,17 @@ ExpressionUtils.EMPTY_CONDITION, new DistributeHint(DistributeType.NONE), Map leftMap = Maps.newHashMap(); leftMap.put(new ExprId(0), 0); leftMap.put(new ExprId(1), 0); - PhysicalProperties left = new PhysicalProperties(new DistributionSpecHash( + DistributionSpecHash leftHash = new DistributionSpecHash( Lists.newArrayList(new ExprId(0)), ShuffleType.NATURAL, 0, + -1L, Sets.newHashSet(0L), ImmutableList.of(Sets.newHashSet(new ExprId(0), new ExprId(1))), - leftMap - )); + leftMap, + ImmutableList.of(new DistributionMapping( + "mapping", ImmutableList.of(new ExprId(1)), ImmutableList.of(0)))); + PhysicalProperties left = naturalProperties(leftHash); PhysicalProperties right = new PhysicalProperties(DistributionSpecReplicated.INSTANCE, new OrderSpec(Lists.newArrayList( @@ -623,6 +638,7 @@ ExpressionUtils.EMPTY_CONDITION, new DistributeHint(DistributeType.NONE), Assertions.assertEquals(ShuffleType.NATURAL, actual.getShuffleType()); // check merged Assertions.assertEquals(3, actual.getExprIdToEquivalenceSet().size()); + assertMappingLocalityCleared(result); } @Test @@ -682,11 +698,13 @@ void testNestedLoopJoin() { Lists.newArrayList(new ExprId(0)), ShuffleType.NATURAL, 0, + -1L, Sets.newHashSet(0L), ImmutableList.of(Sets.newHashSet(new ExprId(0), new ExprId(1))), - leftMap - ); - PhysicalProperties left = new PhysicalProperties(leftHash); + leftMap, + ImmutableList.of(new DistributionMapping( + "mapping", ImmutableList.of(new ExprId(1)), ImmutableList.of(0)))); + PhysicalProperties left = naturalProperties(leftHash); PhysicalProperties right = PhysicalProperties.REPLICATED; List childrenOutputProperties = Lists.newArrayList(left, right); ChildOutputPropertyDeriver deriver = new ChildOutputPropertyDeriver(childrenOutputProperties); @@ -695,7 +713,50 @@ void testNestedLoopJoin() { Assertions.assertTrue(result.getOrderSpec().getOrderKeys().isEmpty()); Assertions.assertInstanceOf(DistributionSpecHash.class, result.getDistributionSpec()); DistributionSpecHash actual = (DistributionSpecHash) result.getDistributionSpec(); - Assertions.assertEquals(leftHash, actual); + Assertions.assertEquals(leftHash.withoutDistributionMappings(), actual); + assertMappingLocalityCleared(result); + } + + @Test + void testRuntimePlacementBarriersDropMappingLocality() { + SlotReference k1 = new SlotReference("k1", IntegerType.INSTANCE); + SlotReference d1 = new SlotReference("d1", IntegerType.INSTANCE); + DistributionSpecHash childHash = new DistributionSpecHash( + ImmutableList.of(k1.getExprId()), ShuffleType.NATURAL, + 1L, 2L, ImmutableSet.of(3L), + ImmutableList.of(new DistributionMapping( + "mapping", ImmutableList.of(d1.getExprId()), ImmutableList.of(0)))); + PhysicalProperties childProperties = naturalProperties(childHash); + + PhysicalGenerate generate = new PhysicalGenerate<>( + ImmutableList.of(), ImmutableList.of(), ImmutableList.of(), logicalProperties, groupPlan); + assertMappingLocalityCleared(deriveOutputProperties(generate, childProperties)); + + PhysicalWindow window = new PhysicalWindow<>( + Mockito.mock(WindowFrameGroup.class), null, ImmutableList.of(), false, + logicalProperties, groupPlan); + assertMappingLocalityCleared(deriveOutputProperties(window, childProperties)); + + for (PartitionTopnPhase phase : PartitionTopnPhase.values()) { + PhysicalPartitionTopN partitionTopN = new PhysicalPartitionTopN<>( + WindowFuncType.ROW_NUMBER, ImmutableList.of(k1), ImmutableList.of(), + false, 1, phase, logicalProperties, groupPlan); + assertMappingLocalityCleared(deriveOutputProperties(partitionTopN, childProperties)); + } + } + + @Test + void testPhysicalPropertiesDoesNotInferNaturalMappingProof() { + SlotReference k1 = new SlotReference("k1", IntegerType.INSTANCE); + SlotReference d1 = new SlotReference("d1", IntegerType.INSTANCE); + DistributionSpecHash hashSpec = new DistributionSpecHash( + ImmutableList.of(k1.getExprId()), ShuffleType.NATURAL, + 1L, 2L, ImmutableSet.of(3L), + ImmutableList.of(new DistributionMapping( + "mapping", ImmutableList.of(d1.getExprId()), ImmutableList.of(0)))); + + Assertions.assertFalse(new PhysicalProperties(hashSpec) + .getNaturalDistributionMappingSpec().isPresent()); } @Test @@ -754,6 +815,329 @@ void testGlobalPhaseAggregate() { actual.getOrderedShuffledColumns()); } + @Test + void testAggregatePropagatesDistributionMappings() { + SlotReference k1 = new SlotReference("k1", IntegerType.INSTANCE); + SlotReference k2 = new SlotReference("k2", IntegerType.INSTANCE); + SlotReference d1 = new SlotReference("d1", IntegerType.INSTANCE); + Alias outputK1 = new Alias(k1, "output_k1"); + Alias outputK2 = new Alias(k2, "output_k2"); + Alias outputD1 = new Alias(d1, "output_d1"); + AggregateParam aggregateParam = new AggregateParam(AggPhase.GLOBAL, AggMode.BUFFER_TO_RESULT); + PhysicalHashAggregate aggregate = new PhysicalHashAggregate<>( + ImmutableList.of(k1, k2, d1), + ImmutableList.of(outputK1, outputK2, outputD1, sumOutput(d1, aggregateParam)), + aggregateParam, + true, + logicalProperties, + false, + groupPlan); + + DistributionSpecHash childHash = new DistributionSpecHash( + ImmutableList.of(k1.getExprId(), k2.getExprId()), ShuffleType.NATURAL, + 1L, 2L, ImmutableSet.of(3L), + ImmutableList.of(new DistributionMapping( + "mapping_1", ImmutableList.of(d1.getExprId()), ImmutableList.of(0)))); + DistributionSpecHash outputHash = deriveAggregateHash(aggregate, childHash); + + Assertions.assertEquals( + ImmutableList.of(outputK1.getExprId(), outputK2.getExprId()), + outputHash.getOrderedShuffledColumns()); + Assertions.assertEquals( + ImmutableList.of(outputD1.getExprId()), + outputHash.getDistributionMappings().get(0).getDeterminantExprIds()); + } + + @Test + void testAggregateKeepsMappingLocalityWhenDistributionKeyIsHidden() { + SlotReference k1 = new SlotReference("k1", IntegerType.INSTANCE); + SlotReference k2 = new SlotReference("k2", IntegerType.INSTANCE); + SlotReference d1 = new SlotReference("d1", IntegerType.INSTANCE); + Alias outputK2 = new Alias(k2, "output_k2"); + Alias outputD1 = new Alias(d1, "output_d1"); + AggregateParam aggregateParam = new AggregateParam(AggPhase.GLOBAL, AggMode.BUFFER_TO_RESULT); + PhysicalHashAggregate aggregate = new PhysicalHashAggregate<>( + ImmutableList.of(d1, k2), + ImmutableList.of(outputK2, outputD1, sumOutput(d1, aggregateParam)), + aggregateParam, + true, + logicalProperties, + false, + groupPlan); + + PhysicalProperties orderedChild = naturalProperties( + naturalHashWithMapping(k1, k2, d1), + new OrderSpec(ImmutableList.of(new OrderKey(d1, true, false)))); + PhysicalProperties output = deriveAggregateProperties(aggregate, orderedChild); + + Assertions.assertSame(DistributionSpecStorageAny.INSTANCE, output.getDistributionSpec()); + Assertions.assertTrue(output.getOrderSpec().getOrderKeys().isEmpty()); + Assertions.assertTrue(output.getNaturalDistributionMappingSpec().isPresent()); + NaturalDistributionMappingSpec mappingSpec = output.getNaturalDistributionMappingSpec().get(); + Assertions.assertEquals( + Integer.valueOf(1), + mappingSpec.getVisibleDistributionExprToIndex().get(outputK2.getExprId())); + Assertions.assertEquals( + ImmutableList.of(outputD1.getExprId()), + mappingSpec.getDistributionMappings().get(0).getDeterminantExprIds()); + Assertions.assertTrue(output.satisfy(new PhysicalProperties(new DistributionSpecHash( + ImmutableList.of(outputD1.getExprId(), outputK2.getExprId()), + ShuffleType.COLOCATE_MAPPING_REQUIRE)))); + } + + @Test + void testAggregateClearsOrderWhenMappingProofFallsBack() { + SlotReference k1 = new SlotReference("k1", IntegerType.INSTANCE); + SlotReference k2 = new SlotReference("k2", IntegerType.INSTANCE); + SlotReference d1 = new SlotReference("d1", IntegerType.INSTANCE); + PhysicalProperties orderedChild = naturalProperties( + naturalHashWithMapping(k1, k2, d1), + new OrderSpec(ImmutableList.of(new OrderKey(d1, true, false)))); + AggregateParam aggregateParam = new AggregateParam(AggPhase.GLOBAL, AggMode.BUFFER_TO_RESULT); + PhysicalHashAggregate repeatAggregate = new PhysicalHashAggregate<>( + ImmutableList.of(k1, k2, d1), + ImmutableList.of(k1, k2, d1, sumOutput(d1, aggregateParam)), + aggregateParam, + true, + logicalProperties, + true, + groupPlan); + PhysicalHashAggregate incompleteGroupByAggregate = + new PhysicalHashAggregate<>( + ImmutableList.of(d1), + ImmutableList.of(d1, sumOutput(d1, aggregateParam)), + aggregateParam, + true, + logicalProperties, + false, + groupPlan); + + PhysicalProperties repeatOutput = + deriveAggregateProperties(repeatAggregate, orderedChild); + PhysicalProperties incompleteOutput = + deriveAggregateProperties(incompleteGroupByAggregate, orderedChild); + + Assertions.assertTrue(repeatOutput.getOrderSpec().getOrderKeys().isEmpty()); + Assertions.assertTrue(incompleteOutput.getOrderSpec().getOrderKeys().isEmpty()); + Assertions.assertFalse( + repeatOutput.getNaturalDistributionMappingSpec().isPresent()); + Assertions.assertFalse( + incompleteOutput.getNaturalDistributionMappingSpec().isPresent()); + } + + @Test + void testAggregateDropsMappingsWithoutRequiredOutputs() { + SlotReference k1 = new SlotReference("k1", IntegerType.INSTANCE); + SlotReference k2 = new SlotReference("k2", IntegerType.INSTANCE); + SlotReference d1 = new SlotReference("d1", IntegerType.INSTANCE); + AggregateParam aggregateParam = new AggregateParam(AggPhase.GLOBAL, AggMode.BUFFER_TO_RESULT); + PhysicalHashAggregate missingDistributionKey = new PhysicalHashAggregate<>( + ImmutableList.of(k1, d1), + ImmutableList.of(k1, d1, sumOutput(d1, aggregateParam)), + aggregateParam, + true, + logicalProperties, + false, + groupPlan); + PhysicalHashAggregate missingDeterminant = new PhysicalHashAggregate<>( + ImmutableList.of(k1, k2, d1), + ImmutableList.of(k1, k2, sumOutput(d1, aggregateParam)), + aggregateParam, + true, + logicalProperties, + false, + groupPlan); + + Assertions.assertTrue(deriveAggregateHash( + missingDistributionKey, + naturalHashWithMapping(k1, k2, d1)).getDistributionMappings().isEmpty()); + Assertions.assertTrue(deriveAggregateHash( + missingDeterminant, + naturalHashWithMapping(k1, k2, d1)).getDistributionMappings().isEmpty()); + } + + @Test + void testAggregateDoesNotUseIncompleteMappingDeterminant() { + SlotReference k1 = new SlotReference("k1", IntegerType.INSTANCE); + SlotReference k2 = new SlotReference("k2", IntegerType.INSTANCE); + SlotReference d1 = new SlotReference("d1", IntegerType.INSTANCE); + SlotReference d2 = new SlotReference("d2", IntegerType.INSTANCE); + AggregateParam aggregateParam = new AggregateParam(AggPhase.GLOBAL, AggMode.BUFFER_TO_RESULT); + PhysicalHashAggregate aggregate = new PhysicalHashAggregate<>( + ImmutableList.of(d1, k2), + ImmutableList.of(d1, k2, sumOutput(d1, aggregateParam)), + aggregateParam, + true, + logicalProperties, + false, + groupPlan); + DistributionSpecHash childHash = new DistributionSpecHash( + ImmutableList.of(k1.getExprId(), k2.getExprId()), ShuffleType.NATURAL, + 1L, 2L, ImmutableSet.of(3L), + ImmutableList.of(new DistributionMapping( + "mapping_1", ImmutableList.of(d1.getExprId(), d2.getExprId()), ImmutableList.of(0)))); + + Assertions.assertTrue(deriveAggregateHash( + aggregate, childHash).getDistributionMappings().isEmpty()); + } + + @Test + void testDistinctAggregateDropsMappingFromNaturalChild() { + SlotReference k1 = new SlotReference("k1", IntegerType.INSTANCE); + SlotReference k2 = new SlotReference("k2", IntegerType.INSTANCE); + SlotReference d1 = new SlotReference("d1", IntegerType.INSTANCE); + SlotReference extra = new SlotReference("extra", IntegerType.INSTANCE); + DistributionSpecHash childHash = naturalHashWithMapping(k1, k2, d1); + ConnectContext connectContext = ConnectContext.get(); + boolean originalEnableBucketedHashAgg = connectContext.getSessionVariable().enableBucketedHashAgg; + try { + connectContext.getSessionVariable().enableBucketedHashAgg = false; + for (AggPhase phase : AggPhase.values()) { + AggregateParam aggregateParam = new AggregateParam(phase, AggMode.INPUT_TO_RESULT); + PhysicalHashAggregate distinctAggregate = new PhysicalHashAggregate<>( + ImmutableList.of(d1, k2), + ImmutableList.of(k2, d1, new Alias( + new AggregateExpression( + new MultiDistinctCount(extra), aggregateParam), "distinct_count")), + aggregateParam, + true, + logicalProperties, + false, + groupPlan); + + PhysicalProperties output = deriveAggregateProperties(distinctAggregate, childHash); + + assertMappingLocalityCleared(output); + } + + AggregateParam distinctPhaseParam = new AggregateParam( + AggPhase.DISTINCT_GLOBAL, AggMode.BUFFER_TO_RESULT); + PhysicalHashAggregate distinctPhaseAggregate = new PhysicalHashAggregate<>( + ImmutableList.of(d1, k2), + ImmutableList.of(k2, d1, sumOutput(extra, distinctPhaseParam)), + distinctPhaseParam, + true, + logicalProperties, + false, + groupPlan); + assertMappingLocalityCleared(deriveAggregateProperties(distinctPhaseAggregate, childHash)); + + AggregateParam aggregateParam = new AggregateParam(AggPhase.GLOBAL, AggMode.INPUT_TO_RESULT); + Alias distinctSum = new Alias( + new AggregateExpression(new Sum(true, extra), aggregateParam), "distinct_sum"); + PhysicalHashAggregate mixedDistinctAggregate = new PhysicalHashAggregate<>( + ImmutableList.of(d1, k2), + ImmutableList.of(k2, d1, sumOutput(extra, aggregateParam), distinctSum), + aggregateParam, + true, + logicalProperties, + false, + groupPlan); + assertMappingLocalityCleared(deriveAggregateProperties(mixedDistinctAggregate, childHash)); + } finally { + connectContext.getSessionVariable().enableBucketedHashAgg = originalEnableBucketedHashAgg; + } + } + + @Test + void testPureDeduplicateAggregateDropsMappingFromNaturalChild() { + SlotReference k1 = new SlotReference("k1", IntegerType.INSTANCE); + SlotReference k2 = new SlotReference("k2", IntegerType.INSTANCE); + SlotReference d1 = new SlotReference("d1", IntegerType.INSTANCE); + PhysicalHashAggregate deduplicate = new PhysicalHashAggregate<>( + ImmutableList.of(d1, k2), + ImmutableList.of(d1, k2), + new AggregateParam(AggPhase.GLOBAL, AggMode.BUFFER_TO_RESULT), + true, + logicalProperties, + false, + groupPlan); + + assertMappingLocalityCleared(deriveAggregateProperties( + deduplicate, naturalHashWithMapping(k1, k2, d1))); + } + + @Test + void testDistinctAggregateDoesNotRestoreMappingAfterRedistribution() { + SlotReference k1 = new SlotReference("k1", IntegerType.INSTANCE); + SlotReference k2 = new SlotReference("k2", IntegerType.INSTANCE); + SlotReference d1 = new SlotReference("d1", IntegerType.INSTANCE); + AggregateParam aggregateParam = new AggregateParam(AggPhase.DISTINCT_GLOBAL, AggMode.INPUT_TO_RESULT); + PhysicalHashAggregate distinctAggregate = new PhysicalHashAggregate<>( + ImmutableList.of(k1, k2, d1), + ImmutableList.of(k1, k2, d1, new Alias( + new AggregateExpression(new MultiDistinctCount(d1), aggregateParam), "distinct_count")), + aggregateParam, + true, + logicalProperties, + false, + groupPlan); + PhysicalProperties redistributedChild = new PhysicalProperties( + new DistributionSpecHash(ImmutableList.of(d1.getExprId()), ShuffleType.REQUIRE)); + + ConnectContext connectContext = ConnectContext.get(); + boolean originalEnableBucketedHashAgg = connectContext.getSessionVariable().enableBucketedHashAgg; + try { + connectContext.getSessionVariable().enableBucketedHashAgg = false; + PhysicalProperties output = deriveAggregateProperties(distinctAggregate, redistributedChild); + Assertions.assertFalse(output.getNaturalDistributionMappingSpec().isPresent()); + } finally { + connectContext.getSessionVariable().enableBucketedHashAgg = originalEnableBucketedHashAgg; + } + } + + @Test + void testAggregateDropsMappingsForRepeat() { + SlotReference k1 = new SlotReference("k1", IntegerType.INSTANCE); + SlotReference k2 = new SlotReference("k2", IntegerType.INSTANCE); + SlotReference d1 = new SlotReference("d1", IntegerType.INSTANCE); + DistributionSpecHash childHash = naturalHashWithMapping(k1, k2, d1); + AggregateParam aggregateParam = new AggregateParam(AggPhase.GLOBAL, AggMode.BUFFER_TO_RESULT); + PhysicalHashAggregate repeatAggregate = new PhysicalHashAggregate<>( + ImmutableList.of(k1, k2, d1), + ImmutableList.of(k1, k2, d1, sumOutput(d1, aggregateParam)), + aggregateParam, + true, + logicalProperties, + true, + groupPlan); + + Assertions.assertTrue(deriveAggregateHash( + repeatAggregate, childHash).getDistributionMappings().isEmpty()); + } + + private DistributionSpecHash naturalHashWithMapping( + SlotReference k1, SlotReference k2, SlotReference d1) { + return new DistributionSpecHash( + ImmutableList.of(k1.getExprId(), k2.getExprId()), ShuffleType.NATURAL, + 1L, 2L, ImmutableSet.of(3L), + ImmutableList.of(new DistributionMapping( + "mapping_1", ImmutableList.of(d1.getExprId()), ImmutableList.of(0)))); + } + + private Alias sumOutput(SlotReference input, AggregateParam aggregateParam) { + return new Alias(new AggregateExpression(new Sum(input), aggregateParam, input), "sum_value"); + } + + private DistributionSpecHash deriveAggregateHash( + PhysicalHashAggregate aggregate, DistributionSpecHash childHash) { + return (DistributionSpecHash) deriveAggregateProperties(aggregate, childHash).getDistributionSpec(); + } + + private PhysicalProperties deriveAggregateProperties( + PhysicalHashAggregate aggregate, DistributionSpecHash childHash) { + return deriveAggregateProperties(aggregate, naturalProperties(childHash)); + } + + private PhysicalProperties deriveAggregateProperties( + PhysicalHashAggregate aggregate, PhysicalProperties childProperties) { + GroupExpression groupExpression = new GroupExpression(aggregate); + new Group(null, groupExpression, null); + ChildOutputPropertyDeriver deriver = new ChildOutputPropertyDeriver( + ImmutableList.of(childProperties)); + return deriver.getOutputProperties(null, groupExpression); + } + @Test void testAggregateWithoutGroupBy() { PhysicalHashAggregate aggregate = new PhysicalHashAggregate<>( @@ -924,6 +1308,40 @@ void testRepeatReturnChild() { Assertions.assertEquals(child, result); } + @Test + void testRepeatDropsNaturalMappingLocality() { + SlotReference k1 = new SlotReference( + new ExprId(1), "k1", TinyIntType.INSTANCE, true, ImmutableList.of()); + SlotReference d1 = new SlotReference( + new ExprId(2), "d1", TinyIntType.INSTANCE, true, ImmutableList.of()); + SlotReference groupingId = new SlotReference( + new ExprId(3), "grouping_id", BigIntType.INSTANCE, false, ImmutableList.of()); + PhysicalRepeat repeat = new PhysicalRepeat<>( + ImmutableList.of(ImmutableList.of(k1, d1), ImmutableList.of(k1)), + ImmutableList.of(k1, d1), + groupingId, + logicalProperties, + groupPlan); + GroupExpression groupExpression = new GroupExpression(repeat); + new Group(null, groupExpression, null); + DistributionSpecHash childHash = new DistributionSpecHash( + ImmutableList.of(k1.getExprId()), ShuffleType.NATURAL, + 1L, 2L, ImmutableSet.of(3L), + ImmutableList.of(new DistributionMapping( + "mapping_1", ImmutableList.of(d1.getExprId()), ImmutableList.of(0)))); + + PhysicalProperties result = new ChildOutputPropertyDeriver( + ImmutableList.of(new PhysicalProperties(childHash))) + .getOutputProperties(null, groupExpression); + + Assertions.assertInstanceOf(DistributionSpecHash.class, result.getDistributionSpec()); + Assertions.assertEquals(ShuffleType.NATURAL, + ((DistributionSpecHash) result.getDistributionSpec()).getShuffleType()); + Assertions.assertTrue(((DistributionSpecHash) result.getDistributionSpec()) + .getDistributionMappings().isEmpty()); + Assertions.assertFalse(result.getNaturalDistributionMappingSpec().isPresent()); + } + @Test void testRepeatReturnChild2() { SlotReference c1 = new SlotReference( @@ -973,6 +1391,62 @@ void testComputeProjectOutputProperties() { Assertions.assertEquals(hashC1, phyProp3); } + @Test + void testProjectOnlyPropagatesMappingThroughHashValuePreservingCast() { + SlotReference distributionKey = new SlotReference("k1", IntegerType.INSTANCE); + SlotReference determinant = new SlotReference("d1", new VarcharType(8)); + DistributionSpecHash childHash = new DistributionSpecHash( + ImmutableList.of(distributionKey.getExprId()), ShuffleType.NATURAL, + 1L, 2L, ImmutableSet.of(3L), + ImmutableList.of(new DistributionMapping( + "mapping_1", ImmutableList.of(determinant.getExprId()), ImmutableList.of(0)))); + PhysicalProperties childProperties = naturalProperties(childHash); + + Alias widened = new Alias(new Cast(determinant, new VarcharType(32)), "widened_d1"); + PhysicalProperties widenedProperties = ChildOutputPropertyDeriver.computeProjectOutputProperties( + ImmutableList.of(widened), childProperties); + Assertions.assertEquals(DistributionSpecStorageAny.INSTANCE, widenedProperties.getDistributionSpec()); + Assertions.assertTrue(widenedProperties.getNaturalDistributionMappingSpec().isPresent()); + Assertions.assertEquals( + ImmutableList.of(widened.getExprId()), + widenedProperties.getNaturalDistributionMappingSpec().get() + .getDistributionMappings().get(0).getDeterminantExprIds()); + + Alias narrowed = new Alias(new Cast(determinant, new VarcharType(4)), "narrowed_d1"); + PhysicalProperties narrowedProperties = ChildOutputPropertyDeriver.computeProjectOutputProperties( + ImmutableList.of(narrowed), childProperties); + Assertions.assertEquals(DistributionSpecStorageAny.INSTANCE, narrowedProperties.getDistributionSpec()); + Assertions.assertFalse(narrowedProperties.getNaturalDistributionMappingSpec().isPresent()); + } + + private PhysicalProperties naturalProperties(DistributionSpecHash hashSpec) { + return naturalProperties(hashSpec, new OrderSpec()); + } + + private PhysicalProperties naturalProperties(DistributionSpecHash hashSpec, OrderSpec orderSpec) { + NaturalDistributionMappingSpec mappingSpec = new NaturalDistributionMappingSpec( + hashSpec.getTableId(), hashSpec.getSelectedIndexId(), hashSpec.getPartitionIds(), + hashSpec.getOrderedShuffledColumns().size(), hashSpec.getExprIdToEquivalenceSet(), + hashSpec.getDistributionMappings()); + return new PhysicalProperties(hashSpec, orderSpec, Optional.of(mappingSpec)); + } + + private PhysicalProperties deriveOutputProperties( + AbstractPhysicalPlan plan, PhysicalProperties... childProperties) { + GroupExpression groupExpression = new GroupExpression(plan); + new Group(null, groupExpression, null); + return new ChildOutputPropertyDeriver(ImmutableList.copyOf(childProperties)) + .getOutputProperties(null, groupExpression); + } + + private void assertMappingLocalityCleared(PhysicalProperties properties) { + Assertions.assertFalse(properties.getNaturalDistributionMappingSpec().isPresent()); + if (properties.getDistributionSpec() instanceof DistributionSpecHash) { + Assertions.assertTrue(((DistributionSpecHash) properties.getDistributionSpec()) + .getDistributionMappings().isEmpty()); + } + } + @Test void testComputeUniformAfterRecomputeLogicalProperties() { // left child has a uniform slot, right child empty diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/properties/ChildrenPropertiesRegulatorTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/properties/ChildrenPropertiesRegulatorTest.java index 2bd445e60c612d..a95e3ae246ce32 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/properties/ChildrenPropertiesRegulatorTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/properties/ChildrenPropertiesRegulatorTest.java @@ -21,24 +21,33 @@ import org.apache.doris.nereids.CascadesContext; import org.apache.doris.nereids.cost.Cost; import org.apache.doris.nereids.cost.CostCalculator; +import org.apache.doris.nereids.hint.DistributeHint; import org.apache.doris.nereids.jobs.JobContext; import org.apache.doris.nereids.memo.Group; import org.apache.doris.nereids.memo.GroupExpression; +import org.apache.doris.nereids.properties.DistributionSpecHash.ShuffleType; +import org.apache.doris.nereids.trees.expressions.EqualTo; import org.apache.doris.nereids.trees.expressions.NamedExpression; import org.apache.doris.nereids.trees.expressions.SlotReference; import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateParam; import org.apache.doris.nereids.trees.plans.AggMode; import org.apache.doris.nereids.trees.plans.AggPhase; +import org.apache.doris.nereids.trees.plans.DistributeType; import org.apache.doris.nereids.trees.plans.GroupPlan; +import org.apache.doris.nereids.trees.plans.JoinType; import org.apache.doris.nereids.trees.plans.Plan; import org.apache.doris.nereids.trees.plans.physical.PhysicalDistribute; import org.apache.doris.nereids.trees.plans.physical.PhysicalFilter; import org.apache.doris.nereids.trees.plans.physical.PhysicalHashAggregate; +import org.apache.doris.nereids.trees.plans.physical.PhysicalHashJoin; import org.apache.doris.nereids.trees.plans.physical.PhysicalLimit; import org.apache.doris.nereids.trees.plans.physical.PhysicalProject; import org.apache.doris.nereids.types.IntegerType; +import org.apache.doris.nereids.util.ExpressionUtils; +import org.apache.doris.nereids.util.JoinUtils; import org.apache.doris.qe.ConnectContext; +import com.google.common.collect.ImmutableList; import com.google.common.collect.Lists; import com.google.common.collect.Maps; import com.google.common.collect.Sets; @@ -177,6 +186,132 @@ public void testSingleExecutionInstanceAllowsOnePhaseAggWithDistribute() { } } + @Test + public void testAggregateColocateMappingRequestRequiresNaturalMappingProof() { + SlotReference k1 = new SlotReference("k1", IntegerType.INSTANCE); + SlotReference k2 = new SlotReference("k2", IntegerType.INSTANCE); + SlotReference d1 = new SlotReference("d1", IntegerType.INSTANCE); + GroupPlan mockedGroupPlan = Mockito.mock(GroupPlan.class); + Mockito.when(mockedGroupPlan.getAllChildrenTypes()).thenReturn(new BitSet()); + PhysicalHashAggregate aggregate = new PhysicalHashAggregate<>( + ImmutableList.of(d1, k2), + ImmutableList.of(d1, k2), + new AggregateParam(AggPhase.GLOBAL, AggMode.BUFFER_TO_RESULT), + true, + null, + false, + mockedGroupPlan); + GroupExpression parent = new GroupExpression(aggregate); + PhysicalProperties required = PhysicalProperties.createHash( + ImmutableList.of(d1.getExprId(), k2.getExprId()), + ShuffleType.COLOCATE_MAPPING_REQUIRE); + PhysicalProperties incompleteRequired = PhysicalProperties.createHash( + ImmutableList.of(d1.getExprId()), + ShuffleType.COLOCATE_MAPPING_REQUIRE); + DistributionSpecHash naturalWithMapping = new DistributionSpecHash( + ImmutableList.of(k1.getExprId(), k2.getExprId()), ShuffleType.NATURAL, + 1L, 2L, Sets.newHashSet(3L), + ImmutableList.of(new DistributionMapping( + "mapping_1", ImmutableList.of(d1.getExprId()), ImmutableList.of(0)))); + DistributionSpecHash naturalWithoutMapping = new DistributionSpecHash( + ImmutableList.of(k1.getExprId(), k2.getExprId()), ShuffleType.NATURAL, + 1L, 2L, Sets.newHashSet(3L), ImmutableList.of()); + + Assertions.assertFalse(adjustAggregateProperties( + parent, naturalProperties(naturalWithMapping), required).isEmpty()); + Assertions.assertTrue(adjustAggregateProperties( + parent, naturalProperties(naturalWithMapping), incompleteRequired).isEmpty()); + Assertions.assertTrue(adjustAggregateProperties( + parent, new PhysicalProperties(naturalWithoutMapping), required).isEmpty()); + } + + private List> adjustAggregateProperties( + GroupExpression parent, PhysicalProperties childOutput, PhysicalProperties required) { + ChildrenPropertiesRegulator regulator = new ChildrenPropertiesRegulator( + parent, + ImmutableList.of(Mockito.mock(GroupExpression.class)), + ImmutableList.of(childOutput), + ImmutableList.of(required), + mockedJobContext); + return regulator.adjustChildrenProperties(); + } + + @Test + public void testRejectPartiallySatisfiedColocateMappingRequest() { + assertColocateMappingRequestRejected(Optional.empty(), Optional.of("mapping_1")); + } + + @Test + public void testRejectColocateMappingRequestWhenFinalProofFails() { + try (MockedStatic mockedJoinUtils = Mockito.mockStatic(JoinUtils.class)) { + mockedJoinUtils.when(() -> JoinUtils.couldColocateJoinByMapping( + Mockito.any(), Mockito.any(), Mockito.anyList())).thenReturn(false); + assertColocateMappingRequestRejected(Optional.of("mapping_1"), Optional.of("mapping_2")); + mockedJoinUtils.verify(() -> JoinUtils.couldColocateJoinByMapping( + Mockito.any(), Mockito.any(), Mockito.anyList())); + } + } + + private void assertColocateMappingRequestRejected( + Optional leftMappingId, Optional rightMappingId) { + SlotReference leftK1 = new SlotReference("left_k1", IntegerType.INSTANCE); + SlotReference leftK2 = new SlotReference("left_k2", IntegerType.INSTANCE); + SlotReference leftD1 = new SlotReference("left_d1", IntegerType.INSTANCE); + SlotReference rightK1 = new SlotReference("right_k1", IntegerType.INSTANCE); + SlotReference rightK2 = new SlotReference("right_k2", IntegerType.INSTANCE); + SlotReference rightD1 = new SlotReference("right_d1", IntegerType.INSTANCE); + GroupPlan leftPlan = Mockito.mock(GroupPlan.class); + GroupPlan rightPlan = Mockito.mock(GroupPlan.class); + Mockito.when(leftPlan.getAllChildrenTypes()).thenReturn(new BitSet()); + Mockito.when(rightPlan.getAllChildrenTypes()).thenReturn(new BitSet()); + PhysicalHashJoin join = new PhysicalHashJoin<>( + JoinType.INNER_JOIN, + ImmutableList.of(new EqualTo(leftD1, rightD1), new EqualTo(leftK2, rightK2)), + ExpressionUtils.EMPTY_CONDITION, + new DistributeHint(DistributeType.NONE), + Optional.empty(), + Mockito.mock(LogicalProperties.class), + leftPlan, + rightPlan); + GroupExpression parent = new GroupExpression(join); + + DistributionSpecHash leftOutput = new DistributionSpecHash( + ImmutableList.of(leftK1.getExprId(), leftK2.getExprId()), + ShuffleType.NATURAL, 1L, 1L, Sets.newHashSet(1L), + leftMappingId.map(mappingId -> ImmutableList.of(new DistributionMapping( + mappingId, ImmutableList.of(leftD1.getExprId()), ImmutableList.of(0)))) + .orElseGet(ImmutableList::of)); + DistributionSpecHash rightOutput = new DistributionSpecHash( + ImmutableList.of(rightK1.getExprId(), rightK2.getExprId()), + ShuffleType.NATURAL, 2L, 1L, Sets.newHashSet(1L), + rightMappingId.map(mappingId -> ImmutableList.of(new DistributionMapping( + mappingId, ImmutableList.of(rightD1.getExprId()), ImmutableList.of(0)))) + .orElseGet(ImmutableList::of)); + DistributionSpecHash leftRequired = new DistributionSpecHash( + ImmutableList.of(leftD1.getExprId(), leftK2.getExprId()), + ShuffleType.COLOCATE_MAPPING_REQUIRE); + DistributionSpecHash rightRequired = new DistributionSpecHash( + ImmutableList.of(rightD1.getExprId(), rightK2.getExprId()), + ShuffleType.COLOCATE_MAPPING_REQUIRE); + + ChildrenPropertiesRegulator regulator = new ChildrenPropertiesRegulator( + parent, + ImmutableList.of(Mockito.mock(GroupExpression.class), Mockito.mock(GroupExpression.class)), + ImmutableList.of(naturalProperties(leftOutput), naturalProperties(rightOutput)), + ImmutableList.of(new PhysicalProperties(leftRequired), new PhysicalProperties(rightRequired)), + mockedJobContext); + + Assertions.assertTrue(regulator.adjustChildrenProperties().isEmpty()); + } + + private PhysicalProperties naturalProperties(DistributionSpecHash hashSpec) { + NaturalDistributionMappingSpec mappingSpec = new NaturalDistributionMappingSpec( + hashSpec.getTableId(), hashSpec.getSelectedIndexId(), hashSpec.getPartitionIds(), + hashSpec.getOrderedShuffledColumns().size(), hashSpec.getExprIdToEquivalenceSet(), + hashSpec.getDistributionMappings()); + return new PhysicalProperties(hashSpec, new OrderSpec(), Optional.of(mappingSpec)); + } + private void testMustShuffleFilter(Class childClazz) { try (MockedStatic mockedCostCalculator = Mockito.mockStatic(CostCalculator.class)) { mockedCostCalculator.when(() -> CostCalculator.calculateCost(Mockito.any(), Mockito.any(), diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/properties/DistributionSpecHashTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/properties/DistributionSpecHashTest.java index da99ec15b6d624..180a7421c6abb6 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/properties/DistributionSpecHashTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/properties/DistributionSpecHashTest.java @@ -29,6 +29,7 @@ import org.junit.jupiter.api.Test; import java.util.Map; +import java.util.Optional; import java.util.Set; public class DistributionSpecHashTest { @@ -387,4 +388,59 @@ public void testHashEqualSatisfyWithDifferentLength() { Assertions.assertFalse(bucketed1.satisfy(bucketed2)); Assertions.assertFalse(bucketed2.satisfy(bucketed1)); } + + @Test + public void testColocateMappingRequireSatisfy() { + ExprId k1 = new ExprId(0); + ExprId k2 = new ExprId(1); + ExprId d1 = new ExprId(2); + ExprId d2 = new ExprId(3); + DistributionSpecHash natural = new DistributionSpecHash( + ImmutableList.of(k1, k2), + ShuffleType.NATURAL, + 1, + 1, + ImmutableSet.of(1L), + ImmutableList.of( + new DistributionMapping("mapping_1", ImmutableList.of(d1), ImmutableList.of(0)), + new DistributionMapping("mapping_2", ImmutableList.of(d2), ImmutableList.of(1)))); + NaturalDistributionMappingSpec mappingSpec = new NaturalDistributionMappingSpec( + natural.getTableId(), natural.getSelectedIndexId(), natural.getPartitionIds(), + natural.getOrderedShuffledColumns().size(), natural.getExprIdToEquivalenceSet(), + natural.getDistributionMappings()); + PhysicalProperties naturalProperties = new PhysicalProperties( + natural, new OrderSpec(), Optional.of(mappingSpec)); + + Assertions.assertTrue(naturalProperties.satisfy( + new PhysicalProperties(new DistributionSpecHash( + ImmutableList.of(d1, k2), ShuffleType.COLOCATE_MAPPING_REQUIRE)))); + Assertions.assertTrue(naturalProperties.satisfy( + new PhysicalProperties(new DistributionSpecHash( + ImmutableList.of(d1, d2), ShuffleType.COLOCATE_MAPPING_REQUIRE)))); + Assertions.assertFalse(naturalProperties.satisfy( + new PhysicalProperties(new DistributionSpecHash( + ImmutableList.of(d1), ShuffleType.COLOCATE_MAPPING_REQUIRE)))); + Assertions.assertFalse(natural.satisfy(new DistributionSpecHash( + ImmutableList.of(d1, k2), ShuffleType.COLOCATE_MAPPING_REQUIRE))); + Assertions.assertFalse(natural.satisfy(new DistributionSpecHash( + ImmutableList.of(d1, k2), ShuffleType.REQUIRE))); + + Map projections = Maps.newHashMap(); + projections.put(d1, d1); + projections.put(k2, k2); + DistributionSpecHash projected = (DistributionSpecHash) natural.project( + projections, ImmutableSet.of(), DistributionSpecAny.INSTANCE); + Assertions.assertTrue(projected.getDistributionMappings().isEmpty()); + } + + @Test + public void testColocateMappingRequireIsNotEnforceable() { + ExprId exprId = new ExprId(0); + + Assertions.assertFalse(new PhysicalProperties(new DistributionSpecHash( + ImmutableList.of(exprId), ShuffleType.COLOCATE_MAPPING_REQUIRE)).isEnforceable()); + Assertions.assertTrue(new PhysicalProperties(new DistributionSpecHash( + ImmutableList.of(exprId), ShuffleType.REQUIRE)).isEnforceable()); + Assertions.assertTrue(PhysicalProperties.GATHER.isEnforceable()); + } } diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/properties/RequestPropertyDeriverTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/properties/RequestPropertyDeriverTest.java index bffe00a0496dee..810dffda6b7539 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/properties/RequestPropertyDeriverTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/properties/RequestPropertyDeriverTest.java @@ -27,8 +27,11 @@ import org.apache.doris.nereids.memo.GroupId; import org.apache.doris.nereids.properties.DistributionSpecHash.ShuffleType; import org.apache.doris.nereids.rules.implementation.LogicalWindowToPhysicalWindow.WindowFrameGroup; +import org.apache.doris.nereids.trees.expressions.Add; +import org.apache.doris.nereids.trees.expressions.AggregateExpression; import org.apache.doris.nereids.trees.expressions.Alias; import org.apache.doris.nereids.trees.expressions.AssertNumRowsElement; +import org.apache.doris.nereids.trees.expressions.Cast; import org.apache.doris.nereids.trees.expressions.EqualTo; import org.apache.doris.nereids.trees.expressions.ExprId; import org.apache.doris.nereids.trees.expressions.Expression; @@ -39,6 +42,8 @@ import org.apache.doris.nereids.trees.expressions.WindowFrame.FrameBoundary; import org.apache.doris.nereids.trees.expressions.WindowFrame.FrameUnitsType; import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateParam; +import org.apache.doris.nereids.trees.expressions.functions.agg.MultiDistinctCount; +import org.apache.doris.nereids.trees.expressions.functions.agg.Sum; import org.apache.doris.nereids.trees.expressions.functions.window.RowNumber; import org.apache.doris.nereids.trees.expressions.literal.Literal; import org.apache.doris.nereids.trees.plans.AggMode; @@ -46,15 +51,22 @@ import org.apache.doris.nereids.trees.plans.DistributeType; import org.apache.doris.nereids.trees.plans.GroupPlan; import org.apache.doris.nereids.trees.plans.JoinType; +import org.apache.doris.nereids.trees.plans.Plan; import org.apache.doris.nereids.trees.plans.RelationId; +import org.apache.doris.nereids.trees.plans.algebra.SetOperation.Qualifier; import org.apache.doris.nereids.trees.plans.logical.LogicalOneRowRelation; import org.apache.doris.nereids.trees.plans.physical.PhysicalAssertNumRows; +import org.apache.doris.nereids.trees.plans.physical.PhysicalExcept; import org.apache.doris.nereids.trees.plans.physical.PhysicalExternalRowLevelMergeSink; import org.apache.doris.nereids.trees.plans.physical.PhysicalHashAggregate; import org.apache.doris.nereids.trees.plans.physical.PhysicalHashJoin; +import org.apache.doris.nereids.trees.plans.physical.PhysicalIntersect; import org.apache.doris.nereids.trees.plans.physical.PhysicalNestedLoopJoin; +import org.apache.doris.nereids.trees.plans.physical.PhysicalSetOperation; +import org.apache.doris.nereids.trees.plans.physical.PhysicalUnion; import org.apache.doris.nereids.trees.plans.physical.PhysicalWindow; import org.apache.doris.nereids.types.IntegerType; +import org.apache.doris.nereids.types.VarcharType; import org.apache.doris.nereids.util.ExpressionUtils; import org.apache.doris.nereids.util.MemoTestUtils; import org.apache.doris.qe.ConnectContext; @@ -211,9 +223,10 @@ void testShuffleOrBroadcastHashJoin() { GroupExpression groupExpression = new GroupExpression(join, Lists.newArrayList(leftGroup, rightGroup)); new Group(null, groupExpression, null); - RequestPropertyDeriver requestPropertyDeriver = new RequestPropertyDeriver(null, jobContext); - List> actual - = requestPropertyDeriver.getRequestChildrenPropertyList(groupExpression); + RequestPropertyDeriver requestPropertyDeriver = + new RequestPropertyDeriver(testConnectContext, jobContext); + List> actual = + requestPropertyDeriver.getRequestChildrenPropertyList(groupExpression); List> expected = Lists.newArrayList(); expected.add(Lists.newArrayList( @@ -224,6 +237,37 @@ void testShuffleOrBroadcastHashJoin() { )); expected.add(Lists.newArrayList(PhysicalProperties.ANY, PhysicalProperties.REPLICATED)); Assertions.assertEquals(expected, actual); + + sessionVariable.enableColocateMappingConstraint = true; + List> enabled = + requestPropertyDeriver.getRequestChildrenPropertyList(groupExpression); + expected.add(1, Lists.newArrayList( + new PhysicalProperties(new DistributionSpecHash( + Lists.newArrayList(leftKey.getExprId()), + ShuffleType.COLOCATE_MAPPING_REQUIRE)), + new PhysicalProperties(new DistributionSpecHash( + Lists.newArrayList(rightKey.getExprId()), + ShuffleType.COLOCATE_MAPPING_REQUIRE)))); + Assertions.assertEquals(expected, enabled); + + PhysicalHashJoin expressionJoin = new PhysicalHashJoin<>( + JoinType.INNER_JOIN, + ImmutableList.of(new EqualTo(leftKey, new Add(rightKey, Literal.of(1)))), + ExpressionUtils.EMPTY_CONDITION, + new DistributeHint(DistributeType.NONE), + Optional.empty(), + logicalProperties, + leftPlan, + rightPlan); + GroupExpression expressionJoinGroup = + new GroupExpression(expressionJoin, Lists.newArrayList(leftGroup, rightGroup)); + new Group(null, expressionJoinGroup, null); + + List> expressionJoinRequests = + requestPropertyDeriver.getRequestChildrenPropertyList(expressionJoinGroup); + Assertions.assertEquals( + Lists.newArrayList(expected.get(0), expected.get(2)), + expressionJoinRequests); } } @@ -274,6 +318,260 @@ void testGlobalAggregate() { Assertions.assertEquals(expected, actual); } + @Test + void testGlobalAggregatePropagatesColocateMappingRequestWhenEnabled() { + ConnectContext testConnectContext = MemoTestUtils.createConnectContext(); + SlotReference d1 = new SlotReference("d1", IntegerType.INSTANCE); + SlotReference k2 = new SlotReference("k2", IntegerType.INSTANCE); + SlotReference extra = new SlotReference("extra", IntegerType.INSTANCE); + Alias outputD1 = new Alias(d1, "output_d1"); + Alias outputK2 = new Alias(k2, "output_k2"); + AggregateParam aggregateParam = new AggregateParam(AggPhase.GLOBAL, AggMode.BUFFER_TO_RESULT); + Alias sum = new Alias(new AggregateExpression(new Sum(extra), aggregateParam), "sum_value"); + PhysicalHashAggregate aggregate = new PhysicalHashAggregate<>( + Lists.newArrayList(d1, k2), + Lists.newArrayList(outputD1, outputK2, sum), + Optional.of(Lists.newArrayList(d1, k2)), + aggregateParam, + true, + logicalProperties, + false, + groupPlan); + GroupExpression groupExpression = new GroupExpression(aggregate); + new Group(null, groupExpression, null); + PhysicalProperties parentProperties = PhysicalProperties.createHash( + Lists.newArrayList(outputD1.getExprId(), outputK2.getExprId(), sum.getExprId()), + ShuffleType.COLOCATE_MAPPING_REQUIRE); + PhysicalProperties mappingRequest = PhysicalProperties.createHash( + Lists.newArrayList(d1.getExprId(), k2.getExprId()), + ShuffleType.COLOCATE_MAPPING_REQUIRE); + PhysicalProperties originalRequest = PhysicalProperties.createHash( + Lists.newArrayList(d1.getExprId(), k2.getExprId()), ShuffleType.REQUIRE); + + testConnectContext.getSessionVariable().enableColocateMappingConstraint = true; + List> enabled = new RequestPropertyDeriver( + testConnectContext, parentProperties).getRequestChildrenPropertyList(groupExpression); + Assertions.assertEquals(ImmutableList.of( + ImmutableList.of(mappingRequest), ImmutableList.of(originalRequest)), enabled); + + testConnectContext.getSessionVariable().enableColocateMappingConstraint = false; + List> disabled = new RequestPropertyDeriver( + testConnectContext, parentProperties).getRequestChildrenPropertyList(groupExpression); + Assertions.assertEquals(ImmutableList.of(ImmutableList.of(originalRequest)), disabled); + + testConnectContext.getSessionVariable().enableColocateMappingConstraint = true; + PhysicalProperties partiallyMappableParent = PhysicalProperties.createHash( + Lists.newArrayList(outputD1.getExprId(), sum.getExprId()), + ShuffleType.COLOCATE_MAPPING_REQUIRE); + PhysicalProperties partialMappingRequest = PhysicalProperties.createHash( + Lists.newArrayList(d1.getExprId()), ShuffleType.COLOCATE_MAPPING_REQUIRE); + List> partiallyMappableRequests = new RequestPropertyDeriver( + testConnectContext, partiallyMappableParent).getRequestChildrenPropertyList(groupExpression); + Assertions.assertEquals(ImmutableList.of( + ImmutableList.of(partialMappingRequest), ImmutableList.of(originalRequest)), + partiallyMappableRequests); + + PhysicalHashAggregate expressionGroupByAggregate = new PhysicalHashAggregate<>( + Lists.newArrayList(new EqualTo(d1, k2), k2), + Lists.newArrayList(outputD1, outputK2, sum), + Optional.of(Lists.newArrayList(d1, k2)), + aggregateParam, + true, + logicalProperties, + false, + groupPlan); + GroupExpression expressionGroupBy = new GroupExpression(expressionGroupByAggregate); + new Group(null, expressionGroupBy, null); + testConnectContext.getSessionVariable().enableColocateMappingConstraint = true; + List> expressionGroupByRequests = new RequestPropertyDeriver( + testConnectContext, parentProperties).getRequestChildrenPropertyList(expressionGroupBy); + Assertions.assertEquals(ImmutableList.of(ImmutableList.of(originalRequest)), expressionGroupByRequests); + } + + @Test + void testDistinctAndDeduplicateAggregatesDoNotPropagateColocateMappingRequest() { + ConnectContext testConnectContext = MemoTestUtils.createConnectContext(); + testConnectContext.getSessionVariable().enableColocateMappingConstraint = true; + SlotReference d1 = new SlotReference("d1", IntegerType.INSTANCE); + SlotReference k2 = new SlotReference("k2", IntegerType.INSTANCE); + SlotReference extra = new SlotReference("extra", IntegerType.INSTANCE); + Alias outputD1 = new Alias(d1, "output_d1"); + Alias outputK2 = new Alias(k2, "output_k2"); + AggregateParam aggregateParam = new AggregateParam(AggPhase.GLOBAL, AggMode.BUFFER_TO_RESULT); + Alias distinctCount = new Alias( + new AggregateExpression(new MultiDistinctCount(extra), aggregateParam), "distinct_count"); + Alias distinctSum = new Alias( + new AggregateExpression(new Sum(true, extra), aggregateParam), "distinct_sum"); + Alias sum = new Alias(new AggregateExpression(new Sum(extra), aggregateParam), "sum_value"); + PhysicalProperties parentProperties = PhysicalProperties.createHash( + Lists.newArrayList(outputD1.getExprId(), outputK2.getExprId()), + ShuffleType.COLOCATE_MAPPING_REQUIRE); + PhysicalProperties originalRequest = PhysicalProperties.createHash( + Lists.newArrayList(d1.getExprId(), k2.getExprId()), ShuffleType.REQUIRE); + + PhysicalHashAggregate distinctAggregate = new PhysicalHashAggregate<>( + Lists.newArrayList(d1, k2), + Lists.newArrayList(outputD1, outputK2, distinctCount), + Optional.of(Lists.newArrayList(d1, k2)), + aggregateParam, + true, + logicalProperties, + false, + groupPlan); + + PhysicalHashAggregate mixedDistinctAggregate = new PhysicalHashAggregate<>( + Lists.newArrayList(d1, k2), + Lists.newArrayList(outputD1, outputK2, sum, distinctSum), + Optional.of(Lists.newArrayList(d1, k2)), + aggregateParam, + true, + logicalProperties, + false, + groupPlan); + + AggregateParam distinctPhaseParam = new AggregateParam( + AggPhase.DISTINCT_GLOBAL, AggMode.BUFFER_TO_RESULT); + Alias distinctPhaseSum = new Alias( + new AggregateExpression(new Sum(extra), distinctPhaseParam), "distinct_phase_sum"); + PhysicalHashAggregate distinctPhaseAggregate = new PhysicalHashAggregate<>( + Lists.newArrayList(d1, k2), + Lists.newArrayList(outputD1, outputK2, distinctPhaseSum), + Optional.of(Lists.newArrayList(d1, k2)), + distinctPhaseParam, + true, + logicalProperties, + false, + groupPlan); + + PhysicalHashAggregate deduplicateAggregate = new PhysicalHashAggregate<>( + Lists.newArrayList(d1, k2), + Lists.newArrayList(outputD1, outputK2), + Optional.of(Lists.newArrayList(d1, k2)), + aggregateParam, + true, + logicalProperties, + false, + groupPlan); + + for (PhysicalHashAggregate barrier + : ImmutableList.of(distinctAggregate, mixedDistinctAggregate, + distinctPhaseAggregate, deduplicateAggregate)) { + GroupExpression groupExpression = new GroupExpression(barrier); + new Group(null, groupExpression, null); + Assertions.assertEquals(ImmutableList.of(ImmutableList.of(originalRequest)), + new RequestPropertyDeriver(testConnectContext, parentProperties) + .getRequestChildrenPropertyList(groupExpression)); + } + } + + @Test + void testAggregateRemapsColocateMappingRequestThroughWideningVarcharCast() { + ConnectContext testConnectContext = MemoTestUtils.createConnectContext(); + testConnectContext.getSessionVariable().enableColocateMappingConstraint = true; + SlotReference determinant = new SlotReference("determinant", new VarcharType(8)); + SlotReference value = new SlotReference("value", IntegerType.INSTANCE); + Alias widenedOutput = + new Alias(new Cast(determinant, new VarcharType(32)), "widened_determinant"); + AggregateParam aggregateParam = + new AggregateParam(AggPhase.GLOBAL, AggMode.BUFFER_TO_RESULT); + Alias sum = new Alias(new AggregateExpression(new Sum(value), aggregateParam), "sum_value"); + PhysicalHashAggregate aggregate = new PhysicalHashAggregate<>( + Lists.newArrayList(determinant), + Lists.newArrayList(widenedOutput, sum), + Optional.of(Lists.newArrayList(determinant)), + aggregateParam, + true, + logicalProperties, + false, + groupPlan); + GroupExpression groupExpression = new GroupExpression(aggregate); + new Group(null, groupExpression, null); + PhysicalProperties parentProperties = PhysicalProperties.createHash( + Lists.newArrayList(widenedOutput.getExprId()), + ShuffleType.COLOCATE_MAPPING_REQUIRE); + + List> requests = new RequestPropertyDeriver( + testConnectContext, parentProperties) + .getRequestChildrenPropertyList(groupExpression); + + Assertions.assertEquals( + PhysicalProperties.createHash( + Lists.newArrayList(determinant.getExprId()), + ShuffleType.COLOCATE_MAPPING_REQUIRE), + requests.get(0).get(0)); + } + + @Test + void testUnionDoesNotPropagateColocateMappingRequest() { + SlotReference outputD1 = new SlotReference("output_d1", IntegerType.INSTANCE); + SlotReference outputK2 = new SlotReference("output_k2", IntegerType.INSTANCE); + SlotReference leftD1 = new SlotReference("left_d1", IntegerType.INSTANCE); + SlotReference leftK2 = new SlotReference("left_k2", IntegerType.INSTANCE); + SlotReference rightD1 = new SlotReference("right_d1", IntegerType.INSTANCE); + SlotReference rightK2 = new SlotReference("right_k2", IntegerType.INSTANCE); + PhysicalUnion union = new PhysicalUnion( + Qualifier.ALL, + ImmutableList.of(outputD1, outputK2), + ImmutableList.of( + ImmutableList.of(leftD1, leftK2), + ImmutableList.of(rightD1, rightK2)), + ImmutableList.of(), + logicalProperties, + ImmutableList.of(groupPlan, groupPlan)); + GroupExpression groupExpression = new GroupExpression(union, Lists.newArrayList(group, group)); + new Group(null, groupExpression, null); + PhysicalProperties parentProperties = PhysicalProperties.createHash( + ImmutableList.of(outputD1.getExprId(), outputK2.getExprId()), + ShuffleType.COLOCATE_MAPPING_REQUIRE); + + List> actual = new RequestPropertyDeriver( + MemoTestUtils.createConnectContext(), parentProperties) + .getRequestChildrenPropertyList(groupExpression); + + Assertions.assertEquals( + ImmutableList.of(ImmutableList.of(PhysicalProperties.ANY, PhysicalProperties.ANY)), + actual); + } + + @Test + void testIntersectAndExceptDoNotPropagateColocateMappingRequest() { + SlotReference outputD1 = new SlotReference("output_d1", IntegerType.INSTANCE); + SlotReference outputK2 = new SlotReference("output_k2", IntegerType.INSTANCE); + SlotReference leftD1 = new SlotReference("left_d1", IntegerType.INSTANCE); + SlotReference leftK2 = new SlotReference("left_k2", IntegerType.INSTANCE); + SlotReference rightD1 = new SlotReference("right_d1", IntegerType.INSTANCE); + SlotReference rightK2 = new SlotReference("right_k2", IntegerType.INSTANCE); + List> childrenOutputs = ImmutableList.of( + ImmutableList.of(leftD1, leftK2), + ImmutableList.of(rightD1, rightK2)); + List children = ImmutableList.of(groupPlan, groupPlan); + List setOperations = ImmutableList.of( + new PhysicalIntersect(Qualifier.DISTINCT, + ImmutableList.of(outputD1, outputK2), childrenOutputs, logicalProperties, children), + new PhysicalExcept(Qualifier.DISTINCT, + ImmutableList.of(outputD1, outputK2), childrenOutputs, logicalProperties, children)); + PhysicalProperties parentProperties = PhysicalProperties.createHash( + ImmutableList.of(outputD1.getExprId(), outputK2.getExprId()), + ShuffleType.COLOCATE_MAPPING_REQUIRE); + + for (PhysicalSetOperation setOperation : setOperations) { + GroupExpression groupExpression = + new GroupExpression(setOperation, Lists.newArrayList(group, group)); + new Group(null, groupExpression, null); + List> actual = new RequestPropertyDeriver( + MemoTestUtils.createConnectContext(), parentProperties) + .getRequestChildrenPropertyList(groupExpression); + + Assertions.assertEquals(1, actual.size()); + Assertions.assertEquals(2, actual.get(0).size()); + for (PhysicalProperties childRequest : actual.get(0)) { + Assertions.assertInstanceOf(DistributionSpecHash.class, childRequest.getDistributionSpec()); + Assertions.assertNotEquals(ShuffleType.COLOCATE_MAPPING_REQUIRE, + ((DistributionSpecHash) childRequest.getDistributionSpec()).getShuffleType()); + } + } + } + @Test void testGlobalAggregateWithoutPartition() { SlotReference key = new SlotReference("col1", IntegerType.INSTANCE); diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/implementation/LogicalOlapScanToPhysicalOlapScanTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/implementation/LogicalOlapScanToPhysicalOlapScanTest.java new file mode 100644 index 00000000000000..d9287b8dbb4d22 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/implementation/LogicalOlapScanToPhysicalOlapScanTest.java @@ -0,0 +1,176 @@ +// 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.rules.implementation; + +import org.apache.doris.catalog.Column; +import org.apache.doris.catalog.Env; +import org.apache.doris.catalog.HashDistributionInfo; +import org.apache.doris.catalog.OlapTable; +import org.apache.doris.catalog.RandomDistributionInfo; +import org.apache.doris.catalog.constraint.ConstraintManager; +import org.apache.doris.catalog.constraint.DistributionMappingConstraint; +import org.apache.doris.nereids.SqlCacheContext; +import org.apache.doris.nereids.StatementContext; +import org.apache.doris.nereids.properties.DistributionMapping; +import org.apache.doris.nereids.properties.DistributionSpecStorageAny; +import org.apache.doris.nereids.trees.expressions.ExprId; +import org.apache.doris.nereids.trees.expressions.Slot; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan; +import org.apache.doris.nereids.util.Utils; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.qe.SessionVariable; + +import com.google.common.collect.ImmutableList; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; +import org.mockito.Mockito; + +import java.util.List; +import java.util.Optional; + +class LogicalOlapScanToPhysicalOlapScanTest { + @Test + void buildDistributionMappingsUsesBaseColumnProvenance() { + ExprId aliasExprId = new ExprId(1); + SlotReference aliasSlot = Mockito.mock(SlotReference.class); + Column aliasColumn = Mockito.mock(Column.class); + Mockito.when(aliasSlot.getExprId()).thenReturn(aliasExprId); + Mockito.when(aliasSlot.getOriginalColumn()).thenReturn(Optional.of(aliasColumn)); + Mockito.when(aliasColumn.getName()).thenReturn("alias_d1"); + Mockito.when(aliasColumn.tryGetBaseColumnName()).thenReturn("d1"); + + Column distributionColumn = Mockito.mock(Column.class); + Mockito.when(distributionColumn.getName()).thenReturn("k1"); + HashDistributionInfo distributionInfo = Mockito.mock(HashDistributionInfo.class); + Mockito.when(distributionInfo.getDistributionColumns()) + .thenReturn(ImmutableList.of(distributionColumn)); + + DistributionMappingConstraint mapping = new DistributionMappingConstraint( + "mapping", "mapping_id", ImmutableList.of("d1"), ImmutableList.of("k1")); + List mappings = + LogicalOlapScanToPhysicalOlapScan.buildDistributionMappings( + distributionInfo, ImmutableList.of(aliasSlot), ImmutableList.of(mapping)); + + Assertions.assertEquals(1, mappings.size()); + Assertions.assertEquals(ImmutableList.of(aliasExprId), mappings.get(0).getDeterminantExprIds()); + Assertions.assertEquals(ImmutableList.of(0), mappings.get(0).getTargetDistributionIndices()); + } + + @Test + void buildDistributionMappingsRejectsMissingDeterminantOrTarget() { + SlotReference slot = Mockito.mock(SlotReference.class); + Column visibleColumn = Mockito.mock(Column.class); + Mockito.when(slot.getExprId()).thenReturn(new ExprId(1)); + Mockito.when(slot.getOriginalColumn()).thenReturn(Optional.of(visibleColumn)); + Mockito.when(visibleColumn.tryGetBaseColumnName()).thenReturn("extra_col"); + + Column distributionColumn = Mockito.mock(Column.class); + Mockito.when(distributionColumn.getName()).thenReturn("k1"); + HashDistributionInfo distributionInfo = Mockito.mock(HashDistributionInfo.class); + Mockito.when(distributionInfo.getDistributionColumns()) + .thenReturn(ImmutableList.of(distributionColumn)); + + DistributionMappingConstraint missingDeterminant = new DistributionMappingConstraint( + "missing_determinant", "mapping_id", ImmutableList.of("d1"), ImmutableList.of("k1")); + DistributionMappingConstraint missingTarget = new DistributionMappingConstraint( + "missing_target", "mapping_id", ImmutableList.of("extra_col"), ImmutableList.of("k2")); + + Assertions.assertTrue(LogicalOlapScanToPhysicalOlapScan.buildDistributionMappings( + distributionInfo, + ImmutableList.of(slot), + ImmutableList.of(missingDeterminant, missingTarget)).isEmpty()); + } + + @Test + void mappingScanDisablesSqlResultCache() { + ExprId determinantExprId = new ExprId(1); + SlotReference determinantSlot = Mockito.mock(SlotReference.class); + Column determinantColumn = Mockito.mock(Column.class); + Mockito.when(determinantSlot.getExprId()).thenReturn(determinantExprId); + Mockito.when(determinantSlot.getOriginalColumn()).thenReturn(Optional.of(determinantColumn)); + Mockito.when(determinantColumn.tryGetBaseColumnName()).thenReturn("d1"); + + Column distributionColumn = Mockito.mock(Column.class); + Mockito.when(distributionColumn.getName()).thenReturn("k1"); + HashDistributionInfo distributionInfo = Mockito.mock(HashDistributionInfo.class); + Mockito.when(distributionInfo.getDistributionColumns()) + .thenReturn(ImmutableList.of(distributionColumn)); + + OlapTable table = Mockito.mock(OlapTable.class); + ConstraintManager constraintManager = Mockito.mock(ConstraintManager.class); + Mockito.when(constraintManager.getDistributionMappingConstraintsForPlanning(table)) + .thenReturn(ImmutableList.of(new DistributionMappingConstraint( + "mapping", "mapping_id", ImmutableList.of("d1"), ImmutableList.of("k1")))); + Env env = Mockito.mock(Env.class); + Mockito.when(env.getConstraintManager()).thenReturn(constraintManager); + + SqlCacheContext sqlCacheContext = Mockito.mock(SqlCacheContext.class); + StatementContext statementContext = Mockito.mock(StatementContext.class); + Mockito.when(statementContext.getSqlCacheContext()).thenReturn(Optional.of(sqlCacheContext)); + SessionVariable sessionVariable = Mockito.mock(SessionVariable.class); + Mockito.when(sessionVariable.isEnableColocateMappingConstraint()).thenReturn(true); + ConnectContext connectContext = Mockito.mock(ConnectContext.class); + Mockito.when(connectContext.getSessionVariable()).thenReturn(sessionVariable); + Mockito.when(connectContext.getStatementContext()).thenReturn(statementContext); + + try (MockedStatic mockedContext = Mockito.mockStatic(ConnectContext.class); + MockedStatic mockedEnv = Mockito.mockStatic(Env.class)) { + mockedContext.when(ConnectContext::get).thenReturn(connectContext); + mockedEnv.when(Env::getCurrentEnv).thenReturn(env); + + List mappings = + LogicalOlapScanToPhysicalOlapScan.buildDistributionMappings( + table, distributionInfo, ImmutableList.of(determinantSlot)); + + Assertions.assertEquals(1, mappings.size()); + Mockito.verify(sqlCacheContext).setHasUnsupportedTables(true); + } + } + + @Test + void randomDistributionWithUnusableMappingFallsBackToRegularPlanning() { + OlapTable table = Mockito.mock(OlapTable.class); + Mockito.when(table.getDefaultDistributionInfo()).thenReturn(new RandomDistributionInfo(4)); + LogicalOlapScan scan = Mockito.mock(LogicalOlapScan.class); + Mockito.when(scan.getTable()).thenReturn(table); + Mockito.when(scan.getSelectedPartitionIds()).thenReturn(ImmutableList.of()); + + ConstraintManager constraintManager = Mockito.mock(ConstraintManager.class); + Mockito.when(constraintManager.getDistributionMappingConstraintsForPlanning(table)) + .thenReturn(ImmutableList.of()); + Env env = Mockito.mock(Env.class); + Mockito.when(env.getConstraintManager()).thenReturn(constraintManager); + SessionVariable sessionVariable = Mockito.mock(SessionVariable.class); + Mockito.when(sessionVariable.isEnableColocateMappingConstraint()).thenReturn(true); + ConnectContext connectContext = Mockito.mock(ConnectContext.class); + Mockito.when(connectContext.getSessionVariable()).thenReturn(sessionVariable); + + try (MockedStatic mockedContext = Mockito.mockStatic(ConnectContext.class); + MockedStatic mockedEnv = Mockito.mockStatic(Env.class); + MockedStatic mockedUtils = Mockito.mockStatic(Utils.class)) { + mockedContext.when(ConnectContext::get).thenReturn(connectContext); + mockedEnv.when(Env::getCurrentEnv).thenReturn(env); + + Assertions.assertSame(DistributionSpecStorageAny.INSTANCE, + LogicalOlapScanToPhysicalOlapScan.convertDistribution(scan)); + Mockito.verify(constraintManager).getDistributionMappingConstraintsForPlanning(table); + } + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/ConstraintTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/ConstraintTest.java index 218ce37f0e0ceb..244c584a5788c4 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/ConstraintTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/ConstraintTest.java @@ -19,9 +19,11 @@ import org.apache.doris.catalog.Column; import org.apache.doris.catalog.Env; +import org.apache.doris.catalog.OlapTable; import org.apache.doris.catalog.TableIf; import org.apache.doris.catalog.constraint.Constraint; import org.apache.doris.catalog.constraint.ConstraintManager; +import org.apache.doris.catalog.constraint.DistributionMappingConstraint; import org.apache.doris.catalog.constraint.ForeignKeyConstraint; import org.apache.doris.catalog.constraint.PrimaryKeyConstraint; import org.apache.doris.catalog.constraint.UniqueConstraint; @@ -32,6 +34,7 @@ import org.apache.doris.nereids.trees.plans.commands.DropConstraintCommand; import org.apache.doris.nereids.util.PlanChecker; import org.apache.doris.nereids.util.PlanPatternMatchSupported; +import org.apache.doris.persist.TableRenameColumnInfo; import org.apache.doris.qe.GlobalVariable; import org.apache.doris.utframe.TestWithFeService; @@ -39,6 +42,8 @@ import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import java.util.HashMap; +import java.util.Map; import java.util.Set; class ConstraintTest extends TestWithFeService implements PlanPatternMatchSupported { @@ -135,6 +140,301 @@ void uniqueConstraintTest() throws Exception { .getConstraints(tableNameInfoOf(o.getTable())).isEmpty())); } + @Test + void distributionMappingConstraintTest() throws Exception { + createTable("create table mapping_basic (\n" + + " k1 int,\n" + + " k2 int\n" + + ")\n" + + "duplicate key(k1)\n" + + "distributed by hash(k1) buckets 4\n" + + "properties(\"replication_num\"=\"1\")"); + try { + Exception duplicateDeterminant = Assertions.assertThrows(Exception.class, () -> addConstraint( + "alter table mapping_basic add constraint duplicate_determinant " + + "colocate mapping mapping_id (k2, K2) " + + "determines distribution key (k1) not enforced")); + Assertions.assertTrue(duplicateDeterminant.getMessage().contains( + "Determinant columns in distribution mapping constraint must be unique")); + + Exception invalidDistribution = Assertions.assertThrows(Exception.class, () -> addConstraint( + "alter table mapping_basic add constraint invalid_distribution " + + "colocate mapping mapping_id (k2) " + + "determines distribution key (k2) not enforced")); + Assertions.assertTrue(invalidDistribution.getMessage().contains( + "must be an ordered subset of table distribution columns")); + + addConstraint("alter table mapping_basic add constraint mapping_constraint " + + "colocate mapping mapping_id (k2) determines distribution key (k1) not enforced"); + OlapTable table = (OlapTable) Env.getCurrentInternalCatalog() + .getDbOrDdlException("test").getTableOrDdlException("mapping_basic"); + TableNameInfo tableNameInfo = tableNameInfoOf(table); + + Assertions.assertNull(getConstraintMgr().getConstraint(tableNameInfo, "mapping_constraint")); + Constraint constraint = getConstraintMgr().getConstraint( + tableNameInfo, table, "mapping_constraint"); + Assertions.assertInstanceOf(DistributionMappingConstraint.class, constraint); + DistributionMappingConstraint mapping = (DistributionMappingConstraint) constraint; + Assertions.assertEquals("mapping_id", mapping.getMappingId()); + Assertions.assertEquals(java.util.List.of("k2"), mapping.getDeterminantColumnNames()); + Assertions.assertEquals(java.util.List.of("k1"), mapping.getDistributionColumnNames()); + Assertions.assertEquals(table.getBaseSchemaVersion(), mapping.getBaseSchemaVersion()); + Assertions.assertEquals(java.util.List.of(table.getColumn("k2").getUniqueId()), + mapping.getDeterminantColumnUniqueIds()); + Assertions.assertEquals(java.util.List.of(table.getColumn("k1").getUniqueId()), + mapping.getDistributionColumnUniqueIds()); + + Exception dropColumn = Assertions.assertThrows(Exception.class, + () -> executeSql("alter table mapping_basic drop column k2")); + Assertions.assertTrue(dropColumn.getMessage().contains("mapping_constraint")); + Exception renameColumn = Assertions.assertThrows(Exception.class, + () -> executeSql("alter table mapping_basic rename column k2 k3")); + Assertions.assertTrue(renameColumn.getMessage().contains("mapping_constraint")); + Exception modifyColumn = Assertions.assertThrows(Exception.class, + () -> executeSql("alter table mapping_basic modify column k2 bigint")); + Assertions.assertTrue(modifyColumn.getMessage().contains("mapping_constraint")); + + dropConstraint("alter table mapping_basic drop constraint mapping_constraint"); + Assertions.assertTrue(getConstraintMgr().getDistributionMappingConstraints(table).isEmpty()); + } finally { + executeSql("drop table if exists mapping_basic force"); + } + } + + @Test + void distributionMappingDropRejectedDuringAtomicRestore() throws Exception { + createTable("create table mapping_atomic_restore (k1 int, k2 int) " + + "duplicate key(k1) distributed by hash(k1) buckets 4 " + + "properties(\"replication_num\"=\"1\")"); + OlapTable table = (OlapTable) Env.getCurrentInternalCatalog() + .getDbOrDdlException("test").getTableOrDdlException("mapping_atomic_restore"); + try { + addConstraint("alter table mapping_atomic_restore add constraint mapping " + + "colocate mapping mapping_id (k2) determines distribution key (k1) not enforced"); + table.writeLock(); + try { + table.setInAtomicRestore(); + } finally { + table.writeUnlock(); + } + + Exception exception = Assertions.assertThrows(Exception.class, + () -> dropConstraint("alter table mapping_atomic_restore drop constraint mapping")); + Assertions.assertTrue(exception.getMessage().contains("atomic restore state")); + Assertions.assertEquals(1, getConstraintMgr().getDistributionMappingConstraints(table).size()); + + table.writeLock(); + try { + table.clearInAtomicRestore(); + } finally { + table.writeUnlock(); + } + dropConstraint("alter table mapping_atomic_restore drop constraint mapping"); + Assertions.assertTrue(getConstraintMgr().getDistributionMappingConstraints(table).isEmpty()); + } finally { + table.writeLock(); + try { + table.clearInAtomicRestore(); + } finally { + table.writeUnlock(); + } + executeSql("drop table if exists mapping_atomic_restore force"); + } + } + + @Test + void distributionMappingFallsBackAfterOldFrontendSchemaReplay() throws Exception { + createTable("create table mapping_schema_binding (k1 int, k2 int) " + + "duplicate key(k1) distributed by hash(k1) buckets 4 " + + "properties(\"replication_num\"=\"1\", \"light_schema_change\"=\"true\")"); + try { + addConstraint("alter table mapping_schema_binding add constraint mapping " + + "colocate mapping mapping_id (k2) determines distribution key (k1) not enforced"); + OlapTable table = (OlapTable) Env.getCurrentInternalCatalog() + .getDbOrDdlException("test").getTableOrDdlException("mapping_schema_binding"); + + executeSql("alter table mapping_schema_binding add column k3 int"); + Assertions.assertEquals(1, + getConstraintMgr().getDistributionMappingConstraintsForPlanning(table).size()); + + Map schemaVersions = new HashMap<>(); + table.getIndexIdToMeta().forEach((indexId, indexMeta) -> + schemaVersions.put(indexId, indexMeta.getSchemaVersion() + 1)); + Env.getCurrentEnv().replayRenameColumn(new TableRenameColumnInfo( + table.getDatabase().getId(), table.getId(), "k2", "renamed_k2", schemaVersions)); + + Assertions.assertNotNull(table.getColumn("renamed_k2")); + Assertions.assertTrue( + getConstraintMgr().getDistributionMappingConstraintsForPlanning(table).isEmpty()); + } finally { + executeSql("drop table if exists mapping_schema_binding force"); + } + } + + @Test + void distributionMappingRejectsTemporaryTable() throws Exception { + createTable("create temporary table mapping_temporary (k1 int, k2 int) " + + "duplicate key(k1) distributed by hash(k1) buckets 4 " + + "properties(\"replication_num\"=\"1\")"); + try { + Exception exception = Assertions.assertThrows(Exception.class, () -> addConstraint( + "alter table mapping_temporary add constraint mapping_constraint " + + "colocate mapping mapping_id (k2) " + + "determines distribution key (k1) not enforced")); + Assertions.assertTrue(exception.getMessage().contains( + "Distribution mapping constraint does not support temporary tables")); + } finally { + executeSql("drop table if exists mapping_temporary force"); + } + } + + @Test + void distributionMappingFollowsTableObjectLifecycle() throws Exception { + createTable("create table mapping_lifecycle_a (k1 int, k2 int) " + + "duplicate key(k1) distributed by hash(k1) buckets 4 " + + "properties(\"replication_num\"=\"1\")"); + createTable("create table mapping_lifecycle_b (k1 int, k2 int) " + + "duplicate key(k1) distributed by hash(k1) buckets 4 " + + "properties(\"replication_num\"=\"1\")"); + createTable("create table mapping_replace_a (k1 int, k2 int) " + + "duplicate key(k1) distributed by hash(k1) buckets 4 " + + "properties(\"replication_num\"=\"1\")"); + createTable("create table mapping_replace_b (k1 int, k2 int) " + + "duplicate key(k1) distributed by hash(k1) buckets 4 " + + "properties(\"replication_num\"=\"1\")"); + try { + addConstraint("alter table mapping_lifecycle_a add constraint mapping_a " + + "colocate mapping mapping_a (k2) determines distribution key (k1) not enforced"); + addConstraint("alter table mapping_lifecycle_b add constraint mapping_b " + + "colocate mapping mapping_b (k2) determines distribution key (k1) not enforced"); + addConstraint("alter table mapping_replace_a add constraint replace_mapping_a " + + "colocate mapping replace_mapping_a (k2) determines distribution key (k1) not enforced"); + addConstraint("alter table mapping_replace_b add constraint replace_mapping_b " + + "colocate mapping replace_mapping_b (k2) determines distribution key (k1) not enforced"); + + OlapTable originalA = (OlapTable) Env.getCurrentInternalCatalog() + .getDbOrDdlException("test").getTableOrDdlException("mapping_lifecycle_a"); + OlapTable originalB = (OlapTable) Env.getCurrentInternalCatalog() + .getDbOrDdlException("test").getTableOrDdlException("mapping_lifecycle_b"); + executeSql("alter table mapping_lifecycle_a rename mapping_lifecycle_a_renamed"); + OlapTable renamedA = (OlapTable) Env.getCurrentInternalCatalog() + .getDbOrDdlException("test").getTableOrDdlException("mapping_lifecycle_a_renamed"); + Assertions.assertSame(originalA, renamedA); + Assertions.assertEquals("mapping_a", + getConstraintMgr().getDistributionMappingConstraints(renamedA).get(0).getName()); + + executeSql("alter table mapping_lifecycle_a_renamed replace with table mapping_lifecycle_b " + + "properties(\"swap\"=\"true\")"); + OlapTable currentA = (OlapTable) Env.getCurrentInternalCatalog() + .getDbOrDdlException("test").getTableOrDdlException("mapping_lifecycle_a_renamed"); + OlapTable currentB = (OlapTable) Env.getCurrentInternalCatalog() + .getDbOrDdlException("test").getTableOrDdlException("mapping_lifecycle_b"); + Assertions.assertSame(originalB, currentA); + Assertions.assertSame(originalA, currentB); + Assertions.assertEquals("mapping_b", + getConstraintMgr().getDistributionMappingConstraints(currentA).get(0).getName()); + Assertions.assertEquals("mapping_a", + getConstraintMgr().getDistributionMappingConstraints(currentB).get(0).getName()); + + executeSql("truncate table mapping_lifecycle_b"); + OlapTable truncatedB = (OlapTable) Env.getCurrentInternalCatalog() + .getDbOrDdlException("test").getTableOrDdlException("mapping_lifecycle_b"); + Assertions.assertSame(originalA, truncatedB); + Assertions.assertEquals("mapping_a", + getConstraintMgr().getDistributionMappingConstraints(truncatedB).get(0).getName()); + + executeSql("drop table mapping_lifecycle_b"); + createTable("create table mapping_lifecycle_b (k1 int, k2 int) " + + "duplicate key(k1) distributed by hash(k1) buckets 4 " + + "properties(\"replication_num\"=\"1\")"); + OlapTable sameNameReplacement = (OlapTable) Env.getCurrentInternalCatalog() + .getDbOrDdlException("test").getTableOrDdlException("mapping_lifecycle_b"); + Assertions.assertNotSame(originalA, sameNameReplacement); + Assertions.assertTrue( + getConstraintMgr().getDistributionMappingConstraints(sameNameReplacement).isEmpty()); + executeSql("drop table mapping_lifecycle_b force"); + executeSql("recover table mapping_lifecycle_b"); + OlapTable recoveredB = (OlapTable) Env.getCurrentInternalCatalog() + .getDbOrDdlException("test").getTableOrDdlException("mapping_lifecycle_b"); + Assertions.assertEquals(originalA.getId(), recoveredB.getId()); + Assertions.assertEquals("mapping_a", + getConstraintMgr().getDistributionMappingConstraints(recoveredB).get(0).getName()); + + OlapTable replacementB = (OlapTable) Env.getCurrentInternalCatalog() + .getDbOrDdlException("test").getTableOrDdlException("mapping_replace_b"); + executeSql("alter table mapping_replace_a replace with table mapping_replace_b " + + "properties(\"swap\"=\"false\")"); + OlapTable currentReplacement = (OlapTable) Env.getCurrentInternalCatalog() + .getDbOrDdlException("test").getTableOrDdlException("mapping_replace_a"); + Assertions.assertSame(replacementB, currentReplacement); + Assertions.assertEquals("replace_mapping_b", + getConstraintMgr().getDistributionMappingConstraints(currentReplacement).get(0).getName()); + } finally { + executeSql("drop table if exists mapping_lifecycle_a_renamed force"); + executeSql("drop table if exists mapping_lifecycle_b force"); + executeSql("drop table if exists mapping_replace_a force"); + executeSql("drop table if exists mapping_replace_b force"); + } + } + + @Test + void distributionMappingFollowsDatabaseRename() throws Exception { + executeSql("drop database if exists mapping_lifecycle_db force"); + executeSql("drop database if exists mapping_lifecycle_db_renamed force"); + createDatabase("mapping_lifecycle_db"); + createTable("create table mapping_lifecycle_db.mapping_table (k1 int, k2 int) " + + "duplicate key(k1) distributed by hash(k1) buckets 4 " + + "properties(\"replication_num\"=\"1\")"); + try { + addConstraint("alter table mapping_lifecycle_db.mapping_table add constraint mapping " + + "colocate mapping mapping_id (k2) determines distribution key (k1) not enforced"); + OlapTable originalTable = (OlapTable) Env.getCurrentInternalCatalog() + .getDbOrDdlException("mapping_lifecycle_db").getTableOrDdlException("mapping_table"); + + executeSql("alter database mapping_lifecycle_db rename mapping_lifecycle_db_renamed"); + + OlapTable renamedTable = (OlapTable) Env.getCurrentInternalCatalog() + .getDbOrDdlException("mapping_lifecycle_db_renamed") + .getTableOrDdlException("mapping_table"); + Assertions.assertSame(originalTable, renamedTable); + Assertions.assertEquals("mapping", + getConstraintMgr().getDistributionMappingConstraints(renamedTable).get(0).getName()); + + executeSql("drop database mapping_lifecycle_db_renamed"); + executeSql("recover database mapping_lifecycle_db_renamed"); + + OlapTable recoveredTable = (OlapTable) Env.getCurrentInternalCatalog() + .getDbOrDdlException("mapping_lifecycle_db_renamed") + .getTableOrDdlException("mapping_table"); + Assertions.assertSame(originalTable, recoveredTable); + Assertions.assertEquals("mapping", + getConstraintMgr().getDistributionMappingConstraints(recoveredTable).get(0).getName()); + } finally { + executeSql("drop database if exists mapping_lifecycle_db force"); + executeSql("drop database if exists mapping_lifecycle_db_renamed force"); + } + } + + @Test + void distributionMappingIsNotCopiedByCreateTableLike() throws Exception { + createTable("create table mapping_like_source (k1 int, k2 int) " + + "duplicate key(k1) distributed by hash(k1) buckets 4 " + + "properties(\"replication_num\"=\"1\")"); + try { + addConstraint("alter table mapping_like_source add constraint mapping " + + "colocate mapping mapping_id (k2) determines distribution key (k1) not enforced"); + + executeSql("create table mapping_like_target like mapping_like_source"); + + TableIf targetTable = Env.getCurrentInternalCatalog() + .getDbOrDdlException("test").getTableOrDdlException("mapping_like_target"); + Assertions.assertTrue(getConstraintMgr().getDistributionMappingConstraints(targetTable).isEmpty()); + } finally { + executeSql("drop table if exists mapping_like_source force"); + executeSql("drop table if exists mapping_like_target force"); + } + } + @Test void foreignKeyConstraintTest() throws Exception { AddConstraintCommand command = (AddConstraintCommand) new NereidsParser().parseSingle( diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/util/JoinUtilsTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/util/JoinUtilsTest.java index 9d634d58ab9fe7..c6b34983bf6635 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/util/JoinUtilsTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/util/JoinUtilsTest.java @@ -20,8 +20,10 @@ import org.apache.doris.catalog.ColocateTableIndex; import org.apache.doris.catalog.ColocateTableIndex.GroupId; import org.apache.doris.catalog.Env; +import org.apache.doris.nereids.properties.DistributionMapping; import org.apache.doris.nereids.properties.DistributionSpecHash; import org.apache.doris.nereids.properties.DistributionSpecHash.ShuffleType; +import org.apache.doris.nereids.properties.NaturalDistributionMappingSpec; import org.apache.doris.nereids.trees.expressions.Add; import org.apache.doris.nereids.trees.expressions.EqualTo; import org.apache.doris.nereids.trees.expressions.ExprId; @@ -31,6 +33,8 @@ import org.apache.doris.nereids.types.TinyIntType; import org.apache.doris.qe.ConnectContext; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import com.google.common.collect.Lists; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; @@ -145,6 +149,154 @@ public void testCouldColocateJoinForDiffTableInSameGroupAndGroupIsStable() { } } + @Test + public void testCouldColocateJoinByDistributionMappings() { + ConnectContext ctx = new ConnectContext(); + ctx.setThreadLocalInfo(); + + try (MockedStatic mockedEnv = Mockito.mockStatic(Env.class)) { + GroupId groupId = new GroupId(1L, 1L); + ColocateTableIndex colocateIndex = Mockito.mock(ColocateTableIndex.class); + Mockito.when(colocateIndex.isSameGroup(1L, 2L)).thenReturn(true); + Mockito.when(colocateIndex.getGroup(1L)).thenReturn(groupId); + Mockito.when(colocateIndex.isGroupUnstable(groupId)).thenReturn(false); + mockedEnv.when(Env::getCurrentColocateIndex).thenReturn(colocateIndex); + + DistributionMapping leftMapping1 = new DistributionMapping( + "mapping_1", ImmutableList.of(new ExprId(5)), ImmutableList.of(0)); + DistributionMapping rightMapping1 = new DistributionMapping( + "mapping_1", ImmutableList.of(new ExprId(6)), ImmutableList.of(0)); + DistributionMapping leftMapping2 = new DistributionMapping( + "mapping_2", ImmutableList.of(new ExprId(9)), ImmutableList.of(1)); + DistributionMapping rightMapping2 = new DistributionMapping( + "mapping_2", ImmutableList.of(new ExprId(10)), ImmutableList.of(1)); + DistributionSpecHash left = new DistributionSpecHash( + ImmutableList.of(new ExprId(1), new ExprId(2)), ShuffleType.NATURAL, + 1L, 1L, Collections.emptySet(), ImmutableList.of(leftMapping1, leftMapping2)); + DistributionSpecHash right = new DistributionSpecHash( + ImmutableList.of(new ExprId(3), new ExprId(4)), ShuffleType.NATURAL, + 2L, 2L, Collections.emptySet(), ImmutableList.of(rightMapping1, rightMapping2)); + + SlotReference leftK2 = slot(2); + SlotReference rightK2 = slot(4); + SlotReference leftD1 = slot(5); + SlotReference rightD1 = slot(6); + SlotReference leftExtra = slot(7); + SlotReference rightExtra = slot(8); + SlotReference leftD2 = slot(9); + SlotReference rightD2 = slot(10); + + List directAndMapping = ImmutableList.of( + new EqualTo(leftD1, rightD1), new EqualTo(leftK2, rightK2)); + Assertions.assertFalse(JoinUtils.couldColocateJoin(left, right, directAndMapping)); + + ctx.getSessionVariable().enableColocateMappingConstraint = true; + Assertions.assertTrue(JoinUtils.couldColocateJoin(left, right, directAndMapping)); + Assertions.assertTrue(JoinUtils.couldColocateJoin(left, right, ImmutableList.of( + new EqualTo(rightD1, leftD1), new EqualTo(leftK2, rightK2)))); + Assertions.assertTrue(JoinUtils.couldColocateJoin(left, right, ImmutableList.of( + new EqualTo(leftD1, rightD1), new EqualTo(leftK2, rightK2), + new EqualTo(leftExtra, rightExtra)))); + DistributionMapping rightMappingWithDifferentId = new DistributionMapping( + "different_mapping", ImmutableList.of(new ExprId(6)), ImmutableList.of(0)); + DistributionSpecHash rightWithDifferentMappingId = new DistributionSpecHash( + ImmutableList.of(new ExprId(3), new ExprId(4)), ShuffleType.NATURAL, + 2L, 2L, Collections.emptySet(), + ImmutableList.of(rightMappingWithDifferentId, rightMapping2)); + Assertions.assertFalse(JoinUtils.couldColocateJoin( + left, rightWithDifferentMappingId, directAndMapping)); + DistributionMapping rightMappingWithWrongDeterminant = new DistributionMapping( + "mapping_1", ImmutableList.of(new ExprId(16)), ImmutableList.of(0)); + DistributionMapping rightMappingWithWrongTarget = new DistributionMapping( + "mapping_1", ImmutableList.of(new ExprId(6)), ImmutableList.of(1)); + DistributionSpecHash rightWithCompatibleCandidates = new DistributionSpecHash( + ImmutableList.of(new ExprId(3), new ExprId(4)), ShuffleType.NATURAL, + 2L, 2L, Collections.emptySet(), + ImmutableList.of( + rightMappingWithWrongDeterminant, + rightMappingWithWrongTarget, + rightMapping1, + rightMapping2)); + Assertions.assertTrue(JoinUtils.couldColocateJoin( + left, rightWithCompatibleCandidates, directAndMapping)); + DistributionSpecHash rightWithoutMatchingCandidate = new DistributionSpecHash( + ImmutableList.of(new ExprId(3), new ExprId(4)), ShuffleType.NATURAL, + 2L, 2L, Collections.emptySet(), + ImmutableList.of( + rightMappingWithWrongDeterminant, + rightMappingWithWrongTarget, + rightMapping2)); + Assertions.assertFalse(JoinUtils.couldColocateJoin( + left, rightWithoutMatchingCandidate, directAndMapping)); + + List mostlyUnrelatedMappings = Lists.newArrayList(); + for (int i = 0; i < 128; i++) { + mostlyUnrelatedMappings.add(new DistributionMapping( + "unrelated_" + i, + ImmutableList.of(new ExprId(1000 + i)), + ImmutableList.of(i % 2))); + } + mostlyUnrelatedMappings.add(rightMapping1); + mostlyUnrelatedMappings.add(rightMapping2); + DistributionSpecHash rightWithMostlyUnrelatedMappings = new DistributionSpecHash( + ImmutableList.of(new ExprId(3), new ExprId(4)), ShuffleType.NATURAL, + 2L, 2L, Collections.emptySet(), mostlyUnrelatedMappings); + Assertions.assertTrue(JoinUtils.couldColocateJoin( + left, rightWithMostlyUnrelatedMappings, directAndMapping)); + + DistributionMapping leftOrderedMapping = new DistributionMapping( + "ordered_mapping", + ImmutableList.of(new ExprId(5), new ExprId(9)), + ImmutableList.of(0, 1)); + DistributionMapping rightReorderedDeterminants = new DistributionMapping( + "ordered_mapping", + ImmutableList.of(new ExprId(10), new ExprId(6)), + ImmutableList.of(0, 1)); + DistributionSpecHash leftWithOrderedMapping = new DistributionSpecHash( + ImmutableList.of(new ExprId(1), new ExprId(2)), ShuffleType.NATURAL, + 1L, 1L, Collections.emptySet(), ImmutableList.of(leftOrderedMapping)); + DistributionSpecHash rightWithReorderedDeterminants = new DistributionSpecHash( + ImmutableList.of(new ExprId(3), new ExprId(4)), ShuffleType.NATURAL, + 2L, 2L, Collections.emptySet(), + ImmutableList.of(rightReorderedDeterminants)); + Assertions.assertFalse(JoinUtils.couldColocateJoin( + leftWithOrderedMapping, rightWithReorderedDeterminants, + ImmutableList.of( + new EqualTo(leftD1, rightD1), + new EqualTo(leftD2, rightD2)))); + Assertions.assertFalse(JoinUtils.couldColocateJoin( + left, right, ImmutableList.of(new EqualTo(leftD1, rightD1)))); + Assertions.assertTrue(JoinUtils.couldColocateJoin(left, right, ImmutableList.of( + new EqualTo(leftD1, rightD1), new EqualTo(leftD2, rightD2)))); + + NaturalDistributionMappingSpec leftWithHiddenK1 = + naturalMappingSpec(left).project(ImmutableMap.of( + new ExprId(2), new ExprId(2), + new ExprId(5), new ExprId(5))).get(); + NaturalDistributionMappingSpec rightWithHiddenK1 = + naturalMappingSpec(right).project(ImmutableMap.of( + new ExprId(4), new ExprId(4), + new ExprId(6), new ExprId(6))).get(); + Assertions.assertTrue(JoinUtils.couldColocateJoinByMapping( + leftWithHiddenK1, rightWithHiddenK1, directAndMapping)); + Assertions.assertFalse(JoinUtils.couldColocateJoinByMapping( + leftWithHiddenK1, rightWithHiddenK1, + ImmutableList.of(new EqualTo(leftD1, rightD1)))); + } + } + + private NaturalDistributionMappingSpec naturalMappingSpec(DistributionSpecHash hashSpec) { + return new NaturalDistributionMappingSpec( + hashSpec.getTableId(), hashSpec.getSelectedIndexId(), hashSpec.getPartitionIds(), + hashSpec.getOrderedShuffledColumns().size(), hashSpec.getExprIdToEquivalenceSet(), + hashSpec.getDistributionMappings()); + } + + private SlotReference slot(int exprId) { + return new SlotReference(new ExprId(exprId), "c" + exprId, + TinyIntType.INSTANCE, false, Lists.newArrayList()); + } + @Test public void testCouldColocateJoinForNotNaturalHashDstribution() { ConnectContext ctx = new ConnectContext(); diff --git a/fe/fe-sql-parser/src/main/antlr4/org/apache/doris/nereids/DorisLexer.g4 b/fe/fe-sql-parser/src/main/antlr4/org/apache/doris/nereids/DorisLexer.g4 index 7217bdb7f7fb9f..2241e5a40b1930 100644 --- a/fe/fe-sql-parser/src/main/antlr4/org/apache/doris/nereids/DorisLexer.g4 +++ b/fe/fe-sql-parser/src/main/antlr4/org/apache/doris/nereids/DorisLexer.g4 @@ -195,6 +195,7 @@ DEFAULT: 'DEFAULT'; DEFERRED: 'DEFERRED'; DELETE: 'DELETE'; DEMAND: 'DEMAND'; +DETERMINES: 'DETERMINES'; DESC: 'DESC'; DESCRIBE: 'DESCRIBE'; DIAGNOSE: 'DIAGNOSE'; @@ -222,6 +223,7 @@ DYNAMIC: 'DYNAMIC'; E:'E'; ELSE: 'ELSE'; ENABLE: 'ENABLE'; +ENFORCED: 'ENFORCED'; ENCRYPTION: 'ENCRYPTION'; ENCRYPTKEY: 'ENCRYPTKEY'; ENCRYPTKEYS: 'ENCRYPTKEYS'; diff --git a/fe/fe-sql-parser/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4 b/fe/fe-sql-parser/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4 index 2176dd3fb84ed0..5be23ea633af4f 100644 --- a/fe/fe-sql-parser/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4 +++ b/fe/fe-sql-parser/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4 @@ -1189,6 +1189,10 @@ constraint | FOREIGN KEY slots=identifierList REFERENCES referenceTable=multipartIdentifier referencedSlots=identifierList + | COLOCATE MAPPING mappingId=identifier + slots=identifierList + DETERMINES DISTRIBUTION KEY distributionSlots=identifierList + NOT ENFORCED ; partitionSpec @@ -2241,6 +2245,7 @@ nonReserved | DECIMALV3 | DEFERRED | DEMAND + | DETERMINES | DIAGNOSE | DIAGNOSIS | DICTIONARIES @@ -2255,6 +2260,7 @@ nonReserved | DYNAMIC | E | ENABLE + | ENFORCED | ENCRYPTION | ENCRYPTKEY | ENCRYPTKEYS diff --git a/regression-test/data/query_p0/colocate/test_colocate_mapping_constraint.out b/regression-test/data/query_p0/colocate/test_colocate_mapping_constraint.out new file mode 100644 index 00000000000000..c917ffd4d6de69 --- /dev/null +++ b/regression-test/data/query_p0/colocate/test_colocate_mapping_constraint.out @@ -0,0 +1,45 @@ +-- This file is automatically generated. You should know what you did if you want to edit this +-- !distinct_aggregate_mapping_barrier -- +10 100 1 1 10 100 +20 200 1 2 20 200 + +-- !pure_deduplicate_mapping_barrier -- +10 100 1 10 100 +20 200 2 20 200 + +-- !mixed_distinct_aggregate_mapping_barrier -- +100 10 7 1 1 10 100 +200 20 8 1 2 20 200 + +-- !generate_barrier_result -- +100 10 0 1 +100 10 1 1 +200 20 0 2 +200 20 1 2 + +-- !window_barrier_result -- +100 10 1 1 +200 20 1 2 + +-- !partition_topn_barrier_result -- +100 10 1 1 +200 20 1 2 + +-- !nested_loop_barrier_result -- +100 10 1 +200 20 2 + +-- !broadcast_barrier_result -- +100 10 1 + +-- !colocate_mapping_result -- +1 10 100 1000 7 1 10 100 1000 7 +2 20 200 2000 8 2 20 200 2000 9 + +-- !aggregate_colocate_mapping_result -- +1 10 100 7 1 10 100 +2 20 200 8 2 20 200 + +-- !hidden_distribution_key_aggregate_result -- +10 100 7 1 10 100 +20 200 8 2 20 200 diff --git a/regression-test/suites/backup_restore/test_backup_restore_atomic_with_alter.groovy b/regression-test/suites/backup_restore/test_backup_restore_atomic_with_alter.groovy index e8f4f9977a19c8..1ce9ea0fbe9ffa 100644 --- a/regression-test/suites/backup_restore/test_backup_restore_atomic_with_alter.groovy +++ b/regression-test/suites/backup_restore/test_backup_restore_atomic_with_alter.groovy @@ -25,9 +25,24 @@ suite("test_backup_restore_atomic_with_alter", "backup_restore,nonConcurrent") { String dbName = "${suiteName}_db" String repoName = "${suiteName}_repo_" + UUID.randomUUID().toString().replace("-", "") String snapshotName = "snapshot_" + UUID.randomUUID().toString().replace("-", "") + String completionSnapshotName = "completion_snapshot_" + UUID.randomUUID().toString().replace("-", "") + String mappingSnapshotName = "mapping_snapshot_" + UUID.randomUUID().toString().replace("-", "") String tableNamePrefix = "${suiteName}_tables" def syncer = getSyncer() + def waitRestorePaused = { label -> + boolean restorePaused = false + for (int k = 0; k < 60; k++) { + def records = sql_return_maparray """ SHOW RESTORE FROM ${dbName} WHERE Label = "${label}" """ + if (records.size() == 1 && records[0].State != 'PENDING' && records[0].State != 'CREATING') { + restorePaused = true + break + } + logger.info("SHOW RESTORE result: ${records}") + sleep(3000) + } + assertTrue(restorePaused) + } syncer.createS3Repository(repoName) sql "DROP DATABASE IF EXISTS ${dbName} FORCE" sql "CREATE DATABASE ${dbName}" @@ -92,6 +107,15 @@ suite("test_backup_restore_atomic_with_alter", "backup_restore,nonConcurrent") { def snapshot = syncer.getSnapshotTimestamp(repoName, snapshotName) assertTrue(snapshot != null) + // Add the mapping after backup so the backup remains eligible for atomic restore while + // the live table exercises the DDL fence during the restore. + sql """ + ALTER TABLE ${dbName}.${tableNamePrefix}_1 + ADD CONSTRAINT atomic_restore_mapping + COLOCATE MAPPING atomic_restore_mapping (`count`) + DETERMINES DISTRIBUTION KEY (`id`) NOT ENFORCED + """ + // drop table_0 sql "DROP TABLE ${dbName}.${tableNamePrefix}_0 FORCE" @@ -111,17 +135,7 @@ suite("test_backup_restore_atomic_with_alter", "backup_restore,nonConcurrent") { """ sql "SYNC" - boolean restore_paused = false - for (int k = 0; k < 60; k++) { - def records = sql_return_maparray """ SHOW RESTORE FROM ${dbName} WHERE Label = "${snapshotName}" """ - if (records.size() == 1 && (records[0].State != 'PENDING' && records[0].State != 'CREATING')) { - restore_paused = true - break - } - logger.info("SHOW RESTORE result: ${records}") - sleep(3000) - } - assertTrue(restore_paused) + waitRestorePaused(snapshotName) sql "SYNC" @@ -200,6 +214,12 @@ suite("test_backup_restore_atomic_with_alter", "backup_restore,nonConcurrent") { RENAME newTableName """ }, "Do not allow doing ALTER ops") + expectExceptionLike({ + sql """ + ALTER TABLE ${dbName}.${tableNamePrefix}_1 + DROP CONSTRAINT atomic_restore_mapping + """ + }, "atomic restore state") // BTW, the tmp table also don't allow rename expectExceptionLike({ sql """ @@ -234,6 +254,106 @@ suite("test_backup_restore_atomic_with_alter", "backup_restore,nonConcurrent") { show_result = master_sql """ SHOW CREATE TABLE ${dbName}.${tableNamePrefix}_1 """ logger.info("SHOW CREATE TABLE ${tableNamePrefix}_1: ${show_result}") assertFalse(show_result[0][1].contains("in_atomic_restore")) + sql """ + ALTER TABLE ${dbName}.${tableNamePrefix}_1 + DROP CONSTRAINT atomic_restore_mapping + """ + + // Also verify the DDL fence while a restore runs to completion, not only while it is cancelled. + sql """ + BACKUP SNAPSHOT ${dbName}.${completionSnapshotName} + TO `${repoName}` + ON (${tableNamePrefix}_1) + """ + syncer.waitSnapshotFinish(dbName) + def completionSnapshot = syncer.getSnapshotTimestamp(repoName, completionSnapshotName) + assertTrue(completionSnapshot != null) + sql """ + ALTER TABLE ${dbName}.${tableNamePrefix}_1 + ADD CONSTRAINT atomic_restore_mapping + COLOCATE MAPPING atomic_restore_mapping (`count`) + DETERMINES DISTRIBUTION KEY (`id`) NOT ENFORCED + """ + GetDebugPoint().enableDebugPointForAllFEs( + "FE.PAUSE_NON_PENDING_RESTORE_JOB", [value:completionSnapshotName]) + sql """ + RESTORE SNAPSHOT ${dbName}.${completionSnapshotName} + FROM `${repoName}` + PROPERTIES + ( + "backup_timestamp" = "${completionSnapshot}", + "reserve_replica" = "true", + "atomic_restore" = "true" + ) + """ + waitRestorePaused(completionSnapshotName) + expectExceptionLike({ + sql """ + ALTER TABLE ${dbName}.${tableNamePrefix}_1 + DROP CONSTRAINT atomic_restore_mapping + """ + }, "atomic restore state") + GetDebugPoint().disableDebugPointForAllFEs("FE.PAUSE_NON_PENDING_RESTORE_JOB") + syncer.waitAllRestoreFinish(dbName) + sql "SYNC" + + def completionRestore = sql_return_maparray """ + SHOW RESTORE FROM ${dbName} WHERE Label = "${completionSnapshotName}" + """ + assertTrue(completionRestore.size() == 1) + assertTrue(completionRestore[0].State == "FINISHED") + show_result = master_sql """ SHOW CREATE TABLE ${dbName}.${tableNamePrefix}_1 """ + assertFalse(show_result[0][1].contains("in_atomic_restore")) + assertTrue((sql "SHOW CONSTRAINTS FROM ${dbName}.${tableNamePrefix}_1").isEmpty()) + + // An atomic restore with mappings is rejected before staging. This keeps a centralized + // constraint with the same name from colliding with a mapping published from the backup. + sql """ + ALTER TABLE ${dbName}.${tableNamePrefix}_1 + ADD CONSTRAINT atomic_restore_collision + COLOCATE MAPPING atomic_restore_collision (`count`) + DETERMINES DISTRIBUTION KEY (`id`) NOT ENFORCED + """ + sql """ + BACKUP SNAPSHOT ${dbName}.${mappingSnapshotName} + TO `${repoName}` + ON (${tableNamePrefix}_1) + """ + syncer.waitSnapshotFinish(dbName) + def mappingSnapshot = syncer.getSnapshotTimestamp(repoName, mappingSnapshotName) + assertTrue(mappingSnapshot != null) + + sql """ + ALTER TABLE ${dbName}.${tableNamePrefix}_1 + DROP CONSTRAINT atomic_restore_collision + """ + sql """ + ALTER TABLE ${dbName}.${tableNamePrefix}_1 + ADD CONSTRAINT atomic_restore_collision UNIQUE (`id`) + """ + sql """ + RESTORE SNAPSHOT ${dbName}.${mappingSnapshotName} + FROM `${repoName}` + PROPERTIES + ( + "backup_timestamp" = "${mappingSnapshot}", + "reserve_replica" = "true", + "atomic_restore" = "true" + ) + """ + syncer.waitRestoreError(dbName, "backup contains distribution mapping constraints") + + def mappingRestore = sql_return_maparray """ + SHOW RESTORE FROM ${dbName} WHERE Label = "${mappingSnapshotName}" + """ + assertTrue(mappingRestore.size() == 1) + assertTrue(mappingRestore[0].State == "CANCELLED") + assertTrue(mappingRestore[0].Status.contains("backup contains distribution mapping constraints")) + def constraints = sql "SHOW CONSTRAINTS FROM ${dbName}.${tableNamePrefix}_1" + assertTrue(constraints.size() == 1) + assertTrue(constraints[0][0] == "atomic_restore_collision") + show_result = master_sql """ SHOW CREATE TABLE ${dbName}.${tableNamePrefix}_1 """ + assertFalse(show_result[0][1].contains("in_atomic_restore")) for (def tableName in tables) { sql "DROP TABLE IF EXISTS ${dbName}.${tableName} FORCE" @@ -244,6 +364,3 @@ suite("test_backup_restore_atomic_with_alter", "backup_restore,nonConcurrent") { GetDebugPoint().disableDebugPointForAllFEs("FE.PAUSE_NON_PENDING_RESTORE_JOB") } } - - - diff --git a/regression-test/suites/ccr_syncer_p1/test_colocate_mapping_constraint.groovy b/regression-test/suites/ccr_syncer_p1/test_colocate_mapping_constraint.groovy new file mode 100644 index 00000000000000..79d09793505346 --- /dev/null +++ b/regression-test/suites/ccr_syncer_p1/test_colocate_mapping_constraint.groovy @@ -0,0 +1,111 @@ +// 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. + +suite("test_colocate_mapping_constraint_ccr") { + def syncer = getSyncer() + if (!syncer.checkEnableFeatureBinlog()) { + logger.info("fe enable_feature_binlog is false, skip case test_colocate_mapping_constraint_ccr") + return + } + + def tableName = "tbl_colocate_mapping_constraint_ccr" + def snapshotName = "snapshot_colocate_mapping_constraint_ccr" + def rowCount = 5 + + sql "DROP TABLE IF EXISTS ${tableName}" + target_sql "DROP TABLE IF EXISTS ${tableName}" + sql """ + CREATE TABLE ${tableName} ( + `test` INT, + `id` INT + ) + ENGINE=OLAP + UNIQUE KEY(`test`, `id`) + DISTRIBUTED BY HASH(id) BUCKETS 4 + PROPERTIES ( + "replication_allocation" = "tag.location.default: 1", + "binlog.enable" = "true" + ) + """ + sql """ + ALTER TABLE ${tableName} + ADD CONSTRAINT ccr_mapping + COLOCATE MAPPING ccr_mapping_id (test) DETERMINES DISTRIBUTION KEY (id) NOT ENFORCED + """ + for (int i = 0; i < rowCount; ++i) { + sql "INSERT INTO ${tableName} VALUES (1, ${i})" + } + sql "SYNC" + + sql """ + BACKUP SNAPSHOT ${context.dbName}.${snapshotName} + TO `__keep_on_local__` + ON (${tableName}) + PROPERTIES ("type" = "full") + """ + syncer.waitSnapshotFinish() + assertTrue(syncer.getSnapshot(snapshotName, tableName)) + assertTrue(syncer.context.getSnapshotResult.isSetCommitSeq()) + syncer.context.seq = syncer.context.getSnapshotResult.getCommitSeq() + assertTrue(syncer.restoreSnapshot(true)) + syncer.waitTargetRestoreFinish() + target_sql "SYNC" + + assertTrue((target_sql "SHOW CONSTRAINTS FROM ${tableName}").isEmpty()) + def targetExplain = target_sql """ + EXPLAIN SELECT /*+ SET_VAR(disable_join_reorder=true, + enable_colocate_mapping_constraint=true, + auto_broadcast_join_threshold=-1, + broadcast_row_count_limit=0) */ COUNT(*) + FROM ${tableName} l JOIN ${tableName} r ON l.test = r.test + """ + assertFalse(targetExplain.toString().contains("COLOCATE")) + assertTrue(syncer.getTargetMeta(tableName)) + + sql """ + ALTER TABLE ${tableName} + DROP CONSTRAINT ccr_mapping + """ + sql "SYNC" + sql "INSERT INTO ${tableName} VALUES (1, ${rowCount})" + boolean foundTableDataBinlog = false + for (int attempt = 0; attempt < 10 && !foundTableDataBinlog; ++attempt) { + assertTrue(syncer.getBinlog(tableName)) + def sourceTable = syncer.context.sourceTableMap.get(tableName) + foundTableDataBinlog = syncer.context.lastBinlog.tableRecords != null + && syncer.context.lastBinlog.tableRecords.containsKey(sourceTable.id) + } + assertTrue(foundTableDataBinlog) + assertTrue(syncer.beginTxn(tableName)) + assertTrue(syncer.getBackendClients()) + assertTrue(syncer.ingestBinlog()) + assertTrue(syncer.commitTxn()) + assertTrue(syncer.checkTargetVersion()) + syncer.closeBackendClients() + + target_sql "SYNC" + def targetRows = target_sql "SELECT * FROM ${tableName}" + assertEquals(rowCount + 1, targetRows.size()) + def joinCount = target_sql """ + SELECT /*+ SET_VAR(disable_join_reorder=true, + enable_colocate_mapping_constraint=true, + auto_broadcast_join_threshold=-1, + broadcast_row_count_limit=0) */ COUNT(*) + FROM ${tableName} l JOIN ${tableName} r ON l.test = r.test + """ + assertEquals((rowCount + 1L) * (rowCount + 1L), joinCount[0][0] as long) +} diff --git a/regression-test/suites/query_p0/colocate/test_colocate_mapping_constraint.groovy b/regression-test/suites/query_p0/colocate/test_colocate_mapping_constraint.groovy new file mode 100644 index 00000000000000..59bcf1e8934ff4 --- /dev/null +++ b/regression-test/suites/query_p0/colocate/test_colocate_mapping_constraint.groovy @@ -0,0 +1,1003 @@ +// 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. + +suite("test_colocate_mapping_constraint") { + sql """ DROP TABLE IF EXISTS test_colocate_mapping_constraint_left """ + sql """ DROP TABLE IF EXISTS test_colocate_mapping_constraint_right """ + sql """ DROP TABLE IF EXISTS test_colocate_mapping_composite_left """ + sql """ DROP TABLE IF EXISTS test_colocate_mapping_composite_right """ + + sql """ + CREATE TABLE test_colocate_mapping_constraint_left ( + k1 INT, + k2 INT, + d1 INT, + d2 INT, + extra_col INT + ) ENGINE=OLAP + DUPLICATE KEY(k1, k2) + DISTRIBUTED BY HASH(k1, k2) BUCKETS 4 + PROPERTIES ( + "replication_num" = "1", + "colocate_with" = "test_colocate_mapping_constraint_group" + ) + """ + sql """ + CREATE TABLE test_colocate_mapping_constraint_right ( + k1 INT, + k2 INT, + d1 INT, + d2 INT, + extra_col INT + ) ENGINE=OLAP + DUPLICATE KEY(k1, k2) + DISTRIBUTED BY HASH(k1, k2) BUCKETS 4 + PROPERTIES ( + "replication_num" = "1", + "colocate_with" = "test_colocate_mapping_constraint_group" + ) + """ + + sql """ + ALTER TABLE test_colocate_mapping_constraint_left + ADD CONSTRAINT left_mapping_1 + COLOCATE MAPPING mapping_1 (d1) DETERMINES DISTRIBUTION KEY (k1) NOT ENFORCED + """ + sql """ + ALTER TABLE test_colocate_mapping_constraint_right + ADD CONSTRAINT right_mapping_1 + COLOCATE MAPPING mapping_1 (d1) DETERMINES DISTRIBUTION KEY (k1) NOT ENFORCED + """ + sql """ + ALTER TABLE test_colocate_mapping_constraint_left + ADD CONSTRAINT left_mapping_2 + COLOCATE MAPPING mapping_2 (d2) DETERMINES DISTRIBUTION KEY (k2) NOT ENFORCED + """ + sql """ + ALTER TABLE test_colocate_mapping_constraint_right + ADD CONSTRAINT right_mapping_2 + COLOCATE MAPPING mapping_2 (d2) DETERMINES DISTRIBUTION KEY (k2) NOT ENFORCED + """ + test { + sql """ + ALTER TABLE test_colocate_mapping_constraint_left + DROP COLUMN D1 + """ + exception "left_mapping_1" + } + test { + sql """ + ALTER TABLE test_colocate_mapping_constraint_left + RENAME COLUMN D1 renamed_d1 + """ + exception "left_mapping_1" + } + sql """ + CREATE TABLE test_colocate_mapping_composite_left ( + k1 INT, + k2 INT, + d1 INT, + d2 INT, + extra_col INT + ) ENGINE=OLAP + DUPLICATE KEY(k1, k2) + DISTRIBUTED BY HASH(k1, k2) BUCKETS 4 + PROPERTIES ( + "replication_num" = "1", + "colocate_with" = "test_colocate_mapping_composite_group" + ) + """ + sql """ + CREATE TABLE test_colocate_mapping_composite_right ( + k1 INT, + k2 INT, + d1 INT, + d2 INT, + extra_col INT + ) ENGINE=OLAP + DUPLICATE KEY(k1, k2) + DISTRIBUTED BY HASH(k1, k2) BUCKETS 4 + PROPERTIES ( + "replication_num" = "1", + "colocate_with" = "test_colocate_mapping_composite_group" + ) + """ + sql """ + ALTER TABLE test_colocate_mapping_composite_left + ADD CONSTRAINT composite_left_mapping + COLOCATE MAPPING composite_mapping (d1, d2) + DETERMINES DISTRIBUTION KEY (k1) NOT ENFORCED + """ + sql """ + ALTER TABLE test_colocate_mapping_composite_right + ADD CONSTRAINT composite_right_mapping + COLOCATE MAPPING composite_mapping (d1, d2) + DETERMINES DISTRIBUTION KEY (k1) NOT ENFORCED + """ + + sql """ INSERT INTO test_colocate_mapping_constraint_left VALUES + (1, 10, 100, 1000, 7), (2, 20, 200, 2000, 8) """ + sql """ INSERT INTO test_colocate_mapping_constraint_right VALUES + (1, 10, 100, 1000, 7), (2, 20, 200, 2000, 9) """ + sql """ INSERT INTO test_colocate_mapping_composite_left VALUES + (1, 10, 100, 1000, 7), (2, 20, 200, 2000, 8) """ + sql """ INSERT INTO test_colocate_mapping_composite_right VALUES + (1, 10, 100, 1000, 7), (2, 20, 200, 2000, 9) """ + sql """ SYNC """ + createMV(""" + CREATE MATERIALIZED VIEW mapping_alias_rollup AS + SELECT k1 AS alias_k1, k2 AS alias_k2, d1 AS alias_d1 + FROM test_colocate_mapping_constraint_left + """) + createMV(""" + CREATE MATERIALIZED VIEW mapping_without_determinant_rollup AS + SELECT k1 AS no_determinant_k1, k2 AS no_determinant_k2, + extra_col AS no_determinant_extra + FROM test_colocate_mapping_constraint_left + """) + waitForColocateGroupStable("test_colocate_mapping_constraint_group") + waitForColocateGroupStable("test_colocate_mapping_composite_group") + + sql """ SET auto_broadcast_join_threshold = -1 """ + sql """ SET broadcast_row_count_limit = 0 """ + // A selected rollup Slot alias is recognized, then conservatively falls back under the + // current selected-rollup boundary. The base-provenance binding is asserted by FE UT. + explain { + sql """ + SELECT /*+ use_mv(test_colocate_mapping_constraint_left.mapping_alias_rollup) */ + l.d1, r.d1 + FROM test_colocate_mapping_constraint_left l + JOIN test_colocate_mapping_constraint_left r + ON l.d1 = r.d1 AND l.k2 = r.k2 + """ + contains "test_colocate_mapping_constraint_left(mapping_alias_rollup)" + notContains "COLOCATE" + } + // A selected rollup without the determinant cannot propagate the mapping proof. + explain { + sql """ + SELECT /*+ use_mv(test_colocate_mapping_constraint_left.mapping_without_determinant_rollup) */ + l.extra_col, r.extra_col + FROM test_colocate_mapping_constraint_left l + JOIN test_colocate_mapping_constraint_left r + ON l.extra_col = r.extra_col AND l.k2 = r.k2 + """ + contains "test_colocate_mapping_constraint_left(mapping_without_determinant_rollup)" + notContains "COLOCATE" + } + def nestedSubqueryJoinQueries = [ + // Parallel Project subqueries on both sides. + """ + SELECT * + FROM ( + SELECT d1, k2, extra_col + FROM test_colocate_mapping_constraint_left + ) l + JOIN ( + SELECT d1, k2, extra_col + FROM test_colocate_mapping_constraint_right + ) r + ON l.d1 = r.d1 AND l.k2 = r.k2 + """, + // Parallel Aggregate subqueries on both sides. + """ + SELECT * + FROM ( + SELECT d1, k2, SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY k1, k2, d1 + ) l + JOIN ( + SELECT d1, k2, SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_right + GROUP BY k1, k2, d1 + ) r + ON l.d1 = r.d1 AND l.k2 = r.k2 + """, + // Multiple nested Project subqueries on both sides. + """ + SELECT * + FROM ( + SELECT inner_l.d1 AS nested_d1, inner_l.k2 AS nested_k2 + FROM ( + SELECT d1, k2 + FROM test_colocate_mapping_constraint_left + ) inner_l + ) l + JOIN ( + SELECT inner_r.d1 AS nested_d1, inner_r.k2 AS nested_k2 + FROM ( + SELECT d1, k2 + FROM test_colocate_mapping_constraint_right + ) inner_r + ) r + ON l.nested_d1 = r.nested_d1 AND l.nested_k2 = r.nested_k2 + """, + // Aggregate subqueries followed by another Project layer on both sides. + """ + SELECT * + FROM ( + SELECT aggregate_l.d1 AS nested_d1, + aggregate_l.k2 AS nested_k2, + aggregate_l.sum_extra + FROM ( + SELECT d1, k2, SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY k1, k2, d1 + ) aggregate_l + ) l + JOIN ( + SELECT aggregate_r.d1 AS nested_d1, + aggregate_r.k2 AS nested_k2, + aggregate_r.sum_extra + FROM ( + SELECT d1, k2, SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_right + GROUP BY k1, k2, d1 + ) aggregate_r + ) r + ON l.nested_d1 = r.nested_d1 AND l.nested_k2 = r.nested_k2 + """ + ] + + sql """ SET enable_colocate_mapping_constraint = false """ + nestedSubqueryJoinQueries.each { query -> + explain { + sql query + notContains "COLOCATE" + } + } + // The feature switch must not affect the original direct distribution-key colocate path. + explain { + sql """ SELECT * + FROM test_colocate_mapping_constraint_left l + JOIN test_colocate_mapping_constraint_right r + ON l.k1 = r.k1 AND l.k2 = r.k2 """ + contains "COLOCATE" + } + explain { + sql """ SELECT * + FROM test_colocate_mapping_constraint_left l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 """ + notContains "COLOCATE" + } + explain { + sql """ SELECT * + FROM test_colocate_mapping_constraint_left l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.d2 = r.d2 """ + notContains "COLOCATE" + } + explain { + sql """ SELECT * + FROM test_colocate_mapping_constraint_left l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 AND l.extra_col = r.extra_col """ + notContains "COLOCATE" + } + explain { + sql """ SELECT l.d1, l.k2, SUM(l.extra_col + r.extra_col) + FROM test_colocate_mapping_constraint_left l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 + GROUP BY l.d1, l.k2 """ + notContains "COLOCATE" + } + explain { + sql """ SELECT l.d1, r.d1 + FROM test_colocate_mapping_constraint_left l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 """ + notContains "COLOCATE" + } + explain { + sql """ SELECT * + FROM test_colocate_mapping_constraint_left l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 """ + notContains "COLOCATE" + } + explain { + sql """ SELECT * + FROM ( + SELECT k1 AS aggregate_k1, + k2 AS aggregate_k2, + d1 AS aggregate_d1, + SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY k1, k2, d1 + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.aggregate_d1 = r.d1 AND l.aggregate_k2 = r.k2 """ + notContains "COLOCATE" + } + explain { + sql """ SELECT * + FROM ( + SELECT k1 AS aggregate_k1, + k2 AS aggregate_k2, + d1 AS aggregate_d1, + d2 AS aggregate_d2, + SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY k1, k2, d1, d2 + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.aggregate_d1 = r.d1 AND l.aggregate_d2 = r.d2 """ + notContains "COLOCATE" + } + explain { + sql """ SELECT * + FROM ( + SELECT k2 AS aggregate_k2, + d1 AS aggregate_d1, + SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY k1, k2, d1 + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.aggregate_d1 = r.d1 AND l.aggregate_k2 = r.k2 """ + notContains "COLOCATE" + } + explain { + sql """ SELECT * + FROM ( + SELECT d1 AS aggregate_d1, + d2 AS aggregate_d2, + SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY k1, k2, d1, d2 + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.aggregate_d1 = r.d1 AND l.aggregate_d2 = r.d2 """ + notContains "COLOCATE" + } + explain { + sql """ SELECT * + FROM ( + SELECT d1, d2, SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY d1, d2 + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.d2 = r.d2 """ + notContains "COLOCATE" + } + explain { + sql """ SELECT * + FROM ( + SELECT d1, d2, k2, SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_composite_left + GROUP BY d1, d2, k2 + ) l + JOIN test_colocate_mapping_composite_right r + ON l.d1 = r.d1 AND l.d2 = r.d2 AND l.k2 = r.k2 """ + notContains "COLOCATE" + } + explain { + sql """ SELECT * + FROM ( + SELECT k2, d1, SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY k2, d1 + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 """ + notContains "COLOCATE" + } + explain { + sql """ SELECT * + FROM ( + SELECT k1, k2, d1, COUNT(DISTINCT extra_col) AS distinct_extra + FROM test_colocate_mapping_constraint_left + GROUP BY k1, k2, d1 + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 """ + notContains "COLOCATE" + } + explain { + sql """ SELECT * + FROM ( + SELECT k2, d1, COUNT(DISTINCT extra_col) AS distinct_extra + FROM test_colocate_mapping_constraint_left + GROUP BY k2, d1 + ) l + JOIN ( + SELECT k2, d1, COUNT(DISTINCT extra_col) AS distinct_extra + FROM test_colocate_mapping_constraint_right + GROUP BY k2, d1 + ) r + ON l.d1 = r.d1 AND l.k2 = r.k2 """ + notContains "COLOCATE" + } + explain { + sql """ SELECT * + FROM ( + SELECT k1, k2, d1, SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY GROUPING SETS ((k1, k2, d1), (k1, k2)) + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 """ + notContains "COLOCATE" + } + explain { + sql """ SELECT * + FROM ( + SELECT k1, k2, d1 + 0 AS d1_expression, SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY k1, k2, d1 + 0 + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.d1_expression = r.d1 AND l.k2 = r.k2 """ + notContains "COLOCATE" + } + explain { + sql """ SELECT * + FROM ( + SELECT k1, k2, d1 + FROM test_colocate_mapping_constraint_left + UNION ALL + SELECT k1, k2, d1 + FROM test_colocate_mapping_constraint_left + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 """ + notContains "COLOCATE" + } + + sql """ SET enable_colocate_mapping_constraint = true """ + nestedSubqueryJoinQueries.each { query -> + explain { + sql query + contains "COLOCATE" + } + } + // Cases supported before Aggregate propagation. + explain { + sql """ SELECT * + FROM test_colocate_mapping_constraint_left l + JOIN test_colocate_mapping_constraint_right r + ON l.k1 = r.k1 AND l.k2 = r.k2 """ + contains "COLOCATE" + } + explain { + sql """ SELECT * + FROM test_colocate_mapping_constraint_left l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 """ + contains "COLOCATE" + } + ["LEFT", "RIGHT", "FULL"].each { outerJoinType -> + explain { + sql """ SELECT * + FROM test_colocate_mapping_constraint_left l + ${outerJoinType} OUTER JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 """ + contains "COLOCATE" + } + } + explain { + sql """ SELECT * + FROM test_colocate_mapping_constraint_left l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.d2 = r.d2 """ + contains "COLOCATE" + } + explain { + sql """ SELECT * + FROM test_colocate_mapping_constraint_left l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 AND l.extra_col = r.extra_col """ + contains "COLOCATE" + } + explain { + sql """ SELECT l.d1, l.k2, SUM(l.extra_col + r.extra_col) + FROM test_colocate_mapping_constraint_left l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 + GROUP BY l.d1, l.k2 """ + contains "COLOCATE" + } + explain { + sql """ SELECT * + FROM test_colocate_mapping_constraint_left l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 """ + notContains "COLOCATE" + } + explain { + sql """ SELECT l.d1, r.d1 + FROM test_colocate_mapping_constraint_left l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 """ + contains "COLOCATE" + } + // Cases supported by the first conservative Aggregate propagation. + explain { + sql """ SELECT * + FROM ( + SELECT k1 AS aggregate_k1, + k2 AS aggregate_k2, + d1 AS aggregate_d1, + SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY k1, k2, d1 + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.aggregate_d1 = r.d1 AND l.aggregate_k2 = r.k2 """ + contains "COLOCATE" + } + explain { + sql """ SELECT * + FROM ( + SELECT k1 AS aggregate_k1, + k2 AS aggregate_k2, + d1 AS aggregate_d1, + d2 AS aggregate_d2, + SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY k1, k2, d1, d2 + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.aggregate_d1 = r.d1 AND l.aggregate_d2 = r.d2 """ + contains "COLOCATE" + } + // Cases supported after carrying hidden natural bucket locality. + explain { + sql """ SELECT * + FROM ( + SELECT k2 AS aggregate_k2, + d1 AS aggregate_d1, + SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY k1, k2, d1 + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.aggregate_d1 = r.d1 AND l.aggregate_k2 = r.k2 """ + contains "COLOCATE" + } + explain { + sql """ SELECT * + FROM ( + SELECT d1 AS aggregate_d1, + d2 AS aggregate_d2, + SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY k1, k2, d1, d2 + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.aggregate_d1 = r.d1 AND l.aggregate_d2 = r.d2 """ + contains "COLOCATE" + } + // Multiple mappings can replace every distribution key in Aggregate Group By. + explain { + sql """ SELECT * + FROM ( + SELECT d1, d2, SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY d1, d2 + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.d2 = r.d2 """ + contains "COLOCATE" + } + // A composite mapping determinant must be complete. + explain { + sql """ SELECT * + FROM ( + SELECT d1, d2, k2, SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_composite_left + GROUP BY d1, d2, k2 + ) l + JOIN test_colocate_mapping_composite_right r + ON l.d1 = r.d1 AND l.d2 = r.d2 AND l.k2 = r.k2 """ + contains "COLOCATE" + } + explain { + sql """ SELECT * + FROM ( + SELECT d1, k2, MAX(d2) AS d2, SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_composite_left + GROUP BY d1, k2 + ) l + JOIN test_colocate_mapping_composite_right r + ON l.d1 = r.d1 AND l.d2 = r.d2 AND l.k2 = r.k2 """ + notContains "COLOCATE" + } + // DISTINCT phases can repartition by deduplication keys, so they must not expose storage bucket locality. + def distinctAggregateBarrierSql = """ + SELECT /*+ SET_VAR(disable_join_reorder=true, enable_local_shuffle_planner=true, + ignore_storage_data_distribution=true, parallel_pipeline_task_num=4) */ + l.k2, l.d1, l.distinct_extra, r.k1, r.k2, r.d1 + FROM ( + SELECT k2, d1, COUNT(DISTINCT extra_col) AS distinct_extra + FROM test_colocate_mapping_constraint_left + GROUP BY k2, d1 + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 + ORDER BY l.k2, l.d1, r.k1 + """ + explain { + sql distinctAggregateBarrierSql + notContains "COLOCATE" + } + order_qt_distinct_aggregate_mapping_barrier distinctAggregateBarrierSql + + explain { + sql """ SELECT * + FROM ( + SELECT k2, d1, COUNT(DISTINCT extra_col) AS distinct_extra + FROM test_colocate_mapping_constraint_left + GROUP BY k2, d1 + ) l + JOIN ( + SELECT k2, d1, COUNT(DISTINCT extra_col) AS distinct_extra + FROM test_colocate_mapping_constraint_right + GROUP BY k2, d1 + ) r + ON l.d1 = r.d1 AND l.k2 = r.k2 """ + notContains "COLOCATE" + } + + def pureDeduplicateBarrierSql = """ + SELECT /*+ SET_VAR(disable_join_reorder=true, enable_local_shuffle_planner=true, + ignore_storage_data_distribution=true, parallel_pipeline_task_num=4) */ + l.k2, l.d1, r.k1, r.k2, r.d1 + FROM ( + SELECT DISTINCT k2, d1 + FROM test_colocate_mapping_constraint_left + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 + ORDER BY l.k2, l.d1, r.k1 + """ + explain { + sql pureDeduplicateBarrierSql + notContains "COLOCATE" + } + order_qt_pure_deduplicate_mapping_barrier pureDeduplicateBarrierSql + def mixedDistinctAggregateBarrierSql = """ + SELECT /*+ SET_VAR(disable_join_reorder=true, enable_local_shuffle_planner=true, + ignore_storage_data_distribution=true, parallel_pipeline_task_num=4) */ + l.d1, l.k2, l.sum_extra, l.distinct_extra, r.k1, r.k2, r.d1 + FROM ( + SELECT d1, MAX(k2) AS k2, SUM(extra_col) AS sum_extra, + COUNT(DISTINCT extra_col) AS distinct_extra + FROM test_colocate_mapping_constraint_left + GROUP BY d1 + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 + ORDER BY l.d1, l.k2, r.k1 + """ + explain { + sql mixedDistinctAggregateBarrierSql + notContains "COLOCATE" + } + order_qt_mixed_distinct_aggregate_mapping_barrier mixedDistinctAggregateBarrierSql + // Mapping determinants can cover distribution keys that are absent from Group By. + explain { + sql """ SELECT /*+ SET_VAR(disable_join_reorder=true, enable_local_shuffle_planner=true, + ignore_storage_data_distribution=true, parallel_pipeline_task_num=4) */ * + FROM ( + SELECT k2, d1, SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY k2, d1 + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 """ + contains "COLOCATE" + } + explain { + sql """ SELECT * + FROM ( + SELECT k2, d1, SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY k2, d1 + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 + AND l.k2 = r.k2 + AND l.sum_extra = r.extra_col """ + contains "COLOCATE" + } + // Unsupported Aggregate shapes must discard mapping locality. + explain { + sql """ SELECT * + FROM ( + SELECT k1, k2, d1, SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY GROUPING SETS ((k1, k2, d1), (k1, k2)) + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 """ + notContains "COLOCATE" + } + explain { + sql """ SELECT * + FROM ( + SELECT k1, k2, d1 + 0 AS d1_expression, SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY k1, k2, d1 + 0 + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.d1_expression = r.d1 AND l.k2 = r.k2 """ + notContains "COLOCATE" + } + // A redistribution before Aggregate cuts the storage bucket locality. + explain { + sql """ SELECT * + FROM ( + SELECT d1, k2, SUM(extra_col) AS sum_extra + FROM ( + SELECT d1, k2, extra_col + FROM test_colocate_mapping_constraint_left + ORDER BY extra_col + LIMIT 10 + ) ordered_l + GROUP BY d1, k2 + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 """ + notContains "COLOCATE" + } + // Removing the determinant from Aggregate output prevents the parent Join proof. + explain { + sql """ SELECT * + FROM ( + SELECT k2, SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY d1, k2 + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.sum_extra = r.d1 AND l.k2 = r.k2 """ + notContains "COLOCATE" + } + explain { + sql """ SELECT * + FROM ( + SELECT k1, k2, d1 + FROM test_colocate_mapping_constraint_left + UNION ALL + SELECT k1, k2, d1 + FROM test_colocate_mapping_constraint_left + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 """ + notContains "COLOCATE" + } + // Union does not merge the natural mapping locality of Aggregate branches. + explain { + sql """ SELECT * + FROM ( + SELECT d1, k2, SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY d1, k2 + UNION ALL + SELECT d1, k2, SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY d1, k2 + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 """ + notContains "COLOCATE" + } + // INTERSECT and EXCEPT require their own hash shuffle and cannot accept a hidden mapping request. + explain { + sql """ SELECT * + FROM ( + SELECT d1, k2, SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY d1, k2 + INTERSECT + SELECT d1, k2, SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY d1, k2 + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 """ + notContains "COLOCATE" + } + explain { + sql """ SELECT * + FROM ( + SELECT d1, k2, SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY d1, k2 + EXCEPT + SELECT d1, k2, SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY d1, k2 + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 """ + notContains "COLOCATE" + } + + // Runtime placement barriers must not expose storage bucket locality to an outer mapping join. + def generateBarrierSql = """ + SELECT /*+ SET_VAR(disable_join_reorder=true, enable_local_shuffle_planner=true, + parallel_pipeline_task_num=4) */ + l.d1, l.k2, l.e, r.k1 + FROM ( + SELECT d1, k2, e + FROM test_colocate_mapping_constraint_left + LATERAL VIEW explode_numbers(2) generated AS e + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 + ORDER BY l.d1, l.k2, l.e + """ + explain { + sql generateBarrierSql + notContains "COLOCATE" + } + order_qt_generate_barrier_result generateBarrierSql + + def windowBarrierSql = """ + SELECT /*+ SET_VAR(disable_join_reorder=true, enable_local_shuffle_planner=true, + parallel_pipeline_task_num=4) */ + l.d1, l.k2, l.rn, r.k1 + FROM ( + SELECT d1, k2, + ROW_NUMBER() OVER (PARTITION BY k1, k2 ORDER BY extra_col) AS rn + FROM test_colocate_mapping_constraint_left + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 + ORDER BY l.d1, l.k2 + """ + explain { + sql windowBarrierSql + notContains "COLOCATE" + } + order_qt_window_barrier_result windowBarrierSql + + def partitionTopNBarrierSql = """ + SELECT /*+ SET_VAR(disable_join_reorder=true, enable_local_shuffle_planner=true, + parallel_pipeline_task_num=4) */ + l.d1, l.k2, l.rn, r.k1 + FROM ( + SELECT d1, k2, + ROW_NUMBER() OVER (PARTITION BY k1, k2 ORDER BY extra_col) AS rn + FROM test_colocate_mapping_constraint_left + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 + WHERE l.rn <= 1 + ORDER BY l.d1, l.k2 + """ + explain { + sql partitionTopNBarrierSql + notContains "COLOCATE" + } + order_qt_partition_topn_barrier_result partitionTopNBarrierSql + + def nestedLoopBarrierSql = """ + SELECT /*+ SET_VAR(disable_join_reorder=true, enable_local_shuffle_planner=true, + parallel_pipeline_task_num=4) */ + l.d1, l.k2, r.k1 + FROM ( + SELECT mapped.d1, mapped.k2 + FROM test_colocate_mapping_constraint_left mapped + CROSS JOIN test_colocate_mapping_constraint_right other + WHERE other.k1 = 1 + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 + ORDER BY l.d1, l.k2 + """ + explain { + sql nestedLoopBarrierSql + notContains "COLOCATE" + } + order_qt_nested_loop_barrier_result nestedLoopBarrierSql + + def broadcastBarrierSql = """ + SELECT /*+ SET_VAR(disable_join_reorder=true, enable_local_shuffle_planner=true, + enable_broadcast_join_force_passthrough=true, + parallel_pipeline_task_num=4) */ + l.d1, l.k2, r.k1 + FROM ( + SELECT mapped.d1, mapped.k2 + FROM test_colocate_mapping_constraint_left mapped + JOIN [broadcast] test_colocate_mapping_constraint_right broadcast_side + ON mapped.extra_col = broadcast_side.extra_col + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 + ORDER BY l.d1, l.k2 + """ + explain { + sql broadcastBarrierSql + notContains "COLOCATE" + } + order_qt_broadcast_barrier_result broadcastBarrierSql + + order_qt_colocate_mapping_result """ + SELECT l.k1, l.k2, l.d1, l.d2, l.extra_col, + r.k1, r.k2, r.d1, r.d2, r.extra_col + FROM test_colocate_mapping_constraint_left l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 + ORDER BY l.k1, l.k2 + """ + + order_qt_aggregate_colocate_mapping_result """ + SELECT l.aggregate_k1, l.aggregate_k2, l.aggregate_d1, l.sum_extra, + r.k1, r.k2, r.d1 + FROM ( + SELECT k1 AS aggregate_k1, + k2 AS aggregate_k2, + d1 AS aggregate_d1, + SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY k1, k2, d1 + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.aggregate_d1 = r.d1 AND l.aggregate_k2 = r.k2 + ORDER BY l.aggregate_k1, l.aggregate_k2 + """ + + order_qt_hidden_distribution_key_aggregate_result """ + SELECT /*+ SET_VAR(disable_join_reorder=true, enable_local_shuffle_planner=true, + ignore_storage_data_distribution=true, parallel_pipeline_task_num=4) */ + l.aggregate_k2, l.aggregate_d1, l.sum_extra, + r.k1, r.k2, r.d1 + FROM ( + SELECT k2 AS aggregate_k2, + d1 AS aggregate_d1, + SUM(extra_col) AS sum_extra + FROM test_colocate_mapping_constraint_left + GROUP BY k1, k2, d1 + ) l + JOIN test_colocate_mapping_constraint_right r + ON l.aggregate_d1 = r.d1 AND l.aggregate_k2 = r.k2 + ORDER BY l.aggregate_k2, l.aggregate_d1 + """ + + sql """ TRUNCATE TABLE test_colocate_mapping_constraint_left """ + sql """ INSERT INTO test_colocate_mapping_constraint_left VALUES (3, 30, 300, 3000, 10) """ + sql """ SYNC """ + waitForColocateGroupStable("test_colocate_mapping_constraint_group") + explain { + sql """ + SELECT * + FROM test_colocate_mapping_constraint_left l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 + """ + contains "COLOCATE" + } + + sql """ DROP TABLE test_colocate_mapping_constraint_left """ + sql """ RECOVER TABLE test_colocate_mapping_constraint_left """ + waitForColocateGroupStable("test_colocate_mapping_constraint_group") + explain { + sql """ + SELECT * + FROM test_colocate_mapping_constraint_left l + JOIN test_colocate_mapping_constraint_right r + ON l.d1 = r.d1 AND l.k2 = r.k2 + """ + contains "COLOCATE" + } +} diff --git a/regression-test/suites/query_p0/colocate/test_colocate_mapping_constraint_being_synced.groovy b/regression-test/suites/query_p0/colocate/test_colocate_mapping_constraint_being_synced.groovy new file mode 100644 index 00000000000000..f282b70ef0ca9d --- /dev/null +++ b/regression-test/suites/query_p0/colocate/test_colocate_mapping_constraint_being_synced.groovy @@ -0,0 +1,73 @@ +// 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. + +suite("test_colocate_mapping_constraint_being_synced") { + sql "DROP TABLE IF EXISTS test_colocate_mapping_constraint_being_synced" + sql """ + CREATE TABLE test_colocate_mapping_constraint_being_synced ( + k1 INT, + d1 INT + ) ENGINE=OLAP + DUPLICATE KEY(k1) + DISTRIBUTED BY HASH(k1) BUCKETS 4 + PROPERTIES ("replication_num" = "1") + """ + sql """ + ALTER TABLE test_colocate_mapping_constraint_being_synced + ADD CONSTRAINT mapping_before_sync + COLOCATE MAPPING being_synced_mapping (d1) + DETERMINES DISTRIBUTION KEY (k1) NOT ENFORCED + """ + sql "INSERT INTO test_colocate_mapping_constraint_being_synced VALUES (1, 1)" + sql "SYNC" + + def selfJoinSql = """ + SELECT /*+ SET_VAR(disable_join_reorder=true, + enable_colocate_mapping_constraint=true, + auto_broadcast_join_threshold=-1, + broadcast_row_count_limit=0) */ l.d1 + FROM test_colocate_mapping_constraint_being_synced l + JOIN test_colocate_mapping_constraint_being_synced r ON l.d1 = r.d1 + """ + explain { + sql selfJoinSql + contains "COLOCATE" + } + + sql """ + ALTER TABLE test_colocate_mapping_constraint_being_synced + SET ("is_being_synced" = "true") + """ + explain { + sql selfJoinSql + notContains "COLOCATE" + } + test { + sql """ + ALTER TABLE test_colocate_mapping_constraint_being_synced + ADD CONSTRAINT mapping_during_sync + COLOCATE MAPPING being_synced_mapping_2 (d1) + DETERMINES DISTRIBUTION KEY (k1) NOT ENFORCED + """ + exception "being synchronized by CCR" + } + + sql """ + ALTER TABLE test_colocate_mapping_constraint_being_synced + DROP CONSTRAINT mapping_before_sync + """ +}