Performance improvement for UDFs - get column name of least value per row in pyspark - python-3.x

Question

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

Answer 1

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:
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.
Find min column using least.
Create the column cnames for storing the column names.
Create the column set, which contains all the values as a comma-separated string.
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.

Answer 2

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|
+----+----+----+----+