fix

Signed-off-by: gengjun-git <gengjun@starrocks.com>

Author: gengjun-git

fe/fe-core/src/main/java/com/starrocks/planner/OlapTableSink.java

@@ -859,6 +859,12 @@ private static DistributionInfo setDistributedColumns(TOlapTablePartitionParam p
                 throw new StarRocksException("different distribute types in two different partitions, type1="
                         + selectedDistInfo.getType() + ", type2=" + distInfo.getType());
             }
+            List<String> selectedDistColumns = getDistColumns(selectedDistInfo, table);
+            List<String> currentDistColumns = getDistColumns(distInfo, table);
+            if (!selectedDistColumns.equals(currentDistColumns)) {
+                throw new StarRocksException("different distribute columns in two different partitions, columns1="
+                        + selectedDistColumns + ", columns2=" + currentDistColumns);
+            }
         }
         return selectedDistInfo;
     }

fe/fe-core/src/test/java/com/starrocks/planner/OlapTableSinkTest2.java

@@ -14,7 +14,11 @@
 
 package com.starrocks.planner;
 
+import com.google.common.collect.Lists;
+import com.starrocks.catalog.Column;
 import com.starrocks.catalog.Database;
+import com.starrocks.catalog.DistributionInfo;
+import com.starrocks.catalog.HashDistributionInfo;
 import com.starrocks.catalog.OlapTable;
 import com.starrocks.catalog.Partition;
 import com.starrocks.catalog.PartitionInfo;
@@ -26,6 +30,7 @@
 import com.starrocks.thrift.TOlapTablePartition;
 import com.starrocks.thrift.TOlapTablePartitionParam;
 import com.starrocks.thrift.TWriteQuorumType;
+import com.starrocks.type.IntegerType;
 import com.starrocks.utframe.StarRocksAssert;
 import com.starrocks.utframe.UtFrameUtils;
 import mockit.Mock;
@@ -34,6 +39,9 @@
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 
+import java.util.List;
+import java.util.stream.Collectors;
+
 public class OlapTableSinkTest2 {
     private static StarRocksAssert starRocksAssert;
     private static ConnectContext connectContext;
@@ -47,6 +55,22 @@ public static void beforeClass() throws Exception {
         starRocksAssert = new StarRocksAssert(connectContext);
         starRocksAssert.withDatabase("db2");
         starRocksAssert.withTable(createTblStmtStr);
+
+        String createRangeTableStmt = "create table db2.tbl_range(k1 int, k2 int, v1 int) " +
+                "DUPLICATE KEY(k1) " +
+                "PARTITION BY RANGE(k1) (" +
+                "  PARTITION p1 VALUES LESS THAN ('10')," +
+                "  PARTITION p2 VALUES LESS THAN ('20')" +
+                ") DISTRIBUTED BY HASH(k1) BUCKETS 3 PROPERTIES('replication_num' = '1');";
+        starRocksAssert.withTable(createRangeTableStmt);
+
+        String createListTableStmt = "create table db2.tbl_list(k1 int, k2 int, v1 int) " +
+                "DUPLICATE KEY(k1) " +
+                "PARTITION BY LIST(k1) (" +
+                "  PARTITION p1 VALUES IN ('1', '2', '3')," +
+                "  PARTITION p2 VALUES IN ('4', '5', '6')" +
+                ") DISTRIBUTED BY HASH(k1) BUCKETS 3 PROPERTIES('replication_num' = '1');";
+        starRocksAssert.withTable(createListTableStmt);
     }
 
     @Test
@@ -79,4 +103,69 @@ public int getQuorumNum(long partitionId, TWriteQuorumType writeQuorum) {
         }
         Assertions.fail("must throw UserException");
     }
+
+    @Test
+    public void testRangePartitionWithDifferentDistributionColumns() {
+        Database db = GlobalStateMgr.getCurrentState().getLocalMetastore().getDb("db2");
+        OlapTable olapTable = (OlapTable) GlobalStateMgr.getCurrentState().getLocalMetastore()
+                .getTable(db.getFullName(), "tbl_range");
+
+        // Get two partitions
+        Partition p1 = olapTable.getPartition("p1");
+        Partition p2 = olapTable.getPartition("p2");
+        Assertions.assertNotNull(p1);
+        Assertions.assertNotNull(p2);
+
+        // Save original distribution info for p2 to restore later
+        DistributionInfo originalDistInfo = p2.getDistributionInfo();
+
+        try {
+            // Change p2's distribution columns to k2 (different from p1's k1)
+            HashDistributionInfo differentDistInfo = new HashDistributionInfo(
+                    3, Lists.newArrayList(new Column("k2", IntegerType.INT)));
+            p2.setDistributionInfo(differentDistInfo);
+
+            List<Long> partitionIds = olapTable.getPartitions().stream()
+                    .map(Partition::getId).collect(Collectors.toList());
+
+            StarRocksException exception = Assertions.assertThrows(StarRocksException.class, () -> {
+                OlapTableSink.createPartition(db.getId(), olapTable, null,
+                        false, 0, partitionIds, null);
+            });
+            Assertions.assertTrue(exception.getMessage().contains("different distribute columns"));
+        } finally {
+            p2.setDistributionInfo(originalDistInfo);
+        }
+    }
+
+    @Test
+    public void testListPartitionWithDifferentDistributionColumns() {
+        Database db = GlobalStateMgr.getCurrentState().getLocalMetastore().getDb("db2");
+        OlapTable olapTable = (OlapTable) GlobalStateMgr.getCurrentState().getLocalMetastore()
+                .getTable(db.getFullName(), "tbl_list");
+
+        Partition p1 = olapTable.getPartition("p1");
+        Partition p2 = olapTable.getPartition("p2");
+        Assertions.assertNotNull(p1);
+        Assertions.assertNotNull(p2);
+
+        DistributionInfo originalDistInfo = p2.getDistributionInfo();
+
+        try {
+            HashDistributionInfo differentDistInfo = new HashDistributionInfo(
+                    3, Lists.newArrayList(new Column("k2", IntegerType.INT)));
+            p2.setDistributionInfo(differentDistInfo);
+
+            List<Long> partitionIds = olapTable.getPartitions().stream()
+                    .map(Partition::getId).collect(Collectors.toList());
+
+            StarRocksException exception = Assertions.assertThrows(StarRocksException.class, () -> {
+                OlapTableSink.createPartition(db.getId(), olapTable, null,
+                        false, 0, partitionIds, null);
+            });
+            Assertions.assertTrue(exception.getMessage().contains("different distribute columns"));
+        } finally {
+            p2.setDistributionInfo(originalDistInfo);
+        }
+    }
 }