Converting a dataframe to dataset retains extra columns - scala

Question

In Spark 2.11, when converting a Dataframe to a Dataset, spark retains the extra columns that are not even referred in the class of the dataset.
scala> case class F(x: String, y: String)
defined class F
scala> import spark.implicits._
import spark.implicits._
scala> val df = Seq(("1a","2a","3a","4a"), ("5a", "6a", "7a","8a")).toDF("w","x","y","z")
df: org.apache.spark.sql.DataFrame = [w: string, x: string ... 2 more fields]
scala> val ds = df.to
toDF toJSON toJavaRDD toLocalIterator toString
scala> val ds = df.as
as asInstanceOf
scala> val ds = df.as[F]
ds: org.apache.spark.sql.Dataset[F] = [w: string, x: string ... 2 more fields]
scala> ds.show()
+---+---+---+---+
| w| x| y| z|
+---+---+---+---+
| 1a| 2a| 3a| 4a|
| 5a| 6a| 7a| 8a|
+---+---+---+---+
This throws the type safety aspect of it down the drain. Is there a way to prevent this from happening ?
I read that using ds.map(identity) because the conversion is supposed to be lazy but it does not trigger when trying to write the dataframe. Plus it's impossible to write a test to cover .map(identity) and prevent deletion.

Answer 1

You can address this behavior in several ways but one of the easiest and most performant option in scala in to define your own to[T] method that selects only the appropriate columns:
import scala.reflect.runtime.universe.TypeTag
import org.apache.spark.sql._
object SparkExtensions {
implicit class ExtendedDataFrame(df: DataFrame) {
def to[T <: Product: TypeTag]: Dataset[T] = {
import df.sparkSession.implicits._
import org.apache.spark.sql.functions.col
df.select(Encoders.product[T].schema.map(f => col(f.name)): _*).as[T]
}
}
}
Which you can than use in the following way:
scala> case class F(x: String, y: String)
defined class F
scala> import SparkExtensions._
import SparkExtensions._
scala> val df = Seq(("1a","2a","3a","4a"), ("5a", "6a", "7a","8a")).toDF("w","x","y","z")
df: org.apache.spark.sql.DataFrame = [w: string, x: string ... 2 more fields]
scala> val ds: Dataset[F] = df.to[F]
ds: org.apache.spark.sql.Dataset[F] = [w: string, x: string]
scala> ds.show()
+---+---+
| w| x|
+---+---+
| 1a| 2a|
| 5a| 6a|
+---+---+
If you don't want to define an implicit class, you can always simply do df.as[F].rdd.toDS to achieve the same result in a slightly less runtime efficient way.