将pyspark Dataframe 转换为字典过滤并从列中收集值

qxsslcnc  于 2022-12-26  发布在  Spark
关注(0)|答案(1)|浏览(113)

我需要将此DataFrame转换为字典:
| 识别号|值-1|值-2|值-3|
| - ------| - ------| - ------| - ------|
| 1A|批准|零|零|
| 2B类|批准|批准|零|
| 3C语言|零|零|批准|
输出:

{'1A': [value-1], '2B': [value-1,value-2], '3C': [value-3]}

注意,我使用DataFrame第一列的值作为字典的键。

vsikbqxv

vsikbqxv1#

基于arrayarray_remove,您可以使用如下代码:

from pyspark.sql import functions as F

# Set column to be used to get keys of the dict 
# and columns to be used to compute the values of the dict
dict_key = df.columns[0]
entry_cols = df.columns[1:]

{
    r[dict_key]: r.dict_entry
     for r in (
        df
        .select(
            dict_key,
            F.array_remove(
                F.array(*[
                    F.when(F.col(c) == 'Approve', F.lit(c)).otherwise('NULL')
                    for c in entry_cols
                ]),
                'NULL',
            ).alias('dict_entry')
        )
        .collect()
    )
}

这就是结果:

{'1A': ['value-1'], '2B': ['value-1', 'value-2'], '3C': ['value-3']}

相关问题