-
Notifications
You must be signed in to change notification settings - Fork 28.3k
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 all 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,61 @@ 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 foldable: Boolean = strExpr.foldable && delimExpr.foldable && countExpr.foldable | ||
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 str.asInstanceOf[UTF8String].subStringIndex( | ||
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"${str.primitive}.subStringIndex(${delim.primitive}, ${count.primitive})" | ||
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. | ||
|
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 ?