Pyspark数据框获取至少一行满足条件的列列表

Question

我有一个 PySpark DataFrame

Col1 Col2 Col3
0.1  0.2  0.3

我想获取至少一行满足条件的列名，例如一行大于 0.1

我的预期结果应该是这种情况：

[Co2 , Co3]

我无法提供任何代码，因为我真的不知道该怎么做。

Answer 1

只需count满足谓词（内部select）的项目并处理结果：

from pyspark.sql.functions import col, count, when

[c for c, v in df.select([
    count(when(col(c) > 0.1, 1)).alias(c) for c in df.columns
]).first().asDict().items() if v]

一步一步：

• 聚合（DataFrame -> DatFrame）：

  df = sc.parallelize([(0.1, 0.2, 0.3)]).toDF()
  counts = df.select([
      count(when(col(c) > 0.1, 1)).alias(c) for c in df.columns
  ])
  

  DataFrame[_1: bigint, _2: bigint, _3: bigint]
  

• collectfirstRow:

  a_row = counts.first()
  

  Row(_1=0, _2=1, _3=1)
  

• 转换成 Python dict:

  a_dict = a_row.asDict()
  

  {'_1': 0, '_2': 1, '_3': 1}
  

• 当值是真的时，迭代它的项目，保持键：

  [c for c, v in a_dict.items() if v]
  

  或显式检查计数：

  [c for c, v in a_dict.items() if v > 0]