Performance improvement for UDFs - get column name of least value per row in pyspark

user_6403976 0
Viewed 231

I use this udf:

mincol = F.udf(lambda row: cols[row.index(min(row))], StringType())
df = df.withColumn("mycol", mincol(F.struct([df[x] for x in cols])))

to get the column name for least value per row as value for another column called 'mycol'.

But this code is very slow. Any suggestions to improve performance? I am using Pyspark 2.3

pyspark-dataframes pyspark apache-spark-sql python-3.x
2 Answers

Here is another solution for Spark 2.3 which uses only built-in functions:

from sys import float_info
from pyspark.sql.functions import array, least, col, lit, concat_ws, expr

cols = df.columns
col_names = array(list(map(lit, cols)))
set_cols = list(map(col, cols))

# replace null with largest python float
df.na.fill(float_info.max) \
  .withColumn("min", least(*cols)) \
  .withColumn("cnames", col_names) \
  .withColumn("set", concat_ws(",", *set_cols)) \
  .withColumn("min_col", expr("cnames[find_in_set(min, set) - 1]")) \
  .select(*[cols + ["min_col"]]) \
  .show()

Steps:

  1. Fill all nulls with the larger possible float number. This is a good candidate for null replacement since is hard to find a larger value.
  2. Find min column using least.
  3. Create the column cnames for storing the column names.
  4. Create the column set, which contains all the values as a comma-separated string.
  5. Create the column min_col using find_in_set. The function handles each string item separately and will return the index of the found item. Finally, we use the index with cnames[indx - 1] to retrieve the column name.
user_7503760

Here is an approach without udf. The idea is to create an array containing the value and name of each column and then sort this array.

df1 = spark.createDataFrame([
    (1., 2., 3.),(3.,2.,1.), (9.,8.,-1.), (1.2, 1.2, 9.1), (3., None, 1.0)], \
        ["col1", "col2", "col3"])

cols = df1.columns
col_string = ', '.join("'{0}'".format(c) for c in cols)

df1 = df1.withColumn("vals", F.array(cols)) \
    .withColumn("cols", F.expr("Array(" + col_string + ")")) \
    .withColumn("zipped", F.arrays_zip("vals", "cols")) \
    .withColumn("without_nulls", F.expr("filter(zipped, x -> not x.vals is null)")) \
    .withColumn("sorted", F.expr("array_sort(without_nulls)")) \
    .withColumn("min", F.col("sorted")[0].cols) \
    .drop("vals", "cols", "zipped", "without_nulls", "sorted")
df1.show(truncate=False)

prints

+----+----+----+----+                                                           
|col1|col2|col3|min |
+----+----+----+----+
|1.0 |2.0 |3.0 |col1|
|3.0 |2.0 |1.0 |col3|
|9.0 |8.0 |-1.0|col3|
|1.2 |1.2 |9.1 |col1|
|3.0 |null|1.0 |col3|
+----+----+----+----+
user_21298010
Related