How to efficiently mean-normalize columns

user_13562335 0
Viewed 82

I have a large dataset (~4M rows by ~3K columns) and I'm currently mean-normalizing each column using the following code in Python/PySpark:

import pyspark.sql.functions as f

means_pd = df.select(*[f.mean(c).alias(c) for c in df.columns]).toPandas()
diffs = df
for c in df.columns:
  mean = means_pd.loc[0,c]
  diffs = diffs.withColumn(c, f.col(c) - f.lit(mean))

This is quite slow, particularly the step to loop over the columns. There must be a better way to do this since there are functions like MinMaxScalar that include a step like this but don't take forever. How can I speed this up?

pyspark apache-spark-sql apache-spark mean normalization
1 Answers

You can calculate the mean over a window and subtract the result.

df.select([(F.col(c) - F.mean(c).over(W.orderBy())).alias(c) for c in df.columns])

This way you avoid the loop (3k withColumn) and do it purely in Spark, without Pandas.

Test:

from pyspark.sql import functions as F, Window as W
df = spark.createDataFrame(
    [(1, 1, 0),
     (2, 3, 3),
     (3, 5, 6)],
    ['c1', 'c2', 'c3'])

df_diffs = df.select([(F.col(c) - F.mean(c).over(W.orderBy())).alias(c) for c in df.columns])

df_diffs.show()
# +----+----+----+
# |  c1|  c2|  c3|
# +----+----+----+
# |-1.0|-2.0|-3.0|
# | 0.0| 0.0| 0.0|
# | 1.0| 2.0| 3.0|
# +----+----+----+
user_27535010
Related