[SPARK-16281][SQL] Implement parse_url SQL function

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala

@@ -284,6 +284,7 @@ object FunctionRegistry {
     expression[StringLPad]("lpad"),
     expression[StringTrimLeft]("ltrim"),
     expression[JsonTuple]("json_tuple"),
+    expression[ParseUrl]("parse_url"),
     expression[FormatString]("printf"),
     expression[RegExpExtract]("regexp_extract"),
     expression[RegExpReplace]("regexp_replace"),

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala

@@ -17,7 +17,9 @@
 
 package org.apache.spark.sql.catalyst.expressions
 
+import java.net.URL
 import java.text.{DecimalFormat, DecimalFormatSymbols}
+import java.util.regex.Pattern
 import java.util.{HashMap, Locale, Map => JMap}
 
 import org.apache.spark.sql.catalyst.InternalRow
@@ -27,6 +29,8 @@ import org.apache.spark.sql.catalyst.util.ArrayData
 import org.apache.spark.sql.types._
 import org.apache.spark.unsafe.types.{ByteArray, UTF8String}
 
+import scala.util.control.NonFatal
+
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 // This file defines expressions for string operations.
 ////////////////////////////////////////////////////////////////////////////////////////////////////
@@ -652,6 +656,125 @@ case class StringRPad(str: Expression, len: Expression, pad: Expression)
   override def prettyName: String = "rpad"
 }
 
+/**
+ * Extracts a part from a URL
+ */
+@ExpressionDescription(
+  usage = "_FUNC_(url, partToExtract[, key]) - extracts a part from a URL",
+  extended = """Parts: HOST, PATH, QUERY, REF, PROTOCOL, AUTHORITY, FILE, USERINFO.
+    Key specifies which query to extract.
+    Examples:
+      > SELECT _FUNC_('http://spark.apache.org/path?query=1', 'HOST')\n 'spark.apache.org'
+      > SELECT _FUNC_('http://spark.apache.org/path?query=1', 'QUERY')\n 'query=1'
+      > SELECT _FUNC_('http://spark.apache.org/path?query=1', 'QUERY', 'query')\n '1'""")
+case class ParseUrl(children: Expression*)
+  extends Expression with ImplicitCastInputTypes with CodegenFallback {
+
+  override def nullable = true
+  override def inputTypes: Seq[DataType] = Seq.fill(children.size)(StringType)
+  override def dataType: DataType = StringType
+
+  @transient private var lastUrlStr: UTF8String = _
+  @transient private var lastUrl: URL = _
+  // last key in string, we will update the pattern if key value changed.
+  @transient private var lastKey: UTF8String = _
+  // last regex pattern, we cache it for performance concern
+  @transient private var pattern: Pattern = _
+
+  private lazy val stringExprs = children.toArray
+  private val HOST = UTF8String.fromString("HOST")
+  private val PATH = UTF8String.fromString("PATH")
+  private val QUERY = UTF8String.fromString("QUERY")
+  private val REF = UTF8String.fromString("REF")
+  private val PROTOCOL = UTF8String.fromString("PROTOCOL")
+  private val FILE = UTF8String.fromString("FILE")
+  private val AUTHORITY = UTF8String.fromString("AUTHORITY")
+  private val USERINFO = UTF8String.fromString("USERINFO")
+  private val REGEXPREFIX = "(&|^)"
+  private val REGEXSUBFIX = "=([^&]*)"
+
+  override def checkInputDataTypes(): TypeCheckResult = {
+    if (children.size > 3 || children.size < 2) {
+      TypeCheckResult.TypeCheckFailure("parse_url function requires two or three arguments")
+    } else {
+      super[ImplicitCastInputTypes].checkInputDataTypes()
+    }
+  }
+
+  def parseUrlWithoutKey(url: Any, partToExtract: Any): Any = {
+    if (url == null || partToExtract == null) {
+      null
+    } else {
+      if (lastUrlStr == null || !url.equals(lastUrlStr)) {
+        try {
+          lastUrlStr = url.asInstanceOf[UTF8String].clone()
+          lastUrl = new URL(url.toString)
+        } catch {
+          case NonFatal(_) => return null
+        }
+      }
+
+      if (partToExtract.equals(HOST)) {
+        UTF8String.fromString(lastUrl.getHost)
+      } else if (partToExtract.equals(PATH)) {
+        UTF8String.fromString(lastUrl.getPath)
+      } else if (partToExtract.equals(QUERY)) {
+        UTF8String.fromString(lastUrl.getQuery)
+      } else if (partToExtract.equals(REF)) {
+        UTF8String.fromString(lastUrl.getRef)
+      } else if (partToExtract.equals(PROTOCOL)) {
+        UTF8String.fromString(lastUrl.getProtocol)
+      } else if (partToExtract.equals(FILE)) {
+        UTF8String.fromString(lastUrl.getFile)
+      } else if (partToExtract.equals(AUTHORITY)) {
+        UTF8String.fromString(lastUrl.getAuthority)
+      } else if (partToExtract.equals(USERINFO)) {
+        UTF8String.fromString(lastUrl.getUserInfo)
+      } else {
+        null
+      }
+    }
+  }
+
+  override def eval(input: InternalRow): Any = {
+    val url = stringExprs(0).eval(input)
+    val partToExtract = stringExprs(1).eval(input)
+    if (stringExprs.size == 2) {
+      parseUrlWithoutKey(url, partToExtract)
+    } else { // QUERY with key
+      if (partToExtract == null || !partToExtract.equals(QUERY)) {
+        null
+      } else {
+        val query = parseUrlWithoutKey(url, partToExtract)
+        if (query == null) {
+          null
+        } else {
+          val key = stringExprs(2).eval(input)
+          if (key == null) {
+            null
+          } else {
+            if (!key.equals(lastKey)) {
+              lastKey = key.asInstanceOf[UTF8String].clone()
+              val sb = new StringBuilder()
+              sb.append(REGEXPREFIX).append(key.toString).append(REGEXSUBFIX)
+              pattern = Pattern.compile(sb.toString())
+            }
+
+            val m = pattern.matcher(query.toString)
+            if (m.find()) {
+              UTF8String.fromString(m.group(2))
+            } else {
+              null
+            }
+          }
+        }
+      }
+    }
+  }
+
+  override def prettyName: String = "parse_url"
+}
+
 /**
  * Returns the input formatted according do printf-style format strings
  */

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/StringExpressionsSuite.scala

@@ -725,4 +725,45 @@ class StringExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
     checkEvaluation(FindInSet(Literal("abf"), Literal("abc,b,ab,c,def")), 0)
     checkEvaluation(FindInSet(Literal("ab,"), Literal("abc,b,ab,c,def")), 0)
   }
+
+  test("ParseUrl") {
+    def checkParseUrl(expected: String, urlStr: String, partToExtract: String): Unit = {
+      checkEvaluation(
+        ParseUrl(Literal.create(urlStr, StringType), Literal.create(partToExtract, StringType)),
+        expected)
+    }
+    def checkParseUrlWithKey(
+        expected: String, urlStr: String,
+        partToExtract: String, key: String): Unit = {
+      checkEvaluation(
+        ParseUrl(Literal.create(urlStr, StringType), Literal.create(partToExtract, StringType),
+          Literal.create(key, StringType)), expected)
+    }
+
+    checkParseUrl("spark.apache.org", "http://spark.apache.org/path?query=1", "HOST")
+    checkParseUrl("/path", "http://spark.apache.org/path?query=1", "PATH")
+    checkParseUrl("query=1", "http://spark.apache.org/path?query=1", "QUERY")
+    checkParseUrl("Ref", "http://spark.apache.org/path?query=1#Ref", "REF")
+    checkParseUrl("http", "http://spark.apache.org/path?query=1", "PROTOCOL")
+    checkParseUrl("/path?query=1", "http://spark.apache.org/path?query=1", "FILE")
+    checkParseUrl("spark.apache.org:8080", "http://spark.apache.org:8080/path?query=1", "AUTHORITY")
+    checkParseUrl("userinfo", "http://userinfo@spark.apache.org/path?query=1", "USERINFO")
+    checkParseUrlWithKey("1", "http://spark.apache.org/path?query=1", "QUERY", "query")
+
+    // Null checking
+    checkParseUrl(null, null, "HOST")
+    checkParseUrl(null, "http://spark.apache.org/path?query=1", null)
+    checkParseUrl(null, null, null)
+    checkParseUrl(null, "test", "HOST")
+    checkParseUrl(null, "http://spark.apache.org/path?query=1", "NO")
+    checkParseUrlWithKey(null, "http://spark.apache.org/path?query=1", "HOST", "query")
+    checkParseUrlWithKey(null, "http://spark.apache.org/path?query=1", "QUERY", "quer")
+    checkParseUrlWithKey(null, "http://spark.apache.org/path?query=1", "QUERY", null)
+    checkParseUrlWithKey(null, "http://spark.apache.org/path?query=1", "QUERY", "")
+
+    // arguments checking
+    assert(ParseUrl(Literal("1")).checkInputDataTypes().isFailure)
+    assert(ParseUrl(Literal("1"), Literal("2"), Literal("3"), Literal("4"))
+      .checkInputDataTypes().isFailure)
+  }
 }

sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala

@@ -226,6 +226,21 @@ class StringFunctionsSuite extends QueryTest with SharedSQLContext {
       Row("???hi", "hi???", "h", "h"))
   }
 
+  test("string parse_url function") {
+    val df = Seq[String](("http://userinfo@spark.apache.org/path?query=1#Ref"))
+      .toDF("url")
+
+    checkAnswer(
+      df.selectExpr(
+        "parse_url(url, 'HOST')", "parse_url(url, 'PATH')",
+        "parse_url(url, 'QUERY')", "parse_url(url, 'REF')",
+        "parse_url(url, 'PROTOCOL')", "parse_url(url, 'FILE')",
+        "parse_url(url, 'AUTHORITY')", "parse_url(url, 'USERINFO')",
+        "parse_url(url, 'QUERY', 'query')"),
+      Row("spark.apache.org", "/path", "query=1", "Ref",
+        "http", "/path?query=1", "userinfo@spark.apache.org", "userinfo", "1"))
+  }
+
   test("string repeat function") {
     val df = Seq(("hi", 2)).toDF("a", "b")
 

sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveSessionCatalog.scala

@@ -240,7 +240,7 @@ private[sql] class HiveSessionCatalog(
   private val hiveFunctions = Seq(
     "hash", "java_method", "histogram_numeric",
     "map_keys", "map_values",
-    "parse_url", "percentile", "percentile_approx", "reflect", "sentences", "stack", "str_to_map",
+    "percentile", "percentile_approx", "reflect", "sentences", "stack", "str_to_map",
     "xpath", "xpath_double", "xpath_float", "xpath_int", "xpath_long",
     "xpath_number", "xpath_short", "xpath_string",