Looking for suggestions on how to unit test a Spark transformation with ScalaTest. The test class generates a DataFrame from static data and passes it to a transformation, then makes assertion on the passing static data generated in the test class. (The transform creates a second column b defined as col("a").plus(5).)

I got this to work but wonder if there's a better way to create the DataFrame? asJava feels awkward, as does defining the row data and schema separately.

class TransformTest extends FlatSpec with Matchers with SharedSparkContext {
  "Transformer" should "add column to dataframe" in {
    val sqlContext = new SQLContext(sc)
    val rows = Seq[Row](
   val schema = StructType(Seq(StructField("a", IntegerType)))

   val df = sqlContext.createDataFrame(rows, schema)
   val df2 = new Transform().addCol(df)
   assert(df2.count() > 0)
   assert(df2.agg(sum("a")).first.getLong(0) == 6)
   assert(df2.agg(sum("b")).first.getLong(0) == 21)

Indeed, asJava looks awkward and completely unnecessary here. You can approach this problem in a few different ways.

  • pass RDD[Row] instead of Seq[Row] to avoid asJava:

    sqlContext.createDataFrame(sc.parallelize(Seq(Row(1), Row(2), Row(3))), schema
  • use product types instead of Row to avoid both asJava and passing schema:

    sqlContext.createDataFrame(Seq(1, 2,  3).map(Tuple1(_))).toDF("a")


    case class Record(x: Int)
    sqlContext.createDataFrame(Seq(1, 2,  3).map(Record(_)))
  • use implicit conversions to avoid asJava and passing schema and get concise syntax (supported for Seq[Int] only in Spark 2.0+):

    import spark.implicits._
    Seq(1, 2, 3).toDF("a")

In general you should prefer product types (tuples, case classes) over Row objects. You get schema for free as well as some static typing.

  • \$\begingroup\$ Sorry, what do you mean by "product type"? \$\endgroup\$ Aug 18 '16 at 0:56
  • \$\begingroup\$ From theoretical perspective see Wikipedia, from practical subclasses of scala.Product. \$\endgroup\$
    – zero323
    Aug 18 '16 at 15:50

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.