Fix overflow

rui-mo · Apr 11, 2024 · 069fb47 · 069fb47
1 parent ecf8118
commit 069fb47
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 10 deletions.
diff --git a/velox/functions/sparksql/tests/SparkCastExprTest.cpp b/velox/functions/sparksql/tests/SparkCastExprTest.cpp
@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+#include "velox/common/base/tests/GTestUtils.h"
 #include "velox/functions/prestosql/tests/CastBaseTest.h"
 #include "velox/functions/sparksql/Register.h"
 #include "velox/parse/TypeResolver.h"
@@ -453,6 +454,29 @@ TEST_F(SparkCastExprTest, overflow) {
   testCast(
       makeNullableFlatVector<int64_t>({214748364890}, DECIMAL(12, 2)),
       makeNullableFlatVector<int64_t>({2147483648}));
+
+  VELOX_ASSERT_THROW(
+      (testCast<std::string, int8_t>("tinyint", {"166"}, {std::nullopt})),
+      "Cannot cast VARCHAR '166' to TINYINT. Value is too large for type.");
+  testTryCast<std::string, int8_t>("tinyint", {"166"}, {std::nullopt});
+
+  VELOX_ASSERT_THROW(
+      (testCast<std::string, int16_t>("smallint", {"52769"}, {std::nullopt})),
+      "Cannot cast VARCHAR '52769' to SMALLINT. Value is too large for type.");
+  testTryCast<std::string, int16_t>("tinyint", {"52769"}, {std::nullopt});
+
+  VELOX_ASSERT_THROW(
+      (testCast<std::string, int32_t>(
+          "integer", {"17515055537"}, {std::nullopt})),
+      "Cannot cast VARCHAR '17515055537' to INTEGER. Value is too large for type.");
+  testTryCast<std::string, int32_t>("integer", {"17515055537"}, {std::nullopt});
+
+  VELOX_ASSERT_THROW(
+      (testCast<std::string, int64_t>(
+          "integer", {"9663372036854775809"}, {std::nullopt})),
+      "Cannot cast VARCHAR '9663372036854775809' to BIGINT. Value is too large for type.");
+  testTryCast<std::string, int64_t>(
+      "integer", {"9663372036854775809"}, {std::nullopt});
 }
 
 TEST_F(SparkCastExprTest, timestampToString) {

diff --git a/velox/type/Conversions.h b/velox/type/Conversions.h
@@ -195,11 +195,14 @@ struct Converter<
           VELOX_USER_FAIL("Encountered a non-digit character");
         }
         if (!decimalPoint) {
-          result = result * 10 - (v[index] - '0');
-        }
-        // Overflow check
-        if (result > 0) {
-          VELOX_USER_FAIL("Value is too large for type");
+          bool overflow = __builtin_mul_overflow(result, 10, &result);
+          if (overflow) {
+            VELOX_USER_FAIL("Value is too large for type.");
+          }
+          overflow = __builtin_sub_overflow(result, v[index] - '0', &result);
+          if (overflow) {
+            VELOX_USER_FAIL("Value is too large for type.");
+          }
         }
       }
     } else {
@@ -215,11 +218,14 @@ struct Converter<
           VELOX_USER_FAIL("Encountered a non-digit character");
         }
         if (!decimalPoint) {
-          result = result * 10 + (v[index] - '0');
-        }
-        // Overflow check
-        if (result < 0) {
-          VELOX_USER_FAIL("Value is too large for type");
+          bool overflow = __builtin_mul_overflow(result, 10, &result);
+          if (overflow) {
+            VELOX_USER_FAIL("Value is too large for type.");
+          }
+          overflow = __builtin_add_overflow(result, v[index] - '0', &result);
+          if (overflow) {
+            VELOX_USER_FAIL("Value is too large for type.");
+          }
         }
       }
     }