[SPARK-21204][SQL] Add support for Scala Set collection types in serialization

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala

@@ -342,6 +342,15 @@ object ScalaReflection extends ScalaReflection {
           mirror.runtimeClass(t.typeSymbol.asClass)
         )
 
+      case t if t <:< localTypeOf[Set[_]] =>
+        val TypeRef(_, _, Seq(elementType)) = t
+
+        CollectObjectsToSet(
+          p => deserializerFor(elementType, Some(p), walkedTypePath),
+          getPath,
+          mirror.runtimeClass(t.typeSymbol.asClass)
+        )
+
       case t if t.typeSymbol.annotations.exists(_.tpe =:= typeOf[SQLUserDefinedType]) =>
         val udt = getClassFromType(t).getAnnotation(classOf[SQLUserDefinedType]).udt().newInstance()
         val obj = NewInstance(
@@ -498,6 +507,17 @@ object ScalaReflection extends ScalaReflection {
           serializerFor(_, valueType, valuePath, seenTypeSet),
           valueNullable = !valueType.typeSymbol.asClass.isPrimitive)
 
+      case t if t <:< localTypeOf[Set[_]] =>
+        val TypeRef(_, _, Seq(elementType)) = t
+        val elementClsName = getClassNameFromType(elementType)
+        val elementPath = s"""- set element class: "$elementClsName"""" +: walkedTypePath
+
+        ExternalSetToCatalystArray(
+          inputObject,
+          dataTypeFor(elementType),
+          serializerFor(_, elementType, elementPath, seenTypeSet),
+          elementNullable = !elementType.typeSymbol.asClass.isPrimitive)
+
       case t if t <:< localTypeOf[String] =>
         StaticInvoke(
           classOf[UTF8String],
@@ -702,6 +722,10 @@ object ScalaReflection extends ScalaReflection {
         val Schema(valueDataType, valueNullable) = schemaFor(valueType)
         Schema(MapType(schemaFor(keyType).dataType,
           valueDataType, valueContainsNull = valueNullable), nullable = true)
+      case t if t <:< localTypeOf[Set[_]] =>
+        val TypeRef(_, _, Seq(elementType)) = t
+        val Schema(dataType, nullable) = schemaFor(elementType)
+        Schema(ArrayType(dataType, containsNull = nullable), nullable = true)
       case t if t <:< localTypeOf[String] => Schema(StringType, nullable = true)
       case t if t <:< localTypeOf[java.sql.Timestamp] => Schema(TimestampType, nullable = true)
       case t if t <:< localTypeOf[java.sql.Date] => Schema(DateType, nullable = true)

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

@@ -834,6 +834,137 @@ case class CollectObjectsToMap private(
   }
 }
 
+object CollectObjectsToSet {
+  private val curId = new java.util.concurrent.atomic.AtomicInteger()
+
+  /**
+   * Construct an instance of CollectObjectsToSet case class.
+   *
+   * @param function The function applied on the collection elements.
+   * @param inputData An expression that when evaluated returns a collection object.
+   * @param collClass The type of the resulting collection.
+   */
+  def apply(
+      function: Expression => Expression,
+      inputData: Expression,
+      collClass: Class[_]): CollectObjectsToSet = {
+    val id = curId.getAndIncrement()
+    val loopValue = s"CollectObjectsToSet_loopValue$id"
+    val loopIsNull = s"CollectObjectsToSet_loopIsNull$id"
+    val arrayType = inputData.dataType.asInstanceOf[ArrayType]
+    val loopVar = LambdaVariable(loopValue, loopIsNull, arrayType.elementType)
+    CollectObjectsToSet(
+      loopValue, loopIsNull, function(loopVar), inputData, collClass)
+  }
+}
+
+/**
+ * Expression used to convert a Catalyst Array to an external Scala Set.
+ * The collection is constructed using the associated builder, obtained by calling `newBuilder`
+ * on the collection's companion object.
+ *
+ * @param loopValue the name of the loop variable that is used when iterating over the value
+ *                       collection, and which is used as input for the `lambdaFunction`
+ * @param loopIsNull the nullability of the loop variable that is used when iterating over
+ *                        the value collection, and which is used as input for the
+ *                        `lambdaFunction`
+ * @param lmbdaFunction A function that takes the `loopValue` as input, and is used as
+ *                            a lambda function to handle collection elements.
+ * @param inputData An expression that when evaluated returns an array object.
+ * @param collClass The type of the resulting collection.
+ */
+case class CollectObjectsToSet private(
+    loopValue: String,
+    loopIsNull: String,
+    lambdaFunction: Expression,
+    inputData: Expression,
+    collClass: Class[_]) extends Expression with NonSQLExpression {
+
+  override def nullable: Boolean = inputData.nullable
+
+  override def children: Seq[Expression] = lambdaFunction :: inputData :: Nil
+
+  override def eval(input: InternalRow): Any =
+    throw new UnsupportedOperationException("Only code-generated evaluation is supported")
+
+  override def dataType: DataType = ObjectType(collClass)
+
+  override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+    // The data with PythonUserDefinedType are actually stored with the data type of its sqlType.
+    def inputDataType(dataType: DataType) = dataType match {
+      case p: PythonUserDefinedType => p.sqlType
+      case _ => dataType
+    }
+
+    val arrayType = inputDataType(inputData.dataType).asInstanceOf[ArrayType]
+    val loopValueJavaType = ctx.javaType(arrayType.elementType)
+    ctx.addMutableState("boolean", loopIsNull, "")
+    ctx.addMutableState(loopValueJavaType, loopValue, "")
+    val genFunction = lambdaFunction.genCode(ctx)
+
+    val genInputData = inputData.genCode(ctx)
+    val dataLength = ctx.freshName("dataLength")
+    val loopIndex = ctx.freshName("loopIndex")
+    val builderValue = ctx.freshName("builderValue")
+
+    val getLength = s"${genInputData.value}.numElements()"
+    val getLoopVar = ctx.getValue(genInputData.value, arrayType.elementType, loopIndex)
+
+    // Make a copy of the data if it's unsafe-backed
+    def makeCopyIfInstanceOf(clazz: Class[_ <: Any], value: String) =
+      s"$value instanceof ${clazz.getSimpleName}? $value.copy() : $value"
+    val genFunctionValue =
+      lambdaFunction.dataType match {
+        case StructType(_) => makeCopyIfInstanceOf(classOf[UnsafeRow], genFunction.value)
+        case ArrayType(_, _) => makeCopyIfInstanceOf(classOf[UnsafeArrayData], genFunction.value)
+        case MapType(_, _, _) => makeCopyIfInstanceOf(classOf[UnsafeMapData], genFunction.value)
+        case _ => genFunction.value
+      }
+
+    val loopNullCheck = s"$loopIsNull = ${genInputData.value}.isNullAt($loopIndex);"
+
+    val builderClass = classOf[Builder[_, _]].getName
+    val constructBuilder = s"""
+      $builderClass $builderValue = ${collClass.getName}$$.MODULE$$.newBuilder();
+      $builderValue.sizeHint($dataLength);
+    """
+
+    val appendToBuilder = s"""
+      if (${genFunction.isNull}) {
+        $builderValue.$$plus$$eq(null);
+      } else {
+        $builderValue.$$plus$$eq($genFunctionValue);
+      }
+     """
+    val getBuilderResult = s"${ev.value} = (${collClass.getName}) $builderValue.result();"
+
+    val code = s"""
+      ${genInputData.code}
+      ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
+
+      if (!${genInputData.isNull}) {
+        int $dataLength = $getLength;
+        $constructBuilder
+
+        int $loopIndex = 0;
+        while ($loopIndex < $dataLength) {
+          $loopValue = ($loopValueJavaType) ($getLoopVar);
+
+          $loopNullCheck
+
+          ${genFunction.code}
+          $appendToBuilder
+
+          $loopIndex += 1;
+        }
+
+        $getBuilderResult
+      }
+    """
+    ev.copy(code = code, isNull = genInputData.isNull)
+  }
+}
+
 object ExternalMapToCatalyst {
   private val curId = new java.util.concurrent.atomic.AtomicInteger()
 
@@ -992,6 +1123,128 @@ case class ExternalMapToCatalyst private(
   }
 }
 
+object ExternalSetToCatalystArray {
+  private val curId = new java.util.concurrent.atomic.AtomicInteger()
+
+  def apply(
+      inputSet: Expression,
+      elementType: DataType,
+      elementConverter: Expression => Expression,
+      elementNullable: Boolean): ExternalSetToCatalystArray = {
+    val id = curId.getAndIncrement()
+    val elementName = "ExternalSetToCatalystArray_element" + id
+    val elementIsNull = "ExternalSetToCatalystArray_element_isNull" + id
+
+    ExternalSetToCatalystArray(
+      elementName,
+      elementIsNull,
+      elementType,
+      elementConverter(LambdaVariable(elementName, elementIsNull, elementType, elementNullable)),
+      inputSet
+    )
+  }
+}
+
+/**
+ * Converts a Scala/Java set object into catalyst array format, by applying the converter when
+ * iterate the set.
+ *
+ * @param element the name of the set element variable that used when iterate the set, and used as
+ *                input for the `elementConverter`
+ * @param elementIsNull the nullability of the element variable that used when iterate the set, and
+ *                    used as input for the `elementConverter`
+ * @param elementType the data type of the element variable that used when iterate the set, and
+ *                  used as input for the `elementConverter`
+ * @param elementConverter A function that take the `element` as input, and converts it to catalyst
+ *                       array format.
+ * @param child An expression that when evaluated returns the input set object.
+ */
+case class ExternalSetToCatalystArray private(
+    element: String,
+    elementIsNull: String,
+    elementType: DataType,
+    elementConverter: Expression,
+    child: Expression)
+  extends UnaryExpression with NonSQLExpression {
+
+  override def foldable: Boolean = false
+
+  override def dataType: ArrayType = ArrayType(
+    elementType = elementConverter.dataType, containsNull = elementConverter.nullable)
+
+  override def eval(input: InternalRow): Any =
+    throw new UnsupportedOperationException("Only code-generated evaluation is supported")
+
+  override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+    val inputSet = child.genCode(ctx)
+    val genElementConverter = elementConverter.genCode(ctx)
+    val length = ctx.freshName("length")
+    val index = ctx.freshName("index")
+
+    val iter = ctx.freshName("iter")
+    val (defineIterator, defineElement) = child.dataType match {
+      case ObjectType(cls) if classOf[java.util.Set[_]].isAssignableFrom(cls) =>
+        val javaIteratorCls = classOf[java.util.Iterator[_]].getName
+        val defineIter = s"final $javaIteratorCls $iter = ${inputSet.value}.iterator();"
+
+        val defineElement =
+          s"""
+            ${ctx.javaType(elementType)} $element = (${ctx.boxedType(elementType)}) $iter.next();
+          """
+
+        defineIter -> defineElement
+
+      case ObjectType(cls) if classOf[scala.collection.Set[_]].isAssignableFrom(cls) =>
+        val scalaIteratorCls = classOf[Iterator[_]].getName
+        val defineIter = s"final $scalaIteratorCls $iter = ${inputSet.value}.iterator();"
+
+        val defineElement =
+          s"""
+            ${ctx.javaType(elementType)} $element = (${ctx.boxedType(elementType)}) $iter.next();
+          """
+
+        defineIter -> defineElement
+    }
+
+    val elementNullCheck = if (ctx.isPrimitiveType(elementType)) {
+      s"boolean $elementIsNull = false;"
+    } else {
+      s"boolean $elementIsNull = $element == null;"
+    }
+
+    val arrayCls = classOf[GenericArrayData].getName
+    val convertedElements = ctx.freshName("convertedElements")
+    val convertedElementType = ctx.boxedType(elementConverter.dataType)
+    val code =
+      s"""
+        ${inputSet.code}
+        ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
+        if (!${inputSet.isNull}) {
+          final int $length = ${inputSet.value}.size();
+          final Object[] $convertedElements = new Object[$length];
+          int $index = 0;
+          $defineIterator
+          while($iter.hasNext()) {
+            $defineElement
+            $elementNullCheck
+
+            ${genElementConverter.code}
+            if (${genElementConverter.isNull}) {
+              $convertedElements[$index] = null;
+            } else {
+              $convertedElements[$index] = ($convertedElementType) ${genElementConverter.value};
+            }
+
+            $index++;
+          }
+
+          ${ev.value} = new $arrayCls($convertedElements);
+        }
+      """
+    ev.copy(code = code, isNull = inputSet.isNull)
+  }
+}
+
 /**
  * Constructs a new external row, using the result of evaluating the specified expressions
  * as content.

sql/core/src/main/scala/org/apache/spark/sql/SQLImplicits.scala

@@ -171,6 +171,10 @@ abstract class SQLImplicits extends LowPrioritySQLImplicits {
   /** @since 2.3.0 */
   implicit def newMapEncoder[T <: Map[_, _] : TypeTag]: Encoder[T] = ExpressionEncoder()
 
+  // Sets
+  /** @since 2.3.0 */
+  implicit def newSetEncoder[T <: Set[_] : TypeTag]: Encoder[T] = ExpressionEncoder()
+
   // Arrays
 
   /** @since 1.6.1 */

sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala

@@ -460,6 +460,15 @@ class DataFrameAggregateSuite extends QueryTest with SharedSQLContext {
       df.select(collect_set($"a"), collect_set($"b")),
       Seq(Row(Seq(1, 2, 3), Seq(2, 4)))
     )
+
+    checkDataset(
+      df.select(collect_set($"a").as("aSet"))
+        .as[Set[Int]],
+      Set(1, 2, 3))
+    checkDataset(
+      df.select(collect_set($"b").as("bSet"))
+        .as[Set[Int]],
+      Set(2, 4))
   }
 
   test("collect functions structs") {

sql/core/src/test/scala/org/apache/spark/sql/DatasetPrimitiveSuite.scala

@@ -17,6 +17,7 @@
 
 package org.apache.spark.sql
 
+import scala.collection.immutable.{HashSet => HSet}
 import scala.collection.immutable.Queue
 import scala.collection.mutable.{LinkedHashMap => LHMap}
 import scala.collection.mutable.ArrayBuffer
@@ -339,6 +340,28 @@ class DatasetPrimitiveSuite extends QueryTest with SharedSQLContext {
       LHMapClass(LHMap(1 -> 2)) -> LHMap("test" -> MapClass(Map(3 -> 4))))
   }
 
+  test("arbitrary sets") {
+    checkDataset(Seq(Set(1, 2, 3, 4)).toDS(), Set(1, 2, 3, 4))
+    checkDataset(Seq(Set(1.toLong, 2.toLong)).toDS(), Set(1.toLong, 2.toLong))
+    checkDataset(Seq(Set(1.toDouble, 2.toDouble)).toDS(), Set(1.toDouble, 2.toDouble))
+    checkDataset(Seq(Set(1.toFloat, 2.toFloat)).toDS(), Set(1.toFloat, 2.toFloat))
+    checkDataset(Seq(Set(1.toByte, 2.toByte)).toDS(), Set(1.toByte, 2.toByte))
+    checkDataset(Seq(Set(1.toShort, 2.toShort)).toDS(), Set(1.toShort, 2.toShort))
+    checkDataset(Seq(Set(true, false)).toDS(), Set(true, false))
+    checkDataset(Seq(Set("test1", "test2")).toDS(), Set("test1", "test2"))
+    checkDataset(Seq(Set(Tuple1(1), Tuple1(2))).toDS(), Set(Tuple1(1), Tuple1(2)))
+
+    checkDataset(Seq(HSet(1, 2)).toDS(), HSet(1, 2))
+    checkDataset(Seq(HSet(1.toLong, 2.toLong)).toDS(), HSet(1.toLong, 2.toLong))
+    checkDataset(Seq(HSet(1.toDouble, 2.toDouble)).toDS(), HSet(1.toDouble, 2.toDouble))
+    checkDataset(Seq(HSet(1.toFloat, 2.toFloat)).toDS(), HSet(1.toFloat, 2.toFloat))
+    checkDataset(Seq(HSet(1.toByte, 2.toByte)).toDS(), HSet(1.toByte, 2.toByte))
+    checkDataset(Seq(HSet(1.toShort, 2.toShort)).toDS(), HSet(1.toShort, 2.toShort))
+    checkDataset(Seq(HSet(true, false)).toDS(), HSet(true, false))
+    checkDataset(Seq(HSet("test1", "test2")).toDS(), HSet("test1", "test2"))
+    checkDataset(Seq(HSet(Tuple1(1), Tuple1(2))).toDS(), HSet(Tuple1(1), Tuple1(2)))
+  }
+
   test("nested sequences") {
     checkDataset(Seq(Seq(Seq(1))).toDS(), Seq(Seq(1)))
     checkDataset(Seq(List(Queue(1))).toDS(), List(Queue(1)))
@@ -349,6 +372,11 @@ class DatasetPrimitiveSuite extends QueryTest with SharedSQLContext {
     checkDataset(Seq(LHMap(Map(1 -> 2) -> 3)).toDS(), LHMap(Map(1 -> 2) -> 3))
   }
 
+  test("nested set") {
+    checkDataset(Seq(Set(HSet(1, 2), HSet(3, 4))).toDS(), Set(HSet(1, 2), HSet(3, 4)))
+    checkDataset(Seq(HSet(Set(1, 2), Set(3, 4))).toDS(), HSet(Set(1, 2), Set(3, 4)))
+  }
+
   test("package objects") {
     import packageobject._
     checkDataset(Seq(PackageClass(1)).toDS(), PackageClass(1))