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

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 should be 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进行操作了。