Improved csv and json readers

restonomer-core/src/main/scala/com/clairvoyant/restonomer/core/converter/CSVResponseToDataFrameConverter.scala

@@ -3,14 +3,16 @@ package com.clairvoyant.restonomer.core.converter
 import com.clairvoyant.restonomer.spark.utils.reader.CSVTextToDataFrameReader
 import org.apache.spark.sql.{DataFrame, SparkSession}
 
-class CSVResponseToDataFrameConverter extends ResponseToDataFrameConverter {
+class CSVResponseToDataFrameConverter(
+ containsHeader: Boolean = true
+) extends ResponseToDataFrameConverter {
 
  def convertResponseToDataFrame(
  restonomerResponseBody: Seq[String]
  )(implicit sparkSession: SparkSession): DataFrame =
  new CSVTextToDataFrameReader(
  sparkSession = sparkSession,
- text = restonomerResponseBody.mkString
- ).read
+ containsHeader = containsHeader
+ ).read(text = restonomerResponseBody.flatMap(_.split("\n")))
 
 }

restonomer-core/src/main/scala/com/clairvoyant/restonomer/core/converter/JSONResponseToDataFrameConverter.scala

@@ -11,9 +11,8 @@ class JSONResponseToDataFrameConverter(dataColumnName: Option[String] = None) ex
  )(implicit sparkSession: SparkSession): DataFrame = {
  val responseDF =
  new JSONTextToDataFrameReader(
- sparkSession = sparkSession,
- text = restonomerResponseBody
- ).read
+ sparkSession = sparkSession
+ ).read(restonomerResponseBody)
 
  dataColumnName
  .map { dataColumn =>

restonomer-core/src/test/scala/com/clairvoyant/restonomer/core/transformation/CastColumnsTransformationSpec.scala

@@ -9,8 +9,9 @@ class CastColumnsTransformationSpec extends CoreSpec {
 
  val restonomerResponseDF: DataFrame =
  new JSONTextToDataFrameReader(
- sparkSession = sparkSession,
- text = Seq(
+ sparkSession = sparkSession
+ ).read(text =
+ Seq(
  """
  |{
  | "col_A": 5,
@@ -19,7 +20,7 @@ class CastColumnsTransformationSpec extends CoreSpec {
  |}
  |""".stripMargin
  )
- ).read
+ )
 
  "transform() - with columnDataTypeMapper" should "cast columns as specified in the mapper" in {
  restonomerResponseDF.schema.fields

restonomer-core/src/test/scala/com/clairvoyant/restonomer/core/transformation/CastNestedColumnTransformationSpec.scala

@@ -9,8 +9,9 @@ class CastNestedColumnTransformationSpec extends CoreSpec with DataFrameMatchers
 
  val restonomerResponseDF: DataFrame =
  new JSONTextToDataFrameReader(
- sparkSession = sparkSession,
- text = Seq(
+ sparkSession = sparkSession
+ ).read(text =
+ Seq(
  """
  |{
  | "col_A": "val_A",
@@ -21,7 +22,7 @@ class CastNestedColumnTransformationSpec extends CoreSpec with DataFrameMatchers
  |}
  |""".stripMargin
  )
- ).read
+ )
 
  "transform() - with valid column name and ddl" should "cast the nested column" in {
  restonomerResponseDF.schema.fields

restonomer-core/src/test/scala/com/clairvoyant/restonomer/core/transformation/ConvertColumnCaseTransformationSpec.scala

@@ -11,15 +11,16 @@ class ConvertColumnCaseTransformationSpec extends CoreSpec with DataFrameMatcher
 
  val restonomerResponseDF: DataFrame =
  new JSONTextToDataFrameReader(
- sparkSession = sparkSession,
- text = Seq(
+ sparkSession = sparkSession
+ ).read(text =
+ Seq(
  """
  |{
  | "col_a": "1",
  | "COL_B": "2"
  |}""".stripMargin
  )
- ).read
+ )
 
  "transform() - with valid column name and case type" should "transform the column case" in {
  val restonomerTransformation = ChangeColumnCase(

restonomer-core/src/test/scala/com/clairvoyant/restonomer/core/transformation/ConvertColumnToJsonTransformationSpec.scala

@@ -10,8 +10,9 @@ class ConvertColumnToJsonTransformationSpec extends CoreSpec with DataFrameMatch
 
  val restonomerResponseDF: DataFrame =
  new JSONTextToDataFrameReader(
- sparkSession = sparkSession,
- text = Seq(
+ sparkSession = sparkSession
+ ).read(text =
+ Seq(
  """
  |{
  | "col_A": "1",
@@ -23,7 +24,7 @@ class ConvertColumnToJsonTransformationSpec extends CoreSpec with DataFrameMatch
  | ]
  |}""".stripMargin
  )
- ).read
+ )
 
  "transform() - with valid column name" should "transform the column to json" in {
  val restonomerTransformation = ConvertColumnToJson(

restonomer-core/src/test/scala/com/clairvoyant/restonomer/core/transformation/ExplodeColumnTransformationSpec.scala

@@ -11,16 +11,17 @@ class ExplodeColumnTransformationSpec extends CoreSpec with DataFrameMatchers {
 
  val restonomerResponseDF: DataFrame =
  new JSONTextToDataFrameReader(
- sparkSession = sparkSession,
- text = Seq(
+ sparkSession = sparkSession
+ ).read(text =
+ Seq(
  """
  |{
  | "col_A": "val_A",
  | "col_B": ["val_1", "val_2", "val_3"]
  |}
  |""".stripMargin
  )
- ).read
+ )
 
  "transform() - with valid column name" should "explodeColumn the column into multiple rows" in {
  val restonomerTransformation = ExplodeColumn(

restonomer-core/src/test/scala/com/clairvoyant/restonomer/core/transformation/FlattenSchemaTransformationSpec.scala

@@ -9,7 +9,8 @@ class FlattenSchemaTransformationSpec extends CoreSpec with DataFrameMatchers {
 
  val restonomerResponseDF: DataFrame =
  new JSONTextToDataFrameReader(
- sparkSession = sparkSession,
+ sparkSession = sparkSession
+ ).read(
  text = Seq(
  """
  |{
@@ -20,7 +21,7 @@ class FlattenSchemaTransformationSpec extends CoreSpec with DataFrameMatchers {
  |}
  |""".stripMargin
  )
- ).read
+ )
 
  "transform()" should "flatten the response dataframe" in {
  val restonomerTransformation = FlattenSchema()
@@ -29,16 +30,17 @@ class FlattenSchemaTransformationSpec extends CoreSpec with DataFrameMatchers {
 
  val expectedRestonomerResponseTransformedDF: DataFrame =
  new JSONTextToDataFrameReader(
- sparkSession = sparkSession,
- text = Seq(
+ sparkSession = sparkSession
+ ).read(text =
+ Seq(
  """
  |{
  | "rewardApprovedMonthPeriod_from": "2021-09",
  | "rewardApprovedMonthPeriod_to": "2021-10"
  |}
  |""".stripMargin
  )
- ).read
+ )
 
  actualRestonomerResponseTransformedDF should matchExpectedDataFrame(expectedRestonomerResponseTransformedDF)
  }

restonomer-spark-utils/src/main/scala/com/clairvoyant/restonomer/spark/utils/reader/CSVTextToDataFrameReader.scala

@@ -4,15 +4,14 @@ import org.apache.spark.sql.{DataFrame, SparkSession}
 
 class CSVTextToDataFrameReader(
  override val sparkSession: SparkSession,
- val text: String
+ containsHeader: Boolean = true
 ) extends DataFrameReader {
 
  import sparkSession.implicits._
 
- override def read: DataFrame =
+ override def read(text: Seq[String]): DataFrame =
  sparkSession.read
- .option("header", "true")
- .option("sep", ",")
- .csv(text.split("\\n").toSeq.toDS())
+ .option("header", containsHeader.toString)
+ .csv(text.toDS())
 
 }

restonomer-spark-utils/src/main/scala/com/clairvoyant/restonomer/spark/utils/reader/DataFrameReader.scala

@@ -6,6 +6,6 @@ trait DataFrameReader {
 
  val sparkSession: SparkSession
 
- def read: DataFrame
+ def read(text: Seq[String]): DataFrame
 
 }

restonomer-spark-utils/src/main/scala/com/clairvoyant/restonomer/spark/utils/reader/JSONTextToDataFrameReader.scala

@@ -3,12 +3,11 @@ package com.clairvoyant.restonomer.spark.utils.reader
 import org.apache.spark.sql.{DataFrame, SparkSession}
 
 class JSONTextToDataFrameReader(
- override val sparkSession: SparkSession,
- val text: Seq[String]
+ override val sparkSession: SparkSession
 ) extends DataFrameReader {
 
  import sparkSession.implicits._
 
- override def read: DataFrame = sparkSession.read.json(text.toDS())
+ override def read(text: Seq[String]): DataFrame = sparkSession.read.json(text.toDS())
 
 }