Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
viirya wants to merge 7 commits into from
Closed

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

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
db1b91e
Add support for Scala Set collection types in serialization.
viirya Jun 25, 2017
53b1dc8
Add more comments.
viirya Jun 26, 2017
31b7812
Add test.
viirya Jun 28, 2017
61f0bb6
Skip null check if the array doesn't contain null.
viirya Jun 28, 2017
4602689
Reuse MapObject for serialize Set.
viirya Jul 5, 2017
56c1298
Address comments.
viirya Jul 6, 2017
e2464c2
Use member instead of declaration to query "newBuilder".
viirya Jul 6, 2017
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 39 additions & 0 deletions sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -342,6 +342,28 @@ object ScalaReflection extends ScalaReflection {
mirror.runtimeClass(t.typeSymbol.asClass)
)

// We serialize a `Set` to Catalyst array. When we deserialize a Catalyst array
// to a `Set`, if there are duplicated elements, the elements will be de-duplicated.
case t if t <:< localTypeOf[scala.collection.Set[_]] =>

@cloud-fan cloud-fan Jul 5, 2017

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we merge it with the seq case? e.g. case t if t <:< localTypeOf[Seq[_]] || t <:< localTypeOf[scala.collection.Set[_]]

@cloud-fan cloud-fan Jul 5, 2017

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we only need to update one place:

val companion = t.normalize.typeSymbol.companionSymbol.typeSignature
val cls = companion.declaration(newTermName("newBuilder")) match {
  case NoSymbol if t <:< localTypeOf[Seq[_]] => classOf[Seq[_]]
  case NoSymbol if t <:< localTypeOf[scala.collection.Set[_]] => classOf[Set[_]]
  case _ => mirror.runtimeClass(t.typeSymbol.asClass)
}

@viirya viirya Jul 6, 2017

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This trick companion.declaration... doesn't work for Set. You always get NoSymbol.

@viirya viirya Jul 6, 2017

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Anyway I will merge them.

val TypeRef(_, _, Seq(elementType)) = t
val Schema(dataType, elementNullable) = schemaFor(elementType)
val className = getClassNameFromType(elementType)
val newTypePath = s"""- set element class: "$className"""" +: walkedTypePath

val mapFunction: Expression => Expression = element => {
// upcast the array element to the data type the encoder expected.
val casted = upCastToExpectedType(element, dataType, newTypePath)
val converter = deserializerFor(elementType, Some(casted), newTypePath)
if (elementNullable) {
converter
} else {
AssertNotNull(converter, newTypePath)
}
}

val cls = mirror.runtimeClass(t.typeSymbol.asClass)
UnresolvedMapObjects(mapFunction, getPath, Some(cls))

case t if t.typeSymbol.annotations.exists(_.tpe =:= typeOf[SQLUserDefinedType]) =>
val udt = getClassFromType(t).getAnnotation(classOf[SQLUserDefinedType]).udt().newInstance()
val obj = NewInstance(
Expand Down Expand Up @@ -498,6 +520,19 @@ object ScalaReflection extends ScalaReflection {
serializerFor(_, valueType, valuePath, seenTypeSet),
valueNullable = !valueType.typeSymbol.asClass.isPrimitive)

case t if t <:< localTypeOf[scala.collection.Set[_]] =>
val TypeRef(_, _, Seq(elementType)) = t

// There's no corresponding Catalyst type for `Set`, we serialize a `Set` to Catalyst array.
// Note that the property of `Set` is only kept when manipulating the data as domain object.
val newInput =
Invoke(
inputObject,
"toSeq",
ObjectType(classOf[Seq[_]]))

@viirya viirya Jul 5, 2017

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For primitive, calling toArray can directly construct UnsafeArrayData in toCatalystArray. However, toArray requires a ClassTag. To generate it with a StaticInvoke might be hacky. So for now I simply use toSeq.


toCatalystArray(newInput, elementType)

case t if t <:< localTypeOf[String] =>
StaticInvoke(
classOf[UTF8String],
Expand Down Expand Up @@ -702,6 +737,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)
Expand Down
5 changes: 3 additions & 2 deletions sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects/objects.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -598,8 +598,9 @@ case class MapObjects private(

val (initCollection, addElement, getResult): (String, String => String, String) =
customCollectionCls match {
case Some(cls) if classOf[Seq[_]].isAssignableFrom(cls) =>
// Scala sequence
case Some(cls) if classOf[Seq[_]].isAssignableFrom(cls) ||
classOf[scala.collection.Set[_]].isAssignableFrom(cls) =>
// Scala sequence or set
val getBuilder = s"${cls.getName}$$.MODULE$$.newBuilder()"
val builder = ctx.freshName("collectionBuilder")
(
Expand Down
10 changes: 10 additions & 0 deletions sql/core/src/main/scala/org/apache/spark/sql/SQLImplicits.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -171,6 +171,16 @@ abstract class SQLImplicits extends LowPrioritySQLImplicits {
/** @since 2.3.0 */
implicit def newMapEncoder[T <: Map[_, _] : TypeTag]: Encoder[T] = ExpressionEncoder()

/**
* Notice that we serialize `Set` to Catalyst array. The set property is only kept when
* manipulating the domain objects. The serialization format doesn't keep the set property.
* When we have a Catalyst array which contains duplicated elements and convert it to
* `Dataset[Set[T]]` by using the encoder, the elements will be de-duplicated.
*
* @since 2.3.0
*/
implicit def newSetEncoder[T <: Set[_] : TypeTag]: Encoder[T] = ExpressionEncoder()

// Arrays

/** @since 1.6.1 */
Expand Down
9 changes: 9 additions & 0 deletions sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))

@cloud-fan cloud-fan Jul 5, 2017

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we add one more case?

      checkAnswer(
        df.select(collect_set($"a"), collect_set($"b")).as[(Set[Int], Set[Int])],
        Seq(Set(1, 2, 3) -> Set(2, 4))
      )

@viirya viirya Jul 6, 2017

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure.

}

test("collect functions structs") {
Expand Down
31 changes: 31 additions & 0 deletions sql/core/src/test/scala/org/apache/spark/sql/DatasetPrimitiveSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -339,6 +340,31 @@ class DatasetPrimitiveSuite extends QueryTest with SharedSQLContext {
LHMapClass(LHMap(1 -> 2)) -> LHMap("test" -> MapClass(Map(3 -> 4))))
}

test("arbitrary sets") {

@maropu maropu Jun 27, 2017

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Better to test null cases?

Seq(Seq(Some(1), None), Seq(Some(2))).toDF("c").as[Set[Int]]

@viirya viirya Jun 29, 2017

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added a test for it.

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)))

checkDataset(Seq(Seq(Some(1), None), Seq(Some(2))).toDF("c").as[Set[Integer]],
Seq(Set[Integer](1, null), Set[Integer](2)): _*)
}

test("nested sequences") {
checkDataset(Seq(Seq(Seq(1))).toDS(), Seq(Seq(1)))
checkDataset(Seq(List(Queue(1))).toDS(), List(Queue(1)))
Expand All @@ -349,6 +375,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))
Expand Down