-
Notifications
You must be signed in to change notification settings - Fork 28.1k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[SPARK-8264][SQL]add substring_index function #7533
Changes from 5 commits
52d7b03
d92951b
12e108f
ac863e9
b19b013
67c253a
9546991
515519b
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -22,7 +22,6 @@ import java.util.Locale | |
import java.util.regex.{MatchResult, Pattern} | ||
|
||
import org.apache.spark.sql.catalyst.InternalRow | ||
import org.apache.spark.sql.catalyst.analysis.UnresolvedException | ||
import org.apache.spark.sql.catalyst.expressions.codegen._ | ||
import org.apache.spark.sql.types._ | ||
import org.apache.spark.unsafe.types.UTF8String | ||
|
@@ -355,6 +354,63 @@ case class StringInstr(str: Expression, substr: Expression) | |
} | ||
} | ||
|
||
/** | ||
* Returns the substring from string str before count occurrences of the delimiter delim. | ||
* If count is positive, everything the left of the final delimiter (counting from left) is | ||
* returned. If count is negative, every to the right of the final delimiter (counting from the | ||
* right) is returned. substring_index performs a case-sensitive match when searching for delim. | ||
*/ | ||
case class Substring_index(strExpr: Expression, delimExpr: Expression, countExpr: Expression) | ||
extends Expression with ImplicitCastInputTypes { | ||
|
||
override def dataType: DataType = StringType | ||
override def inputTypes: Seq[DataType] = Seq(StringType, StringType, IntegerType) | ||
override def nullable: Boolean = strExpr.nullable || delimExpr.nullable || countExpr.nullable | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Need to override the There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I guess it's abstract in Expression which Substring_index inherited from? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Sorry, my bad, actually I mean the function There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Sorry, actually I mean we need to override the function |
||
override def children: Seq[Expression] = Seq(strExpr, delimExpr, countExpr) | ||
override def prettyName: String = "substring_index" | ||
|
||
override def eval(input: InternalRow): Any = { | ||
val str = strExpr.eval(input) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. let's do that like: val str = ...
if (str != null) {
val delim = ...
if (delim != null) {
val count = ... That's kind of short circuit evaluation. |
||
if (str != null) { | ||
val delim = delimExpr.eval(input) | ||
if (delim != null) { | ||
val count = countExpr.eval(input) | ||
if (count != null) { | ||
return UTF8String.subStringIndex( | ||
str.asInstanceOf[UTF8String], | ||
delim.asInstanceOf[UTF8String], | ||
count.asInstanceOf[Int]) | ||
} | ||
} | ||
} | ||
null | ||
} | ||
|
||
override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = { | ||
val str = strExpr.gen(ctx) | ||
val delim = delimExpr.gen(ctx) | ||
val count = countExpr.gen(ctx) | ||
val resultCode = | ||
s"""org.apache.spark.unsafe.types.UTF8String.subStringIndex( | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. if the
|
||
|${str.primitive}, ${delim.primitive}, ${count.primitive})""".stripMargin | ||
s""" | ||
${str.code} | ||
boolean ${ev.isNull} = true; | ||
${ctx.javaType(dataType)} ${ev.primitive} = ${ctx.defaultValue(dataType)}; | ||
if (!${str.isNull}) { | ||
${delim.code} | ||
if (!${delim.isNull}) { | ||
${count.code} | ||
if (!${count.isNull}) { | ||
${ev.isNull} = false; | ||
${ev.primitive} = $resultCode; | ||
} | ||
} | ||
} | ||
""" | ||
} | ||
} | ||
|
||
/** | ||
* A function that returns the position of the first occurrence of substr | ||
* in given string after position pos. | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -1777,8 +1777,29 @@ object functions { | |
def instr(str: Column, substring: String): Column = StringInstr(str.expr, lit(substring).expr) | ||
|
||
/** | ||
* Locate the position of the first occurrence of substr in a string column. | ||
* Returns the substring from string str before count occurrences of the delimiter delim. | ||
* If count is positive, everything the left of the final delimiter (counting from left) is | ||
* returned. If count is negative, every to the right of the final delimiter (counting from the | ||
* right) is returned. substring_index performs a case-sensitive match when searching for delim. | ||
* | ||
* @group string_funcs | ||
*/ | ||
def substring_index(str: String, delim: String, count: Int): Column = | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Let's remove this version of API. @rxin actually made some clean up, and removed the |
||
substring_index(Column(str), delim, count) | ||
|
||
/** | ||
* Returns the substring from string str before count occurrences of the delimiter delim. | ||
* If count is positive, everything the left of the final delimiter (counting from left) is | ||
* returned. If count is negative, every to the right of the final delimiter (counting from the | ||
* right) is returned. substring_index performs a case-sensitive match when searching for delim. | ||
* | ||
* @group string_funcs | ||
*/ | ||
def substring_index(str: Column, delim: String, count: Int): Column = | ||
Substring_index(str.expr, lit(delim).expr, lit(count).expr) | ||
|
||
/** | ||
* Locate the position of the first occurrence of substr. | ||
* NOTE: The position is not zero based, but 1 based index, returns 0 if substr | ||
* could not be found in str. | ||
* | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Could you use TernaryExpression ?