将多列与另一列进行比较时立即选择更小/更大的值

Question

我有可变数量的列，假设在这个例子中，我们有 4 列要比较 (textX) 与具有不同值的单个列 (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)

我想根据“id”的值对textX 列中的最小和更大的值进行操作。 例如，对于id value = 2500，我想对值 2000 和 3000 进行操作。对于值为 500 的“id”，它将是 null 和 1000。 我试图将这些作为附加列，例如获取较低的列值

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

但是得到一个错误：

col 应该是 Column

我猜这是因为有多个列符合条件，但无法在此处插入。

有没有人有任何解决方案可以根据第三列将操作转换为 2 个值，或者只是如何正确获得这些界限？实际上，textX 列的数量会有所不同。由于性能问题，我尽可能远离 Pandas 和 UDF。

Answer 1

您可以使用least和greatest获取相关列：

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

然后你就可以对col1和col2进行操作了。

Answer 2

这是一种对 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   |
+----+-----+-----+-----+-----+--------+-------+

Answer 3

这是使用array_max 和array_min 函数以及when 表达式的另一种方式：

• lowerBound = max thresh_cols 满足条件thresh_col < id
• upperBound = min thresh_cols 满足条件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|
#+----+-----+-----+-----+-----+----------+----------+