fix column mismatch

pingcap · Jul 17, 2024 · 63c5c18 · 63c5c18
1 parent 07632bc
commit 63c5c18
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 1 deletion.
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,14 @@ 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, `SRC_KEY` varchar(100) NOT NULL, PRIMARY KEY (`SRC_KEY`,`CI_NO`,`AC_DT`) /*T![clustered_index] CLUSTERED */, KEY `IDX_FLOW_01` (`CI_NO`,`AC_DT`))")
+    spark.sql(s"select ci_no,ac_dt from $dbTable ").show()
+  }
+
   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
@@ -227,7 +227,14 @@ public DAGRequest buildDAGGetIndexData() {
     indexScanBuilder.setTableId(getPhysicalId()).setIndexId(indexInfo.getId());
     if (tableInfo.isCommonHandle()) {
       for (TiIndexColumn col : tableInfo.getPrimaryKey().getIndexColumns()) {
-        indexScanBuilder.addPrimaryColumnIds(tableInfo.getColumn(col.getName()).getId());
+        long primaryColumnId = tableInfo.getColumn(col.getName()).getId();
+        indexScanBuilder.addPrimaryColumnIds(primaryColumnId);
+        // add primary columns to columns:
+        // https://github.com/pingcap/tidb/blob/ca7ba14da87abd060922b32c893b121c4c3d6373/pkg/executor/builder.go#L569
+        if (!isPrimaryInColumns(indexScanBuilder.getColumnsList(), primaryColumnId)) {
+          TiColumnInfo columnInfo = tableInfo.getColumn(col.getName());
+          indexScanBuilder.addColumns(ColumnInfo.newBuilder(columnInfo.toProto(tableInfo)));
+        }
       }
     }
 
@@ -238,6 +245,16 @@ public DAGRequest buildDAGGetIndexData() {
     return buildRequest(dagRequestBuilder, outputOffsets);
   }
 
+  private boolean isPrimaryInColumns(
+      java.util.List<com.pingcap.tidb.tipb.ColumnInfo> columnsList, long primaryColumnId) {
+    for (com.pingcap.tidb.tipb.ColumnInfo column : columnsList) {
+      if (column.getColumnId() == primaryColumnId) {
+        return true;
+      }
+    }
+    return false;
+  }
+
   private void addIndexLookUpIndexRangeScanExecutorCols(
       IndexScan.Builder indexScanBuilder,
       List<Integer> outputOffsets,