从 Python 列表向 PySpark DataFrame 添加新列

Question

我有一个清单：

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

我尝试将其添加到的数据帧长度相同（没有问题）。

我试过：

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

我也试过：

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

但这也行不通。

我看到了这个问题：如何在 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

你的错误来自于你需要传递一个事实，withColumn一个Column物体。

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

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是连续的。如果您有另一种方法将每条记录映射到一个值dates（如先前构建的 id 列），则可以避免重新分区到单个分区。在这个用例中，由于我们期望 Python 列表对象非常小，这意味着您的 DataFrame 也非常小，因此这种重新分区不是什么大问题。

/!\如果数据框和 python 列表太大，为什么它不会缩放：

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

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

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

可以这样做：

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

重要提示：使用 Apache Arrow可以更快地从熊猫到熊猫的转换