Repository: pingcap/tispark (language: python)

Issue:
Fix column mismatch
### What problem does this PR solve? <!--add issue link with summary if exists-->

close #2750 

fix column mismatch reported by tikv. https://github.com/tikv/tikv/blob/254ceaaa654a70d78919ee5d85f5b4cbe3c444a3/components/tidb_query_executors/src/index_scan_executor.rs#L113

This occurs when the `primary_column_ids` number is more than  columns number in the index scan of DAGRequest. Example:
```
start_ts_fallback: 451202703803809793
executors {
  tp: TypeIndexScan
  idx_scan {
    table_id: 93
    index_id: 2
    columns {
      column_id: 1
      tp: 15
      collation: -46
      columnLen: 64
      decimal: 0
      flag: 4107
      default_val: "\000"
      pk_handle: false
    }
    columns {
      column_id: 2
      tp: 8
      collation: -63
      columnLen: 20
      decimal: 0
      flag: 4099
      default_val: "\000"
      pk_handle: false
    }
    primary_column_ids: 3
    primary_column_ids: 1
    primary_column_ids: 2
  }
}
executors {
  tp: TypeLimit
  limit {
    limit: 21
  }
}
time_zone_offset: -25200
flags: 0
output_offsets: 1
output_offsets: 0
encode_type: TypeChunk
```

### What is changed and how it works?

Add primary columns to columns in index scan of DAGRequest when pk is command handle.

Add this because TiDB also do like this. Ref: https://github.com/pingcap/tidb/blob/ddcaadbb856f0890e91e4c77991f0d2aa5aa93d0/pkg/planner/core/planbuilder.go#L1515


### Check List <!--REMOVE the items that are not applicable-->

Tests <!-- At least one of them must be included. -->

 - Unit test
 - Integration test
 - Manual test (add detailed scripts or steps below)
 - No code

Code changes

 - Has exported function/method change
 - Has exported variable/fields change
 - Has interface methods change
 - Has persistent data change

Side effects

 - Possible performance regression
 - Increased code complexity
 - Breaking backward compatibility

Related changes

 - Need to cherry-pick to the release branch
 - Need to update the documentation
 - Need to update the `tidb-ansible` repository
 - Need to be included in the release note

The number of handle columns exceeds the length of `columns_info`
## General Question
org.apache.spark.SparkException: Job aborted due to stage failure: Task 0 in stage 51.0 failed 4 times, most recent failure: Lost task 0.3 in stage 51.0 (TID 2116) : com.pingcap.tikv.exception.TiClientInternalException: Error reading region:
	at com.pingcap.tikv.operation.iterator.DAGIterator.doReadNextRegionChunks(DAGIterator.java:190)
	at com.pingcap.tikv.operation.iterator.DAGIterator.readNextRegionChunks(DAGIterator.java:167)
	at com.pingcap.tikv.operation.iterator.DAGIterator.hasNext(DAGIterator.java:113)
	at org.apache.spark.sql.tispark.TiRowRDD$$anon$1.hasNext(TiRowRDD.scala:70)
	at org.apache.spark.sql.catalyst.expressions.GeneratedClass$GeneratedIteratorForCodegenStage2.columnartorow_nextBatch_0$(Unknown Source)
	at org.apache.spark.sql.catalyst.expressions.GeneratedClass$GeneratedIteratorForCodegenStage2.processNext(Unknown Source)
	at org.apache.spark.sql.execution.BufferedRowIterator.hasNext(BufferedRowIterator.java:43)
	at org.apache.spark.sql.execution.WholeStageCodegenExec$$anon$1.hasNext(WholeStageCodegenExec.scala:755)
	at scala.collection.Iterator$$anon$10.hasNext(Iterator.scala:458)
	at org.apache.spark.shuffle.sort.UnsafeShuffleWriter.write(UnsafeShuffleWriter.java:177)
	at org.apache.spark.shuffle.ShuffleWriteProcessor.write(ShuffleWriteProcessor.scala:59)
	at org.apache.spark.scheduler.ShuffleMapTask.runTask(ShuffleMapTask.scala:99)
	at org.apache.spark.scheduler.ShuffleMapTask.runTask(ShuffleMapTask.scala:52)
	at org.apache.spark.scheduler.Task.run(Task.scala:131)
	at org.apache.spark.executor.Executor$TaskRunner.$anonfun$run$3(Executor.scala:497)
	at org.apache.spark.util.Utils$.tryWithSafeFinally(Utils.scala:1439)
	at org.apache.spark.executor.Executor$TaskRunner.run(Executor.scala:500)
	at java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1149)
	at java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:624)
	at java.lang.Thread.run(Thread.java:748)
Caused by: java.util.concurrent.ExecutionException: com.pingcap.tikv.exception.RegionTaskException: Handle region task failed:
	at java.util.concurrent.FutureTask.report(FutureTask.java:122)
	at java.util.concurrent.FutureTask.get(FutureTask.java:192)
	at com.pingcap.tikv.operation.iterator.DAGIterator.doReadNextRegionChunks(DAGIterator.java:185)
	... 19 more
Caused by: com.pingcap.tikv.exception.RegionTaskException: Handle region task failed:
	at com.pingcap.tikv.operation.iterator.DAGIterator.process(DAGIterator.java:233)
	at com.pingcap.tikv.operation.iterator.DAGIterator.lambda$submitTasks$1(DAGIterator.java:91)
	at java.util.concurrent.FutureTask.run(FutureTask.java:266)
	at java.util.concurrent.Executors$RunnableAdapter.call(Executors.java:511)
	at java.util.concurrent.FutureTask.run(FutureTask.java:266)
	... 3 more
Caused by: com.pingcap.tikv.exception.GrpcException: [components/tidb_query_executors/src/index_scan_executor.rs:109]: The number of handle columns exceeds the length of `columns_info`
	at com.pingcap.tikv.region.RegionStoreClient.handleCopResponse(RegionStoreClient.java:734)
	at com.pingcap.tikv.region.RegionStoreClient.coprocess(RegionStoreClient.java:681)
	at com.pingcap.tikv.operation.iterator.DAGIterator.process(DAGIterator.java:220)
	... 7 more


Discussion / hints:

/run-all-tests tidb=release-6.1 tikv=release-6.1 pd=release-6.1



/run-all-tests tidb=release-6.1 tikv=release-6.1 pd=release-6.1



/run-all-tests tidb=release-6.1 tikv=release-6.1 pd=release-6.1



/run-all-tests tidb=release-6.1 tikv=release-6.1 pd=release-6.1



/run-all-tests tidb=release-6.1 tikv=release-6.1 pd=release-6.1



/run-all-tests tidb=release-6.1 tikv=release-6.1 pd=release-6.1



/run-all-tests tidb=release-6.1 tikv=release-6.1 pd=release-6.1



/run-all-tests tidb=release-6.1 tikv=release-6.1 pd=release-6.1



/run-all-tests tidb=release-6.1 tikv=release-6.1 pd=release-6.1



/run-all-tests tidb=release-6.1 tikv=release-6.1 pd=release-6.1


[APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/pingcap/tispark/pull/2785#pullrequestreview-2193769579" title="Approved">xuanyu66</a>*, *<a href="https://github.com/pingcap/tispark/pull/2785#pullrequestreview-2198210997" title="Approved">zhangyangyu</a>*

The full list of commands accepted by this bot can be found [here](https://prow.tidb.net/command-help?repo=pingcap%2Ftispark).

The pull request process is described [here](https://book.prow.tidb.net/#/workflows/pr)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/pingcap/tispark/blob/master/OWNERS)~~ [xuanyu66]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->
[LGTM Timeline notifier]
---
Timeline:

- `2024-07-23 11:56:33.072751912 +0000 UTC m=+959815.063693367`: :ballot_box_with_check: agreed by [xuanyu66](https://github.com/xuanyu66).
- `2024-07-25 03:34:16.270161046 +0000 UTC m=+1102478.261102516`: :ballot_box_with_check: agreed by [zhangyangyu](https://github.com/zhangyangyu).
This issue is stale because it has been open for 30 days with no activity.
This issue was closed because it has been inactive for 14 days since being marked as stale.
Also meet the issue. reproduce steps:
1. DDL
```
CREATE TABLE `test`.`t` (
  `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`)
)
```
2. spark.sql("select ci_no,ac_dt from tidb_catalog.test.t1").show()

Fix patch:
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
index a62d2c5be..70fe906dd 100644
--- a/tikv-client/src/main/java/com/pingcap/tikv/meta/TiDAGRequest.java
+++ b/tikv-client/src/main/java/com/pingcap/tikv/meta/TiDAGRequest.java
@@ -298,6 +298,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(


Test patch:
diff --git a/core/src/test/scala/org/apache/spark/sql/IssueTestSuite.scala b/core/src/test/scala/org/apache/spark/sql/IssueTestSuite.scala
index e510bb17a..270a2dab0 100644
--- 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 tiflash overflow in unsigned bigint") {
     if (!enableTiFlashTest) {
       cancel("tiflash test not enabled")