Skip to content

Commit

Permalink
[SPARK-9156][SQL] codegen StringSplit
Browse files Browse the repository at this point in the history
Jira: https://issues.apache.org/jira/browse/SPARK-9156

Author: Tarek Auel <tarek.auel@googlemail.com>

Closes #7547 from tarekauel/SPARK-9156 and squashes the following commits:

0be2700 [Tarek Auel] [SPARK-9156][SQL] indention fix
b860eaf [Tarek Auel] [SPARK-9156][SQL] codegen StringSplit
5ad6a1f [Tarek Auel] [SPARK-9156] codegen StringSplit
  • Loading branch information
tarekbecker authored and rxin committed Jul 21, 2015
1 parent 047ccc8 commit 6853ac7
Show file tree
Hide file tree
Showing 3 changed files with 28 additions and 4 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -615,17 +615,21 @@ case class StringSpace(child: Expression)
* Splits str around pat (pattern is a regular expression).
*/
case class StringSplit(str: Expression, pattern: Expression)
extends BinaryExpression with ImplicitCastInputTypes with CodegenFallback {
extends BinaryExpression with ImplicitCastInputTypes {

override def left: Expression = str
override def right: Expression = pattern
override def dataType: DataType = ArrayType(StringType)
override def inputTypes: Seq[DataType] = Seq(StringType, StringType)

override def nullSafeEval(string: Any, regex: Any): Any = {
val splits =
string.asInstanceOf[UTF8String].toString.split(regex.asInstanceOf[UTF8String].toString, -1)
splits.toSeq.map(UTF8String.fromString)
string.asInstanceOf[UTF8String].split(regex.asInstanceOf[UTF8String], -1).toSeq
}

override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
nullSafeCodeGen(ctx, ev, (str, pattern) =>
s"""${ev.primitive} = scala.collection.JavaConversions.asScalaBuffer(
java.util.Arrays.asList($str.split($pattern, -1)));""")
}

override def prettyName: String = "split"
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -487,6 +487,15 @@ public static UTF8String concatWs(UTF8String separator, UTF8String... inputs) {
return fromBytes(result);
}

public UTF8String[] split(UTF8String pattern, int limit) {
String[] splits = toString().split(pattern.toString(), limit);
UTF8String[] res = new UTF8String[splits.length];
for (int i = 0; i < res.length; i++) {
res[i] = fromString(splits[i]);
}
return res;
}

@Override
public String toString() {
try {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
package org.apache.spark.unsafe.types;

import java.io.UnsupportedEncodingException;
import java.util.Arrays;

import org.junit.Test;

Expand Down Expand Up @@ -270,6 +271,16 @@ public void pad() {
fromString("数据砖头孙行者孙行者孙行"),
fromString("数据砖头").rpad(12, fromString("孙行者")));
}

@Test
public void split() {
assertTrue(Arrays.equals(fromString("ab,def,ghi").split(fromString(","), -1),
new UTF8String[]{fromString("ab"), fromString("def"), fromString("ghi")}));
assertTrue(Arrays.equals(fromString("ab,def,ghi").split(fromString(","), 2),
new UTF8String[]{fromString("ab"), fromString("def,ghi")}));
assertTrue(Arrays.equals(fromString("ab,def,ghi").split(fromString(","), 2),
new UTF8String[]{fromString("ab"), fromString("def,ghi")}));
}

@Test
public void levenshteinDistance() {
Expand Down

0 comments on commit 6853ac7

Please sign in to comment.