将新列添加到 Python 列表中的 PySpark DataFrame

Question

我有一个清单：

dates = [2017, 2018, 2018, 2018, 2019, 2019, 2019, 2020, 2020, 2020]

我尝试添加的 dataframe 长度相同（没有问题）。

我试过了：

df = df.withColumn("YEARS", dates)
Error: Column needs to be col

我也试过：

df = df.withColumn("YEARS", f.lit(dates))

但这也行不通。

我看到了这个问题： How to add a constant column in a Spark DataFrame?

但是对于这种情况没有任何用处。

更新：预期的结果是：

df_columns...   | dates_from_list
---------------------------------
original_df_data| 2017
original_df_data| 2018
original_df_data| 2018
original_df_data| 2018
original_df_data| 2019
original_df_data| 2019
original_df_data| 2019
original_df_data| 2020
original_df_data| 2020
original_df_data| 2020

Answer 1

您的错误来自您需要将Column object 传递给withColumn的事实。

这里有两种方法可以将日期添加为 Spark DataFrame上的新列（使用每个记录的顺序进行连接），具体取决于日期数据的大小。

1）如果你操作一个小数据集

实现它的一种简洁方法是将 UDF 应用于单调递增的 id：

from pyspark.sql.functions import udf, monotonically_increasing_id

df = [...]  # 10 records

dates = [2017, 2018, 2018, 2018, 2019, 2019, 2019, 2020, 2020, 2020]

df = df.repartition(1).withColumn(
    "YEARS", 
    udf(lambda id: dates[id])(monotonically_increasing_id()))

df.show()

输出：

+---+-----+
|...|YEARS|
+---+-----+
|...| 2017|
|...| 2018|
|...| 2018|
|...| 2018|
|...| 2019|
|...| 2019|
|...| 2019|
|...| 2020|
|...| 2020|
|...| 2020|
+---+-----+

注意： .repartition(1)确保生成的 id 是连续的。 如果您有另一种方法将 map 每个记录转换为dates值（如先前构建的 id 列），则可以避免重新分区到单个分区。 在这个用例中，正如我们预期的 Python 列表 object 很小，这意味着您的 DataFrame 也很小，所以这个重新分区不是什么大问题。

/!\ 如果 dataframe 和 python 列表太大，为什么它不会扩展：

• 需要对 dataframe 进行重新分区，从而导致昂贵的洗牌/交换
• .repartition(1)可能会导致生成一个非常大的分区，该分区的处理速度可能非常慢（因为它很大，而且如果它不适合执行 memory 它可能意味着许多额外的磁盘 I/O 会溢出 RDD块到磁盘），或使用OutOfMemoryError使作业崩溃。
• python 列表由 udf 捕获（通过 lambda 闭包），这意味着它将被广播到集群的每个执行程序

2) 如果您操作的数据集大小 > 百万行

这是另一种方法，可以通过使用 pandas 操作 id 和日期列来更好地处理数百万行，并避免对 Spark DataFrame进行任何重新分区。

可以这样做：

import pandas as pd
from pyspark.sql.functions import monotonically_increasing_id
from pyspark.sql.session import SparkSession

spark = SparkSession.builder.getOrCreate()

# some spark DataFrame of length N
df = [...]  

# generate monotically increasing ids (not consecutive) without repartitioning the Spark DataFrame.
df = df.withColumn("id", monotonically_increasing_id())

# get generated ids (not consecutive) as a mono-column pandas DataFrame
spark_df_ids = df.select("id").toPandas()

# some python list of length N
dates = [2017, 2018, 2018, 2018, 2019, ..., 2019, 2019, 2020, 2020, 2020]

# build pandas DataFrame from dates
dates_pandas_df = pd.DataFrame(dates, columns=["YEARS"])

# append the id column to the dates in pandas
dates_and_ids_pandas_df = dates_pandas_df.join(spark_df_ids)

# convert from pandas DataFrame to spark DataFrame
dates_and_ids_spark_df = spark.createDataFrame(dates_and_ids_pandas_df)

# Perform the final adding of the dates column to the Spark DataFrame with a join in Spark
df.join(dates_and_ids_spark_df, ["id"]).show()

重要提示： 使用 Apache 箭头可以更快地从和到 pandas 的转换

Answer 2

你可以试试这个：

dates = [2017, 2018, 2018, 2018, 2019, 2019, 2019, 2020, 2020, 2020]
df = spark.createDataFrame([Row(a=1)])
df = df.withColumn("YEARS",  array( [lit(x) for x in dates]  ))


df.show(truncate=False)
+---+------------------------------------------------------------+
|a  |YEARS                                                       |
+---+------------------------------------------------------------+
|1  |[2017, 2018, 2018, 2018, 2019, 2019, 2019, 2020, 2020, 2020]|
+---+------------------------------------------------------------+

df.select("a", explode("YEARS")).show()
+---+----+
|  a| col|
+---+----+
|  1|2017|
|  1|2018|
|  1|2018|
|  1|2018|
|  1|2019|
|  1|2019|
|  1|2019|
|  1|2020|
|  1|2020|
|  1|2020|
+---+----+