forked from apache/spark
-
Notifications
You must be signed in to change notification settings - Fork 4
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[SPARK-50056][SQL] Codegen Support for ParseUrl (by Invoke & RuntimeR…
…eplaceable) ### What changes were proposed in this pull request? The pr aims to add `Codegen` Support for `parse_url`. ### Why are the changes needed? - improve codegen coverage. - simplified code. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Pass GA & Existed UT (eg: UrlFunctionsSuite#`*parse_url*`) ### Was this patch authored or co-authored using generative AI tooling? No. Closes apache#48579 from panbingkun/SPARK-50056. Authored-by: panbingkun <panbingkun@baidu.com> Signed-off-by: Wenchen Fan <wenchen@databricks.com>
- Loading branch information
1 parent
4f95a7f
commit 00bff28
Showing
7 changed files
with
194 additions
and
145 deletions.
There are no files selected for viewing
148 changes: 148 additions & 0 deletions
148
...src/main/scala/org/apache/spark/sql/catalyst/expressions/url/UrlExpressionEvalUtils.scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,148 @@ | ||
/* | ||
* Licensed to the Apache Software Foundation (ASF) under one or more | ||
* contributor license agreements. See the NOTICE file distributed with | ||
* this work for additional information regarding copyright ownership. | ||
* The ASF licenses this file to You under the Apache License, Version 2.0 | ||
* (the "License"); you may not use this file except in compliance with | ||
* the License. You may obtain a copy of the License at | ||
* | ||
* http://www.apache.org/licenses/LICENSE-2.0 | ||
* | ||
* Unless required by applicable law or agreed to in writing, software | ||
* distributed under the License is distributed on an "AS IS" BASIS, | ||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
* See the License for the specific language governing permissions and | ||
* limitations under the License. | ||
*/ | ||
package org.apache.spark.sql.catalyst.expressions.url | ||
|
||
import java.net.{URI, URISyntaxException} | ||
import java.util.regex.Pattern | ||
|
||
import org.apache.spark.sql.errors.QueryExecutionErrors | ||
import org.apache.spark.unsafe.types.UTF8String | ||
|
||
case class ParseUrlEvaluator( | ||
url: UTF8String, | ||
extractPart: UTF8String, | ||
pattern: UTF8String, | ||
failOnError: Boolean) { | ||
|
||
import ParseUrlEvaluator._ | ||
|
||
private lazy val cachedUrl: URI = | ||
if (url != null) getUrl(url, failOnError) else null | ||
|
||
private lazy val cachedExtractPartFunc: URI => String = | ||
if (extractPart != null) getExtractPartFunc(extractPart) else null | ||
|
||
private lazy val cachedPattern: Pattern = | ||
if (pattern != null) getPattern(pattern) else null | ||
|
||
private def extractValueFromQuery(query: UTF8String, pattern: Pattern): UTF8String = { | ||
val m = pattern.matcher(query.toString) | ||
if (m.find()) { | ||
UTF8String.fromString(m.group(2)) | ||
} else { | ||
null | ||
} | ||
} | ||
|
||
private def extractFromUrl(url: URI, partToExtract: UTF8String): UTF8String = { | ||
if (cachedExtractPartFunc ne null) { | ||
UTF8String.fromString(cachedExtractPartFunc(url)) | ||
} else { | ||
UTF8String.fromString(getExtractPartFunc(partToExtract)(url)) | ||
} | ||
} | ||
|
||
private def parseUrlWithoutKey(url: UTF8String, partToExtract: UTF8String): UTF8String = { | ||
if (cachedUrl ne null) { | ||
extractFromUrl(cachedUrl, partToExtract) | ||
} else { | ||
val currentUrl = getUrl(url, failOnError) | ||
if (currentUrl ne null) { | ||
extractFromUrl(currentUrl, partToExtract) | ||
} else { | ||
null | ||
} | ||
} | ||
} | ||
|
||
final def evaluate(url: UTF8String, path: UTF8String): Any = { | ||
parseUrlWithoutKey(url, path) | ||
} | ||
|
||
final def evaluate(url: UTF8String, path: UTF8String, key: UTF8String): Any = { | ||
if (path != QUERY) return null | ||
|
||
val query = parseUrlWithoutKey(url, path) | ||
if (query eq null) return null | ||
|
||
if (cachedPattern ne null) { | ||
extractValueFromQuery(query, cachedPattern) | ||
} else { | ||
extractValueFromQuery(query, getPattern(key)) | ||
} | ||
} | ||
} | ||
|
||
object ParseUrlEvaluator { | ||
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 = "=([^&]*)" | ||
|
||
private def getPattern(key: UTF8String): Pattern = { | ||
Pattern.compile(REGEXPREFIX + key.toString + REGEXSUBFIX) | ||
} | ||
|
||
private def getUrl(url: UTF8String, failOnError: Boolean): URI = { | ||
try { | ||
new URI(url.toString) | ||
} catch { | ||
case e: URISyntaxException if failOnError => | ||
throw QueryExecutionErrors.invalidUrlError(url, e) | ||
case _: URISyntaxException => null | ||
} | ||
} | ||
|
||
private def getExtractPartFunc(partToExtract: UTF8String): URI => String = { | ||
|
||
// partToExtract match { | ||
// case HOST => _.toURL().getHost | ||
// case PATH => _.toURL().getPath | ||
// case QUERY => _.toURL().getQuery | ||
// case REF => _.toURL().getRef | ||
// case PROTOCOL => _.toURL().getProtocol | ||
// case FILE => _.toURL().getFile | ||
// case AUTHORITY => _.toURL().getAuthority | ||
// case USERINFO => _.toURL().getUserInfo | ||
// case _ => (url: URI) => null | ||
// } | ||
|
||
partToExtract match { | ||
case HOST => _.getHost | ||
case PATH => _.getRawPath | ||
case QUERY => _.getRawQuery | ||
case REF => _.getRawFragment | ||
case PROTOCOL => _.getScheme | ||
case FILE => | ||
(url: URI) => | ||
if (url.getRawQuery ne null) { | ||
url.getRawPath + "?" + url.getRawQuery | ||
} else { | ||
url.getRawPath | ||
} | ||
case AUTHORITY => _.getRawAuthority | ||
case USERINFO => _.getRawUserInfo | ||
case _ => (_: URI) => null | ||
} | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
2 changes: 1 addition & 1 deletion
2
sql/connect/common/src/test/resources/query-tests/explain-results/function_parse_url.explain
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,2 +1,2 @@ | ||
Project [parse_url(g#0, g#0, false) AS parse_url(g, g)#0] | ||
Project [invoke(ParseUrlEvaluator(null,null,null,false).evaluate(g#0, g#0)) AS parse_url(g, g)#0] | ||
+- LocalRelation <empty>, [id#0L, a#0, b#0, d#0, e#0, f#0, g#0] |
Oops, something went wrong.