Fix column mismatch (#2785)

pingcap · Jul 25, 2024 · fb08132 · fb08132
1 parent 07632bc
commit fb08132
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/core/src/test/scala/org/apache/spark/sql/IssueTestSuite.scala b/core/src/test/scala/org/apache/spark/sql/IssueTestSuite.scala
@@ -22,6 +22,18 @@ import org.apache.spark.sql.functions.{col, sum}
 
 class IssueTestSuite extends BaseTiSparkTest {
 
+  test("test column mismatch, issue 2750") {
+    val dbTable = "tispark_test.column_mismatch"
+    tidbStmt.execute(s"drop table if exists $dbTable")
+    tidbStmt.execute(
+      s"CREATE TABLE $dbTable (`CI_NO` varchar(64) NOT NULL, `AC_DT` bigint(20) NOT NULL, `YM_DT` bigint(20) NOT NULL,`SYS_TYPE` varchar(20) NOT NULL, PRIMARY KEY (`CI_NO`,`AC_DT`,`YM_DT`,`SYS_TYPE`) /*T![clustered_index] CLUSTERED */, KEY `IDX_FLOW_01` (`CI_NO`,`AC_DT`))")
+    tidbStmt.execute(s"insert into $dbTable values('1',1,1,'1')")
+    spark.sql(s"select ci_no,ac_dt from $dbTable").explain()
+    spark.sqlContext.setConf(TiConfigConst.USE_INDEX_SCAN_FIRST, "true")
+    spark.sql(s"select * from $dbTable").show()
+    spark.sqlContext.setConf(TiConfigConst.USE_INDEX_SCAN_FIRST, "false")
+  }
+
   test("test like escape") {
     val dbTable = "tispark_test.like_escape"
     tidbStmt.execute(s"drop table if exists $dbTable")

diff --git a/tikv-client/src/main/java/com/pingcap/tikv/meta/TiDAGRequest.java b/tikv-client/src/main/java/com/pingcap/tikv/meta/TiDAGRequest.java
@@ -300,6 +300,14 @@ private void addIndexReaderIndexRangeScanExecutorCols(
                 + columnInfo.getName());
       }
     }
+    if (tableInfo.isCommonHandle()) {
+      for (TiIndexColumn col : tableInfo.getPrimaryKey().getIndexColumns()) {
+        TiColumnInfo columnInfo = tableInfo.getColumn(col.getName());
+        // add primary columns to columns:
+        // https://github.com/pingcap/tidb/blob/ddcaadbb856f0890e91e4c77991f0d2aa5aa93d0/pkg/planner/core/planbuilder.go#L1515
+        indexScanBuilder.addColumns(ColumnInfo.newBuilder(columnInfo.toProto(tableInfo)));
+      }
+    }
   }
 
   private void addIndexColsToScanBuilder(