12
\$\begingroup\$

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?
\$\endgroup\$
11
\$\begingroup\$

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()
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.