Skip to content
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.

Sign up for GitHub

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

Jump to bottom

[SPARK-8264][SQL]add substring_index function #7533

Closed
wants to merge 8 commits into from
Closed

[SPARK-8264][SQL]add substring_index function #7533

Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
52d7b03
add substring_index function
zhichao-li Jul 20, 2015
d92951b
add lastIndexOf
zhichao-li Jul 21, 2015
12e108f
refine unittest
zhichao-li Jul 22, 2015
ac863e9
reduce the calling of numChars
zhichao-li Jul 22, 2015
b19b013
add codegen and clean code
zhichao-li Jul 23, 2015
67c253a
hide some apis and clean code
zhichao-li Jul 24, 2015
9546991
scala style
zhichao-li Jul 24, 2015
515519b
add foldable and remove null checking
zhichao-li Jul 24, 2015
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -177,6 +177,7 @@ object FunctionRegistry {
expression[StringSplit]("split"),
expression[Substring]("substr"),
expression[Substring]("substring"),
expression[Substring_index]("substring_index"),
expression[StringTrim]("trim"),
expression[UnBase64]("unbase64"),
expression[Upper]("ucase"),
Expand Down
55 changes: 54 additions & 1 deletion sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringOperations.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -355,6 +354,60 @@ 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 {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you use TernaryExpression ?


override def dataType: DataType = StringType
override def inputTypes: Seq[DataType] = Seq(StringType, StringType, IntegerType)
override def nullable: Boolean = strExpr.nullable || delimExpr.nullable || countExpr.nullable
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Need to override the nullable method, as the default value is false.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess it's abstract in Expression which Substring_index inherited from?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry, my bad, actually I mean the function foldable, not the nullable.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry, actually I mean we need to override the function foldable, which is false by default, but obviously in this expression, which depends on its children.

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)
Copy link
Contributor

Choose a reason for hiding this comment

The 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.
Expand Down
31 changes: 31 additions & 0 deletions ...yst/src/test/scala/org/apache/spark/sql/catalyst/expressions/StringExpressionsSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ package org.apache.spark.sql.catalyst.expressions
import org.apache.spark.SparkFunSuite
import org.apache.spark.sql.catalyst.dsl.expressions._
import org.apache.spark.sql.types._
import org.apache.spark.unsafe.types.UTF8String


class StringExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
Expand Down Expand Up @@ -187,6 +188,36 @@ class StringExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
checkEvaluation(s.substring(0), "example", row)
}

test("string substring_index function") {
checkEvaluation(
Substring_index(Literal("www.apache.org"), Literal("."), Literal(3)), "www.apache.org")
checkEvaluation(
Substring_index(Literal("www.apache.org"), Literal("."), Literal(2)), "www.apache")
checkEvaluation(
Substring_index(Literal("www.apache.org"), Literal("."), Literal(1)), "www")
checkEvaluation(
Substring_index(Literal("www.apache.org"), Literal("."), Literal(0)), "")
checkEvaluation(
Substring_index(Literal("www.apache.org"), Literal("."), Literal(-3)), "www.apache.org")
checkEvaluation(
Substring_index(Literal("www.apache.org"), Literal("."), Literal(-2)), "apache.org")
checkEvaluation(
Substring_index(Literal("www.apache.org"), Literal("."), Literal(-1)), "org")
checkEvaluation(
Substring_index(Literal(""), Literal("."), Literal(-2)), "")
checkEvaluation(
Substring_index(Literal.create(null, StringType), Literal("."), Literal(-2)), null)
checkEvaluation(Substring_index(
Literal("www.apache.org"), Literal.create(null, StringType), Literal(-2)), null)
// non ascii chars
// scalastyle:off
checkEvaluation(
Substring_index(Literal("大千世界大千世界"), Literal( "千"), Literal(2)), "大千世界大")
// scalastyle:on
checkEvaluation(
Substring_index(Literal("www||apache||org"), Literal( "||"), Literal(2)), "www||apache")
}

test("LIKE literal Regular Expression") {
checkEvaluation(Literal.create(null, StringType).like("a"), null)
checkEvaluation(Literal.create("a", StringType).like(Literal.create(null, StringType)), null)
Expand Down
12 changes: 11 additions & 1 deletion sql/core/src/main/scala/org/apache/spark/sql/functions.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1777,8 +1777,18 @@ 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: 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.
*
Expand Down
57 changes: 57 additions & 0 deletions sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,63 @@ class StringFunctionsSuite extends QueryTest {
Row(1))
}

test("string substring_index function") {
val df = Seq(("www.apache.org", ".", "zz")).toDF("a", "b", "c")
checkAnswer(
df.select(substring_index($"a", ".", 3)),
Row("www.apache.org"))
checkAnswer(
df.select(substring_index($"a", ".", 2)),
Row("www.apache"))
checkAnswer(
df.select(substring_index($"a", ".", 1)),
Row("www"))
checkAnswer(
df.select(substring_index($"a", ".", 0)),
Row(""))
checkAnswer(
df.select(substring_index(lit("www.apache.org"), ".", -1)),
Row("org"))
checkAnswer(
df.select(substring_index(lit("www.apache.org"), ".", -2)),
Row("apache.org"))
checkAnswer(
df.select(substring_index(lit("www.apache.org"), ".", -3)),
Row("www.apache.org"))
// str is empty string
checkAnswer(
df.select(substring_index(lit(""), ".", 1)),
Row(""))
// empty string delim
checkAnswer(
df.select(substring_index(lit("www.apache.org"), "", 1)),
Row(""))
// delim does not exist in str
checkAnswer(
df.select(substring_index(lit("www.apache.org"), "#", 1)),
Row("www.apache.org"))
// delim is 2 chars
checkAnswer(
df.select(substring_index(lit("www||apache||org"), "||", 2)),
Row("www||apache"))
checkAnswer(
df.select(substring_index(lit("www||apache||org"), "||", -2)),
Row("apache||org"))
// null
checkAnswer(
df.select(substring_index(lit(null), "||", 2)),
Row(null))
checkAnswer(
df.select(substring_index(lit("www.apache.org"), null, 2)),
Row(null))
// non ascii chars
// scalastyle:off
checkAnswer(
df.selectExpr("""substring_index("大千世界大千世界", "千", 2)"""),
Row("大千世界大"))
// scalastyle:on
}

test("string locate function") {
val df = Seq(("aaads", "aa", "zz", 1)).toDF("a", "b", "c", "d")

Expand Down
Loading