Fix TiFlash overflow (#2740) (#2743)

* This is an automated cherry-pick of #2740 Signed-off-by: ti-chi-bot <[email protected]> * Update TiBlockColumnVector.java * fmt --------- Signed-off-by: ti-chi-bot <[email protected]> Co-authored-by: shi yuhang <[email protected]> Co-authored-by: shiyuhang <[email protected]>
pingcap · Aug 2, 2023 · fe87946 · fe87946
1 parent 466db63
commit fe87946
Show file tree

Hide file tree

Showing 2 changed files with 33 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,33 @@ import org.apache.spark.sql.functions.{col, sum}
 
 class IssueTestSuite extends BaseTiSparkTest {
 
+  test("test tiflash overflow in unsigned bigint") {
+    if (!enableTiFlashTest) {
+      cancel("tiflash test not enabled")
+    }
+    val dbTable = "tispark_test.tiflash_overflow"
+    tidbStmt.execute(s"drop table if exists $dbTable")
+    tidbStmt.execute(
+      s"CREATE TABLE $dbTable (`a` bigint(20) UNSIGNED NOT NULL,`b` bigint(20) UNSIGNED NOT NULL)")
+    tidbStmt.execute(s"insert into $dbTable values(16717361816800086255, 16717361816800086255)")
+    tidbStmt.execute(s"ALTER TABLE $dbTable SET TIFLASH REPLICA 1")
+
+    Thread.sleep(5 * 1000)
+
+    val prev = spark.conf.getOption(TiConfigConst.ISOLATION_READ_ENGINES)
+    try {
+      spark.conf
+        .set(TiConfigConst.ISOLATION_READ_ENGINES, TiConfigConst.TIFLASH_STORAGE_ENGINE)
+      val df = spark.sql(s"select * from $dbTable")
+      val row = Row(BigDecimal("16717361816800086255"), BigDecimal("16717361816800086255"))
+      checkAnswer(df, Seq(row))
+    } finally {
+      spark.conf.set(
+        TiConfigConst.ISOLATION_READ_ENGINES,
+        prev.getOrElse(TiConfigConst.DEFAULT_STORAGE_ENGINES))
+    }
+  }
+
   test("test issue 2649") {
     val dbTable = "tispark_test.mutil_uniq"
     tidbStmt.execute(s"drop table if exists $dbTable")

diff --git a/tikv-client/src/main/java/com/pingcap/tikv/columnar/TiBlockColumnVector.java b/tikv-client/src/main/java/com/pingcap/tikv/columnar/TiBlockColumnVector.java
@@ -18,13 +18,15 @@
 
 import static com.pingcap.tikv.util.MemoryUtil.EMPTY_BYTE_BUFFER_DIRECT;
 
+import com.google.common.primitives.UnsignedLong;
 import com.pingcap.tikv.ExtendedDateTime;
 import com.pingcap.tikv.codec.Codec.DateCodec;
 import com.pingcap.tikv.codec.Codec.DateTimeCodec;
 import com.pingcap.tikv.columnar.datatypes.CHType;
 import com.pingcap.tikv.types.AbstractDateTimeType;
 import com.pingcap.tikv.types.BytesType;
 import com.pingcap.tikv.types.DateType;
+import com.pingcap.tikv.types.DecimalType;
 import com.pingcap.tikv.util.MemoryUtil;
 import java.math.BigDecimal;
 import java.nio.ByteBuffer;
@@ -240,6 +242,10 @@ public double getDouble(int rowId) {
   @Override
   public BigDecimal getDecimal(int rowId, int precision, int scale) {
     long rowIdAddr = (long) rowId * fixedLength + dataAddr;
+    // avoid unsigned long overflow here
+    if (type == DecimalType.BIG_INT_DECIMAL) {
+      return new BigDecimal(UnsignedLong.fromLongBits(this.getLong(rowId)).bigIntegerValue());
+    }
     if (fixedLength == 4) {
       return MemoryUtil.getDecimal32(rowIdAddr, scale);
     } else if (fixedLength == 8) {