Fix bugs in CAST string to integer

sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuCast.scala

@@ -175,6 +175,59 @@ object GpuCast extends Arm {
       cv.stringReplaceWithBackrefs(rule.search, rule.replace)
     })
   }
+
+  def sanitizeStringToIntegralType(input: ColumnVector, ansiEnabled: Boolean): ColumnVector = {
+    // Convert any strings containing whitespace to null values. The input is assumed to already
+    // have been stripped of leading and trailing whitespace
+    val sanitized = withResource(input.containsRe("\\s")) { hasWhitespace =>
+      withResource(hasWhitespace.any()) { any =>
+        if (any.getBoolean) {
+          if (ansiEnabled) {
+            throw new NumberFormatException(GpuCast.INVALID_INPUT_MESSAGE)
+          } else {
+            withResource(GpuScalar.from(null, DataTypes.StringType)) { nullVal =>
+              hasWhitespace.ifElse(nullVal, input)
+            }
+          }
+        } else {
+          input.incRefCount()
+        }
+      }
+    }
+
+    // truncate strings that represent decimals, so that we just look at the string before the dot
+    if (ansiEnabled) {
+      // ansi mode does not support casting decimal strings to integers
+      sanitized
+    } else {
+      withResource(sanitized) { _ =>
+        withResource(Scalar.fromString(".")) { dot =>
+          withResource(sanitized.stringContains(dot)) { hasDot =>
+            // only do the decimal sanitization if any strings do contain dot
+            withResource(hasDot.any(DType.BOOL8)) { anyDot =>
+              if (anyDot.getBoolean) {
+                // Special handling for strings that have no numeric value before the dot, such
+                // as "." and ".1" because extractsRe returns null for the capture group
+                // for these values and it also returns null for invalid inputs so we need this
+                // explicit check
+                withResource(sanitized.matchesRe("^[+\\-]?\\.[0-9]*$")) { startsWithDot =>
+                  withResource(sanitized.extractRe("^([+\\-]?[0-9]*)\\.[0-9]*$")) { table =>
+                    withResource(Scalar.fromString("0")) { zero =>
+                      withResource(startsWithDot.ifElse(zero, table.getColumn(0))) {
+                        decimal => hasDot.ifElse(decimal, sanitized)
+                      }
+                    }
+                  }
+                }
+              } else {
+                sanitized.incRefCount()
+              }
+            }
+          }
+        }
+      }
+    }
+  }
 }
 
 /**
@@ -681,46 +734,19 @@ case class GpuCast(
       input: ColumnVector,
       ansiEnabled: Boolean,
       dType: DType): ColumnVector = {
-    val cleaned = if (!ansiEnabled) {
-      // TODO would be great to get rid of this regex, but the overflow checks don't work
-      //  on the more lenient pattern.
-      // To avoid doing the expensive regex all the time, we will first check to see if we need
-      // to do it. The only time we do need to do it is when we have a '.' in any of the strings.
-      val data = input.getData
-      val hasDot = if (data != null) {
-        withResource(
-          ColumnView.fromDeviceBuffer(data, 0, DType.INT8, data.getLength.toInt)) { childData =>
-          withResource(GpuScalar.from('.'.toByte, ByteType)) { dot =>
-            childData.contains(dot)
-          }
-        }
-      } else {
-        false
-      }
 
-      if (hasDot) {
-        withResource(input.extractRe("^([+\\-]?[0-9]+)(?:\\.[0-9]*)?$")) { table =>
-          table.getColumn(0).incRefCount()
-        }
-      } else {
-        input.incRefCount()
-      }
-    } else {
-      input.incRefCount()
-    }
-    withResource(cleaned) { cleaned =>
-      withResource(cleaned.isInteger(dType)) { isInt =>
-        if (ansiEnabled) {
-          withResource(isInt.all()) { allInts =>
-            if (!allInts.getBoolean) {
-              throw new NumberFormatException(GpuCast.INVALID_INPUT_MESSAGE)
-            }
-          }
-          cleaned.castTo(dType)
-        } else {
-          withResource(cleaned.castTo(dType)) { parsedInt =>
-            withResource(GpuScalar.from(null, dataType)) { nullVal =>
-              isInt.ifElse(parsedInt, nullVal)
+    withResource(GpuCast.sanitizeStringToIntegralType(input, ansiEnabled)) { sanitized =>
+      withResource(sanitized.isInteger(dType)) { isInt =>
+        withResource(isInt.all()) { allInts =>
+          if (allInts.getBoolean) {
+            sanitized.castTo(dType)
+          } else if (ansiEnabled) {
+            throw new NumberFormatException(GpuCast.INVALID_INPUT_MESSAGE)
+          } else {
+            withResource(sanitized.castTo(dType)) { parsedInt =>
+              withResource(GpuScalar.from(null, dataType)) { nullVal =>
+                isInt.ifElse(parsedInt, nullVal)
+              }
             }
           }
         }

tests/src/test/scala/com/nvidia/spark/rapids/CastOpSuite.scala

@@ -78,27 +78,24 @@ class CastOpSuite extends GpuExpressionTestSuite {
       "f", "F", "True", "TRUE", "true", "tRuE", "t", "T", "Y", "y", "10", "01", "0", "1"))
   }
 
-  ignore("Cast from string to byte using random inputs") {
-    // Test ignored due to known issues
-    // https://github.com/NVIDIA/spark-rapids/issues/2899
+  test("Cast from string to byte using random inputs") {
     testCastStringTo(DataTypes.ByteType, generateRandomStrings(Some(NUMERIC_CHARS)))
   }
 
-  ignore("Cast from string to short using random inputs") {
-    // Test ignored due to known issues
-    // https://github.com/NVIDIA/spark-rapids/issues/2899
+  test("Cast from string to short using random inputs") {
     testCastStringTo(DataTypes.ShortType, generateRandomStrings(Some(NUMERIC_CHARS)))
   }
 
-  ignore("Cast from string to int using random inputs") {
-    // Test ignored due to known issues
-    // https://github.com/NVIDIA/spark-rapids/issues/2899
+  test("Cast from string to int using random inputs") {
     testCastStringTo(DataTypes.IntegerType, generateRandomStrings(Some(NUMERIC_CHARS)))
   }
 
-  ignore("Cast from string to long using random inputs") {
-    // Test ignored due to known issues
-    // https://github.com/NVIDIA/spark-rapids/issues/2899
+  test("Cast from string to int using hand-picked values") {
+    testCastStringTo(DataTypes.IntegerType, Seq(".--e-37602.n", "\r\r\t\n11.12380", "-.2", ".3",
+      ".", "+1.2", "\n123\n456\n"))
+  }
+
+  test("Cast from string to long using random inputs") {
     testCastStringTo(DataTypes.LongType, generateRandomStrings(Some(NUMERIC_CHARS)))
   }
 
@@ -831,6 +828,28 @@ class CastOpSuite extends GpuExpressionTestSuite {
     }
   }
 
+  test("CAST string to integer - sanitize step") {
+    val testPairs: Seq[(String, String)] = Seq(
+      ("123", "123"),
+      (".", "0"),
+      (".2", "0"),
+      ("-.2", "0"),
+      ("0.123", "0"),
+      ("321.123", "321"),
+      ("0.123\r123", null),
+      (".\r123", null)
+    )
+    val inputs = testPairs.map(_._1)
+    val expected = testPairs.map(_._2)
+    withResource(ColumnVector.fromStrings(inputs: _*)) { v =>
+      withResource(ColumnVector.fromStrings(expected: _*)) { expected =>
+        withResource(GpuCast.sanitizeStringToIntegralType(v, ansiEnabled = false)) { actual =>
+          CudfTestHelper.assertColumnsAreEqual(expected, actual)
+        }
+      }
+    }
+  }
+
   test("CAST string to date - sanitize step") {
     val testPairs = Seq(
       ("2001-1", "2001-01"),