You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ue...@apache.org on 2018/06/01 05:04:32 UTC
spark git commit: [SPARK-23920][SQL] add array_remove to remove all
elements that equal element from array
Repository: spark
Updated Branches:
refs/heads/master 22df953f6 -> 98909c398
[SPARK-23920][SQL] add array_remove to remove all elements that equal element from array
## What changes were proposed in this pull request?
add array_remove to remove all elements that equal element from array
## How was this patch tested?
add unit tests
Author: Huaxin Gao <hu...@us.ibm.com>
Closes #21069 from huaxingao/spark-23920.
Project: http://git-wip-us.apache.org/repos/asf/spark/repo
Commit: http://git-wip-us.apache.org/repos/asf/spark/commit/98909c39
Tree: http://git-wip-us.apache.org/repos/asf/spark/tree/98909c39
Diff: http://git-wip-us.apache.org/repos/asf/spark/diff/98909c39
Branch: refs/heads/master
Commit: 98909c398dbcbffcae8015d36f44f185d3280af6
Parents: 22df953
Author: Huaxin Gao <hu...@us.ibm.com>
Authored: Thu May 31 22:04:26 2018 -0700
Committer: Takuya UESHIN <ue...@databricks.com>
Committed: Thu May 31 22:04:26 2018 -0700
----------------------------------------------------------------------
python/pyspark/sql/functions.py | 16 +++
.../catalyst/analysis/FunctionRegistry.scala | 1 +
.../expressions/collectionOperations.scala | 123 +++++++++++++++++++
.../CollectionExpressionsSuite.scala | 58 +++++++++
.../scala/org/apache/spark/sql/functions.scala | 9 ++
.../spark/sql/DataFrameFunctionsSuite.scala | 29 +++++
6 files changed, 236 insertions(+)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/spark/blob/98909c39/python/pyspark/sql/functions.py
----------------------------------------------------------------------
diff --git a/python/pyspark/sql/functions.py b/python/pyspark/sql/functions.py
index fd656c5..1759195 100644
--- a/python/pyspark/sql/functions.py
+++ b/python/pyspark/sql/functions.py
@@ -1964,6 +1964,22 @@ def element_at(col, extraction):
return Column(sc._jvm.functions.element_at(_to_java_column(col), extraction))
+@since(2.4)
+def array_remove(col, element):
+ """
+ Collection function: Remove all elements that equal to element from the given array.
+
+ :param col: name of column containing array
+ :param element: element to be removed from the array
+
+ >>> df = spark.createDataFrame([([1, 2, 3, 1, 1],), ([],)], ['data'])
+ >>> df.select(array_remove(df.data, 1)).collect()
+ [Row(array_remove(data, 1)=[2, 3]), Row(array_remove(data, 1)=[])]
+ """
+ sc = SparkContext._active_spark_context
+ return Column(sc._jvm.functions.array_remove(_to_java_column(col), element))
+
+
@since(1.4)
def explode(col):
"""Returns a new row for each element in the given array or map.
http://git-wip-us.apache.org/repos/asf/spark/blob/98909c39/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
----------------------------------------------------------------------
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
index 23a4a44..49fb35b 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
@@ -430,6 +430,7 @@ object FunctionRegistry {
expression[Concat]("concat"),
expression[Flatten]("flatten"),
expression[ArrayRepeat]("array_repeat"),
+ expression[ArrayRemove]("array_remove"),
CreateStruct.registryEntry,
// mask functions
http://git-wip-us.apache.org/repos/asf/spark/blob/98909c39/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala
----------------------------------------------------------------------
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala
index 8a877b0..176995a 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala
@@ -2066,3 +2066,126 @@ case class ArrayRepeat(left: Expression, right: Expression)
}
}
+
+/**
+ * Remove all elements that equal to element from the given array
+ */
+@ExpressionDescription(
+ usage = "_FUNC_(array, element) - Remove all elements that equal to element from array.",
+ examples = """
+ Examples:
+ > SELECT _FUNC_(array(1, 2, 3, null, 3), 3);
+ [1,2,null]
+ """, since = "2.4.0")
+case class ArrayRemove(left: Expression, right: Expression)
+ extends BinaryExpression with ImplicitCastInputTypes {
+
+ override def dataType: DataType = left.dataType
+
+ override def inputTypes: Seq[AbstractDataType] = {
+ val elementType = left.dataType match {
+ case t: ArrayType => t.elementType
+ case _ => AnyDataType
+ }
+ Seq(ArrayType, elementType)
+ }
+
+ lazy val elementType: DataType = left.dataType.asInstanceOf[ArrayType].elementType
+
+ @transient private lazy val ordering: Ordering[Any] =
+ TypeUtils.getInterpretedOrdering(right.dataType)
+
+ override def checkInputDataTypes(): TypeCheckResult = {
+ super.checkInputDataTypes() match {
+ case f: TypeCheckResult.TypeCheckFailure => f
+ case TypeCheckResult.TypeCheckSuccess =>
+ TypeUtils.checkForOrderingExpr(right.dataType, s"function $prettyName")
+ }
+ }
+
+ override def nullSafeEval(arr: Any, value: Any): Any = {
+ val newArray = new Array[Any](arr.asInstanceOf[ArrayData].numElements())
+ var pos = 0
+ arr.asInstanceOf[ArrayData].foreach(right.dataType, (i, v) =>
+ if (v == null || !ordering.equiv(v, value)) {
+ newArray(pos) = v
+ pos += 1
+ }
+ )
+ new GenericArrayData(newArray.slice(0, pos))
+ }
+
+ override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ nullSafeCodeGen(ctx, ev, (arr, value) => {
+ val numsToRemove = ctx.freshName("numsToRemove")
+ val newArraySize = ctx.freshName("newArraySize")
+ val i = ctx.freshName("i")
+ val getValue = CodeGenerator.getValue(arr, elementType, i)
+ val isEqual = ctx.genEqual(elementType, value, getValue)
+ s"""
+ |int $numsToRemove = 0;
+ |for (int $i = 0; $i < $arr.numElements(); $i ++) {
+ | if (!$arr.isNullAt($i) && $isEqual) {
+ | $numsToRemove = $numsToRemove + 1;
+ | }
+ |}
+ |int $newArraySize = $arr.numElements() - $numsToRemove;
+ |${genCodeForResult(ctx, ev, arr, value, newArraySize)}
+ """.stripMargin
+ })
+ }
+
+ def genCodeForResult(
+ ctx: CodegenContext,
+ ev: ExprCode,
+ inputArray: String,
+ value: String,
+ newArraySize: String): String = {
+ val values = ctx.freshName("values")
+ val i = ctx.freshName("i")
+ val pos = ctx.freshName("pos")
+ val getValue = CodeGenerator.getValue(inputArray, elementType, i)
+ val isEqual = ctx.genEqual(elementType, value, getValue)
+ if (!CodeGenerator.isPrimitiveType(elementType)) {
+ val arrayClass = classOf[GenericArrayData].getName
+ s"""
+ |int $pos = 0;
+ |Object[] $values = new Object[$newArraySize];
+ |for (int $i = 0; $i < $inputArray.numElements(); $i ++) {
+ | if ($inputArray.isNullAt($i)) {
+ | $values[$pos] = null;
+ | $pos = $pos + 1;
+ | }
+ | else {
+ | if (!($isEqual)) {
+ | $values[$pos] = $getValue;
+ | $pos = $pos + 1;
+ | }
+ | }
+ |}
+ |${ev.value} = new $arrayClass($values);
+ """.stripMargin
+ } else {
+ val primitiveValueTypeName = CodeGenerator.primitiveTypeName(elementType)
+ s"""
+ |${ctx.createUnsafeArray(values, newArraySize, elementType, s" $prettyName failed.")}
+ |int $pos = 0;
+ |for (int $i = 0; $i < $inputArray.numElements(); $i ++) {
+ | if ($inputArray.isNullAt($i)) {
+ | $values.setNullAt($pos);
+ | $pos = $pos + 1;
+ | }
+ | else {
+ | if (!($isEqual)) {
+ | $values.set$primitiveValueTypeName($pos, $getValue);
+ | $pos = $pos + 1;
+ | }
+ | }
+ |}
+ |${ev.value} = $values;
+ """.stripMargin
+ }
+ }
+
+ override def prettyName: String = "array_remove"
+}
http://git-wip-us.apache.org/repos/asf/spark/blob/98909c39/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala
----------------------------------------------------------------------
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala
index 3fc0b08..f8ad624 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala
@@ -622,4 +622,62 @@ class CollectionExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper
checkEvaluation(ArrayRepeat(strArray, Literal(2)), Seq(Seq("hi", "hola"), Seq("hi", "hola")))
checkEvaluation(ArrayRepeat(Literal("hi"), Literal(null, IntegerType)), null)
}
+
+ test("Array remove") {
+ val a0 = Literal.create(Seq(1, 2, 3, 2, 2, 5), ArrayType(IntegerType))
+ val a1 = Literal.create(Seq("b", "a", "a", "c", "b"), ArrayType(StringType))
+ val a2 = Literal.create(Seq[String](null, "", null, ""), ArrayType(StringType))
+ val a3 = Literal.create(Seq.empty[Integer], ArrayType(IntegerType))
+ val a4 = Literal.create(null, ArrayType(StringType))
+ val a5 = Literal.create(Seq(1, null, 8, 9, null), ArrayType(IntegerType))
+ val a6 = Literal.create(Seq(true, false, false, true), ArrayType(BooleanType))
+
+ checkEvaluation(ArrayRemove(a0, Literal(0)), Seq(1, 2, 3, 2, 2, 5))
+ checkEvaluation(ArrayRemove(a0, Literal(1)), Seq(2, 3, 2, 2, 5))
+ checkEvaluation(ArrayRemove(a0, Literal(2)), Seq(1, 3, 5))
+ checkEvaluation(ArrayRemove(a0, Literal(3)), Seq(1, 2, 2, 2, 5))
+ checkEvaluation(ArrayRemove(a0, Literal(5)), Seq(1, 2, 3, 2, 2))
+ checkEvaluation(ArrayRemove(a0, Literal(null, IntegerType)), null)
+
+ checkEvaluation(ArrayRemove(a1, Literal("")), Seq("b", "a", "a", "c", "b"))
+ checkEvaluation(ArrayRemove(a1, Literal("a")), Seq("b", "c", "b"))
+ checkEvaluation(ArrayRemove(a1, Literal("b")), Seq("a", "a", "c"))
+ checkEvaluation(ArrayRemove(a1, Literal("c")), Seq("b", "a", "a", "b"))
+
+ checkEvaluation(ArrayRemove(a2, Literal("")), Seq(null, null))
+ checkEvaluation(ArrayRemove(a2, Literal(null, StringType)), null)
+
+ checkEvaluation(ArrayRemove(a3, Literal(1)), Seq.empty[Integer])
+
+ checkEvaluation(ArrayRemove(a4, Literal("a")), null)
+
+ checkEvaluation(ArrayRemove(a5, Literal(9)), Seq(1, null, 8, null))
+ checkEvaluation(ArrayRemove(a6, Literal(false)), Seq(true, true))
+
+ // complex data types
+ val b0 = Literal.create(Seq[Array[Byte]](Array[Byte](5, 6), Array[Byte](1, 2),
+ Array[Byte](1, 2), Array[Byte](5, 6)), ArrayType(BinaryType))
+ val b1 = Literal.create(Seq[Array[Byte]](Array[Byte](2, 1), null),
+ ArrayType(BinaryType))
+ val b2 = Literal.create(Seq[Array[Byte]](null, Array[Byte](1, 2)),
+ ArrayType(BinaryType))
+ val nullBinary = Literal.create(null, BinaryType)
+
+ val dataToRemove1 = Literal.create(Array[Byte](5, 6), BinaryType)
+ checkEvaluation(ArrayRemove(b0, dataToRemove1),
+ Seq[Array[Byte]](Array[Byte](1, 2), Array[Byte](1, 2)))
+ checkEvaluation(ArrayRemove(b0, nullBinary), null)
+ checkEvaluation(ArrayRemove(b1, dataToRemove1), Seq[Array[Byte]](Array[Byte](2, 1), null))
+ checkEvaluation(ArrayRemove(b2, dataToRemove1), Seq[Array[Byte]](null, Array[Byte](1, 2)))
+
+ val c0 = Literal.create(Seq[Seq[Int]](Seq[Int](1, 2), Seq[Int](3, 4)),
+ ArrayType(ArrayType(IntegerType)))
+ val c1 = Literal.create(Seq[Seq[Int]](Seq[Int](5, 6), Seq[Int](2, 1)),
+ ArrayType(ArrayType(IntegerType)))
+ val c2 = Literal.create(Seq[Seq[Int]](null, Seq[Int](2, 1)), ArrayType(ArrayType(IntegerType)))
+ val dataToRemove2 = Literal.create(Seq[Int](1, 2), ArrayType(IntegerType))
+ checkEvaluation(ArrayRemove(c0, dataToRemove2), Seq[Seq[Int]](Seq[Int](3, 4)))
+ checkEvaluation(ArrayRemove(c1, dataToRemove2), Seq[Seq[Int]](Seq[Int](5, 6), Seq[Int](2, 1)))
+ checkEvaluation(ArrayRemove(c2, dataToRemove2), Seq[Seq[Int]](null, Seq[Int](2, 1)))
+ }
}
http://git-wip-us.apache.org/repos/asf/spark/blob/98909c39/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
----------------------------------------------------------------------
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 443ba2a..a2aae9a 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
@@ -3170,6 +3170,15 @@ object functions {
def array_sort(e: Column): Column = withExpr { ArraySort(e.expr) }
/**
+ * Remove all elements that equal to element from the given array.
+ * @group collection_funcs
+ * @since 2.4.0
+ */
+ def array_remove(column: Column, element: Any): Column = withExpr {
+ ArrayRemove(column.expr, Literal(element))
+ }
+
+ /**
* Creates a new row for each element in the given array or map column.
*
* @group collection_funcs
http://git-wip-us.apache.org/repos/asf/spark/blob/98909c39/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala
----------------------------------------------------------------------
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala
index cc8bad4..59119bb 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala
@@ -1110,6 +1110,35 @@ class DataFrameFunctionsSuite extends QueryTest with SharedSQLContext {
}
+ test("array remove") {
+ val df = Seq(
+ (Array[Int](2, 1, 2, 3), Array("a", "b", "c", "a"), Array("", "")),
+ (Array.empty[Int], Array.empty[String], Array.empty[String]),
+ (null, null, null)
+ ).toDF("a", "b", "c")
+ checkAnswer(
+ df.select(array_remove($"a", 2), array_remove($"b", "a"), array_remove($"c", "")),
+ Seq(
+ Row(Seq(1, 3), Seq("b", "c"), Seq.empty[String]),
+ Row(Seq.empty[Int], Seq.empty[String], Seq.empty[String]),
+ Row(null, null, null))
+ )
+
+ checkAnswer(
+ df.selectExpr("array_remove(a, 2)", "array_remove(b, \"a\")",
+ "array_remove(c, \"\")"),
+ Seq(
+ Row(Seq(1, 3), Seq("b", "c"), Seq.empty[String]),
+ Row(Seq.empty[Int], Seq.empty[String], Seq.empty[String]),
+ Row(null, null, null))
+ )
+
+ val e = intercept[AnalysisException] {
+ Seq(("a string element", "a")).toDF().selectExpr("array_remove(_1, _2)")
+ }
+ assert(e.message.contains("argument 1 requires array type, however, '`_1`' is of string type"))
+ }
+
private def assertValuesDoNotChangeAfterCoalesceOrUnion(v: Column): Unit = {
import DataFrameFunctionsSuite.CodegenFallbackExpr
for ((codegenFallback, wholeStage) <- Seq((true, false), (false, false), (false, true))) {
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org