address review comments

Author: kiszk

sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeArrayData.java

@@ -477,7 +477,7 @@ public static UnsafeArrayData forPrimitiveArray(int offset, int length, int elem
     return fromPrimitiveArray(null, offset, length, elementSize);
   }
 
-  public static boolean useGenericArrayData(int elementSize, int length) {
+  public static boolean canUseGenericArrayData(int elementSize, int length) {
     final long headerInBytes = calculateHeaderPortionInBytes(length);
     final long valueRegionInBytes = (long)elementSize * length;
     final long totalSizeInLongs = (headerInBytes + valueRegionInBytes + 7) / 8;

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

@@ -3262,23 +3262,13 @@ case class ArrayDistinct(child: Expression)
   override def prettyName: String = "array_distinct"
 }
 
-object ArraySetLike {
-  def throwUnionLengthOverflowException(length: Int): Unit = {
-    throw new RuntimeException(s"Unsuccessful try to union arrays with $length " +
-      s"elements due to exceeding the array size limit " +
-      s"${ByteArrayMethods.MAX_ROUNDED_ARRAY_LENGTH}.")
-  }
-}
-
-
+/**
+ * Will become common base class for [[ArrayUnion]], ArrayIntersect, and ArrayExcept.
+ */
 abstract class ArraySetLike extends BinaryArrayExpressionWithImplicitCast {
   override def dataType: DataType = {
-    val dataTypes = children.map(_.dataType)
-    dataTypes.headOption.map {
-      case ArrayType(et, _) =>
-        ArrayType(et, dataTypes.exists(_.asInstanceOf[ArrayType].containsNull))
-      case dt => dt
-    }.getOrElse(StringType)
+    val dataTypes = children.map(_.dataType.asInstanceOf[ArrayType])
+    ArrayType(elementType, dataTypes.exists(_.containsNull))
   }
 
   override def checkInputDataTypes(): TypeCheckResult = {
@@ -3301,6 +3291,15 @@ abstract class ArraySetLike extends BinaryArrayExpressionWithImplicitCast {
   }
 }
 
+object ArraySetLike {
+  def throwUnionLengthOverflowException(length: Int): Unit = {
+    throw new RuntimeException(s"Unsuccessful try to union arrays with $length " +
+      s"elements due to exceeding the array size limit " +
+      s"${ByteArrayMethods.MAX_ROUNDED_ARRAY_LENGTH}.")
+  }
+}
+
+
 /**
  * Returns an array of the elements in the union of x and y, without duplicates
  */
@@ -3353,7 +3352,7 @@ case class ArrayUnion(left: Expression, right: Expression) extends ArraySetLike
     // store elements into resultArray
     var nullElementSize = 0
     var pos = 0
-    Seq(array1, array2).foreach(array => {
+    Seq(array1, array2).foreach { array =>
       var i = 0
       while (i < array.numElements()) {
         val size = if (!isLongType) hsInt.size else hsLong.size
@@ -3380,7 +3379,7 @@ case class ArrayUnion(left: Expression, right: Expression) extends ArraySetLike
         }
         i += 1
       }
-    })
+    }
     pos
   }
 
@@ -3396,7 +3395,7 @@ case class ArrayUnion(left: Expression, right: Expression) extends ArraySetLike
           hsInt = new OpenHashSet[Int]
           val elements = evalIntLongPrimitiveType(array1, array2, null, false)
           hsInt = new OpenHashSet[Int]
-          val resultArray = if (UnsafeArrayData.useGenericArrayData(
+          val resultArray = if (UnsafeArrayData.canUseGenericArrayData(
             IntegerType.defaultSize, elements)) {
             new GenericArrayData(new Array[Any](elements))
           } else {
@@ -3411,7 +3410,7 @@ case class ArrayUnion(left: Expression, right: Expression) extends ArraySetLike
           hsLong = new OpenHashSet[Long]
           val elements = evalIntLongPrimitiveType(array1, array2, null, true)
           hsLong = new OpenHashSet[Long]
-          val resultArray = if (UnsafeArrayData.useGenericArrayData(
+          val resultArray = if (UnsafeArrayData.canUseGenericArrayData(
             LongType.defaultSize, elements)) {
             new GenericArrayData(new Array[Any](elements))
           } else {

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

@@ -1171,7 +1171,7 @@ class CollectionExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper
     val a00 = Literal.create(Seq(1, 2, 3), ArrayType(IntegerType, containsNull = false))
     val a01 = Literal.create(Seq(4, 2), ArrayType(IntegerType, containsNull = false))
     val a02 = Literal.create(Seq(1, 2, null, 4, 5), ArrayType(IntegerType, containsNull = true))
-    val a03 = Literal.create(Seq(-5, 4, -3, 2, -1), ArrayType(IntegerType, containsNull = false))
+    val a03 = Literal.create(Seq(-5, 4, -3, 2, 4), ArrayType(IntegerType, containsNull = false))
     val a04 = Literal.create(Seq.empty[Int], ArrayType(IntegerType, containsNull = false))
     val a05 = Literal.create(Seq[Byte](1, 2, 3), ArrayType(ByteType, containsNull = false))
     val a06 = Literal.create(Seq[Byte](4, 2), ArrayType(ByteType, containsNull = false))
@@ -1191,17 +1191,14 @@ class CollectionExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper
     val a30 = Literal.create(Seq(null, null), ArrayType(IntegerType))
     val a31 = Literal.create(null, ArrayType(StringType))
 
-    checkEvaluation(ArrayUnion(a00, a01), UnsafeArrayData.fromPrimitiveArray(Array(1, 2, 3, 4)))
-    checkEvaluation(ArrayUnion(a02, a03), Seq(1, 2, null, 4, 5, -5, -3, -1))
-    checkEvaluation(ArrayUnion(a03, a02), Seq(-5, 4, -3, 2, -1, 1, null, 5))
+    checkEvaluation(ArrayUnion(a00, a01), Seq(1, 2, 3, 4))
+    checkEvaluation(ArrayUnion(a02, a03), Seq(1, 2, null, 4, 5, -5, -3))
+    checkEvaluation(ArrayUnion(a03, a02), Seq(-5, 4, -3, 2, 1, null, 5))
     checkEvaluation(ArrayUnion(a02, a04), Seq(1, 2, null, 4, 5))
-    checkEvaluation(
-      ArrayUnion(a05, a06), UnsafeArrayData.fromPrimitiveArray(Array[Byte](1, 2, 3, 4)))
-    checkEvaluation(
-      ArrayUnion(a07, a08), UnsafeArrayData.fromPrimitiveArray(Array[Short](1, 2, 3, 4)))
+    checkEvaluation(ArrayUnion(a05, a06), Seq[Byte](1, 2, 3, 4))
+    checkEvaluation(ArrayUnion(a07, a08), Seq[Short](1, 2, 3, 4))
 
-    checkEvaluation(
-      ArrayUnion(a10, a11), UnsafeArrayData.fromPrimitiveArray(Array(1L, 2L, 3L, 4L)))
+    checkEvaluation(ArrayUnion(a10, a11), Seq(1L, 2L, 3L, 4L))
     checkEvaluation(ArrayUnion(a12, a13), Seq(1L, 2L, null, 4L, 5L, -5L, -3L, -1L))
     checkEvaluation(ArrayUnion(a13, a12), Seq(-5L, 4L, -3L, 2L, -1L, 1L, null, 5L))
     checkEvaluation(ArrayUnion(a12, a14), Seq(1L, 2L, null, 4L, 5L))