Pull null handling into TypeCast, add tests

Author: dtpeacock

src/main/scala/com/databricks/spark/csv/CsvRelation.scala

@@ -108,14 +108,8 @@ case class CsvRelation protected[spark] (
         try {
           index = 0
           while (index < schemaFields.length) {
-            rowArray(index) = if (schemaFields(index).nullable && tokens(index) == ""){
-              schemaFields(index).dataType match {
-                case StringType => ""
-                case _ => null
-              }
-            } else {
-              TypeCast.castTo(tokens(index), schemaFields(index).dataType)
-            }
+            val field = schemaFields(index)
+            rowArray(index) = TypeCast.castTo(tokens(index), field.dataType, field.nullable)
             index = index + 1
           }
           Some(Row.fromSeq(rowArray))

src/main/scala/com/databricks/spark/csv/util/TypeCast.scala

@@ -29,25 +29,33 @@ object TypeCast {
    * Casts given string datum to specified type.
    * Currently we do not support complex types (ArrayType, MapType, StructType).
    *
+   * For string types, this is simply the datum. For other types.
+   * For other nullable types, this is null if the string datum is empty.
+   *
    * @param datum string value
    * @param castType SparkSQL type
    */
-  private[csv] def castTo(datum: String, castType: DataType): Any = {
-    castType match {
-      case _: ByteType => datum.toByte
-      case _: ShortType => datum.toShort
-      case _: IntegerType => datum.toInt
-      case _: LongType => datum.toLong
-      case _: FloatType => datum.toFloat
-      case _: DoubleType => datum.toDouble
-      case _: BooleanType => datum.toBoolean
-      case _: DecimalType => new BigDecimal(datum.replaceAll(",", ""))
-      // TODO(hossein): would be good to support other common timestamp formats
-      case _: TimestampType => Timestamp.valueOf(datum)
-      // TODO(hossein): would be good to support other common date formats
-      case _: DateType => Date.valueOf(datum)
-      case _: StringType => datum
-      case _ => throw new RuntimeException(s"Unsupported type: ${castType.typeName}")
+  private[csv] def castTo(datum: String, castType: DataType, nullable: Boolean = true): Any = {
+    if (castType.isInstanceOf[StringType]){
+      datum
+    } else if (nullable && datum == ""){
+      null
+    } else {
+      castType match {
+        case _: ByteType => datum.toByte
+        case _: ShortType => datum.toShort
+        case _: IntegerType => datum.toInt
+        case _: LongType => datum.toLong
+        case _: FloatType => datum.toFloat
+        case _: DoubleType => datum.toDouble
+        case _: BooleanType => datum.toBoolean
+        case _: DecimalType => new BigDecimal(datum.replaceAll(",", ""))
+        // TODO(hossein): would be good to support other common timestamp formats
+        case _: TimestampType => Timestamp.valueOf(datum)
+        // TODO(hossein): would be good to support other common date formats
+        case _: DateType => Date.valueOf(datum)
+        case _ => throw new RuntimeException(s"Unsupported type: ${castType.typeName}")
+      }
     }
   }
 

src/test/scala/com/databricks/spark/csv/util/TypeCastSuite.scala

@@ -19,7 +19,7 @@ import java.math.BigDecimal
 
 import org.scalatest.FunSuite
 
-import org.apache.spark.sql.types.DecimalType
+import org.apache.spark.sql.types.{StringType, IntegerType, DecimalType}
 
 class TypeCastSuite extends FunSuite {
 
@@ -56,4 +56,20 @@ class TypeCastSuite extends FunSuite {
     }
     assert(exception.getMessage.contains("Unsupported special character for delimiter"))
   }
+
+  test("Nullable types are handled"){
+    assert(TypeCast.castTo("", IntegerType, nullable = true) == null)
+  }
+
+  test("String type should always return the same as the input"){
+    assert(TypeCast.castTo("", StringType, nullable = true) == "")
+    assert(TypeCast.castTo("", StringType, nullable = false) == "")
+  }
+
+  test("Throws exception for empty string with non null type"){
+    val exception = intercept[NumberFormatException]{
+      TypeCast.castTo("", IntegerType, nullable = false)
+    }
+    assert(exception.getMessage.contains("For input string: \"\""))
+  }
 }