使用 PySpark 将数组 <string> 类型导出到 csv 而不一一指定它们？

Question

我有一个包含很多列的 DataFrame。其中一些列的类型为 array<string>.
我需要将示例导出到 csv，而 csv 不支持数组。现在我正在为每个数组列执行此操作（有时会错过一个或多个）

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

有没有一种方法可以使用循环并为每个数组列执行此操作而无需一一指定它们？

Answer 1

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

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])

Answer 2

您还可以使用 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
])

使用 PySpark 将数组 <string> 类型导出到 csv 而不一一指定它们？

Export array<string> type to csv using PySpark without specifying them one by one?

arrays

export-to-csv

apache-spark

apache-spark-sql

pyspark