Passing column information to SparkSession within withColumn

user_5029295 0
Viewed 27

I have a simple df with 2 columns, as shown below,

+------------+---+
|file_name   |id |
+------------+---+
|file1.csv   |1  |
|file2.csv   |2  |
+------------+---+

root
 |-- file_name: string (nullable = true)
 |-- id: string (nullable = true)

I wish to add a 3rd column with the count() from each file specified in the file_name column These are large files so I wish to go for a Spark based approach for getting the count() from each file.

Assuming originalDF is the above df,

I have tried:

 dfWithCounts =  originalDF.withColumn("counts", lit(spark.read.csv(lit(col('file_name'))).count))

but this seems to be throwing error.

Column is not iterable

Is there way I can achieve this?

I'm using Spark 2.4.

pyspark apache-spark scala
1 Answers

You can't run a Spark job from within another Spark job. Assuming that the file list is not super huge you can collect originalDF to the driver and spawn individual jobs to count lines from there.

val dfWithCounts = originalDF.collect.map { r =>
  (r.getString(0), r.getInt(1), spark.read.csv(r.getString(0)).count)
}.toSeq.toDF("file_name", "id", "count")

Optionally you can use Scala parallel collections to run these jobs in parallel.

val dfWithCounts = originalDF.collect.par.map { r =>
  (r.getString(0), r.getInt(1), spark.read.csv(r.getString(0)).count)
}.toSeq.seq.toDF("file_name", "id", "count")
user_28900930
Related