planner: Fix BatchPointGetPlan.HandleColOffset (#54686)

close #54667
pingcap · Jul 22, 2024 · 669bf4d · 669bf4d
1 parent 0311052
commit 669bf4d
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 1 deletion.
diff --git a/pkg/planner/core/find_best_task.go b/pkg/planner/core/find_best_task.go
@@ -2699,7 +2699,18 @@ func (ds *DataSource) convertToBatchPointGet(prop *property.PhysicalProperty, ca
 			batchPointGetPlan.Handles = append(batchPointGetPlan.Handles, kv.IntHandle(ran.LowVal[0].GetInt64()))
 		}
 		batchPointGetPlan.accessCols = ds.TblCols
-		batchPointGetPlan.HandleColOffset = ds.HandleCols.GetCol(0).Index
+		found := false
+		for i := range ds.Columns {
+			if ds.Columns[i].ID == ds.HandleCols.GetCol(0).ID {
+				batchPointGetPlan.HandleColOffset = ds.Columns[i].Offset
+				found = true
+				break
+			}
+		}
+		if !found {
+			return base.InvalidTask
+		}
+
 		// Add filter condition to table plan now.
 		if len(candidate.path.TableFilters) > 0 {
 			batchPointGetPlan.Init(ds.SCtx(), ds.TableStats.ScaleByExpectCnt(accessCnt), ds.Schema().Clone(), ds.OutputNames(), ds.QueryBlockOffset())

diff --git a/tests/integrationtest/r/executor/partition/issues.result b/tests/integrationtest/r/executor/partition/issues.result
@@ -433,3 +433,16 @@ select space(1), b from issue52198 where b in (1);
 space(1)	b
  	1
 drop table issue52198;
+drop table if exists t;
+set tidb_partition_prune_mode=static;
+CREATE TABLE t (
+a text COLLATE utf8mb4_unicode_ci DEFAULT NULL,
+b mediumint(8) unsigned NOT NULL DEFAULT '11075363',
+c tinyblob NOT NULL,
+PRIMARY KEY (b) /*T![clustered_index] CLUSTERED */
+) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
+PARTITION BY KEY (b) PARTITIONS 7;
+insert into t values ('a' ,6970066, 'a');
+update t set c = 'AH6' where b in ( 7691699 ,11807884 ,10523838 ,15662349 ,6970066 );
+drop table t;
+set tidb_partition_prune_mode=default;
diff --git a/tests/integrationtest/t/executor/partition/issues.test b/tests/integrationtest/t/executor/partition/issues.test
@@ -342,3 +342,18 @@ create table issue52198 (a int, b int, primary key (b)) partition by hash(b) par
 insert into issue52198 values (1,1);
 select space(1), b from issue52198 where b in (1);
 drop table issue52198;
+
+# TestIssues54667
+drop table if exists t;
+set tidb_partition_prune_mode=static;
+CREATE TABLE t (
+ a text COLLATE utf8mb4_unicode_ci DEFAULT NULL,
+ b mediumint(8) unsigned NOT NULL DEFAULT '11075363',
+ c tinyblob NOT NULL,
+ PRIMARY KEY (b) /*T![clustered_index] CLUSTERED */
+) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
+PARTITION BY KEY (b) PARTITIONS 7;
+insert into t values ('a' ,6970066, 'a');
+update t set c = 'AH6' where b in ( 7691699 ,11807884 ,10523838 ,15662349 ,6970066 );
+drop table t;
+set tidb_partition_prune_mode=default;