pandas 如何为pyspark Dataframe 中具有多个虚拟变量的每个元素获取虚拟变量

uqjltbpv  于 2022-12-16  发布在  Spark
关注(0)|答案(1)|浏览(156)

下面是我的原始 Dataframe :

+---+---+---+
| ID| P1| P2|
+---+---+---+
|  0|447| O1|
|  0|448| O2|
|  1|447| O2|
|  1|450| O3|
|  2|450| O3|
|  3|451| O4|
|  3|452| O5|
+---+---+---+

我想要的是这样的 Dataframe :

+---+------+------+------+------+------+-----+-----+-----+-----+-----+
| ID|P1_447|P1_448|P1_450|P1_451|P1_452|P2_O1|P2_O2|P2_O3|P2_O4|P2_O5|
+---+------+------+------+------+------+-----+-----+-----+-----+-----+
|  0|     1|     1|     0|     0|     0|    1|    1|    0|    0|    0|
|  1|     1|     0|     1|     0|     0|    0|    1|    1|    0|    0|
|  2|     0|     0|     1|     0|     0|    0|    0|    1|    0|    0|
|  3|     0|     0|     0|     1|     1|    0|    0|    0|    1|    1|
+---+------+------+------+------+------+-----+-----+-----+-----+-----+

我试过了

df.groupby('ID').any().astype(int)

但没有成功。
谢谢你!

nimxete2

nimxete21#

import pyspark.sql.functions as F

df = spark.createDataFrame(
    [('0','447','O1')
    ,('0','448','O2')
    ,('1','447','O2')
    ,('1','450','O3')
    ,('2','450','O3')
    ,('3','451','O4')
    ,('3','452','O5')
    ],
    ['ID','P1','P2']
)

for c in df.columns:
    if not c == 'ID':       
        df = df\
                .withColumn(c, F.concat(F.lit(c),F.lit('_'),F.col(c)))\
                .groupBy(*[x for x in df.columns if x != c]).pivot(c).agg(F.lit(1))
df = df\
        .groupBy('ID').agg(*[F.sum(x).alias(x) for x in df.columns if x!='ID'])\
        .fillna(0)\
        .orderBy('ID')

df.show()

# +---+------+------+------+------+------+-----+-----+-----+-----+-----+
# | ID|P1_447|P1_448|P1_450|P1_451|P1_452|P2_O1|P2_O2|P2_O3|P2_O4|P2_O5|
# +---+------+------+------+------+------+-----+-----+-----+-----+-----+
# |  0|     1|     1|     0|     0|     0|    1|    1|    0|    0|    0|
# |  1|     1|     0|     1|     0|     0|    0|    1|    1|    0|    0|
# |  2|     0|     0|     1|     0|     0|    0|    0|    1|    0|    0|
# |  3|     0|     0|     0|     1|     1|    0|    0|    0|    1|    1|
# +---+------+------+------+------+------+-----+-----+-----+-----+-----+

相关问题