使用pyspark将array< string>类型导出到csv而不逐个指定它们?

sqougxex  于 2021-07-13  发布在  Spark
关注(0)|答案(2)|浏览(350)

我有一个有很多列的数据框。其中一些列属于 array<string> .
我需要导出一个样本到csv和csv不支持数组。现在我对每个数组列都这样做(有时会丢失一个或多个)

df_write = df\
.withColumn('col_a', F.concat_ws(',', 'col_a'))\
.withColumn('col_g', F.concat_ws(',', 'col_g'))\
....

有没有一种方法可以使用循环并对每个数组列执行此操作,而不逐个指定它们?

j2qf4p5b

j2qf4p5b1#

您可以检查每一列的类型并进行列表理解:

import pyspark.sql.functions as F
from pyspark.sql.types import ArrayType

arr_col = [
    i.name
    for i in df.schema
    if isinstance(i.dataType, ArrayType)
]

df_write = df.select([
    F.concat_ws(',', c)
    if c in arr_col
    else F.col(c)
    for c in df.columns
])

实际上,你不需要使用 concat_ws . 您可以在写入csv之前将所有列强制转换为字符串类型,例如。

df_write = df.select([F.col(c).cast('string') for c in df.columns])
ifsvaxew

ifsvaxew2#

您还可以使用 df.dtypes :

from pyspark.sql import functions as F

array_cols = [c for c, t in df.dtypes if t == "array<string>"]

df.select(*[
    F.array_join(c, ",").alias(c) if c in array_cols else F.col(c)
    for c in df.columns
])

相关问题