Hello all,

I've been tasked to convert a Scala Spark code to PySpark code with minimal changes (kinda literal translation).

I've come across some code that claims to be a list comprehension. Look below for code snippet:

%scala
val desiredColumn = Seq("firstName", "middleName", "lastName")
val colSize = desiredColumn.size
 
val columnList = for (i <- 0 until colSize) yield $"elements".getItem(i).alias(desiredColumn(i))
 
print(columnList)
 
// df_nameSplit.select(columnList: _ *).show(false)

Output for this code snippet:

Vector(elements[0] AS firstName, elements[1] AS middleName, elements[2] AS lastName)desiredColumn: Seq[String] = List(firstName, middleName, lastName)
colSize: Int = 3
columnList: scala.collection.immutable.IndexedSeq[org.apache.spark.sql.Column] = Vector(elements[0] AS firstName, elements[1] AS middleName, elements[2] AS lastName)

Also, the schema of the `df_nameSplit` data frame is as below and the elements column is a split version of the `name` column:

root
 |-- name: string (nullable = true)
 |-- dob_year: string (nullable = true)
 |-- gender: string (nullable = true)
 |-- salary: long (nullable = true)
 |-- elements: array (nullable = true)
 |    |-- element: string (containsNull = false)

The PySpark version of the code I was able to come-up with:

desired_columns = ["firstName", "middleName", "lastName"]
 
col_size = len(desired_columns)
 
col_list = [df_nameSplit.select(col("elements").getItem(i).alias(desired_columns[i])) for i in range(col_size)]
 
print(col_list)
 
# df_nameSplit.select(*col_list).display()

Output for PySpark code:

[DataFrame[firstName: string], DataFrame[middleName: string], DataFrame[lastName: string]]

Could someone help me with where I'm going wrong?

Tagging @Kaniz Fatma​ for better reach!