Pyspark show the column that has the lowest value for each row

Viewed 1036

I have the following dataframe

enter image description here

df_old_list= [
  { "Col1":"0", "Col2" : "7","Col3": "8", "Col4" : "","Col5": "20"},
{"Col1":"5", "Col2" : "5","Col3": "5", "Col4" : "","Col5": "28"},
 { "Col1":"-1", "Col2" : "-1","Col3": "13", "Col4" : "","Col5": "83"},

 {"Col1":"-1", "Col2" : "6","Col3": "6", "Col4" : "","Col5": "18"},

 { "Col1":"5", "Col2" : "4","Col3": "2", "Col4" : "","Col5": "84"},

  { "Col1":"0", "Col2" : "0","Col3": "14", "Col4" : "7","Col5": "86"}
]

spark = SparkSession.builder.getOrCreate()
df_old_list = spark.createDataFrame(Row(**x) for x in df_old_list)
df_old_list.show()

+----+----+----+----+----+
|Col1|Col2|Col3|Col4|Col5|
+----+----+----+----+----+
|   0|   7|   8|    |  20|
|   5|   5|   5|    |  28|
|  -1|  -1|  13|    |  83|
|  -1|   6|   6|    |  18|
|   5|   4|   2|    |  84|
|   0|   0|  14|   7|  86|
+----+----+----+----+----+

I want to get the lowest value from each column for each row.

This is what i was able to achieve so far

df1=df_old_list.selectExpr("*","array_sort(split(concat_ws(',',*),','))[0] lowest_col")

df1.show()

+----+----+----+----+----+----------+
|Col1|Col2|Col3|Col4|Col5|lowest_col|
+----+----+----+----+----+----------+
|   0|   7|   8|    |  20|          |
|   5|   5|   5|    |  28|          |
|  -1|  -1|  13|    |  83|          |
|  -1|   6|   6|    |  18|          |
|   5|   4|   2|    |  84|          |
|   0|   0|  14|   7|  86|         0|
+----+----+----+----+----+----------+

The problem is Col4 is blank and therefore its not able to compute the lowest value. What i am looking for is something like this : where i get the lowest value irrespective of blank columns and also if there are more than one lowest numbers then that column fields names should be shown in the lowest_col_title as concatenated.

+-----------------+----------+----+----+----+----+----+
|lowest_cols_title|lowest_col|Col1|Col2|Col3|Col4|Col5|
+-----------------+----------+----+----+----+----+----+
|             Col1|         0|   0|   7|   8|    |  20|
|   Col1;Col2;Col3|         5|   5|   5|   5|    |  28|
|        Col1;Col2|        -1|  -1|  -1|  13|    |  83|
|             Col1|        -1|  -1|   6|   6|    |  18|
|             Col3|         5|   5|   4|   2|    |  84|
|        Col1;Col2|         0|   0|   0|  14|   7|  86|
+-----------------+----------+----+----+----+----+----+
3 Answers

You can use pyspark.sql.functions.least

Returns the least value of the list of column names, skipping null values. This function takes at least 2 parameters. It will return null iff all parameters are null.

Once we have the minimum column we can compare the min value against all columns and create another column.

Create DataFrame:

from pyspark.sql import Row
from pyspark.sql.functions import col,least,when,array,concat_ws
df_old_list= [
  { "Col1":"0", "Col2" : "7","Col3": "8", "Col4" : "","Col5": "20"},  {"Col1":"5", "Col2" : "5","Col3": "5", "Col4" : "","Col5": "28"},
  { "Col1":"-1", "Col2" : "-1","Col3": "13", "Col4" : "","Col5": "83"},  {"Col1":"-1", "Col2" : "6","Col3": "6", "Col4" : "","Col5": "18"},
  { "Col1":"5", "Col2" : "4","Col3": "2", "Col4" : "","Col5": "84"},  { "Col1":"0", "Col2" : "0","Col3": "14", "Col4" : "7","Col5": "86"}]
df = spark.createDataFrame(Row(**x) for x in df_old_list)
from pyspark.sql.functions import least, when

Calculate row wise minimum and all columns which have the minimum value.

collist = df.columns
min_ = least(*[
    when(col(c) == "", float("inf")).otherwise(col(c).cast('int'))
    for c in df.columns
]).alias("lowest_col")

df = df.select("*", min_)
df = df.select("*",concat_ws(";",array([
       when(col(c)==col("lowest_col") ,c).otherwise(None) 
       for c in collist
     ])).alias("lowest_cols_title") )

df.show(10,False)

Output:

+----+----+----+----+----+----------+-----------------+
|Col1|Col2|Col3|Col4|Col5|lowest_col|lowest_cols_title|
+----+----+----+----+----+----------+-----------------+
|0   |7   |8   |    |20  |0.0       |Col1             |
|5   |5   |5   |    |28  |5.0       |Col1;Col2;Col3   |
|-1  |-1  |13  |    |83  |-1.0      |Col1;Col2        |
|-1  |6   |6   |    |18  |-1.0      |Col1             |
|5   |4   |2   |    |84  |2.0       |Col3             |
|0   |0   |14  |7   |86  |0.0       |Col1;Col2        |
+----+----+----+----+----+----------+-----------------+

There are several ways to avoid your empty cols.

  1. The problem is that your columns have string type:
df_old_list.dtypes
Out[13]: 
[('Col1', 'string'),
 ('Col2', 'string'),
 ('Col3', 'string'),
 ('Col4', 'string'),
 ('Col5', 'string')]

So you can just cast them to int:

df_old_list = df_old_list.withColumn('Col1', F.col('Col1').cast('int'))
df_old_list = df_old_list.withColumn('Col2', F.col('Col2').cast('int'))
df_old_list = df_old_list.withColumn('Col3', F.col('Col3').cast('int'))
df_old_list = df_old_list.withColumn('Col4', F.col('Col4').cast('int'))
df_old_list = df_old_list.withColumn('Col5', F.col('Col5').cast('int'))
df_old_list.show()
+----+----+----+----+----+
|Col1|Col2|Col3|Col4|Col5|
+----+----+----+----+----+
|   0|   7|   8|null|  20|
|   5|   5|   5|null|  28|
|  -1|  -1|  13|null|  83|
|  -1|   6|   6|null|  18|
|   5|   4|   2|null|  84|
|   0|   0|  14|   7|  86|
+----+----+----+----+----+

Now your code results in following:

df1=df_old_list.selectExpr("*","array_sort(split(concat_ws(',',*),','))[0] lowest_col")
   ...: 
   ...: df1.show()
+----+----+----+----+----+----------+
|Col1|Col2|Col3|Col4|Col5|lowest_col|
+----+----+----+----+----+----------+
|   0|   7|   8|null|  20|         0|
|   5|   5|   5|null|  28|        28|
|  -1|  -1|  13|null|  83|        -1|
|  -1|   6|   6|null|  18|        -1|
|   5|   4|   2|null|  84|         2|
|   0|   0|  14|   7|  86|         0|
+----+----+----+----+----+----------+
  1. You can replace empty values with some big int, for example
df_old_list = df_old_list.replace('', str((1 << 31) - 1))
df_old_list.show()
+----+----+----+----------+----+
|Col1|Col2|Col3|      Col4|Col5|
+----+----+----+----------+----+
|   0|   7|   8|2147483647|  20|
|   5|   5|   5|2147483647|  28|
|  -1|  -1|  13|2147483647|  83|
|  -1|   6|   6|2147483647|  18|
|   5|   4|   2|2147483647|  84|
|   0|   0|  14|         7|  86|
+----+----+----+----------+----+

And then do your manipulation:

df1=df_old_list.selectExpr("*","array_sort(split(concat_ws(',',*),','))[0] lowest_col")

df1.show()
+----+----+----+----------+----+----------+
|Col1|Col2|Col3|      Col4|Col5|lowest_col|
+----+----+----+----------+----+----------+
|   0|   7|   8|2147483647|  20|         0|
|   5|   5|   5|2147483647|  28|2147483647|
|  -1|  -1|  13|2147483647|  83|        -1|
|  -1|   6|   6|2147483647|  18|        -1|
|   5|   4|   2|2147483647|  84|         2|
|   0|   0|  14|         7|  86|         0|
+----+----+----+----------+----+----------+

You can notice that your way of finding min value comes to false result in the second row. So I can also recommend you to try to make this manipulation with transposing your dataframe and finding min value in column rather than in row. This can be done with pandas_udf very fast if you use arrow.

Using @venky__ answer, I found a solution for you:

from pyspark.sql import functions as F

join_key = df_old_list.columns

min_ = F.least(
    *[F.when(F.col(c).isNull() | (F.col(c) == ""), float("inf")).otherwise(F.col(c).cast('int')) 
      for c in join_key]
).alias("lowest_col")

df_with_lowest_col = df_old_list.select("*", min_.cast('int'))

df_exploded = df_old_list.withColumn(
    'vars_and_vals',
    F.explode(F.array(
        *(F.struct(F.lit(c).alias('var'), F.col(c).alias('val')) for c in join_key)
    )))

cols = join_key + [F.col("vars_and_vals")[x].alias(x) for x in ['var', 'val']]

df_exploded = df_exploded.select(*cols)

df = df_exploded.join(df_with_lowest_col, join_key)

df = df.filter('val = lowest_col')

df_with_col_names = df.groupby(*join_key).agg(
    F.array_join(F.collect_list('var'), ';').alias('lowest_cols_title')
)

res_df = df_with_lowest_col.join(df_with_col_names, join_key)

result:

res_df.show()

+----+----+----+----+----+----------+-----------------+
|Col1|Col2|Col3|Col4|Col5|lowest_col|lowest_cols_title|
+----+----+----+----+----+----------+-----------------+
|   0|   0|  14|   7|  86|         0|        Col1;Col2|
|  -1|   6|   6|    |  18|        -1|             Col1|
|   5|   5|   5|    |  28|         5|   Col1;Col2;Col3|
|   0|   7|   8|    |  20|         0|             Col1|
|  -1|  -1|  13|    |  83|        -1|        Col1;Col2|
|   5|   4|   2|    |  84|         2|             Col3|
+----+----+----+----+----+----------+-----------------+

It looks complicated and can be optimized I think but it works.

Related