PySpark: running the same operation on multiple columns in one go

user_2352319 0
Viewed 993

My DataFrame table contains rows such as

['row1', 'col_1', 'col_2', 'col_3', ..., 'col_N', 'alpha']

N (the number of columns except the first and the last ones) is relatively large.

Now, I need to create another DataFrame out of this by multiplying each of these columns named col_i by column alpha. Is there a smarter way than to do a manual multiplication per each of these columns, as in:

sc = SparkContext()
sqlc = SQLContext(sc)

sqlc.sql('SELECT col_1 * alpha, col_2 * alpha, ..., col_N * alpha FROM table')

So I'd like to know whether it's possible to do the same operation on each column without specifically writing it for each one.

pyspark dataframe select sql python
3 Answers

Old question... but I ran into the same issue and found this question while looking for the answer.

As @Katya Willard suspected, her answer is not the most efficient. The pyspark documentation notes about withColumn():

(...) calling it multiple times, for instance, via loops in order to add multiple columns can generate big plans which can cause performance issues and even StackOverflowException. To avoid this, use select() with the multiple columns at once. https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.sql.DataFrame.withColumn.html

This introduces the question of how to do the operation on columns with a select statement. You want to build the select statement in a loop, and only then send it to your dataframe in one go:

from pyspark.sql.functions import col
# We need the 'col' function to use string column names without the df.colname syntax

def col_transform_select_list(col_name_list_for_transform, full_col_list):
  '''Prepare SELECT column list where some columns have the value multiplied 
     by the value from 'alpha` into a renamed column.
     The other columns are returned untouched and the column order is preserved.
  '''
  selection_col_list = list()
  alias_pattern = "{}_*_alpha"
  for col_name in full_col_list:
    if col_name in col_name_list_for_transform:
      selection_col_list.append((col(col_name)*col('alpha')).alias(alias_pattern.format(col_name)))
    else:
      selection_col_list.append(col_name)
  return selection_col_list

# Now run the select query directly on your table:
cols_to_multiply = [colname for colname in table.columns if colname.startswith('col_')]
prepared_column_list = col_transform_select_list(cols_to_multiply, table.columns)
multiplied_table = table.select(prepared_column_list)
user_32146690

You can use the list comprehension:

df.select(*[(F.col(col) * F.col('alpha')).alias(col) for col in df.columns if col != 'alpha'])
user_89736200
Related