[BugFix] Name a partition column's key range by column id, not by column name (#77988)

(cherry picked from commit e843b3c)

# Conflicts:
#	fe/fe-core/src/main/java/com/starrocks/sql/plan/PlanFragmentBuilder.java

Author: Hyper-FF

fe/fe-core/src/main/java/com/starrocks/sql/plan/PlanFragmentBuilder.java

@@ -1086,6 +1086,193 @@ public PlanFragment visitPhysicalOlapScan(OptExpression optExpr, ExecPlan contex
             return fragment;
         }
 
+<<<<<<< HEAD
+=======
+        /**
+         * Compute partition key ranges for dynamic partition pruning.
+         * Returns a list of TKeyRange describing the value ranges of each used partition column.
+         * Returns empty if the partition has too many values (exceeding the session limit) or
+         * if the partition type is unsupported.
+         */
+        private List<TKeyRange> computePartitionRange(OlapTable table, Partition partition,
+                                                      Collection<Column> usedPartitionCols, SessionVariable session) {
+            PartitionInfo partitionInfo = table.getPartitionInfo();
+            if (usedPartitionCols.isEmpty() || !partition.hasData()
+                    || !(partitionInfo.isRangePartition() || partitionInfo.isListPartition())) {
+                return List.of();
+            }
+            List<Column> partitionCols = partitionInfo.getPartitionColumns(table.getIdToColumn());
+            Preconditions.checkState(partitionCols.containsAll(usedPartitionCols));
+
+            long limit = session.getDynamicPartitionPruneValuesLimit();
+            if (partitionInfo.isRangePartition()) {
+                return computeRangePartitionKeyRanges(partitionInfo, partition, partitionCols, usedPartitionCols, limit);
+            } else {
+                return computeListPartitionKeyRanges(partitionInfo, partition, partitionCols, usedPartitionCols, limit);
+            }
+        }
+    
+        private List<TKeyRange> computeRangePartitionKeyRanges(PartitionInfo partitionInfo, Partition partition,
+                                                               List<Column> partitionCols,
+                                                               Collection<Column> usedPartitionCols, long limit) {
+            RangePartitionInfo rangeInfo = (RangePartitionInfo) partitionInfo;
+            Range<PartitionKey> keyRange = rangeInfo.getRange(partition.getId());
+            if (!keyRange.hasLowerBound() || !keyRange.hasUpperBound()) {
+                return List.of();
+            }
+
+            boolean isNullPartition = keyRange.lowerEndpoint().isMinValue();
+            long partitionValues = 1;
+            List<TKeyRange> result = Lists.newArrayList();
+    
+            for (int i = 0; i < partitionCols.size(); i++) {
+                Column col = partitionCols.get(i);
+                if (!usedPartitionCols.contains(col)) {
+                    continue;
+                }
+    
+                TKeyRange kr = new TKeyRange();
+                long rangeSize;
+
+                if (col.getType().isDate()) {
+                    LiteralExpr lowerExpr = keyRange.lowerEndpoint().getKeys().get(i);
+                    LiteralExpr upperExpr = keyRange.upperEndpoint().getKeys().get(i);
+                    if (!(lowerExpr instanceof DateLiteral lower) || !(upperExpr instanceof DateLiteral upper)) {
+                        continue;
+                    }
+                    kr.setBegin_key(lower.getYear() * 10000 + lower.getMonth() * 100 + lower.getDay());
+                    kr.setEnd_key(upper.getYear() * 10000 + upper.getMonth() * 100 + upper.getDay());
+                    rangeSize = upper.toLocalDateTime().toLocalDate().toEpochDay()
+                            - lower.toLocalDateTime().toLocalDate().toEpochDay();
+                } else if (col.getType().isIntegerType()) {
+                    long lowerVal = keyRange.lowerEndpoint().getKeys().get(i).getLongValue();
+                    long upperVal = keyRange.upperEndpoint().getKeys().get(i).getLongValue();
+                    kr.setBegin_key(lowerVal);
+                    kr.setEnd_key(upperVal);
+                    rangeSize = upperVal - lowerVal;
+                } else {
+                    continue;
+                }
+
+                if (rangeSize <= 0 || wouldOverflowOrExceedLimit(partitionValues, rangeSize, limit)) {
+                    break;
+                }
+                partitionValues *= rangeSize;
+    
+                kr.setColumn_type(TypeSerializer.toThrift(col.getType().getPrimitiveType()));
+                // BE indexes the tuple's slots by col_name, which is the column id, so name the range by
+                // the id as well: a renamed partition column would otherwise be skipped and its
+                // scan ranges never pruned.
+                kr.setColumn_name(col.getColumnId().getId());
+                if (isNullPartition) {
+                    kr.setHas_null(true);
+                }
+                result.add(kr);
+            }
+
+            return result;
+        }
+    
+        private List<TKeyRange> computeListPartitionKeyRanges(PartitionInfo partitionInfo, Partition partition,
+                                                              List<Column> partitionCols,
+                                                              Collection<Column> usedPartitionCols, long limit) {
+            ListPartitionInfo listInfo = (ListPartitionInfo) partitionInfo;
+            if (listInfo.getLiteralExprValues().containsKey(partition.getId())) {
+                return computeSingleColumnListKeyRanges(listInfo, partition, partitionCols, usedPartitionCols, limit);
+            } else if (listInfo.getMultiLiteralExprValues().containsKey(partition.getId())) {
+                return computeMultiColumnListKeyRanges(listInfo, partition, partitionCols, usedPartitionCols, limit);
+            } else {
+                return List.of();
+            }
+        }
+    
+        private List<TKeyRange> computeSingleColumnListKeyRanges(ListPartitionInfo listInfo, Partition partition,
+                                                                 List<Column> partitionCols,
+                                                                 Collection<Column> usedPartitionCols, long limit) {
+            Preconditions.checkState(partitionCols.size() == 1);
+            List<LiteralExpr> partitionValuesList = listInfo.getLiteralExprValues().get(partition.getId());
+            long partitionValues = 1;
+            List<TKeyRange> result = Lists.newArrayList();
+    
+            for (Column col : partitionCols) {
+                if (!usedPartitionCols.contains(col)) {
+                    continue;
+                }
+
+                long listSize = partitionValuesList.size();
+                if (wouldOverflowOrExceedLimit(partitionValues, listSize, limit)) {
+                    break;
+                }
+                partitionValues *= listSize;
+
+                TKeyRange kr = new TKeyRange();
+                kr.setColumn_type(TypeSerializer.toThrift(col.getType().getPrimitiveType()));
+                // BE indexes the tuple's slots by col_name, which is the column id, so name the range by
+                // the id as well: a renamed partition column would otherwise be skipped and its
+                // scan ranges never pruned.
+                kr.setColumn_name(col.getColumnId().getId());
+                List<TExpr> l = Lists.newArrayList();
+                partitionValuesList.forEach(v -> l.add(ExprToThrift.treeToThrift(v)));
+                kr.setList_values(l);
+                result.add(kr);
+            }
+
+            return result;
+        }
+    
+        private List<TKeyRange> computeMultiColumnListKeyRanges(ListPartitionInfo listInfo, Partition partition,
+                                                                List<Column> partitionCols,
+                                                                Collection<Column> usedPartitionCols, long limit) {
+            List<List<LiteralExpr>> partitionValuesList = listInfo.getMultiLiteralExprValues().get(partition.getId());
+            long partitionValues = 1;
+            List<TKeyRange> result = Lists.newArrayList();
+
+            for (int i = 0; i < partitionCols.size(); i++) {
+                Column col = partitionCols.get(i);
+                if (!usedPartitionCols.contains(col)) {
+                    continue;
+                }
+
+                long listSize = partitionValuesList.size();
+                if (wouldOverflowOrExceedLimit(partitionValues, listSize, limit)) {
+                    break;
+                }
+                partitionValues *= listSize;
+
+                TKeyRange kr = new TKeyRange();
+                kr.setColumn_type(TypeSerializer.toThrift(col.getType().getPrimitiveType()));
+                // BE indexes the tuple's slots by col_name, which is the column id, so name the range by
+                // the id as well: a renamed partition column would otherwise be skipped and its
+                // scan ranges never pruned.
+                kr.setColumn_name(col.getColumnId().getId());
+                List<TExpr> l = Lists.newArrayList();
+                for (var values : partitionValuesList) {
+                    Preconditions.checkState(values.size() == partitionCols.size());
+                    l.add(ExprToThrift.treeToThrift(values.get(i)));
+                }
+                kr.setList_values(l);
+                result.add(kr);
+            }
+
+            return result;
+        }
+
+        /**
+         * Check if multiplying {@code current} by {@code factor} would overflow {@code long}
+         * or if the product would exceed the given {@code limit}.
+         */
+        private static boolean wouldOverflowOrExceedLimit(long current, long factor, long limit) {
+            if (factor == 0) {
+                return false;
+            }
+            // Overflow check: current * factor > Long.MAX_VALUE
+            if (current > limit / factor) {
+                return true;
+            }
+            return current * factor > limit;
+        }
+
+>>>>>>> e843b3cab2 ([BugFix] Name a partition column's key range by column id, not by column name (#77988))
         @NotNull
         private static Map<Integer, Expr> getGlobalDictsExprs(Map<Integer, ScalarOperator> dictExprs,
                                                               ExecPlan context) {

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

@@ -0,0 +1,95 @@
+// Copyright 2021-present StarRocks, Inc. All rights reserved.
+//
+// Licensed 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
+//
+//     https://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 com.starrocks.planner;
+
+import com.google.common.collect.Lists;
+import com.starrocks.common.FeConstants;
+import com.starrocks.common.Pair;
+import com.starrocks.common.util.UUIDUtil;
+import com.starrocks.qe.ConnectContext;
+import com.starrocks.sql.plan.ExecPlan;
+import com.starrocks.thrift.TInternalScanRange;
+import com.starrocks.thrift.TKeyRange;
+import com.starrocks.thrift.TScanRangeLocations;
+import com.starrocks.utframe.StarRocksAssert;
+import com.starrocks.utframe.UtFrameUtils;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+
+/**
+ * A scan range carries the value range of each partition column it was cut by, and BE's dynamic
+ * partition pruning evaluates the partition conjuncts against those values to drop scan ranges the
+ * predicate cannot match. It finds the column those values belong to in a map keyed by the slot's
+ * col_name - the storage-side column id - so the range has to name the column by its id as well.
+ */
+public class PartitionKeyRangeColumnIdTest {
+    private static ConnectContext connectContext;
+    private static StarRocksAssert starRocksAssert;
+
+    @BeforeAll
+    public static void setUp() throws Exception {
+        FeConstants.runningUnitTest = true;
+        UtFrameUtils.createMinStarRocksCluster();
+        connectContext = UtFrameUtils.createDefaultCtx();
+        connectContext.setQueryId(UUIDUtil.genUUID());
+        starRocksAssert = new StarRocksAssert(connectContext);
+        starRocksAssert.withDatabase("test_pkr").useDatabase("test_pkr")
+                .withTable("create table pkr(dt date not null, v int) duplicate key(dt)" +
+                        " partition by range(dt) (" +
+                        "  partition p1 values [('2026-01-01'), ('2026-01-02'))," +
+                        "  partition p2 values [('2026-01-02'), ('2026-01-03')))" +
+                        " distributed by hash(v) buckets 1 properties('replication_num' = '1');");
+    }
+
+    @Test
+    public void testPartitionKeyRangeNamesTheColumnById() throws Exception {
+        // dayofmonth() is not something FE can prune partitions by, so both partitions survive into
+        // the plan and BE gets a chance to drop their scan ranges - the case this naming decides.
+        String sql = "select v from test_pkr.pkr where dayofmonth(%s) = 1";
+        Assertions.assertEquals(List.of("dt", "dt"), rangeColumnNames(String.format(sql, "dt")));
+
+        starRocksAssert.ddl("alter table pkr rename column dt to dt_new");
+        try {
+            Assertions.assertEquals(List.of("dt", "dt"), rangeColumnNames(String.format(sql, "dt_new")),
+                    "a partition column range must name its column by the id, which the rename left "
+                            + "alone, or BE's column_name_to_slot lookup misses it and prunes nothing");
+        } finally {
+            starRocksAssert.ddl("alter table pkr rename column dt_new to dt");
+        }
+    }
+
+    private static List<String> rangeColumnNames(String sql) throws Exception {
+        connectContext.setQueryId(UUIDUtil.genUUID());
+        connectContext.setExecutionId(UUIDUtil.toTUniqueId(connectContext.getQueryId()));
+        Pair<String, ExecPlan> plan = UtFrameUtils.getPlanAndFragment(connectContext, sql);
+        List<String> names = Lists.newArrayList();
+        for (ScanNode scanNode : plan.second.getScanNodes()) {
+            for (TScanRangeLocations locations : scanNode.getScanRangeLocations(0)) {
+                TInternalScanRange range = locations.getScan_range().getInternal_scan_range();
+                if (range == null || range.getPartition_column_ranges() == null) {
+                    continue;
+                }
+                for (TKeyRange keyRange : range.getPartition_column_ranges()) {
+                    names.add(keyRange.getColumn_name());
+                }
+            }
+        }
+        Assertions.assertFalse(names.isEmpty(), "the plan must carry partition column ranges at all");
+        return names;
+    }
+}