address comments(3)

Author: huaxingao

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala

@@ -19,6 +19,7 @@ package org.apache.spark.sql.catalyst.expressions
 import java.util.Comparator
 
 import scala.collection.mutable
+
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.analysis.{TypeCheckResult, TypeCoercion}
 import org.apache.spark.sql.catalyst.expressions.ArraySortLike.NullOrder
@@ -2386,22 +2387,20 @@ case class ArrayDistinct(child: Expression)
       val i = ctx.freshName("i")
       val j = ctx.freshName("j")
       val hs = ctx.freshName("hs")
+      val foundNullElement = ctx.freshName("foundNullElement")
       val distinctArrayLen = ctx.freshName("distinctArrayLen")
       val getValue = CodeGenerator.getValue(array, elementType, i)
       val openHashSet = classOf[OpenHashSet[_]].getName
       val classTag = s"scala.reflect.ClassTag$$.MODULE$$.Object()"
       s"""
          |int $distinctArrayLen = 0;
+         |boolean $foundNullElement = false;
          |$openHashSet $hs = new $openHashSet($classTag);
          |for (int $i = 0; $i < $array.numElements(); $i++) {
          |  if ($array.isNullAt($i)) {
-         |    int $j;
-         |    for ($j = 0; $j < $i; $j ++) {
-         |      if ($array.isNullAt($j))
-         |        break;
-         |    }
-         |    if ($i == $j) {
+         |    if (!($foundNullElement)) {
          |      $distinctArrayLen = $distinctArrayLen + 1;
+         |      $foundNullElement = true;
          |    }
          |  }
          |  else {
@@ -2427,6 +2426,7 @@ case class ArrayDistinct(child: Expression)
     val i = ctx.freshName("i")
     val j = ctx.freshName("j")
     val pos = ctx.freshName("pos")
+    val foundNullElement = ctx.freshName("foundNullElement")
     val genericArrayData = classOf[GenericArrayData].getName
     val getValue = CodeGenerator.getValue(inputArray, elementType, i)
 
@@ -2436,17 +2436,14 @@ case class ArrayDistinct(child: Expression)
       s"""
          |Object[] $distinctArr = new Object[$newArraySize];
          |int $pos = 0;
+         |boolean $foundNullElement = false;
          |$openHashSet $hs = new $openHashSet($classTag);
          |for (int $i = 0; $i < $inputArray.numElements(); $i++) {
          |  if ($inputArray.isNullAt($i)) {
-         |    int $j;
-         |    for ($j = 0; $j < $i; $j ++) {
-         |      if ($inputArray.isNullAt($j))
-         |        break;
-         |    }
-         |    if ($i == $j) {
+         |    if (!($foundNullElement)) {
          |      $distinctArr[$pos] = null;
          |      $pos = $pos + 1;
+         |      $foundNullElement = true;
          |    }
          |  }
          |  else {
@@ -2465,17 +2462,14 @@ case class ArrayDistinct(child: Expression)
       s"""
          |${ctx.createUnsafeArray(distinctArr, newArraySize, elementType, s" $prettyName failed.")}
          |int $pos = 0;
+         |boolean $foundNullElement = false;
          |$openHashSet $hs = new $openHashSet($classTag);
          |for (int $i = 0; $i < $inputArray.numElements(); $i++) {
          |  if ($inputArray.isNullAt($i)) {
-         |    int $j;
-         |    for ($j = 0; $j < $i; $j ++) {
-         |      if ($inputArray.isNullAt($j))
-         |        break;
-         |    }
-         |    if ($i == $j) {
+         |    if (!($foundNullElement)) {
          |      $distinctArr.setNullAt($pos);
          |      $pos = $pos + 1;
+         |      $foundNullElement = true;
          |    }
          |  }
          |  else {

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala

@@ -771,13 +771,21 @@ class CollectionExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper
     val a0 = Literal.create(Seq(2, 1, 2, 3, 4, 4, 5), ArrayType(IntegerType))
     val a1 = Literal.create(Seq.empty[Integer], ArrayType(IntegerType))
     val a2 = Literal.create(Seq("b", "a", "a", "c", "b"), ArrayType(StringType))
-    val a3 = Literal.create(Seq("b", null, "a", "a"), ArrayType(StringType))
-    val a4 = Literal.create(Seq(null, null), ArrayType(NullType))
+    val a3 = Literal.create(Seq("b", null, "a", null, "a", null), ArrayType(StringType))
+    val a4 = Literal.create(Seq(null, null, null), ArrayType(NullType))
+    val a5 = Literal.create(Seq(true, false, false, true), ArrayType(BooleanType))
+    val a6 = Literal.create(Seq(1.123, 0.1234, 1.121, 1.123, 1.1230, 1.121, 0.1234),
+      ArrayType(DoubleType))
+    val a7 = Literal.create(Seq(1.123f, 0.1234f, 1.121f, 1.123f, 1.1230f, 1.121f, 0.1234f),
+      ArrayType(FloatType))
 
     checkEvaluation(new ArrayDistinct(a0), Seq(2, 1, 3, 4, 5))
     checkEvaluation(new ArrayDistinct(a1), Seq.empty[Integer])
     checkEvaluation(new ArrayDistinct(a2), Seq("b", "a", "c"))
     checkEvaluation(new ArrayDistinct(a3), Seq("b", null, "a"))
     checkEvaluation(new ArrayDistinct(a4), Seq(null))
+    checkEvaluation(new ArrayDistinct(a5), Seq(true, false))
+    checkEvaluation(new ArrayDistinct(a6), Seq(1.123, 0.1234, 1.121))
+    checkEvaluation(new ArrayDistinct(a7), Seq(1.123f, 0.1234f, 1.121f))
   }
 }