Re-enable dollar ($) line anchor in regular expressions in find mode

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

@@ -467,7 +467,7 @@ class CudfRegexTranspiler(mode: RegexMode) {
     // parse the source regular expression
     val regex = new RegexParser(pattern).parse()
     // validate that the regex is supported by cuDF
-    val cudfRegex = rewrite(regex)
+    val cudfRegex = rewrite(regex, None)
     // write out to regex string, performing minor transformations
     // such as adding additional escaping
     cudfRegex.toRegexString
@@ -535,19 +535,79 @@ class CudfRegexTranspiler(mode: RegexMode) {
     }
   }
 
+  private val lineTerminatorChars = Seq('\n', '\r', '\u0085', '\u2028', '\u2029')
+
+  // from Java 8 documention: a line terminator is a 1 to 2 character sequence that marks
+  // the end of a line of an input character sequence. 
+  // this method produces a RegexAST which outputs a regular expression to match any possible
+  // combination of line terminators
+  private def lineTerminatorMatcher(exclude: Set[Char], excludeCRLF: Boolean):RegexAST = {
+    val terminatorChars = new ListBuffer[RegexCharacterClassComponent]()
+    terminatorChars ++= lineTerminatorChars.filter(!exclude.contains(_)).map(RegexChar)
+
+    if (terminatorChars.size == 0 && excludeCRLF) {
+      RegexEmpty()
+    } else if (terminatorChars.size == 0) {
+      RegexGroup(capture = false, RegexSequence(ListBuffer(RegexChar('\r'), RegexChar('\n'))))
+    } else if (excludeCRLF) {
+      RegexCharacterClass(negated = false, characters = terminatorChars)
+    } else {
+      RegexGroup(capture = false,
+        RegexChoice(
+          RegexCharacterClass(negated = false, characters = terminatorChars),
+          RegexSequence(ListBuffer(RegexChar('\r'), RegexChar('\n')))))
+    }
+  }
 
-  private def rewrite(regex: RegexAST): RegexAST = {
+  private def rewrite(regex: RegexAST, previous: Option[RegexAST]): RegexAST = {
     regex match {
 
       case RegexChar(ch) => ch match {
         case '.' =>
           // workaround for https://github.com/rapidsai/cudf/issues/9619
           RegexCharacterClass(negated = true, ListBuffer(RegexChar('\r'), RegexChar('\n')))
-        case '$' =>
+        case '$' if mode == RegexSplitMode || mode == RegexReplaceMode =>
           // see https://github.com/NVIDIA/spark-rapids/issues/4533
-          throw new RegexUnsupportedException("line anchor $ is not supported")
+          throw new RegexUnsupportedException("line anchor $ is not supported in split or replace")
+        case '$' =>
+          // in the case of the line anchor $, the JVM has special conditions when handling line 
+          // terminators in and around the anchor
+          // this handles cases where the line terminator characters are *before* the anchor ($)
+          // NOTE: this applies to when using *standard* mode. In multiline mode, all these 
+          // conditions will change. Currently Spark does not use multiline mode.
+          previous match {
+            case Some(RegexChar('$')) =>
+              // repeating the line anchor in cuDF (for example b$$) causes matches to fail, but in 
+              // Java, it's treated as a single (b$ and b$$ are synonymous), so we create 
+              // an empty RegexAST that outputs to empty string
+              RegexEmpty()
+            case Some(RegexChar(ch)) if ch == '\r' =>
+              // when using the the CR (\r), it prevents the line anchor from handling any other 
+              // line terminator sequences, so we just output the anchor and we are finished
+              // for example: \r$ -> \r$ (no transpilation)
+              RegexChar('$')
+            case Some(RegexChar(ch)) if lineTerminatorChars.contains(ch) =>
+              // when using any other line terminator character, you can match any of the other
+              // line terminator characters individually as part of the line anchor match.
+              // for example: \n$ -> \n[\r\u0085\u2028\u2029]?$
+              RegexSequence(ListBuffer(
+                RegexRepetition(lineTerminatorMatcher(Set(ch), true), SimpleQuantifier('?')),
+                RegexChar('$')))
+            case _ =>
+              // otherwise by default we can match any or none the full set of line terminators
+              RegexSequence(ListBuffer(
+                RegexRepetition(lineTerminatorMatcher(Set.empty, false), SimpleQuantifier('?')),
+                RegexChar('$')))
+          }
         case '^' if mode == RegexSplitMode =>
           throw new RegexUnsupportedException("line anchor ^ is not supported in split mode")
+        case '\r' | '\n' if mode == RegexFindMode =>
+          previous match {
+            case Some(RegexChar('$')) =>
+              RegexEmpty()
+            case _ =>
+              regex
+          }
         case _ =>
           regex
       }
@@ -634,7 +694,10 @@ class CudfRegexTranspiler(mode: RegexMode) {
           case _ =>
         }
         val components: Seq[RegexCharacterClassComponent] = characters
-          .map(x => rewrite(x).asInstanceOf[RegexCharacterClassComponent])
+          .map(x => x match {
+            case RegexChar(ch) if "^$".contains(ch) => x
+            case _ => rewrite(x, None).asInstanceOf[RegexCharacterClassComponent]
+          })
 
         if (negated) {
           // There are differences between cuDF and Java handling of newlines
@@ -699,7 +762,52 @@ class CudfRegexTranspiler(mode: RegexMode) {
           throw new RegexUnsupportedException(
             "sequences that only contain '^' or '$' are not supported")
         }
-        RegexSequence(parts.map(rewrite))
+
+        // Special handling for line anchor ($)
+        // This code is implemented here because to make it work in cuDF, we have to reorder 
+        // the items in the regex.
+        // In the JVM, regexes like "\n$" and "$\n" have similar treatment
+        RegexSequence(parts.foldLeft((new ListBuffer[RegexAST](),
+          Option.empty[RegexAST]))((m, part) => {
+            val (r, last) = m
+            last match {
+              // when the previous character is a line anchor ($), the JVM has special handling
+              // when matching against line terminator characters
+              case Some(RegexChar('$')) => 
+                val j = r.lastIndexWhere {
+                  case RegexEmpty() => false
+                  case _ => true
+                }
+                part match {
+                  case RegexCharacterClass(true, parts)
+                      if parts.forall(!isBeginOrEndLineAnchor(_)) =>
+                    r(j) = RegexSequence(
+                      ListBuffer(lineTerminatorMatcher(Set.empty, true), RegexChar('$')))
+                  case RegexChar(ch) if ch == '\n' =>
+                    // what's really needed here is negative lookahead, but that is not 
+                    // supported by cuDF
+                    // in this case: $\n would transpile to (?!\r)\n$
+                    throw new RegexUnsupportedException("regex sequence $\\n is not supported")
+                  case RegexChar(ch) if "\r\u0085\u2028\u2029".contains(ch) =>
+                    r(j) = RegexSequence(
+                      ListBuffer(
+                        rewrite(part, None),
+                        RegexSequence(ListBuffer(
+                          RegexRepetition(lineTerminatorMatcher(Set(ch), true),
+                            SimpleQuantifier('?')), RegexChar('$')))))
+                  case _ =>
+                    r.append(rewrite(part, last))
+                }
+              case _ =>
+                r.append(rewrite(part, last))
+            }
+            r.last match {
+              case RegexEmpty() =>
+                (r, last)
+              case _ =>
+                (r, Some(part))
+            }
+        })._1)
 
       case RegexRepetition(base, quantifier) => (base, quantifier) match {
         case (_, SimpleQuantifier(ch)) if mode == RegexReplaceMode && "?*".contains(ch) =>
@@ -742,17 +850,17 @@ class CudfRegexTranspiler(mode: RegexMode) {
           // specifically this variable length repetition: \A{2,}
           throw new RegexUnsupportedException(nothingToRepeat)
         case (RegexGroup(_, _), SimpleQuantifier(ch)) if ch == '?' =>
-          RegexRepetition(rewrite(base), quantifier)
+          RegexRepetition(rewrite(base, None), quantifier)
         case _ if isSupportedRepetitionBase(base) =>
-          RegexRepetition(rewrite(base), quantifier)
+          RegexRepetition(rewrite(base, None), quantifier)
         case _ =>
           throw new RegexUnsupportedException(nothingToRepeat)
 
       }
 
       case RegexChoice(l, r) =>
-        val ll = rewrite(l)
-        val rr = rewrite(r)
+        val ll = rewrite(l, None)
+        val rr = rewrite(r, None)
 
         // cuDF does not support repetition on one side of a choice, such as "a*|a"
         if (isRepetition(ll) || isRepetition(rr)) {
@@ -776,7 +884,7 @@ class CudfRegexTranspiler(mode: RegexMode) {
         RegexChoice(ll, rr)
 
       case RegexGroup(capture, term) =>
-        RegexGroup(capture, rewrite(term))
+        RegexGroup(capture, rewrite(term, None))
 
       case other =>
         throw new RegexUnsupportedException(s"Unhandled expression in transpiler: $other")
@@ -804,6 +912,11 @@ sealed trait RegexAST {
   def toRegexString: String
 }
 
+sealed case class RegexEmpty() extends RegexAST {
+  override def children(): Seq[RegexAST] = Seq.empty
+  override def toRegexString: String = ""
+}
+
 sealed case class RegexSequence(parts: ListBuffer[RegexAST]) extends RegexAST {
   override def children(): Seq[RegexAST] = parts
   override def toRegexString: String = parts.map(_.toRegexString).mkString

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

@@ -93,17 +93,13 @@ class RegularExpressionSuite extends SparkQueryCompareTestSuite {
     frame => frame.selectExpr("regexp_replace(strings,'\\(foo\\)','D')")
   }
 
-  testGpuFallback("String regexp_extract regex 1",
-    "RegExpExtract",
-    extractStrings, execsAllowedNonGpu = Seq("ProjectExec", "ShuffleExchangeExec", "Alias",
-      "RegExpExtract", "AttributeReference", "Literal"),conf = conf) {
+  testSparkResultsAreEqual("String regexp_extract regex 1",
+    extractStrings, conf = conf) {
     frame => frame.selectExpr("regexp_extract(strings, '^([a-z]*)([0-9]*)([a-z]*)$', 1)")
   }
 
-  testGpuFallback("String regexp_extract regex 2",
-    "RegExpExtract",
-    extractStrings, execsAllowedNonGpu = Seq("ProjectExec", "ShuffleExchangeExec", "Alias",
-      "RegExpExtract", "AttributeReference", "Literal"),conf = conf) {
+  testSparkResultsAreEqual("String regexp_extract regex 2",
+    extractStrings, conf = conf) {
     frame => frame.selectExpr("regexp_extract(strings, '^([a-z]*)([0-9]*)([a-z]*)$', 2)")
   }
 

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

@@ -215,12 +215,25 @@ class RegularExpressionTranspilerSuite extends FunSuite with Arm {
       "\ntest", "test\n", "\ntest\n", "\ntest\r\ntest\n"))
   }
 
-  test("line anchor $ fall back to CPU") {
-    for (mode <- Seq(RegexFindMode, RegexReplaceMode)) {
-      assertUnsupported("a$b", mode, "line anchor $ is not supported")
+  test("line anchor $ fall back to CPU - split and replace") {
+    for (mode <- Seq(RegexSplitMode, RegexReplaceMode)) {
+      assertUnsupported("a$b", mode, "line anchor $ is not supported in split or replace")
     }
   }
 
+  test("line anchor sequence $\\n fall back to CPU") {
+    assertUnsupported("a$\n", RegexFindMode, "regex sequence $\\n is not supported")
+  }
+
+  test("line anchor $ - find") {
+    val patterns = Seq("$\r", "a$", "\r$", "\f$", "$\f", "\u0085$", "\u2028$", "\u2029$", "\n$",
+      "\r\n$", "[\r\n]?$", "\\00*[D$3]$", "a$b")
+    val inputs = Seq("a", "a\n", "a\r", "a\r\n", "a\u0085\n", "a\f", "\f", "\r", "\u0085", "\u2028",
+      "\u2029", "\n", "\r\n", "\r\n\r", "\r\n\u0085", "\u0085\r", "\u2028\n", "\u2029\n", "\n\r",
+      "\n\u0085", "\n\u2028", "\n\u2029", "2+|+??wD\n", "a\r\nb")
+    assertCpuGpuMatchesRegexpFind(patterns, inputs)
+  }
+
   test("whitespace boundaries - replace") {
     assertCpuGpuMatchesRegexpReplace(
       Seq("\\s", "\\S"),
@@ -280,6 +293,13 @@ class RegularExpressionTranspilerSuite extends FunSuite with Arm {
     doTranspileTest("abc\\z", "abc$")
   }
 
+  test("transpile $") {
+    doTranspileTest("a$", "a(?:[\n\r\u0085\u2028\u2029]|\r\n)?$")
+    doTranspileTest("$$\r", "\r[\n\u0085\u2028\u2029]?$")
+    doTranspileTest("]$\r", "]\r[\n\u0085\u2028\u2029]?$")
+    doTranspileTest("^$[^*A-ZA-Z]", "^[\n\r\u0085\u2028\u2029]$")
+  }
+
   test("compare CPU and GPU: character range including unescaped + and -") {
     val patterns = Seq("a[-]+", "a[a-b-]+", "a[-a-b]", "a[-+]", "a[+-]")
     val inputs = Seq("a+", "a-", "a", "a-+", "a[a-b-]")