Select immediate smaller/larger values when comparing multiple columns to another single column

user_4595093 0
Viewed 92

I have a variable number of columns, lets say for this example, we have 4 columns to be compared (textX) with a single column with different values (id):

d =     [
  {'id':  500, 'text1': 1000 ,'text2': 2000 ,'text3': 3000, 'text4': 5000}, 
  {'id': 1500, 'text1': 1000 ,'text2': 2000 ,'text3': 3000, 'text4': 5000},
  {'id': 2500, 'text1': 1000 ,'text2': 2000 ,'text3': 3000, 'text4': 5000},
  {'id': 3500, 'text1': 1000 ,'text2': 2000 ,'text3': 3000, 'text4': 5000},
  {'id': 4500, 'text1': 1000 ,'text2': 2000 ,'text3': 3000, 'text4': 5000},
  {'id': 5500, 'text1': 1000 ,'text2': 2000 ,'text3': 3000, 'text4': 5000}
] 
data = spark.createDataFrame(d)

I want to operate on both the immediate smallest and larger values from the textX columns depending on the value of 'id'. For example, for id value = 2500, I would like to operate on the values 2000 and 3000. In the case of 'id' of value 500, it would be a null and 1000. I have tried to get these as additional columns for example to get the lower column value

df_cols = data.columns
thresh_list = [x for x in df_cols if x.startswith('text')]

data.withColumn('inic_th', (col(x) for x in thresh_list if col('id') > col(x)))

But get an error :

col should be Column

I guess this is because there are more than one columns that match the condition but cannot get to insert that here.

Does anyone have any solution to either get an operation into 2 values depending on a third column or just how to properly get these bounds? In reality, the number of textX columns will vary. Due to performance issues, I am steering away from Pandas and UDFs as possible.

pyspark apache-spark-sql apache-spark python
3 Answers

You can use least and greatest to get the relevant columns:

import pyspark.sql.functions as F

df = data.withColumn(
    'col1',
    F.greatest(*[
        F.when(F.col(c) < F.col('id'), F.col(c))
        for c in data.columns
    ])
).withColumn(
    'col2',
    F.least(*[
        F.when(F.col(c) > F.col('id'), F.col(c))
        for c in data.columns
    ])
)

df.show()
+----+-----+-----+-----+-----+----+----+
|  id|text1|text2|text3|text4|col1|col2|
+----+-----+-----+-----+-----+----+----+
| 500| 1000| 2000| 3000| 5000|null|1000|
|1500| 1000| 2000| 3000| 5000|1000|2000|
|2500| 1000| 2000| 3000| 5000|2000|3000|
|3500| 1000| 2000| 3000| 5000|3000|5000|
|4500| 1000| 2000| 3000| 5000|3000|5000|
|5500| 1000| 2000| 3000| 5000|5000|null|
+----+-----+-----+-----+-----+----+----+

And then you can operate on col1 and col2.

user_141657300

Here is a way using higher order functions for spark >=2.4:

df_cols = data.columns
thresh_list = [x for x in df_cols if x.startswith('text')]

out = (data.select("*",F.sort_array(F.array(*thresh_list)).alias("Arr"))
.withColumn("FirstVal",F.expr('element_at(filter (Arr, x-> x<id),-1)'))
.withColumn("LastVal",F.expr('filter (Arr, x->x>id)[0]')).drop("Arr")
)

out.show(truncate=False)

+----+-----+-----+-----+-----+--------+-------+
|id  |text1|text2|text3|text4|FirstVal|LastVal|
+----+-----+-----+-----+-----+--------+-------+
|500 |1000 |2000 |3000 |5000 |null    |1000   |
|1500|1000 |2000 |3000 |5000 |1000    |2000   |
|2500|1000 |2000 |3000 |5000 |2000    |3000   |
|3500|1000 |2000 |3000 |5000 |3000    |5000   |
|4500|1000 |2000 |3000 |5000 |3000    |5000   |
|5500|1000 |2000 |3000 |5000 |5000    |null   |
+----+-----+-----+-----+-----+--------+-------+
user_98406370

Here's another way using array_max and array_min functions along with when expressions :

  • lowerBound = max thresh_cols that statisfy condition thresh_col < id
  • upperBound = min thresh_cols that statisfy condition thresh_col > id
from pyspark.sql import functions as F

result = data.withColumn(
    'lowerBound',
    F.array_max(F.array(*[F.when(F.col(c) < F.col('id'), F.col(c)) for c in thresh_cols]))
).withColumn(
    'upperBound',
    F.array_min(F.array(*[F.when(F.col(c) > F.col('id'), F.col(c)) for c in thresh_cols]))
)

result.show()

#+----+-----+-----+-----+-----+----------+----------+
#|  id|text1|text2|text3|text4|lowerBound|upperBound|
#+----+-----+-----+-----+-----+----------+----------+
#| 500| 1000| 2000| 3000| 5000|      null|      1000|
#|1500| 1000| 2000| 3000| 5000|      1000|      2000|
#|2500| 1000| 2000| 3000| 5000|      2000|      3000|
#|3500| 1000| 2000| 3000| 5000|      3000|      5000|
#|4500| 1000| 2000| 3000| 5000|      3000|      5000|
#|5500| 1000| 2000| 3000| 5000|      5000|      null|
#+----+-----+-----+-----+-----+----------+----------+
user_13865510
Related