PySpark code that turns columns into rows

Question

I am new to PySpark. My goal is to translate this table:

|Key|Value1|Value2|Value3| ...
+---+------+------+------+
|  1|    v1|    v2|    v3|

into a table like this:

|Key|ColumnName|ColumnValue|
+---+----------+-----------+
|  1|    Value1|         v1|
|  1|    Value2|         v2|
|  1|    Value3|         v3|      

And here is my code for doing so:

from pyspark.sql import *

sample = spark.read.format("csv").options(header='true', delimiter = ',').load("/FileStore/tables/sample.csv")

class Closure:  
  def __init__(self, columnNames):
    self.columnNames = columnNames

  def flatMapFunction(self, columnValues): 
    result = []
    columnIndex = 0
    for columnValue in columnValues:
      if not columnIndex == 0:
        result.append(Row(Key = columnValues[0], ColumnName = self.columnNames[columnIndex], ColumnValue = columnValue))
      columnIndex = columnIndex + 1
    return result

closure = Closure(sample.columns)
sample.rdd.flatMap(closure.flatMapFunction).toDF().show()

I know it is weird that I have to explicitly created the closure, but it looks like it is required, without which the execution would fail with a weird exception.

Questions:

• Is it good code?
• Is it going to scale well? There will be a few hundred columns but millions of rows.
• Is there anything I should do to improve the code?

Answer 1

There is a function in the standard library to create closure for you: functools.partial. This mean you can focus on writting your function as naturally as possible and bother of binding parameters later on.

As for your function:

1. Don't explicitly increment a running index, use enumerate instead;
2. You can use zip to iterate over two iterables at the same time;
3. Prefer using a list-comprehension to using [] + for + append;
4. You can use next on an iterator to retrieve an element and advance it outside of a for loop;
5. Avoid wildcard imports, they clutter the namespace and may lead to name collisions.
6. Use an if __name__ == '__main__': guard for your top-level code.

Proposed improvements

from functools import partial

from pyspark.sql import spark, Row


def flatten_table(column_names, column_values):
    row = zip(column_names, column_values)
    _, key = next(row)  # Special casing retrieving the first column
    return [
        Row(Key=key, ColumnName=column, ColumnValue=value)
        for column, value in row
    ]


if __name__ == '__main__':
    sample = spark.read.format("csv").options(header='true', delimiter = ',').load("/FileStore/tables/sample.csv")
    sample.rdd.flatMap(partial(flatten_table, sample.columns)).toDF().show()