Read, transform and write data within each partition in the DataFrame - scala

Question

Language - Scala
Spark version - 2.4
I am new to both Scala and Spark. (I am from python background, so the whole JVM ecosystem is quite new to me)
I want to write a spark program to parallelize following steps:
Read data from S3 in dataframe
Transform each row of this dataframe
Write the updated dataframe back to S3 at a new location
Let's say I have 3 items, A, B and C. For each of these items, I want to do the above 3 steps.
I want to do this in parallel for all these 3 items.
I tried creating an RDD with 3 partition, where each partition has one item, A, B and C, respectively.
Then I tried to use mapPartition method to write my logic for each partition (the 3 steps mentioned above).
I am getting Task not serializable errors. Although I understand the meaning of this error, I don't know how to solve it.
val items = Array[String]("A", "B", "C")
val rdd = sc.parallelize(items, 3)
rdd.mapPartitions(
partition => {
val item = partition.next()
val filePath = new ListBuffer[String]()
filePath += s"$basePath/item=$item/*"
val df = spark.read.format("parquet").option("basePath",s"$basePath").schema(schema).load(filePaths: _*)
//Transform this dataframe
val newDF = df.rdd.mapPartitions(partition => {partition.map(row =>{methodToTransformAndReturnRow(row)})})
newDf.write.mode(SaveMode.Overwrite).parquet(path)
})
My use case is, for each item, read the data from S3, transform it (I am adding new columns directly to each row for our use case), and write the final result, for each item, back to S3.
Note - I can read the whole data, repartition by items, transform and write it back, but repartition results in a shuffle, which I am trying to avoid, and the way I am trying is, reading data for each item in an executor itself, so that it can work on whatever data it gets, and there is no need for shuffle.

Answer 1

I am not sure what you are trying to achieve using the approach that you have shown, but I feel that you might be going about it the hard way. Unless there is a good reason to do so, it is often best to let Spark (especially spark 2.0+) to let it do it's own thing. In this case, just process all three partitions using a single operation. Spark will usually manage your dataset quite well. It may also automatically introduce optimisations that you did not think of, or optimisations that it can not do if you try to control the process too much. Having said that, if it doesn't manage the process well, then you can start arguing with it by trying to take more control and do things more manually. At least that is my experience so far.
For example, I once had a complex series of transformations that added more logic to each step/DataFrame. If I forced spark to evaluate each intermediate one (e.g. run a count or a show on the intermediate dataframes) I would eventually get to a point where it couldn't evaluate one DataFrame (i.e. it couldn't do the count) because of insufficient resources. However, if I ignored that and added in more transformations, Spark was able to push some optimisations to earlier steps (from the later steps). This meant that the subsequent DataFrames (and importantly my final DataFrame) could be evaluated correctly. This final evaluation was possible despite the fact that the intermediate DataFrame, that on it's own could not be evaluated, was still in the overall process.
Consider the following. I use CSV, but it will work the same for parquet.
Here is my input:
data
├── tag=A
│   └── data.csv
├── tag=B
│   └── data.csv
└── tag=C
└── data.csv
Here is an example of one of the data files (tag=A/data.csv)
id,name,amount
1,Fred,100
2,Jane,200
Here is a script that recognises the partitions within this structure (i.e. tag is one of the columns).
scala> val inDataDF = spark.read.option("header","true").option("inferSchema","true").csv("data")
inDataDF: org.apache.spark.sql.DataFrame = [id: int, name: string ... 2 more fields]
scala> inDataDF.show
+---+-------+------+---+
| id| name|amount|tag|
+---+-------+------+---+
| 31| Scott| 3100| C|
| 32|Barnaby| 3200| C|
| 20| Bill| 2000| B|
| 21| Julia| 2100| B|
| 1| Fred| 100| A|
| 2| Jane| 200| A|
+---+-------+------+---+
scala> inDataDF.printSchema
root
|-- id: integer (nullable = true)
|-- name: string (nullable = true)
|-- amount: integer (nullable = true)
|-- tag: string (nullable = true)
scala> inDataDF.write.partitionBy("tag").csv("outData")
scala>
Again I used csv rather than parquet so you can dispense with the options to read a header and infer the schema (parquet will do that automatically). But apart from that, it will work the same way.
The above produces the following directory structure:
outData/
├── _SUCCESS
├── tag=A
│   └── part-00002-9e13ec13-7c63-4cda-b5af-e2d69cb76278.c000.csv
├── tag=B
│   └── part-00001-9e13ec13-7c63-4cda-b5af-e2d69cb76278.c000.csv
└── tag=C
└── part-00000-9e13ec13-7c63-4cda-b5af-e2d69cb76278.c000.csv
If you wanted to manipulate the contenst, by all means add whatever map operation, join, filter or whatever else you need between the read and write.
For example, add 500 to the amount:
scala> val outDataDF = inDataDF.withColumn("amount", $"amount" + 500)
outDataDF: org.apache.spark.sql.DataFrame = [id: int, name: string ... 2 more fields]
scala> outDataDF.show(false)
+---+-------+------+---+
|id |name |amount|tag|
+---+-------+------+---+
|31 |Scott |3600 |C |
|32 |Barnaby|3700 |C |
|20 |Bill |2500 |B |
|21 |Julia |2600 |B |
|1 |Fred |600 |A |
|2 |Jane |700 |A |
+---+-------+------+---+
Then simply write outDataDF instead of inDataDF.