[feature](fe) Add constraint-based colocate join inference with distribution mappings

fe/fe-core/src/main/java/org/apache/doris/alter/SchemaChangeHandler.java

@@ -543,6 +543,17 @@ private void checkSeqMapConditionMet(OlapTable olapTable, Map<String, String> 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<Column> newSchema)
             throws DdlException {
         String dropColName = dropColumnOp.getColName();
@@ -594,19 +605,21 @@ 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'. "
                             + "Drop the constraint first.",
                     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();

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()

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<Long, MaterializedIndexMeta> indexIdToMeta = table.getIndexIdToMeta();
         for (Map.Entry<Long, MaterializedIndexMeta> 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<String, String> 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");
             }

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<String, String> properties;
 
+    // Derived from DISTRIBUTION_MAPPING_CONSTRAINTS_PROPERTY. The property is the persistent source of truth.
+    private Map<String, DistributionMappingConstraint> 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<String, String> getProperties() {
         return properties;
     }
 
+    public Map<String, DistributionMappingConstraint> getDistributionMappingConstraints() {
+        return distributionMappingConstraints;
+    }
+
+    public void addDistributionMappingConstraint(DistributionMappingConstraint constraint) {
+        Map<String, DistributionMappingConstraint> updated = new HashMap<>(distributionMappingConstraints);
+        updated.put(constraint.getName(), constraint);
+        updateDistributionMappingConstraints(updated);
+    }
+
+    public DistributionMappingConstraint removeDistributionMappingConstraint(String constraintName) {
+        Map<String, DistributionMappingConstraint> updated = new HashMap<>(distributionMappingConstraints);
+        DistributionMappingConstraint removed = updated.remove(constraintName);
+        if (removed != null) {
+            updateDistributionMappingConstraints(updated);
+        }
+        return removed;
+    }
+
+    public Map<String, String> getDistributionMappingConstraintProperties() {
+        return ImmutableMap.of(DISTRIBUTION_MAPPING_CONSTRAINTS_PROPERTY,
+                properties.get(DISTRIBUTION_MAPPING_CONSTRAINTS_PROPERTY));
+    }
+
+    private void updateDistributionMappingConstraints(
+            Map<String, DistributionMappingConstraint> updatedConstraints) {
+        List<DistributionMappingConstraint> 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<String, DistributionMappingConstraint> 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<String, DistributionMappingConstraint> 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();

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");