Github user mn-mikke commented on a diff in the pull request:
https://github.com/apache/spark/pull/21103#discussion_r205397226
--- Diff:
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala
---
@@ -3805,3 +3799,330 @@ object ArrayUnion {
new GenericArrayData(arrayBuffer)
}
}
+
+/**
+ * Returns an array of the elements in the intersect of x and y, without
duplicates
+ */
+@ExpressionDescription(
+ usage = """
+ _FUNC_(array1, array2) - Returns an array of the elements in array1 but
not in array2,
+ without duplicates.
+ """,
+ examples = """
+ Examples:
+ > SELECT _FUNC_(array(1, 2, 3), array(1, 3, 5));
+ array(2)
+ """,
+ since = "2.4.0")
+case class ArrayExcept(left: Expression, right: Expression) extends
ArraySetLike {
+ override def dataType: DataType = left.dataType
+
+ var hsInt: OpenHashSet[Int] = _
+ var hsLong: OpenHashSet[Long] = _
+
+ def assignInt(array: ArrayData, idx: Int, resultArray: ArrayData, pos:
Int): Boolean = {
+ val elem = array.getInt(idx)
+ if (!hsInt.contains(elem)) {
+ if (resultArray != null) {
+ resultArray.setInt(pos, elem)
+ }
+ hsInt.add(elem)
+ true
+ } else {
+ false
+ }
+ }
+
+ def assignLong(array: ArrayData, idx: Int, resultArray: ArrayData, pos:
Int): Boolean = {
+ val elem = array.getLong(idx)
+ if (!hsLong.contains(elem)) {
+ if (resultArray != null) {
+ resultArray.setLong(pos, elem)
+ }
+ hsLong.add(elem)
+ true
+ } else {
+ false
+ }
+ }
+
+ def evalIntLongPrimitiveType(
+ array1: ArrayData,
+ array2: ArrayData,
+ resultArray: ArrayData,
+ isLongType: Boolean): Int = {
+ // store elements into resultArray
+ var notFoundNullElement = true
+ var i = 0
+ while (i < array2.numElements()) {
+ if (array2.isNullAt(i)) {
+ notFoundNullElement = false
+ } else {
+ val assigned = if (!isLongType) {
+ hsInt.add(array2.getInt(i))
+ } else {
+ hsLong.add(array2.getLong(i))
+ }
+ }
+ i += 1
+ }
+ var pos = 0
+ i = 0
+ while (i < array1.numElements()) {
+ if (array1.isNullAt(i)) {
+ if (notFoundNullElement) {
+ if (resultArray != null) {
+ resultArray.setNullAt(pos)
+ }
+ pos += 1
+ notFoundNullElement = false
+ }
+ } else {
+ val assigned = if (!isLongType) {
+ assignInt(array1, i, resultArray, pos)
+ } else {
+ assignLong(array1, i, resultArray, pos)
+ }
+ if (assigned) {
+ pos += 1
+ }
+ }
+ i += 1
+ }
+ pos
+ }
+
+ override def nullSafeEval(input1: Any, input2: Any): Any = {
+ val array1 = input1.asInstanceOf[ArrayData]
+ val array2 = input2.asInstanceOf[ArrayData]
+
+ if (elementTypeSupportEquals) {
+ elementType match {
+ case IntegerType =>
+ // avoid boxing of primitive int array elements
+ // calculate result array size
+ hsInt = new OpenHashSet[Int]
+ val elements = evalIntLongPrimitiveType(array1, array2, null,
false)
+ // allocate result array
+ hsInt = new OpenHashSet[Int]
+ val resultArray = if (UnsafeArrayData.shouldUseGenericArrayData(
+ IntegerType.defaultSize, elements)) {
+ new GenericArrayData(new Array[Any](elements))
+ } else {
+ UnsafeArrayData.forPrimitiveArray(
+ Platform.INT_ARRAY_OFFSET, elements, IntegerType.defaultSize)
+ }
+ // assign elements into the result array
+ evalIntLongPrimitiveType(array1, array2, resultArray, false)
+ resultArray
+ case LongType =>
+ // avoid boxing of primitive long array elements
+ // calculate result array size
+ hsLong = new OpenHashSet[Long]
+ val elements = evalIntLongPrimitiveType(array1, array2, null,
true)
+ // allocate result array
+ hsLong = new OpenHashSet[Long]
+ val resultArray = if (UnsafeArrayData.shouldUseGenericArrayData(
+ LongType.defaultSize, elements)) {
+ new GenericArrayData(new Array[Any](elements))
+ } else {
+ UnsafeArrayData.forPrimitiveArray(
+ Platform.LONG_ARRAY_OFFSET, elements, LongType.defaultSize)
+ }
+ // assign elements into the result array
+ evalIntLongPrimitiveType(array1, array2, resultArray, true)
+ resultArray
+ case _ =>
+ val hs = new OpenHashSet[Any]
+ var notFoundNullElement = true
+ var i = 0
+ while (i < array2.numElements()) {
+ if (array2.isNullAt(i)) {
+ notFoundNullElement = false
+ } else {
+ val elem = array2.get(i, elementType)
+ hs.add(elem)
+ }
+ i += 1
+ }
+ val arrayBuffer = new scala.collection.mutable.ArrayBuffer[Any]
+ i = 0
+ while (i < array1.numElements()) {
+ if (array1.isNullAt(i)) {
+ if (notFoundNullElement) {
+ arrayBuffer += null
+ notFoundNullElement = false
+ }
+ } else {
+ val elem = array1.get(i, elementType)
+ if (!hs.contains(elem)) {
+ arrayBuffer += elem
+ hs.add(elem)
+ }
+ }
+ i += 1
+ }
+ new GenericArrayData(arrayBuffer)
+ }
+ } else {
+ ArrayExcept.exceptOrdering(array1, array2, elementType, ordering)
+ }
+ }
+
+ override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ val i = ctx.freshName("i")
+ val pos = ctx.freshName("pos")
+ val value = ctx.freshName("value")
+ val size = ctx.freshName("size")
+ val (postFix, openHashElementType, getter, setter, javaTypeName,
castOp, arrayBuilder) =
+ if (elementTypeSupportEquals) {
+ elementType match {
+ case ByteType | ShortType | IntegerType | LongType =>
+ val ptName = CodeGenerator.primitiveTypeName(elementType)
+ val unsafeArray = ctx.freshName("unsafeArray")
+ (if (elementType == LongType) s"$$mcJ$$sp" else s"$$mcI$$sp",
+ if (elementType == LongType) "Long" else "Int",
+ s"get$ptName($i)", s"set$ptName($pos, $value)",
CodeGenerator.javaType(elementType),
+ if (elementType == LongType) "(long)" else "(int)",
+ s"""
+ |${ctx.createUnsafeArray(unsafeArray, size, elementType,
s" $prettyName failed.")}
+ |${ev.value} = $unsafeArray;
+ """.stripMargin)
+ case _ =>
+ val genericArrayData = classOf[GenericArrayData].getName
+ val et = ctx.addReferenceObj("elementType", elementType)
+ ("", "Object",
+ s"get($i, $et)", s"update($pos, $value)", "Object", "",
+ s"${ev.value} = new $genericArrayData(new Object[$size]);")
+ }
+ } else {
+ ("", "", "", "", "", "", "")
+ }
+
+ nullSafeCodeGen(ctx, ev, (array1, array2) => {
+ if (openHashElementType != "") {
+ // Here, we ensure elementTypeSupportEquals is true
+ val notFoundNullElement = ctx.freshName("notFoundNullElement")
+ val openHashSet = classOf[OpenHashSet[_]].getName
+ val classTag =
s"scala.reflect.ClassTag$$.MODULE$$.$openHashElementType()"
+ val hs = ctx.freshName("hs")
+ val arrayData = classOf[ArrayData].getName
+ val arrays = ctx.freshName("arrays")
+ val array = ctx.freshName("array")
+ val arrayDataIdx = ctx.freshName("arrayDataIdx")
+ s"""
+ |$openHashSet $hs = new $openHashSet$postFix($classTag);
+ |boolean $notFoundNullElement = true;
+ |int $size = 0;
+ |for (int $i = 0; $i < $array2.numElements(); $i++) {
+ | if ($array2.isNullAt($i)) {
+ | $notFoundNullElement = false;
+ | } else {
+ | $hs.add$postFix($array2.$getter);
+ | }
+ |}
+ |for (int $i = 0; $i < $array1.numElements(); $i++) {
+ | if ($array1.isNullAt($i)) {
+ | if ($notFoundNullElement) {
+ | $size++;
+ | $notFoundNullElement = false;
+ | }
+ | } else {
+ | $javaTypeName $value = $array1.$getter;
+ | if (!$hs.contains($castOp $value)) {
+ | $hs.add$postFix($value);
+ | $size++;
+ | }
+ | }
+ |}
+ |$arrayBuilder
+ |$hs = new $openHashSet$postFix($classTag);
+ |$notFoundNullElement = true;
+ |int $pos = 0;
+ |for (int $i = 0; $i < $array2.numElements(); $i++) {
+ | if ($array2.isNullAt($i)) {
+ | $notFoundNullElement = false;
+ | } else {
+ | $hs.add$postFix($array2.$getter);
+ | }
+ |}
+ |for (int $i = 0; $i < $array1.numElements(); $i++) {
+ | if ($array1.isNullAt($i)) {
+ | if ($notFoundNullElement) {
+ | ${ev.value}.setNullAt($pos++);
+ | $notFoundNullElement = false;
+ | }
+ | } else {
+ | $javaTypeName $value = $array1.$getter;
+ | if (!$hs.contains($castOp $value)) {
+ | $hs.add$postFix($value);
+ | ${ev.value}.$setter;
+ | $pos++;
+ | }
+ | }
+ |}
+ """.stripMargin
+ } else {
+ val arrayExcept = classOf[ArrayExcept].getName
+ val et = ctx.addReferenceObj("elementTypeIntersect", elementType)
+ val order = ctx.addReferenceObj("orderingIntersect", ordering)
+ val method = "exceptOrdering"
+ s"${ev.value} = $arrayExcept$$.MODULE$$.$method($array1, $array2,
$et, $order);"
+ }
+ })
+ }
+
+ override def prettyName: String = "array_except"
+}
+
+object ArrayExcept {
+ def exceptOrdering(
+ array1: ArrayData,
+ array2: ArrayData,
+ elementType: DataType,
+ ordering: Ordering[Any]): ArrayData = {
+ val arrayBuffer = new scala.collection.mutable.ArrayBuffer[Any]
+ var scannedNullElements = false
+ var i = 0
+ while (i < array1.numElements()) {
+ var found = false
+ var elem1 = array1.get(i, elementType)
+ if (array1.isNullAt(i)) {
+ if (!scannedNullElements) {
+ var j = 0
+ while (!found && j < array2.numElements()) {
+ found = array2.isNullAt(j)
+ j += 1
+ }
+ // array2 is scanned only once for null element
+ scannedNullElements = true
+ } else {
+ found = true
+ }
+ } else {
+ var j = 0
+ while (!found && j < array2.numElements()) {
+ if (!array2.isNullAt(j)) {
--- End diff --
nit: Couldn't we shorten the null check as on the line
[4116](https://github.com/apache/spark/pull/21103/files#diff-9853dcf5ce3d2ac1e94d473197ff5768R4116)?
---
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
