Add checks for format option

tgravescs · Jun 3, 2021 · 6f5271c · 6f5271c
1 parent ef5cf4c
commit 6f5271c
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 19 deletions.
diff --git a/...spark-tools/src/main/scala/com/nvidia/spark/rapids/tool/qualification/Qualification.scala b/...spark-tools/src/main/scala/com/nvidia/spark/rapids/tool/qualification/Qualification.scala
@@ -105,25 +105,24 @@ object Qualification extends Logging {
 
   def writeQualification(df: DataFrame,
       outputDir: String, format: String): Unit = {
-    // val fileWriter = apps.head.fileWriter
-    // val dfRenamed = apps.head.renameQualificationColumns(df)
-    if (format.equals("csv")) {
-      df.repartition(1).write.option("header", "true").
-        mode("overwrite").csv(s"$outputDir/rapids_4_spark_qualification_output")
-      logInfo(s"Output log location:  $outputDir")
-    } else {
-      // This tool's output log file name
-      val logFileName = "rapids_4_spark_qualification_output.log"
-      val outputFilePath = new Path(s"$outputDir/$logFileName")
-      val fs = FileSystem.get(outputFilePath.toUri, new Configuration())
-      val outFile = fs.create(outputFilePath)
-      outFile.writeUTF(ToolUtils.showString(df, 1000))
-      outFile.flush()
-      outFile.close()
-      logInfo(s"Output log location: $outputFilePath")
+    format match {
+      case "csv" =>
+        df.repartition(1).write.option("header", "true").
+          mode("overwrite").csv(s"$outputDir/rapids_4_spark_qualification_output")
+        logInfo(s"Output log location:  $outputDir")
+      case "text" =>
+        // This tool's output log file name
+        val logFileName = "rapids_4_spark_qualification_output.log"
+        val outputFilePath = new Path(s"$outputDir/$logFileName")
+        val fs = FileSystem.get(outputFilePath.toUri, new Configuration())
+        val outFile = fs.create(outputFilePath)
+        // outFile.writeUTF(ToolUtils.showString(df, 1000))
+        df.repartition(1).write.option("header", "true").
+          mode("overwrite").text(s"$outputDir/rapids_4_spark_qualification_output_text")
+        outFile.flush()
+        outFile.close()
+        logInfo(s"Output log location: $outputFilePath")
+      case _ => logError("Invalid format")
     }
-
-    // fileWriter.write("\n" + ToolUtils.showString(dfRenamed,
-    //   apps(0).numOutputRows))
   }
 }
diff --git a/...k-tools/src/main/scala/com/nvidia/spark/rapids/tool/qualification/QualificationArgs.scala b/...k-tools/src/main/scala/com/nvidia/spark/rapids/tool/qualification/QualificationArgs.scala
@@ -68,5 +68,10 @@ For usage see below:
       required = false,
       default = Some(false),
       descr = "Include the executor CPU time percent. It will take longer with this option.")
+
+  validateOpt(outputFormat) {
+    case Some("text") | Some("csv") => Right(Unit)
+    case _ => Left(s"Invalid format - must be 'csv' or 'text'.")
+  }
   verify()
 }