Spark Schema DSL

SparkSchemaDsl.scala

import org.apache.spark.sql.types._
import org.apache.spark.sql._

object SchemaDsl {
  case class ScalaToSparkType[ScalaType](sparkType: DataType, isNullable: Boolean = false) {
    def toField(name: String) = StructField(name     = name, dataType = sparkType, nullable = isNullable)
  }

  implicit val stringType: ScalaToSparkType[String] = ScalaToSparkType(StringType)
  implicit val intType: ScalaToSparkType[Int] = ScalaToSparkType(IntegerType)
  implicit val longType: ScalaToSparkType[Long] = ScalaToSparkType(LongType)

  implicit def optionType[A](implicit ev: ScalaToSparkType[A]): ScalaToSparkType[Option[A]] =
    ev.copy(isNullable = false)

  implicit class ColumnDsl(s: String) {
    def ofType[A](implicit ev: ScalaToSparkType[A]): StructField = ev.toField(s)
  }

  implicit class DataFrameReaderExtension(reader: DataFrameReader) {
    def withSchema(fields: StructField*) = reader.schema(StructType(fields))
  }
}

object SchemaDslTest {
  import SchemaDsl._

  def read(path: String)(implicit spark: SparkSession) = {
    spark.read
      .format("csv")
      .withSchema(
        "id".ofType[Long],
        "name".ofType[Option[String]]
      )
  }
}