fix mutil_uniq (#2650) (#2666)

pingcap · Mar 28, 2023 · e6bdee3 · e6bdee3
1 parent b317437
commit e6bdee3
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 3 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,17 @@ import org.apache.spark.sql.functions.{col, sum}
 
 class IssueTestSuite extends BaseTiSparkTest {
 
+  test("test issue 2649") {
+    val dbTable = "tispark_test.mutil_uniq"
+    tidbStmt.execute(s"drop table if exists $dbTable")
+    tidbStmt.execute(
+      s"CREATE TABLE $dbTable (`a` int(5) NOT NULL,`b` int(5) NOT NULL,UNIQUE KEY `idx_ab` (`a`,`b`))")
+    tidbStmt.execute(s"insert into $dbTable values(0, 0),(1,1)")
+    val df = spark.sql(s"select * from $dbTable where a=1")
+    val row1 = Row(1, 1)
+    checkAnswer(df, Seq(row1))
+  }
+
   //https://github.com/pingcap/tispark/issues/2268
   test("show rowid in commonhandle") {
     spark.sqlContext.setConf(TiConfigConst.SHOW_ROWID, "true")

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
@@ -304,10 +304,10 @@ private void addIndexReaderIndexRangeScanExecutorCols(
 
   private void addIndexColsToScanBuilder(
       IndexScan.Builder indexScanBuilder, Map<String, Integer> colOffsetInFieldMap) {
-    if (indexInfo.isUnique()) {
-      indexScanBuilder.setUnique(true);
-    }
     for (TiIndexColumn indexColumn : indexInfo.getIndexColumns()) {
+      // We can optimize performance by pass unique into index scan.
+      // Refer to
+      // https://github.com/pingcap/tidb/blob/4ae5be190b0017675deedd9af8b80d8868e03f0e/planner/core/plan_to_pb.go#L271 to see if we can pass unique into index scan.
       TiColumnInfo tableInfoColumn = tableInfo.getColumn(indexColumn.getName());
       // already add this col before.
       ColumnInfo.Builder colBuild = ColumnInfo.newBuilder(tableInfoColumn.toProto(tableInfo));