support column query and address comments

Signed-off-by: Haoyang Li <[email protected]>

docs/compatibility.md

@@ -537,6 +537,10 @@ The following regular expression patterns are not yet supported on the GPU and w
 
 Work is ongoing to increase the range of regular expressions that can run on the GPU.
 
+## URL Parsing
+
+`parse_url` QUERY with a column key could produce different results on CPU and GPU. In Spark, the `key` in `parse_url` could act like a regex, but GPU will match the key exactly. If key is literal, GPU will check if key contains regex special characters and fallback to CPU if it does, but if key is column, it will not be able to fallback. For example, `parse_url("http://foo/bar?abc=BAD&a.c=GOOD", QUERY, "a.c")` will return "BAD" on CPU, but "GOOD" on GPU. See the Spark issue: https://issues.apache.org/jira/browse/SPARK-44500
+
 ## Timestamps
 
 Spark stores timestamps internally relative to the JVM time zone.  Converting an arbitrary timestamp

docs/supported_ops.md

@@ -10817,7 +10817,7 @@ are limited.
 <td> </td>
 <td> </td>
 <td> </td>
-<td><em>PS<br/>only support partToExtract = PROTOCOL | HOST;<br/>Literal value only</em></td>
+<td><em>PS<br/>only support partToExtract = PROTOCOL | HOST | QUERY;<br/>Literal value only</em></td>
 <td> </td>
 <td> </td>
 <td> </td>
@@ -10838,7 +10838,7 @@ are limited.
 <td> </td>
 <td> </td>
 <td> </td>
-<td><em>PS<br/>Literal value only</em></td>
+<td>S</td>
 <td> </td>
 <td> </td>
 <td> </td>

integration_tests/src/main/python/url_test.py

@@ -172,13 +172,22 @@ def test_parse_url_query_with_key():
         lambda spark: unary_op_df(spark, url_gen)
             .selectExpr("a", "parse_url(a, 'QUERY', 'abc')", "parse_url(a, 'QUERY', 'a')")
         )
-
+
+def test_parse_url_query_with_key_column():
+    url_gen = StringGen(url_pattern_with_key)
+    key_gen = StringGen('[a-d]{1,3}')
+    assert_gpu_and_cpu_are_equal_collect(
+        lambda spark: two_col_df(spark, url_gen, key_gen)
+            .selectExpr("a", "parse_url(a, 'QUERY', b)")
+        )    
+
+@pytest.mark.parametrize('key', ['a?c', '*'], ids=idfn)
 @allow_non_gpu('ProjectExec', 'ParseUrl')
-def test_parse_url_query_with_key_regex_fallback():
+def test_parse_url_query_with_key_regex_fallback(key):
     url_gen = StringGen(url_pattern_with_key)
     assert_gpu_fallback_collect(
         lambda spark: unary_op_df(spark, url_gen)
-            .selectExpr("a", "parse_url(a, 'QUERY', 'a?c')", "parse_url(a, 'QUERY', '*')"),
+            .selectExpr("a", "parse_url(a, 'QUERY', '" + key + "')"),
             'ParseUrl')
 
 @pytest.mark.parametrize('part', supported_with_key_parts, ids=idfn)

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

@@ -445,7 +445,7 @@ object GpuOverrides extends Logging {
     "\f", "\\a", "\\e", "\\cx", "[", "]", "^", "&", ".", "*", "\\d", "\\D", "\\h", "\\H", "\\s",
     "\\S", "\\v", "\\V", "\\w", "\\w", "\\p", "$", "\\b", "\\B", "\\A", "\\G", "\\Z", "\\z", "\\R",
     "?", "|", "(", ")", "{", "}", "\\k", "\\Q", "\\E", ":", "!", "<=", ">")
-
+  val regexMetaChars = ".$^[]\\|?*+(){}"
   /**
    * Provides a way to log an info message about how long an operation took in milliseconds.
    */
@@ -3252,26 +3252,26 @@ object GpuOverrides extends Logging {
       ExprChecks.projectOnly(TypeSig.STRING, TypeSig.STRING,
         Seq(ParamCheck("url", TypeSig.STRING, TypeSig.STRING),
           ParamCheck("partToExtract", TypeSig.lit(TypeEnum.STRING).withPsNote(
-            TypeEnum.STRING, "only support partToExtract = PROTOCOL | HOST"), TypeSig.STRING)),
+            TypeEnum.STRING, "only support partToExtract = PROTOCOL | HOST | QUERY"), 
+            TypeSig.STRING)),
           // Should really be an OptionalParam
-          Some(RepeatingParamCheck("key", TypeSig.lit(TypeEnum.STRING), TypeSig.STRING))),
+          Some(RepeatingParamCheck("key", TypeSig.STRING, TypeSig.STRING))),
       (a, conf, p, r) => new ExprMeta[ParseUrl](a, conf, p, r) {
         override def tagExprForGpu(): Unit = {
           if (a.failOnError) {
             willNotWorkOnGpu("Fail on error is not supported on GPU when parsing urls.")
           }
 
-          // In Spark, the key in parse_url could act like a regex, but on GPU, we only support
-          // literal keys. So we need to fallback if the key contains regex special characters.
-          // See Spark issue: https://issues.apache.org/jira/browse/SPARK-44500
           extractStringLit(a.children(1)).map(_.toUpperCase) match {
+            // In Spark, the key in parse_url could act like a regex, but GPU will match the key 
+            // exactly. When key is literal, GPU will check if the key contains regex special and
+            // fallbcak to CPU if it does, but we are not able to fallback when key is column.
+            // see Spark issue: https://issues.apache.org/jira/browse/SPARK-44500
             case Some("QUERY") if (a.children.size == 3) => {
               extractLit(a.children(2)).foreach { key =>
                 if (key.value != null) {
                   val keyStr = key.value.asInstanceOf[UTF8String].toString
-                  val specialCharacters = List("\\", "[", "]", "{", "}", "^", "-", "$", 
-                      ".", "+", "*", "?", "|")
-                  if (specialCharacters.exists(keyStr.contains(_))) {
+                  if (regexMetaChars.exists(keyStr.contains(_))) {
                     willNotWorkOnGpu(s"Key $keyStr could act like a regex which is not " + 
                         "supported on GPU")
                   }

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2021-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -19,6 +19,7 @@ import java.sql.SQLException
 
 import scala.collection.mutable.ListBuffer
 
+import com.nvidia.spark.rapids.GpuOverrides.regexMetaChars
 import com.nvidia.spark.rapids.RegexParser.toReadableString
 
 /**
@@ -684,7 +685,6 @@ sealed class RegexRewriteFlags(val emptyRepetition: Boolean)
                 RegexSplitMode   if performing a split (string_split)
  */
 class CudfRegexTranspiler(mode: RegexMode) {
-  private val regexMetaChars = ".$^[]\\|?*+(){}"
   private val regexPunct = "!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~"
   private val escapeChars = Map('n' -> '\n', 'r' -> '\r', 't' -> '\t', 'f' -> '\f', 'a' -> '\u0007',
       'b' -> '\b', 'e' -> '\u001b')
@@ -747,7 +747,7 @@ class CudfRegexTranspiler(mode: RegexMode) {
     e match {
       case RegexEscaped(ch) if escapeChars.contains(ch) => Some(escapeChars(ch).toString)
       case RegexEscaped(ch) if regexPunct.contains(ch) => Some(ch.toString)
-      case RegexChar(ch) if !regexMetaChars.contains(ch) => Some(ch.toString)
+      case RegexChar(ch) if regexMetaChars.contains(ch) => Some(ch.toString)
       case RegexSequence(parts) =>
         parts.foldLeft[Option[String]](Some("")) { (all, x) =>
           all match {

sql-plugin/src/main/scala/org/apache/spark/sql/rapids/GpuParseUrl.scala

@@ -85,6 +85,16 @@ case class GpuParseUrl(children: Seq[Expression])
     ParseURI.parseURIQueryWithLiteral(col.getBase, keyStr)
   }
 
+  def doColumnar(col: GpuColumnVector, partToExtract: GpuScalar, 
+      key: GpuColumnVector): ColumnVector = {
+    val part = partToExtract.getValue.asInstanceOf[UTF8String].toString
+    if (part != QUERY) {
+      // return a null columnvector
+      return GpuColumnVector.columnVectorFromNull(col.getRowCount.toInt, StringType)
+    }
+    ParseURI.parseURIQueryWithColumn(col.getBase, key.getBase)
+  }
+
   override def columnarEval(batch: ColumnarBatch): GpuColumnVector = {
     if (children.size == 2) {
       val Seq(url, partToExtract) = children
@@ -109,6 +119,8 @@ case class GpuParseUrl(children: Seq[Expression])
             (urls, parts, keys) match {
               case (urlCv: GpuColumnVector, partScalar: GpuScalar, keyScalar: GpuScalar) =>
                 GpuColumnVector.from(doColumnar(urlCv, partScalar, keyScalar), dataType)
+              case (urlCv: GpuColumnVector, partScalar: GpuScalar, keyCv: GpuColumnVector) =>
+                GpuColumnVector.from(doColumnar(urlCv, partScalar, keyCv), dataType)
               case _ =>
                 throw new 
                     UnsupportedOperationException(s"Cannot columnar evaluate expression: $this")

tools/generated_files/supportedExprs.csv

@@ -386,7 +386,7 @@ Or,S,`or`,None,AST,rhs,S,NA,NA,NA,NA,NA,NA,NA,NA,NA,NA,NA,NA,NA,NA,NA,NA,NA
 Or,S,`or`,None,AST,result,S,NA,NA,NA,NA,NA,NA,NA,NA,NA,NA,NA,NA,NA,NA,NA,NA,NA
 ParseUrl,S,`parse_url`,None,project,url,NA,NA,NA,NA,NA,NA,NA,NA,NA,S,NA,NA,NA,NA,NA,NA,NA,NA
 ParseUrl,S,`parse_url`,None,project,partToExtract,NA,NA,NA,NA,NA,NA,NA,NA,NA,PS,NA,NA,NA,NA,NA,NA,NA,NA
-ParseUrl,S,`parse_url`,None,project,key,NA,NA,NA,NA,NA,NA,NA,NA,NA,PS,NA,NA,NA,NA,NA,NA,NA,NA
+ParseUrl,S,`parse_url`,None,project,key,NA,NA,NA,NA,NA,NA,NA,NA,NA,S,NA,NA,NA,NA,NA,NA,NA,NA
 ParseUrl,S,`parse_url`,None,project,result,NA,NA,NA,NA,NA,NA,NA,NA,NA,S,NA,NA,NA,NA,NA,NA,NA,NA
 PercentRank,S,`percent_rank`,None,window,ordering,S,S,S,S,S,S,S,S,PS,S,S,S,NS,NS,NS,NS,NS,NS
 PercentRank,S,`percent_rank`,None,window,result,NA,NA,NA,NA,NA,NA,S,NA,NA,NA,NA,NA,NA,NA,NA,NA,NA,NA