python - Pyspark dataframe 获取至少一行满足条件的列列表

Question

我有一个 PySpark DataFrame

Col1 Col2 Col3
0.1  0.2  0.3

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

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

[Co2 , Co3]

我无法提供任何代码，因为我真的不知道如何做到这一点。

Answer 1

只需计数个满足谓词的项目(内部选择)并处理结果:

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]
  

• 收集第一个行:

  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]