diff --git a/python/pyspark/sql/functions.py b/python/pyspark/sql/functions.py
index 81dc7d832eaf9506209e3d64738cd3e3dc29c1b8..96975f54ff695a07eef47416c54ba05175d3883b 100644
--- a/python/pyspark/sql/functions.py
+++ b/python/pyspark/sql/functions.py
@@ -64,7 +64,7 @@ __all__ += [
'year', 'quarter', 'month', 'hour', 'minute', 'second',
'dayofmonth', 'dayofyear', 'weekofyear']
-__all__ += ['soundex']
+__all__ += ['soundex', 'substring', 'substring_index']
def _create_function(name, doc=""):
@@ -923,6 +923,22 @@ def trunc(date, format):
return Column(sc._jvm.functions.trunc(_to_java_column(date), format))
+@since(1.5)
+@ignore_unicode_prefix
+def substring(str, pos, len):
+ """
+ Substring starts at `pos` and is of length `len` when str is String type or
+ returns the slice of byte array that starts at `pos` in byte and is of length `len`
+ when str is Binary type
+
+ >>> df = sqlContext.createDataFrame([('abcd',)], ['s',])
+ >>> df.select(substring(df.s, 1, 2).alias('s')).collect()
+ [Row(s=u'ab')]
+ """
+ sc = SparkContext._active_spark_context
+ return Column(sc._jvm.functions.substring(_to_java_column(str), pos, len))
+
+
@since(1.5)
@ignore_unicode_prefix
def substring_index(str, delim, count):
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringOperations.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringOperations.scala
index 3ce5d6a9c7e02c0b5c37fa10b56cb59e49049f86..4d78c55497a3464148b707d3d35aaf44e36a6bc7 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringOperations.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringOperations.scala
@@ -18,6 +18,7 @@
package org.apache.spark.sql.catalyst.expressions
import java.text.DecimalFormat
+import java.util.Arrays
import java.util.Locale
import java.util.regex.{MatchResult, Pattern}
@@ -679,6 +680,34 @@ case class StringSplit(str: Expression, pattern: Expression)
override def prettyName: String = "split"
}
+object Substring {
+ def subStringBinarySQL(bytes: Array[Byte], pos: Int, len: Int): Array[Byte] = {
+ if (pos > bytes.length) {
+ return Array[Byte]()
+ }
+
+ var start = if (pos > 0) {
+ pos - 1
+ } else if (pos < 0) {
+ bytes.length + pos
+ } else {
+ 0
+ }
+
+ val end = if ((bytes.length - start) < len) {
+ bytes.length
+ } else {
+ start + len
+ }
+
+ start = Math.max(start, 0) // underflow
+ if (start < end) {
+ Arrays.copyOfRange(bytes, start, end)
+ } else {
+ Array[Byte]()
+ }
+ }
+}
/**
* A function that takes a substring of its first argument starting at a given position.
* Defined for String and Binary types.
@@ -690,18 +719,31 @@ case class Substring(str: Expression, pos: Expression, len: Expression)
this(str, pos, Literal(Integer.MAX_VALUE))
}
- override def dataType: DataType = StringType
+ override def dataType: DataType = str.dataType
- override def inputTypes: Seq[DataType] = Seq(StringType, IntegerType, IntegerType)
+ override def inputTypes: Seq[AbstractDataType] =
+ Seq(TypeCollection(StringType, BinaryType), IntegerType, IntegerType)
override def children: Seq[Expression] = str :: pos :: len :: Nil
override def nullSafeEval(string: Any, pos: Any, len: Any): Any = {
- string.asInstanceOf[UTF8String].substringSQL(pos.asInstanceOf[Int], len.asInstanceOf[Int])
+ str.dataType match {
+ case StringType => string.asInstanceOf[UTF8String]
+ .substringSQL(pos.asInstanceOf[Int], len.asInstanceOf[Int])
+ case BinaryType => Substring.subStringBinarySQL(string.asInstanceOf[Array[Byte]],
+ pos.asInstanceOf[Int], len.asInstanceOf[Int])
+ }
}
override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
- defineCodeGen(ctx, ev, (str, pos, len) => s"$str.substringSQL($pos, $len)")
+
+ val cls = classOf[Substring].getName
+ defineCodeGen(ctx, ev, (string, pos, len) => {
+ str.dataType match {
+ case StringType => s"$string.substringSQL($pos, $len)"
+ case BinaryType => s"$cls.subStringBinarySQL($string, $pos, $len)"
+ }
+ })
}
}
@@ -1161,4 +1203,3 @@ case class FormatNumber(x: Expression, d: Expression)
override def prettyName: String = "format_number"
}
-
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/StringExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/StringExpressionsSuite.scala
index ad87ab36fd64552c0d0184e7c71d42fdeb43f1db..89c1e33420f09dd760de798aa58b366351fa01b1 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/StringExpressionsSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/StringExpressionsSuite.scala
@@ -20,7 +20,6 @@ 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 {
@@ -186,6 +185,20 @@ class StringExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
checkEvaluation(s.substr(0), "example", row)
checkEvaluation(s.substring(0, 2), "ex", row)
checkEvaluation(s.substring(0), "example", row)
+
+ val bytes = Array[Byte](1, 2, 3, 4)
+ checkEvaluation(Substring(bytes, 0, 2), Array[Byte](1, 2))
+ checkEvaluation(Substring(bytes, 1, 2), Array[Byte](1, 2))
+ checkEvaluation(Substring(bytes, 2, 2), Array[Byte](2, 3))
+ checkEvaluation(Substring(bytes, 3, 2), Array[Byte](3, 4))
+ checkEvaluation(Substring(bytes, 4, 2), Array[Byte](4))
+ checkEvaluation(Substring(bytes, 8, 2), Array[Byte]())
+ checkEvaluation(Substring(bytes, -1, 2), Array[Byte](4))
+ checkEvaluation(Substring(bytes, -2, 2), Array[Byte](3, 4))
+ checkEvaluation(Substring(bytes, -3, 2), Array[Byte](2, 3))
+ checkEvaluation(Substring(bytes, -4, 2), Array[Byte](1, 2))
+ checkEvaluation(Substring(bytes, -5, 2), Array[Byte](1))
+ checkEvaluation(Substring(bytes, -8, 2), Array[Byte]())
}
test("string substring_index function") {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
index 3c9421f5cd14be8e544ad3a71f79c28634ca2787..babfe2187999c6d18870f0052df44d715045bd52 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
@@ -1725,6 +1725,17 @@ object functions {
*/
def format_number(x: Column, d: Int): Column = FormatNumber(x.expr, lit(d).expr)
+ /**
+ * Substring starts at `pos` and is of length `len` when str is String type or
+ * returns the slice of byte array that starts at `pos` in byte and is of length `len`
+ * when str is Binary type
+ *
+ * @group string_funcs
+ * @since 1.5.0
+ */
+ def substring(str: Column, pos: Int, len: Int): Column =
+ Substring(str.expr, lit(pos).expr, lit(len).expr)
+
/**
* Computes the Levenshtein distance of the two given string columns.
* @group string_funcs
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala
index 628da95298aa9bdaf65475bf05c0c34547517ee5..f40233db0aebdc6090dca54256cd0c0fec29de73 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala
@@ -103,6 +103,16 @@ class StringFunctionsSuite extends QueryTest {
Row("AQIDBA==", bytes))
}
+ test("string / binary substring function") {
+ // scalastyle:off
+ // non ascii characters are not allowed in the code, so we disable the scalastyle here.
+ val df = Seq(("1世3", Array[Byte](1, 2, 3, 4))).toDF("a", "b")
+ checkAnswer(df.select(substring($"a", 1, 2)), Row("1世"))
+ checkAnswer(df.select(substring($"b", 2, 2)), Row(Array[Byte](2,3)))
+ checkAnswer(df.selectExpr("substring(a, 1, 2)"), Row("1世"))
+ // scalastyle:on
+ }
+
test("string encode/decode function") {
val bytes = Array[Byte](-27, -92, -89, -27, -115, -125, -28, -72, -106, -25, -107, -116)
// scalastyle:off
diff --git a/unsafe/src/main/java/org/apache/spark/unsafe/types/UTF8String.java b/unsafe/src/main/java/org/apache/spark/unsafe/types/UTF8String.java
index f6dafe94c617a3a46d549479e952e64bce12fb3c..208503d2fd5dc283b8bf60c9a5b28248ff4ccc73 100644
--- a/unsafe/src/main/java/org/apache/spark/unsafe/types/UTF8String.java
+++ b/unsafe/src/main/java/org/apache/spark/unsafe/types/UTF8String.java
@@ -198,7 +198,7 @@ public final class UTF8String implements Comparable<UTF8String>, Serializable {
*/
public UTF8String substring(final int start, final int until) {
if (until <= start || start >= numBytes) {
- return UTF8String.EMPTY_UTF8;
+ return EMPTY_UTF8;
}
int i = 0;
@@ -214,9 +214,13 @@ public final class UTF8String implements Comparable<UTF8String>, Serializable {
c += 1;
}
- byte[] bytes = new byte[i - j];
- copyMemory(base, offset + j, bytes, BYTE_ARRAY_OFFSET, i - j);
- return fromBytes(bytes);
+ if (i > j) {
+ byte[] bytes = new byte[i - j];
+ copyMemory(base, offset + j, bytes, BYTE_ARRAY_OFFSET, i - j);
+ return fromBytes(bytes);
+ } else {
+ return EMPTY_UTF8;
+ }
}
public UTF8String substringSQL(int pos, int length) {
@@ -226,8 +230,9 @@ public final class UTF8String implements Comparable<UTF8String>, Serializable {
// refers to element i-1 in the sequence. If a start index i is less than 0, it refers
// to the -ith element before the end of the sequence. If a start index i is 0, it
// refers to the first element.
- int start = (pos > 0) ? pos -1 : ((pos < 0) ? numChars() + pos : 0);
- int end = (length == Integer.MAX_VALUE) ? Integer.MAX_VALUE : start + length;
+ int len = numChars();
+ int start = (pos > 0) ? pos -1 : ((pos < 0) ? len + pos : 0);
+ int end = (length == Integer.MAX_VALUE) ? len : start + length;
return substring(start, end);
}